wip

Author: jayshrivastava

datafusion/physical-plan/src/aggregates/mod.rs

@@ -46,7 +46,8 @@ use arrow_schema::FieldRef;
 use datafusion_common::stats::Precision;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common::{
-    Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err, not_impl_err,
+    Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err,
+    internal_err, not_impl_err,
 };
 use datafusion_execution::TaskContext;
 use datafusion_expr::{Accumulator, Aggregate};
@@ -1052,26 +1053,39 @@ impl AggregateExec {
         self.dynamic_filter.as_ref().map(|df| &df.filter)
     }
 
-    /// Replace the dynamic filter expression, recomputing any internal state
-    /// which may depend on the previous dynamic filter.
-    ///
-    /// This is a no-op if the aggregate does not support dynamic filtering.
-    ///
-    /// If dynamic filtering is supported, this method returns an error if the filter's
-    /// children reference invalid columns in the aggregate's input schema.
+    /// Replace the dynamic filter expression. This method errors if the aggregate does not
+    /// support dynamic filtering or if the filter expression is incompatible with this
+    /// [`AggregateExec`].
     pub fn with_dynamic_filter(
         mut self,
         filter: Arc<DynamicFilterPhysicalExpr>,
     ) -> Result<Self> {
-        if let Some(supported_accumulators_info) = self.supported_accumulators_info() {
-            for child in filter.children() {
-                child.data_type(&self.input_schema)?;
+        // If there is no dynamic filter state initialized via `try_new`, then
+        // we can safely assume that the aggregate does not support dynamic filtering.
+        let Some(dyn_filter) = self.dynamic_filter.as_ref() else {
+            return internal_err!("Aggregate does not support dynamic filtering");
+        };
+
+        // Validate that the filter is compatible with the aggregation columns.
+        let cols = self.cols_for_dynamic_filter(&dyn_filter.supported_accumulators_info);
+        if cols.len() != filter.children().len() {
+            return internal_err!(
+                "Dynamic filter expression is incompatible with aggregate due to mismatched number of columns"
+            );
+        }
+        for (col, child) in cols.iter().zip(filter.children()) {
+            if !col.eq(child) {
+                return internal_err!(
+                    "Dynamic filter expression is incompatible with aggregate due to mismatched column references {col} != {child}"
+                );
             }
-            self.dynamic_filter = Some(Arc::new(AggrDynFilter {
-                filter,
-                supported_accumulators_info,
-            }));
         }
+
+        // Overwrite our filter
+        self.dynamic_filter = Some(Arc::new(AggrDynFilter {
+            filter,
+            supported_accumulators_info: dyn_filter.supported_accumulators_info.clone(),
+        }));
         Ok(self)
     }
 
@@ -1257,40 +1271,27 @@ impl AggregateExec {
     /// - If yes, init one inside `AggregateExec`'s `dynamic_filter` field.
     /// - If not supported, `self.dynamic_filter` should be kept `None`
     fn init_dynamic_filter(&mut self) {
-        // Already initialized.
-        if self.dynamic_filter.is_some() {
+        if (!self.group_by.is_empty()) || (self.mode != AggregateMode::Partial) {
+            debug_assert!(
+                self.dynamic_filter.is_none(),
+                "The current operator node does not support dynamic filter"
+            );
             return;
         }
 
-        if let Some(supported_accumulators_info) = self.supported_accumulators_info() {
-            // Collect column references for the dynamic filter expression.
-            let all_cols: Vec<Arc<dyn PhysicalExpr>> = supported_accumulators_info
-                .iter()
-                .map(|info| Arc::clone(&self.aggr_expr[info.aggr_index].expressions()[0]))
-                .collect();
-
-            self.dynamic_filter = Some(Arc::new(AggrDynFilter {
-                filter: Arc::new(DynamicFilterPhysicalExpr::new(all_cols, lit(true))),
-                supported_accumulators_info,
-            }));
-        }
-    }
-
-    /// Returns the supported accumulator info if this aggregate supports
-    /// dynamic filtering, or `None` otherwise.
-    ///
-    /// Dynamic filtering requires:
-    /// - `Partial` aggregation mode with no group-by expressions
-    /// - All aggregate functions are `min` or `max` with a single column arg
-    fn supported_accumulators_info(&self) -> Option<Vec<PerAccumulatorDynFilter>> {
-        if !self.group_by.is_empty() || !matches!(self.mode, AggregateMode::Partial) {
-            return None;
+        // Already initialized.
+        if self.dynamic_filter.is_some() {
+            return;
         }
 
-        // Collect supported accumulators.
-        // It is assumed the order of aggregate expressions are not changed
-        // from `AggregateExec` to `AggregateStream`.
+        // Collect supported accumulators
+        // It is assumed the order of aggregate expressions are not changed from `AggregateExec`
+        // to `AggregateStream`
         let mut aggr_dyn_filters = Vec::new();
+        // All column references in the dynamic filter, used when initializing the dynamic
+        // filter, and it's used to decide if this dynamic filter is able to get push
+        // through certain node during optimization.
+        let mut all_cols: Vec<Arc<dyn PhysicalExpr>> = Vec::new();
         for (i, aggr_expr) in self.aggr_expr.iter().enumerate() {
             // 1. Only `min` or `max` aggregate function
             let fun_name = aggr_expr.fun().name();
@@ -1301,13 +1302,14 @@ impl AggregateExec {
             } else if fun_name.eq_ignore_ascii_case("max") {
                 DynamicFilterAggregateType::Max
             } else {
-                return None;
+                return;
             };
 
             // 2. arg should be only 1 column reference
             if let [arg] = aggr_expr.expressions().as_slice()
                 && arg.is::<Column>()
             {
+                all_cols.push(Arc::clone(arg));
                 aggr_dyn_filters.push(PerAccumulatorDynFilter {
                     aggr_type,
                     aggr_index: i,
@@ -1316,13 +1318,36 @@ impl AggregateExec {
             }
         }
 
-        if aggr_dyn_filters.is_empty() {
-            None
-        } else {
-            Some(aggr_dyn_filters)
+        if !aggr_dyn_filters.is_empty() {
+            self.dynamic_filter = Some(Arc::new(AggrDynFilter {
+                filter: Arc::new(DynamicFilterPhysicalExpr::new(all_cols, lit(true))),
+                supported_accumulators_info: aggr_dyn_filters,
+            }))
         }
     }
 
+    // Collect column references for the dynamic filter expression from the supported accumulators.
+    fn cols_for_dynamic_filter(
+        &self,
+        supported_accumulators_info: &[PerAccumulatorDynFilter],
+    ) -> Vec<Arc<dyn PhysicalExpr>> {
+        let all_cols: Vec<Arc<dyn PhysicalExpr>> = supported_accumulators_info
+            .iter()
+            .filter_map(|info| {
+                // This should always be true due to how the supported accumulators
+                // are constructed. See `init_dynamic_filter` for more details.
+                if let [arg] = &self.aggr_expr[info.aggr_index].expressions().as_slice()
+                    && arg.is::<Column>()
+                {
+                    return Some(Arc::clone(arg));
+                }
+                None
+            })
+            .collect();
+        debug_assert!(all_cols.len() == supported_accumulators_info.len());
+        all_cols
+    }
+
     /// Calculate scaled byte size based on row count ratio.
     /// Returns `Precision::Absent` if input statistics are insufficient.
     /// Returns `Precision::Inexact` with the scaled value otherwise.
@@ -4817,12 +4842,13 @@ mod tests {
         Ok(())
     }
 
+    /// Test that [`AggregateExec::with_dynamic_filter`] overrides the existing dynamic filter
     #[test]
     fn test_with_dynamic_filter() -> Result<()> {
         let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
         let child = Arc::new(EmptyExec::new(Arc::clone(&schema)));
 
-        // Partial min triggers init_dynamic_filter.
+        // Partial min aggregate supports dynamic filtering
         let agg = AggregateExec::try_new(
             AggregateMode::Partial,
             PhysicalGroupBy::new_single(vec![]),
@@ -4836,39 +4862,40 @@ mod tests {
             child,
             Arc::clone(&schema),
         )?;
-        let original_inner_id = agg
-            .dynamic_filter()
-            .expect("should have dynamic filter after init")
-            .expression_id()
-            .expect("DynamicFilterPhysicalExpr always has an expression_id");
 
+        // Assertion 1: A filter with the same children can override the existing
+        // dynamic filter.
         let new_df = Arc::new(DynamicFilterPhysicalExpr::new(
             vec![col("a", &schema)?],
-            lit(true),
+            lit(false),
         ));
         let agg = agg.with_dynamic_filter(Arc::clone(&new_df))?;
-        let restored = agg
+
+        // The aggregate's filter should now resolve to the new inner expression.
+        let swapped = agg
             .dynamic_filter()
-            .expect("should still have dynamic filter");
-        assert_eq!(
-            restored
-                .expression_id()
-                .expect("DynamicFilterPhysicalExpr always has an expression_id"),
-            new_df
-                .expression_id()
-                .expect("DynamicFilterPhysicalExpr always has an expression_id"),
-        );
-        assert_ne!(
-            restored
-                .expression_id()
-                .expect("DynamicFilterPhysicalExpr always has an expression_id"),
-            original_inner_id,
-        );
+            .expect("should still have dynamic filter")
+            .current()?;
+        assert_eq!(format!("{swapped}"), format!("{}", lit(false)));
+
+        // Assertion 2: A filter that has been through `PhysicalExpr::with_new_children`
+        // should still be accepted when the new children are equivalent to the originals.
+        let new_df_as_pexpr: Arc<dyn PhysicalExpr> = new_df.clone();
+        let remapped_pexpr =
+            new_df_as_pexpr.with_new_children(vec![col("a", &schema)?])?;
+        let remapped_df = (remapped_pexpr as Arc<dyn std::any::Any + Send + Sync>)
+            .downcast::<DynamicFilterPhysicalExpr>()
+            .ok()
+            .expect("should be DynamicFilterPhysicalExpr after with_new_children");
+        // Hard to assert this because the filter is identical. No error means
+        // the filter was accepted. That's a good enough assertion for now.
+        let _agg = agg.with_dynamic_filter(remapped_df)?;
         Ok(())
     }
 
+    /// Test that [`AggregateExec::with_dynamic_filter`] errors when the aggregate does not support dynamic filtering
     #[test]
-    fn test_with_dynamic_filter_noop_when_unsupported() -> Result<()> {
+    fn test_with_dynamic_filter_error_unsupported() -> Result<()> {
         let schema = Arc::new(Schema::new(vec![
             Field::new("a", DataType::Int64, false),
             Field::new("b", DataType::Int64, false),
@@ -4891,18 +4918,17 @@ mod tests {
         )?;
         assert!(agg.dynamic_filter().is_none());
 
-        // with_dynamic_filter should be a no-op.
         let df = Arc::new(DynamicFilterPhysicalExpr::new(
             vec![col("a", &schema)?],
             lit(true),
         ));
-        let agg = agg.with_dynamic_filter(df)?;
-        assert!(agg.dynamic_filter().is_none());
+        assert!(agg.with_dynamic_filter(df).is_err());
         Ok(())
     }
 
+    /// Test that [`AggregateExec::with_dynamic_filter`] errors when the column is not in the schema
     #[test]
-    fn test_with_dynamic_filter_rejects_invalid_columns() -> Result<()> {
+    fn test_with_dynamic_filter_error_column_mismatch() -> Result<()> {
         let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
         let child = Arc::new(EmptyExec::new(Arc::clone(&schema)));
 
@@ -4920,7 +4946,6 @@ mod tests {
             Arc::clone(&schema),
         )?;
 
-        // Column index 99 is out of bounds for the input schema.
         let df = Arc::new(DynamicFilterPhysicalExpr::new(
             vec![Arc::new(Column::new("bad", 99)) as _],
             lit(true),

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -910,7 +910,7 @@ impl HashJoinExec {
 
     /// Set the dynamic filter on this hash join.
     ///
-    /// Resets any internal state that depends on any previous dynamic filter.
+    /// Resets any internal state that depends on any existing dynamic filter.
     ///
     /// Validates that the filter's children reference valid columns in
     /// the probe (right) side's schema.