You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "yyy1000 (via GitHub)" <gi...@apache.org> on 2024/02/29 02:23:17 UTC
[PR] Handle serde for ScalarUDF [arrow-datafusion]
yyy1000 opened a new pull request, #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395
## 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?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "thinkharderdev (via GitHub)" <gi...@apache.org>.
thinkharderdev commented on code in PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#discussion_r1509507556
##########
datafusion/proto/Cargo.toml:
##########
@@ -55,5 +55,6 @@ serde_json = { workspace = true, optional = true }
[dev-dependencies]
doc-comment = { workspace = true }
+regex = "1.10.3"
Review Comment:
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.
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "yyy1000 (via GitHub)" <gi...@apache.org>.
yyy1000 commented on code in PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#discussion_r1509460509
##########
datafusion/proto/Cargo.toml:
##########
@@ -55,5 +55,6 @@ serde_json = { workspace = true, optional = true }
[dev-dependencies]
doc-comment = { workspace = true }
+regex = "1.10.3"
Review Comment:
I think add this dependency for test is OK?
##########
datafusion/proto/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -756,6 +767,113 @@ impl LogicalExtensionCodec for TopKExtensionCodec {
}
}
+#[derive(Debug)]
+struct MyRegexUdf {
+ signature: Signature,
+ // compiled regex
+ _regex: Regex,
+ // regex as original string
+ pattern: String,
+}
+
+impl MyRegexUdf {
+ fn new(_regex: Regex, pattern: String) -> Self {
+ Self {
+ signature: Signature::uniform(
+ 1,
+ vec![DataType::Int32],
+ Volatility::Immutable,
+ ),
+ _regex,
+ pattern,
+ }
+ }
+}
+
+/// Implement the ScalarUDFImpl trait for MyRegexUdf
+impl ScalarUDFImpl for MyRegexUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+ fn name(&self) -> &str {
+ "regex_udf"
+ }
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+ fn return_type(&self, args: &[DataType]) -> Result<DataType> {
+ if !matches!(args.first(), Some(&DataType::Utf8)) {
+ return plan_err!("regex_udf only accepts Utf8 arguments");
+ }
+ Ok(DataType::Int32)
+ }
+ fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ unimplemented!()
+ }
+}
+
+#[derive(Debug)]
+pub struct ScalarUDFExtensionCodec {}
+
+impl LogicalExtensionCodec for ScalarUDFExtensionCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[LogicalPlan],
+ _ctx: &SessionContext,
+ ) -> Result<Extension> {
+ not_impl_err!("No extension codec provided")
+ }
+
+ fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
+ not_impl_err!("No extension codec provided")
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ _buf: &[u8],
+ _schema: SchemaRef,
+ _ctx: &SessionContext,
+ ) -> Result<Arc<dyn TableProvider>> {
+ internal_err!("unsupported plan type")
+ }
+
+ fn try_encode_table_provider(
+ &self,
+ _node: Arc<dyn TableProvider>,
+ _buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ internal_err!("unsupported plan type")
+ }
+
+ 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(())
+ }
Review Comment:
This is the encode and decode logic
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "yyy1000 (via GitHub)" <gi...@apache.org>.
yyy1000 commented on PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#issuecomment-1974825842
> @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.
1. This PR will lead conflict for the moving of functions. So merging this first will not need resolving conflicts regularly. :)
2. 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!
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "thinkharderdev (via GitHub)" <gi...@apache.org>.
thinkharderdev commented on PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#issuecomment-1974834031
> > @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.
>
> 1. This PR will lead conflict for the moving of functions. So merging this first will not need resolving conflicts regularly. :)
>
> 2. 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!
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "thinkharderdev (via GitHub)" <gi...@apache.org>.
thinkharderdev merged PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "yyy1000 (via GitHub)" <gi...@apache.org>.
yyy1000 commented on code in PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#discussion_r1509530223
##########
datafusion/proto/Cargo.toml:
##########
@@ -55,5 +55,6 @@ serde_json = { workspace = true, optional = true }
[dev-dependencies]
doc-comment = { workspace = true }
+regex = "1.10.3"
Review Comment:
Yeah, I think so.
Fixed it.
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "yyy1000 (via GitHub)" <gi...@apache.org>.
yyy1000 commented on code in PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#discussion_r1508504101
##########
datafusion/proto/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -631,6 +638,12 @@ pub mod proto {
#[prost(uint64, tag = "1")]
pub k: u64,
}
+
+ #[derive(Clone, PartialEq, ::prost::Message)]
+ pub struct ScalarUDFProto {
+ #[prost(message, tag = "1")]
+ pub expr: ::core::option::Option<datafusion_proto::protobuf::LogicalExprNode>,
+ }
Review Comment:
Here is my test proto
##########
datafusion/proto/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -756,6 +769,82 @@ impl LogicalExtensionCodec for TopKExtensionCodec {
}
}
+#[derive(Debug)]
+pub struct ScalarUDFExtensionCodec {}
+
+impl LogicalExtensionCodec for ScalarUDFExtensionCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[LogicalPlan],
+ _ctx: &SessionContext,
+ ) -> Result<Extension> {
+ not_impl_err!("No extension codec provided")
+ }
+
+ fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
+ not_impl_err!("No extension codec provided")
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ _buf: &[u8],
+ _schema: SchemaRef,
+ _ctx: &SessionContext,
+ ) -> Result<Arc<dyn TableProvider>> {
+ internal_err!("unsupported plan type")
+ }
+
+ fn try_encode_table_provider(
+ &self,
+ _node: Arc<dyn TableProvider>,
+ _buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ internal_err!("unsupported plan type")
+ }
+
+ fn try_decode_udf(
+ &self,
+ _name: &str,
+ buf: &[u8],
+ ctx: &dyn FunctionRegistry,
+ ) -> Result<Arc<ScalarUDF>> {
+ let msg = ScalarUDFProto::decode(buf).map_err(|_| {
+ 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)?;
Review Comment:
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.
##########
datafusion/proto/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -756,6 +769,82 @@ impl LogicalExtensionCodec for TopKExtensionCodec {
}
}
+#[derive(Debug)]
+pub struct ScalarUDFExtensionCodec {}
+
+impl LogicalExtensionCodec for ScalarUDFExtensionCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[LogicalPlan],
+ _ctx: &SessionContext,
+ ) -> Result<Extension> {
+ not_impl_err!("No extension codec provided")
+ }
+
+ fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
+ not_impl_err!("No extension codec provided")
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ _buf: &[u8],
+ _schema: SchemaRef,
+ _ctx: &SessionContext,
+ ) -> Result<Arc<dyn TableProvider>> {
+ internal_err!("unsupported plan type")
+ }
+
+ fn try_encode_table_provider(
+ &self,
+ _node: Arc<dyn TableProvider>,
+ _buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ internal_err!("unsupported plan type")
+ }
+
+ fn try_decode_udf(
+ &self,
+ _name: &str,
+ buf: &[u8],
+ ctx: &dyn FunctionRegistry,
+ ) -> Result<Arc<ScalarUDF>> {
+ let msg = ScalarUDFProto::decode(buf).map_err(|_| {
+ 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)?;
+ match node {
+ Expr::ScalarFunction(ScalarFunction { func_def, args: _ }) => {
+ match func_def {
+ ScalarFunctionDefinition::UDF(fun) => Ok(fun),
+ _ => internal_err!("invalid plan, no udf"),
+ }
+ }
+ _ => internal_err!("invalid plan, no udf"),
+ }
+ } else {
+ internal_err!("invalid plan, no expr")
+ }
+ }
+
+ fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
+ let func = Expr::ScalarFunction(ScalarFunction {
+ func_def: ScalarFunctionDefinition::UDF(Arc::new(node.clone())),
+ args: vec![],
+ });
+
+ let proto = proto::ScalarUDFProto {
+ expr: Some(serialize_expr(&func, self)?),
Review Comment:
Also here.
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "thinkharderdev (via GitHub)" <gi...@apache.org>.
thinkharderdev commented on code in PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#discussion_r1507418283
##########
datafusion/proto/src/logical_plan/mod.rs:
##########
@@ -168,6 +175,14 @@ impl LogicalExtensionCodec for DefaultLogicalExtensionCodec {
) -> Result<()> {
not_impl_err!("LogicalExtensionCodec is not provided")
}
+
+ 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")
+ }
Review Comment:
These shouldn't be neccessary if we add the default impl in the trait
##########
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>>;
Review Comment:
```suggestion
fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result<Arc<ScalarUDF>> {
not_impl_err!("LogicalExtensionCodec is not provided for scalar function {name}")
}
```
##########
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>>;
+
+ fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()>;
Review Comment:
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(())
}
```
##########
datafusion/proto/src/bytes/mod.rs:
##########
@@ -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)
Review Comment:
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
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "yyy1000 (via GitHub)" <gi...@apache.org>.
yyy1000 commented on code in PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#discussion_r1506916545
##########
datafusion/proto/src/bytes/mod.rs:
##########
@@ -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)
Review Comment:
Would this be OK?
I think `DefaultLogicalExtensionCodec` may need to implement some functions it didn't implement yet. I will also look into it.
##########
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)
Review Comment:
And also here
##########
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))
Review Comment:
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.
##########
datafusion/proto/src/physical_plan/mod.rs:
##########
@@ -1911,6 +1912,30 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync {
) -> Result<Arc<dyn ExecutionPlan>>;
fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()>;
+
+ // 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_udaf(
+ // &self,
+ // name: &str,
+ // buf: &[u8],
+ // ) -> Result<Arc<AggregateUDF>>;
+
+ // fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> Result<()>;
+
+ // fn try_decode_udwf(
+ // &self,
+ // name: &str,
+ // buf: &[u8],
+ // ) -> Result<Arc<WindowUDF>>;
+
+ // fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec<u8>) -> Result<()>;
Review Comment:
Plan to finish parts related to `PhysicalExtensionCodec` after `LogicalExtensionCodec`
##########
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>>;
+
+ fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()>;
Review Comment:
I plan to add aggregate_udf and window_udf after scalar_udf.
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#issuecomment-1974834446
🎉
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "thinkharderdev (via GitHub)" <gi...@apache.org>.
thinkharderdev commented on code in PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#discussion_r1509371775
##########
datafusion/proto/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -631,6 +638,12 @@ pub mod proto {
#[prost(uint64, tag = "1")]
pub k: u64,
}
+
+ #[derive(Clone, PartialEq, ::prost::Message)]
+ pub struct ScalarUDFProto {
+ #[prost(message, tag = "1")]
+ pub expr: ::core::option::Option<datafusion_proto::protobuf::LogicalExprNode>,
+ }
Review Comment:
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")
}
}
```
--
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
Re: [PR] Handle serde for ScalarUDF [arrow-datafusion]
Posted by "thinkharderdev (via GitHub)" <gi...@apache.org>.
thinkharderdev commented on code in PR #9395:
URL: https://github.com/apache/arrow-datafusion/pull/9395#discussion_r1509951518
##########
datafusion/proto/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -756,6 +766,105 @@ impl LogicalExtensionCodec for TopKExtensionCodec {
}
}
+#[derive(Debug)]
+struct MyRegexUdf {
+ signature: Signature,
+ // regex as original string
+ pattern: String,
+}
+
+impl MyRegexUdf {
+ fn new(pattern: String) -> Self {
+ Self {
+ signature: Signature::uniform(
+ 1,
+ vec![DataType::Int32],
+ Volatility::Immutable,
+ ),
+ pattern,
+ }
+ }
+}
+
+/// Implement the ScalarUDFImpl trait for MyRegexUdf
+impl ScalarUDFImpl for MyRegexUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+ fn name(&self) -> &str {
+ "regex_udf"
+ }
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+ fn return_type(&self, args: &[DataType]) -> Result<DataType> {
+ if !matches!(args.first(), Some(&DataType::Utf8)) {
+ return plan_err!("regex_udf only accepts Utf8 arguments");
+ }
+ Ok(DataType::Int32)
+ }
+ fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ unimplemented!()
+ }
+}
+
+#[derive(Debug)]
+pub struct ScalarUDFExtensionCodec {}
+
+impl LogicalExtensionCodec for ScalarUDFExtensionCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[LogicalPlan],
+ _ctx: &SessionContext,
+ ) -> Result<Extension> {
+ not_impl_err!("No extension codec provided")
+ }
+
+ fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
+ not_impl_err!("No extension codec provided")
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ _buf: &[u8],
+ _schema: SchemaRef,
+ _ctx: &SessionContext,
+ ) -> Result<Arc<dyn TableProvider>> {
+ internal_err!("unsupported plan type")
+ }
+
+ fn try_encode_table_provider(
+ &self,
+ _node: Arc<dyn TableProvider>,
+ _buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ internal_err!("unsupported plan type")
+ }
+
+ fn try_decode_udf(&self, _name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+ if let Ok(proto) = proto::MyRegexUdfNode::decode(buf) {
+ Ok(Arc::new(ScalarUDF::new_from_impl(MyRegexUdf::new(
+ proto.pattern,
+ ))))
+ } else {
+ not_impl_err!("unrecognized scalar UDF implementation, cannot decode")
+ }
+ }
Review Comment:
small tweak here to use the `name` argument for demo purposes
```suggestion
fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
if name == "regex_udf" {
let proto = proto::MyRegexUdfNode::decode(buf).map_err(|err| {
DataFusionError::Internal(format!("failed to decode regex_udf: {}", err))
})?;
Ok(Arc::new(ScalarUDF::new_from_impl(MyRegexUdf::new(
proto.pattern,
))))
} else {
not_impl_err!("unrecognized scalar UDF implementation, cannot decode")
}
}
```
--
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