Handle serde for ScalarUDF #9395

yyy1000 · 2024-02-29T02:23:15Z

Which issue does this PR close?

Closes #8706 .

Rationale for this change

See the #8706

What changes are included in this PR?

Add a new field in ScalarUDFExprNode in proto
Extend LogicalExtensionCodec to have methods like try_decode_udf and try_encode_udf
Add a serialize_expr function which accept a LogicalExtensionCodec param to serialize Expr
In parse_expr, add a LogicalExtensionCodec param to serialize Expr to deserialize

Are these changes tested?

Are there any user-facing changes?

yyy1000

Hi @thinkharderdev
I want to help #8706 in this PR but I have some questions.
I have put them in this review :)

yyy1000 · 2024-02-29T02:25:42Z

datafusion/proto/src/bytes/mod.rs

    fn to_bytes(&self) -> Result<Bytes> {
        let mut buffer = BytesMut::new();
-        let protobuf: protobuf::LogicalExprNode = self
-            .try_into()
+        let extension_codec = DefaultLogicalExtensionCodec {};
+        let protobuf: protobuf::LogicalExprNode = serialize_expr(self, &extension_codec)


Would this be OK?
I think DefaultLogicalExtensionCodec may need to implement some functions it didn't implement yet. I will also look into it.

Yes I think this is fine. These extension methods can be used for vanilla serialization and then if user needs custom stuff with a custom codec they can use the lower-level apis

yyy1000 · 2024-02-29T02:26:24Z

datafusion/proto/src/bytes/mod.rs

@@ -177,7 +178,8 @@ impl Serializeable for Expr {
        let protobuf = protobuf::LogicalExprNode::decode(bytes)
            .map_err(|e| plan_datafusion_err!("Error decoding expr as protobuf: {e}"))?;

-        logical_plan::from_proto::parse_expr(&protobuf, registry)
+        let extension_codec = DefaultLogicalExtensionCodec {};
+        logical_plan::from_proto::parse_expr(&protobuf, registry, &extension_codec)


And also here

yyy1000 · 2024-02-29T02:27:01Z

datafusion/proto/src/logical_plan/mod.rs

+
+    fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>>;
+
+    fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()>;


I plan to add aggregate_udf and window_udf after scalar_udf.

Here I would include default implementations for all the new methods so it doesn't break backwards compatibility:

fn try_encode_udf(&self, _node: &ScalarUDF, _buf: &mut Vec<u8>) -> Result<()> { Ok(()) }

yyy1000 · 2024-02-29T02:30:58Z

datafusion/proto/src/physical_plan/mod.rs

+    //     buf: &[u8],
+    // ) -> Result<Arc<WindowUDF>>;
+
+    // fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()>;


Plan to finish parts related to PhysicalExtensionCodec after LogicalExtensionCodec

yyy1000 · 2024-02-29T02:32:39Z

datafusion/proto/src/logical_plan/from_proto.rs

@@ -999,7 +1002,7 @@ pub fn parse_expr(
            let operands = binary_expr
                .operands
                .iter()
-                .map(|expr| parse_expr(expr, registry))
+                .map(|expr| parse_expr(expr, registry, codec))


In from_proto.rs file, a lot of places need to change like this line.
I haven't finish it but submit this draft PR to verify the road.

thinkharderdev

Looking good

thinkharderdev · 2024-02-29T11:11:31Z

datafusion/proto/src/logical_plan/mod.rs

+
+    fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>>;
+
+    fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()>;


Here I would include default implementations for all the new methods so it doesn't break backwards compatibility:

fn try_encode_udf(&self, _node: &ScalarUDF, _buf: &mut Vec<u8>) -> Result<()> { Ok(()) }

thinkharderdev · 2024-02-29T11:20:46Z

datafusion/proto/src/logical_plan/mod.rs

+    fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+        not_impl_err!("LogicalExtensionCodec is not provided")
+    }
+
+    fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
+        not_impl_err!("LogicalExtensionCodec is not provided")
+    }


These shouldn't be neccessary if we add the default impl in the trait

thinkharderdev · 2024-02-29T11:21:03Z

datafusion/proto/src/logical_plan/mod.rs

@@ -133,6 +136,10 @@ pub trait LogicalExtensionCodec: Debug + Send + Sync {
        node: Arc<dyn TableProvider>,
        buf: &mut Vec<u8>,
    ) -> Result<()>;
+
+    fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>>;


Suggested change

fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>>;

fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result<Arc<ScalarUDF>> {

not_impl_err!("LogicalExtensionCodec is not provided for scalar function {name}")

}

thinkharderdev · 2024-02-29T11:25:28Z

datafusion/proto/src/bytes/mod.rs

    fn to_bytes(&self) -> Result<Bytes> {
        let mut buffer = BytesMut::new();
-        let protobuf: protobuf::LogicalExprNode = self
-            .try_into()
+        let extension_codec = DefaultLogicalExtensionCodec {};
+        let protobuf: protobuf::LogicalExprNode = serialize_expr(self, &extension_codec)


Yes I think this is fine. These extension methods can be used for vanilla serialization and then if user needs custom stuff with a custom codec they can use the lower-level apis

yyy1000

Hi @thinkharderdev
Thanks for your review and I think this can pass all CIs and is close to merge.
A question I have is when designing the test case, I put the details below.

yyy1000 · 2024-03-01T05:38:02Z

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

+    pub struct ScalarUDFProto {
+        #[prost(message, tag = "1")]
+        pub expr: ::core::option::Option<datafusion_proto::protobuf::LogicalExprNode>,
+    }


Here is my test proto

This is not how this would work. Expressions would be passed to the udf as arguments. What would potentially get serialized into ScalarUDFProto would just be some non-expr state. For example, a regex UDF might look like:

struct MyRegexUdf { // compiled regex regex: Regex, // regex as original string pattern: String, } impl ScalarUdfImpl for MyRegexUdf { // implementation here, not important for example }

This would take a Utf8 expression in as an argument when evaluated but that will be serialized and deserialized separately. In the proto we only need to serialized the regex. So the proto would be something like

message MyRegexUdfNode { string pattern = 1; }

Then in our custom codec we have

fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> { if let Ok(proto) = MyRegexUdfNode::decode(buf) { let regex = Regex::new(&proto.pattern).map_err(|err| ...)?; Ok(Arc::new( ScalarUDF::new_from_impl(MyRegexUdf { regex, pattern: proto.pattern, }) )) } else { not_impl_err!("unrecognized scalar UDF implementation, cannot decode") } }

yyy1000 · 2024-03-01T05:39:58Z

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

+            DataFusionError::Internal("Error decoding test table".to_string())
+        })?;
+        if let Some(expr) = msg.expr.as_ref() {
+            let node = from_proto::parse_expr(expr, ctx, self)?;


Here if I want to convert a LogicalExprNode to a Expr, I have to call parse_expr.
However, parse_expr will call codec. try_decode_udf, which lead circle call.
That makes me a little confused how to implement this method.

yyy1000 · 2024-03-01T05:40:08Z

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

+        });
+
+        let proto = proto::ScalarUDFProto {
+            expr: Some(serialize_expr(&func, self)?),


…nto serde

yyy1000

Hi @thinkharderdev
Thanks for your guide and this PR can work now. :)
Can you review it when you are available?

yyy1000 · 2024-03-01T19:46:19Z

datafusion/proto/Cargo.toml

@@ -55,5 +55,6 @@ serde_json = { workspace = true, optional = true }

 [dev-dependencies]
 doc-comment = { workspace = true }
+regex = "1.10.3"


I think add this dependency for test is OK?

I think the test works just fine without actually having regex. We can just make the example like

struct MyRegexUdf { pattern: String }

so we don't need to add an extra dependency.

Yeah, I think so.
Fixed it.

yyy1000 · 2024-03-01T19:52:19Z

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

+    fn try_decode_udf(&self, _name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+        if let Ok(proto) = proto::MyRegexUdfNode::decode(buf) {
+            let regex = Regex::new(&proto.pattern).map_err(|_err| ..);
+            match regex {
+                Ok(regex) => Ok(Arc::new(ScalarUDF::new_from_impl(MyRegexUdf::new(
+                    regex,
+                    proto.pattern,
+                )))),
+                Err(e) => internal_err!("unsupported regex pattern {e:?}"),
+            }
+        } else {
+            not_impl_err!("unrecognized scalar UDF implementation, cannot decode")
+        }
+    }
+
+    fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
+        let binding = node.inner();
+        let udf = binding.as_any().downcast_ref::<MyRegexUdf>().unwrap();
+        let proto = proto::MyRegexUdfNode {
+            pattern: udf.pattern.clone(),
+        };
+        proto.encode(buf).map_err(|e| {
+            DataFusionError::Internal(format!("failed to encode udf: {e:?}"))
+        })?;
+        Ok(())
+    }


This is the encode and decode logic

thinkharderdev

One small suggestion on the test but this looks great. Really nice work!

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

thinkharderdev

@yyy1000 We do still need to handle serde for physical expressions as well. This would follow the same pattern as for logical expressions.

We can either do that here or in a follow up PR. If you'd rather do it in follow-up PR then I wouldn't object to merging this one as is as I think being able to serialize the logical plan is a useful feature on it's own

Co-authored-by: Dan Harris <[email protected]>

yyy1000 · 2024-03-02T15:12:03Z

@yyy1000 We do still need to handle serde for physical expressions as well. This would follow the same pattern as for logical expressions.

We can either do that here or in a follow up PR. If you'd rather do it in follow-up PR then I wouldn't object to merging this one as is as I think being able to serialize the logical plan is a useful feature on it's own

Hi, @thinkharderdev
I'd like to work on a follow-up PR for physical-expr and merge this first, for two reasons.

This PR will lead conflict for the moving of functions. So merging this first will not need resolving conflicts regularly. :)
The change lines would be too large if this PR also includes for physical-expr, which may not be good for future review.

Thank you again for your guidance and review, it's so helpful!

thinkharderdev · 2024-03-02T15:46:14Z

@yyy1000 We do still need to handle serde for physical expressions as well. This would follow the same pattern as for logical expressions.

We can either do that here or in a follow up PR. If you'd rather do it in follow-up PR then I wouldn't object to merging this one as is as I think being able to serialize the logical plan is a useful feature on it's own

Hi, @thinkharderdev

I'd like to work on a follow-up PR for physical-expr and merge this first, for two reasons.

This PR will lead conflict for the moving of functions. So merging this first will not need resolving conflicts regularly. :)

The change lines would be too large if this PR also includes for physical-expr, which may not be good for future review.

Thank you again for your guidance and review, it's so helpful!

Yep makes sense. No worries! Thanks for your work on this!

alamb · 2024-03-02T15:47:57Z

🎉

feat: support serde for ScalarUDF

8917593

yyy1000 marked this pull request as draft February 29, 2024 02:23

yyy1000 commented Feb 29, 2024

View reviewed changes

thinkharderdev reviewed Feb 29, 2024

View reviewed changes

feat: change other places leveraing codec

c8463a7

yyy1000 marked this pull request as ready for review March 1, 2024 02:09

yyy1000 added 3 commits February 29, 2024 20:10

Merge branch 'main' into serde

47607c1

fix conflict

982d7c8

feat: add a test case(no working now)

c7db879

yyy1000 commented Mar 1, 2024

View reviewed changes

yyy1000 added 6 commits March 1, 2024 13:13

feat: complete test case

5675568

Merge branch 'main' into serde

eb11755

fix cargo fmt

88fd1bc

Merge branch 'serde' of https://github.com/yyy1000/arrow-datafusion i…

0aa6d1a

…nto serde

fix typo

ecb3141

fix typo

2ea00ee

yyy1000 commented Mar 1, 2024

View reviewed changes

remove regex dep

ef31c6f

thinkharderdev approved these changes Mar 2, 2024

View reviewed changes

datafusion/proto/tests/cases/roundtrip_logical_plan.rs Outdated Show resolved Hide resolved

thinkharderdev reviewed Mar 2, 2024

View reviewed changes

Update datafusion/proto/tests/cases/roundtrip_logical_plan.rs

2f7f269

Co-authored-by: Dan Harris <[email protected]>

thinkharderdev merged commit b87cde1 into apache:main Mar 2, 2024

yyy1000 deleted the serde branch March 2, 2024 15:50

yyy1000 mentioned this pull request Mar 2, 2024

Support Serde for ScalarUDF in Physical Expressions #9428

Closed


		fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>>;

		fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()>;

-    fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>>;
+    fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+        not_impl_err!("LogicalExtensionCodec is not provided for scalar function {name}")
+    }

Handle serde for ScalarUDF #9395

Handle serde for ScalarUDF #9395

Uh oh!

Conversation

yyy1000 commented Feb 29, 2024

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

yyy1000 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thinkharderdev left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yyy1000 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yyy1000 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thinkharderdev left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

thinkharderdev left a comment

Choose a reason for hiding this comment

Uh oh!

yyy1000 commented Mar 2, 2024

Uh oh!

thinkharderdev commented Mar 2, 2024

Uh oh!

alamb commented Mar 2, 2024

Uh oh!

Uh oh!