You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/06/10 23:48:20 UTC
(datafusion) branch main updated: fix: preserve async UDF return field metadata (#22663)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new db2d21e094 fix: preserve async UDF return field metadata (#22663)
db2d21e094 is described below
commit db2d21e094b1be2ff27653ce58ba1317e2237b01
Author: Kristin Cowalcijk <kr...@gmail.com>
AuthorDate: Thu Jun 11 07:34:18 2026 +0800
fix: preserve async UDF return field metadata (#22663)
## Which issue does this PR close?
Closes https://github.com/apache/datafusion/issues/22662.
## Rationale for this change
Async scalar UDFs can compute output field metadata in
`return_field_from_args(...)`, but `AsyncFuncExpr` rebuilt the output
field from only name, data type, and nullability. This dropped metadata
from async UDF result fields.
## What changes are included in this PR?
This PR updates `AsyncFuncExpr::field(...)` to preserve the planned
`return_field` metadata and only rename the field for the async
expression output.
It also adds a regression test that verifies an async UDF result batch
preserves metadata attached by `return_field_from_args(...)`.
## Are these changes tested?
Yes.
Added a regression test in:
-
`datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs`
The test fails without the fix and passes with the fix.
## Are there any user-facing changes?
Yes.
Async scalar UDF result fields now preserve metadata attached by
`return_field_from_args(...)`.
---------
Co-authored-by: Jeffrey Vo <je...@gmail.com>
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
.../user_defined_async_scalar_functions.rs | 109 ++++++++++++++++++++-
.../physical-expr/src/async_scalar_function.rs | 15 +--
datafusion/physical-plan/src/async_func.rs | 8 +-
3 files changed, 120 insertions(+), 12 deletions(-)
diff --git a/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs
index 58a5cb8039..dd91267d58 100644
--- a/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs
+++ b/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs
@@ -18,14 +18,15 @@
use std::sync::Arc;
use arrow::array::{Int32Array, RecordBatch, StringArray};
-use arrow::datatypes::{DataType, Field, Schema};
+use arrow::datatypes::{DataType, Field, FieldRef, Schema};
use async_trait::async_trait;
use datafusion::prelude::*;
use datafusion_common::test_util::format_batches;
use datafusion_common::{Result, assert_batches_eq};
use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
use datafusion_expr::{
- ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+ ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature,
+ Volatility,
};
fn register_table_and_udf() -> Result<SessionContext> {
@@ -113,6 +114,110 @@ async fn test_async_udf_metrics() -> Result<()> {
Ok(())
}
+#[tokio::test]
+async fn test_async_udf_preserves_result_field_metadata() -> Result<()> {
+ #[derive(Debug, PartialEq, Eq, Hash, Clone)]
+ struct AsyncExtensionUDF {
+ signature: Signature,
+ }
+
+ impl Default for AsyncExtensionUDF {
+ fn default() -> Self {
+ Self {
+ signature: Signature::exact(vec![DataType::Utf8], Volatility::Volatile),
+ }
+ }
+ }
+
+ impl ScalarUDFImpl for AsyncExtensionUDF {
+ fn name(&self) -> &str {
+ "async_extension"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Utf8)
+ }
+
+ fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
+ Ok(args.arg_fields[0]
+ .as_ref()
+ .clone()
+ .with_name(self.name())
+ .with_metadata(std::collections::HashMap::from([(
+ "ARROW:extension:name".to_string(),
+ "test.async.extension".to_string(),
+ )]))
+ .into())
+ }
+
+ fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ panic!("Call invoke_async_with_args instead")
+ }
+ }
+
+ #[async_trait]
+ impl AsyncScalarUDFImpl for AsyncExtensionUDF {
+ async fn invoke_async_with_args(
+ &self,
+ args: ScalarFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ Ok(args.args[0].clone())
+ }
+ }
+
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("value", DataType::Utf8, false),
+ ])),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2, 3])),
+ Arc::new(StringArray::from(vec!["one", "two", "three"])),
+ ],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("test_table", batch)?;
+ ctx.register_udf(
+ AsyncScalarUDF::new(Arc::new(AsyncExtensionUDF::default())).into_scalar_udf(),
+ );
+
+ let result = ctx
+ .sql("SELECT async_extension(value) AS result FROM test_table")
+ .await?
+ .collect()
+ .await?;
+
+ assert_eq!(result[0].schema().field(0).name(), "result");
+ assert_eq!(
+ result[0]
+ .schema()
+ .field(0)
+ .metadata()
+ .get("ARROW:extension:name"),
+ Some(&"test.async.extension".to_string())
+ );
+
+ assert_batches_eq!(
+ &[
+ "+--------+",
+ "| result |",
+ "+--------+",
+ "| one |",
+ "| two |",
+ "| three |",
+ "+--------+",
+ ],
+ &result
+ );
+
+ Ok(())
+}
+
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
struct TestAsyncUDFImpl {
batch_size: usize,
diff --git a/datafusion/physical-expr/src/async_scalar_function.rs b/datafusion/physical-expr/src/async_scalar_function.rs
index 5612e63b53..e8ee9a6948 100644
--- a/datafusion/physical-expr/src/async_scalar_function.rs
+++ b/datafusion/physical-expr/src/async_scalar_function.rs
@@ -88,12 +88,9 @@ impl AsyncFuncExpr {
}
/// Return the output field generated by evaluating this function
- pub fn field(&self, input_schema: &Schema) -> Result<Field> {
- Ok(Field::new(
- &self.name,
- self.func.data_type(input_schema)?,
- self.func.nullable(input_schema)?,
- ))
+ #[deprecated(since = "55.0.0", note = "Use return_field instead")]
+ pub fn field(&self, _input_schema: &Schema) -> Result<Field> {
+ Ok(self.return_field.as_ref().clone().with_name(&self.name))
}
/// Return the ideal batch size for this function
@@ -211,6 +208,12 @@ impl PhysicalExpr for AsyncFuncExpr {
self.func.data_type(input_schema)
}
+ fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
+ Ok(Arc::new(
+ self.return_field.as_ref().clone().with_name(&self.name),
+ ))
+ }
+
fn nullable(&self, input_schema: &Schema) -> Result<bool> {
self.func.nullable(input_schema)
}
diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs
index 1b15bf27e7..94860e54ca 100644
--- a/datafusion/physical-plan/src/async_func.rs
+++ b/datafusion/physical-plan/src/async_func.rs
@@ -23,7 +23,7 @@ use crate::{
check_if_same_properties,
};
use arrow::array::RecordBatch;
-use arrow_schema::{Fields, Schema, SchemaRef};
+use arrow_schema::{FieldRef, Fields, Schema, SchemaRef};
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
use datafusion_common::{Result, assert_eq_or_internal_err};
use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
@@ -61,8 +61,8 @@ impl AsyncFuncExec {
) -> Result<Self> {
let async_fields = async_exprs
.iter()
- .map(|async_expr| async_expr.field(input.schema().as_ref()))
- .collect::<Result<Vec<_>>>()?;
+ .map(|async_expr| async_expr.return_field(input.schema().as_ref()))
+ .collect::<Result<Vec<FieldRef>>>()?;
// compute the output schema: input schema then async expressions
let fields: Fields = input
@@ -70,7 +70,7 @@ impl AsyncFuncExec {
.fields()
.iter()
.cloned()
- .chain(async_fields.into_iter().map(Arc::new))
+ .chain(async_fields)
.collect();
let schema = Arc::new(Schema::new(fields));
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org