diff --git a/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs index d53522f5de7c8..02408fdae17ec 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,20 @@ 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<()> { + self.update(lit(true))?; + let mut current = self.inner.write(); + current.is_complete = false; + let generation = current.generation; + drop(current); + let _ = self + .state_watch + .send(FilterState::InProgress { generation }); + Ok(()) + } + /// Mark this dynamic filter as complete and broadcast to all waiters. /// /// This signals that all expected updates have been received. @@ -1032,6 +1047,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..b4706a4af872a 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, 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;