proto: serialize and dedupe dynamic filters

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

Consider you have a plan with a `HashJoinExec` and `DataSourceExec`
```
HashJoinExec(dynamic_filter_1 on a@0)
  (...left side of join)
  ProjectionExec(a := Column("a", source_index))
    DataSourceExec
      ParquetSource(predicate = dynamic_filter_2)
```

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`) due to the `PhysicalExpr::snapshot()` API
2. Even if `DynamicFilterPhysicalExpr` survives round-tripping, the one pushed down to the `DataSourceExec` often has different children. In this case, you have two `DynamicFilterPhysicalExpr` which
do not survive deduping, causing referential integrity to be lost.

This PR aims to fix those problems by:
1. Removing the `snapshot()` call from the serialization process
2. Adding protos for `DynamicFilterPhysicalExpr` so it can be serialized and deserialized
3. Adding a new concept, a `PhysicalExprId`, which has two identifiers,
   a "shallow" identifier to indicate two equal expressions which may
   have different children, and an "exact" identifier to indicate two
   exprs that are exactly the same.
4. Updating the deduping deserializer and protos to now be aware of the
   new "shallow" id, deduping exprs which are the same but have
   different children accordingly.

This change adds tests which roundtrip dynamic filters and assert that
referential integrity is maintained.

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,118 @@ 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 {
+        // 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(),
+            generation,
+            inner_expr,
+            is_complete,
+        }
+    }
+}
+
 impl Inner {
     fn new(expr: Arc<dyn PhysicalExpr>) -> Self {
         Self {
@@ -448,6 +562,15 @@ 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(
+            // Capture the outer arc, which contains children and the expr.
+            expr_id_from_arc(&self, salt),
+            // Capture the inner arc, which contains the expr only.
+            Some(expr_id_from_arc(&self.inner, salt)),
+        ))
+    }
 }
 
 #[cfg(test)]
@@ -867,4 +990,111 @@ 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();
+
+        // Change the children of the expr.
+        let reassigned_schema = Arc::new(Schema::new(vec![
+            Field::new("b", DataType::Int32, false),
+            Field::new("a", DataType::Int32, false),
+        ]));
+        let reassigned = reassign_expr_columns(
+            Arc::new(filter) as Arc<dyn PhysicalExpr>,
+            &reassigned_schema,
+        )
+        .expect("reassign_expr_columns should succeed");
+        let reassigned = reassigned
+            .as_any()
+            .downcast_ref::<DynamicFilterPhysicalExpr>()
+            .expect("Expected dynamic filter after reassignment");
+
+        // Take a snapshot and reconstruct
+        let snapshot = DynamicFilterSnapshot::from(reassigned);
+        let reconstructed = DynamicFilterPhysicalExpr::from(snapshot);
+
+        // Assert snapshots are equal.
+        assert_eq!(
+            DynamicFilterSnapshot::from(reassigned).to_string(),
+            DynamicFilterSnapshot::from(&reconstructed).to_string(),
+        );
+    }
+
+    #[tokio::test]
+    async fn test_expr_id() {
+        let source_schema =
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+        let col_a = col("a", &source_schema).unwrap();
+
+        // Create a source filter
+        let source = Arc::new(DynamicFilterPhysicalExpr::new(
+            vec![Arc::clone(&col_a)],
+            lit(true) as Arc<dyn PhysicalExpr>,
+        ));
+        let source_clone = Arc::clone(&source);
+
+        // Create a derived filter by reassigning the source filter to a different schema.
+        let derived_schema = Arc::new(Schema::new(vec![
+            Field::new("x", DataType::Int32, false),
+            Field::new("a", DataType::Int32, false),
+        ]));
+        let derived = reassign_expr_columns(
+            Arc::clone(&source) as Arc<dyn PhysicalExpr>,
+            &derived_schema,
+        )
+        .expect("reassign_expr_columns should succeed");
+
+        let derived_expr_id = Arc::clone(&derived)
+            .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");
+        let source_clone_expr_id = Arc::clone(&source_clone)
+            .expr_id(&[])
+            .expect("source clone should have an expr_id");
+
+        assert_eq!(
+            source_clone_expr_id.exact(),
+            source_expr_id.exact(),
+            "cloned filter should have the same exact id because the children are the same",
+        );
+
+        assert_eq!(
+            source_clone_expr_id.shallow(),
+            source_expr_id.shallow(),
+            "cloned filter should have the same shallow id because the exprs are the same",
+        );
+
+        assert_ne!(
+            derived_expr_id.exact(),
+            source_expr_id.exact(),
+            "filters should have different exact ids because the children are different",
+        );
+
+        assert_eq!(
+            derived_expr_id.shallow(),
+            source_expr_id.shallow(),
+            "filters should have the same shallow id because the exprs are the same",
+        );
+    }
 }

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

@@ -870,18 +870,18 @@ message PhysicalExtensionNode {
   repeated PhysicalPlanNode inputs = 2;
 }
 
-// physical expressions
 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.
+  // Unique identifiers for this expression used during deserialization to restore
+  // referential integrity across serde roundtrips.
+  // 
+  // expr_id: if two exprs have the same expr_id, they are identical (including children)
+  // shallow_expr_id: if two exprs have the same shallow_expr_id, they are identical but may 
+  //   have different children
   optional uint64 expr_id = 30;
+  optional uint64 shallow_expr_id = 31;
 
   oneof ExprType {
     // column references
@@ -920,9 +920,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 +1480,4 @@ message AsyncFuncExecNode {
 message BufferExecNode {
   PhysicalPlanNode input = 1;
   uint64 capacity = 2;
-}
+}