diff --git a/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs index d53522f5de7c8..1c384589ef421 100644 --- a/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs +++ b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs @@ -29,6 +29,7 @@ use datafusion_common::{ tree_node::{Transformed, TransformedResult, TreeNode}, }; +use crate::expressions::lit; use datafusion_expr::ColumnarValue; use datafusion_expr_common::dyn_eq::DynHash; @@ -343,6 +344,35 @@ impl DynamicFilterPhysicalExpr { Ok(()) } + /// Reset this dynamic filter so its producer can publish a new expression + /// during a later execution of the same plan. + pub fn reset(&self) -> Result<()> { + let new_expr = Self::remap_children( + &self.children, + self.remapped_children.as_ref(), + lit(true), + )?; + + let mut current = self.inner.write(); + // Keep the generation increasing: `current()` caches by generation, so going + // back would let a later update hit the previous run's cached expression. + let new_generation = current.generation + 1; + *current = Inner { + expression_id: current.expression_id, + generation: new_generation, + expr: new_expr, + is_complete: false, + }; + drop(current); + + // The channel still holds `Complete` from the previous run; publish + // `InProgress` so `wait_complete()` waits for this run's filter. + let _ = self.state_watch.send(FilterState::InProgress { + generation: new_generation, + }); + Ok(()) + } + /// Mark this dynamic filter as complete and broadcast to all waiters. /// /// This signals that all expected updates have been received. @@ -1032,6 +1062,21 @@ mod test { dynamic_filter.wait_complete().await; } + #[test] + fn test_reset_reopens_completed_filter() { + let dynamic_filter = + DynamicFilterPhysicalExpr::new(vec![], lit(false) as Arc); + dynamic_filter.mark_complete(); + let generation_before_reset = dynamic_filter.inner.read().generation; + + dynamic_filter.reset().unwrap(); + + let inner = dynamic_filter.inner.read(); + assert!(!inner.is_complete); + assert!(inner.generation > generation_before_reset); + assert_eq!(inner.expr.to_string(), "true"); + } + #[test] fn test_with_new_children_independence() { // Create a schema with columns a, b, c, d diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index e528cf467ad9c..faff968040965 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -2617,6 +2617,19 @@ impl ExecutionPlan for AggregateExec { ) } + fn reset_state(self: Arc) -> Result> { + if let Some(dynamic_filter) = &self.dynamic_filter { + dynamic_filter.filter.reset()?; + for accumulator_filter in &dynamic_filter.accumulator_dyn_filter_info { + *accumulator_filter.shared_bound.lock() = ScalarValue::Null; + } + } + + let mut aggregate = (*self).clone(); + aggregate.metrics = ExecutionPlanMetricsSet::new(); + Ok(Arc::new(aggregate)) + } + fn execute( &self, partition: usize, @@ -9814,4 +9827,58 @@ mod tests { assert!(agg.set_dynamic_filter(df).is_err()); Ok(()) } + + #[test] + fn test_aggregate_exec_reset_state() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let agg = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![]), + vec![Arc::new( + AggregateExprBuilder::new(min_udaf(), vec![col("a", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("min_a") + .build()?, + )], + vec![None], + child, + Arc::clone(&schema), + )?; + + let df = Arc::new(DynamicFilterPhysicalExpr::new( + vec![col("a", &schema)?], + lit(false), + )); + let agg = Arc::new(agg.set_dynamic_filter(Arc::clone(&df))?); + + // Simulate an iteration setting a shared bound and completing the filter + if let Some(dyn_filter) = &agg.dynamic_filter { + for info in &dyn_filter.accumulator_dyn_filter_info { + *info.shared_bound.lock() = ScalarValue::Int64(Some(42)); + } + dyn_filter.filter.mark_complete(); + } + + // Call reset_state + let reset_plan = agg.reset_state()?; + let reset_agg = reset_plan + .downcast_ref::() + .expect("should downcast back to AggregateExec"); + + let dyn_filter = reset_agg + .dynamic_filter + .as_ref() + .expect("dynamic filter should be retained"); + + for info in &dyn_filter.accumulator_dyn_filter_info { + assert_eq!(*info.shared_bound.lock(), ScalarValue::Null); + } + + let current_expr = dyn_filter.filter.current()?; + assert_eq!(format!("{current_expr}"), format!("{}", lit(true))); + + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/cte.slt b/datafusion/sqllogictest/test_files/cte.slt index 376fb82a27b3d..c470a23ed6f98 100644 --- a/datafusion/sqllogictest/test_files/cte.slt +++ b/datafusion/sqllogictest/test_files/cte.slt @@ -1727,3 +1727,34 @@ order by 1, 2; statement ok drop table test; + +# Dynamic filters produced by a recursive term must be reset between +# iterations. Otherwise the second execution of this uncorrelated aggregate +# reuses the first iteration's MAX bound and scans no rows. +statement ok +COPY (SELECT CAST(value AS DOUBLE) AS r FROM generate_series(0, 128)) TO 'test_files/scratch/cte/dynamic_filter_recursive.parquet' STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE dynamic_filter_recursive STORED AS PARQUET LOCATION 'test_files/scratch/cte/dynamic_filter_recursive.parquet'; + +query IR rowsort +WITH RECURSIVE w(n, m) AS ( + SELECT 0, CAST(NULL AS DOUBLE) + UNION ALL + SELECT w.n + 1, x.m + FROM w CROSS JOIN ( + SELECT MAX(r) AS m + FROM dynamic_filter_recursive + WHERE r % 2 = 0 + ) x + WHERE w.n < 3 +) +SELECT n, m FROM w; +---- +0 NULL +1 128 +2 128 +3 128 + +statement ok +DROP TABLE dynamic_filter_recursive;