restore snapshot method

Author: jayshrivastava

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

@@ -308,7 +308,7 @@ impl 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 {
+        if self.is_used() {
             return internal_err!(
                 "Cannot replace the inner state of a DynamicFilterPhysicalExpr that is in use"
             );
@@ -587,6 +587,16 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr {
     fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         self.render(f, |expr, f| expr.fmt_sql(f))
     }
+
+    fn snapshot(&self) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+        // Return the current expression as a snapshot.
+        Ok(Some(self.current()?))
+    }
+
+    fn snapshot_generation(&self) -> u64 {
+        // Return the current generation of the expression.
+        self.inner.read().generation
+    }
 }
 
 #[cfg(test)]
@@ -716,6 +726,23 @@ mod test {
         assert!(arr_1.eq(&expected));
     }
 
+    #[test]
+    fn test_snapshot() {
+        let expr = lit(42) as Arc<dyn PhysicalExpr>;
+        let dynamic_filter = DynamicFilterPhysicalExpr::new(vec![], Arc::clone(&expr));
+
+        // Take a snapshot of the current expression
+        let snapshot = dynamic_filter.snapshot().unwrap();
+        assert_eq!(snapshot, Some(expr));
+
+        // Update the current expression
+        let new_expr = lit(100) as Arc<dyn PhysicalExpr>;
+        dynamic_filter.update(Arc::clone(&new_expr)).unwrap();
+        // Take another snapshot
+        let snapshot = dynamic_filter.snapshot().unwrap();
+        assert_eq!(snapshot, Some(new_expr));
+    }
+
     #[test]
     fn test_dynamic_filter_physical_expr_misbehaves_data_type_nullable() {
         let dynamic_filter =

datafusion/proto/src/physical_plan/to_proto.rs

@@ -258,6 +258,47 @@ pub fn serialize_physical_expr_with_converter(
     codec: &dyn PhysicalExtensionCodec,
     proto_converter: &dyn PhysicalProtoConverterExtension,
 ) -> Result<protobuf::PhysicalExprNode> {
+    // Check for DynamicFilterPhysicalExpr before snapshotting.
+    // We need to handle it before snapshot_physical_expr because snapshot()
+    // replaces the DynamicFilterPhysicalExpr with its inner expression.
+    if let Some(df) = value.as_any().downcast_ref::<DynamicFilterPhysicalExpr>() {
+        // Capture all state atomically
+        let snapshot = DynamicFilterSnapshot::from(df);
+
+        let children = snapshot
+            .children()
+            .iter()
+            .map(|child| proto_converter.physical_expr_to_proto(child, codec))
+            .collect::<Result<Vec<_>>>()?;
+
+        let remapped_children = if let Some(remapped) = snapshot.remapped_children() {
+            remapped
+                .iter()
+                .map(|child| proto_converter.physical_expr_to_proto(child, codec))
+                .collect::<Result<Vec<_>>>()?
+        } else {
+            vec![]
+        };
+
+        let inner_expr = Box::new(
+            proto_converter.physical_expr_to_proto(snapshot.inner_expr(), codec)?,
+        );
+
+        return Ok(protobuf::PhysicalExprNode {
+            expr_id: None,
+            dynamic_filter_inner_id: None,
+            expr_type: Some(protobuf::physical_expr_node::ExprType::DynamicFilter(
+                Box::new(protobuf::PhysicalDynamicFilterNode {
+                    children,
+                    remapped_children,
+                    generation: snapshot.generation(),
+                    inner_expr: Some(inner_expr),
+                    is_complete: snapshot.is_complete(),
+                }),
+            )),
+        });
+    }
+
     // Snapshot the expr in case it has dynamic predicate state so
     // it can be serialized
     let value = snapshot_physical_expr(Arc::clone(value))?;
@@ -328,42 +369,6 @@ pub fn serialize_physical_expr_with_converter(
                 binary_expr,
             )),
         })
-    } else if let Some(df) = expr.downcast_ref::<DynamicFilterPhysicalExpr>() {
-        // Capture all state atomically
-        let snapshot = DynamicFilterSnapshot::from(df);
-
-        let children = snapshot
-            .children()
-            .iter()
-            .map(|child| proto_converter.physical_expr_to_proto(child, codec))
-            .collect::<Result<Vec<_>>>()?;
-
-        let remapped_children = if let Some(remapped) = snapshot.remapped_children() {
-            remapped
-                .iter()
-                .map(|child| proto_converter.physical_expr_to_proto(child, codec))
-                .collect::<Result<Vec<_>>>()?
-        } else {
-            vec![]
-        };
-
-        let inner_expr = Box::new(
-            proto_converter.physical_expr_to_proto(snapshot.inner_expr(), codec)?,
-        );
-
-        Ok(protobuf::PhysicalExprNode {
-            expr_id: None,
-            dynamic_filter_inner_id: None,
-            expr_type: Some(protobuf::physical_expr_node::ExprType::DynamicFilter(
-                Box::new(protobuf::PhysicalDynamicFilterNode {
-                    children,
-                    remapped_children,
-                    generation: snapshot.generation(),
-                    inner_expr: Some(inner_expr),
-                    is_complete: snapshot.is_complete(),
-                }),
-            )),
-        })
     } else if let Some(expr) = expr.downcast_ref::<CaseExpr>() {
         Ok(protobuf::PhysicalExprNode {
             expr_id: None,