Support serialization/deserialization for custom physical exprs in proto (#11387)

* Add PhysicalExtensionExprNode

* regen proto

* Add ser/de extension expr logic

* Add test and fix clippy lint

Author: lewiszlw

datafusion/proto/proto/datafusion.proto

@@ -836,6 +836,8 @@ message PhysicalExprNode {
     // was PhysicalDateTimeIntervalExprNode date_time_interval_expr = 17;
 
     PhysicalLikeExprNode like_expr = 18;
+
+    PhysicalExtensionExprNode extension = 19;
   }
 }
 
@@ -942,6 +944,11 @@ message PhysicalNegativeNode {
   PhysicalExprNode expr = 1;
 }
 
+message PhysicalExtensionExprNode {
+  bytes expr = 1;
+  repeated PhysicalExprNode inputs = 2;
+}
+
 message FilterExecNode {
   PhysicalPlanNode input = 1;
   PhysicalExprNode expr = 2;

datafusion/proto/src/generated/pbjson.rs

@@ -394,6 +394,14 @@ pub fn parse_physical_expr(
                 codec,
             )?,
         )),
+        ExprType::Extension(extension) => {
+            let inputs: Vec<Arc<dyn PhysicalExpr>> = extension
+                .inputs
+                .iter()
+                .map(|e| parse_physical_expr(e, registry, input_schema, codec))
+                .collect::<Result<_>>()?;
+            (codec.try_decode_expr(extension.expr.as_slice(), &inputs)?) as _
+        }
     };
 
     Ok(pexpr)

datafusion/proto/src/generated/prost.rs

@@ -2018,6 +2018,22 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync {
     fn try_encode_udf(&self, _node: &ScalarUDF, _buf: &mut Vec<u8>) -> Result<()> {
         Ok(())
     }
+
+    fn try_decode_expr(
+        &self,
+        _buf: &[u8],
+        _inputs: &[Arc<dyn PhysicalExpr>],
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        not_impl_err!("PhysicalExtensionCodec is not provided")
+    }
+
+    fn try_encode_expr(
+        &self,
+        _node: Arc<dyn PhysicalExpr>,
+        _buf: &mut Vec<u8>,
+    ) -> Result<()> {
+        not_impl_err!("PhysicalExtensionCodec is not provided")
+    }
 }
 
 #[derive(Debug)]

datafusion/proto/src/physical_plan/from_proto.rs

@@ -495,7 +495,24 @@ pub fn serialize_physical_expr(
             ))),
         })
     } else {
-        internal_err!("physical_plan::to_proto() unsupported expression {value:?}")
+        let mut buf: Vec<u8> = vec![];
+        match codec.try_encode_expr(Arc::clone(&value), &mut buf) {
+            Ok(_) => {
+                let inputs: Vec<protobuf::PhysicalExprNode> = value
+                    .children()
+                    .into_iter()
+                    .map(|e| serialize_physical_expr(Arc::clone(e), codec))
+                    .collect::<Result<_>>()?;
+                Ok(protobuf::PhysicalExprNode {
+                    expr_type: Some(protobuf::physical_expr_node::ExprType::Extension(
+                        protobuf::PhysicalExtensionExprNode { expr: buf, inputs },
+                    )),
+                })
+            }
+            Err(e) => internal_err!(
+                "Unsupported physical expr and extension codec failed with [{e}]. Expr: {value:?}"
+            ),
+        }
     }
 }