Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle serde for ScalarUDF #9395

Merged
merged 13 commits into from
Mar 2, 2024
1 change: 1 addition & 0 deletions datafusion/proto/proto/datafusion.proto
Original file line number Diff line number Diff line change
Expand Up @@ -751,6 +751,7 @@ message AggregateUDFExprNode {
message ScalarUDFExprNode {
string fun_name = 1;
repeated LogicalExprNode args = 2;
optional bytes fun_definition = 3;
}

enum BuiltInWindowFunction {
Expand Down
8 changes: 5 additions & 3 deletions datafusion/proto/src/bytes/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
// under the License.

//! Serialization / Deserialization to Bytes
use crate::logical_plan::to_proto::serialize_expr;
use crate::logical_plan::{
self, AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec,
};
Expand Down Expand Up @@ -87,8 +88,8 @@ pub trait Serializeable: Sized {
impl Serializeable for Expr {
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)
Comment on lines 89 to +92
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

.map_err(|e| plan_datafusion_err!("Error encoding expr as protobuf: {e}"))?;

protobuf
Expand Down Expand Up @@ -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)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And also here

.map_err(|e| plan_datafusion_err!("Error parsing protobuf into Expr: {e}"))
}
}
Expand Down
21 changes: 21 additions & 0 deletions datafusion/proto/src/generated/pbjson.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions datafusion/proto/src/generated/prost.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading