proto: serialize and dedupe dynamic filters

Informs: datafusion-contrib/datafusion-distributed#180
Closes: apache#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/src/expressions/dynamic_filters.rs

@@ -23,7 +23,7 @@ use crate::PhysicalExpr;
 use crate::expressions::lit;
 use arrow::datatypes::{DataType, Schema};
 use datafusion_common::{
-    Result,
+    Result, internal_err,
     tree_node::{Transformed, TransformedResult, TreeNode},
 };
 use datafusion_expr::ColumnarValue;
@@ -136,6 +136,117 @@ struct Inner {
     partitioned_exprs: PartitionedFilters,
 }
 
+/// 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 {
@@ -232,18 +343,52 @@ impl DynamicFilterPhysicalExpr {
         }
     }
 
+    /// Create a new [`DynamicFilterPhysicalExpr`] from `self`, except it overwrites the
+    /// internal state with the source filter's state.
+    ///
+    /// This is a low-level API intended for use by the proto deserialization layer.
+    ///
+    /// # Safety
+    ///
+    /// The dynamic filter should not be in use when calling this method, otherwise there
+    /// 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: &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() {
+            return internal_err!(
+                "Cannot replace the inner state of a DynamicFilterPhysicalExpr that is in use"
+            );
+        };
+
+        Ok(Self {
+            children: self.children.clone(),
+            remapped_children: self.remapped_children.clone(),
+            inner: Arc::clone(&source.inner),
+            state_watch: self.state_watch.clone(),
+            data_type: Arc::clone(&self.data_type),
+            nullable: Arc::clone(&self.nullable),
+        })
+    }
+
     fn remap_children(
-        children: &[Arc<dyn PhysicalExpr>],
-        remapped_children: Option<&Vec<Arc<dyn PhysicalExpr>>>,
+        &self,
         expr: Arc<dyn PhysicalExpr>,
     ) -> Result<Arc<dyn PhysicalExpr>> {
-        if let Some(remapped_children) = remapped_children {
+        if let Some(remapped_children) = &self.remapped_children {
             // Remap the children to the new children
             // of the expression.
             expr.transform_up(|child| {
                 // Check if this is any of our original children
-                if let Some(pos) =
-                    children.iter().position(|c| c.as_ref() == child.as_ref())
+                if let Some(pos) = self
+                    .children
+                    .iter()
+                    .position(|c| c.as_ref() == child.as_ref())
                 {
                     // If so, remap it to the current children
                     // of the expression.
@@ -433,6 +578,14 @@ impl DynamicFilterPhysicalExpr {
         Arc::strong_count(self) > 1 || Arc::strong_count(&self.inner) > 1
     }
 
+    /// Returns a unique identifier for the inner shared state.
+    ///
+    /// Useful for checking if two [`Arc<PhysicalExpr>`] with the same
+    /// underlying [`DynamicFilterPhysicalExpr`] are the same.
+    pub fn inner_id(&self) -> u64 {
+        Arc::as_ptr(&self.inner) as *const () as u64
+    }
+
     fn render(
         &self,
         f: &mut std::fmt::Formatter<'_>,
@@ -1154,7 +1307,6 @@ mod test {
         );
     }
 
-    #[test]
     fn test_bind_dynamic_filters_for_partition_without_partitioned_data() {
         let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(
             vec![],
@@ -1283,4 +1435,59 @@ mod test {
             "Bound dynamic filter should route to partition-local payload after update"
         );
     }
+
+    #[test]
+    fn test_new_from_source() {
+        // Create a source filter
+        let source = Arc::new(DynamicFilterPhysicalExpr::new(
+            vec![],
+            lit(42) as Arc<dyn PhysicalExpr>,
+        ));
+
+        // Update and mark complete
+        source
+            .update(DynamicFilterUpdate::Global(
+                lit(100) as Arc<dyn PhysicalExpr>,
+            ))
+            .unwrap();
+        source.mark_complete();
+
+        // 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 = 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();
+
+        // Verify inner state is shared (same inner_id)
+        assert_eq!(
+            combined.inner_id(),
+            source.inner_id(),
+            "new_from_source should share inner state with source"
+        );
+
+        // Verify children are from target, not source
+        let combined_snapshot = DynamicFilterSnapshot::from(&combined);
+        assert_eq!(
+            combined_snapshot.children().len(),
+            1,
+            "Combined filter should have target's children"
+        );
+        assert_eq!(
+            format!("{:?}", combined_snapshot.children()[0]),
+            format!("{:?}", col_x),
+            "Combined filter should have target's children"
+        );
+
+        // Verify inner expression comes from source
+        assert_eq!(
+            format!("{:?}", combined_snapshot.inner_expr()),
+            format!("{:?}", lit(100)),
+            "Combined filter should have source's inner expression"
+        );
+    }
 }

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

@@ -47,6 +47,7 @@ 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::DynamicFilterRuntimeContext;
+pub use dynamic_filters::DynamicFilterSnapshot;
 pub use dynamic_filters::DynamicFilterUpdate;
 pub use in_list::{InListExpr, in_list};
 pub use is_not_null::{IsNotNullExpr, is_not_null};

datafusion/proto/proto/datafusion.proto

@@ -845,6 +845,12 @@ message PhysicalExprNode {
   // across serde roundtrips.
   optional uint64 expr_id = 30;
 
+  // For DynamicFilterPhysicalExpr, this identifies the shared inner state.
+  // Multiple expressions may have different expr_id values (different outer Arc wrappers)
+  // but the same dynamic_filter_inner_id (shared inner state).
+  // Used to reconstruct shared inner state during deserialization.
+  optional uint64 dynamic_filter_inner_id = 31;
+
   oneof ExprType {
     // column references
     PhysicalColumn column = 1;
@@ -882,9 +888,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;