fix: handle WindowFunction in ExprFunctionExt filter and distinct on Expr

Author: duc12111

datafusion/core/tests/expr_api/mod.rs

@@ -24,7 +24,7 @@ use arrow::util::pretty::{pretty_format_batches, pretty_format_columns};
 use datafusion::prelude::*;
 use datafusion_common::{DFSchema, ScalarValue};
 use datafusion_expr::ExprFunctionExt;
-use datafusion_expr::expr::NullTreatment;
+use datafusion_expr::expr::{NullTreatment, WindowFunction, WindowFunctionDefinition};
 use datafusion_expr::simplify::SimplifyContext;
 use datafusion_functions::core::expr_ext::FieldAccessor;
 use datafusion_functions_aggregate::first_last::first_value_udaf;
@@ -263,6 +263,101 @@ async fn test_aggregate_ext_distinct() {
     .await;
 }
 
+/// Test that `filter()` works as the first chained call on a window function.
+/// This verifies the fix for https://github.com/apache/datafusion/issues/21697
+#[tokio::test]
+async fn test_window_ext_filter_first() {
+    // Build a window function with filter as the FIRST chained method.
+    // Before the fix, this would fail because filter() on Expr only handled
+    // AggregateFunction, not WindowFunction.
+    let window_expr = Expr::from(WindowFunction::new(
+        WindowFunctionDefinition::AggregateUDF(sum_udaf()),
+        vec![col("i")],
+    ))
+    .filter(col("i").is_not_null())
+    .order_by(vec![col("id").sort(true, true)])
+    .build()
+    .unwrap()
+    .alias("sum_filtered");
+
+    let ctx = SessionContext::new();
+    let result = ctx
+        .read_batch(TEST_BATCH.clone())
+        .unwrap()
+        .select(vec![col("id"), col("i"), window_expr])
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+
+    let result = pretty_format_batches(&result).unwrap().to_string();
+    let actual_lines = result.lines().collect::<Vec<_>>();
+
+    // TEST_BATCH: id=["1","2","3"], i=[10, NULL, 5]
+    // Ordered by id ASC: row1(id=1,i=10), row2(id=2,i=NULL), row3(id=3,i=5)
+    // FILTER(i IS NOT NULL) excludes row2 from the sum
+    // Running sum: row1=10, row2=10 (NULL filtered out), row3=15
+    let expected_lines = vec![
+        "+----+----+--------------+",
+        "| id | i  | sum_filtered |",
+        "+----+----+--------------+",
+        "| 1  | 10 | 10           |",
+        "| 2  |    | 10           |",
+        "| 3  | 5  | 15           |",
+        "+----+----+--------------+",
+    ];
+
+    assert_eq!(
+        expected_lines, actual_lines,
+        "\n\nexpected:\n\n{expected_lines:#?}\nactual:\n\n{actual_lines:#?}\n\n"
+    );
+}
+
+/// Test that `distinct()` works as the first chained call on a window function.
+#[tokio::test]
+async fn test_window_ext_distinct_first() {
+    let window_expr = Expr::from(WindowFunction::new(
+        WindowFunctionDefinition::AggregateUDF(sum_udaf()),
+        vec![col("i")],
+    ))
+    .distinct()
+    .order_by(vec![col("id").sort(true, true)])
+    .build()
+    .unwrap()
+    .alias("sum_distinct");
+
+    let ctx = SessionContext::new();
+    let result = ctx
+        .read_batch(TEST_BATCH.clone())
+        .unwrap()
+        .select(vec![col("id"), col("i"), window_expr])
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+
+    let result = pretty_format_batches(&result).unwrap().to_string();
+    let actual_lines = result.lines().collect::<Vec<_>>();
+
+    // TEST_BATCH: id=["1","2","3"], i=[10, NULL, 5]
+    // Ordered by id ASC: row1(id=1,i=10), row2(id=2,i=NULL), row3(id=3,i=5)
+    // DISTINCT running sum: row1=10, row2=10 (NULL skipped), row3=15
+    let expected_lines = vec![
+        "+----+----+--------------+",
+        "| id | i  | sum_distinct |",
+        "+----+----+--------------+",
+        "| 1  | 10 | 10           |",
+        "| 2  |    | 10           |",
+        "| 3  | 5  | 15           |",
+        "+----+----+--------------+",
+    ];
+
+    assert_eq!(
+        expected_lines, actual_lines,
+        "\n\nexpected:\n\n{expected_lines:#?}\nactual:\n\n{actual_lines:#?}\n\n"
+    );
+}
+
 #[tokio::test]
 async fn test_aggregate_ext_null_treatment() {
     let agg = first_value_udaf()

datafusion/expr/src/expr_fn.rs

@@ -932,26 +932,34 @@ impl ExprFunctionExt for Expr {
         builder
     }
     fn filter(self, filter: Expr) -> ExprFuncBuilder {
-        match self {
+        let mut builder = match self {
             Expr::AggregateFunction(udaf) => {
-                let mut builder =
-                    ExprFuncBuilder::new(Some(ExprFuncKind::Aggregate(udaf)));
-                builder.filter = Some(filter);
-                builder
+                ExprFuncBuilder::new(Some(ExprFuncKind::Aggregate(udaf)))
+            }
+            Expr::WindowFunction(udwf) => {
+                ExprFuncBuilder::new(Some(ExprFuncKind::Window(udwf)))
             }
             _ => ExprFuncBuilder::new(None),
+        };
+        if builder.fun.is_some() {
+            builder.filter = Some(filter);
         }
+        builder
     }
     fn distinct(self) -> ExprFuncBuilder {
-        match self {
+        let mut builder = match self {
             Expr::AggregateFunction(udaf) => {
-                let mut builder =
-                    ExprFuncBuilder::new(Some(ExprFuncKind::Aggregate(udaf)));
-                builder.distinct = true;
-                builder
+                ExprFuncBuilder::new(Some(ExprFuncKind::Aggregate(udaf)))
+            }
+            Expr::WindowFunction(udwf) => {
+                ExprFuncBuilder::new(Some(ExprFuncKind::Window(udwf)))
             }
             _ => ExprFuncBuilder::new(None),
+        };
+        if builder.fun.is_some() {
+            builder.distinct = true;
         }
+        builder
     }
     fn null_treatment(
         self,
@@ -998,6 +1006,10 @@ impl ExprFunctionExt for Expr {
 #[cfg(test)]
 mod test {
     use super::*;
+    use crate::WindowFunctionDefinition;
+    use crate::expr::{AggregateFunction, WindowFunction};
+    use crate::lit;
+    use crate::test::function_stub::sum_udaf;
 
     #[test]
     fn filter_is_null_and_is_not_null() {
@@ -1009,4 +1021,91 @@ mod test {
             "col2 IS NOT NULL"
         );
     }
+
+    /// Create a window function expression for testing
+    fn test_window_expr() -> Expr {
+        Expr::WindowFunction(Box::new(WindowFunction::new(
+            WindowFunctionDefinition::AggregateUDF(sum_udaf()),
+            vec![col("x")],
+        )))
+    }
+
+    #[test]
+    fn test_window_filter_first() {
+        let result = test_window_expr().filter(col("a").gt_eq(lit(5))).build();
+        assert!(
+            result.is_ok(),
+            "filter as first call on WindowFunction should succeed"
+        );
+        let expr = result.unwrap();
+        match expr {
+            Expr::WindowFunction(wf) => {
+                assert!(wf.params.filter.is_some(), "filter should be set");
+            }
+            other => panic!("expected WindowFunction, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn test_window_distinct_first() {
+        let result = test_window_expr().distinct().build();
+        assert!(
+            result.is_ok(),
+            "distinct as first call on WindowFunction should succeed"
+        );
+        let expr = result.unwrap();
+        match expr {
+            Expr::WindowFunction(wf) => {
+                assert!(wf.params.distinct, "distinct should be true");
+            }
+            other => panic!("expected WindowFunction, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn test_window_filter_then_partition_by() {
+        let result = test_window_expr()
+            .filter(col("a").gt_eq(lit(5)))
+            .partition_by(vec![col("y")])
+            .build();
+        assert!(
+            result.is_ok(),
+            "filter then partition_by on WindowFunction should succeed"
+        );
+        let expr = result.unwrap();
+        match expr {
+            Expr::WindowFunction(wf) => {
+                assert!(wf.params.filter.is_some(), "filter should be set");
+                assert_eq!(
+                    wf.params.partition_by.len(),
+                    1,
+                    "partition_by should have one entry"
+                );
+            }
+            other => panic!("expected WindowFunction, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn test_aggregate_filter_still_works() {
+        let agg = Expr::AggregateFunction(AggregateFunction::new_udf(
+            sum_udaf(),
+            vec![col("x")],
+            false,
+            None,
+            vec![],
+            None,
+        ));
+        let result = agg.filter(col("a").gt_eq(lit(5))).build();
+        assert!(
+            result.is_ok(),
+            "filter on AggregateFunction should still work"
+        );
+        match result.unwrap() {
+            Expr::AggregateFunction(af) => {
+                assert!(af.params.filter.is_some(), "filter should be set");
+            }
+            other => panic!("expected AggregateFunction, got {other:?}"),
+        }
+    }
 }