fix window function regression

src/functions.rs

@@ -22,6 +22,7 @@ use datafusion_expr::expr::{Sort, WindowFunction};
 use datafusion_expr::window_function::find_df_window_func;
 use datafusion_expr::{lit, BuiltinScalarFunction, WindowFrame};
 
+use crate::errors::DataFusionError;
 use crate::expression::PyExpr;
 
 #[pyfunction]
@@ -81,34 +82,37 @@ fn alias(expr: PyExpr, name: &str) -> PyResult<PyExpr> {
 }
 
 /// Creates a new Window function expression
-// #[pyfunction]
-// fn window(
-//     name: &str,
-//     args: Vec<PyExpr>,
-//     partition_by: Option<Vec<PyExpr>>,
-//     order_by: Option<Vec<PyExpr>>,
-//     window_frame: PyWindowFrame,
-// ) -> PyResult<PyExpr> {
-//     let fun = find_df_window_func(name).unwrap();
-//         //.ok_or_else(errors::DataFusionError { e.into() })?;
-//     Ok(PyExpr {
-//         expr: datafusion_expr::Expr::WindowFunction(WindowFunction {
-//             fun,
-//             args: args.into_iter().map(|x| x.expr).collect::<Vec<_>>(),
-//             partition_by: partition_by
-//                 .unwrap_or_default()
-//                 .into_iter()
-//                 .map(|x| x.expr)
-//                 .collect::<Vec<_>>(),
-//             order_by: order_by
-//                 .unwrap_or_default()
-//                 .into_iter()
-//                 .map(|x| x.expr)
-//                 .collect::<Vec<_>>(),
-//             window_frame,
-//         }),
-//     })
-// }
+#[pyfunction]
+fn window(
+    name: &str,
+    args: Vec<PyExpr>,
+    partition_by: Option<Vec<PyExpr>>,
+    order_by: Option<Vec<PyExpr>>,
+) -> PyResult<PyExpr> {
+    let fun = find_df_window_func(name);
+    if fun.is_none() {
+        return Err(DataFusionError::Common("window function not found".to_string()).into());
+    }
+    let fun = fun.unwrap();
+    let window_frame = WindowFrame::new(order_by.is_some());
+    Ok(PyExpr {
+        expr: datafusion_expr::Expr::WindowFunction(WindowFunction {
+            fun,
+            args: args.into_iter().map(|x| x.expr).collect::<Vec<_>>(),
+            partition_by: partition_by
+                .unwrap_or_default()
+                .into_iter()
+                .map(|x| x.expr)
+                .collect::<Vec<_>>(),
+            order_by: order_by
+                .unwrap_or_default()
+                .into_iter()
+                .map(|x| x.expr)
+                .collect::<Vec<_>>(),
+            window_frame,
+        }),
+    })
+}
 
 macro_rules! scalar_function {
     ($NAME: ident, $FUNC: ident) => {
@@ -360,6 +364,6 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(trunc))?;
     m.add_wrapped(wrap_pyfunction!(upper))?;
     //m.add_wrapped(wrap_pyfunction!(uuid))?;
-    // m.add_wrapped(wrap_pyfunction!(window))?;
+    m.add_wrapped(wrap_pyfunction!(window))?;
     Ok(())
 }