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-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
