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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
