plumb Reduce through

pola-rs · May 30, 2024 · 001628f · 001628f
1 parent 3a1bf85
commit 001628f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/py-polars/src/lazyframe/visitor/nodes.rs b/py-polars/src/lazyframe/visitor/nodes.rs
@@ -202,6 +202,15 @@ pub struct HStack {
     options: (), // ProjectionOptions,
 }
 
+#[pyclass]
+/// Like Select, but all operations produce a single row.
+pub struct Reduce {
+    #[pyo3(get)]
+    input: usize,
+    #[pyo3(get)]
+    exprs: Vec<PyExprIR>,
+}
+
 #[pyclass]
 /// Remove duplicates from the table
 pub struct Distinct {
@@ -436,11 +445,9 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
             input,
             exprs,
             schema: _,
-        } => Select {
+        } => Reduce {
             input: input.0,
-            expr: exprs.iter().map(|e| e.into()).collect(),
-            cse_expr: vec![],
-            options: (),
+            exprs: exprs.iter().map(|e| e.into()).collect(),
         }
         .into_py(py),
         IR::Distinct { input, options } => Distinct {

diff --git a/py-polars/src/lib.rs b/py-polars/src/lib.rs
@@ -75,6 +75,7 @@ fn _ir_nodes(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_class::<GroupBy>().unwrap();
     m.add_class::<Join>().unwrap();
     m.add_class::<HStack>().unwrap();
+    m.add_class::<Reduce>().unwrap();
     m.add_class::<Distinct>().unwrap();
     m.add_class::<MapFunction>().unwrap();
     m.add_class::<Union>().unwrap();