address pr comments - add From trait and improve safety check on new_from_source

Author: jayshrivastava

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

@@ -151,6 +151,54 @@ impl Display for DynamicFilterSnapshot {
     }
 }
 
+impl From<DynamicFilterSnapshot> for DynamicFilterPhysicalExpr {
+    fn from(snapshot: DynamicFilterSnapshot) -> Self {
+        let DynamicFilterSnapshot {
+            children,
+            remapped_children,
+            generation,
+            inner_expr,
+            is_complete,
+        } = snapshot;
+
+        let state = if is_complete {
+            FilterState::Complete { generation }
+        } else {
+            FilterState::InProgress { generation }
+        };
+        let (state_watch, _) = watch::channel(state);
+
+        Self {
+            children,
+            remapped_children,
+            inner: Arc::new(RwLock::new(Inner {
+                generation,
+                expr: inner_expr,
+                is_complete,
+            })),
+            state_watch,
+            data_type: Arc::new(RwLock::new(None)),
+            nullable: Arc::new(RwLock::new(None)),
+        }
+    }
+}
+
+impl From<&DynamicFilterPhysicalExpr> for DynamicFilterSnapshot {
+    fn from(expr: &DynamicFilterPhysicalExpr) -> Self {
+        let (generation, inner_expr, is_complete) = {
+            let inner = expr.inner.read();
+            (inner.generation, Arc::clone(&inner.expr), inner.is_complete)
+        };
+        DynamicFilterSnapshot {
+            children: expr.children.clone(),
+            remapped_children: expr.remapped_children.clone(),
+            generation,
+            inner_expr,
+            is_complete,
+        }
+    }
+}
+
 impl Inner {
     fn new(expr: Arc<dyn PhysicalExpr>) -> Self {
         Self {
@@ -243,58 +291,6 @@ impl DynamicFilterPhysicalExpr {
         }
     }
 
-    /// Reconstructs a [`DynamicFilterPhysicalExpr`] from a snapshot.
-    ///
-    /// This is a low-level API intended for use by the proto deserialization layer.
-    #[doc(hidden)]
-    pub fn new_from_snapshot(snapshot: DynamicFilterSnapshot) -> Self {
-        let DynamicFilterSnapshot {
-            children,
-            remapped_children,
-            generation,
-            inner_expr,
-            is_complete,
-        } = snapshot;
-
-        let state = if is_complete {
-            FilterState::Complete { generation }
-        } else {
-            FilterState::InProgress { generation }
-        };
-        let (state_watch, _) = watch::channel(state);
-
-        Self {
-            children,
-            remapped_children,
-            inner: Arc::new(RwLock::new(Inner {
-                generation,
-                expr: inner_expr,
-                is_complete,
-            })),
-            state_watch,
-            data_type: Arc::new(RwLock::new(None)),
-            nullable: Arc::new(RwLock::new(None)),
-        }
-    }
-
-    /// Atomically captures all state needed for serialization into a [`DynamicFilterSnapshot`].
-    ///
-    /// This is a low-level API intended for use by the proto deserialization layer.
-    #[doc(hidden)]
-    pub fn current_snapshot(&self) -> DynamicFilterSnapshot {
-        let (generation, inner_expr, is_complete) = {
-            let inner = self.inner.read();
-            (inner.generation, Arc::clone(&inner.expr), inner.is_complete)
-        };
-        DynamicFilterSnapshot {
-            children: self.children.clone(),
-            remapped_children: self.remapped_children.clone(),
-            generation,
-            inner_expr,
-            is_complete,
-        }
-    }
-
     /// Create a new [`DynamicFilterPhysicalExpr`] from `self`, except it overwrites the
     /// internal state with the source filter's state.
     ///
@@ -303,16 +299,21 @@ impl DynamicFilterPhysicalExpr {
     /// # Safety
     ///
     /// The dynamic filter should not be in use when calling this method, otherwise there
-    /// may be undefined behavior. This method may do the following or worse:
+    /// may be undefined behavior. Changing the inner state of a filter may do the following:
     /// - transition the state to complete without notifying the watch
     /// - cause a generation number to be emitted which is out of order
-    pub fn new_from_source(&self, source: &DynamicFilterPhysicalExpr) -> Result<Self> {
-        // Best effort check that no one is subscribed.
-        if self.state_watch.receiver_count() > 0 {
+    pub fn new_from_source(
+        self: &Arc<Self>,
+        source: &DynamicFilterPhysicalExpr,
+    ) -> Result<Self> {
+        // If there's any references to this filter or any watchers, we should not replace the
+        // inner state.
+        if self.is_used() || self.state_watch.receiver_count() > 0 {
             return internal_err!(
-                "Cannot replace the inner state of a DynamicFilterPhysicalExpr that has subscribers"
+                "Cannot replace the inner state of a DynamicFilterPhysicalExpr that is in use"
             );
         };
+
         Ok(Self {
             children: self.children.clone(),
             remapped_children: self.remapped_children.clone(),
@@ -1020,13 +1021,13 @@ mod test {
         filter.mark_complete();
 
         // Take a snapshot and reconstruct
-        let snapshot = filter.current_snapshot();
-        let reconstructed = DynamicFilterPhysicalExpr::new_from_snapshot(snapshot);
+        let snapshot = DynamicFilterSnapshot::from(&filter);
+        let reconstructed = DynamicFilterPhysicalExpr::from(snapshot);
 
         // String representations should be equal
         assert_eq!(
-            filter.current_snapshot().to_string(),
-            reconstructed.current_snapshot().to_string(),
+            DynamicFilterSnapshot::from(&filter).to_string(),
+            DynamicFilterSnapshot::from(&reconstructed).to_string(),
         );
     }
 
@@ -1045,10 +1046,10 @@ mod test {
         // Create a target filter with different children
         let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
         let col_x = col("x", &schema).unwrap();
-        let target = DynamicFilterPhysicalExpr::new(
+        let target = Arc::new(DynamicFilterPhysicalExpr::new(
             vec![Arc::clone(&col_x)],
             lit(0) as Arc<dyn PhysicalExpr>,
-        );
+        ));
 
         // Create new filter from source's inner state
         let combined = target.new_from_source(&source).unwrap();
@@ -1061,7 +1062,7 @@ mod test {
         );
 
         // Verify children are from target, not source
-        let combined_snapshot = combined.current_snapshot();
+        let combined_snapshot = DynamicFilterSnapshot::from(&combined);
         assert_eq!(
             combined_snapshot.children().len(),
             1,

datafusion/proto/src/physical_plan/from_proto.rs

@@ -529,16 +529,16 @@ pub fn parse_physical_expr_with_converter(
             )?;
 
             // Recreate filter from snapshot
-            let base_filter = Arc::new(DynamicFilterPhysicalExpr::new_from_snapshot(
-                DynamicFilterSnapshot::new(
-                    children,
-                    remapped_children,
-                    dynamic_filter.generation,
-                    inner_expr,
-                    dynamic_filter.is_complete,
-                ),
-            ));
-            base_filter as Arc<dyn PhysicalExpr>
+            let snapshot = DynamicFilterSnapshot::new(
+                children,
+                remapped_children,
+                dynamic_filter.generation,
+                inner_expr,
+                dynamic_filter.is_complete,
+            );
+            let base_filter: Arc<dyn PhysicalExpr> =
+                Arc::new(DynamicFilterPhysicalExpr::from(snapshot));
+            base_filter
         }
         ExprType::Extension(extension) => {
             let inputs: Vec<Arc<dyn PhysicalExpr>> = extension

datafusion/proto/src/physical_plan/mod.rs

@@ -1,3 +1,4 @@
+use std::any::Any;
 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
@@ -3962,13 +3963,8 @@ impl PhysicalProtoConverterExtension for DeduplicatingDeserializer {
                 };
 
                 // Get the base filter's structure
-                let Some(dynamic_filter_expr) =
-                    expr.as_any().downcast_ref::<DynamicFilterPhysicalExpr>()
-                else {
-                    return internal_err!(
-                        "dynamic_filter_id present in proto, but the expression was not a DynamicFilterPhysicalExpr"
-                    );
-                };
+                let dynamic_filter_expr = (expr as Arc<dyn Any + Send + Sync>).downcast::<DynamicFilterPhysicalExpr>()
+                    .map_err(|_| internal_datafusion_err!("dynamic_filter_id present in proto, but the expression was not a DynamicFilterPhysicalExpr"))?;
                 expr = Arc::new(dynamic_filter_expr.new_from_source(cached_df)?)
                     as Arc<dyn PhysicalExpr>;
             } else {

datafusion/proto/src/physical_plan/to_proto.rs

@@ -36,9 +36,9 @@ use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindo
 use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr;
 use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
 use datafusion_physical_plan::expressions::{
-    BinaryExpr, CaseExpr, CastExpr, Column, DynamicFilterPhysicalExpr, InListExpr,
-    IsNotNullExpr, IsNullExpr, LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr,
-    UnKnownColumn,
+    BinaryExpr, CaseExpr, CastExpr, Column, DynamicFilterPhysicalExpr,
+    DynamicFilterSnapshot, InListExpr, IsNotNullExpr, IsNullExpr, LikeExpr, Literal,
+    NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn,
 };
 use datafusion_physical_plan::joins::{HashExpr, HashTableLookupExpr};
 use datafusion_physical_plan::udaf::AggregateFunctionExpr;
@@ -330,7 +330,7 @@ pub fn serialize_physical_expr_with_converter(
         })
     } else if let Some(df) = expr.downcast_ref::<DynamicFilterPhysicalExpr>() {
         // Capture all state atomically
-        let snapshot = df.current_snapshot();
+        let snapshot = DynamicFilterSnapshot::from(df);
 
         let children = snapshot
             .children()

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

@@ -130,7 +130,9 @@ use crate::cases::{
     MyRegexUdfNode,
 };
 
-use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr;
+use datafusion_physical_expr::expressions::{
+    DynamicFilterPhysicalExpr, DynamicFilterSnapshot,
+};
 use datafusion_physical_expr::utils::reassign_expr_columns;
 
 /// Perform a serde roundtrip and assert that the string representation of the before and after plans
@@ -3066,8 +3068,8 @@ fn test_dynamic_filter_roundtrip() -> Result<()> {
         .expect("Should be DynamicFilterPhysicalExpr");
 
     assert_eq!(
-        df.current_snapshot().to_string(),
-        deserialized_df.current_snapshot().to_string(),
+        DynamicFilterSnapshot::from(df).to_string(),
+        DynamicFilterSnapshot::from(deserialized_df).to_string(),
         "Snapshots should be equal"
     );
 
@@ -3221,12 +3223,12 @@ fn test_deduplication_of_dynamic_filter_expression(
         .unwrap();
 
     assert_eq!(
-        filter_1_before_roundtrip.current_snapshot().to_string(),
-        df1.current_snapshot().to_string()
+        DynamicFilterSnapshot::from(filter_1_before_roundtrip).to_string(),
+        DynamicFilterSnapshot::from(df1).to_string()
     );
     assert_eq!(
-        filter_2_before_roundtrip.current_snapshot().to_string(),
-        df2.current_snapshot().to_string()
+        DynamicFilterSnapshot::from(filter_2_before_roundtrip).to_string(),
+        DynamicFilterSnapshot::from(df2).to_string()
     );
 
     Ok(())