address feedback

Author: jayshrivastava

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

@@ -80,72 +80,21 @@ pub struct DynamicFilterPhysicalExpr {
     nullable: Arc<RwLock<Option<bool>>>,
 }
 
-#[derive(Debug)]
-struct Inner {
+/// Atomic internal state of a [`DynamicFilterPhysicalExpr`].
+///
+/// **Warning:** exposed publicly solely so that proto (de)serialization in
+/// `datafusion-proto` can read and rebuild this state. Do not treat this type
+/// or its layout as a stable API.
+#[derive(Debug, Clone)]
+pub 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::snapshot_generation`] to have a cheap check for changes.
-    generation: u64,
-    expr: Arc<dyn PhysicalExpr>,
+    pub generation: u64,
+    pub expr: Arc<dyn PhysicalExpr>,
     /// Flag for quick synchronous check if filter is complete.
     /// This is redundant with the watch channel state, but allows us to return immediately
     /// from `wait_complete()` without subscribing if already complete.
-    is_complete: bool,
-}
-
-/// An atomic snapshot of a [`DynamicFilterPhysicalExpr`] used to reconstruct the expression during
-/// serialization / deserialization.
-pub struct DynamicFilterSnapshot {
-    children: Vec<Arc<dyn PhysicalExpr>>,
-    remapped_children: Option<Vec<Arc<dyn PhysicalExpr>>>,
-    expression_id: u64,
-    // Inner state.
-    generation: u64,
-    inner_expr: Arc<dyn PhysicalExpr>,
-    is_complete: bool,
-}
-
-impl DynamicFilterSnapshot {
-    pub fn new(
-        expression_id: u64,
-        children: Vec<Arc<dyn PhysicalExpr>>,
-        remapped_children: Option<Vec<Arc<dyn PhysicalExpr>>>,
-        generation: u64,
-        inner_expr: Arc<dyn PhysicalExpr>,
-        is_complete: bool,
-    ) -> Self {
-        Self {
-            children,
-            remapped_children,
-            expression_id,
-            generation,
-            inner_expr,
-            is_complete,
-        }
-    }
-
-    pub fn children(&self) -> &[Arc<dyn PhysicalExpr>] {
-        &self.children
-    }
-
-    pub fn expression_id(&self) -> u64 {
-        self.expression_id
-    }
-
-    pub fn remapped_children(&self) -> Option<&[Arc<dyn PhysicalExpr>]> {
-        self.remapped_children.as_deref()
-    }
-
-    pub fn generation(&self) -> u64 {
-        self.generation
-    }
-
-    pub fn inner_expr(&self) -> &Arc<dyn PhysicalExpr> {
-        &self.inner_expr
-    }
-
-    pub fn is_complete(&self) -> bool {
-        self.is_complete
-    }
+    pub is_complete: bool,
 }
 
 // TODO: Include expression_id in debug output.
@@ -167,73 +116,6 @@ impl std::fmt::Debug for DynamicFilterPhysicalExpr {
     }
 }
 
-impl Display for DynamicFilterSnapshot {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "DynamicFilterSnapshot {{ children: {:?}, remapped_children: {:?}, expression_id: {}, generation: {}, inner_expr: {:?}, is_complete: {} }}",
-            self.children,
-            self.remapped_children,
-            self.expression_id,
-            self.generation,
-            self.inner_expr,
-            self.is_complete
-        )
-    }
-}
-
-impl From<DynamicFilterSnapshot> for DynamicFilterPhysicalExpr {
-    fn from(snapshot: DynamicFilterSnapshot) -> Self {
-        let DynamicFilterSnapshot {
-            children,
-            remapped_children,
-            expression_id,
-            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,
-            expression_id,
-            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 {
-        // Snapshot everything in the mutex atomically.
-        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(),
-            expression_id: expr.expression_id,
-            generation,
-            inner_expr,
-            is_complete,
-        }
-    }
-}
-
 impl Inner {
     fn new(expr: Arc<dyn PhysicalExpr>) -> Self {
         Self {
@@ -495,9 +377,67 @@ impl DynamicFilterPhysicalExpr {
     }
 
     /// Generate a new expression id for this filter.
-    pub fn new_expression_id() -> u64 {
+    fn new_expression_id() -> u64 {
         random::<u64>()
     }
+
+    /// Return the filter's original children (before any remapping).
+    pub fn original_children(&self) -> &[Arc<dyn PhysicalExpr>] {
+        &self.children
+    }
+
+    /// Return the filter's remapped children, if any have been set via
+    /// [`PhysicalExpr::with_new_children`].
+    pub fn remapped_children(&self) -> Option<&[Arc<dyn PhysicalExpr>]> {
+        self.remapped_children.as_deref()
+    }
+
+    /// Rebuild a `DynamicFilterPhysicalExpr` from its stored parts. Used by
+    /// proto deserialization to preserve `expression_id` across a roundtrip
+    /// rather than minting a fresh one.
+    ///
+    /// **Warning:** intended only for `datafusion-proto` (de)serialization.
+    /// Not a stable API.
+    pub fn from_parts(
+        expression_id: u64,
+        children: Vec<Arc<dyn PhysicalExpr>>,
+        remapped_children: Option<Vec<Arc<dyn PhysicalExpr>>>,
+        inner: Inner,
+    ) -> Self {
+        let state = if inner.is_complete {
+            FilterState::Complete {
+                generation: inner.generation,
+            }
+        } else {
+            FilterState::InProgress {
+                generation: inner.generation,
+            }
+        };
+        let (state_watch, _) = watch::channel(state);
+
+        Self {
+            children,
+            remapped_children,
+            expression_id,
+            inner: Arc::new(RwLock::new(inner)),
+            state_watch,
+            data_type: Arc::new(RwLock::new(None)),
+            nullable: Arc::new(RwLock::new(None)),
+        }
+    }
+
+    /// Return a clone of the atomically-captured `Inner` state.
+    ///
+    /// **Warning:** intended only for `datafusion-proto` (de)serialization.
+    /// Not a stable API.
+    pub fn inner(&self) -> Inner {
+        let guard = self.inner.read();
+        Inner {
+            generation: guard.generation,
+            expr: Arc::clone(&guard.expr),
+            is_complete: guard.is_complete,
+        }
+    }
 }
 
 impl PhysicalExpr for DynamicFilterPhysicalExpr {
@@ -1021,6 +961,10 @@ mod test {
         );
     }
 
+    /// Verifies that reassigning (which triggers `with_new_children`) followed
+    /// by rebuilding via `from_parts` preserves the filter's observable state
+    /// (original children, remapped children, expression id, and inner
+    /// generation/expr/is_complete).
     #[test]
     fn test_current_snapshot_roundtrip() {
         let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
@@ -1057,14 +1001,37 @@ mod test {
             .downcast_ref::<DynamicFilterPhysicalExpr>()
             .expect("Expected dynamic filter after reassignment");
 
-        // Take a snapshot and reconstruct
-        let snapshot = DynamicFilterSnapshot::from(reassigned);
-        let reconstructed = DynamicFilterPhysicalExpr::from(snapshot);
+        // Capture the parts and reconstruct.
+        let reconstructed = DynamicFilterPhysicalExpr::from_parts(
+            reassigned.expression_id,
+            reassigned.original_children().to_vec(),
+            reassigned.remapped_children().map(|r| r.to_vec()),
+            reassigned.inner(),
+        );
 
-        // Assert snapshots are equal.
+        // Assert the reconstructed filter matches field-by-field.
+        assert_eq!(
+            reassigned.original_children(),
+            reconstructed.original_children(),
+        );
+        assert_eq!(
+            reassigned.remapped_children(),
+            reconstructed.remapped_children(),
+        );
+        assert_eq!(
+            reassigned.expression_id().unwrap(),
+            reconstructed.expression_id().unwrap(),
+        );
+        let reassigned_inner = reassigned.inner();
+        let reconstructed_inner = reconstructed.inner();
+        assert_eq!(reassigned_inner.generation, reconstructed_inner.generation);
+        assert_eq!(
+            reassigned_inner.is_complete,
+            reconstructed_inner.is_complete
+        );
         assert_eq!(
-            DynamicFilterSnapshot::from(reassigned).to_string(),
-            DynamicFilterSnapshot::from(&reconstructed).to_string(),
+            format!("{:?}", reassigned_inner.expr),
+            format!("{:?}", reconstructed_inner.expr),
         );
     }
 

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

@@ -43,7 +43,7 @@ pub use case::{CaseExpr, case};
 pub use cast::{CastExpr, cast};
 pub use column::{Column, col, with_new_schema};
 pub use datafusion_expr::utils::format_state_name;
-pub use dynamic_filters::{DynamicFilterPhysicalExpr, DynamicFilterSnapshot};
+pub use dynamic_filters::{DynamicFilterPhysicalExpr, Inner as DynamicFilterInner};
 pub use in_list::{InListExpr, in_list};
 pub use is_not_null::{IsNotNullExpr, is_not_null};
 pub use is_null::{IsNullExpr, is_null};

datafusion/proto/proto/datafusion.proto

@@ -875,8 +875,11 @@ message PhysicalExprNode {
   // Was date_time_interval_expr
   reserved 17;
 
-  // Was expr_id
-  reserved 30;
+  // Unique identifier for this expression to do deduplication during deserialization.
+  // When serializing, this is set via `PhysicalExpr::expression_id`. When deserializing,
+  // this id is used by the `DeduplicatingProtoConverter` to preserve referential
+  // integrity across serde roundtrips for different expressions with the same id.
+  optional uint64 expr_id = 30;
 
   oneof ExprType {
     // column references
@@ -926,7 +929,6 @@ message PhysicalDynamicFilterNode {
   uint64 generation = 3;
   PhysicalExprNode inner_expr = 4;
   bool is_complete = 5;
-  uint64 expression_id = 6;
 }
 
 message PhysicalScalarUdfNode {