feat: Support FILTER clause in aggregate window functions

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,13 @@ 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 = match filter {
+                Some(f) => {
+                    Some(create_physical_expr(f, logical_schema, execution_props)?)
+                }
+                None => None,
+            };
+
             windows::create_window_expr(
                 fun,
                 name,
@@ -1679,6 +1687,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,8 @@
 
 use crate::expr::{
     AggregateFunction, BinaryExpr, Cast, Exists, GroupingSet, InList, InSubquery,
-    Placeholder, TryCast, Unnest, WildcardOptions, WindowFunction, WindowFunctionParams,
+    Placeholder, TryCast, Unnest, WildcardOptions, WindowFunction,
+    WindowFunctionDefinition, WindowFunctionParams,
 };
 use crate::function::{
     AccumulatorArgs, AccumulatorFactoryFunction, PartitionEvaluatorFactory,
@@ -836,6 +837,14 @@ impl ExprFuncBuilder {
                 fun,
                 params: WindowFunctionParams { args, .. },
             }) => {
+                // FILTER is only supported for aggregate window functions
+                if filter.is_some()
+                    && matches!(fun, WindowFunctionDefinition::WindowUDF(_))
+                {
+                    return plan_err!(
+                        "FILTER clause is only permitted for aggregate window functions"
+                    );
+                }
                 let has_order_by = order_by.as_ref().map(|o| !o.is_empty());
                 Expr::from(WindowFunction {
                     fun,
@@ -845,6 +854,7 @@ impl ExprFuncBuilder {
                         order_by: order_by.unwrap_or_default(),
                         window_frame: window_frame
                             .unwrap_or_else(|| WindowFrame::new(has_order_by)),
+                        filter: filter.map(Box::new),
                         null_treatment,
                         distinct,
                     },

datafusion/expr/src/planner.rs

@@ -307,6 +307,7 @@ pub struct RawWindowExpr {
     pub partition_by: Vec<Expr>,
     pub order_by: Vec<SortExpr>,
     pub window_frame: WindowFrame,
+    pub filter: Option<Box<Expr>>,
     pub null_treatment: Option<NullTreatment>,
     pub distinct: bool,
 }