feat: Support FILTER clause in aggregate window functions

WIP

Author: geoffreyclaude

datafusion-examples/examples/advanced_udwf.rs

@@ -200,6 +200,7 @@ impl WindowUDFImpl for SimplifySmoothItUdf {
                     window_frame: window_function.params.window_frame,
                     null_treatment: window_function.params.null_treatment,
                     distinct: window_function.params.distinct,
+                    filter: window_function.params.filter,
                 },
             }))
         };

datafusion/common/src/tree_node.rs

@@ -990,6 +990,48 @@ impl<
     }
 }
 
+impl<
+        'a,
+        T: 'a,
+        C0: TreeNodeContainer<'a, T>,
+        C1: TreeNodeContainer<'a, T>,
+        C2: TreeNodeContainer<'a, T>,
+        C3: TreeNodeContainer<'a, T>,
+    > TreeNodeContainer<'a, T> for (C0, C1, C2, C3)
+{
+    fn apply_elements<F: FnMut(&'a T) -> Result<TreeNodeRecursion>>(
+        &'a self,
+        mut f: F,
+    ) -> Result<TreeNodeRecursion> {
+        self.0
+            .apply_elements(&mut f)?
+            .visit_sibling(|| self.1.apply_elements(&mut f))?
+            .visit_sibling(|| self.2.apply_elements(&mut f))?
+            .visit_sibling(|| self.3.apply_elements(&mut f))
+    }
+
+    fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
+        self,
+        mut f: F,
+    ) -> Result<Transformed<Self>> {
+        self.0
+            .map_elements(&mut f)?
+            .map_data(|new_c0| Ok((new_c0, self.1, self.2, self.3)))?
+            .transform_sibling(|(new_c0, c1, c2, c3)| {
+                c1.map_elements(&mut f)?
+                    .map_data(|new_c1| Ok((new_c0, new_c1, c2, c3)))
+            })?
+            .transform_sibling(|(new_c0, new_c1, c2, c3)| {
+                c2.map_elements(&mut f)?
+                    .map_data(|new_c2| Ok((new_c0, new_c1, new_c2, c3)))
+            })?
+            .transform_sibling(|(new_c0, new_c1, new_c2, c3)| {
+                c3.map_elements(&mut f)?
+                    .map_data(|new_c3| Ok((new_c0, new_c1, new_c2, new_c3)))
+            })
+    }
+}
+
 /// [`TreeNodeRefContainer`] contains references to elements that a function can be
 /// applied on. The elements of the container are siblings so the continuation rules are
 /// similar to [`TreeNodeRecursion::visit_sibling`].
@@ -1065,6 +1107,27 @@ impl<
     }
 }
 
+impl<
+        'a,
+        T: 'a,
+        C0: TreeNodeContainer<'a, T>,
+        C1: TreeNodeContainer<'a, T>,
+        C2: TreeNodeContainer<'a, T>,
+        C3: TreeNodeContainer<'a, T>,
+    > TreeNodeRefContainer<'a, T> for (&'a C0, &'a C1, &'a C2, &'a C3)
+{
+    fn apply_ref_elements<F: FnMut(&'a T) -> Result<TreeNodeRecursion>>(
+        &self,
+        mut f: F,
+    ) -> Result<TreeNodeRecursion> {
+        self.0
+            .apply_elements(&mut f)?
+            .visit_sibling(|| self.1.apply_elements(&mut f))?
+            .visit_sibling(|| self.2.apply_elements(&mut f))?
+            .visit_sibling(|| self.3.apply_elements(&mut f))
+    }
+}
+
 /// Transformation helper to process a sequence of iterable tree nodes that are siblings.
 pub trait TreeNodeIterator: Iterator {
     /// Apples `f` to each item in this iterator

datafusion/core/src/physical_planner.rs

@@ -1650,6 +1650,7 @@ pub fn create_window_expr_with_name(
                         window_frame,
                         null_treatment,
                         distinct,
+                        filter,
                     },
             } = window_fun.as_ref();
             let physical_args =
@@ -1669,6 +1670,11 @@ pub fn create_window_expr_with_name(
             let window_frame = Arc::new(window_frame.clone());
             let ignore_nulls = null_treatment.unwrap_or(NullTreatment::RespectNulls)
                 == NullTreatment::IgnoreNulls;
+            let physical_filter = filter
+                .as_ref()
+                .map(|f| create_physical_expr(f, logical_schema, execution_props))
+                .transpose()?;
+
             windows::create_window_expr(
                 fun,
                 name,
@@ -1679,6 +1685,7 @@ pub fn create_window_expr_with_name(
                 physical_schema,
                 ignore_nulls,
                 *distinct,
+                physical_filter,
             )
         }
         other => plan_err!("Invalid window expression '{other:?}'"),

datafusion/core/tests/fuzz_cases/window_fuzz.rs

@@ -289,6 +289,7 @@ async fn bounded_window_causal_non_causal() -> Result<()> {
                     &extended_schema,
                     false,
                     false,
+                    None,
                 )?;
                 let running_window_exec = Arc::new(BoundedWindowAggExec::try_new(
                     vec![window_expr],
@@ -662,6 +663,7 @@ async fn run_window_test(
             &extended_schema,
             false,
             false,
+            None,
         )?],
         exec1,
         false,
@@ -681,6 +683,7 @@ async fn run_window_test(
             &extended_schema,
             false,
             false,
+            None,
         )?],
         exec2,
         search_mode.clone(),

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -3686,6 +3686,7 @@ async fn test_window_partial_constant_and_set_monotonicity() -> Result<()> {
             input_schema.as_ref(),
             false,
             false,
+            None,
         )?;
         let window_exec = if window_expr.uses_bounded_memory() {
             Arc::new(BoundedWindowAggExec::try_new(

datafusion/core/tests/physical_optimizer/test_utils.rs

@@ -266,6 +266,7 @@ pub fn bounded_window_exec_with_partition(
         schema.as_ref(),
         false,
         false,
+        None,
     )
     .unwrap();
 

datafusion/core/tests/physical_optimizer/window_optimize.rs

@@ -49,6 +49,7 @@ mod test {
             &partition,
             &[],
             Arc::new(frame),
+            None,
         );
 
         let bounded_agg_exec = BoundedWindowAggExec::try_new(

datafusion/expr/src/expr.rs

@@ -1230,6 +1230,8 @@ pub struct WindowFunctionParams {
     pub order_by: Vec<Sort>,
     /// Window frame
     pub window_frame: WindowFrame,
+    /// Optional filter expression (FILTER (WHERE ...))
+    pub filter: Option<Box<Expr>>,
     /// Specifies how NULL value is treated: ignore or respect
     pub null_treatment: Option<NullTreatment>,
     /// Distinct flag
@@ -1247,6 +1249,7 @@ impl WindowFunction {
                 partition_by: Vec::default(),
                 order_by: Vec::default(),
                 window_frame: WindowFrame::new(None),
+                filter: None,
                 null_treatment: None,
                 distinct: false,
             },
@@ -2388,6 +2391,7 @@ impl NormalizeEq for Expr {
                             window_frame: self_window_frame,
                             partition_by: self_partition_by,
                             order_by: self_order_by,
+                            filter: self_filter,
                             null_treatment: self_null_treatment,
                             distinct: self_distinct,
                         },
@@ -2400,13 +2404,19 @@ impl NormalizeEq for Expr {
                             window_frame: other_window_frame,
                             partition_by: other_partition_by,
                             order_by: other_order_by,
+                            filter: other_filter,
                             null_treatment: other_null_treatment,
                             distinct: other_distinct,
                         },
                 } = other.as_ref();
 
                 self_fun.name() == other_fun.name()
                     && self_window_frame == other_window_frame
+                    && match (self_filter, other_filter) {
+                        (Some(a), Some(b)) => a.normalize_eq(b),
+                        (None, None) => true,
+                        _ => false,
+                    }
                     && self_null_treatment == other_null_treatment
                     && self_args.len() == other_args.len()
                     && self_args
@@ -2658,12 +2668,14 @@ impl HashNode for Expr {
                             partition_by: _,
                             order_by: _,
                             window_frame,
+                            filter,
                             null_treatment,
                             distinct,
                         },
                 } = window_fun.as_ref();
                 fun.hash(state);
                 window_frame.hash(state);
+                filter.hash(state);
                 null_treatment.hash(state);
                 distinct.hash(state);
             }
@@ -2967,6 +2979,7 @@ impl Display for SchemaDisplay<'_> {
                             partition_by,
                             order_by,
                             window_frame,
+                            filter,
                             null_treatment,
                             distinct,
                         } = params;
@@ -2993,6 +3006,10 @@ impl Display for SchemaDisplay<'_> {
                             write!(f, " {null_treatment}")?;
                         }
 
+                        if let Some(filter) = filter {
+                            write!(f, " FILTER (WHERE {filter})")?;
+                        }
+
                         if !partition_by.is_empty() {
                             write!(
                                 f,
@@ -3370,6 +3387,7 @@ impl Display for Expr {
                             partition_by,
                             order_by,
                             window_frame,
+                            filter,
                             null_treatment,
                             distinct,
                         } = params;
@@ -3380,6 +3398,10 @@ impl Display for Expr {
                             write!(f, "{nt}")?;
                         }
 
+                        if let Some(fe) = filter {
+                            write!(f, " FILTER (WHERE {fe})")?;
+                        }
+
                         if !partition_by.is_empty() {
                             write!(f, " PARTITION BY [{}]", expr_vec_fmt!(partition_by))?;
                         }

datafusion/expr/src/expr_fn.rs

@@ -19,7 +19,7 @@
 
 use crate::expr::{
     AggregateFunction, BinaryExpr, Cast, Exists, GroupingSet, InList, InSubquery,
-    Placeholder, TryCast, Unnest, WildcardOptions, WindowFunction, WindowFunctionParams,
+    Placeholder, TryCast, Unnest, WildcardOptions, WindowFunction,
 };
 use crate::function::{
     AccumulatorArgs, AccumulatorFactoryFunction, PartitionEvaluatorFactory,
@@ -832,23 +832,17 @@ impl ExprFuncBuilder {
                 udaf.params.null_treatment = null_treatment;
                 Expr::AggregateFunction(udaf)
             }
-            ExprFuncKind::Window(WindowFunction {
-                fun,
-                params: WindowFunctionParams { args, .. },
-            }) => {
+            ExprFuncKind::Window(mut udwf) => {
                 let has_order_by = order_by.as_ref().map(|o| !o.is_empty());
-                Expr::from(WindowFunction {
-                    fun,
-                    params: WindowFunctionParams {
-                        args,
-                        partition_by: partition_by.unwrap_or_default(),
-                        order_by: order_by.unwrap_or_default(),
-                        window_frame: window_frame
-                            .unwrap_or_else(|| WindowFrame::new(has_order_by)),
-                        null_treatment,
-                        distinct,
-                    },
-                })
+
+                udwf.params.partition_by = partition_by.unwrap_or_default();
+                udwf.params.order_by = order_by.unwrap_or_default();
+                udwf.params.window_frame =
+                    window_frame.unwrap_or_else(|| WindowFrame::new(has_order_by));
+                udwf.params.filter = filter.map(Box::new);
+                udwf.params.null_treatment = null_treatment;
+                udwf.params.distinct = distinct;
+                Expr::WindowFunction(Box::new(udwf))
             }
         };
 

datafusion/expr/src/logical_plan/plan.rs

@@ -2469,6 +2469,20 @@ impl Window {
             window_func_dependencies.extend(new_deps);
         }
 
+        // Validate that FILTER if present is only used with aggregate window functions
+        if let Some(e) = window_expr.iter().find(|e| {
+            matches!(
+                e,
+                Expr::WindowFunction(wf)
+                    if !matches!(wf.fun, WindowFunctionDefinition::AggregateUDF(_))
+                        && wf.params.filter.is_some()
+            )
+        }) {
+            return plan_err!(
+                "FILTER clause can only be used with aggregate window functions. Found in '{e}'"
+            );
+        }
+
         Self::try_new_with_schema(
             window_expr,
             input,