[GitHub] [arrow-datafusion] tustvold commented on a diff in pull request #2670: MINOR: Improve `datafusion-proto` API

[GitBox <gi...@apache.org>]

tustvold commented on code in PR #2670:
URL: https://github.com/apache/arrow-datafusion/pull/2670#discussion_r886658567


##########
datafusion/proto/src/lib.rs:
##########
@@ -33,6 +42,122 @@ pub mod to_proto;
 #[cfg(doctest)]
 doc_comment::doctest!("../README.md", readme_example_test);
 
+/// Serialization for `LogicalPlan` and `Expr`
+#[derive(Default)]
+pub struct Serializer<'a> {
+    /// Optional function registry
+    function_registry: Option<&'a dyn FunctionRegistry>,
+    /// Optional extension codec
+    extension_codec: Option<&'a dyn LogicalExtensionCodec>,
+}
+
+impl<'a> Serializer<'a> {
+    /// Create a Serializer
+    pub fn new() -> Self {
+        Self {
+            function_registry: None,
+            extension_codec: None,
+        }
+    }
+
+    /// Create a new Serializer with the provided function registry
+    pub fn with_function_registry(
+        self,
+        function_registry: &'a dyn FunctionRegistry,
+    ) -> Self {
+        Self {
+            function_registry: Some(function_registry),
+            extension_codec: self.extension_codec,
+        }
+    }
+
+    /// Create a new Serializer with the provided function registry
+    pub fn with_extension_codec(
+        self,
+        extension_codec: &'a dyn LogicalExtensionCodec,
+    ) -> Self {
+        Self {
+            function_registry: self.function_registry,
+            extension_codec: Some(extension_codec),
+        }
+    }
+
+    /// Serialize a logical expression
+    pub fn serialize_expr(&self, expr: &Expr) -> Result<Vec<u8>> {
+        expr.to_bytes().map(|b| b.to_vec())
+    }
+
+    /// Deserialize a logical expression
+    pub fn deserialize_expr(&self, bytes: &[u8]) -> Result<Expr> {
+        //TODO note that we will need to also pass the extension_codec along once we
+        // add support for subquery expressions in https://github.com/apache/arrow-datafusion/issues/2640
+        match self.function_registry {
+            Some(r) => Serializeable::from_bytes_with_registry(bytes, r),
+            _ => Serializeable::from_bytes(bytes),
+        }
+    }
+
+    /// Serialize a logical plan
+    pub fn serialize_plan(&self, plan: &LogicalPlan) -> Result<Vec<u8>> {
+        //TODO note that we will need to also pass the function_registry along once we
+        // add support for subquery expressions in https://github.com/apache/arrow-datafusion/issues/2640
+        let protobuf = match self.extension_codec {
+            Some(codec) => protobuf::LogicalPlanNode::try_from_logical_plan(plan, codec)?,

Review Comment:
   Is your plan to move this logic onto Serializer or to pass function_registry down?



##########
datafusion/proto/src/logical_plan.rs:
##########
@@ -69,7 +69,10 @@ pub(crate) fn proto_error<S: Into<String>>(message: S) -> DataFusionError {
     DataFusionError::Internal(message.into())
 }
 
-pub trait AsLogicalPlan: Debug + Send + Sync + Clone {
+// this trait was inherited from Ballista and was intended to support custom logical plans
+// but that does not make sense in the context of this crate so would be good to remove
+// this at some point

Review Comment:
   ```suggestion
   /// This trait was inherited from Ballista and was intended to support custom logical plans
   /// but that does not make sense in the context of this crate so would be good to remove
   /// this at some point
   ```
   
   Perhaps a ticket also?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org