################################################################################
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
"""
User-defined function helpers for the DataFrame API.
"""
import copy
import functools
import inspect
from collections.abc import Generator as AbcGenerator
from collections.abc import Iterable as AbcIterable
from collections.abc import Iterator as AbcIterator
from collections.abc import Mapping as AbcMapping
from typing import (
Any,
Callable,
Dict,
Generator,
Iterable,
Iterator,
List,
Optional,
Tuple,
Type,
Union,
get_args,
get_origin,
get_type_hints,
overload,
)
from pyflink.common import Row
from pyflink.dataframe.datatype import DataType, DataTypeLike
from pyflink.table.expression import Expression
from pyflink.table.expressions import call as table_call
from pyflink.table.expressions import col as table_col
from pyflink.table.expressions import with_columns
from pyflink.table.types import MapType, RowType
from pyflink.table.udf import (
udf as table_udf,
AsyncScalarFunction,
ScalarFunction,
TableFunction,
UserDefinedFunctionWrapper,
UserDefinedTableFunctionWrapper,
udtf as table_udtf,
)
__all__ = [
"DataFrameFunctionWrapper",
"DataFrameUDFWrapper",
"DataFrameUDTFCall",
"DataFrameUDTFWrapper",
"udf",
"udtf",
]
# ============= General =============
class DataFrameFunctionWrapper:
"""Base wrapper for DataFrame SQL-bindable Python functions."""
def __init__(self, table_wrapper: Any, default_name: str) -> None:
self._table_wrapper = table_wrapper
self.__name__ = getattr(table_wrapper, "_name", default_name)
self.__doc__ = getattr(getattr(table_wrapper, "_func", None), "__doc__", None)
# ============= UDF =============
class DataFrameUDFWrapper(DataFrameFunctionWrapper):
"""
Wrapper for DataFrame UDF that provides a callable interface.
This wrapper holds the underlying Table UDF wrapper and return type,
allowing the UDF to be called with Expression arguments.
"""
def __init__(
self,
table_udf_wrapper: UserDefinedFunctionWrapper,
return_dtype: DataType,
):
"""
Initialize DataFrame UDF wrapper.
Args:
table_udf_wrapper: The underlying Table API UDF wrapper.
return_dtype: The return DataType.
"""
super().__init__(table_udf_wrapper, "udf")
self._return_dtype = return_dtype
def __call__(self, *args: Any) -> Expression:
"""
Call the UDF with Expression arguments.
Args:
*args: Expression arguments.
Returns:
An Expression representing the UDF call.
"""
return table_call(self._table_udf_wrapper, *args)
@property
def _table_udf_wrapper(self) -> UserDefinedFunctionWrapper:
return self._table_wrapper
@property
def return_dtype(self) -> DataType:
"""Get the return DataType."""
return self._return_dtype
# Overload 1: function is supplied at the call site, so udf returns the
# wrapped UDF directly. Covers:
# @udf
# def fn(...): ... # bare decorator: udf(fn)
# udf(MyScalarFunction) # direct call with a class
# udf(MyScalarFunction(), # direct call with an instance and kwargs
# return_dtype=...)
@overload
def udf(
func: Union[Callable, ScalarFunction, AsyncScalarFunction, Type],
*,
return_dtype: Optional[DataTypeLike] = ...,
deterministic: bool = ...,
name: Optional[str] = ...,
func_type: Optional[str] = ...,
concurrency: Optional[int] = ...,
batch_size: Optional[int] = ...,
num_gpus: Optional[float] = ...,
gpu_type: Optional[str] = ...,
) -> "DataFrameUDFWrapper":
...
# Overload 2: no function supplied — udf is being used as a decorator
# factory and must return a decorator that will receive the function next.
# Covers the parameterized decorator form:
# @udf(concurrency=32)
# async def fn(...): ... # evaluated as udf(concurrency=32)(fn)
@overload
def udf(
func: None = ...,
*,
return_dtype: Optional[DataTypeLike] = ...,
deterministic: bool = ...,
name: Optional[str] = ...,
func_type: Optional[str] = ...,
concurrency: Optional[int] = ...,
batch_size: Optional[int] = ...,
num_gpus: Optional[float] = ...,
gpu_type: Optional[str] = ...,
) -> Callable[
[Union[Callable, ScalarFunction, AsyncScalarFunction, Type]],
"DataFrameUDFWrapper",
]:
...
[docs]def udf(
func: Optional[Union[Callable, ScalarFunction, AsyncScalarFunction, Type]] = None,
*,
return_dtype: Optional[DataTypeLike] = None,
deterministic: bool = True,
name: Optional[str] = None,
func_type: Optional[str] = None,
concurrency: Optional[int] = None,
batch_size: Optional[int] = None,
num_gpus: Optional[float] = None,
gpu_type: Optional[str] = None,
):
"""
Create a user-defined function for DataFrame operations.
This decorator supports:
- Plain Python functions and lambdas
- Async functions defined with ``async def`` (executed asynchronously)
- ScalarFunction subclasses (class type or instance), using eval method
- AsyncScalarFunction subclasses (class type or instance), using async eval method
- Plain callable classes with __call__ method (class type or instance)
- Automatic type inference from Python type hints
Args:
func: The Python function, ScalarFunction instance/subclass, or
callable class instance/type to wrap as a UDF.
return_dtype: Optional return type. Can be:
- DataType instance (e.g., DataType.int64())
- Python type (e.g., int, str, float)
- String (e.g., 'INT', 'BIGINT')
If not specified, inferred from function's return type hint
(eval for ScalarFunction, __call__ for plain classes).
deterministic: Whether the function is deterministic (default: True).
name: Optional name for the UDF.
func_type: Optional execution format. Supported values are
``"general"``, ``"pandas"``, and ``"arrow"``. If omitted,
pandas and arrow formats are detected from pandas/pyarrow
type hints when possible, otherwise ``"general"`` is used.
concurrency: Optional concurrency (parallelism) for the UDF operator.
If specified, the operator running this UDF will use this
parallelism. UDFs with different concurrency values will be
split into separate operators.
batch_size: Optional maximum number of elements per batch.
Only applies to batch-wise UDFs. If not specified,
the global configuration 'python.fn-execution.arrow.batch.size'
is used (default 1000). When multiple UDFs are fused,
the maximum batch_size among them is used.
num_gpus: Optional number of GPUs requested for this UDF (e.g., 0.5, 1.0).
Each GPU UDF will run in its own operator and will not be fused
with other UDFs (including other GPU UDFs).
gpu_type: Optional GPU type (e.g., 'A10', 'V100').
Returns:
A DataFrameUDFWrapper that can be called with Expressions.
Example::
>>> from pyflink.dataframe import udf, DataType, col
>>> from pyflink.table.udf import ScalarFunction
>>>
>>> # 1. Plain function with explicit return_dtype
>>> @udf(return_dtype=DataType.string())
... def to_string(x):
... return str(x)
>>>
>>> df.with_columns(s=to_string(col("a")))
>>>
>>> # 2. Plain function with type hint inference
>>> @udf
... def add(x: int, y: int) -> int:
... return x + y
>>>
>>> df.with_columns(sum=add(col("a"), col("b")))
>>>
>>> # 3. ScalarFunction with explicit return_dtype (instance or class)
>>> class AddOne(ScalarFunction):
... def eval(self, x):
... return x + 1
>>>
>>> add_one = udf(AddOne(), return_dtype=DataType.int64())
>>> add_one = udf(AddOne, return_dtype=DataType.int64()) # auto-instantiated
>>> df.with_columns(b=add_one(col("a")))
>>>
>>> # 4. ScalarFunction instance with type hint inference
>>> class Double(ScalarFunction):
... def eval(self, x: int) -> int:
... return x * 2
>>>
>>> double = udf(Double())
>>> df.with_columns(doubled=double(col("a")))
>>>
>>> # 5. @udf decorator on ScalarFunction class (with return_dtype or type hints)
>>> @udf(return_dtype=DataType.int64())
... class AddTwo(ScalarFunction):
... def eval(self, x):
... return x + 2
>>>
>>> df.with_columns(b=AddTwo(col("a")))
>>>
>>> @udf
... class Triple(ScalarFunction):
... def eval(self, x: int) -> int:
... return x * 3
>>>
>>> df.with_columns(tripled=Triple(col("a")))
>>>
>>> # 6. Plain callable class instance
>>> class MultiplyBy:
... def __init__(self, factor):
... self._factor = factor
...
... def __call__(self, x: int) -> int:
... return x * self._factor
>>>
>>> times_three = udf(MultiplyBy(3))
>>> df.with_columns(result=times_three(col("a")))
>>>
>>> # 7. @udf decorator on plain callable class
>>> @udf(return_dtype=DataType.int64())
... class AddOne:
... def __call__(self, x):
... return x + 1
>>>
>>> df.with_columns(b=AddOne(col("a")))
>>>
>>> @udf
... class DoubleIt:
... def __call__(self, x: int) -> int:
... return x * 2
>>>
>>> df.with_columns(doubled=DoubleIt(col("a")))
>>>
>>> # 8. Async function (executed asynchronously, e.g. for I/O-bound work)
>>> import asyncio
>>> @udf
... async def async_lookup(key: str) -> str:
... await asyncio.sleep(0.01)
... return f"value_for_{key}"
>>>
>>> df.with_columns(v=async_lookup(col("a")))
>>>
>>> # 9. AsyncScalarFunction subclass
>>> from pyflink.table.udf import AsyncScalarFunction
>>> class AsyncLookup(AsyncScalarFunction):
... async def eval(self, key: str) -> str:
... await asyncio.sleep(0.01)
... return f"value_for_{key}"
>>>
>>> async_lookup = udf(AsyncLookup())
>>> df.with_columns(v=async_lookup(col("a")))
>>>
>>> # 10. Pandas scalar function (executed on pandas.Series batches)
>>> import pandas as pd
>>> @udf(return_dtype=DataType.int64(), func_type="pandas")
... def add_one_pandas(x: pd.Series) -> pd.Series:
... return x + 1
>>>
>>> df.with_columns(v=add_one_pandas(col("a")))
>>>
>>> # 11. Pandas async scalar function (executed on pandas.Series batches)
>>> import pandas as pd
>>> @udf(return_dtype=DataType.string(), batch_size=128)
... async def async_batch_lookup(keys: pd.Series) -> pd.Series:
... await asyncio.sleep(0.01)
... return keys.map(lambda key: f"value_for_{key}")
>>>
>>> df.with_columns(v=async_batch_lookup(col("a")))
>>>
>>> # 12. Arrow scalar function (executed on pyarrow.Array batches)
>>> import pyarrow as pa
>>> @udf(return_dtype=DataType.int64(), func_type="arrow")
... def add_one_arrow(x):
... return pa.array([v + 1 for v in x.to_pylist()])
>>>
>>> df.with_columns(v=add_one_arrow(col("a")))
>>>
>>> # 13. Arrow async scalar function (executed on pyarrow.Array batches)
>>> import pyarrow as pa
>>> @udf(return_dtype=DataType.string(), func_type="arrow", batch_size=128)
... async def async_arrow_lookup(keys):
... await asyncio.sleep(0.01)
... return pa.array([f"value_for_{k}" for k in keys.to_pylist()])
>>>
>>> df.with_columns(v=async_arrow_lookup(col("a")))
"""
if concurrency is not None and (not isinstance(concurrency, int) or concurrency <= 0):
raise ValueError("concurrency must be a positive integer, got: {}".format(concurrency))
if batch_size is not None and (not isinstance(batch_size, int) or batch_size <= 0):
raise ValueError("batch_size must be a positive integer, got: {}".format(batch_size))
if num_gpus is not None and (not isinstance(num_gpus, (int, float)) or num_gpus <= 0):
raise ValueError("num_gpus must be a positive number, got: {}".format(num_gpus))
if num_gpus is not None and gpu_type is None:
raise ValueError("gpu_type must be specified when num_gpus is set")
# Handle class-based UDFs (class type or instance)
if func is not None and _is_class_udf(func):
return _create_class_udf(func, return_dtype, deterministic, name, concurrency, batch_size,
num_gpus, gpu_type, func_type)
def decorator(f) -> DataFrameUDFWrapper:
# Handle class-based UDFs passed through @udf(...) decorator
if _is_class_udf(f):
return _create_class_udf(f, return_dtype, deterministic, name, concurrency, batch_size,
num_gpus, gpu_type, func_type)
# Infer return type
actual_return_dtype = _infer_return_dtype(f, return_dtype)
# Create Table API UDF
table_result_type = actual_return_dtype._table_type
# Determine func_type
actual_func_type = func_type or _detect_func_type(f)
_validate_scalar_udf_options(f, actual_func_type, batch_size)
# Create Table API UDF wrapper
table_wrapper = table_udf(
f,
result_type=table_result_type,
deterministic=deterministic,
name=name,
func_type=actual_func_type,
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
)
# Create DataFrame UDF wrapper
return DataFrameUDFWrapper(table_wrapper, actual_return_dtype)
# Handle both @udf and @udf(...) syntax
if func is not None:
# @udf without parentheses
return decorator(func)
else:
# @udf(...) with parentheses
return decorator
def _validate_scalar_udf_options(func, func_type: str,
batch_size: Optional[int]) -> None:
if batch_size is not None and func_type not in ("pandas", "arrow"):
raise ValueError(
"batch_size is only supported for batch-wise UDFs. "
"The function '{}' is detected as a general UDF. "
"Use it with map_batches() or annotate parameters with pandas/pyarrow types."
.format(getattr(func, '__name__', type(func).__name__)))
def _has_custom_call(cls) -> bool:
"""Check if cls defines __call__ in its MRO (excluding object)."""
return any('__call__' in c.__dict__ for c in cls.__mro__ if c is not object)
def _is_class_udf(func) -> bool:
"""
Check if func is a class-based UDF (ScalarFunction, AsyncScalarFunction, or callable class).
"""
# Case 1: ScalarFunction / AsyncScalarFunction instance
if isinstance(func, (ScalarFunction, AsyncScalarFunction)):
return True
if inspect.isclass(func):
# Case 2: ScalarFunction / AsyncScalarFunction class (will be auto-instantiated)
if issubclass(func, (ScalarFunction, AsyncScalarFunction)):
return True
# Case 3: Plain class with __call__ (will be auto-instantiated),
# e.g. udf(MyCallable, return_dtype=...) or @udf class MyCallable
# Also handles inherited __call__ from parent class.
if _has_custom_call(func):
return True
# Case 4: Instance of a plain callable class (not a function, not a class),
# e.g. udf(MultiplyBy(3), return_dtype=...)
# Note: cannot use callable() here as it would also match lambdas etc.
if not inspect.isfunction(func) and not inspect.isclass(func) and hasattr(func, '__call__'):
if _has_custom_call(type(func)):
return True
return False
def _create_class_udf(
func: Union[ScalarFunction, AsyncScalarFunction, Type],
return_dtype: Optional[DataTypeLike],
deterministic: bool,
name: Optional[str],
concurrency: Optional[int] = None,
batch_size: Optional[int] = None,
num_gpus: Optional[float] = None,
gpu_type: Optional[str] = None,
func_type: Optional[str] = None,
) -> DataFrameUDFWrapper:
"""Create a DataFrameUDFWrapper from a class-based UDF."""
# Instantiate if a class is passed
if inspect.isclass(func):
instance = func()
else:
instance = func
is_scalar = isinstance(instance, (ScalarFunction, AsyncScalarFunction))
# Determine which method to use for type hint inference
hint_method = instance.eval if is_scalar else instance.__call__
# Infer return type
actual_return_dtype = _infer_return_dtype(hint_method, return_dtype)
table_result_type = actual_return_dtype._table_type
# Detect func type (general, pandas, or arrow)
actual_func_type = func_type or _detect_func_type(hint_method)
_validate_scalar_udf_options(func, actual_func_type, batch_size)
# For ScalarFunction, pass the instance directly.
# For plain callable classes, pass the instance as a callable.
table_wrapper = table_udf(
instance,
result_type=table_result_type,
deterministic=deterministic,
name=name,
func_type=actual_func_type,
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
)
return DataFrameUDFWrapper(table_wrapper, actual_return_dtype)
def _infer_return_dtype(func: Callable, return_dtype: Optional[DataTypeLike]) -> DataType:
"""
Infer return DataType from function or explicit specification.
Args:
func: The Python function.
return_dtype: Explicitly specified return type (can be None).
Returns:
The inferred or specified DataType.
Raises:
ValueError: If no return type hint and no explicit return_dtype.
TypeError: If the type cannot be inferred.
"""
if return_dtype is not None:
return _convert_to_dtype(return_dtype)
# Try to infer from type hints
try:
hints = get_type_hints(func)
except Exception:
hints = {}
if "return" not in hints:
raise ValueError(
f"Function '{func.__name__}' must have a return type hint "
f"or specify return_dtype parameter explicitly.\n"
f"Example:\n"
f" @udf\n"
f" def {func.__name__}(x: int) -> int:\n"
f" return x + 1\n\n"
f"Or:\n"
f" @udf(return_dtype=DataType.int64())\n"
f" def {func.__name__}(x):\n"
f" return x + 1"
)
return DataType._infer_from_type(hints["return"])
def _convert_to_dtype(dtype_like: DataTypeLike) -> DataType:
"""
Convert DataTypeLike to DataType.
Args:
dtype_like: DataType instance, Python type, or SQL string.
Returns:
A DataType instance.
"""
if isinstance(dtype_like, DataType):
return dtype_like
elif isinstance(dtype_like, str):
return DataType._from_sql(dtype_like)
else:
# Assume it's a Python type
return DataType._infer_from_type(dtype_like)
def _detect_func_type(func: Callable) -> str:
"""
Detect the function type (general, pandas, or arrow).
This checks if the function signature uses pandas or pyarrow types.
Args:
func: The Python function.
Returns:
'general', 'pandas', or 'arrow'.
"""
try:
import pandas as pd
func_globals = getattr(func, '__globals__', {})
hints = get_type_hints(func, globalns={**func_globals, "pandas": pd, "pd": pd})
# Check if any parameter or return type is pandas Series/DataFrame
for hint in hints.values():
origin = getattr(hint, "__origin__", None)
if origin is not None:
# Check generic types
pass
elif hint in (pd.Series, pd.DataFrame):
return "pandas"
except (ImportError, NameError):
pass
try:
import pyarrow as pa
func_globals = getattr(func, '__globals__', {})
hints = get_type_hints(func, globalns={**func_globals, "pyarrow": pa, "pa": pa})
# Check if any parameter or return type is an Arrow batch/column type.
for hint in hints.values():
if hint in (pa.Array, pa.ChunkedArray):
return "arrow"
except (ImportError, NameError):
pass
return "general"
def _is_typeddict(tp) -> bool:
"""Check if tp is a TypedDict class."""
# typing.is_typeddict available since Python 3.10
try:
from typing import is_typeddict
return is_typeddict(tp)
except (ImportError, NameError):
pass
# Fallback: TypedDict classes have __required_keys__ (Python 3.9+)
return (
isinstance(tp, type)
and issubclass(tp, dict)
and hasattr(tp, '__required_keys__')
)
def _infer_struct_from_typeddict(td) -> DataType:
"""Infer a DataType.struct from a TypedDict class."""
hints = get_type_hints(td)
fields = {name: DataType._infer_from_type(hint) for name, hint in hints.items()}
return DataType.struct(fields)
def _resolve_map_udf(func, return_dtype, func_type, input_columns,
concurrency: Optional[int] = None, batch_size: Optional[int] = None,
num_gpus: Optional[float] = None, gpu_type: Optional[str] = None):
"""
Resolve a function into a Table API UDF wrapper for map/map_batches.
The user function uses dict as input/output:
- For map (general): func(dict[str, scalar]) -> dict[str, scalar]
- For map_batches (pandas): func(dict[str, pd.Series]) -> dict[str, pd.Series]
- For map_batches (arrow): func(dict[str, pa.Array]) -> dict[str, pa.Array]
If func is already a DataFrameUDFWrapper, extracts the underlying wrapper directly.
TypedDict return type hints are supported for automatic return_dtype inference.
Args:
func: A DataFrameUDFWrapper or raw callable.
return_dtype: The return DataType (can be None if TypedDict hint is present).
func_type: 'general' for map, 'pandas' or 'arrow' for map_batches.
input_columns: List of input column names from the source DataFrame.
concurrency: Optional concurrency (parallelism) for the UDF operator.
num_gpus: Optional number of GPUs for the UDF operator.
gpu_type: Optional GPU type for the UDF operator.
Returns:
A UserDefinedScalarFunctionWrapper for Table.map().
"""
if isinstance(func, DataFrameUDFWrapper):
wrapper = func._table_udf_wrapper
if concurrency is not None or batch_size is not None or num_gpus is not None \
or gpu_type is not None:
# Override settings if explicitly specified in map/map_batches call.
# Use a shallow copy to avoid mutating the original wrapper.
wrapper = copy.copy(wrapper)
wrapper._judf_placeholder = None
if concurrency is not None:
wrapper._concurrency = concurrency
if batch_size is not None:
wrapper._batch_size = batch_size
if num_gpus is not None:
wrapper._num_gpus = num_gpus
if gpu_type is not None:
wrapper._gpu_type = gpu_type
return wrapper
if return_dtype is None:
if func_type == "arrow":
raise ValueError(
"return_dtype is required for arrow map_batches because the "
"output Arrow schema cannot be inferred from type hints.")
# Try to infer return_dtype from TypedDict return type hint
try:
hints = get_type_hints(func)
except NameError as e:
import warnings
warnings.warn(
f"Failed to resolve type hints for {func}: {e}. "
"Please specify return_dtype explicitly.",
stacklevel=2,
)
hints = {}
except Exception:
hints = {}
ret = hints.get("return")
if ret is not None and _is_typeddict(ret):
return_dtype = _infer_struct_from_typeddict(ret)
else:
raise ValueError(
"return_dtype is required when passing a raw function to "
"map/map_batches.\n"
"You can either specify return_dtype explicitly:\n"
" df.map(func, return_dtype=DataType.struct("
"{'a': DataType.int64()}))\n"
"Or use a TypedDict return type hint:\n"
" class Output(TypedDict):\n"
" a: int\n"
" def func(row: ...) -> Output: ..."
)
dtype = _convert_to_dtype(return_dtype)
if func_type == "general":
wrapped = _wrap_map_general(func, input_columns, dtype)
elif func_type == "arrow":
wrapped = _wrap_map_arrow(func, input_columns, dtype)
else:
wrapped = _wrap_map_pandas(func, input_columns, dtype)
return table_udf(
wrapped,
result_type=dtype._table_type,
func_type=func_type,
name=func.__name__ if hasattr(func, '__name__') else type(func).__name__,
concurrency=concurrency,
batch_size=batch_size,
num_gpus=num_gpus,
gpu_type=gpu_type,
)
def _wrap_map_general(func, input_columns, return_dtype):
"""Wrap a dict-based map function for general UDF (Row -> Row)."""
output_names = return_dtype._table_type.field_names()
@functools.wraps(func)
def wrapper(row):
input_dict = {name: row[i] for i, name in enumerate(input_columns)}
result_dict = func(input_dict)
return Row(*[result_dict[name] for name in output_names])
return wrapper
def _wrap_map_pandas(func, input_columns, return_dtype):
"""Wrap a dict-based map_batches function for pandas UDF (pd.DataFrame -> Row of Series)."""
import pandas as pd
output_names = return_dtype._table_type.field_names()
@functools.wraps(func)
def wrapper(pdf):
input_dict = {name: pdf.iloc[:, i] for i, name in enumerate(input_columns)}
result_dict = func(input_dict)
return pd.concat([result_dict[name] for name in output_names], axis=1)
return wrapper
def _wrap_map_arrow(func, input_columns, return_dtype):
"""Wrap a dict-based map_batches function for Arrow UDF."""
output_names = return_dtype._table_type.field_names()
def wrapper(batch):
input_dict = {name: batch.column(i) for i, name in enumerate(input_columns)}
result_dict = func(input_dict)
return _record_batch_from_column_dict(result_dict, output_names)
return wrapper
def _record_batch_from_column_dict(result_dict, output_names):
import pyarrow as pa
import pyarrow_hotfix # noqa # pylint: disable=unused-import
columns = []
for name in output_names:
if name not in result_dict:
raise KeyError(
"Arrow batch UDF result is missing output column %r." % name)
column = result_dict[name]
if isinstance(column, pa.ChunkedArray):
if column.num_chunks == 1:
column = column.chunk(0)
else:
column = column.combine_chunks()
columns.append(column)
return pa.record_batch(columns, names=output_names)
# ============= UDTF =============
class DataFrameUDTFCall:
"""A DataFrame UDTF call plus Python-side output metadata."""
def __init__(
self,
expression: Expression,
return_dtype: DataType,
has_named_fields: bool,
output_aliases: Optional[List[str]] = None,
) -> None:
self.expression: Expression = expression
self.return_dtype: DataType = return_dtype
self.has_named_fields: bool = has_named_fields
self.output_aliases: Optional[List[str]] = output_aliases
def alias(self, name: str, *extra_names: str) -> "DataFrameUDTFCall":
"""Return a UDTF call with output fields named by aliases."""
output_aliases = [name] + list(extra_names)
_validate_output_aliases(
output_aliases, _get_output_arity(self.return_dtype))
return DataFrameUDTFCall(
self.expression.alias(name, *extra_names),
self.return_dtype,
self.has_named_fields,
output_aliases,
)
class DataFrameUDTFWrapper(DataFrameFunctionWrapper):
"""Wrapper for DataFrame UDTFs that can be called with expressions."""
def __init__(
self,
table_udtf_wrapper: UserDefinedTableFunctionWrapper,
return_dtype: DataType,
has_named_fields: bool,
input_func: Optional[Callable[..., Any]] = None,
) -> None:
super().__init__(table_udtf_wrapper, "udtf")
self._return_dtype: DataType = return_dtype
self._has_named_fields: bool = has_named_fields
self._input_func: Optional[Callable[..., Any]] = input_func
def __call__(self, *args: Any) -> DataFrameUDTFCall:
return DataFrameUDTFCall(
self._table_udtf_wrapper(*args),
self._return_dtype,
self._has_named_fields,
)
@property
def _table_udtf_wrapper(self) -> UserDefinedTableFunctionWrapper:
return self._table_wrapper
@property
def return_dtype(self) -> DataType:
return self._return_dtype
@overload
def udtf(
func: Union[Callable[..., Any], TableFunction, Type[TableFunction]],
*,
return_dtype: Optional[DataTypeLike] = ...,
deterministic: bool = ...,
name: Optional[str] = ...,
concurrency: Optional[int] = ...,
num_gpus: Optional[float] = ...,
gpu_type: Optional[str] = ...,
) -> DataFrameUDTFWrapper:
...
@overload
def udtf(
func: None = ...,
*,
return_dtype: Optional[DataTypeLike] = ...,
deterministic: bool = ...,
name: Optional[str] = ...,
concurrency: Optional[int] = ...,
num_gpus: Optional[float] = ...,
gpu_type: Optional[str] = ...,
) -> Callable[
[Union[Callable[..., Any], TableFunction, Type[TableFunction]]],
DataFrameUDTFWrapper,
]:
...
[docs]def udtf(
func: Optional[Union[Callable[..., Any], TableFunction, Type[TableFunction]]] = None,
*,
return_dtype: Optional[DataTypeLike] = None,
deterministic: bool = True,
name: Optional[str] = None,
concurrency: Optional[int] = None,
num_gpus: Optional[float] = None,
gpu_type: Optional[str] = None,
) -> Union[
DataFrameUDTFWrapper,
Callable[[Union[Callable[..., Any], TableFunction, Type[TableFunction]]],
DataFrameUDTFWrapper],
]:
"""
Create a user-defined table function for DataFrame operations.
This decorator wraps:
- Plain Python functions and lambdas
- TableFunction subclasses or instances, using the eval method
- Plain callable class instances, using the __call__ method
The wrapped object must emit zero or more rows by returning an
``Iterable``/``Iterator``/list, or by yielding from a generator. Each
emitted row may be a scalar value, tuple/Row, or dict when output fields are
named.
Args:
func: The Python function, TableFunction instance/subclass, or
callable class instance to wrap as a UDTF.
return_dtype: Optional emitted row type. Can be:
- DataType instance (e.g., DataType.string())
- Python type (e.g., str, int)
- Struct DataType for multi-column output
If not specified, inferred from the function's return type
hint (eval for TableFunction, __call__ for callable
class instances).
deterministic: Whether the function is deterministic (default: True).
name: Optional name for the UDTF.
concurrency: Optional concurrency (parallelism) for the UDTF operator.
If specified, the operator running this UDTF will use this
parallelism.
num_gpus: Optional number of GPUs requested for this UDTF (e.g., 0.5, 1.0).
gpu_type: Optional GPU type (e.g., 'A10', 'V100'). Required when
num_gpus is set.
Returns:
A DataFrameUDTFWrapper that can be called with Expressions and passed to
join_lateral, or passed directly to flat_map.
Example::
>>> from typing import Iterator, TypedDict
>>> from pyflink.dataframe import DataType, col, udtf
>>> from pyflink.table.udf import TableFunction
>>>
>>> # 1. Plain generator function with explicit return_dtype
>>> @udtf(return_dtype=DataType.string())
... def split_words(text):
... for word in text.split():
... yield word
>>>
>>> df.join_lateral(split_words(col("text")).alias("word"))
>>>
>>> # 2. Plain function with Iterator[T] type hint inference
>>> @udtf
... def chars(text: str) -> Iterator[str]:
... for ch in text:
... yield ch
>>>
>>> df.join_lateral(chars(col("text")).alias("ch"))
>>>
>>> # 3. TypedDict output supplies output column names
>>> class Token(TypedDict):
... word: str
... length: int
>>>
>>> @udtf
... def tokenize(text: str) -> Iterator[Token]:
... for word in text.split():
... yield {"word": word, "length": len(word)}
>>>
>>> df.join_lateral(tokenize(col("text")))
>>>
>>> # 4. Tuple output infers field types; provide output aliases
>>> from typing import Tuple
>>> @udtf
... def pair(x: int) -> Iterator[Tuple[int, str]]:
... yield x, str(x)
>>>
>>> df.join_lateral(pair(col("a")).alias("value", "text"))
>>>
>>> # 5. Parameterized decorator with resources
>>> @udtf(return_dtype=DataType.string(), concurrency=4,
... num_gpus=0.5, gpu_type="A10")
... def gpu_split(text):
... return text.split()
>>>
>>> df.join_lateral(gpu_split(col("text")).alias("word"))
>>>
>>> # 6. Direct call form
>>> split_words = udtf(lambda text: text.split(),
... return_dtype=DataType.string())
>>> df.join_lateral(split_words(col("text")).alias("word"))
>>>
>>> # 7. TableFunction instance or class
>>> class Split(TableFunction):
... def eval(self, text: str) -> Iterator[str]:
... for word in text.split():
... yield word
>>>
>>> split = udtf(Split())
>>> split = udtf(Split) # auto-instantiated
>>> df.join_lateral(split(col("text")).alias("word"))
>>>
>>> # 8. Callable class instance
>>> class Repeat:
... def __call__(self, value: str) -> Iterator[str]:
... yield value
... yield value
>>>
>>> repeat = udtf(Repeat())
>>> df.join_lateral(repeat(col("text")).alias("value"))
>>>
>>> # 9. Left outer lateral join semantics
>>> df.join_lateral(chars(col("text")).alias("ch"), ignore_empty=False)
"""
_validate_resource_overrides(concurrency, num_gpus, gpu_type)
def decorator(
f: Union[Callable[..., Any], TableFunction, Type[TableFunction]]
) -> DataFrameUDTFWrapper:
actual_func = _instantiate_table_function_if_needed(f)
hint_func = _get_type_hint_function(actual_func)
actual_return_dtype, has_named_fields = _infer_udtf_return_dtype(
hint_func, return_dtype)
table_func = (
_wrap_table_function_output(
actual_func, actual_return_dtype, has_named_fields)
if isinstance(actual_func, TableFunction)
else _wrap_udtf_function_output(
actual_func, actual_return_dtype, has_named_fields)
)
table_wrapper = _create_table_udtf(
table_func,
result_types=_to_table_udtf_result_types(actual_return_dtype),
deterministic=deterministic,
name=name,
concurrency=concurrency,
num_gpus=num_gpus,
gpu_type=gpu_type,
)
return DataFrameUDTFWrapper(
table_wrapper, actual_return_dtype, has_named_fields, hint_func)
if func is not None:
return decorator(func)
return decorator
# ======================== Function Normalization and Type Inference ========================
def _instantiate_table_function_if_needed(
func: Union[Callable[..., Any], TableFunction, Type[TableFunction]]
) -> Union[Callable[..., Any], TableFunction]:
"""Return a callable UDTF object, instantiating TableFunction classes."""
if inspect.isclass(func):
if issubclass(func, TableFunction):
return func()
raise TypeError(
"DataFrame udtf only supports plain functions or TableFunction "
"classes, got class {}".format(func.__name__))
if isinstance(func, TableFunction) or callable(func):
return func
raise TypeError(
"Invalid function: expected a callable or TableFunction, got {}".format(
type(func).__name__))
def _get_type_hint_function(
func: Union[Callable[..., Any], TableFunction]
) -> Callable[..., Any]:
"""Return the callable whose annotations describe the UDTF output."""
if isinstance(func, TableFunction):
return func.eval
if inspect.isfunction(func) or inspect.ismethod(func):
return func
call = getattr(func, "__call__", None)
if call is not None:
return call
return func
def _infer_udtf_return_dtype(
func: Callable[..., Any],
return_dtype: Optional[DataTypeLike],
) -> Tuple[DataType, bool]:
"""Resolve the UDTF output type and whether it has named fields."""
if return_dtype is not None:
dtype = _convert_to_dtype(return_dtype)
return dtype, isinstance(dtype._table_type, RowType)
try:
hints = get_type_hints(func)
except Exception:
hints = {}
if "return" not in hints:
name = getattr(func, "__name__", type(func).__name__)
raise ValueError(
"Function '{}' must have a return type hint or specify "
"return_dtype parameter explicitly.".format(name))
return _infer_udtf_dtype_from_hint(hints["return"])
def _infer_udtf_dtype_from_hint(type_hint: Any) -> Tuple[DataType, bool]:
"""Infer the emitted row DataType and naming flag from a return hint."""
emitted_type = _extract_emitted_type(type_hint)
return _infer_emitted_dtype(emitted_type)
def _extract_emitted_type(type_hint: Any) -> Any:
"""Return the item type emitted by Iterable/Iterator/Generator hints."""
origin = get_origin(type_hint)
args = get_args(type_hint)
iterable_origins = {Iterator, Iterable, AbcIterator, AbcIterable, list}
generator_origins = {Generator, AbcGenerator}
if origin in iterable_origins:
if not args:
raise TypeError(
"Cannot infer UDTF return type from an unparameterized "
"iterable.")
return args[0]
if origin in generator_origins:
if not args:
raise TypeError(
"Cannot infer UDTF return type from an unparameterized "
"generator.")
return args[0]
return type_hint
def _infer_emitted_dtype(emitted_type: Any) -> Tuple[DataType, bool]:
"""Convert one emitted item type into a DataType and naming flag."""
if _is_typeddict(emitted_type):
hints = get_type_hints(emitted_type)
fields = {
name: DataType._infer_from_type(hint)
for name, hint in hints.items()
}
return DataType.struct(fields), True
origin = get_origin(emitted_type)
if origin is dict:
raise TypeError(
"Cannot infer UDTF output fields from dict[K, V]. Use a "
"TypedDict return type hint or specify return_dtype explicitly.")
if origin in (tuple, Tuple):
args = get_args(emitted_type)
if not args:
raise TypeError(
"Cannot infer UDTF return type from tuple without type "
"arguments.")
if len(args) == 2 and args[1] is Ellipsis:
raise TypeError(
"Cannot infer UDTF return type from variable-length tuple.")
fields = [
("f{}".format(index), DataType._infer_from_type(arg))
for index, arg in enumerate(args)
]
return DataType.struct(fields), False
return DataType._infer_from_type(emitted_type), False
# ======================== Table Wrapper and Resource Metadata ========================
def _to_table_udtf_result_types(return_dtype: DataType):
"""Return the Table API result type representation for a DataFrame UDTF."""
return return_dtype._table_type
def _create_table_udtf(
func: Union[Callable[..., Any], TableFunction],
result_types: Any,
deterministic: bool,
name: Optional[str],
concurrency: Optional[int],
num_gpus: Optional[float],
gpu_type: Optional[str],
) -> UserDefinedTableFunctionWrapper:
"""Create the Table API UDTF wrapper and attach resource metadata."""
wrapper = table_udtf(
func,
result_types=result_types,
deterministic=deterministic,
name=name,
)
if concurrency is not None:
wrapper._concurrency = concurrency
if num_gpus is not None:
wrapper._num_gpus = num_gpus
if gpu_type is not None:
wrapper._gpu_type = gpu_type
return wrapper
def _copy_table_wrapper_with_resources(
wrapper: UserDefinedTableFunctionWrapper,
concurrency: Optional[int],
num_gpus: Optional[float],
gpu_type: Optional[str],
) -> UserDefinedTableFunctionWrapper:
"""Return a shallow wrapper copy with optional resource overrides."""
copied = copy.copy(wrapper)
copied._judf_placeholder = None
if concurrency is not None:
copied._concurrency = concurrency
if num_gpus is not None:
copied._num_gpus = num_gpus
if gpu_type is not None:
copied._gpu_type = gpu_type
return copied
def _validate_resource_overrides(
concurrency: Optional[int],
num_gpus: Optional[float],
gpu_type: Optional[str],
) -> None:
"""Validate UDTF resource overrides, raising ValueError when invalid."""
if concurrency is not None and (
not isinstance(concurrency, int) or concurrency <= 0):
raise ValueError(
"concurrency must be a positive integer, got: {}".format(
concurrency))
if num_gpus is not None and (
not isinstance(num_gpus, (int, float)) or num_gpus <= 0):
raise ValueError(
"num_gpus must be a positive number, got: {}".format(num_gpus))
if num_gpus is not None and gpu_type is None:
raise ValueError("gpu_type must be specified when num_gpus is set")
# ======================== Output Metadata and Validation ========================
def _get_output_arity(return_dtype: DataType) -> int:
"""Return the number of columns emitted by a UDTF output type."""
table_type = return_dtype._table_type
if isinstance(table_type, RowType):
return len(table_type.fields)
return 1
def _get_named_output_columns(return_dtype: DataType) -> List[str]:
"""Return named output columns from a RowType, or an empty list."""
table_type = return_dtype._table_type
if isinstance(table_type, RowType):
return list(table_type.field_names())
return []
def _get_adapter_output_columns(
return_dtype: DataType,
has_named_fields: bool,
) -> List[str]:
"""Return output column names used when adapting Python result rows."""
if has_named_fields:
return _get_named_output_columns(return_dtype)
return ["f{}".format(index) for index in range(_get_output_arity(return_dtype))]
def _resolve_flat_map_output_names(
return_dtype: DataType,
has_named_fields: bool,
) -> List[str]:
"""Return flat_map output names, rejecting unnamed multi-field outputs."""
arity = _get_output_arity(return_dtype)
if has_named_fields:
return _get_named_output_columns(return_dtype)
if arity == 1:
return ["f0"]
raise ValueError(
"flat_map requires named output fields for multi-field UDTFs. "
"Use a TypedDict return type hint or "
"return_dtype=DataType.struct({...}).")
def _validate_output_aliases(
output_aliases: List[str],
arity: Optional[int],
existing_columns: Optional[List[str]] = None,
) -> None:
"""Validate UDTF output aliases against arity and optional input columns."""
if not output_aliases:
raise ValueError("UDTF output aliases must not be empty")
if not all(isinstance(name, str) for name in output_aliases):
raise TypeError("UDTF output aliases must contain only strings")
if len(set(output_aliases)) != len(output_aliases):
raise ValueError("UDTF output aliases contain duplicate column names")
if existing_columns is not None:
conflicts = set(existing_columns).intersection(output_aliases)
if conflicts:
raise ValueError(
"UDTF output aliases conflict with existing DataFrame columns: "
"{}".format(sorted(conflicts)))
if arity is not None and len(output_aliases) != arity:
raise ValueError(
"UDTF output alias count must match output column count: "
"expected {}, got {}".format(arity, len(output_aliases)))
def _validate_flat_map_row_input_function(
func: Optional[Callable[..., Any]],
row_hint_validator: Callable[[Any], bool],
row_hint_name: str,
context: str,
) -> None:
"""Validate that a flat_map function can accept one row argument."""
if func is None:
return
try:
signature = inspect.signature(func)
except (TypeError, ValueError):
return
positional_params = [
param
for param in signature.parameters.values()
if param.kind in (
inspect.Parameter.POSITIONAL_ONLY,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
)
]
has_varargs = any(
param.kind == inspect.Parameter.VAR_POSITIONAL
for param in signature.parameters.values()
)
required_positional_params = [
param
for param in positional_params
if param.default is inspect.Parameter.empty
]
if not positional_params and not has_varargs:
raise ValueError(
"{} used with DataFrame.flat_map must accept one row argument."
.format(context))
if not has_varargs and len(required_positional_params) > 1:
raise ValueError(
"{} used with DataFrame.flat_map must accept one row argument; "
"call the UDTF with columns and use join_lateral for element-wise "
"UDTFs.".format(context))
if not positional_params:
return
hint = _get_first_parameter_type_hint(func, positional_params[0].name)
if hint is not None and not row_hint_validator(hint):
raise ValueError(
"{} used with DataFrame.flat_map receives one {} row argument, "
"but its first parameter is annotated as {}. Use a row-based "
"function for flat_map, or call the UDTF with columns and use "
"join_lateral for element-wise UDTFs.".format(
context, row_hint_name, _format_type_hint(hint)))
def _get_first_parameter_type_hint(
func: Callable[..., Any],
parameter_name: str,
) -> Optional[Any]:
"""Return the resolved first-parameter type hint, if one is present."""
try:
hints = get_type_hints(func)
except Exception:
hints = {}
hint = hints.get(parameter_name)
if hint is not None:
return hint
annotation = inspect.signature(func).parameters[parameter_name].annotation
if annotation is inspect.Parameter.empty:
return None
return annotation
def _is_wrapper_flat_map_row_hint(type_hint: Any) -> bool:
"""Return whether a type hint is compatible with wrapper flat_map Row input."""
origin = get_origin(type_hint)
return type_hint in (Any, object, Row) or origin is tuple
def _is_raw_flat_map_row_hint(type_hint: Any) -> bool:
"""Return whether a type hint is compatible with raw flat_map dict input."""
origin = get_origin(type_hint)
return type_hint in (Any, object, dict) or origin in (dict, AbcMapping)
def _format_type_hint(type_hint: Any) -> str:
"""Return a compact display string for a type hint."""
return getattr(type_hint, "__name__", str(type_hint))
def _allows_dict_as_scalar(
return_dtype: DataType,
has_named_fields: bool,
) -> bool:
"""Return whether dict output should be treated as one scalar map field."""
if has_named_fields:
return False
table_type = return_dtype._table_type
if isinstance(table_type, MapType):
return True
if isinstance(table_type, RowType) and len(table_type.fields) == 1:
return isinstance(table_type.field_types()[0], MapType)
return False
# ======================== User Result Adapters ========================
def _iter_user_results(result: Any):
"""Yield user-emitted rows from scalar, iterable, or empty results."""
if result is None:
return
if isinstance(result, (Row, tuple, dict, str, bytes, bytearray)):
yield result
return
if isinstance(result, AbcIterable):
for item in result:
yield item
return
yield result
def _row_from_user_result(
item: Any,
output_names: List[str],
has_named_fields: bool,
allow_dict_as_scalar: bool,
) -> Row:
"""Convert one user-emitted item into a Table API Row."""
if isinstance(item, Row):
return item
if isinstance(item, dict) and has_named_fields:
return Row(*[item[name] for name in output_names])
if isinstance(item, dict) and not allow_dict_as_scalar:
raise TypeError(
"UDTF emitted a dict, but the return type does not define named "
"fields or a map-typed scalar output. Use a TypedDict return type, "
"return_dtype=DataType.struct({...}), or DataType.map(...) for "
"map scalar output.")
if isinstance(item, tuple):
return Row(*item)
return Row(item)
def _wrap_udtf_function_output(
func: Callable[..., Any],
return_dtype: DataType,
has_named_fields: bool,
) -> Callable[..., Any]:
"""Return a callable that adapts plain-function UDTF results to Rows."""
output_names = _get_adapter_output_columns(return_dtype, has_named_fields)
allow_dict_as_scalar = _allows_dict_as_scalar(return_dtype, has_named_fields)
@functools.wraps(func)
def wrapper(*args: Any):
result = func(*args)
for item in _iter_user_results(result):
yield _row_from_user_result(
item, output_names, has_named_fields, allow_dict_as_scalar)
return wrapper
class _DataFrameTableFunctionAdapter(TableFunction):
"""Adapt TableFunction rows while preserving its lifecycle hooks.
DataFrame UDTFs accept dict/TypedDict-style emitted rows and normalize them
into positional Row values before they reach the Table runtime. Plain
functions can use a simple wrapper, but TableFunction instances may also
define open, close, finish, and is_deterministic. This adapter keeps those
hooks delegated to the original TableFunction while applying the same row
normalization to eval and finish results.
"""
def __init__(
self,
func: TableFunction,
output_names: List[str],
has_named_fields: bool,
allow_dict_as_scalar: bool,
) -> None:
self._func = func
self._output_names = output_names
self._has_named_fields = has_named_fields
self._allow_dict_as_scalar = allow_dict_as_scalar
self.__name__ = getattr(func, "__name__", func.__class__.__name__)
def open(self, function_context):
return self._func.open(function_context)
def close(self):
return self._func.close()
def is_deterministic(self) -> bool:
return self._func.is_deterministic()
def eval(self, *args):
return self._adapt_result(self._func.eval(*args))
def finish(self):
if not hasattr(self._func, "finish"):
return None
return self._adapt_result(self._func.finish())
def _adapt_result(self, result):
for item in _iter_user_results(result):
yield _row_from_user_result(
item,
self._output_names,
self._has_named_fields,
self._allow_dict_as_scalar,
)
def _wrap_table_function_output(
func: TableFunction,
return_dtype: DataType,
has_named_fields: bool,
) -> TableFunction:
"""Return a TableFunction adapter that normalizes emitted rows."""
output_names = _get_adapter_output_columns(return_dtype, has_named_fields)
allow_dict_as_scalar = _allows_dict_as_scalar(return_dtype, has_named_fields)
return _DataFrameTableFunctionAdapter(
func, output_names, has_named_fields, allow_dict_as_scalar)
def _wrap_flat_map_general(
func: Callable[[Dict[str, Any]], Any],
input_columns: List[str],
output_names: List[str],
has_named_fields: bool,
allow_dict_as_scalar: bool,
) -> Callable[..., Any]:
"""Return a row-input wrapper for DataFrame.flat_map callables."""
@functools.wraps(func)
def wrapper(row: Row):
input_dict = {name: row[index] for index, name in enumerate(input_columns)}
result = func(input_dict)
for item in _iter_user_results(result):
yield _row_from_user_result(
item, output_names, has_named_fields, allow_dict_as_scalar)
return wrapper
# ======================== DataFrame Operation Resolvers ========================
def _resolve_flat_map_udtf(
func: Union[Callable[[Dict[str, Any]], Any], DataFrameUDTFWrapper],
return_dtype: Optional[DataTypeLike],
input_columns: List[str],
concurrency: Optional[int] = None,
num_gpus: Optional[float] = None,
gpu_type: Optional[str] = None,
) -> Tuple[Expression, List[str]]:
"""Return the flat_map UDTF expression and its output column names."""
_validate_resource_overrides(concurrency, num_gpus, gpu_type)
if isinstance(func, DataFrameUDTFWrapper):
if return_dtype is not None:
raise ValueError(
"return_dtype must not be specified when func is a "
"DataFrameUDTFWrapper because the wrapper already has a "
"return type.")
_validate_flat_map_row_input_function(
func._input_func,
_is_wrapper_flat_map_row_hint,
"Row",
"DataFrameUDTFWrapper",
)
wrapper = _copy_table_wrapper_with_resources(
func._table_udtf_wrapper, concurrency, num_gpus, gpu_type)
output_names = _resolve_flat_map_output_names(
func.return_dtype, func._has_named_fields)
return (
_row_input_expression(wrapper),
output_names,
)
actual_func = _instantiate_table_function_if_needed(func)
if isinstance(actual_func, TableFunction):
raise TypeError("flat_map does not support TableFunction instances directly")
_validate_flat_map_row_input_function(
_get_type_hint_function(actual_func),
_is_raw_flat_map_row_hint,
"dict",
"callable",
)
actual_return_dtype, has_named_fields = _infer_udtf_return_dtype(
_get_type_hint_function(actual_func), return_dtype)
output_names = _resolve_flat_map_output_names(
actual_return_dtype, has_named_fields)
wrapped = _wrap_flat_map_general(
actual_func,
input_columns,
output_names,
has_named_fields,
_allows_dict_as_scalar(actual_return_dtype, has_named_fields),
)
table_wrapper = _create_table_udtf(
wrapped,
result_types=_to_table_udtf_result_types(actual_return_dtype),
deterministic=True,
name=(
actual_func.__name__
if hasattr(actual_func, "__name__") else type(actual_func).__name__
),
concurrency=concurrency,
num_gpus=num_gpus,
gpu_type=gpu_type,
)
return (
_row_input_expression(table_wrapper),
output_names,
)
def _row_input_expression(
wrapper: UserDefinedTableFunctionWrapper,
) -> Expression:
"""Return a row-input UDTF expression for flat_map."""
wrapper._set_takes_row_as_input()
return wrapper(with_columns(table_col("*")))
def _resolve_lateral_expression(
call: Union[Expression, DataFrameUDTFCall],
existing_columns: List[str],
) -> Expression:
"""Return a lateral UDTF expression, validating output column names.
Raw Expression arguments must carry aliases recorded by Expression.alias.
"""
if isinstance(call, DataFrameUDTFCall):
if call.output_aliases is not None:
output_names = call.output_aliases
expression = call.expression
elif call.has_named_fields:
output_names = _get_named_output_columns(call.return_dtype)
else:
raise ValueError(
"DataFrame UDTF calls passed to join_lateral must define "
"output column names with alias or named return_dtype fields, "
"for example: "
"df.join_lateral(my_udtf(col('x')).alias('out')).")
_validate_output_aliases(
output_names,
_get_output_arity(call.return_dtype),
existing_columns,
)
if call.output_aliases is None:
expression = call.expression.alias(
output_names[0], *output_names[1:])
return expression
elif isinstance(call, Expression):
alias_names = getattr(call, "_alias_names", None)
if alias_names is None:
raise ValueError(
"Expression arguments passed to join_lateral must define "
"output column names with alias, for example: "
"df.join_lateral(expr.call('split', col('x')).alias('out')).")
_validate_output_aliases(list(alias_names), None, existing_columns)
return call
else:
raise TypeError(
"table_function_call must be a DataFrameUDTFCall or Expression, "
"got {}".format(type(call).__name__))