You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "timsaucer (via GitHub)" <gi...@apache.org> on 2025/04/08 23:29:32 UTC
[PR] Add metadata support for Scalar UDFs [datafusion]
timsaucer opened a new pull request, #15646:
URL: https://github.com/apache/datafusion/pull/15646
## Which issue does this PR close?
- Closes #14247 but it is a different solution than requested
## Rationale for this change
We have many users who wish to use extension data or other metadata when writing user defined functions. This PR enables two features for Scalar UDFs
- When invoking the UDF, the metadata of each input column will be available, if it exists
- The UDF can also specify output metadata that is attached to the schema of the record batch
## What changes are included in this PR?
This is a fairly large change, but at a high level we add in a vector of argument metadata to the `ScalarFunctionArgs` that gets passed when invoked. Additionally all Physical Expressions are now required to implement a new function to output their metadata. The rest of the work included is around plumbing these changes through the system, extracting the metadata from the input schema, and setting it as output.
## Are these changes tested?
- All existing unit tests pass.
- Additional unit test exercising this feature is added.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2059245151
##########
datafusion/physical-expr-common/src/physical_expr.rs:
##########
@@ -71,11 +71,23 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + DynEq + DynHash {
/// downcast to a specific implementation.
fn as_any(&self) -> &dyn Any;
/// Get the data type of this expression, given the schema of the input
- fn data_type(&self, input_schema: &Schema) -> Result<DataType>;
+ fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+ Ok(self.return_field(input_schema)?.data_type().to_owned())
+ }
/// Determine whether this expression is nullable, given the schema of the input
- fn nullable(&self, input_schema: &Schema) -> Result<bool>;
+ fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+ Ok(self.return_field(input_schema)?.is_nullable())
+ }
/// Evaluate an expression against a RecordBatch
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue>;
+ /// The output field associated with this expression
+ fn return_field(&self, input_schema: &Schema) -> Result<Field> {
+ Ok(Field::new(
+ format!("{self}"),
Review Comment:
I believe the place it happens is in https://github.com/apache/datafusion/blob/e4d5846dddefed5bf6934acc63da7f4bd4644868/datafusion/expr/src/expr_schema.rs#L375
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer merged PR #15646:
URL: https://github.com/apache/datafusion/pull/15646
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2059245151
##########
datafusion/physical-expr-common/src/physical_expr.rs:
##########
@@ -71,11 +71,23 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + DynEq + DynHash {
/// downcast to a specific implementation.
fn as_any(&self) -> &dyn Any;
/// Get the data type of this expression, given the schema of the input
- fn data_type(&self, input_schema: &Schema) -> Result<DataType>;
+ fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+ Ok(self.return_field(input_schema)?.data_type().to_owned())
+ }
/// Determine whether this expression is nullable, given the schema of the input
- fn nullable(&self, input_schema: &Schema) -> Result<bool>;
+ fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+ Ok(self.return_field(input_schema)?.is_nullable())
+ }
/// Evaluate an expression against a RecordBatch
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue>;
+ /// The output field associated with this expression
+ fn return_field(&self, input_schema: &Schema) -> Result<Field> {
+ Ok(Field::new(
+ format!("{self}"),
Review Comment:
I believe the place it is computed is here and then used later down in the same function. https://github.com/apache/datafusion/blob/e4d5846dddefed5bf6934acc63da7f4bd4644868/datafusion/expr/src/expr_schema.rs#L375
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052889243
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
Review Comment:
I can try to take a look. My sense is that it make UDFs with metadata ineligible for constant folding but I am still new to all of this!
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
Review Comment:
(Agreed that it's follow-up PR material!)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "crystalxyz (via GitHub)" <gi...@apache.org>.
crystalxyz commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2059575045
##########
docs/source/library-user-guide/upgrading.md:
##########
@@ -19,6 +19,33 @@
# Upgrade Guides
+## DataFusion `48.0.0`
+
+### Processing `Field` instead of `DataType` for user defined functions
+
+In order to support metadata handling and extension types, user defined functions are
+now switching to traits which use `Field` rather than a `DataType` and nullability.
+This gives a single interface to both of these parameters and additionally allows
+access to metadata fields, which can be used for extension types.
+
+To upgrade structs which implement `ScalarUDFImpl`, if you have implemented
+`return_type_from_args` you need instead to implement `return_field_from_args`.
+If your functions do not need to handle metadata, this should be straightforward
+repackaging of the output data into a `Field`. The name you specify on the
+field is not important. It will be overwritten during planning. `ReturnInfo`
Review Comment:
Yeah that's true. Sounds good to me
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052893480
##########
datafusion/expr/src/udf.rs:
##########
@@ -293,14 +293,17 @@ where
/// Arguments passed to [`ScalarUDFImpl::invoke_with_args`] when invoking a
/// scalar function.
-pub struct ScalarFunctionArgs<'a> {
+pub struct ScalarFunctionArgs<'a, 'b> {
/// The evaluated arguments to the function
pub args: Vec<ColumnarValue>,
+ /// Field associated with each arg, if it exists
Review Comment:
Fixed
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "berkaysynnada (via GitHub)" <gi...@apache.org>.
berkaysynnada commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2821563176
Hi @timsaucer. I'd like to review this before getting merged. If it's not blocking anything, would it be okay to wait until weekend?
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "crystalxyz (via GitHub)" <gi...@apache.org>.
crystalxyz commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2058241720
##########
docs/source/library-user-guide/upgrading.md:
##########
@@ -19,6 +19,33 @@
# Upgrade Guides
+## DataFusion `48.0.0`
+
+### Processing `Field` instead of `DataType` for user defined functions
+
+In order to support metadata handling and extension types, user defined functions are
+now switching to traits which use `Field` rather than a `DataType` and nullability.
+This gives a single interface to both of these parameters and additionally allows
+access to metadata fields, which can be used for extension types.
+
+To upgrade structs which implement `ScalarUDFImpl`, if you have implemented
+`return_type_from_args` you need instead to implement `return_field_from_args`.
+If your functions do not need to handle metadata, this should be straightforward
+repackaging of the output data into a `Field`. The name you specify on the
+field is not important. It will be overwritten during planning. `ReturnInfo`
Review Comment:
Just wondering, it seems a little weird to me here that we ask users to provide a name value but then overwrite it. Is it considered acceptable (because it is part of arrow's Field class), or could we maybe hide this field from the user?
##########
datafusion/physical-expr-common/src/physical_expr.rs:
##########
@@ -71,11 +71,23 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + DynEq + DynHash {
/// downcast to a specific implementation.
fn as_any(&self) -> &dyn Any;
/// Get the data type of this expression, given the schema of the input
- fn data_type(&self, input_schema: &Schema) -> Result<DataType>;
+ fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+ Ok(self.return_field(input_schema)?.data_type().to_owned())
+ }
/// Determine whether this expression is nullable, given the schema of the input
- fn nullable(&self, input_schema: &Schema) -> Result<bool>;
+ fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+ Ok(self.return_field(input_schema)?.is_nullable())
+ }
/// Evaluate an expression against a RecordBatch
fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue>;
+ /// The output field associated with this expression
+ fn return_field(&self, input_schema: &Schema) -> Result<Field> {
+ Ok(Field::new(
+ format!("{self}"),
Review Comment:
If I understand correctly, this is one of the places where we overwrite the name of the field?
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2802933973
> > FYI @rluvaton as I think this is very related to
> >
> > * [deprecating `return_type` in favor of `return_type_from_args` #15662](https://github.com/apache/datafusion/issues/15662)
> >
> > Shall we just change functions to return `Field`? I think that would be the cleanest (though disruptive) solution
>
> @alamb Not sure how this is solving the problem that I said in the issue
What I was thinking is that if we are going to be changing ScalarFunction again so that instead of returning a `DataType` for `return_type_from_args` it would return a `Field`
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052897307
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1372,342 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ metadata: HashMap<String, String>,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ metadata,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ unimplemented!(
+ "this should never be called since return_field_from_args is implemented"
+ );
+ }
+
+ fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Field> {
+ Ok(Field::new(self.name(), DataType::UInt64, true)
+ .with_metadata(self.metadata.clone()))
+ }
Review Comment:
Thank you for 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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052891382
##########
datafusion/common/src/dfschema.rs:
##########
@@ -969,16 +969,28 @@ impl Display for DFSchema {
/// widely used in the DataFusion codebase.
pub trait ExprSchema: std::fmt::Debug {
/// Is this column reference nullable?
- fn nullable(&self, col: &Column) -> Result<bool>;
+ fn nullable(&self, col: &Column) -> Result<bool> {
+ Ok(self.to_field(col)?.is_nullable())
+ }
/// What is the datatype of this column?
- fn data_type(&self, col: &Column) -> Result<&DataType>;
+ fn data_type(&self, col: &Column) -> Result<&DataType> {
+ Ok(self.to_field(col)?.data_type())
+ }
/// Returns the column's optional metadata.
- fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>>;
+ fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>> {
+ Ok(self.to_field(col)?.metadata())
+ }
/// Return the column's datatype and nullability
- fn data_type_and_nullable(&self, col: &Column) -> Result<(&DataType, bool)>;
+ fn data_type_and_nullable(&self, col: &Column) -> Result<(&DataType, bool)> {
+ let field = self.to_field(col)?;
+ Ok((field.data_type(), field.is_nullable()))
+ }
+
+ // Return the column's field
+ fn to_field(&self, col: &Column) -> Result<&Field>;
Review Comment:
I changed this to `field_from_column` to be consistent with the rest of the API on this struct. If we used `field` then we would keep running into places where it conflicts with `field()` that takes the column index.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "berkaysynnada (via GitHub)" <gi...@apache.org>.
berkaysynnada commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2837754661
> @berkaysynnada do you think you'll be able to review soon? I know we wanted to get this in earlier in the 48 cycle to shake out any bugs since it is a big change
Sorry, I wasn't able to finish it yet, and it doesn't look like I can do in the next few days either. I don't want to block progress, so I'll do a post-merge review.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add metadata support for Scalar UDFs [datafusion]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2790143960
Just a link to my experiments starting from the Expr enum in case they are useful! https://github.com/apache/datafusion/pull/15036
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052888136
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
Review Comment:
https://github.com/apache/datafusion/issues/15797 for tracking
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2840222682
The march to supporting user defined types has begun!
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2059247381
##########
docs/source/library-user-guide/upgrading.md:
##########
@@ -19,6 +19,33 @@
# Upgrade Guides
+## DataFusion `48.0.0`
+
+### Processing `Field` instead of `DataType` for user defined functions
+
+In order to support metadata handling and extension types, user defined functions are
+now switching to traits which use `Field` rather than a `DataType` and nullability.
+This gives a single interface to both of these parameters and additionally allows
+access to metadata fields, which can be used for extension types.
+
+To upgrade structs which implement `ScalarUDFImpl`, if you have implemented
+`return_type_from_args` you need instead to implement `return_field_from_args`.
+If your functions do not need to handle metadata, this should be straightforward
+repackaging of the output data into a `Field`. The name you specify on the
+field is not important. It will be overwritten during planning. `ReturnInfo`
Review Comment:
I think it's more work to ask them to pass in all of the other parts rather than reuse the `Field`
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2790756798
I think Aggregate and Window UDFs should come as a separate PR. I did notice however that for Aggregates the input portion is already viable with this PR. Since `AccumulatorArgs` already passes in the input physical expression and input schema we would be able to compute the input fields. I've tested this locally with success. For Window functions we will want to add in the input schema.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2054292174
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1372,342 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ metadata: HashMap<String, String>,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ metadata,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ unimplemented!(
+ "this should never be called since return_field_from_args is implemented"
+ );
+ }
+
+ fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Field> {
+ Ok(Field::new(self.name(), DataType::UInt64, true)
+ .with_metadata(self.metadata.clone()))
+ }
Review Comment:
Thanks - I'll do testing as well on our suite when this merges and see if we can remove our internal workaround 🙂
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2819299161
I'm sorry the scope of this PR keeps growing, but I've tried to address the comments and I believe this is a cleaner interface. Now all physical expressions a required to implement `return_field` (which I renamed from `output_field` to be consistent). It meant updating a lot of unit tests.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2790515046
I have updated the PR to use `Field` instead of metadata `HashMap<String, String>`. In doing so we can now use extension types directly. I've added a second unit test that uses extension types, both canonical and a user defined.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2796307540
>
> FYI @rluvaton as I think this is very related to
>
> - https://github.com/apache/datafusion/issues/15662
>
>
>
> Shall we just change functions to return `Field`? I think that would be the cleanest (though disruptive) solution
@alamb Not sure how this is solving the problem that I said in the issue
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2051449542
##########
datafusion/common/src/dfschema.rs:
##########
@@ -969,16 +969,28 @@ impl Display for DFSchema {
/// widely used in the DataFusion codebase.
pub trait ExprSchema: std::fmt::Debug {
/// Is this column reference nullable?
- fn nullable(&self, col: &Column) -> Result<bool>;
+ fn nullable(&self, col: &Column) -> Result<bool> {
Review Comment:
It seems like we could (perhaps as a follow on ticket) deprecate all the other methods on `ExprSchema` as `to_field` supercedes all of them. It would make a good first issue ticket
##########
datafusion/expr/src/expr_schema.rs:
##########
@@ -762,29 +806,25 @@ mod tests {
#[derive(Debug)]
struct MockExprSchema {
- nullable: bool,
- data_type: DataType,
+ field: Field,
Review Comment:
the fact this is simpler but now also more full featured suggests to me we are on the right track
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -803,17 +806,17 @@ impl ScalarUDFImpl for TakeUDF {
&self.signature
}
fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
- not_impl_err!("Not called because the return_type_from_args is implemented")
+ not_impl_err!("Not called because the return_field_from_args is implemented")
}
/// This function returns the type of the first or second argument based on
/// the third argument:
///
/// 1. If the third argument is '0', return the type of the first argument
/// 2. If the third argument is '1', return the type of the second argument
- fn return_type_from_args(&self, args: ReturnTypeArgs) -> Result<ReturnInfo> {
- if args.arg_types.len() != 3 {
- return plan_err!("Expected 3 arguments, got {}.", args.arg_types.len());
+ fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<Field> {
Review Comment:
This is the core change, as I understand.
##########
datafusion/expr/src/udf.rs:
##########
@@ -719,6 +682,12 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
fn documentation(&self) -> Option<&Documentation> {
None
}
+
+ /// This describes the output field associated with this UDF.
+ /// Input field is handled through `ScalarFunctionArgs`
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ None
Review Comment:
Doesn't this function have to be consistent with `return_field_from_args`? It seems like if we leave it as `None` it means many function implementation will not implement it and thus it will be inconsistent
Also, can you please document what `Non` vs `Some` means (it isn't clear what returning `None` means here -- does it mean the Field information is not known? If so, what are the implications of that?)
##########
datafusion/common/src/dfschema.rs:
##########
@@ -969,16 +969,28 @@ impl Display for DFSchema {
/// widely used in the DataFusion codebase.
pub trait ExprSchema: std::fmt::Debug {
/// Is this column reference nullable?
- fn nullable(&self, col: &Column) -> Result<bool>;
+ fn nullable(&self, col: &Column) -> Result<bool> {
+ Ok(self.to_field(col)?.is_nullable())
+ }
/// What is the datatype of this column?
- fn data_type(&self, col: &Column) -> Result<&DataType>;
+ fn data_type(&self, col: &Column) -> Result<&DataType> {
+ Ok(self.to_field(col)?.data_type())
+ }
/// Returns the column's optional metadata.
- fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>>;
+ fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>> {
+ Ok(self.to_field(col)?.metadata())
+ }
/// Return the column's datatype and nullability
- fn data_type_and_nullable(&self, col: &Column) -> Result<(&DataType, bool)>;
+ fn data_type_and_nullable(&self, col: &Column) -> Result<(&DataType, bool)> {
+ let field = self.to_field(col)?;
+ Ok((field.data_type(), field.is_nullable()))
+ }
+
+ // Return the column's field
+ fn to_field(&self, col: &Column) -> Result<&Field>;
Review Comment:
I would personally find `field()` (rather than `to_field`) clearer here as I normally think of methods like `to_` doing some sort of owned concersion where as this is just accessing a field
##########
datafusion/expr/src/udf.rs:
##########
@@ -293,14 +293,17 @@ where
/// Arguments passed to [`ScalarUDFImpl::invoke_with_args`] when invoking a
/// scalar function.
-pub struct ScalarFunctionArgs<'a> {
+pub struct ScalarFunctionArgs<'a, 'b> {
/// The evaluated arguments to the function
pub args: Vec<ColumnarValue>,
+ /// Field associated with each arg, if it exists
Review Comment:
very minor but i found it strange that `a` was used second (as in I would expect
```rust
pub arg_fields: Vec<Option<&'a Field>>,
...
pub return_field: &'b Field,
```
##########
datafusion/expr/src/udf.rs:
##########
@@ -309,64 +312,18 @@ pub struct ScalarFunctionArgs<'a> {
/// such as the type of the arguments, any scalar arguments and if the
/// arguments can (ever) be null
///
-/// See [`ScalarUDFImpl::return_type_from_args`] for more information
+/// See [`ScalarUDFImpl::return_field_from_args`] for more information
#[derive(Debug)]
-pub struct ReturnTypeArgs<'a> {
+pub struct ReturnFieldArgs<'a> {
/// The data types of the arguments to the function
- pub arg_types: &'a [DataType],
+ pub arg_fields: &'a [Field],
/// Is argument `i` to the function a scalar (constant)
///
/// If argument `i` is not a scalar, it will be None
///
/// For example, if a function is called like `my_function(column_a, 5)`
/// this field will be `[None, Some(ScalarValue::Int32(Some(5)))]`
pub scalar_arguments: &'a [Option<&'a ScalarValue>],
- /// Can argument `i` (ever) null?
- pub nullables: &'a [bool],
-}
-
-/// Return metadata for this function.
-///
-/// See [`ScalarUDFImpl::return_type_from_args`] for more information
-#[derive(Debug)]
-pub struct ReturnInfo {
Review Comment:
It is sort of interesting to see that DataFUsion reinvented Field in several places (ReturnInfo, ExprSchamable, etc)
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
Review Comment:
I think this particular function `invoke_with_args`, gets run during the execuiton phase on each batch
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
Review Comment:
It is a good call we'll need to add the metadata somewhere. I think we'll have to figure out the best place as a follow on PR (we can defer until we figure out the right place)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2838939271
Thank you, everyone, for the thoughtful discussions and reviews.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "tobixdev (via GitHub)" <gi...@apache.org>.
tobixdev commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2794906165
Just chiming in here. I also encountered the similar issues with UDFs not having access to (and being able to return) fields during my experiments.
Thanks for working on this 🚀! I'd really love to help out here, as I am craving for better UDTs support in DF. Unfortunately, I am a bit swamped at the moment but I'll find a few hours if there is something I can help with.
I also believe that we should update the return API to Field as well.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2050655982
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
Review Comment:
I don't think that's a real problem because this happens during the planning stages, which is generally a smaller fraction of the runtime for any non-trivially sized data.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2050653344
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
Review Comment:
I *think* it's okay to leave signature as is, because that's mostly going to be used to determine if we need type coercion. In the call to `return_field_from_args` the user will have the opportunity to check the full metadata *after* any coercion has happened and can then fail the plan at that point.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052893257
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -803,17 +806,17 @@ impl ScalarUDFImpl for TakeUDF {
&self.signature
}
fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
- not_impl_err!("Not called because the return_type_from_args is implemented")
+ not_impl_err!("Not called because the return_field_from_args is implemented")
}
/// This function returns the type of the first or second argument based on
/// the third argument:
///
/// 1. If the third argument is '0', return the type of the first argument
/// 2. If the third argument is '1', return the type of the second argument
- fn return_type_from_args(&self, args: ReturnTypeArgs) -> Result<ReturnInfo> {
- if args.arg_types.len() != 3 {
- return plan_err!("Expected 3 arguments, got {}.", args.arg_types.len());
+ fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<Field> {
Review Comment:
This and also `return_field` in `PhysicalExpr`
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2038019878
##########
datafusion/expr/src/udf.rs:
##########
@@ -293,9 +293,11 @@ where
/// Arguments passed to [`ScalarUDFImpl::invoke_with_args`] when invoking a
/// scalar function.
-pub struct ScalarFunctionArgs<'a> {
+pub struct ScalarFunctionArgs<'a, 'b> {
/// The evaluated arguments to the function
pub args: Vec<ColumnarValue>,
+ /// Field associated with each arg, if it exists
+ pub arg_fields: Vec<Option<&'b Field>>,
/// The number of rows in record batch being evaluated
pub number_rows: usize,
Review Comment:
I think we should also remove `ReturnTypeInfo` and `return_type` and have a single way to return values -- `output_field` (or `return_field`?)
##########
datafusion/expr/src/udf.rs:
##########
@@ -293,9 +293,11 @@ where
/// Arguments passed to [`ScalarUDFImpl::invoke_with_args`] when invoking a
/// scalar function.
-pub struct ScalarFunctionArgs<'a> {
+pub struct ScalarFunctionArgs<'a, 'b> {
/// The evaluated arguments to the function
pub args: Vec<ColumnarValue>,
+ /// Field associated with each arg, if it exists
+ pub arg_fields: Vec<Option<&'b Field>>,
/// The number of rows in record batch being evaluated
pub number_rows: usize,
Review Comment:
Perhaps we should change the return type to also be `pub return_field: &'a Field` to be consistetn
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052889243
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
Review Comment:
I can try to take a look. My sense is that it make UDFs with metadata ineligible for constant folding but I am still new to all of 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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2050662012
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::Int8(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::Utf8(print_value(
+ *value,
+ output_as_bool,
+ ))))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(
+ Field::new("canonical_extension_udf", DataType::Utf8, true)
+ .with_extension_type(MyUserExtentionType {}),
Review Comment:
I agree the name is ignored, but the nullability *will* be supported with the recent changes. It's now plumbed through so this field is not thrown away.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052889464
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
Review Comment:
(Agreed that it's follow-up PR material!)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2819327271
@alamb Please let me know if you want a re-review or to merge in after CI finishes. I did make some decent sized changes per your suggestions. A lot of what changed are unit tests, just for the way we handle borrowed `Field` in the invocation arguments.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2053073575
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1372,342 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ metadata: HashMap<String, String>,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ metadata,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ unimplemented!(
+ "this should never be called since return_field_from_args is implemented"
+ );
+ }
+
+ fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Field> {
+ Ok(Field::new(self.name(), DataType::UInt64, true)
+ .with_metadata(self.metadata.clone()))
+ }
Review Comment:
@paleolimbot I have an end to end example working on my company's code, so if there are still gaps from what you need please let me know!
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2790938185
I'll take a look this evening!
It's mostly updating to Arrow 55, but https://github.com/apache/datafusion/pull/15663 (in particular https://github.com/apache/datafusion/pull/15663/commits/e37ef60d68d74bf62f251d3f2f11021b83a4effb ) is the change required to support Extension types if arrow-rs supported them (minus reading them from files, which is admittedly a big minus).
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2797745297
I need to take some time to review these comments and think more about it, likely next week. Also I'm dropping a note for myself that the current implementation isn't sufficient for my needs because I need the UDF to be able to compute the output field based on the arguments and my function signature for `output_field` will not enable that.
Case in point: `GetFieldFunc` should return the field with meta data for the subfield of the struct, so it needs to know which field it was called for.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2835919583
@berkaysynnada do you think you'll be able to review soon? I know we wanted to get this in earlier in the 48 cycle to shake out any bugs since it is a big change
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2050660795
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
Review Comment:
Interesting. I don't know. As a a work around someone could make a light UDF that returns a scalar value but with metadata. Alternatively we could add methods to the literal expression to give the ability to add metadata to it. I don't have a clear use case in my head, though, so I'm mostly spitballing 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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2815465523
I've been working through switching over to returning field instead of data type. This had many touch points. The PR has grown very large, but a significant portion of that is updating unit tests to pass in the appropriate values.
One big change occurs in `datafusion/expr/src/expr_schema.rs` where I moved over the logic that was in `data_type_and_nullable` *and* `metadata` into `to_field`. Then the other two functions call `to_field` and pick out the parts they want.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2050650946
##########
datafusion/expr/src/udf.rs:
##########
@@ -293,9 +293,11 @@ where
/// Arguments passed to [`ScalarUDFImpl::invoke_with_args`] when invoking a
/// scalar function.
-pub struct ScalarFunctionArgs<'a> {
+pub struct ScalarFunctionArgs<'a, 'b> {
/// The evaluated arguments to the function
pub args: Vec<ColumnarValue>,
+ /// Field associated with each arg, if it exists
+ pub arg_fields: Vec<Option<&'b Field>>,
/// The number of rows in record batch being evaluated
pub number_rows: usize,
Review Comment:
Changed in most recent push
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052893897
##########
datafusion/expr/src/udf.rs:
##########
@@ -719,6 +682,12 @@ pub trait ScalarUDFImpl: Debug + Send + Sync {
fn documentation(&self) -> Option<&Documentation> {
None
}
+
+ /// This describes the output field associated with this UDF.
+ /// Input field is handled through `ScalarFunctionArgs`
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ None
Review Comment:
Yes! I took your suggestion and actually removed this function entirely.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2052890614
##########
datafusion/common/src/dfschema.rs:
##########
@@ -969,16 +969,28 @@ impl Display for DFSchema {
/// widely used in the DataFusion codebase.
pub trait ExprSchema: std::fmt::Debug {
/// Is this column reference nullable?
- fn nullable(&self, col: &Column) -> Result<bool>;
+ fn nullable(&self, col: &Column) -> Result<bool> {
Review Comment:
https://github.com/apache/datafusion/issues/15798
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2821605182
> Hi @timsaucer. I'd like to review this before getting merged. If it's not blocking anything, would it be okay to wait until weekend?
I would greatly appreciate the review, and I think it is worth waiting to get another look.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2794766525
This sounds like a good idea to me, but I suggest we wait until DF 47 is shipped to merge 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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add Extension Type / Metadata support for Scalar UDFs [datafusion]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#discussion_r2038115265
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
Review Comment:
Would the `Signature` also need additional options at some point such that it would have an opportunity to consider the metadata of the arguments?
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
Review Comment:
One of the things that I've worried about (but haven't benchmarked) is whether the parsing of the extension metadata (usually small amounts of JSON) will add up here (my sense is that it won't but it may be worth checking).
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::Int8(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::Utf8(print_value(
+ *value,
+ output_as_bool,
+ ))))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(
+ Field::new("canonical_extension_udf", DataType::Utf8, true)
+ .with_extension_type(MyUserExtentionType {}),
Review Comment:
I wonder what the name and/or nullability should be in these cases (this also came up in my adventures with the `Expr` enum). If we don't actually need them/their value is always ignored, I wonder if a dedicated structure would be better (e.g., closer to your original version that just used a `HashMap`).
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -1367,3 +1370,346 @@ async fn register_alltypes_parquet(ctx: &SessionContext) -> Result<()> {
async fn plan_and_collect(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
ctx.sql(sql).await?.collect().await
}
+
+#[derive(Debug)]
+struct MetadataBasedUdf {
+ name: String,
+ signature: Signature,
+ output_field: Field,
+}
+
+impl MetadataBasedUdf {
+ fn new(metadata: HashMap<String, String>) -> Self {
+ // The name we return must be unique. Otherwise we will not call distinct
+ // instances of this UDF. This is a small hack for the unit tests to get unique
+ // names, but you could do something more elegant with the metadata.
+ let name = format!("metadata_based_udf_{}", metadata.len());
+ let output_field =
+ Field::new(&name, DataType::UInt64, true).with_metadata(metadata);
+ Self {
+ name,
+ signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+ output_field,
+ }
+ }
+}
+
+impl ScalarUDFImpl for MetadataBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let should_double = match &args.arg_fields[0] {
+ Some(field) => field
+ .metadata()
+ .get("modify_values")
+ .map(|v| v == "double_output")
+ .unwrap_or(false),
+ None => false,
+ };
+ let mulitplier = if should_double { 2 } else { 1 };
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap()
+ .iter()
+ .map(|v| v.map(|x| x * mulitplier))
+ .collect();
+ let array_ref = Arc::new(UInt64Array::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
+ let ScalarValue::UInt64(value) = value else {
+ return exec_err!("incorrect data type");
+ };
+
+ Ok(ColumnarValue::Scalar(ScalarValue::UInt64(
+ value.map(|v| v * mulitplier),
+ )))
+ }
+ }
+ }
+
+ fn equals(&self, other: &dyn ScalarUDFImpl) -> bool {
+ self.name == other.name()
+ }
+
+ fn output_field(&self, _input_schema: &Schema) -> Option<Field> {
+ Some(self.output_field.clone())
+ }
+}
+
+#[tokio::test]
+async fn test_metadata_based_udf() -> Result<()> {
+ let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("no_metadata", DataType::UInt64, true),
+ Field::new("with_metadata", DataType::UInt64, true).with_metadata(
+ [("modify_values".to_string(), "double_output".to_string())]
+ .into_iter()
+ .collect(),
+ ),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![Arc::clone(&data_array), Arc::clone(&data_array)],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+ let t = ctx.table("t").await?;
+ let no_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(HashMap::new()));
+ let with_output_meta_udf = ScalarUDF::from(MetadataBasedUdf::new(
+ [("output_metatype".to_string(), "custom_value".to_string())]
+ .into_iter()
+ .collect(),
+ ));
+
+ let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
+ .project(vec![
+ no_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_no_out"),
+ no_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_no_out"),
+ with_output_meta_udf
+ .call(vec![col("no_metadata")])
+ .alias("meta_no_in_with_out"),
+ with_output_meta_udf
+ .call(vec![col("with_metadata")])
+ .alias("meta_with_in_with_out"),
+ ])?
+ .build()?;
+
+ let actual = DataFrame::new(ctx.state(), plan).collect().await?;
+
+ // To test for output metadata handling, we set the expected values on the result
+ // To test for input metadata handling, we check the numbers returned
+ let mut output_meta = HashMap::new();
+ let _ = output_meta.insert("output_metatype".to_string(), "custom_value".to_string());
+ let expected_schema = Schema::new(vec![
+ Field::new("meta_no_in_no_out", DataType::UInt64, true),
+ Field::new("meta_with_in_no_out", DataType::UInt64, true),
+ Field::new("meta_no_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ Field::new("meta_with_in_with_out", DataType::UInt64, true)
+ .with_metadata(output_meta.clone()),
+ ]);
+
+ let expected = record_batch!(
+ ("meta_no_in_no_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_no_out", UInt64, [0, 10, 20, 30, 40]),
+ ("meta_no_in_with_out", UInt64, [0, 5, 10, 15, 20]),
+ ("meta_with_in_with_out", UInt64, [0, 10, 20, 30, 40])
+ )?
+ .with_schema(Arc::new(expected_schema))?;
+
+ assert_eq!(expected, actual[0]);
+
+ ctx.deregister_table("t")?;
+ Ok(())
+}
+
+/// This UDF is to test extension handling, both on the input and output
+/// sides. For the input, we will handle the data differently if there is
+/// the canonical extension type Bool8. For the output we will add a
+/// user defined extension type.
+#[derive(Debug)]
+struct ExtensionBasedUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl Default for ExtensionBasedUdf {
+ fn default() -> Self {
+ Self {
+ name: "canonical_extension_udf".to_string(),
+ signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),
+ }
+ }
+}
+impl ScalarUDFImpl for ExtensionBasedUdf {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ assert_eq!(args.arg_fields.len(), 1);
+ let input_field = args.arg_fields[0].unwrap();
+
+ let output_as_bool = matches!(
+ CanonicalExtensionType::try_from(input_field),
+ Ok(CanonicalExtensionType::Bool8(_))
+ );
+
+ // If we have the extension type set, we are outputting a boolean value.
+ // Otherwise we output a string representation of the numeric value.
+ fn print_value(v: Option<i8>, as_bool: bool) -> Option<String> {
+ v.map(|x| match as_bool {
+ true => format!("{}", x != 0),
+ false => format!("{x}"),
+ })
+ }
+
+ match &args.args[0] {
+ ColumnarValue::Array(array) => {
+ let array_values: Vec<_> = array
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .iter()
+ .map(|v| print_value(v, output_as_bool))
+ .collect();
+ let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+ Ok(ColumnarValue::Array(array_ref))
+ }
+ ColumnarValue::Scalar(value) => {
Review Comment:
I wonder how extension scalars/literals would fit in here (perhaps `Literal` will also need a `lit_field()` or `ScalarValue::Extension` would be needed)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add metadata support for Scalar UDFs [datafusion]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #15646:
URL: https://github.com/apache/datafusion/pull/15646#issuecomment-2789600400
Some of the concerns I have:
- Canonical extension types in `arrow-rs` have an implementation of ` TryFrom<&Field>` which lends towards the original Issue's suggestion of passing `Field` rather than just the metadata.
- I need to add unit tests to make sure other operations don't throw away the metadata. I don't know if there is a one size fits all solution for which should an which should not pass through metadata. For example, `alias` should definitely keep the metadata IMO but what about functions that take two inputs? And all of the other single input functions - it's not clear if we need to handle this one a case by case basis or not.
- If we do switch over to `Field` then it isn't immediately obvious to me what the `name` should be for the field. I need to investigate this a little more because there might already be a solution.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org