remove snapshot method from dynamic filter physical expr

Author: jayshrivastava

datafusion/ffi/src/physical_expr/mod.rs

@@ -112,8 +112,6 @@ pub struct FFI_PhysicalExpr {
 
     pub fmt_sql: unsafe extern "C" fn(&Self) -> FFI_Result<SString>,
 
-    pub snapshot: unsafe extern "C" fn(&Self) -> FFI_Result<FFI_Option<FFI_PhysicalExpr>>,
-
     pub snapshot_generation: unsafe extern "C" fn(&Self) -> u64,
 
     pub is_volatile_node: unsafe extern "C" fn(&Self) -> bool,
@@ -364,16 +362,6 @@ unsafe extern "C" fn fmt_sql_fn_wrapper(expr: &FFI_PhysicalExpr) -> FFI_Result<S
     FFI_Result::Ok(result.into())
 }
 
-unsafe extern "C" fn snapshot_fn_wrapper(
-    expr: &FFI_PhysicalExpr,
-) -> FFI_Result<FFI_Option<FFI_PhysicalExpr>> {
-    let expr = expr.inner();
-    sresult!(
-        expr.snapshot()
-            .map(|snapshot| snapshot.map(FFI_PhysicalExpr::from).into())
-    )
-}
-
 unsafe extern "C" fn snapshot_generation_fn_wrapper(expr: &FFI_PhysicalExpr) -> u64 {
     let expr = expr.inner();
     expr.snapshot_generation()
@@ -427,7 +415,6 @@ unsafe extern "C" fn clone_fn_wrapper(expr: &FFI_PhysicalExpr) -> FFI_PhysicalEx
             propagate_statistics: propagate_statistics_fn_wrapper,
             get_properties: get_properties_fn_wrapper,
             fmt_sql: fmt_sql_fn_wrapper,
-            snapshot: snapshot_fn_wrapper,
             snapshot_generation: snapshot_generation_fn_wrapper,
             is_volatile_node: is_volatile_node_fn_wrapper,
             display: display_fn_wrapper,
@@ -470,7 +457,6 @@ impl From<Arc<dyn PhysicalExpr>> for FFI_PhysicalExpr {
             propagate_statistics: propagate_statistics_fn_wrapper,
             get_properties: get_properties_fn_wrapper,
             fmt_sql: fmt_sql_fn_wrapper,
-            snapshot: snapshot_fn_wrapper,
             snapshot_generation: snapshot_generation_fn_wrapper,
             is_volatile_node: is_volatile_node_fn_wrapper,
             display: display_fn_wrapper,
@@ -691,15 +677,6 @@ impl PhysicalExpr for ForeignPhysicalExpr {
         }
     }
 
-    fn snapshot(&self) -> Result<Option<Arc<dyn PhysicalExpr>>> {
-        unsafe {
-            let result = df_result!((self.expr.snapshot)(&self.expr))?;
-            Ok(result
-                .map(|expr| <Arc<dyn PhysicalExpr>>::from(&expr))
-                .into())
-        }
-    }
-
     fn snapshot_generation(&self) -> u64 {
         unsafe { (self.expr.snapshot_generation)(&self.expr) }
     }
@@ -940,19 +917,13 @@ mod tests {
     }
 
     #[test]
-    fn ffi_physical_expr_snapshots() -> Result<(), DataFusionError> {
+    fn ffi_physical_expr_snapshot_generation() {
         let (original, foreign_expr) = create_test_expr();
 
-        let left = original.snapshot()?;
-        let right = foreign_expr.snapshot()?;
-        assert_eq!(left, right);
-
         assert_eq!(
             original.snapshot_generation(),
             foreign_expr.snapshot_generation()
         );
-
-        Ok(())
     }
 
     #[test]

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

@@ -27,9 +27,7 @@ use arrow::array::{ArrayRef, BooleanArray, new_empty_array};
 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, TreeNodeRecursion,
-};
+use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_common::{
     Result, ScalarValue, assert_eq_or_internal_err, exec_err, not_impl_err,
 };
@@ -357,55 +355,6 @@ pub trait PhysicalExpr: Any + Send + Sync + Display + Debug + DynEq + DynHash {
     /// See the [`fmt_sql`] function for an example of printing `PhysicalExpr`s as SQL.
     fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result;
 
-    /// Take a snapshot of this `PhysicalExpr`, if it is dynamic.
-    ///
-    /// "Dynamic" in this case means containing references to structures that may change
-    /// during plan execution, such as hash tables.
-    ///
-    /// This method is used to capture the current state of `PhysicalExpr`s that may contain
-    /// dynamic references to other operators in order to serialize it over the wire
-    /// or treat it via downcast matching.
-    ///
-    /// You should not call this method directly as it does not handle recursion.
-    /// Instead use [`snapshot_physical_expr`] to handle recursion and capture the
-    /// full state of the `PhysicalExpr`.
-    ///
-    /// This is expected to return "simple" expressions that do not have mutable state
-    /// and are composed of DataFusion's built-in `PhysicalExpr` implementations.
-    /// Callers however should *not* assume anything about the returned expressions
-    /// since callers and implementers may not agree on what "simple" or "built-in"
-    /// means.
-    /// In other words, if you need to serialize a `PhysicalExpr` across the wire
-    /// you should call this method and then try to serialize the result,
-    /// but you should handle unknown or unexpected `PhysicalExpr` implementations gracefully
-    /// just as if you had not called this method at all.
-    ///
-    /// In particular, consider:
-    /// * A `PhysicalExpr` that references the current state of a `datafusion::physical_plan::TopK`
-    ///   that is involved in a query with `SELECT * FROM t1 ORDER BY a LIMIT 10`.
-    ///   This function may return something like `a >= 12`.
-    /// * A `PhysicalExpr` that references the current state of a `datafusion::physical_plan::joins::HashJoinExec`
-    ///   from a query such as `SELECT * FROM t1 JOIN t2 ON t1.a = t2.b`.
-    ///   This function may return something like `t2.b IN (1, 5, 7)`.
-    ///
-    /// A system or function that can only deal with a hardcoded set of `PhysicalExpr` implementations
-    /// or needs to serialize this state to bytes may not be able to handle these dynamic references.
-    /// In such cases, we should return a simplified version of the `PhysicalExpr` that does not
-    /// contain these dynamic references.
-    ///
-    /// Systems that implement remote execution of plans, e.g. serialize a portion of the query plan
-    /// and send it across the wire to a remote executor may want to call this method after
-    /// every batch on the source side and broadcast / update the current snapshot to the remote executor.
-    ///
-    /// Note for implementers: this method should *not* handle recursion.
-    /// Recursion is handled in [`snapshot_physical_expr`].
-    fn snapshot(&self) -> Result<Option<Arc<dyn PhysicalExpr>>> {
-        // By default, we return None to indicate that this PhysicalExpr does not
-        // have any dynamic references or state.
-        // 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.
@@ -618,50 +567,6 @@ pub fn fmt_sql(expr: &dyn PhysicalExpr) -> impl Display + '_ {
     Wrapper { expr }
 }
 
-/// Take a snapshot of the given `PhysicalExpr` if it is dynamic.
-///
-/// Take a snapshot of this `PhysicalExpr` if it is dynamic.
-/// This is used to capture the current state of `PhysicalExpr`s that may contain
-/// dynamic references to other operators in order to serialize it over the wire
-/// or treat it via downcast matching.
-///
-/// See the documentation of [`PhysicalExpr::snapshot`] for more details.
-///
-/// # Returns
-///
-/// Returns a snapshot of the `PhysicalExpr` if it is dynamic, otherwise
-/// returns itself.
-pub fn snapshot_physical_expr(
-    expr: Arc<dyn PhysicalExpr>,
-) -> Result<Arc<dyn PhysicalExpr>> {
-    snapshot_physical_expr_opt(expr).data()
-}
-
-/// Take a snapshot of the given `PhysicalExpr` if it is dynamic.
-///
-/// Take a snapshot of this `PhysicalExpr` if it is dynamic.
-/// This is used to capture the current state of `PhysicalExpr`s that may contain
-/// dynamic references to other operators in order to serialize it over the wire
-/// or treat it via downcast matching.
-///
-/// See the documentation of [`PhysicalExpr::snapshot`] for more details.
-///
-/// # Returns
-///
-/// Returns a `[`Transformed`] indicating whether a snapshot was taken,
-/// along with the resulting `PhysicalExpr`.
-pub fn snapshot_physical_expr_opt(
-    expr: Arc<dyn PhysicalExpr>,
-) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
-    expr.transform_up(|e| {
-        if let Some(snapshot) = e.snapshot()? {
-            Ok(Transformed::yes(snapshot))
-        } else {
-            Ok(Transformed::no(Arc::clone(&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.

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

@@ -525,11 +525,6 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr {
         self.render(f, |expr, f| expr.fmt_sql(f))
     }
 
-    fn snapshot(&self) -> Result<Option<Arc<dyn PhysicalExpr>>> {
-        // 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().generation
@@ -611,14 +606,20 @@ mod test {
             &filter_schema_1,
         )
         .unwrap();
-        let snap = dynamic_filter_1.snapshot().unwrap().unwrap();
+        let dynamic_filter_1_df = dynamic_filter_1
+            .downcast_ref::<DynamicFilterPhysicalExpr>()
+            .expect("Expected DynamicFilterPhysicalExpr");
+        let snap = dynamic_filter_1_df.current().unwrap();
         insta::assert_snapshot!(format!("{snap:?}"), @r#"BinaryExpr { left: Column { name: "a", index: 0 }, op: Eq, right: Literal { value: Int32(42), field: Field { name: "lit", data_type: Int32 } }, fail_on_overflow: false }"#);
         let dynamic_filter_2 = reassign_expr_columns(
             Arc::clone(&dynamic_filter) as Arc<dyn PhysicalExpr>,
             &filter_schema_2,
         )
         .unwrap();
-        let snap = dynamic_filter_2.snapshot().unwrap().unwrap();
+        let dynamic_filter_2_df = dynamic_filter_2
+            .downcast_ref::<DynamicFilterPhysicalExpr>()
+            .expect("Expected DynamicFilterPhysicalExpr");
+        let snap = dynamic_filter_2_df.current().unwrap();
         insta::assert_snapshot!(format!("{snap:?}"), @r#"BinaryExpr { left: Column { name: "a", index: 1 }, op: Eq, right: Literal { value: Int32(42), field: Field { name: "lit", data_type: Int32 } }, fail_on_overflow: false }"#);
         // Both filters allow evaluating the same expression
         let batch_1 = RecordBatch::try_new(
@@ -684,20 +685,17 @@ mod test {
     }
 
     #[test]
-    fn test_snapshot() {
+    fn test_current() {
         let expr = lit(42) as Arc<dyn PhysicalExpr>;
         let dynamic_filter = DynamicFilterPhysicalExpr::new(vec![], Arc::clone(&expr));
 
-        // Take a snapshot of the current expression
-        let snapshot = dynamic_filter.snapshot().unwrap();
-        assert_eq!(snapshot, Some(expr));
+        // Read the current expression.
+        assert_eq!(&dynamic_filter.current().unwrap(), &expr);
 
-        // Update the current expression
+        // Update the current expression.
         let new_expr = lit(100) as Arc<dyn PhysicalExpr>;
         dynamic_filter.update(Arc::clone(&new_expr)).unwrap();
-        // Take another snapshot
-        let snapshot = dynamic_filter.snapshot().unwrap();
-        assert_eq!(snapshot, Some(new_expr));
+        assert_eq!(&dynamic_filter.current().unwrap(), &new_expr);
     }
 
     #[test]

datafusion/pruning/src/pruning_predicate.rs

@@ -42,9 +42,9 @@ use datafusion_common::{
     tree_node::{Transformed, TreeNode},
 };
 use datafusion_expr_common::operator::Operator;
+use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr;
 use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee};
 use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr};
-use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr_opt;
 use datafusion_physical_plan::{ColumnarValue, PhysicalExpr};
 
 /// Used to prove that arbitrary predicates (boolean expression) can not
@@ -456,22 +456,28 @@ impl PruningPredicate {
     /// details.
     ///
     /// Note that `PruningPredicate` does not attempt to normalize or simplify
-    /// the input expression unless calling [`snapshot_physical_expr_opt`]
-    /// returns a new expression.
+    /// the input expression unless any [`DynamicFilterPhysicalExpr`] nodes
+    /// are replaced with their current value.
     /// It is recommended that you pass the expressions through [`PhysicalExprSimplifier`]
     /// before calling this method to make sure the expressions can be used for pruning.
     pub fn try_new(mut expr: Arc<dyn PhysicalExpr>, schema: SchemaRef) -> Result<Self> {
-        // Get a (simpler) snapshot of the physical expr here to use with `PruningPredicate`.
-        // In particular this unravels any `DynamicFilterPhysicalExpr`s by snapshotting them
-        // so that PruningPredicate can work with a static expression.
-        let tf = snapshot_physical_expr_opt(expr)?;
+        // Replace any `DynamicFilterPhysicalExpr` nodes with their current value so
+        // that `PruningPredicate` can work with a static expression.
+        let tf = expr.transform_up(|e| {
+            if let Some(df) = e.downcast_ref::<DynamicFilterPhysicalExpr>() {
+                Ok(Transformed::yes(df.current()?))
+            } else {
+                Ok(Transformed::no(e))
+            }
+        })?;
         if tf.transformed {
             // If we had an expression such as Dynamic(part_col < 5 and col < 10)
             // (this could come from something like `select * from t order by part_col, col, limit 10`)
-            // after snapshotting and because `DynamicFilterPhysicalExpr` applies child replacements to its
-            // children after snapshotting and previously `replace_columns_with_literals` may have been called with partition values
-            // the expression we have now is `8 < 5 and col < 10`.
-            // Thus we need as simplifier pass to get `false and col < 10` => `false` here.
+            // after replacing the dynamic filter with its current value and because
+            // `DynamicFilterPhysicalExpr` applies child replacements to its children,
+            // and previously `replace_columns_with_literals` may have been called with
+            // partition values, the expression we have now is `8 < 5 and col < 10`.
+            // Thus we need a simplifier pass to get `false and col < 10` => `false` here.
             let simplifier = PhysicalExprSimplifier::new(&schema);
             expr = simplifier.simplify(tf.data)?;
         } else {