fix: preserve async UDF return field metadata (#22663)

## Which issue does this PR close?

Closes #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 <jeffrey.vo.australia@gmail.com>
Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: 3 people

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,

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)
     }

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,16 +61,16 @@ 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
             .schema()
             .fields()
             .iter()
             .cloned()
-            .chain(async_fields.into_iter().map(Arc::new))
+            .chain(async_fields)
             .collect();
 
         let schema = Arc::new(Schema::new(fields));