Add conversion from physical expr to proto in converter. Doing some renaming

Author: timsaucer

datafusion-examples/examples/custom_data_source/adapter_serialization.rs

@@ -37,8 +37,7 @@ use std::sync::Arc;
 use arrow::array::record_batch;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion::assert_batches_eq;
-use datafusion::common::Result;
-use datafusion::common::not_impl_err;
+use datafusion::common::{Result, not_impl_err};
 use datafusion::datasource::listing::{
     ListingTable, ListingTableConfig, ListingTableConfigExt, ListingTableUrl,
 };
@@ -55,15 +54,17 @@ use datafusion_physical_expr_adapter::{
     DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory,
 };
 use datafusion_proto::bytes::{
-    physical_plan_from_bytes_with_extension_codec,
-    physical_plan_to_bytes_with_extension_codec,
+    physical_plan_from_bytes_with_proto_converter,
+    physical_plan_to_bytes_with_proto_converter,
 };
+use datafusion_proto::physical_plan::from_proto::proto_to_physical_expr;
+use datafusion_proto::physical_plan::to_proto::serialize_physical_expr;
 use datafusion_proto::physical_plan::{
-    AsExecutionPlan, PhysicalExtensionCodec, PhysicalExtensionProtoCodec,
+    AsExecutionPlan, PhysicalExtensionCodec, PhysicalProtoConverterExtension,
 };
+use datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType;
 use datafusion_proto::protobuf::{
     PhysicalExprNode, PhysicalExtensionNode, PhysicalPlanNode,
-    physical_plan_node::PhysicalPlanType,
 };
 use object_store::memory::InMemory;
 use object_store::path::Path;
@@ -126,7 +127,7 @@ pub async fn adapter_serialization() -> Result<()> {
     // Step 4: Serialize with our custom codec
     println!("\nStep 4: Serializing plan with AdapterPreservingCodec...");
     let codec = AdapterPreservingCodec;
-    let bytes = physical_plan_to_bytes_with_extension_codec(
+    let bytes = physical_plan_to_bytes_with_proto_converter(
         Arc::clone(&original_plan),
         &codec,
         &codec,
@@ -138,7 +139,7 @@ pub async fn adapter_serialization() -> Result<()> {
     println!("\nStep 5: Deserializing plan with AdapterPreservingCodec...");
     let task_ctx = ctx.task_ctx();
     let restored_plan =
-        physical_plan_from_bytes_with_extension_codec(&bytes, &task_ctx, &codec, &codec)?;
+        physical_plan_from_bytes_with_proto_converter(&bytes, &task_ctx, &codec, &codec)?;
 
     // Verify adapter is restored
     let has_adapter_after = verify_adapter_in_plan(&restored_plan, "restored");
@@ -314,10 +315,10 @@ impl PhysicalExtensionCodec for AdapterPreservingCodec {
     }
 }
 
-impl PhysicalExtensionProtoCodec for AdapterPreservingCodec {
+impl PhysicalProtoConverterExtension for AdapterPreservingCodec {
     fn execution_plan_to_proto(
         &self,
-        plan: Arc<dyn ExecutionPlan>,
+        plan: &Arc<dyn ExecutionPlan>,
         extension_codec: &dyn PhysicalExtensionCodec,
     ) -> Result<PhysicalPlanNode> {
         // Check if this is a DataSourceExec with adapter
@@ -378,7 +379,7 @@ impl PhysicalExtensionProtoCodec for AdapterPreservingCodec {
         }
 
         // No adapter found - use default serialization
-        PhysicalPlanNode::try_from_physical_plan(plan, extension_codec, self)
+        PhysicalPlanNode::try_from_physical_plan(Arc::clone(plan), extension_codec, self)
     }
 
     // Interception point: override deserialization to unwrap adapters
@@ -424,12 +425,20 @@ impl PhysicalExtensionProtoCodec for AdapterPreservingCodec {
 
     fn proto_to_physical_expr(
         &self,
-        _proto: &PhysicalExprNode,
-        _ctx: &TaskContext,
-        _input_schema: &Schema,
-        _codec: &dyn PhysicalExtensionCodec,
+        proto: &PhysicalExprNode,
+        ctx: &TaskContext,
+        input_schema: &Schema,
+        codec: &dyn PhysicalExtensionCodec,
     ) -> Result<Arc<dyn PhysicalExpr>> {
-        todo!()
+        proto_to_physical_expr(proto, ctx, input_schema, codec, self)
+    }
+
+    fn physical_expr_to_proto(
+        &self,
+        expr: &Arc<dyn PhysicalExpr>,
+        codec: &dyn PhysicalExtensionCodec,
+    ) -> Result<PhysicalExprNode> {
+        serialize_physical_expr(expr, codec, self)
     }
 }
 

datafusion-examples/examples/proto/expression_deduplication.rs

@@ -47,9 +47,10 @@ use datafusion::physical_plan::filter::FilterExec;
 use datafusion::physical_plan::placeholder_row::PlaceholderRowExec;
 use datafusion::prelude::SessionContext;
 use datafusion_proto::physical_plan::from_proto::proto_to_physical_expr;
+use datafusion_proto::physical_plan::to_proto::serialize_physical_expr;
 use datafusion_proto::physical_plan::{
     AsExecutionPlan, DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
-    PhysicalExtensionProtoCodec,
+    PhysicalProtoConverterExtension,
 };
 use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode};
 use prost::Message;
@@ -103,13 +104,7 @@ pub async fn expression_deduplication() -> Result<()> {
 
     let extension_codec = DefaultPhysicalExtensionCodec {};
     let caching_codec = CachingCodec::new();
-    // let proto = PhysicalPlanNode::try_from_physical_plan(
-    //     filter_plan.clone(),
-    //     &extension_codec,
-    //     &caching_codec,
-    // )?;
-    let proto =
-        caching_codec.execution_plan_to_proto(filter_plan.clone(), &extension_codec)?;
+    let proto = caching_codec.execution_plan_to_proto(&filter_plan, &extension_codec)?;
 
     // Serialize to bytes
     let mut bytes = Vec::new();
@@ -205,7 +200,7 @@ impl PhysicalExtensionCodec for CachingCodec {
     }
 }
 
-impl PhysicalExtensionProtoCodec for CachingCodec {
+impl PhysicalProtoConverterExtension for CachingCodec {
     fn proto_to_execution_plan(
         &self,
         _ctx: &TaskContext,
@@ -217,10 +212,10 @@ impl PhysicalExtensionProtoCodec for CachingCodec {
 
     fn execution_plan_to_proto(
         &self,
-        plan: Arc<dyn ExecutionPlan>,
+        plan: &Arc<dyn ExecutionPlan>,
         extension_codec: &dyn PhysicalExtensionCodec,
     ) -> Result<PhysicalPlanNode> {
-        PhysicalPlanNode::try_from_physical_plan(plan, extension_codec, self)
+        PhysicalPlanNode::try_from_physical_plan(Arc::clone(plan), extension_codec, self)
     }
 
     // CACHING IMPLEMENTATION: Intercept expression deserialization
@@ -263,4 +258,12 @@ impl PhysicalExtensionProtoCodec for CachingCodec {
 
         Ok(expr)
     }
+
+    fn physical_expr_to_proto(
+        &self,
+        expr: &Arc<dyn PhysicalExpr>,
+        codec: &dyn PhysicalExtensionCodec,
+    ) -> Result<PhysicalExprNode> {
+        serialize_physical_expr(expr, codec, self)
+    }
 }

datafusion/proto/src/bytes/mod.rs

@@ -22,7 +22,7 @@ use crate::logical_plan::{
 };
 use crate::physical_plan::{
     AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalExtensionProtoCodec,
-    PhysicalExtensionCodec, PhysicalExtensionProtoCodec,
+    PhysicalExtensionCodec, PhysicalProtoConverterExtension,
 };
 use crate::protobuf;
 use datafusion_common::{Result, plan_datafusion_err};
@@ -90,8 +90,8 @@ pub trait Serializeable: Sized {
 impl Serializeable for Expr {
     fn to_bytes(&self) -> Result<Bytes> {
         let mut buffer = BytesMut::new();
-        let extension_codec = DefaultLogicalExtensionCodec {};
-        let protobuf: protobuf::LogicalExprNode = serialize_expr(self, &extension_codec)
+        let codec = DefaultLogicalExtensionCodec {};
+        let protobuf: protobuf::LogicalExprNode = serialize_expr(self, &codec)
             .map_err(|e| plan_datafusion_err!("Error encoding expr as protobuf: {e}"))?;
 
         protobuf
@@ -192,8 +192,8 @@ impl Serializeable for Expr {
         let protobuf = protobuf::LogicalExprNode::decode(bytes)
             .map_err(|e| plan_datafusion_err!("Error decoding expr as protobuf: {e}"))?;
 
-        let extension_codec = DefaultLogicalExtensionCodec {};
-        logical_plan::from_proto::parse_expr(&protobuf, registry, &extension_codec)
+        let codec = DefaultLogicalExtensionCodec {};
+        logical_plan::from_proto::parse_expr(&protobuf, registry, &codec)
             .map_err(|e| plan_datafusion_err!("Error parsing protobuf into Expr: {e}"))
     }
 }
@@ -277,17 +277,20 @@ pub fn logical_plan_from_json_with_extension_codec(
 /// Serialize a PhysicalPlan as bytes
 pub fn physical_plan_to_bytes(plan: Arc<dyn ExecutionPlan>) -> Result<Bytes> {
     let extension_codec = DefaultPhysicalExtensionCodec {};
-    let proto_codec = DefaultPhysicalExtensionProtoCodec {};
-    physical_plan_to_bytes_with_extension_codec(plan, &extension_codec, &proto_codec)
+    physical_plan_to_bytes_with_extension_codec(plan, &extension_codec)
 }
 
 /// Serialize a PhysicalPlan as JSON
 #[cfg(feature = "json")]
 pub fn physical_plan_to_json(plan: Arc<dyn ExecutionPlan>) -> Result<String> {
     let extension_codec = DefaultPhysicalExtensionCodec {};
-    let protobuf =
-        protobuf::PhysicalPlanNode::try_from_physical_plan(plan, &extension_codec)
-            .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))?;
+    let proto_converter = DefaultPhysicalExtensionProtoCodec {};
+    let protobuf = protobuf::PhysicalPlanNode::try_from_physical_plan(
+        plan,
+        &extension_codec,
+        &proto_converter,
+    )
+    .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))?;
     serde_json::to_string(&protobuf)
         .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))
 }
@@ -296,9 +299,19 @@ pub fn physical_plan_to_json(plan: Arc<dyn ExecutionPlan>) -> Result<String> {
 pub fn physical_plan_to_bytes_with_extension_codec(
     plan: Arc<dyn ExecutionPlan>,
     extension_codec: &dyn PhysicalExtensionCodec,
-    proto_codec: &dyn PhysicalExtensionProtoCodec,
 ) -> Result<Bytes> {
-    let protobuf = proto_codec.execution_plan_to_proto(plan, extension_codec)?;
+    let proto_conveter = DefaultPhysicalExtensionProtoCodec {};
+    physical_plan_to_bytes_with_proto_converter(plan, extension_codec, &proto_conveter)
+}
+
+/// Serialize a PhysicalPlan as bytes, using the provided extension codec
+/// and protobuf converter.
+pub fn physical_plan_to_bytes_with_proto_converter(
+    plan: Arc<dyn ExecutionPlan>,
+    extension_codec: &dyn PhysicalExtensionCodec,
+    proto_converter: &dyn PhysicalProtoConverterExtension,
+) -> Result<Bytes> {
+    let protobuf = proto_converter.execution_plan_to_proto(&plan, extension_codec)?;
     let mut buffer = BytesMut::new();
     protobuf
         .encode(&mut buffer)
@@ -315,7 +328,8 @@ pub fn physical_plan_from_json(
     let back: protobuf::PhysicalPlanNode = serde_json::from_str(json)
         .map_err(|e| plan_datafusion_err!("Error serializing plan: {e}"))?;
     let extension_codec = DefaultPhysicalExtensionCodec {};
-    back.try_into_physical_plan(ctx, &extension_codec)
+    let proto_converter = DefaultPhysicalExtensionProtoCodec {};
+    back.try_into_physical_plan(ctx, &extension_codec, &proto_converter)
 }
 
 /// Deserialize a PhysicalPlan from bytes
@@ -324,12 +338,12 @@ pub fn physical_plan_from_bytes(
     ctx: &TaskContext,
 ) -> Result<Arc<dyn ExecutionPlan>> {
     let extension_codec = DefaultPhysicalExtensionCodec {};
-    let proto_codec = DefaultPhysicalExtensionProtoCodec {};
-    physical_plan_from_bytes_with_extension_codec(
+    let proto_converter = DefaultPhysicalExtensionProtoCodec {};
+    physical_plan_from_bytes_with_proto_converter(
         bytes,
         ctx,
         &extension_codec,
-        &proto_codec,
+        &proto_converter,
     )
 }
 
@@ -338,9 +352,24 @@ pub fn physical_plan_from_bytes_with_extension_codec(
     bytes: &[u8],
     ctx: &TaskContext,
     extension_codec: &dyn PhysicalExtensionCodec,
-    proto_codec: &dyn PhysicalExtensionProtoCodec,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let proto_converter = DefaultPhysicalExtensionProtoCodec {};
+    physical_plan_from_bytes_with_proto_converter(
+        bytes,
+        ctx,
+        extension_codec,
+        &proto_converter,
+    )
+}
+
+/// Deserialize a PhysicalPlan from bytes
+pub fn physical_plan_from_bytes_with_proto_converter(
+    bytes: &[u8],
+    ctx: &TaskContext,
+    extension_codec: &dyn PhysicalExtensionCodec,
+    proto_converter: &dyn PhysicalProtoConverterExtension,
 ) -> Result<Arc<dyn ExecutionPlan>> {
     let protobuf = protobuf::PhysicalPlanNode::decode(bytes)
         .map_err(|e| plan_datafusion_err!("Error decoding expr as protobuf: {e}"))?;
-    protobuf.try_into_physical_plan(ctx, extension_codec, proto_codec)
+    protobuf.try_into_physical_plan(ctx, extension_codec, proto_converter)
 }