Skip to content

Commit f068caa

Browse files
committed
Update return_type to raise error for metadata source
Ensure return_field_from_args is the only metadata source by having PySimpleScalarUDF::return_type raise an internal error. This aligns with DataFusion guidance. Enhance Python UDF helper documentation to clarify how callers can declare extension metadata on both arguments and results.
1 parent 0f28465 commit f068caa

File tree

2 files changed

+9
-4
lines changed

2 files changed

+9
-4
lines changed

python/datafusion/user_defined.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -274,9 +274,11 @@ def udf(*args: Any, **kwargs: Any): # noqa: D417
274274
:class:`pyarrow.ChunkedArray` values.
275275
input_types (list[pa.DataType | pa.Field]): The argument types for ``func``.
276276
This list must be of the same length as the number of arguments. Pass
277-
:class:`pyarrow.Field` instances to preserve extension metadata.
278-
return_type (pa.DataType | pa.Field): The return type of the function. Use a
279-
:class:`pyarrow.Field` to preserve metadata on extension arrays.
277+
:class:`pyarrow.Field` instances when you need to declare extension
278+
metadata for an argument.
279+
return_type (pa.DataType | pa.Field): The return type of the function. Supply
280+
a :class:`pyarrow.Field` when the result should expose extension metadata
281+
to downstream consumers.
280282
volatility (Volatility | str): See `Volatility` for allowed values.
281283
name (Optional[str]): A descriptive name for the function.
282284

src/udf.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,10 @@ impl ScalarUDFImpl for PySimpleScalarUDF {
135135
}
136136

137137
fn return_type(&self, _arg_types: &[DataType]) -> datafusion::error::Result<DataType> {
138-
Ok(self.return_field.data_type().clone())
138+
Err(DataFusionError::Internal(
139+
"return_type should be unreachable when return_field_from_args is implemented"
140+
.to_string(),
141+
))
139142
}
140143

141144
fn return_field_from_args(

0 commit comments

Comments
 (0)