proto: serialize and dedupe dynamic filters

Informs: datafusion-contrib/datafusion-distributed#180
Closes: #20418

Consider this scenario
1. You have a plan with a `HashJoinExec` and `DataSourceExec`
2. You run the physical optimizer and the `DataSourceExec` accepts `DynamicFilterPhysicalExpr` pushdown from the `HashJoinExec`
3. You serialize the plan, deserialize it, and execute it

What should happen is that the dynamic filter should "work", meaning
1. When you deserialize the plan, both the `HashJoinExec` and `DataSourceExec` should have pointers to the same `DynamicFilterPhysicalExpr`
2. The `DynamicFilterPhysicalExpr` should be updated during execution by the `HashJoinExec`  and the `DataSourceExec` should filter out rows

This does not happen today for a few reasons, a couple of which this PR aims to address
1. `DynamicFilterPhysicalExpr` is not survive round-tripping. The internal exprs get inlined (ex. it may be serialized as `Literal`)
2. Even if `DynamicFilterPhysicalExpr` survives round-tripping, during pushdown, it's often the case that the `DynamicFilterPhysicalExpr` is rewritten. In this case, you have two `DynamicFilterPhysicalExpr` which are different `Arc`s but share the same `Inner` dynamic filter state. The current `DeduplicatingProtoConverter` does not handle this specific form of deduping.

This PR aims to fix those problems by adding serde for `DynamicFilterPhysicalExpr` and deduping logic for the inner state of dynamic filters.

It does not yet add a test for the `HashJoinExec` and `DataSourceExec` filter pushdown case, but this is relevant follow up work. I tried to keep the PR small for reviewers.

Yes, via unit tests.

`DynamicFilterPhysicalExpr` are now serialized by the default codec

Author: jayshrivastava

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

@@ -18,7 +18,7 @@
 use std::any::Any;
 use std::fmt;
 use std::fmt::{Debug, Display, Formatter};
-use std::hash::{Hash, Hasher};
+use std::hash::{DefaultHasher, Hash, Hasher};
 use std::sync::Arc;
 
 use crate::utils::scatter;
@@ -441,6 +441,50 @@ pub trait PhysicalExpr: Any + Send + Sync + Display + Debug + DynEq + DynHash {
     fn placement(&self) -> ExpressionPlacement {
         ExpressionPlacement::KeepInPlace
     }
+
+    /// Returns a composite identifier for a [`PhysicalExpr`]. Note that if the expression
+    /// is dropped, then the returned id is no longer valid.
+    fn expr_id(self: Arc<Self>, salt: &[u64]) -> Option<PhysicalExprId> {
+        Some(PhysicalExprId::new(expr_id_from_arc(&self, salt), None))
+    }
+}
+
+/// A composite identifier for [`PhysicalExpr`].
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct PhysicalExprId {
+    exact: u64,
+    shallow: Option<u64>,
+}
+
+impl PhysicalExprId {
+    /// Create a new [`PhysicalExprId`]. Both ids must be globally unique within
+    /// a process.
+    pub fn new(exact: u64, shallow: Option<u64>) -> Self {
+        Self { exact, shallow }
+    }
+
+    /// Returns the identifier for the full expression tree, including children.
+    pub fn exact(&self) -> u64 {
+        self.exact
+    }
+
+    /// Returns the identifier for just the expression root, ignoring children.
+    pub fn shallow(&self) -> Option<u64> {
+        self.shallow
+    }
+}
+
+/// Computes a unique identifier for a type contained within an [`Arc`]. It hashes
+/// the [`Arc`] pointer to create a process-local identifier that remains valid
+/// only while that allocation is still alive.
+pub fn expr_id_from_arc<T: ?Sized>(expr: &Arc<T>, salt: &[u64]) -> u64 {
+    let mut hasher = DefaultHasher::new();
+    let ptr = Arc::as_ptr(expr) as *const () as u64;
+    ptr.hash(&mut hasher);
+    for &salt in salt {
+        salt.hash(&mut hasher);
+    }
+    hasher.finish()
 }
 
 #[deprecated(

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

@@ -26,7 +26,9 @@ use datafusion_common::{
     tree_node::{Transformed, TransformedResult, TreeNode},
 };
 use datafusion_expr::ColumnarValue;
-use datafusion_physical_expr_common::physical_expr::DynHash;
+use datafusion_physical_expr_common::physical_expr::{
+    DynHash, PhysicalExprId, expr_id_from_arc,
+};
 
 /// State of a dynamic filter, tracking both updates and completion.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -88,6 +90,116 @@ struct Inner {
     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>>>,
+    // Inner state.
+    generation: u64,
+    inner_expr: Arc<dyn PhysicalExpr>,
+    is_complete: bool,
+}
+
+impl DynamicFilterSnapshot {
+    pub fn new(
+        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,
+            generation,
+            inner_expr,
+            is_complete,
+        }
+    }
+
+    pub fn children(&self) -> &[Arc<dyn PhysicalExpr>] {
+        &self.children
+    }
+
+    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
+    }
+}
+impl Display for DynamicFilterSnapshot {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "DynamicFilterSnapshot {{ children: {:?}, remapped_children: {:?}, generation: {}, inner_expr: {:?}, is_complete: {} }}",
+            self.children,
+            self.remapped_children,
+            self.generation,
+            self.inner_expr,
+            self.is_complete
+        )
+    }
+}
+
+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 {
@@ -448,6 +560,13 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr {
         // Return the current generation of the expression.
         self.inner.read().generation
     }
+
+    fn expr_id(self: Arc<Self>, salt: &[u64]) -> Option<PhysicalExprId> {
+        Some(PhysicalExprId::new(
+            expr_id_from_arc(&self, salt),
+            Some(expr_id_from_arc(&self.inner, salt)),
+        ))
+    }
 }
 
 #[cfg(test)]
@@ -867,4 +986,120 @@ mod test {
             "Hash should be stable after update (identity-based)"
         );
     }
+
+    #[test]
+    fn test_current_snapshot_roundtrip() {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+        let col_a = col("a", &schema).unwrap();
+
+        // Create a dynamic filter with children
+        let expr = Arc::new(BinaryExpr::new(
+            Arc::clone(&col_a),
+            datafusion_expr::Operator::Gt,
+            lit(10) as Arc<dyn PhysicalExpr>,
+        ));
+        let filter = DynamicFilterPhysicalExpr::new(
+            vec![Arc::clone(&col_a)],
+            expr as Arc<dyn PhysicalExpr>,
+        );
+
+        // Update expression and mark complete
+        filter
+            .update(lit(42) as Arc<dyn PhysicalExpr>)
+            .expect("Update should succeed");
+        filter.mark_complete();
+
+        // Take a snapshot and reconstruct
+        let snapshot = DynamicFilterSnapshot::from(&filter);
+        let reconstructed = DynamicFilterPhysicalExpr::from(snapshot);
+
+        // String representations should be equal
+        assert_eq!(
+            DynamicFilterSnapshot::from(&filter).to_string(),
+            DynamicFilterSnapshot::from(&reconstructed).to_string(),
+        );
+    }
+
+    #[tokio::test]
+    async fn test_with_new_children_preserves_shared_inner_state() {
+        // Create a source filter
+        let source = Arc::new(DynamicFilterPhysicalExpr::new(
+            vec![],
+            lit(42) as Arc<dyn PhysicalExpr>,
+        ));
+
+        // 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();
+        // Rebuild a wrapper from the source filter using the target's children.
+        let combined = Arc::clone(&source)
+            .with_new_children(vec![Arc::clone(&col_x)])
+            .unwrap();
+
+        let combined_expr_id = Arc::clone(&combined)
+            .expr_id(&[])
+            .expect("combined filter should have an expr_id");
+        let source_expr_id = Arc::clone(&source)
+            .expr_id(&[])
+            .expect("source filter should have an expr_id");
+        assert_ne!(
+            combined_expr_id.exact(),
+            source_expr_id.exact(),
+            "with_new_children should produce a distinct outer Arc identity",
+        );
+        assert_eq!(
+            combined_expr_id.shallow(),
+            source_expr_id.shallow(),
+            "dynamic filters with shared inner state should have the same shallow identity",
+        );
+
+        let combined_dyn_filter = combined
+            .as_any()
+            .downcast_ref::<DynamicFilterPhysicalExpr>()
+            .unwrap();
+
+        // Verify the children are unchanged.
+        assert_eq!(
+            format!("{:?}", combined.children()),
+            format!("{:?}", vec![col_x]),
+            "Combined filter's children should be unchanged"
+        );
+
+        // Verify inner expression changed to the one from source.
+        assert_eq!(
+            format!("{:?}", combined_dyn_filter.current().unwrap()),
+            format!("{:?}", lit(42)),
+            "Combined filter should have inner expression from source filter"
+        );
+
+        // Verify that completing one filter also completes the other.
+        let combined_binding = Arc::clone(&combined) as Arc<dyn PhysicalExpr>;
+        #[expect(clippy::disallowed_methods)]
+        let wait_handle = tokio::task::spawn({
+            async move {
+                let df = combined_binding
+                    .as_any()
+                    .downcast_ref::<DynamicFilterPhysicalExpr>()
+                    .unwrap();
+                df.wait_complete().await;
+                format!("{:?}", df.current().unwrap())
+            }
+        });
+        source.update(lit(999) as Arc<dyn PhysicalExpr>).unwrap();
+        source.mark_complete();
+
+        // The linked filter should be notified via the shared watch
+        let result = tokio::time::timeout(std::time::Duration::from_secs(1), wait_handle)
+            .await
+            .expect(
+                "linked filter should have been notified of completion within timeout",
+            )
+            .expect("task should not panic");
+
+        assert_eq!(
+            result,
+            format!("{:?}", lit(999)),
+            "linked filter should see source's updated inner expression"
+        );
+    }
 }

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

@@ -45,7 +45,7 @@ pub use cast::{CastExpr, cast};
 pub use cast_column::CastColumnExpr;
 pub use column::{Column, col, with_new_schema};
 pub use datafusion_expr::utils::format_state_name;
-pub use dynamic_filters::DynamicFilterPhysicalExpr;
+pub use dynamic_filters::{DynamicFilterPhysicalExpr, DynamicFilterSnapshot};
 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

@@ -871,17 +871,18 @@ message PhysicalExtensionNode {
 }
 
 // physical expressions
+message PhysicalExprId {
+  uint64 exact = 1;
+  optional uint64 shallow = 2;
+}
+
 message PhysicalExprNode {
   // Was date_time_interval_expr
   reserved 17;
 
-  // Unique identifier for this expression to do deduplication during deserialization.
-  // When serializing, this is set to a unique identifier for each combination of
-  // expression, process and serialization run.
-  // When deserializing, if this ID has been seen before, the cached Arc is returned
-  // instead of creating a new one, enabling reconstruction of referential integrity
-  // across serde roundtrips.
-  optional uint64 expr_id = 30;
+  // Unique identifier for this expression used during deserialization to restore
+  // referential integrity across serde roundtrips.
+  PhysicalExprId expr_id = 31;
 
   oneof ExprType {
     // column references
@@ -920,9 +921,19 @@ message PhysicalExprNode {
     UnknownColumn unknown_column = 20;
 
     PhysicalHashExprNode hash_expr = 21;
+
+    PhysicalDynamicFilterNode dynamic_filter = 22;
   }
 }
 
+message PhysicalDynamicFilterNode {
+  repeated PhysicalExprNode children = 1;
+  repeated PhysicalExprNode remapped_children = 2;
+  uint64 generation = 3;
+  PhysicalExprNode inner_expr = 4;
+  bool is_complete = 5;
+}
+
 message PhysicalScalarUdfNode {
   string name = 1;
   repeated PhysicalExprNode args = 2;
@@ -1470,4 +1481,4 @@ message AsyncFuncExecNode {
 message BufferExecNode {
   PhysicalPlanNode input = 1;
   uint64 capacity = 2;
-}
+}