raw

Author: askalt

datafusion/datasource-parquet/src/source.rs

@@ -734,6 +734,13 @@ impl FileSource for ParquetSource {
         .with_updated_node(source))
     }
 
+    fn with_new_state(
+        &self,
+        _state: Arc<dyn Any + Send + Sync>,
+    ) -> Option<Arc<dyn FileSource>> {
+        todo!()
+    }
+
     /// Try to optimize the scan to produce data in the requested sort order.
     ///
     /// This method receives:

datafusion/datasource/src/file.rs

@@ -148,6 +148,13 @@ pub trait FileSource: Send + Sync {
         ))
     }
 
+    fn with_new_state(
+        &self,
+        _state: Arc<dyn Any + Send + Sync>,
+    ) -> Option<Arc<dyn FileSource>> {
+        None
+    }
+
     /// Try to create a new FileSource that can produce data in the specified sort order.
     ///
     /// This method attempts to optimize data retrieval to match the requested ordering.

datafusion/datasource/src/source.rs

@@ -74,8 +74,8 @@ use datafusion_physical_plan::filter_pushdown::{
 /// ```text
 ///                       ┌─────────────────────┐                              -----► execute path
 ///                       │                     │                              ┄┄┄┄┄► init path
-///                       │   DataSourceExec    │  
-///                       │                     │    
+///                       │   DataSourceExec    │
+///                       │                     │
 ///                       └───────▲─────────────┘
 ///                               ┊  │
 ///                               ┊  │
@@ -192,6 +192,13 @@ pub trait DataSource: Send + Sync + Debug {
         ))
     }
 
+    fn with_new_state(
+        &self,
+        _state: Arc<dyn Any + Send + Sync>,
+    ) -> Option<Arc<dyn DataSource>> {
+        None
+    }
+
     /// Try to create a new DataSource that produces data in the specified sort order.
     ///
     /// # Arguments

datafusion/physical-expr/src/expressions/dynamic_filters.rs

@@ -74,6 +74,17 @@ pub struct DynamicFilterPhysicalExpr {
     /// But this can have overhead in production, so it's only included in our tests.
     data_type: Arc<RwLock<Option<DataType>>>,
     nullable: Arc<RwLock<Option<bool>>>,
+    /// Used to:
+    ///
+    /// * Reset this filter state to the initial.
+    /// * Compare origins of two different filters.
+    ///
+    /// Typically filter is created by some node and then pushed into children during plan optimization.
+    /// Then, when the plan is re-executed (e.g., a recursive query) the filter must be re-set and a new
+    /// indepenent [`Self::inner`] must be created and pushed into child node. To do it, parent node keeps
+    /// [`Self::identity`] within a new filter instance to provide the ability for child node to establish
+    /// a correspondence between stored filter and its new pushed version.
+    origin: Arc<dyn PhysicalExpr>,
 }
 
 #[derive(Debug)]
@@ -173,10 +184,11 @@ impl DynamicFilterPhysicalExpr {
         Self {
             children,
             remapped_children: None, // Initially no remapped children
-            inner: Arc::new(RwLock::new(Inner::new(inner))),
+            inner: Arc::new(RwLock::new(Inner::new(Arc::clone(&inner)))),
             state_watch,
             data_type: Arc::new(RwLock::new(None)),
             nullable: Arc::new(RwLock::new(None)),
+            origin: inner,
         }
     }
 
@@ -324,6 +336,33 @@ impl DynamicFilterPhysicalExpr {
         Arc::strong_count(self) > 1 || Arc::strong_count(&self.inner) > 1
     }
 
+    /// Check if this filter and `other` has the same origin.
+    ///
+    /// Please, check [`Self::origin`] doc-comment for the details.
+    pub fn has_same_origin(&self, other: &DynamicFilterPhysicalExpr) -> bool {
+        Arc::ptr_eq(&self.origin, &other.origin)
+    }
+
+    /// Reset this filter to the initial state.
+    ///
+    /// # Origin
+    ///
+    /// Returned filter has the same origin, check [`Self::origin`] doc-comment
+    /// for the details.
+    ///
+    pub fn reset_state(&self) -> Arc<DynamicFilterPhysicalExpr> {
+        let (state_watch, _) = watch::channel(FilterState::InProgress { generation: 1 });
+        Arc::new(Self {
+            children: self.children.clone(),
+            remapped_children: self.remapped_children.clone(),
+            inner: Arc::new(RwLock::new(Inner::new(Arc::clone(&self.origin)))),
+            state_watch,
+            data_type: Arc::new(RwLock::new(None)),
+            nullable: Arc::new(RwLock::new(None)),
+            origin: Arc::clone(&self.origin),
+        })
+    }
+
     fn render(
         &self,
         f: &mut std::fmt::Formatter<'_>,
@@ -369,6 +408,7 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr {
             state_watch: self.state_watch.clone(),
             data_type: Arc::clone(&self.data_type),
             nullable: Arc::clone(&self.nullable),
+            origin: Arc::clone(&self.origin),
         }))
     }
 
@@ -864,4 +904,17 @@ mod test {
             "Hash should be stable after update (identity-based)"
         );
     }
+
+    #[test]
+    fn test_dynamic_filter_reset_state() {
+        let origin = lit(true) as Arc<dyn PhysicalExpr>;
+        let filter = DynamicFilterPhysicalExpr::new(vec![], Arc::clone(&origin));
+        filter.update(lit(false) as _).unwrap();
+        filter.mark_complete();
+        let reset_filter = filter.reset_state();
+        assert!(!Arc::ptr_eq(&filter.inner, &reset_filter.inner));
+        assert!(filter.has_same_origin(&reset_filter));
+        assert!(Arc::ptr_eq(&reset_filter.current().unwrap(), &origin));
+        assert!(!reset_filter.inner.read().is_complete);
+    }
 }

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

@@ -1104,6 +1104,15 @@ impl AggregateExec {
             _ => Precision::Absent,
         }
     }
+
+    fn reset_dynamic_filter(&self) -> Option<Arc<AggrDynFilter>> {
+        self.dynamic_filter.as_ref().map(|f| {
+            Arc::new(AggrDynFilter {
+                filter: f.filter.reset_state(),
+                supported_accumulators_info: f.supported_accumulators_info.clone(),
+            })
+        })
+    }
 }
 
 impl DisplayAs for AggregateExec {
@@ -1300,6 +1309,13 @@ impl ExecutionPlan for AggregateExec {
         Ok(Arc::new(me))
     }
 
+    fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> {
+        let mut me = AggregateExec::clone(&self);
+        me.dynamic_filter = self.reset_dynamic_filter();
+        me.metrics = ExecutionPlanMetricsSet::new();
+        Ok(Arc::new(me))
+    }
+
     fn execute(
         &self,
         partition: usize,
@@ -1466,6 +1482,14 @@ impl ExecutionPlan for AggregateExec {
 
         Ok(result)
     }
+
+    fn dynamic_filters(&self) -> Vec<Arc<DynamicFilterPhysicalExpr>> {
+        if let Some(expr) = self.dynamic_filter.as_ref() {
+            vec![Arc::clone(&expr.filter)]
+        } else {
+            vec![]
+        }
+    }
 }
 
 fn create_schema(

datafusion/physical-plan/src/execution_plan.rs

@@ -22,6 +22,7 @@ use crate::filter_pushdown::{
 };
 pub use crate::metrics::Metric;
 pub use crate::ordering::InputOrderMode;
+use crate::recursive_query::RecursiveQueryExec;
 use crate::sort_pushdown::SortOrderPushdownResult;
 pub use crate::stream::EmptyRecordBatchStream;
 
@@ -31,12 +32,15 @@ pub use datafusion_common::utils::project_schema;
 pub use datafusion_common::{ColumnStatistics, Statistics, internal_err};
 pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream};
 pub use datafusion_expr::{Accumulator, ColumnarValue};
+use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr;
 pub use datafusion_physical_expr::window::WindowExpr;
 pub use datafusion_physical_expr::{
     Distribution, Partitioning, PhysicalExpr, expressions,
 };
+use parking_lot::Mutex;
 
 use std::any::Any;
+use std::cell::RefCell;
 use std::fmt::Debug;
 use std::sync::Arc;
 
@@ -680,6 +684,12 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync {
     /// in order to wire up the working table used during recursive-CTE execution.
     /// Similar patterns can be followed by custom nodes that need late-bound
     /// dependencies or shared state.
+    ///
+    /// Also, this method is used to update dynamic filters in the plan when
+    /// its state is re-set. So if the node supports dynamic filtering it must
+    /// implement this method to be reusable. Please, check [`reset_plan_states`]
+    /// for the details.
+    ///
     fn with_new_state(
         &self,
         _state: Arc<dyn Any + Send + Sync>,
@@ -722,6 +732,11 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync {
     ) -> Option<Arc<dyn ExecutionPlan>> {
         None
     }
+
+    /// Returns dynamic filters owned by the plan.
+    fn dynamic_filters(&self) -> Vec<Arc<DynamicFilterPhysicalExpr>> {
+        vec![]
+    }
 }
 
 /// [`ExecutionPlan`] Invariant Level
@@ -1398,6 +1413,43 @@ pub fn check_not_null_constraints(
     Ok(batch)
 }
 
+/// Unions dynamic filters collected from plan nodes during state re-set bypass.
+/// Passed into [`ExecutionPlan::with_new_state`] to actualize node filters.
+#[derive(Default)]
+pub struct DynamicFilters {
+    filters: Mutex<Vec<Arc<DynamicFilterPhysicalExpr>>>,
+}
+
+impl DynamicFilters {
+    fn extend(&self, iter: impl Iterator<Item = Arc<DynamicFilterPhysicalExpr>>) {
+        self.filters.lock().extend(iter);
+    }
+
+    /// Lookup for the filter with the same origin as a passed filter `expr`.
+    fn filter_with_same_origin_as(
+        &self,
+        expr: &DynamicFilterPhysicalExpr,
+    ) -> Result<Arc<DynamicFilterPhysicalExpr>> {
+        for filter in self.filters.lock().iter() {
+            if filter.has_same_origin(expr) {
+                return Ok(Arc::clone(filter));
+            }
+        }
+        internal_err!("updated dynamic filter is not found")
+    }
+
+    /// Update each dynamic filter sub-expression to an actual version.
+    ///
+    /// * If dynamic filter with the same origin is not found, then an error is returned.
+    /// * If `expr` does not contain dynamic filters and there are no updates, then `Ok(None)` is returned.
+    pub fn actualize_dynamic_filter(
+        &self,
+        expr: &DynamicFilterPhysicalExpr,
+    ) -> Result<Option<Arc<DynamicFilterPhysicalExpr>>> {
+        todo!()
+    }
+}
+
 /// Make plan ready to be re-executed returning its clone with state reset for all nodes.
 ///
 /// Some plans will change their internal states after execution, making them unable to be executed again.
@@ -1407,18 +1459,61 @@ pub fn check_not_null_constraints(
 /// However, if the data of the left table is derived from the work table, it will become outdated
 /// as the work table changes. When the next iteration executes this plan again, we must clear the left table.
 ///
-/// # Limitations
+/// # Dynamic filters
+///
+/// Dynamic filters are re-set during state re-set bypass and must be updated in nodes which poll them.
+/// To be able to do it the method [`ExecutionPlan::with_new_state`] should be implemented for the node
+/// that supports dynamic filtering. The node should down-cast input state into [`DynamicFilters`] and
+/// lookup for the filter with the same origin as it has to acquire an actual filter version.
 ///
-/// While this function enables plan reuse, it does not allow the same plan to be executed if it (OR):
+/// # Work table
 ///
-/// * uses dynamic filters,
-/// * represents a recursive query.
+/// Work table is re-set during state re-set bypass and must be updated in nodes which write into it.
+/// To be able to do it the method [`ExecutionPlan::with_new_state`] should be implemented for the node
+/// that uses work table.
 ///
 pub fn reset_plan_states(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
-    plan.transform_up(|plan| {
-        let new_plan = Arc::clone(&plan).reset_state()?;
-        Ok(Transformed::yes(new_plan))
-    })
+    let dynamic_filters = Arc::new(DynamicFilters::default());
+    let work_tables = RefCell::new(vec![]);
+
+    plan.transform_down_up(
+        // Collect dynamic filters from nodes.
+        |plan| {
+            // Note: here we re-set state prior to children state is re-set to be able
+            // to collect and push new versions of the dynamic filters/work table.
+            // Children actually will be replacedo on the way up from the bottom.
+            let mut new_plan = Arc::clone(&plan).reset_state()?;
+
+            // Try to inject updated dynamic filters.
+            if let Some(plan) = new_plan.with_new_state(Arc::clone(&dynamic_filters) as _)
+            {
+                new_plan = plan;
+            }
+
+            // Try to inject updated work table.
+            if let Some(work_table) = work_tables.borrow_mut().last()
+                && let Some(plan) = new_plan.with_new_state(Arc::clone(work_table))
+            {
+                new_plan = plan;
+            }
+
+            dynamic_filters.extend(new_plan.dynamic_filters().into_iter());
+            if let Some(plan) = new_plan.as_any().downcast_ref::<RecursiveQueryExec>() {
+                work_tables
+                    .borrow_mut()
+                    .push(Arc::clone(plan.work_table()) as _)
+            }
+
+            Ok(Transformed::yes(new_plan))
+        },
+        |plan| {
+            if plan.as_any().is::<RecursiveQueryExec>() {
+                work_tables.borrow_mut().pop();
+            }
+            // Here we must return [`Transformed::yes`] to actually replace children.
+            Ok(Transformed::yes(plan))
+        },
+    )
     .data()
 }
 

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

@@ -847,6 +847,15 @@ impl HashJoinExec {
             reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema())
         }
     }
+
+    fn reset_dynamic_filter(&self) -> Option<HashJoinExecDynamicFilter> {
+        self.dynamic_filter
+            .as_ref()
+            .map(|f| HashJoinExecDynamicFilter {
+                filter: f.filter.reset_state(),
+                build_accumulator: OnceLock::new(),
+            })
+    }
 }
 
 impl DisplayAs for HashJoinExec {
@@ -1046,7 +1055,7 @@ impl ExecutionPlan for HashJoinExec {
             null_aware: self.null_aware,
             cache: self.cache.clone(),
             // Reset dynamic filter and bounds accumulator to initial state
-            dynamic_filter: None,
+            dynamic_filter: self.reset_dynamic_filter(),
         }))
     }
 
@@ -1384,6 +1393,14 @@ impl ExecutionPlan for HashJoinExec {
         }
         Ok(result)
     }
+
+    fn dynamic_filters(&self) -> Vec<Arc<DynamicFilterPhysicalExpr>> {
+        if let Some(expr) = self.dynamic_filter.as_ref() {
+            vec![Arc::clone(&expr.filter)]
+        } else {
+            vec![]
+        }
+    }
 }
 
 /// Accumulator for collecting min/max bounds from build-side data during hash join.