add generation checks

Author: adriangb

datafusion/datasource-parquet/src/opener.rs

@@ -38,7 +38,7 @@ use datafusion_common::pruning::{
 use datafusion_common::{exec_err, DataFusionError, Result};
 use datafusion_datasource::PartitionedFile;
 use datafusion_physical_expr_common::physical_expr::{
-    is_dynamic_physical_expr, PhysicalExpr,
+    is_dynamic_physical_expr, snapshot_generation, PhysicalExpr,
 };
 use datafusion_physical_optimizer::pruning::PruningPredicate;
 use datafusion_physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder};
@@ -144,25 +144,23 @@ impl FileOpener for ParquetOpener {
             // We'll also check this after every record batch we read,
             // and if at some point we are able to prove we can prune the file using just the file level statistics
             // we can end the stream early.
-            let file_pruner = predicate
+            let mut file_pruner = predicate
                 .as_ref()
                 .map(|p| {
-                    Ok::<_, DataFusionError>(
-                        is_dynamic_physical_expr(Arc::clone(p))?.then_some(Arc::new(
-                            FilePruner::new(
-                                Arc::clone(p),
-                                &logical_file_schema,
-                                partition_fields.clone(),
-                                file.clone(),
-                                predicate_creation_errors.clone(),
-                            )?,
-                        )),
-                    )
+                    Ok::<_, DataFusionError>(is_dynamic_physical_expr(&p).then_some(
+                        FilePruner::new(
+                            Arc::clone(p),
+                            &logical_file_schema,
+                            partition_fields.clone(),
+                            file.clone(),
+                            predicate_creation_errors.clone(),
+                        )?,
+                    ))
                 })
                 .transpose()?
                 .flatten();
 
-            if let Some(file_pruner) = &file_pruner {
+            if let Some(file_pruner) = &mut file_pruner {
                 if file_pruner.should_prune()? {
                     // Return an empty stream immediately to skip the work of setting up the actual stream
                     file_metrics.files_pruned_statistics.add(1);
@@ -362,7 +360,7 @@ impl FileOpener for ParquetOpener {
                         .and_then(|b| schema_mapping.map_batch(b).map_err(Into::into))
                 })
                 .take_while(move |_| {
-                    if let Some(file_pruner) = file_pruner.as_ref() {
+                    if let Some(file_pruner) = file_pruner.as_mut() {
                         match file_pruner.should_prune() {
                             Ok(false) => futures::future::ready(true),
                             Ok(true) => {
@@ -520,6 +518,7 @@ fn should_enable_page_index(
 
 /// Prune based on partition values and file-level statistics.
 pub struct FilePruner {
+    predicate_generation: u64,
     predicate: Arc<dyn PhysicalExpr>,
     /// Schema used for pruning, which combines the file schema and partition fields.
     /// Partition fields are always at the end, as they are during scans.
@@ -551,6 +550,7 @@ impl FilePruner {
             .with_metadata(logical_file_schema.metadata().clone()),
         );
         Ok(Self {
+            predicate_generation: snapshot_generation(&predicate),
             predicate,
             pruning_schema,
             file,
@@ -559,7 +559,14 @@ impl FilePruner {
         })
     }
 
-    pub fn should_prune(&self) -> Result<bool> {
+    pub fn should_prune(&mut self) -> Result<bool> {
+        let current_generation = self.predicate_generation;
+        let new_generation = snapshot_generation(&self.predicate);
+        // If the predicate has not changed since the last time we checked, we can skip pruning
+        if current_generation == new_generation {
+            return Ok(false);
+        }
+        self.predicate_generation = new_generation;
         let pruning_predicate = build_pruning_predicate(
             Arc::clone(&self.predicate),
             &self.pruning_schema,

datafusion/physical-expr-common/src/physical_expr.rs

@@ -27,7 +27,9 @@ use arrow::array::BooleanArray;
 use arrow::compute::filter_record_batch;
 use arrow::datatypes::{DataType, Field, FieldRef, Schema};
 use arrow::record_batch::RecordBatch;
-use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::tree_node::{
+    Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
+};
 use datafusion_common::{internal_err, not_impl_err, Result, ScalarValue};
 use datafusion_expr_common::columnar_value::ColumnarValue;
 use datafusion_expr_common::interval_arithmetic::Interval;
@@ -345,6 +347,24 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + DynEq + DynHash {
         // This is a safe default behavior.
         Ok(None)
     }
+
+    /// Returns the generation of this `PhysicalExpr` for snapshotting purposes.
+    /// The generation is an arbitrary u64 that can be used to track changes
+    /// in the state of the `PhysicalExpr` over time without having to do an exhaustive comparison.
+    /// This is useful to avoid unecessary computation or serialization if there are no changes to the expression.
+    /// In particular, dynamic expressions that may change over time; this allows cheap checks for changes.
+    /// Static expressions that do not change over time should return 0, as does the default implementation.
+    /// You should not call this method directly as it does not handle recursion.
+    /// Instead use [`snapshot_generation`] to handle recursion and capture the
+    /// full state of the `PhysicalExpr`.
+    fn snapshot_generation(&self) -> u64 {
+        // By default, we return 0 to indicate that this PhysicalExpr does not
+        // have any dynamic references or state.
+        // Since the recursive algorithm XORs the generations of all children the overall
+        // generation will be 0 if no children have a non-zero generation, meaning that
+        // static expressions will always return 0.
+        0
+    }
 }
 
 /// [`PhysicalExpr`] can't be constrained by [`Eq`] directly because it must remain object
@@ -538,16 +558,30 @@ pub fn snapshot_physical_expr(
     .data()
 }
 
-/// Check if the given `PhysicalExpr` is dynamic.
-/// See the documentation of [`PhysicalExpr::snapshot`] for more details.
-pub fn is_dynamic_physical_expr(expr: Arc<dyn PhysicalExpr>) -> Result<bool> {
-    let mut is_dynamic = false;
-    expr.transform_up(|e| {
-        if e.snapshot()?.is_some() {
-            is_dynamic = true;
-        }
-        Ok(Transformed::no(e))
+/// Check the generation of this `PhysicalExpr`.
+/// Dynamic `PhysicalExpr`s may have a generation that is incremented
+/// every time the state of the `PhysicalExpr` changes.
+/// If the generation changes that means this `PhysicalExpr` or one of its children
+/// has changed since the last time it was evaluated.
+///
+/// This algorithm will not produce collisions as long as the structure of the
+/// `PhysicalExpr` does not change and no `PhysicalExpr` decrements its own generation.
+pub fn snapshot_generation(expr: &Arc<dyn PhysicalExpr>) -> u64 {
+    let mut generation = 0u64;
+    expr.apply(|e| {
+        // Add the current generation of the `PhysicalExpr` to our global generation.
+        generation = generation.wrapping_add(e.snapshot_generation());
+        Ok(TreeNodeRecursion::Continue)
     })
-    .data()?;
-    Ok(is_dynamic)
+    .expect("this traversal is infallible");
+
+    generation
+}
+
+/// Check if the given `PhysicalExpr` is dynamic.
+/// Internally this calls [`snapshot_generation`] to check if the generation is non-zero,
+/// any dynamic `PhysicalExpr` should have a non-zero generation.
+pub fn is_dynamic_physical_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
+    // If the generation is non-zero, then this `PhysicalExpr` is dynamic.
+    snapshot_generation(expr) != 0
 }

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

@@ -43,14 +43,33 @@ pub struct DynamicFilterPhysicalExpr {
     /// so that when we update `current()` in subsequent iterations we can re-apply the replacements.
     remapped_children: Option<Vec<Arc<dyn PhysicalExpr>>>,
     /// The source of dynamic filters.
-    inner: Arc<RwLock<Arc<dyn PhysicalExpr>>>,
+    inner: Arc<RwLock<Inner>>,
     /// For testing purposes track the data type and nullability to make sure they don't change.
     /// If they do, there's a bug in the implementation.
     /// 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>>>,
 }
 
+#[derive(Debug)]
+struct Inner {
+    /// A counter that gets incremented every time the expression is updated so that we can track changes cheaply.
+    /// This is used for [`PhysicalExpr::generation`] to have a cheap check for changes.
+    generation: u64,
+    expr: Arc<dyn PhysicalExpr>,
+}
+
+impl Inner {
+    fn new(expr: Arc<dyn PhysicalExpr>) -> Self {
+        Self {
+            // Start with generation 1 which gives us a different result for [`PhysicalExpr::generation`] than the default 0.
+            // This is not currently used anywhere but it seems useful to have this simple distinction.
+            generation: 1,
+            expr,
+        }
+    }
+}
+
 impl Hash for DynamicFilterPhysicalExpr {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         let inner = self.current().expect("Failed to get current expression");
@@ -111,7 +130,7 @@ impl DynamicFilterPhysicalExpr {
         Self {
             children,
             remapped_children: None, // Initially no remapped children
-            inner: Arc::new(RwLock::new(inner)),
+            inner: Arc::new(RwLock::new(Inner::new(inner))),
             data_type: Arc::new(RwLock::new(None)),
             nullable: Arc::new(RwLock::new(None)),
         }
@@ -158,6 +177,7 @@ impl DynamicFilterPhysicalExpr {
                     "Failed to acquire read lock for inner".to_string(),
                 )
             })?
+            .expr
             .clone();
         let inner =
             Self::remap_children(&self.children, self.remapped_children.as_ref(), inner)?;
@@ -186,7 +206,9 @@ impl DynamicFilterPhysicalExpr {
             self.remapped_children.as_ref(),
             new_expr,
         )?;
-        *current = new_expr;
+        current.expr = new_expr;
+        current.generation += 1; // Increment the generation to indicate that the expression has changed.
+                                 // Increment the generation to indicate that the expression has changed.
         Ok(())
     }
 }
@@ -291,6 +313,14 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr {
         // Return the current expression as a snapshot.
         Ok(Some(self.current()?))
     }
+
+    fn snapshot_generation(&self) -> u64 {
+        // Return the current generation of the expression.
+        self.inner
+            .read()
+            .expect("Failed to acquire read lock for inner")
+            .generation
+    }
 }
 
 #[cfg(test)]