From 1fc5d5a88282e80c45760a612ea8d3235fe8f4f6 Mon Sep 17 00:00:00 2001 From: auroflow Date: Wed, 26 Aug 2026 20:32:30 +0800 Subject: [PATCH 01/18] [FLINK-40431][python] Support scalar UDFs in DataFrame API Add expression-oriented general, async, and pandas scalar UDF support using the release-11 implementation structure adapted for community PyFlink. Generated-by: OpenAI Codex (GPT-5) --- .../reference/pyflink.dataframe/index.rst | 1 + .../docs/reference/pyflink.dataframe/udf.rst | 48 + flink-python/pyflink/dataframe/__init__.py | 2 + .../pyflink/dataframe/tests/test_udf.py | 741 ++++++++++++++ flink-python/pyflink/dataframe/udf.py | 902 ++++++++++++++++++ 5 files changed, 1694 insertions(+) create mode 100644 flink-python/docs/reference/pyflink.dataframe/udf.rst create mode 100644 flink-python/pyflink/dataframe/tests/test_udf.py create mode 100644 flink-python/pyflink/dataframe/udf.py diff --git a/flink-python/docs/reference/pyflink.dataframe/index.rst b/flink-python/docs/reference/pyflink.dataframe/index.rst index 5f0645c2baa38..197cfccdaca8a 100644 --- a/flink-python/docs/reference/pyflink.dataframe/index.rst +++ b/flink-python/docs/reference/pyflink.dataframe/index.rst @@ -26,6 +26,7 @@ This page gives an overview of all public PyFlink DataFrame APIs. :maxdepth: 1 dataframe + udf creation io sql diff --git a/flink-python/docs/reference/pyflink.dataframe/udf.rst b/flink-python/docs/reference/pyflink.dataframe/udf.rst new file mode 100644 index 0000000000000..4d046bd61b2c8 --- /dev/null +++ b/flink-python/docs/reference/pyflink.dataframe/udf.rst @@ -0,0 +1,48 @@ +.. ################################################################################ + 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 Scalar Functions +============================= + +Use :func:`pyflink.dataframe.udf` to apply Python code to one or more DataFrame +columns. A scalar UDF produces one logical output column and can be used in +:meth:`~pyflink.dataframe.DataFrame.with_column`, +:meth:`~pyflink.dataframe.DataFrame.with_columns`, and +:meth:`~pyflink.dataframe.DataFrame.select`. + +DataFrame scalar UDFs support synchronous, asynchronous, and pandas-vectorized +callables. See :func:`pyflink.dataframe.udf` for declaration forms, type +inference, execution modes, and examples. + +API Reference +============= + +.. currentmodule:: pyflink.dataframe + +.. autosummary:: + :toctree: api/ + + udf + +.. currentmodule:: pyflink.dataframe.udf + +.. autosummary:: + :toctree: api/ + + DataFrameUDFWrapper diff --git a/flink-python/pyflink/dataframe/__init__.py b/flink-python/pyflink/dataframe/__init__.py index 326b139efc9da..50496a3fa505b 100644 --- a/flink-python/pyflink/dataframe/__init__.py +++ b/flink-python/pyflink/dataframe/__init__.py @@ -55,6 +55,7 @@ from pyflink.dataframe.datatype import DataType from pyflink.dataframe.io import read_generic from pyflink.dataframe.sql import sql +from pyflink.dataframe.udf import udf __all__ = [ "DataFrame", @@ -62,6 +63,7 @@ "DataType", "col", "lit", + "udf", "from_arrow", "from_dict", "from_pandas", diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py new file mode 100644 index 0000000000000..78e0cb30cde43 --- /dev/null +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -0,0 +1,741 @@ +################################################################################ +# 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. +################################################################################ + +import functools +import inspect +import unittest +from dataclasses import dataclass +from typing import TypedDict + +import pandas as pd +import pyarrow as pa +import pyflink.dataframe as pf +from pyflink.common import Row +from pyflink.table import DataTypes as TableDataTypes +from pyflink.table.types import RowType +from pyflink.table.udf import AsyncScalarFunction, ScalarFunction +from pyflink.testing.test_case_utils import ( + PyFlinkDataFrameUTTestCase, + PyFlinkStreamDataFrameTestCase, +) + + +class DataFrameUDFDeclarationTests(unittest.TestCase): + def test_function_declarations_return_types_and_metadata(self): + class Details(TypedDict): + label: str + scores: list[int] + + class Result(TypedDict): + id: int + details: Details + + def add_one(value: int) -> int: + """Add one to a value.""" + return value + 1 + + def identity(value): + return value + + def describe(value: int) -> Result: + return { + "id": value, + "details": {"label": str(value), "scores": [value]}, + } + + decorated = pf.udf(add_one) + + from pyflink.dataframe.udf import DataFrameUDFWrapper + + self.assertIsInstance(decorated, DataFrameUDFWrapper) + self.assertFalse(hasattr(pf, "DataFrameUDFWrapper")) + self.assertEqual(decorated.return_dtype, pf.DataType.int64()) + self.assertEqual(decorated.__name__, "add_one") + self.assertEqual(decorated.__doc__, "Add one to a value.") + self.assertIs(decorated.__wrapped__, add_one) + + configured = pf.udf(return_dtype=pf.DataType.string())( + lambda value: str(value) + ) + direct = pf.udf(functools.partial(add_one), name="partial_add_one") + + self.assertEqual(configured.return_dtype, pf.DataType.string()) + self.assertEqual(direct.return_dtype, pf.DataType.int64()) + self.assertEqual(direct.__name__, "partial_add_one") + + declarations = [ + ( + "Python type", + lambda: pf.udf(identity, return_dtype=int), + pf.DataType.int64(), + ), + ( + "nested TypedDict annotation", + lambda: pf.udf(describe), + pf.DataType.struct( + { + "id": pf.DataType.int64(), + "details": pf.DataType.struct( + { + "label": pf.DataType.string(), + "scores": pf.DataType.list(pf.DataType.int64()), + } + ), + } + ), + ), + ] + for case_name, declare, expected in declarations: + with self.subTest(case=case_name): + self.assertEqual(declare().return_dtype, expected) + + def test_callable_classes_and_instances_infer_from_invocation_method(self): + plain_constructor_calls = [] + scalar_constructor_calls = [] + + class AddOne: + def __init__(self): + plain_constructor_calls.append("AddOne") + + def __call__(self, value: int) -> int: + return value + 1 + + class AddOffset: + def __init__(self, offset): + self.offset = offset + + def __call__(self, value: int) -> int: + return value + self.offset + + class NamedCallable: + __name__ = "configured_add" + + def __call__(self, value: int) -> int: + return value + 1 + + class Double(ScalarFunction): + def __init__(self): + scalar_constructor_calls.append("Double") + + def eval(self, value: int) -> int: + return value * 2 + + class AsyncDouble(AsyncScalarFunction): + def __init__(self): + scalar_constructor_calls.append("AsyncDouble") + + async def eval(self, value: int) -> int: + return value * 2 + + named_callable = NamedCallable() + double_instance = Double() + async_double_instance = AsyncDouble() + scalar_constructor_calls.clear() + callables = [ + AddOne, + AddOffset(2), + named_callable, + Double, + double_instance, + AsyncDouble, + async_double_instance, + ] + for source in callables: + with self.subTest(source=source): + decorated = pf.udf(source) + self.assertEqual(decorated.return_dtype, pf.DataType.int64()) + + self.assertEqual(plain_constructor_calls, []) + self.assertEqual(scalar_constructor_calls, ["Double", "AsyncDouble"]) + self.assertEqual(pf.udf(named_callable).__name__, "configured_add") + + decorated_class = pf.udf(Double) + self.assertIs(decorated_class.__wrapped__, Double) + self.assertEqual(decorated_class.__qualname__, Double.__qualname__) + + def test_func_type_resolution_and_async_detection(self): + def pandas_add_one(values: pd.Series) -> pd.Series: + return values + 1 + + def with_pandas_context(context: pd.Series, value: int) -> int: + return value + + def pandas_forward_reference(values): + return values + + pandas_forward_reference.__annotations__["values"] = "pandas.Series" + + def mixed(values: pd.Series, offset: int): + return values + offset + + def arrow_add_one(values: pa.Array) -> pa.Array: + return pa.array([value.as_py() + 1 for value in values]) + + async def async_add_one(value: int) -> int: + return value + 1 + + async def async_pandas(values: pd.Series) -> pd.Series: + return values + 1 + + declarations = [ + ( + "inferred pandas", + lambda: pf.udf(pandas_add_one, return_dtype=pf.DataType.int64()), + "pandas", + False, + ), + ( + "bound pandas annotation is ignored", + lambda: pf.udf( + functools.partial(with_pandas_context, pd.Series([1])), + ), + "general", + False, + ), + ( + "pandas forward reference", + lambda: pf.udf( + pandas_forward_reference, + return_dtype=pf.DataType.int64(), + ), + "pandas", + False, + ), + ( + "any pandas annotation selects pandas", + lambda: pf.udf(mixed, return_dtype=pf.DataType.int64()), + "pandas", + False, + ), + ( + "explicit general wins", + lambda: pf.udf( + pandas_add_one, + return_dtype=pf.DataType.int64(), + func_type="general", + ), + "general", + False, + ), + ( + "pyarrow annotations remain general", + lambda: pf.udf(arrow_add_one, return_dtype=pf.DataType.int64()), + "general", + False, + ), + ( + "async general", + lambda: pf.udf(async_add_one), + "general", + True, + ), + ] + for case_name, declare, expected_type, expected_async in declarations: + with self.subTest(case=case_name): + wrapped = declare() + self.assertEqual(wrapped._func_type, expected_type) + self.assertEqual(wrapped._is_async, expected_async) + + invalid_declarations = [ + ( + "async inferred pandas", + lambda: pf.udf(async_pandas, return_dtype=pf.DataType.int64()), + ValueError, + "Async scalar functions", + ), + ( + "async explicit pandas", + lambda: pf.udf( + async_add_one, + return_dtype=pf.DataType.int64(), + func_type="pandas", + ), + ValueError, + "Async scalar functions", + ), + ] + for case_name, declare, error_type, message in invalid_declarations: + with self.subTest(case=case_name): + with self.assertRaisesRegex(error_type, message): + declare() + + def test_determinism_and_name_metadata(self): + class NonDeterministic(ScalarFunction): + def eval(self, value: int) -> int: + return value + + def is_deterministic(self): + return False + + class DefaultDeterministic(ScalarFunction): + def eval(self, value: int) -> int: + return value + + instance = NonDeterministic() + declarations = [ + ( + "matching instance metadata", + lambda: pf.udf(instance, deterministic=False), + False, + ), + ("class default", lambda: pf.udf(DefaultDeterministic), True), + ( + "class matching metadata", + lambda: pf.udf(NonDeterministic, deterministic=False), + False, + ), + ] + for case_name, declare, expected in declarations: + with self.subTest(case=case_name): + self.assertEqual(declare()._deterministic, expected) + + self.assertIs( + inspect.signature(pf.udf).parameters["deterministic"].default, + True, + ) + with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"): + pf.udf(instance) + with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"): + pf.udf(NonDeterministic) + + named = pf.udf(instance, deterministic=False, name="identity") + self.assertEqual(named.__name__, "identity") + self.assertEqual(named._table_udf_wrapper._name, "identity") + + def test_general_structured_results_are_normalized_recursively(self): + from pyflink.dataframe.udf import _normalize_user_value + + class Details: + __slots__ = ("label", "scores") + + def __init__(self, label, scores): + self.label = label + self.scores = scores + + class ItemsOnly: + def __init__(self, items): + self._items = items + + def items(self): + return self._items + + @dataclass + class Result: + id: int + details: Details + attributes: dict + + return_dtype = pf.DataType.struct( + { + "id": pf.DataType.int64(), + "details": pf.DataType.struct( + { + "label": pf.DataType.string(), + "scores": pf.DataType.list(pf.DataType.int64()), + } + ), + "attributes": pf.DataType.map( + pf.DataType.string(), pf.DataType.int64() + ), + } + ) + table_type = return_dtype._to_table_data_type() + self.assertIsInstance(table_type, RowType) + + cases = [ + ( + "mapping", + { + "id": 1, + "details": {"scores": (2, 3), "ignored": "extra"}, + "attributes": [("answer", 42)], + "ignored": "extra", + }, + Row( + id=1, + details=Row(label=None, scores=[2, 3]), + attributes={"answer": 42}, + ), + ), + ( + "named row", + Row( + id=4, + details=Row(label="named", scores=[5]), + attributes={"count": 6}, + ), + Row( + id=4, + details=Row(label="named", scores=[5]), + attributes={"count": 6}, + ), + ), + ( + "positional list and tuple", + [7, ("positional", (8, 9)), {"count": 10}], + Row( + id=7, + details=Row(label="positional", scores=[8, 9]), + attributes={"count": 10}, + ), + ), + ( + "dataclass and attribute objects", + Result( + id=11, + details=Details(label="object", scores=[12]), + attributes=ItemsOnly([("count", 13)]), + ), + Row( + id=11, + details=Row(label="object", scores=[12]), + attributes={"count": 13}, + ), + ), + ] + for case_name, value, expected in cases: + with self.subTest(case=case_name): + self.assertEqual( + _normalize_user_value(value, table_type), expected + ) + + with self.assertRaisesRegex(ValueError, "Expected 3 value"): + _normalize_user_value((1, 2), table_type) + with self.assertRaisesRegex(TypeError, "Expected a Mapping"): + _normalize_user_value(object(), table_type) + + def test_invalid_declarations_fail_eagerly(self): + def missing_return(value): + return value + + def pandas_identity(values: pd.Series) -> pd.Series: + return values + + class RequiresArgument: + def __init__(self, value): + self.value = value + + def __call__(self, other: int) -> int: + return other + self.value + + class NotCallable: + pass + + invalid_declarations = [ + ( + "not callable", + lambda: pf.udf(42, return_dtype=pf.DataType.int64()), + TypeError, + "func must be callable", + ), + ( + "non-callable class", + lambda: pf.udf(NotCallable, return_dtype=pf.DataType.int64()), + TypeError, + "func must be callable", + ), + ( + "missing return", + lambda: pf.udf(missing_return), + TypeError, + "Cannot infer return_dtype", + ), + ( + "Table return type", + lambda: pf.udf( + missing_return, return_dtype=TableDataTypes.BIGINT() + ), + TypeError, + "return_dtype must be", + ), + ( + "required constructor argument", + lambda: pf.udf(RequiresArgument), + TypeError, + "zero-argument constructor", + ), + ( + "invalid determinism", + lambda: pf.udf( + missing_return, + return_dtype=pf.DataType.int64(), + deterministic=1, + ), + TypeError, + "deterministic must be", + ), + ( + "invalid name", + lambda: pf.udf( + missing_return, + return_dtype=pf.DataType.int64(), + name=1, + ), + TypeError, + "name must be", + ), + ( + "empty name", + lambda: pf.udf( + missing_return, + return_dtype=pf.DataType.int64(), + name="", + ), + ValueError, + "name must not be empty", + ), + ( + "arrow func type", + lambda: pf.udf( + missing_return, + return_dtype=pf.DataType.int64(), + func_type="arrow", + ), + ValueError, + "func_type must be one of", + ), + ( + "pandas return type required", + lambda: pf.udf(pandas_identity), + TypeError, + "return_dtype is required", + ), + ] + for case_name, declare, error_type, message in invalid_declarations: + with self.subTest(case=case_name): + with self.assertRaisesRegex(error_type, message): + declare() + + +class DataFrameUDFAdapterTests(unittest.TestCase): + def test_scalar_function_lifecycle_and_cleanup(self): + from pyflink.dataframe.udf import ( + _DataFrameScalarFunctionAdapter, + _UDFUsage, + ) + + events = [] + + class LifecycleFunction(ScalarFunction): + def __init__(self): + events.append("init") + + def open(self, function_context): + events.append(("open", function_context)) + + def eval(self, value): + return value + 1 + + def close(self): + events.append("close") + + context = object() + adapter = _DataFrameScalarFunctionAdapter( + LifecycleFunction(), + pf.DataType.int64(), + True, + _UDFUsage.EXPRESSION, + "general", + ) + + with self.assertRaisesRegex(RuntimeError, "before open"): + adapter.eval(1) + + adapter.open(context) + self.assertEqual(adapter.eval(1), 2) + adapter.close() + + with self.assertRaisesRegex(RuntimeError, "before open"): + adapter.eval(1) + + adapter.open(context) + self.assertEqual(adapter.eval(2), 3) + adapter.close() + self.assertEqual( + events, + [ + "init", + ("open", context), + "close", + ("open", context), + "close", + ], + ) + + deferred_constructor_calls = [] + + class DeferredCallable: + def __init__(self): + deferred_constructor_calls.append("init") + + def __call__(self, value): + return value + 1 + + deferred_adapter = _DataFrameScalarFunctionAdapter( + DeferredCallable, + pf.DataType.int64(), + True, + _UDFUsage.EXPRESSION, + "general", + ) + deferred_adapter.open(context) + self.assertEqual(deferred_adapter.eval(1), 2) + deferred_adapter.close() + deferred_adapter.open(context) + self.assertEqual(deferred_adapter.eval(2), 3) + deferred_adapter.close() + self.assertEqual(deferred_constructor_calls, ["init", "init"]) + + class FailingCloseFunction(ScalarFunction): + def eval(self, value): + return value + + def close(self): + raise RuntimeError("close failed") + + failing_adapter = _DataFrameScalarFunctionAdapter( + FailingCloseFunction(), + pf.DataType.int64(), + True, + _UDFUsage.EXPRESSION, + "general", + ) + failing_adapter.open(context) + with self.assertRaisesRegex(RuntimeError, "close failed"): + failing_adapter.close() + with self.assertRaisesRegex(RuntimeError, "before open"): + failing_adapter.eval(1) + + +class DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase): + def test_with_columns_binds_expressions_and_resolves_output_schema(self): + sql_typed = pf.udf(lambda value: value, return_dtype="BIGINT") + + @pf.udf(name="render_value") + def render(value: int, suffix: str) -> str: + return f"{value}{suffix}" + + @pf.udf( + return_dtype=pf.DataType.struct( + { + "value": pf.DataType.int64(), + "tags": pf.DataType.list(pf.DataType.string()), + } + ) + ) + def describe(value): + return {"value": value, "tags": [str(value)]} + + result = pf.from_records([(1,)], schema=["id"]).with_columns( + rendered=render(pf.col("id"), "-literal"), + description=describe(pf.col("id")), + sql_value=sql_typed(pf.col("id")), + ) + + self.assert_dataframe_schema( + result, + ["id", "rendered", "description", "sql_value"], + [ + TableDataTypes.BIGINT(), + TableDataTypes.STRING(), + TableDataTypes.ROW( + [ + TableDataTypes.FIELD("value", TableDataTypes.BIGINT()), + TableDataTypes.FIELD( + "tags", TableDataTypes.ARRAY(TableDataTypes.STRING()) + ), + ] + ), + TableDataTypes.BIGINT(), + ], + ) + + +class DataFrameUDFITCase(PyFlinkStreamDataFrameTestCase): + def test_supported_scalar_udfs_in_one_job(self): + @dataclass + class Details: + doubled: int + labels: list + + @pf.udf + def add_one(value: int) -> int: + return value + 1 + + @pf.udf + async def add_two(value: int) -> int: + return value + 2 + + @pf.udf(return_dtype=pf.DataType.int64(), func_type="pandas") + def add_three(values: pd.Series) -> pd.Series: + return values + 3 + + @pf.udf( + return_dtype=pf.DataType.struct( + { + "doubled": pf.DataType.int64(), + "labels": pf.DataType.list(pf.DataType.string()), + } + ) + ) + def details(value): + return Details(doubled=value * 2, labels=[str(value)]) + + class DeferredCallable: + def __call__(self, value: int) -> int: + return value + 4 + + class OpenedScalarFunction(ScalarFunction): + def open(self, function_context): + self._increment = 5 + + def eval(self, value: int) -> int: + return value + self._increment + + class ClassNonDeterministic(ScalarFunction): + def eval(self, value: int) -> int: + return value + 6 + + def is_deterministic(self): + return False + + deferred = pf.udf(DeferredCallable) + scalar_instance = pf.udf(OpenedScalarFunction()) + scalar_class = pf.udf(ClassNonDeterministic, deterministic=False) + + result = ( + pf.from_records([(1,)], schema=["id"]) + .with_columns(async_value=add_two(pf.col("id"))) + .with_columns( + sync_value=add_one(pf.col("id")), + pandas_value=add_three(pf.col("id")), + details=details(pf.col("id")), + deferred_value=deferred(pf.col("id")), + scalar_value=scalar_instance(pf.col("id")), + scalar_class_value=scalar_class(pf.col("id")), + ) + ) + + self.assertEqual( + result.collect(), + [Row(1, 3, 2, 4, Row(2, ["1"]), 5, 6, 7)], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py new file mode 100644 index 0000000000000..525e2b3f3f1ff --- /dev/null +++ b/flink-python/pyflink/dataframe/udf.py @@ -0,0 +1,902 @@ +################################################################################ +# 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 scalar functions for the DataFrame API.""" + +import functools +import inspect +from collections.abc import Mapping +from enum import Enum +from typing import ( + Any, + Callable, + Dict, + Optional, + Set, + Tuple, + Type, + Union, + cast, + get_type_hints, + overload, +) + +from pyflink.common import Row +from pyflink.dataframe.datatype import DataType +from pyflink.table.expression import Expression +from pyflink.table.expressions import call as table_call +from pyflink.table.types import ArrayType, MapType, RowType +from pyflink.table.udf import ( + AsyncScalarFunction, + ScalarFunction, + UserDefinedFunction, + UserDefinedFunctionWrapper, + udf as table_udf, +) +from pyflink.util.api_stability_decorators import PublicEvolving + +__all__ = ["DataFrameUDFWrapper", "udf"] + +_UDFInput = Union[Callable[..., Any], ScalarFunction, AsyncScalarFunction, Type] +_DataTypeLike = Union[DataType, Type, str] + + +class _UDFUsage(Enum): + EXPRESSION = "expression" + MAP = "map" + MAP_BATCHES = "map_batches" + + +@PublicEvolving() +class DataFrameUDFWrapper: + """ + A callable DataFrame scalar UDF declaration. + + Instances are created with :func:`udf` and can be called with DataFrame + expressions or Python literals to produce an expression. + + Example:: + + >>> import pyflink.dataframe as pf + >>> @pf.udf + ... def add_one(value: int) -> int: + ... return value + 1 + >>> expression = add_one(pf.col("value")) + + .. versionadded:: 2.4.0 + """ + + _func: _UDFInput + _return_dtype: DataType + _deterministic: bool + _func_type: str + _is_async: bool + _cached_table_udf_wrapper: Optional[UserDefinedFunctionWrapper] + _frozen: bool + __name__: str + + def __init__( + self, + func: _UDFInput, + return_dtype: DataType, + deterministic: bool, + name: str, + func_type: str, + is_async: bool, + metadata_source: Optional[_UDFInput] = None, + ) -> None: + object.__setattr__(self, "_func", func) + object.__setattr__(self, "_return_dtype", return_dtype) + object.__setattr__(self, "_deterministic", deterministic) + object.__setattr__(self, "_func_type", func_type) + object.__setattr__(self, "_is_async", is_async) + object.__setattr__(self, "_cached_table_udf_wrapper", None) + + declaration = func if metadata_source is None else metadata_source + declaration_metadata = _unwrap_partial(declaration) + functools.update_wrapper(self, declaration_metadata, updated=()) + object.__setattr__(self, "__name__", name) + object.__setattr__(self, "__wrapped__", declaration) + object.__setattr__(self, "_frozen", True) + + def __setattr__(self, name: str, value: Any) -> None: + if getattr(self, "_frozen", False): + raise AttributeError("DataFrameUDFWrapper declarations are immutable.") + object.__setattr__(self, name, value) + + @PublicEvolving() + def __call__(self, *args: Any) -> Expression: + """ + Create an expression that calls this UDF. + + Example:: + + >>> import pyflink.dataframe as pf + >>> @pf.udf + ... def add_one(value: int) -> int: + ... return value + 1 + >>> expression = add_one(pf.col("value")) + """ + return table_call(self._table_udf_wrapper, *args) + + @property + def _table_udf_wrapper(self) -> UserDefinedFunctionWrapper: + if self._cached_table_udf_wrapper is None: + object.__setattr__( + self, + "_cached_table_udf_wrapper", + self._create_table_udf_wrapper(_UDFUsage.EXPRESSION), + ) + return cast(UserDefinedFunctionWrapper, self._cached_table_udf_wrapper) + + def _create_table_udf_wrapper( + self, usage: _UDFUsage + ) -> UserDefinedFunctionWrapper: + adapter_type = ( + _DataFrameAsyncScalarFunctionAdapter + if self._is_async + else _DataFrameScalarFunctionAdapter + ) + actual_func = cast( + Union[ScalarFunction, AsyncScalarFunction], + adapter_type( + self._func, + self._return_dtype, + self._deterministic, + usage, + self._func_type, + ), + ) + return cast( + UserDefinedFunctionWrapper, + table_udf( + actual_func, + result_type=self._return_dtype._to_table_data_type(), + deterministic=self._deterministic, + name=self.__name__, + func_type=self._func_type, + ), + ) + + @property + @PublicEvolving() + def return_dtype(self) -> DataType: + """ + The logical result type of this UDF. + + Example:: + + >>> import pyflink.dataframe as pf + >>> @pf.udf + ... def add_one(value: int) -> int: + ... return value + 1 + >>> add_one.return_dtype == pf.DataType.int64() + True + """ + return self._return_dtype + + +@overload +def udf( + func: _UDFInput, + *, + return_dtype: Optional[_DataTypeLike] = ..., + deterministic: bool = ..., + name: Optional[str] = ..., + func_type: Optional[str] = ..., +) -> DataFrameUDFWrapper: + ... + + +@overload +def udf( + func: None = ..., + *, + return_dtype: Optional[_DataTypeLike] = ..., + deterministic: bool = ..., + name: Optional[str] = ..., + func_type: Optional[str] = ..., +) -> Callable[[_UDFInput], DataFrameUDFWrapper]: + ... + + +@PublicEvolving() +def udf( + func: Optional[_UDFInput] = None, + *, + return_dtype: Optional[_DataTypeLike] = None, + deterministic: bool = True, + name: Optional[str] = None, + func_type: Optional[str] = None, +) -> Union[DataFrameUDFWrapper, Callable[[_UDFInput], DataFrameUDFWrapper]]: + """ + Create a scalar UDF for DataFrame expressions. + + The function may be synchronous or asynchronous. Pandas UDFs operate on + ``pandas.Series`` or ``pandas.DataFrame`` batches and must declare + ``return_dtype``. Plain callable class objects must have a zero-argument + constructor and are instantiated on the worker. + + A UDF can be declared with a bare decorator, a configured decorator, or a + direct call. General UDFs may infer ``return_dtype`` from the return + annotation of the function, ``__call__``, or ``eval``. A ``TypedDict`` + return annotation becomes a struct column:: + + >>> import pyflink.dataframe as pf + + >>> @pf.udf + ... def add_one(value: int) -> int: + ... return value + 1 + + >>> @pf.udf(return_dtype=str) + ... def as_text(value): + ... return str(value) + + >>> increment = pf.udf( + ... lambda value, amount: value + amount, + ... return_dtype="BIGINT", + ... ) + + >>> from typing import TypedDict + + >>> class LabeledValue(TypedDict): + ... value: int + ... label: str + + >>> @pf.udf + ... def describe(value: int) -> LabeledValue: + ... return {"value": value, "label": str(value)} + + Plain callable classes can be supplied as zero-argument class objects or + as configured instances. Class objects, including their ``__init__``, are + constructed during worker initialization, so expensive initialization is + deferred to the TaskManager:: + + >>> class AddOne: + ... def __call__(self, value: int) -> int: + ... return value + 1 + + >>> add_one_from_class = pf.udf(AddOne) + >>> add_one_from_instance = pf.udf(AddOne()) + + >>> @pf.udf + ... class ModelInference: + ... def __init__(self): + ... self.model = load_model() + ... def __call__(self, features: list[float]) -> float: + ... return self.model.predict(features) + + :class:`~pyflink.table.udf.ScalarFunction` and + :class:`~pyflink.table.udf.AsyncScalarFunction` class objects and instances + are also supported. Their logical result type is inferred from ``eval`` + when it is not given explicitly. Class objects are instantiated on the + client; their ``open`` and ``close`` methods still run on the worker:: + + >>> from pyflink.table.udf import AsyncScalarFunction, ScalarFunction + + >>> class AddOneFunction(ScalarFunction): + ... def eval(self, value: int) -> int: + ... return value + 1 + + >>> add_one_class = pf.udf(AddOneFunction) + >>> add_one_instance = pf.udf(AddOneFunction()) + + >>> class AsyncLookup(AsyncScalarFunction): + ... async def eval(self, key: int) -> str: + ... return await lookup(key) + + >>> async_lookup = pf.udf(AsyncLookup) + + Plain ``async def`` functions and callable objects with an asynchronous + ``__call__`` use general asynchronous execution:: + + >>> @pf.udf + ... async def async_add_one(value: int) -> int: + ... return value + 1 + + Pandas UDFs receive and return ``pandas.Series`` or ``pandas.DataFrame`` + batches and always require an explicit logical ``return_dtype``. Pandas + mode can be selected explicitly, or inferred from a pandas container + annotation on any unbound parameter or the return value:: + + >>> import pandas as pd + + >>> @pf.udf(return_dtype=pf.DataType.int64(), func_type="pandas") + ... def pandas_add_one(values): + ... return values + 1 + + >>> @pf.udf(return_dtype=pf.DataType.int64()) + ... def inferred_pandas_add_one(values: pd.Series) -> pd.Series: + ... return values + 1 + + A declared UDF is called with DataFrame expressions or Python literals to + produce a single-column expression:: + + >>> df = pf.from_records([(1,), (2,)], schema=["value"]) + + >>> result = df.with_columns( + ... next_value=add_one(pf.col("value")), + ... incremented=increment(pf.col("value"), 2), + ... ) + + :param func: Function, callable object, scalar UDF instance, or zero-argument + callable/scalar-UDF class. + :param return_dtype: DataFrame logical type, Python type, or SQL type string. + General UDFs may infer it from a return annotation; + pandas UDFs require it. + :param deterministic: Whether equal inputs always produce equal results. + Must agree with scalar-function metadata. + :param name: Non-empty function identity used by the Table planner. + :param func_type: ``"general"`` or ``"pandas"``. If omitted, any unbound + pandas container annotation selects pandas mode. + :return: A :class:`DataFrameUDFWrapper`, or a decorator when ``func`` is omitted. + + .. versionadded:: 2.4.0 + """ + + def decorator(f: _UDFInput) -> DataFrameUDFWrapper: + _validate_scalar_udf_source(f) + ( + actual_func, + inspection_target, + skip_first_parameter, + is_async, + ) = _resolve_udf_source(f) + actual_func_type = ( + func_type + if func_type is not None + else _detect_func_type( + inspection_target, skip_first=skip_first_parameter + ) + ) + _validate_scalar_udf_options(actual_func_type, return_dtype, is_async) + actual_return_dtype = _infer_return_dtype(inspection_target, return_dtype) + actual_deterministic = _resolve_deterministic(actual_func, deterministic) + actual_name = _resolve_name(actual_func, name) + + return DataFrameUDFWrapper( + actual_func, + actual_return_dtype, + actual_deterministic, + actual_name, + actual_func_type, + is_async, + metadata_source=f, + ) + + return decorator if func is None else decorator(func) + + +# ======================== Declaration Validation ======================== + + +def _validate_scalar_udf_source(func: Any) -> None: + if inspect.isclass(func): + if issubclass(func, UserDefinedFunction) and not issubclass( + func, (ScalarFunction, AsyncScalarFunction) + ): + raise TypeError(f"func must be a scalar UDF, got {func.__name__}.") + if not issubclass( + func, (ScalarFunction, AsyncScalarFunction) + ) and not _has_custom_call(func): + raise TypeError(f"func must be callable, got {func.__name__}.") + return + if isinstance(func, UserDefinedFunction) and not isinstance( + func, (ScalarFunction, AsyncScalarFunction) + ): + raise TypeError(f"func must be a scalar UDF, got {type(func).__name__}.") + if not isinstance(func, (ScalarFunction, AsyncScalarFunction)) and not callable(func): + raise TypeError(f"func must be callable, got {type(func).__name__}.") + + +def _validate_scalar_udf_options( + func_type: str, + return_dtype: Optional[_DataTypeLike], + is_async: bool, +) -> None: + if func_type not in ("general", "pandas"): + raise ValueError( + f"The func_type must be one of 'general, pandas', got {func_type}." + ) + if return_dtype is None and func_type == "pandas": + raise TypeError( + "return_dtype is required for pandas UDFs because pandas container " + "annotations do not describe the logical result type." + ) + if is_async and func_type == "pandas": + raise ValueError( + "Async scalar functions do not support pandas func_type. " + "Use func_type='general'." + ) + + +# ======================== Callable Inspection and Resolution ======================== + + +def _has_custom_call(cls: Type) -> bool: + """Check whether a class defines ``__call__`` in its MRO.""" + return any("__call__" in base.__dict__ for base in cls.__mro__ if base is not object) + + +def _get_callable_class_hint_method( + func_class: Type, method_name: str = "__call__" +) -> Tuple[Optional[Callable[..., Any]], bool]: + """Return a class method that can be inspected without constructing the class.""" + descriptor = inspect.getattr_static(func_class, method_name, None) + if isinstance(descriptor, staticmethod): + return cast(Callable[..., Any], descriptor.__func__), False + if isinstance(descriptor, classmethod): + return cast(Callable[..., Any], descriptor.__func__), True + if inspect.isroutine(descriptor): + return cast(Callable[..., Any], descriptor), True + return None, False + + +def _is_class_udf(func: Any) -> bool: + """Check whether ``func`` uses a class-based scalar UDF declaration form.""" + if isinstance(func, functools.partial): + return False + if isinstance(func, (ScalarFunction, AsyncScalarFunction)): + return True + if inspect.isclass(func): + if issubclass(func, (ScalarFunction, AsyncScalarFunction)): + return True + return _has_custom_call(func) + if not inspect.isroutine(func) and hasattr(func, "__call__"): + return _has_custom_call(type(func)) + return False + + +def _resolve_udf_source( + func: _UDFInput, +) -> Tuple[_UDFInput, Callable[..., Any], bool, bool]: + """Resolve the runtime source and callable that describes its declaration.""" + if not _is_class_udf(func): + callable_func = cast(Callable[..., Any], func) + return func, callable_func, False, _is_async_callable(callable_func) + + skip_first_parameter = False + if inspect.isclass(func): + _validate_zero_argument_class(func) + if issubclass(func, (ScalarFunction, AsyncScalarFunction)): + actual_func = func() + hint_method = actual_func.eval + is_async = isinstance( + actual_func, AsyncScalarFunction + ) or inspect.iscoroutinefunction(hint_method) + return actual_func, hint_method, False, is_async + + class_hint_method, skip_first_parameter = _get_callable_class_hint_method( + func + ) + if class_hint_method is None: + raise TypeError( + f"Callable class '{func.__name__}': __call__ must be defined as a method." + ) + hint_method = class_hint_method + is_async = inspect.iscoroutinefunction(hint_method) + elif isinstance(func, (ScalarFunction, AsyncScalarFunction)): + hint_method = func.eval + is_async = isinstance(func, AsyncScalarFunction) or inspect.iscoroutinefunction( + hint_method + ) + else: + hint_method = cast(Callable[..., Any], getattr(func, "__call__")) + is_async = inspect.iscoroutinefunction(hint_method) + return func, hint_method, skip_first_parameter, is_async + + +def _validate_zero_argument_class(func_class: Type) -> None: + if inspect.isabstract(func_class): + raise TypeError(f"UDF class '{func_class.__name__}' must not be abstract.") + try: + constructor_signature = inspect.signature(func_class) + except (TypeError, ValueError) as exc: + raise TypeError( + f"Cannot verify that UDF class '{func_class.__name__}' has a zero-argument " + "constructor; pass a configured instance instead." + ) from exc + try: + constructor_signature.bind() + except TypeError as exc: + raise TypeError( + f"UDF class '{func_class.__name__}' must have a zero-argument constructor; " + "pass a configured instance instead." + ) from exc + + +def _infer_return_dtype( + func: Callable[..., Any], return_dtype: Optional[_DataTypeLike] +) -> DataType: + """Infer the DataFrame return type or validate its explicit declaration.""" + if return_dtype is not None: + return _convert_to_dtype(return_dtype) + + hint_func = _get_callable_inspection_target(func) + hints = _get_callable_type_hints(hint_func) + if "return" not in hints: + func_name = _default_udf_name(func) + raise TypeError( + f"Cannot infer return_dtype for '{func_name}': add a return annotation " + "or specify return_dtype explicitly." + ) + return _data_type_from_type_hint(hints["return"]) + + +def _convert_to_dtype(dtype_like: _DataTypeLike) -> DataType: + if isinstance(dtype_like, DataType): + return dtype_like + if isinstance(dtype_like, str): + return DataType._from_sql(dtype_like) + try: + return DataType._from_type_hint(dtype_like) + except TypeError as exc: + raise TypeError( + "return_dtype must be a DataFrame DataType, Python type, or SQL " + f"type string, got {type(dtype_like).__name__}." + ) from exc + + +def _is_typed_dict(type_hint: Any) -> bool: + try: + from typing import is_typeddict + + if is_typeddict(type_hint): + return True + except ImportError: + pass + return ( + isinstance(type_hint, type) + and issubclass(type_hint, dict) + and hasattr(type_hint, "__required_keys__") + ) + + +def _data_type_from_type_hint(type_hint: Any) -> DataType: + if _is_typed_dict(type_hint): + return DataType.struct( + { + name: _data_type_from_type_hint(field_hint) + for name, field_hint in get_type_hints(type_hint).items() + } + ) + return DataType._from_type_hint(type_hint) + + +def _detect_func_type(func: Callable[..., Any], skip_first: bool = False) -> str: + """Detect pandas mode from an unbound pandas container annotation.""" + hint_func = _get_callable_inspection_target(func) + try: + import pandas as pd + except ImportError: + return "general" + hints = _get_callable_type_hints( + hint_func, fallback_globals={"pandas": pd, "pd": pd} + ) + + try: + parameters = list(inspect.signature(hint_func).parameters) + except (TypeError, ValueError): + parameters = [] + ignored_hint_names: Set[str] = set() + if skip_first and parameters: + ignored_hint_names.add(parameters[0]) + if isinstance(func, functools.partial): + try: + ignored_hint_names.update( + inspect.signature(hint_func) + .bind_partial(*func.args, **(func.keywords or {})) + .arguments + ) + except (TypeError, ValueError): + pass + + pandas_types = (pd.Series, pd.DataFrame) + return ( + "pandas" + if any( + name not in ignored_hint_names and hint in pandas_types + for name, hint in hints.items() + ) + else "general" + ) + + +def _unwrap_partial(func: Any) -> Any: + while isinstance(func, functools.partial): + func = func.func + return func + + +def _get_callable_inspection_target( + func: Callable[..., Any], +) -> Callable[..., Any]: + target = _unwrap_partial(func) + if callable(target) and not inspect.isroutine(target) and not inspect.isclass(target): + return cast(Callable[..., Any], getattr(target, "__call__")) + return cast(Callable[..., Any], target) + + +def _get_callable_type_hints( + func: Callable[..., Any], fallback_globals: Optional[Dict[str, Any]] = None +) -> Dict[str, Any]: + try: + if fallback_globals is None: + return get_type_hints(func) + func_globals = getattr(func, "__globals__", None) + if func_globals is None: + func_globals = getattr( + getattr(func, "__func__", None), "__globals__", {} + ) + return get_type_hints( + func, + globalns={**fallback_globals, **func_globals}, + ) + except (NameError, TypeError): + return {} + + +def _is_async_callable(func: Callable[..., Any]) -> bool: + return inspect.iscoroutinefunction(_get_callable_inspection_target(func)) + + +def _resolve_deterministic(func: _UDFInput, deterministic: bool) -> bool: + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool.") + if isinstance(func, (ScalarFunction, AsyncScalarFunction)): + _validate_deterministic(deterministic, func.is_deterministic()) + return deterministic + + +def _validate_deterministic(declared: bool, actual: bool) -> None: + if declared != actual: + raise ValueError(f"Inconsistent deterministic: {declared} and {actual}.") + + +def _resolve_name(func: _UDFInput, name: Optional[str]) -> str: + actual_name = _default_udf_name(func) if name is None else name + if not isinstance(actual_name, str): + raise TypeError("name must be a str or None.") + if not actual_name: + raise ValueError("name must not be empty.") + return actual_name + + +def _default_udf_name(func: _UDFInput) -> str: + target = _unwrap_partial(func) + name = getattr(target, "__name__", None) + return name if isinstance(name, str) else type(target).__name__ + + +# ======================== Worker Adapters ======================== + + +def _wrap_scalar_general_result( + func: Callable[..., Any], return_dtype: DataType +) -> Callable[..., Any]: + result_type = return_dtype._to_table_data_type() + + if _is_async_callable(func): + + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + return _normalize_user_value(await func(*args, **kwargs), result_type) + + wrapper = async_wrapper + else: + + @functools.wraps(func) + def sync_wrapper(*args: Any, **kwargs: Any) -> Any: + return _normalize_user_value(func(*args, **kwargs), result_type) + + wrapper = sync_wrapper + + if not hasattr(func, "__name__"): + wrapper.__name__ = type(func).__name__ + return wrapper + + +class _DataFrameUDFAdapterBase: + """Bind a lazy DataFrame UDF source to one worker invocation protocol.""" + + def __init__( + self, + func: _UDFInput, + return_dtype: DataType, + deterministic: bool, + usage: _UDFUsage, + func_type: str, + ) -> None: + self._func_class: Optional[Type] = func if inspect.isclass(func) else None + self._func: Optional[_UDFInput] = None if self._func_class is not None else func + self._return_dtype = return_dtype if func_type == "general" else None + self._deterministic = deterministic + self._usage = usage + self._func_type = func_type + self._bound_func: Optional[Callable[..., Any]] = None + self.__name__ = _default_udf_name(func) + self.__doc__ = getattr(func, "__doc__", None) + + def open(self, function_context: Any) -> None: + if self._func_class is not None: + self._func = self._func_class() + + func = self._func + if func is None: + raise RuntimeError("DataFrame UDF source was not initialized.") + if not isinstance(func, (ScalarFunction, AsyncScalarFunction)) and not callable(func): + class_name = self._func_class.__name__ if self._func_class else type(func).__name__ + raise TypeError( + f"Callable class '{class_name}' constructed a non-callable " + f"object of type '{type(func).__name__}'." + ) + + if isinstance(func, (ScalarFunction, AsyncScalarFunction)): + _validate_deterministic(self._deterministic, func.is_deterministic()) + func.open(function_context) + invoke_func = func.eval + elif inspect.isroutine(func) or isinstance(func, functools.partial): + invoke_func = cast(Callable[..., Any], func) + else: + invoke_func = cast(Callable[..., Any], getattr(func, "__call__")) + self._bound_func = self._bind_func(invoke_func) + + def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: + if self._usage is not _UDFUsage.EXPRESSION: + raise NotImplementedError( + f"DataFrame UDF usage {self._usage.value!r} is not supported yet." + ) + if self._func_type == "general": + return _wrap_scalar_general_result( + invoke_func, cast(DataType, self._return_dtype) + ) + return invoke_func + + def close(self) -> None: + func = self._func + try: + if isinstance(func, (ScalarFunction, AsyncScalarFunction)): + func.close() + finally: + self._bound_func = None + if self._func_class is not None: + self._func = None + + def is_deterministic(self) -> bool: + return self._deterministic + + def _invocation(self) -> Callable[..., Any]: + if self._bound_func is None: + raise RuntimeError("DataFrame UDF was invoked before open().") + return self._bound_func + + +class _DataFrameScalarFunctionAdapter(_DataFrameUDFAdapterBase, ScalarFunction): + """Synchronous terminal adapter for a bound DataFrame UDF.""" + + def eval(self, *args: Any) -> Any: + invoke_func = self._invocation() + if self._func_type == "pandas": + from pyflink.fn_execution.utils.operation_utils import ( + check_pandas_udf_result, + ) + + return check_pandas_udf_result(invoke_func, *args) + return invoke_func(*args) + + +class _DataFrameAsyncScalarFunctionAdapter( + _DataFrameUDFAdapterBase, AsyncScalarFunction +): + """Asynchronous terminal adapter for a bound DataFrame UDF.""" + + async def eval(self, *args: Any) -> Any: + return await self._invocation()(*args) + + +# ======================== Result Normalization ======================== + + +def _row_value_by_type(value: Any, row_type: RowType, index: int) -> Any: + field_name = row_type.field_names()[index] + if isinstance(value, Mapping): + return value.get(field_name) + if isinstance(value, Row) and hasattr(value, "_fields"): + return _named_row_field_value(value, field_name) + if isinstance(value, (Row, tuple, list)): + if len(value) != len(row_type.fields): + raise ValueError( + f"Expected {len(row_type.fields)} value(s) for RowType " + f"{row_type.field_names()}, got {len(value)}." + ) + return value[index] + attributes = getattr(value, "__dict__", None) + if isinstance(attributes, Mapping): + return attributes.get(field_name) + try: + return getattr(value, field_name) + except AttributeError: + raise TypeError( + f"Expected a Mapping, Row, tuple, list, or object with fields for RowType " + f"{row_type.field_names()}, got {type(value).__name__}." + ) from None + + +def _named_row_field_value(row: Row, field_name: str) -> Any: + if field_name not in row._fields: + raise ValueError( + f"Field name {field_name!r} does not exist in Row fields {row._fields}." + ) + field_index = row._fields.index(field_name) + if field_index >= len(row): + raise ValueError( + f"Field name {field_name!r} is declared in Row fields {row._fields} " + "but has no value." + ) + return row[field_name] + + +def _normalize_user_value(value: Any, data_type: Any) -> Any: + """Normalize nested user values to the Python shape expected by Table coders.""" + if value is None: + return None + if isinstance(data_type, RowType): + row = Row( + *[ + _normalize_user_value( + _row_value_by_type(value, data_type, index), field.data_type + ) + for index, field in enumerate(data_type) + ] + ) + row.set_field_names(data_type.field_names()) + if isinstance(value, Row): + row.set_row_kind(value.get_row_kind()) + return row + if isinstance(data_type, ArrayType): + return [ + _normalize_user_value(item, data_type.element_type) for item in value + ] + if isinstance(data_type, MapType): + items_method = getattr(value, "items", None) + if callable(items_method): + items = list(items_method()) + else: + try: + items = list(value) + except TypeError as exc: + raise TypeError( + f"Expected a Mapping or iterable of key/value pairs for {data_type}, " + f"got {type(value).__name__}." + ) from exc + if any( + not isinstance(item, (tuple, list)) or len(item) != 2 for item in items + ): + raise TypeError( + f"Expected a Mapping or iterable of key/value pairs for {data_type}, " + f"got {type(value).__name__}." + ) + if any(item[0] is None for item in items): + raise TypeError(f"MapType keys must not be null for {data_type}.") + return { + _normalize_user_value(key, data_type.key_type): _normalize_user_value( + item_value, data_type.value_type + ) + for key, item_value in items + } + return value From 3100ebd205dc8e3219425e163723516e7b3897cd Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 27 Aug 2026 10:49:40 +0800 Subject: [PATCH 02/18] [FLINK-40431][python] Centralize DataFrame UDF source resolution Introduce a resolved source descriptor that centralizes callable classification, construction, invocation, lifecycle, async detection, and annotation inspection. Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 24 +- flink-python/pyflink/dataframe/udf.py | 372 +++++++++++------- 2 files changed, 251 insertions(+), 145 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 78e0cb30cde43..732da8d22fc92 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -28,7 +28,7 @@ from pyflink.common import Row from pyflink.table import DataTypes as TableDataTypes from pyflink.table.types import RowType -from pyflink.table.udf import AsyncScalarFunction, ScalarFunction +from pyflink.table.udf import AsyncScalarFunction, ScalarFunction, TableFunction from pyflink.testing.test_case_utils import ( PyFlinkDataFrameUTTestCase, PyFlinkStreamDataFrameTestCase, @@ -249,7 +249,7 @@ async def async_pandas(values: pd.Series) -> pd.Series: with self.subTest(case=case_name): wrapped = declare() self.assertEqual(wrapped._func_type, expected_type) - self.assertEqual(wrapped._is_async, expected_async) + self.assertEqual(wrapped._source.is_async, expected_async) invalid_declarations = [ ( @@ -436,6 +436,10 @@ def __call__(self, other: int) -> int: class NotCallable: pass + class NonScalarFunction(TableFunction): + def eval(self, value): + return value + invalid_declarations = [ ( "not callable", @@ -449,6 +453,15 @@ class NotCallable: TypeError, "func must be callable", ), + ( + "non-scalar UDF class", + lambda: pf.udf( + NonScalarFunction, + return_dtype=pf.DataType.int64(), + ), + TypeError, + "func must be a scalar UDF", + ), ( "missing return", lambda: pf.udf(missing_return), @@ -527,6 +540,7 @@ def test_scalar_function_lifecycle_and_cleanup(self): from pyflink.dataframe.udf import ( _DataFrameScalarFunctionAdapter, _UDFUsage, + _resolve_udf_source, ) events = [] @@ -546,7 +560,7 @@ def close(self): context = object() adapter = _DataFrameScalarFunctionAdapter( - LifecycleFunction(), + _resolve_udf_source(LifecycleFunction()), pf.DataType.int64(), True, _UDFUsage.EXPRESSION, @@ -587,7 +601,7 @@ def __call__(self, value): return value + 1 deferred_adapter = _DataFrameScalarFunctionAdapter( - DeferredCallable, + _resolve_udf_source(DeferredCallable), pf.DataType.int64(), True, _UDFUsage.EXPRESSION, @@ -609,7 +623,7 @@ def close(self): raise RuntimeError("close failed") failing_adapter = _DataFrameScalarFunctionAdapter( - FailingCloseFunction(), + _resolve_udf_source(FailingCloseFunction()), pf.DataType.int64(), True, _UDFUsage.EXPRESSION, diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 525e2b3f3f1ff..71451722aa151 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -21,13 +21,14 @@ import functools import inspect from collections.abc import Mapping +from dataclasses import dataclass from enum import Enum from typing import ( Any, Callable, Dict, + FrozenSet, Optional, - Set, Tuple, Type, Union, @@ -62,6 +63,107 @@ class _UDFUsage(Enum): MAP_BATCHES = "map_batches" +class _UDFSourceKind(Enum): + """How a resolved UDF source is initialized and invoked on a worker.""" + + DIRECT_CALLABLE = "direct_callable" + CALLABLE_INSTANCE = "callable_instance" + CALLABLE_CLASS = "callable_class" + SCALAR_FUNCTION = "scalar_function" + + +@dataclass(frozen=True) +class _ResolvedUDFSource: + """Callable metadata resolved once on the client and reused on workers.""" + + declaration_source: _UDFInput + runtime_source: _UDFInput + kind: _UDFSourceKind + is_async: bool + ignored_hint_names: FrozenSet[str] = frozenset() + + @property + def inspection_target(self) -> Callable[..., Any]: + if self.kind is _UDFSourceKind.DIRECT_CALLABLE: + return _get_callable_inspection_target( + cast(Callable[..., Any], self.runtime_source) + ) + if self.kind is _UDFSourceKind.SCALAR_FUNCTION: + return cast( + Union[ScalarFunction, AsyncScalarFunction], self.runtime_source + ).eval + if self.kind is _UDFSourceKind.CALLABLE_CLASS: + hint_method, _ = _get_callable_class_hint_method( + cast(Type, self.runtime_source) + ) + if hint_method is None: + raise RuntimeError("Resolved callable class has no inspection target.") + return hint_method + return cast( + Callable[..., Any], getattr(self.runtime_source, "__call__") + ) + + @property + def default_name(self) -> str: + return _default_udf_name(self.declaration_source) + + @property + def constructs_on_worker(self) -> bool: + return self.kind is _UDFSourceKind.CALLABLE_CLASS + + def create_worker_source(self) -> _UDFInput: + if not self.constructs_on_worker: + return self.runtime_source + source_class = cast(Type, self.runtime_source) + source = source_class() + if not callable(source): + raise TypeError( + f"Callable class '{source_class.__name__}' constructed a non-callable " + f"object of type '{type(source).__name__}'." + ) + return cast(_UDFInput, source) + + def validate_deterministic( + self, declared: bool, worker_source: Optional[_UDFInput] = None + ) -> None: + if self.kind is _UDFSourceKind.SCALAR_FUNCTION: + source = self.runtime_source if worker_source is None else worker_source + _validate_deterministic( + declared, + cast( + Union[ScalarFunction, AsyncScalarFunction], source + ).is_deterministic(), + ) + + def open_worker_source( + self, worker_source: _UDFInput, function_context: Any + ) -> None: + if self.kind is _UDFSourceKind.SCALAR_FUNCTION: + cast( + Union[ScalarFunction, AsyncScalarFunction], worker_source + ).open(function_context) + + def worker_invocation( + self, worker_source: _UDFInput + ) -> Callable[..., Any]: + if self.kind is _UDFSourceKind.DIRECT_CALLABLE: + return cast(Callable[..., Any], worker_source) + if self.kind is _UDFSourceKind.SCALAR_FUNCTION: + return cast( + Union[ScalarFunction, AsyncScalarFunction], worker_source + ).eval + return cast(Callable[..., Any], getattr(worker_source, "__call__")) + + def close_worker_source(self, worker_source: Optional[_UDFInput]) -> None: + if ( + self.kind is _UDFSourceKind.SCALAR_FUNCTION + and worker_source is not None + ): + cast( + Union[ScalarFunction, AsyncScalarFunction], worker_source + ).close() + + @PublicEvolving() class DataFrameUDFWrapper: """ @@ -81,37 +183,32 @@ class DataFrameUDFWrapper: .. versionadded:: 2.4.0 """ - _func: _UDFInput + _source: _ResolvedUDFSource _return_dtype: DataType _deterministic: bool _func_type: str - _is_async: bool _cached_table_udf_wrapper: Optional[UserDefinedFunctionWrapper] _frozen: bool __name__: str def __init__( self, - func: _UDFInput, + source: _ResolvedUDFSource, return_dtype: DataType, deterministic: bool, name: str, func_type: str, - is_async: bool, - metadata_source: Optional[_UDFInput] = None, ) -> None: - object.__setattr__(self, "_func", func) + object.__setattr__(self, "_source", source) object.__setattr__(self, "_return_dtype", return_dtype) object.__setattr__(self, "_deterministic", deterministic) object.__setattr__(self, "_func_type", func_type) - object.__setattr__(self, "_is_async", is_async) object.__setattr__(self, "_cached_table_udf_wrapper", None) - declaration = func if metadata_source is None else metadata_source - declaration_metadata = _unwrap_partial(declaration) + declaration_metadata = _unwrap_partial(source.declaration_source) functools.update_wrapper(self, declaration_metadata, updated=()) object.__setattr__(self, "__name__", name) - object.__setattr__(self, "__wrapped__", declaration) + object.__setattr__(self, "__wrapped__", source.declaration_source) object.__setattr__(self, "_frozen", True) def __setattr__(self, name: str, value: Any) -> None: @@ -149,13 +246,13 @@ def _create_table_udf_wrapper( ) -> UserDefinedFunctionWrapper: adapter_type = ( _DataFrameAsyncScalarFunctionAdapter - if self._is_async + if self._source.is_async else _DataFrameScalarFunctionAdapter ) actual_func = cast( Union[ScalarFunction, AsyncScalarFunction], adapter_type( - self._func, + self._source, self._return_dtype, self._deterministic, usage, @@ -350,33 +447,27 @@ def udf( """ def decorator(f: _UDFInput) -> DataFrameUDFWrapper: - _validate_scalar_udf_source(f) - ( - actual_func, - inspection_target, - skip_first_parameter, - is_async, - ) = _resolve_udf_source(f) + source = _resolve_udf_source(f) actual_func_type = ( func_type if func_type is not None - else _detect_func_type( - inspection_target, skip_first=skip_first_parameter - ) + else _detect_func_type(source) ) - _validate_scalar_udf_options(actual_func_type, return_dtype, is_async) - actual_return_dtype = _infer_return_dtype(inspection_target, return_dtype) - actual_deterministic = _resolve_deterministic(actual_func, deterministic) - actual_name = _resolve_name(actual_func, name) + _validate_scalar_udf_options( + actual_func_type, return_dtype, source.is_async + ) + actual_return_dtype = _infer_return_dtype( + source.inspection_target, return_dtype + ) + actual_deterministic = _resolve_deterministic(source, deterministic) + actual_name = _resolve_name(source, name) return DataFrameUDFWrapper( - actual_func, + source, actual_return_dtype, actual_deterministic, actual_name, actual_func_type, - is_async, - metadata_source=f, ) return decorator if func is None else decorator(func) @@ -385,25 +476,6 @@ def decorator(f: _UDFInput) -> DataFrameUDFWrapper: # ======================== Declaration Validation ======================== -def _validate_scalar_udf_source(func: Any) -> None: - if inspect.isclass(func): - if issubclass(func, UserDefinedFunction) and not issubclass( - func, (ScalarFunction, AsyncScalarFunction) - ): - raise TypeError(f"func must be a scalar UDF, got {func.__name__}.") - if not issubclass( - func, (ScalarFunction, AsyncScalarFunction) - ) and not _has_custom_call(func): - raise TypeError(f"func must be callable, got {func.__name__}.") - return - if isinstance(func, UserDefinedFunction) and not isinstance( - func, (ScalarFunction, AsyncScalarFunction) - ): - raise TypeError(f"func must be a scalar UDF, got {type(func).__name__}.") - if not isinstance(func, (ScalarFunction, AsyncScalarFunction)) and not callable(func): - raise TypeError(f"func must be callable, got {type(func).__name__}.") - - def _validate_scalar_udf_options( func_type: str, return_dtype: Optional[_DataTypeLike], @@ -447,58 +519,108 @@ def _get_callable_class_hint_method( return None, False -def _is_class_udf(func: Any) -> bool: - """Check whether ``func`` uses a class-based scalar UDF declaration form.""" - if isinstance(func, functools.partial): - return False - if isinstance(func, (ScalarFunction, AsyncScalarFunction)): - return True - if inspect.isclass(func): - if issubclass(func, (ScalarFunction, AsyncScalarFunction)): - return True - return _has_custom_call(func) - if not inspect.isroutine(func) and hasattr(func, "__call__"): - return _has_custom_call(type(func)) - return False - - def _resolve_udf_source( func: _UDFInput, -) -> Tuple[_UDFInput, Callable[..., Any], bool, bool]: - """Resolve the runtime source and callable that describes its declaration.""" - if not _is_class_udf(func): - callable_func = cast(Callable[..., Any], func) - return func, callable_func, False, _is_async_callable(callable_func) +) -> _ResolvedUDFSource: + """Validate and classify one callable declaration.""" + if isinstance(func, functools.partial) or inspect.isroutine(func): + inspection_target = _get_callable_inspection_target( + cast(Callable[..., Any], func) + ) + return _ResolvedUDFSource( + func, + func, + _UDFSourceKind.DIRECT_CALLABLE, + inspect.iscoroutinefunction(inspection_target), + _ignored_hint_names(func, inspection_target), + ) - skip_first_parameter = False if inspect.isclass(func): - _validate_zero_argument_class(func) + if issubclass(func, UserDefinedFunction) and not issubclass( + func, (ScalarFunction, AsyncScalarFunction) + ): + raise TypeError(f"func must be a scalar UDF, got {func.__name__}.") if issubclass(func, (ScalarFunction, AsyncScalarFunction)): + _validate_zero_argument_class(func) actual_func = func() hint_method = actual_func.eval is_async = isinstance( actual_func, AsyncScalarFunction ) or inspect.iscoroutinefunction(hint_method) - return actual_func, hint_method, False, is_async + return _ResolvedUDFSource( + func, + actual_func, + _UDFSourceKind.SCALAR_FUNCTION, + is_async, + ) - class_hint_method, skip_first_parameter = _get_callable_class_hint_method( - func + if not _has_custom_call(func): + raise TypeError(f"func must be callable, got {func.__name__}.") + _validate_zero_argument_class(func) + class_hint_method, skip_first_parameter = ( + _get_callable_class_hint_method(func) ) if class_hint_method is None: raise TypeError( f"Callable class '{func.__name__}': __call__ must be defined as a method." ) - hint_method = class_hint_method - is_async = inspect.iscoroutinefunction(hint_method) - elif isinstance(func, (ScalarFunction, AsyncScalarFunction)): + return _ResolvedUDFSource( + func, + func, + _UDFSourceKind.CALLABLE_CLASS, + inspect.iscoroutinefunction(class_hint_method), + _ignored_hint_names( + func, class_hint_method, skip_first=skip_first_parameter + ), + ) + + if isinstance(func, UserDefinedFunction) and not isinstance( + func, (ScalarFunction, AsyncScalarFunction) + ): + raise TypeError(f"func must be a scalar UDF, got {type(func).__name__}.") + if isinstance(func, (ScalarFunction, AsyncScalarFunction)): hint_method = func.eval is_async = isinstance(func, AsyncScalarFunction) or inspect.iscoroutinefunction( hint_method ) - else: - hint_method = cast(Callable[..., Any], getattr(func, "__call__")) - is_async = inspect.iscoroutinefunction(hint_method) - return func, hint_method, skip_first_parameter, is_async + return _ResolvedUDFSource( + func, + func, + _UDFSourceKind.SCALAR_FUNCTION, + is_async, + ) + + if not callable(func): + raise TypeError(f"func must be callable, got {type(func).__name__}.") + hint_method = cast(Callable[..., Any], getattr(func, "__call__")) + return _ResolvedUDFSource( + func, + func, + _UDFSourceKind.CALLABLE_INSTANCE, + inspect.iscoroutinefunction(hint_method), + ) + + +def _ignored_hint_names( + func: _UDFInput, + inspection_target: Callable[..., Any], + skip_first: bool = False, +) -> FrozenSet[str]: + try: + parameters = list(inspect.signature(inspection_target).parameters) + except (TypeError, ValueError): + parameters = [] + ignored_names = set(parameters[:1]) if skip_first else set() + if isinstance(func, functools.partial): + try: + ignored_names.update( + inspect.signature(inspection_target) + .bind_partial(*func.args, **(func.keywords or {})) + .arguments + ) + except (TypeError, ValueError): + pass + return frozenset(ignored_names) def _validate_zero_argument_class(func_class: Type) -> None: @@ -527,8 +649,7 @@ def _infer_return_dtype( if return_dtype is not None: return _convert_to_dtype(return_dtype) - hint_func = _get_callable_inspection_target(func) - hints = _get_callable_type_hints(hint_func) + hints = _get_callable_type_hints(func) if "return" not in hints: func_name = _default_udf_name(func) raise TypeError( @@ -578,9 +699,9 @@ def _data_type_from_type_hint(type_hint: Any) -> DataType: return DataType._from_type_hint(type_hint) -def _detect_func_type(func: Callable[..., Any], skip_first: bool = False) -> str: +def _detect_func_type(source: _ResolvedUDFSource) -> str: """Detect pandas mode from an unbound pandas container annotation.""" - hint_func = _get_callable_inspection_target(func) + hint_func = source.inspection_target try: import pandas as pd except ImportError: @@ -589,28 +710,11 @@ def _detect_func_type(func: Callable[..., Any], skip_first: bool = False) -> str hint_func, fallback_globals={"pandas": pd, "pd": pd} ) - try: - parameters = list(inspect.signature(hint_func).parameters) - except (TypeError, ValueError): - parameters = [] - ignored_hint_names: Set[str] = set() - if skip_first and parameters: - ignored_hint_names.add(parameters[0]) - if isinstance(func, functools.partial): - try: - ignored_hint_names.update( - inspect.signature(hint_func) - .bind_partial(*func.args, **(func.keywords or {})) - .arguments - ) - except (TypeError, ValueError): - pass - pandas_types = (pd.Series, pd.DataFrame) return ( "pandas" if any( - name not in ignored_hint_names and hint in pandas_types + name not in source.ignored_hint_names and hint in pandas_types for name, hint in hints.items() ) else "general" @@ -651,15 +755,12 @@ def _get_callable_type_hints( return {} -def _is_async_callable(func: Callable[..., Any]) -> bool: - return inspect.iscoroutinefunction(_get_callable_inspection_target(func)) - - -def _resolve_deterministic(func: _UDFInput, deterministic: bool) -> bool: +def _resolve_deterministic( + source: _ResolvedUDFSource, deterministic: bool +) -> bool: if not isinstance(deterministic, bool): raise TypeError("deterministic must be a bool.") - if isinstance(func, (ScalarFunction, AsyncScalarFunction)): - _validate_deterministic(deterministic, func.is_deterministic()) + source.validate_deterministic(deterministic) return deterministic @@ -668,8 +769,8 @@ def _validate_deterministic(declared: bool, actual: bool) -> None: raise ValueError(f"Inconsistent deterministic: {declared} and {actual}.") -def _resolve_name(func: _UDFInput, name: Optional[str]) -> str: - actual_name = _default_udf_name(func) if name is None else name +def _resolve_name(source: _ResolvedUDFSource, name: Optional[str]) -> str: + actual_name = source.default_name if name is None else name if not isinstance(actual_name, str): raise TypeError("name must be a str or None.") if not actual_name: @@ -687,11 +788,11 @@ def _default_udf_name(func: _UDFInput) -> str: def _wrap_scalar_general_result( - func: Callable[..., Any], return_dtype: DataType + func: Callable[..., Any], return_dtype: DataType, is_async: bool ) -> Callable[..., Any]: result_type = return_dtype._to_table_data_type() - if _is_async_callable(func): + if is_async: @functools.wraps(func) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: @@ -716,44 +817,34 @@ class _DataFrameUDFAdapterBase: def __init__( self, - func: _UDFInput, + source: _ResolvedUDFSource, return_dtype: DataType, deterministic: bool, usage: _UDFUsage, func_type: str, ) -> None: - self._func_class: Optional[Type] = func if inspect.isclass(func) else None - self._func: Optional[_UDFInput] = None if self._func_class is not None else func + self._source = source + self._func: Optional[_UDFInput] = ( + None if source.constructs_on_worker else source.runtime_source + ) self._return_dtype = return_dtype if func_type == "general" else None self._deterministic = deterministic self._usage = usage self._func_type = func_type self._bound_func: Optional[Callable[..., Any]] = None - self.__name__ = _default_udf_name(func) - self.__doc__ = getattr(func, "__doc__", None) + self.__name__ = source.default_name + self.__doc__ = getattr(source.declaration_source, "__doc__", None) def open(self, function_context: Any) -> None: - if self._func_class is not None: - self._func = self._func_class() + if self._source.constructs_on_worker: + self._func = self._source.create_worker_source() func = self._func if func is None: raise RuntimeError("DataFrame UDF source was not initialized.") - if not isinstance(func, (ScalarFunction, AsyncScalarFunction)) and not callable(func): - class_name = self._func_class.__name__ if self._func_class else type(func).__name__ - raise TypeError( - f"Callable class '{class_name}' constructed a non-callable " - f"object of type '{type(func).__name__}'." - ) - - if isinstance(func, (ScalarFunction, AsyncScalarFunction)): - _validate_deterministic(self._deterministic, func.is_deterministic()) - func.open(function_context) - invoke_func = func.eval - elif inspect.isroutine(func) or isinstance(func, functools.partial): - invoke_func = cast(Callable[..., Any], func) - else: - invoke_func = cast(Callable[..., Any], getattr(func, "__call__")) + self._source.validate_deterministic(self._deterministic, func) + self._source.open_worker_source(func, function_context) + invoke_func = self._source.worker_invocation(func) self._bound_func = self._bind_func(invoke_func) def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: @@ -763,18 +854,19 @@ def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: ) if self._func_type == "general": return _wrap_scalar_general_result( - invoke_func, cast(DataType, self._return_dtype) + invoke_func, + cast(DataType, self._return_dtype), + self._source.is_async, ) return invoke_func def close(self) -> None: func = self._func try: - if isinstance(func, (ScalarFunction, AsyncScalarFunction)): - func.close() + self._source.close_worker_source(func) finally: self._bound_func = None - if self._func_class is not None: + if self._source.constructs_on_worker: self._func = None def is_deterministic(self) -> bool: From 094cc1849ecc2343370382cc736375ba4aa56e38 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 27 Aug 2026 14:16:47 +0800 Subject: [PATCH 03/18] [FLINK-40431][python] Initialize UDF class declarations on TaskManagers Defer zero-argument callable, ScalarFunction, and AsyncScalarFunction class construction while keeping configured instances client-created. Resolve class annotations statically and clean up partial lifecycle initialization. Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 246 ++++++++++++++++-- flink-python/pyflink/dataframe/udf.py | 160 +++++++----- 2 files changed, 321 insertions(+), 85 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 732da8d22fc92..88dea6ed8cded 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -16,6 +16,7 @@ # limitations under the License. ################################################################################ +import asyncio import functools import inspect import unittest @@ -142,6 +143,13 @@ def __init__(self): async def eval(self, value: int) -> int: return value * 2 + class AddScalarOffset(ScalarFunction): + def __init__(self, offset): + self.offset = offset + + def eval(self, value: int) -> int: + return value + self.offset + named_callable = NamedCallable() double_instance = Double() async_double_instance = AsyncDouble() @@ -154,6 +162,7 @@ async def eval(self, value: int) -> int: double_instance, AsyncDouble, async_double_instance, + AddScalarOffset(2), ] for source in callables: with self.subTest(source=source): @@ -161,7 +170,7 @@ async def eval(self, value: int) -> int: self.assertEqual(decorated.return_dtype, pf.DataType.int64()) self.assertEqual(plain_constructor_calls, []) - self.assertEqual(scalar_constructor_calls, ["Double", "AsyncDouble"]) + self.assertEqual(scalar_constructor_calls, []) self.assertEqual(pf.udf(named_callable).__name__, "configured_add") decorated_class = pf.udf(Double) @@ -192,6 +201,18 @@ async def async_add_one(value: int) -> int: async def async_pandas(values: pd.Series) -> pd.Series: return values + 1 + class PandasCallable: + def __call__(self, values: pd.Series) -> pd.Series: + return values + 1 + + class PandasScalarFunction(ScalarFunction): + def eval(self, values: pd.Series) -> pd.Series: + return values + 1 + + class AsyncScalarClass(AsyncScalarFunction): + async def eval(self, value: int) -> int: + return value + 1 + declarations = [ ( "inferred pandas", @@ -244,6 +265,30 @@ async def async_pandas(values: pd.Series) -> pd.Series: "general", True, ), + ( + "pandas callable class", + lambda: pf.udf( + PandasCallable, + return_dtype=pf.DataType.int64(), + ), + "pandas", + False, + ), + ( + "pandas scalar-function class", + lambda: pf.udf( + PandasScalarFunction, + return_dtype=pf.DataType.int64(), + ), + "pandas", + False, + ), + ( + "async scalar-function class", + lambda: pf.udf(AsyncScalarClass), + "general", + True, + ), ] for case_name, declare, expected_type, expected_async in declarations: with self.subTest(case=case_name): @@ -310,8 +355,7 @@ def eval(self, value: int) -> int: ) with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"): pf.udf(instance) - with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"): - pf.udf(NonDeterministic) + self.assertTrue(pf.udf(NonDeterministic)._deterministic) named = pf.udf(instance, deterministic=False, name="identity") self.assertEqual(named.__name__, "identity") @@ -433,6 +477,20 @@ def __init__(self, value): def __call__(self, other: int) -> int: return other + self.value + class RequiresScalarArgument(ScalarFunction): + def __init__(self, value): + self.value = value + + def eval(self, other: int) -> int: + return other + self.value + + class RequiresAsyncScalarArgument(AsyncScalarFunction): + def __init__(self, value): + self.value = value + + async def eval(self, other: int) -> int: + return other + self.value + class NotCallable: pass @@ -482,6 +540,18 @@ def eval(self, value): TypeError, "zero-argument constructor", ), + ( + "required scalar constructor argument", + lambda: pf.udf(RequiresScalarArgument), + TypeError, + "zero-argument constructor", + ), + ( + "required async scalar constructor argument", + lambda: pf.udf(RequiresAsyncScalarArgument), + TypeError, + "zero-argument constructor", + ), ( "invalid determinism", lambda: pf.udf( @@ -538,6 +608,7 @@ def eval(self, value): class DataFrameUDFAdapterTests(unittest.TestCase): def test_scalar_function_lifecycle_and_cleanup(self): from pyflink.dataframe.udf import ( + _DataFrameAsyncScalarFunctionAdapter, _DataFrameScalarFunctionAdapter, _UDFUsage, _resolve_udf_source, @@ -545,6 +616,20 @@ def test_scalar_function_lifecycle_and_cleanup(self): events = [] + def create_adapter(source, deterministic=True, async_mode=False): + adapter_type = ( + _DataFrameAsyncScalarFunctionAdapter + if async_mode + else _DataFrameScalarFunctionAdapter + ) + return adapter_type( + _resolve_udf_source(source), + pf.DataType.int64(), + deterministic, + _UDFUsage.EXPRESSION, + "general", + ) + class LifecycleFunction(ScalarFunction): def __init__(self): events.append("init") @@ -559,13 +644,8 @@ def close(self): events.append("close") context = object() - adapter = _DataFrameScalarFunctionAdapter( - _resolve_udf_source(LifecycleFunction()), - pf.DataType.int64(), - True, - _UDFUsage.EXPRESSION, - "general", - ) + adapter = create_adapter(LifecycleFunction) + self.assertEqual(events, []) with self.assertRaisesRegex(RuntimeError, "before open"): adapter.eval(1) @@ -586,11 +666,97 @@ def close(self): "init", ("open", context), "close", + "init", ("open", context), "close", ], ) + failed_lifecycle_events = [] + + class NonDeterministicFunction(ScalarFunction): + def __init__(self): + failed_lifecycle_events.append("init") + + def eval(self, value): + return value + + def is_deterministic(self): + return False + + def close(self): + failed_lifecycle_events.append("close") + + mismatched_adapter = create_adapter(NonDeterministicFunction) + with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"): + mismatched_adapter.open(context) + mismatched_adapter.close() + self.assertEqual(failed_lifecycle_events, ["init"]) + + async_events = [] + + class AsyncLifecycleFunction(AsyncScalarFunction): + def __init__(self): + async_events.append("init") + + def open(self, function_context): + async_events.append(("open", function_context)) + + async def eval(self, value): + return value + 1 + + def close(self): + async_events.append("close") + + async_adapter = create_adapter( + AsyncLifecycleFunction, + async_mode=True, + ) + self.assertEqual(async_events, []) + async_adapter.open(context) + self.assertEqual(asyncio.run(async_adapter.eval(1)), 2) + async_adapter.close() + self.assertEqual(async_events, ["init", ("open", context), "close"]) + + initialization_failure_events = [] + + class ConstructorFailureFunction(ScalarFunction): + def __init__(self): + initialization_failure_events.append("init") + raise RuntimeError("constructor failed") + + def eval(self, value): + return value + + constructor_failure_adapter = create_adapter(ConstructorFailureFunction) + with self.assertRaisesRegex(RuntimeError, "constructor failed"): + constructor_failure_adapter.open(context) + constructor_failure_adapter.close() + self.assertEqual(initialization_failure_events, ["init"]) + + class OpenFailureFunction(ScalarFunction): + def __init__(self): + initialization_failure_events.append("second init") + + def open(self, function_context): + initialization_failure_events.append("open") + raise RuntimeError("open failed") + + def eval(self, value): + return value + + def close(self): + initialization_failure_events.append("close") + + open_failure_adapter = create_adapter(OpenFailureFunction) + with self.assertRaisesRegex(RuntimeError, "open failed"): + open_failure_adapter.open(context) + open_failure_adapter.close() + self.assertEqual( + initialization_failure_events, + ["init", "second init", "open"], + ) + deferred_constructor_calls = [] class DeferredCallable: @@ -600,13 +766,7 @@ def __init__(self): def __call__(self, value): return value + 1 - deferred_adapter = _DataFrameScalarFunctionAdapter( - _resolve_udf_source(DeferredCallable), - pf.DataType.int64(), - True, - _UDFUsage.EXPRESSION, - "general", - ) + deferred_adapter = create_adapter(DeferredCallable) deferred_adapter.open(context) self.assertEqual(deferred_adapter.eval(1), 2) deferred_adapter.close() @@ -622,19 +782,53 @@ def eval(self, value): def close(self): raise RuntimeError("close failed") - failing_adapter = _DataFrameScalarFunctionAdapter( - _resolve_udf_source(FailingCloseFunction()), - pf.DataType.int64(), - True, - _UDFUsage.EXPRESSION, - "general", - ) + failing_adapter = create_adapter(FailingCloseFunction()) failing_adapter.open(context) with self.assertRaisesRegex(RuntimeError, "close failed"): failing_adapter.close() with self.assertRaisesRegex(RuntimeError, "before open"): failing_adapter.eval(1) + def test_binding_failure_closes_and_resets_deferred_scalar_class(self): + from pyflink.dataframe.udf import ( + _DataFrameScalarFunctionAdapter, + _UDFUsage, + _resolve_udf_source, + ) + + events = [] + + class BindingFailureFunction(ScalarFunction): + def __init__(self): + events.append("init") + + def open(self, function_context): + events.append("open") + + def eval(self, value): + return value + + def close(self): + events.append("close") + raise RuntimeError("close failed") + + adapter = _DataFrameScalarFunctionAdapter( + _resolve_udf_source(BindingFailureFunction), + pf.DataType.int64(), + True, + _UDFUsage.MAP, + "general", + ) + for _ in range(2): + with self.assertRaisesRegex(NotImplementedError, "'map'"): + adapter.open(object()) + adapter.close() + + self.assertEqual( + events, + ["init", "open", "close", "init", "open", "close"], + ) + class DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase): def test_with_columns_binds_expressions_and_resolves_output_schema(self): @@ -729,7 +923,7 @@ def is_deterministic(self): return False deferred = pf.udf(DeferredCallable) - scalar_instance = pf.udf(OpenedScalarFunction()) + opened_scalar_class = pf.udf(OpenedScalarFunction) scalar_class = pf.udf(ClassNonDeterministic, deterministic=False) result = ( @@ -740,7 +934,7 @@ def is_deterministic(self): pandas_value=add_three(pf.col("id")), details=details(pf.col("id")), deferred_value=deferred(pf.col("id")), - scalar_value=scalar_instance(pf.col("id")), + scalar_value=opened_scalar_class(pf.col("id")), scalar_class_value=scalar_class(pf.col("id")), ) ) diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 71451722aa151..cda1f420785bc 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -69,15 +69,15 @@ class _UDFSourceKind(Enum): DIRECT_CALLABLE = "direct_callable" CALLABLE_INSTANCE = "callable_instance" CALLABLE_CLASS = "callable_class" - SCALAR_FUNCTION = "scalar_function" + SCALAR_FUNCTION_INSTANCE = "scalar_function_instance" + SCALAR_FUNCTION_CLASS = "scalar_function_class" @dataclass(frozen=True) class _ResolvedUDFSource: """Callable metadata resolved once on the client and reused on workers.""" - declaration_source: _UDFInput - runtime_source: _UDFInput + source: _UDFInput kind: _UDFSourceKind is_async: bool ignored_hint_names: FrozenSet[str] = frozenset() @@ -86,37 +86,57 @@ class _ResolvedUDFSource: def inspection_target(self) -> Callable[..., Any]: if self.kind is _UDFSourceKind.DIRECT_CALLABLE: return _get_callable_inspection_target( - cast(Callable[..., Any], self.runtime_source) + cast(Callable[..., Any], self.source) ) - if self.kind is _UDFSourceKind.SCALAR_FUNCTION: + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: return cast( - Union[ScalarFunction, AsyncScalarFunction], self.runtime_source + Union[ScalarFunction, AsyncScalarFunction], self.source ).eval - if self.kind is _UDFSourceKind.CALLABLE_CLASS: + if self.kind in ( + _UDFSourceKind.CALLABLE_CLASS, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ): hint_method, _ = _get_callable_class_hint_method( - cast(Type, self.runtime_source) + cast(Type, self.source), + "eval" + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS + else "__call__", ) if hint_method is None: - raise RuntimeError("Resolved callable class has no inspection target.") + raise RuntimeError("Resolved UDF class has no inspection target.") return hint_method - return cast( - Callable[..., Any], getattr(self.runtime_source, "__call__") - ) + return cast(Callable[..., Any], getattr(self.source, "__call__")) @property def default_name(self) -> str: - return _default_udf_name(self.declaration_source) + return _default_udf_name(self.source) + + @property + def is_scalar_function(self) -> bool: + return self.kind in ( + _UDFSourceKind.SCALAR_FUNCTION_INSTANCE, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ) @property def constructs_on_worker(self) -> bool: - return self.kind is _UDFSourceKind.CALLABLE_CLASS + return self.kind in ( + _UDFSourceKind.CALLABLE_CLASS, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ) def create_worker_source(self) -> _UDFInput: if not self.constructs_on_worker: - return self.runtime_source - source_class = cast(Type, self.runtime_source) + return self.source + source_class = cast(Type, self.source) source = source_class() - if not callable(source): + if self.is_scalar_function: + if not isinstance(source, (ScalarFunction, AsyncScalarFunction)): + raise TypeError( + f"Scalar UDF class '{source_class.__name__}' constructed an " + f"unsupported object of type '{type(source).__name__}'." + ) + elif not callable(source): raise TypeError( f"Callable class '{source_class.__name__}' constructed a non-callable " f"object of type '{type(source).__name__}'." @@ -126,8 +146,14 @@ def create_worker_source(self) -> _UDFInput: def validate_deterministic( self, declared: bool, worker_source: Optional[_UDFInput] = None ) -> None: - if self.kind is _UDFSourceKind.SCALAR_FUNCTION: - source = self.runtime_source if worker_source is None else worker_source + source: Optional[_UDFInput] + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: + source = self.source + elif self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS: + source = worker_source + else: + source = None + if source is not None: _validate_deterministic( declared, cast( @@ -138,7 +164,7 @@ def validate_deterministic( def open_worker_source( self, worker_source: _UDFInput, function_context: Any ) -> None: - if self.kind is _UDFSourceKind.SCALAR_FUNCTION: + if self.is_scalar_function: cast( Union[ScalarFunction, AsyncScalarFunction], worker_source ).open(function_context) @@ -148,17 +174,14 @@ def worker_invocation( ) -> Callable[..., Any]: if self.kind is _UDFSourceKind.DIRECT_CALLABLE: return cast(Callable[..., Any], worker_source) - if self.kind is _UDFSourceKind.SCALAR_FUNCTION: + if self.is_scalar_function: return cast( Union[ScalarFunction, AsyncScalarFunction], worker_source ).eval return cast(Callable[..., Any], getattr(worker_source, "__call__")) def close_worker_source(self, worker_source: Optional[_UDFInput]) -> None: - if ( - self.kind is _UDFSourceKind.SCALAR_FUNCTION - and worker_source is not None - ): + if self.is_scalar_function and worker_source is not None: cast( Union[ScalarFunction, AsyncScalarFunction], worker_source ).close() @@ -205,10 +228,10 @@ def __init__( object.__setattr__(self, "_func_type", func_type) object.__setattr__(self, "_cached_table_udf_wrapper", None) - declaration_metadata = _unwrap_partial(source.declaration_source) + declaration_metadata = _unwrap_partial(source.source) functools.update_wrapper(self, declaration_metadata, updated=()) object.__setattr__(self, "__name__", name) - object.__setattr__(self, "__wrapped__", source.declaration_source) + object.__setattr__(self, "__wrapped__", source.source) object.__setattr__(self, "_frozen", True) def __setattr__(self, name: str, value: Any) -> None: @@ -326,8 +349,8 @@ def udf( The function may be synchronous or asynchronous. Pandas UDFs operate on ``pandas.Series`` or ``pandas.DataFrame`` batches and must declare - ``return_dtype``. Plain callable class objects must have a zero-argument - constructor and are instantiated on the worker. + ``return_dtype``. Callable and scalar-function class objects must have a + zero-argument constructor and are initialized on the TaskManager. A UDF can be declared with a bare decorator, a configured decorator, or a direct call. General UDFs may infer ``return_dtype`` from the return @@ -361,8 +384,7 @@ def udf( Plain callable classes can be supplied as zero-argument class objects or as configured instances. Class objects, including their ``__init__``, are - constructed during worker initialization, so expensive initialization is - deferred to the TaskManager:: + initialized on the TaskManager, so expensive initialization is deferred:: >>> class AddOne: ... def __call__(self, value: int) -> int: @@ -381,8 +403,8 @@ def udf( :class:`~pyflink.table.udf.ScalarFunction` and :class:`~pyflink.table.udf.AsyncScalarFunction` class objects and instances are also supported. Their logical result type is inferred from ``eval`` - when it is not given explicitly. Class objects are instantiated on the - client; their ``open`` and ``close`` methods still run on the worker:: + when it is not given explicitly. Class objects are initialized on the + TaskManager, where their ``open`` and ``close`` methods also run:: >>> from pyflink.table.udf import AsyncScalarFunction, ScalarFunction @@ -528,7 +550,6 @@ def _resolve_udf_source( cast(Callable[..., Any], func) ) return _ResolvedUDFSource( - func, func, _UDFSourceKind.DIRECT_CALLABLE, inspect.iscoroutinefunction(inspection_target), @@ -542,16 +563,22 @@ def _resolve_udf_source( raise TypeError(f"func must be a scalar UDF, got {func.__name__}.") if issubclass(func, (ScalarFunction, AsyncScalarFunction)): _validate_zero_argument_class(func) - actual_func = func() - hint_method = actual_func.eval - is_async = isinstance( - actual_func, AsyncScalarFunction - ) or inspect.iscoroutinefunction(hint_method) + hint_method, skip_first_parameter = _get_callable_class_hint_method( + func, "eval" + ) + if hint_method is None: + raise TypeError( + f"Scalar UDF class '{func.__name__}': eval must be defined as a " + "method." + ) return _ResolvedUDFSource( func, - actual_func, - _UDFSourceKind.SCALAR_FUNCTION, - is_async, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + issubclass(func, AsyncScalarFunction) + or inspect.iscoroutinefunction(hint_method), + _ignored_hint_names( + func, hint_method, skip_first=skip_first_parameter + ), ) if not _has_custom_call(func): @@ -565,7 +592,6 @@ def _resolve_udf_source( f"Callable class '{func.__name__}': __call__ must be defined as a method." ) return _ResolvedUDFSource( - func, func, _UDFSourceKind.CALLABLE_CLASS, inspect.iscoroutinefunction(class_hint_method), @@ -585,8 +611,7 @@ def _resolve_udf_source( ) return _ResolvedUDFSource( func, - func, - _UDFSourceKind.SCALAR_FUNCTION, + _UDFSourceKind.SCALAR_FUNCTION_INSTANCE, is_async, ) @@ -594,7 +619,6 @@ def _resolve_udf_source( raise TypeError(f"func must be callable, got {type(func).__name__}.") hint_method = cast(Callable[..., Any], getattr(func, "__call__")) return _ResolvedUDFSource( - func, func, _UDFSourceKind.CALLABLE_INSTANCE, inspect.iscoroutinefunction(hint_method), @@ -825,27 +849,43 @@ def __init__( ) -> None: self._source = source self._func: Optional[_UDFInput] = ( - None if source.constructs_on_worker else source.runtime_source + None if source.constructs_on_worker else source.source ) self._return_dtype = return_dtype if func_type == "general" else None self._deterministic = deterministic self._usage = usage self._func_type = func_type self._bound_func: Optional[Callable[..., Any]] = None + self._lifecycle_opened = False self.__name__ = source.default_name - self.__doc__ = getattr(source.declaration_source, "__doc__", None) + self.__doc__ = getattr(source.source, "__doc__", None) def open(self, function_context: Any) -> None: - if self._source.constructs_on_worker: - self._func = self._source.create_worker_source() - - func = self._func - if func is None: - raise RuntimeError("DataFrame UDF source was not initialized.") - self._source.validate_deterministic(self._deterministic, func) - self._source.open_worker_source(func, function_context) - invoke_func = self._source.worker_invocation(func) - self._bound_func = self._bind_func(invoke_func) + lifecycle_opened = False + try: + if self._source.constructs_on_worker: + self._func = self._source.create_worker_source() + + func = self._func + if func is None: + raise RuntimeError("DataFrame UDF source was not initialized.") + self._source.validate_deterministic(self._deterministic, func) + self._source.open_worker_source(func, function_context) + lifecycle_opened = self._source.is_scalar_function + invoke_func = self._source.worker_invocation(func) + self._bound_func = self._bind_func(invoke_func) + self._lifecycle_opened = lifecycle_opened + except Exception: + if lifecycle_opened: + try: + self._source.close_worker_source(self._func) + except Exception: + pass + self._bound_func = None + self._lifecycle_opened = False + if self._source.constructs_on_worker: + self._func = None + raise def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: if self._usage is not _UDFUsage.EXPRESSION: @@ -863,9 +903,11 @@ def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: def close(self) -> None: func = self._func try: - self._source.close_worker_source(func) + if self._lifecycle_opened: + self._source.close_worker_source(func) finally: self._bound_func = None + self._lifecycle_opened = False if self._source.constructs_on_worker: self._func = None From be9f1796bfa619372b0a85031aa16cb715aadf49 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 27 Aug 2026 14:43:49 +0800 Subject: [PATCH 04/18] [FLINK-40431][python] Fix DataFrame UDF test override signatures Keep scalar-function test fixtures compatible with the variadic Table API eval contract while preserving unary behavior and type-hint inference. Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 39 ++++++++++++------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 88dea6ed8cded..6af60d8abe536 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -133,21 +133,24 @@ class Double(ScalarFunction): def __init__(self): scalar_constructor_calls.append("Double") - def eval(self, value: int) -> int: + def eval(self, *values: int) -> int: + value, = values return value * 2 class AsyncDouble(AsyncScalarFunction): def __init__(self): scalar_constructor_calls.append("AsyncDouble") - async def eval(self, value: int) -> int: + async def eval(self, *values: int) -> int: + value, = values return value * 2 class AddScalarOffset(ScalarFunction): def __init__(self, offset): self.offset = offset - def eval(self, value: int) -> int: + def eval(self, *values: int) -> int: + value, = values return value + self.offset named_callable = NamedCallable() @@ -206,11 +209,13 @@ def __call__(self, values: pd.Series) -> pd.Series: return values + 1 class PandasScalarFunction(ScalarFunction): - def eval(self, values: pd.Series) -> pd.Series: - return values + 1 + def eval(self, *values: pd.Series) -> pd.Series: + value, = values + return value + 1 class AsyncScalarClass(AsyncScalarFunction): - async def eval(self, value: int) -> int: + async def eval(self, *values: int) -> int: + value, = values return value + 1 declarations = [ @@ -321,14 +326,16 @@ async def eval(self, value: int) -> int: def test_determinism_and_name_metadata(self): class NonDeterministic(ScalarFunction): - def eval(self, value: int) -> int: + def eval(self, *values: int) -> int: + value, = values return value def is_deterministic(self): return False class DefaultDeterministic(ScalarFunction): - def eval(self, value: int) -> int: + def eval(self, *values: int) -> int: + value, = values return value instance = NonDeterministic() @@ -481,15 +488,17 @@ class RequiresScalarArgument(ScalarFunction): def __init__(self, value): self.value = value - def eval(self, other: int) -> int: - return other + self.value + def eval(self, *values: int) -> int: + value, = values + return value + self.value class RequiresAsyncScalarArgument(AsyncScalarFunction): def __init__(self, value): self.value = value - async def eval(self, other: int) -> int: - return other + self.value + async def eval(self, *values: int) -> int: + value, = values + return value + self.value class NotCallable: pass @@ -912,11 +921,13 @@ class OpenedScalarFunction(ScalarFunction): def open(self, function_context): self._increment = 5 - def eval(self, value: int) -> int: + def eval(self, *values: int) -> int: + value, = values return value + self._increment class ClassNonDeterministic(ScalarFunction): - def eval(self, value: int) -> int: + def eval(self, *values: int) -> int: + value, = values return value + 6 def is_deterministic(self): From 57bd73709f666c1d77f76c55ca25356a79228a1c Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 15:26:34 +0800 Subject: [PATCH 05/18] [FLINK-40431][python] Address DataFrame UDF review feedback Generated-by: OpenAI Codex (GPT-5) --- flink-python/pyflink/dataframe/dataframe.py | 13 +- .../pyflink/dataframe/tests/test_udf.py | 249 ++++++++++++++++- flink-python/pyflink/dataframe/udf.py | 261 +++++++++++++----- 3 files changed, 430 insertions(+), 93 deletions(-) diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 0e542d0ab2924..6cfe09f662d9c 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -200,7 +200,7 @@ def with_column( Add a column, or replace an existing column with the same name. ``expr`` may be an expression or a callable that receives this DataFrame and returns an - expression. + expression. See :func:`~pyflink.dataframe.udf.udf` for supported UDF declaration forms. :param name: Name of the added or replaced column. :param expr: Expression or callable used to compute the column value. @@ -211,10 +211,19 @@ def with_column( >>> import pyflink.dataframe as pf >>> df = pf.from_records([{"left": 1, "right": 2}]) - >>> result = df.with_column( + + >>> with_expression = df.with_column( ... "total", lambda current: current["left"] + current["right"] ... ) + >>> @pf.udf + ... def add(left: int, right: int) -> int: + ... return left + right + + >>> with_udf = df.with_column( + ... "total", add(pf.col("left"), pf.col("right")) + ... ) + .. versionadded:: 2.4.0 """ if not isinstance(name, str): diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 6af60d8abe536..098625df85d7b 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -19,6 +19,7 @@ import asyncio import functools import inspect +import operator import unittest from dataclasses import dataclass from typing import TypedDict @@ -26,7 +27,7 @@ import pandas as pd import pyarrow as pa import pyflink.dataframe as pf -from pyflink.common import Row +from pyflink.common import Row, RowKind from pyflink.table import DataTypes as TableDataTypes from pyflink.table.types import RowType from pyflink.table.udf import AsyncScalarFunction, ScalarFunction, TableFunction @@ -59,6 +60,22 @@ def describe(value: int) -> Result: "details": {"label": str(value), "scores": [value]}, } + def concrete_return_with_unresolved_input(value): + return value + + concrete_return_with_unresolved_input.__annotations__ = { + "value": "UnavailableInput", + "return": int, + } + + def postponed_return_with_unresolved_input(value): + return value + + postponed_return_with_unresolved_input.__annotations__ = { + "value": "UnavailableInput", + "return": "int", + } + decorated = pf.udf(add_one) from pyflink.dataframe.udf import DataFrameUDFWrapper @@ -100,6 +117,16 @@ def describe(value: int) -> Result: } ), ), + ( + "concrete return with unresolved input", + lambda: pf.udf(concrete_return_with_unresolved_input), + pf.DataType.int64(), + ), + ( + "postponed return with unresolved input", + lambda: pf.udf(postponed_return_with_unresolved_input), + pf.DataType.int64(), + ), ] for case_name, declare, expected in declarations: with self.subTest(case=case_name): @@ -180,6 +207,70 @@ def eval(self, *values: int) -> int: self.assertIs(decorated_class.__wrapped__, Double) self.assertEqual(decorated_class.__qualname__, Double.__qualname__) + def test_wrapped_signature_describes_udf_invocation(self): + def add(value: int, amount: int = 1) -> int: + return value + amount + + class CallableClass: + def __call__(self, value: int, amount: int = 1) -> int: + return value + amount + + class StaticCallableClass: + @staticmethod + def __call__(value: int, amount: int = 1) -> int: + return value + amount + + class ClassMethodCallableClass: + @classmethod + def __call__(cls, value: int, amount: int = 1) -> int: + return value + amount + + class AddFunction(ScalarFunction): + def eval(self, *values: int) -> int: + return sum(values) + + class AsyncAddFunction(AsyncScalarFunction): + async def eval(self, *values: int) -> int: + return sum(values) + + class ExplodingSignature: + @property + def __signature__(self): + raise RuntimeError("signature lookup failed") + + def __call__(self, value): + return value + + def variadic_add(*values: int) -> int: + return sum(values) + + expected_signature = inspect.signature(add) + variadic_signature = inspect.signature(variadic_add) + declarations = [ + (add, expected_signature), + (CallableClass, expected_signature), + (CallableClass(), expected_signature), + (StaticCallableClass, expected_signature), + (ClassMethodCallableClass, expected_signature), + (AddFunction, variadic_signature), + (AddFunction(), variadic_signature), + (AsyncAddFunction, variadic_signature), + (AsyncAddFunction(), variadic_signature), + ] + for source, expected in declarations: + with self.subTest(source=source): + self.assertEqual(inspect.signature(pf.udf(source)), expected) + + partial_add = functools.partial(add, 1) + self.assertEqual( + inspect.signature(pf.udf(partial_add)), inspect.signature(partial_add) + ) + + uninspectable = pf.udf(operator.itemgetter(0), return_dtype=int) + self.assertEqual(uninspectable.return_dtype, pf.DataType.int64()) + exploding_signature = pf.udf(ExplodingSignature(), return_dtype=int) + self.assertEqual(exploding_signature.return_dtype, pf.DataType.int64()) + def test_func_type_resolution_and_async_detection(self): def pandas_add_one(values: pd.Series) -> pd.Series: return values + 1 @@ -369,7 +460,7 @@ def eval(self, *values: int) -> int: self.assertEqual(named._table_udf_wrapper._name, "identity") def test_general_structured_results_are_normalized_recursively(self): - from pyflink.dataframe.udf import _normalize_user_value + from pyflink.dataframe.udf import _create_result_normalizer class Details: __slots__ = ("label", "scores") @@ -385,6 +476,26 @@ def __init__(self, items): def items(self): return self._items + class PropertyDetails: + def __init__(self, label, scores): + self._label = label + self.scores = scores + + @property + def label(self): + return self._label.upper() + + class MissingLabelDetails: + def __init__(self, scores): + self.scores = scores + + class FailingPropertyDetails: + scores = [17] + + @property + def label(self): + raise AttributeError("label lookup failed") + @dataclass class Result: id: int @@ -407,6 +518,21 @@ class Result: ) table_type = return_dtype._to_table_data_type() self.assertIsInstance(table_type, RowType) + result_normalizer = _create_result_normalizer(table_type) + self.assertIsNotNone(result_normalizer) + + named_row = Row( + id=4, + details=Row(label="named", scores=[5]), + attributes={"count": 6}, + ) + named_row.set_row_kind(RowKind.DELETE) + expected_named_row = Row( + id=4, + details=Row(label="named", scores=[5]), + attributes={"count": 6}, + ) + expected_named_row.set_row_kind(RowKind.DELETE) cases = [ ( @@ -425,16 +551,8 @@ class Result: ), ( "named row", - Row( - id=4, - details=Row(label="named", scores=[5]), - attributes={"count": 6}, - ), - Row( - id=4, - details=Row(label="named", scores=[5]), - attributes={"count": 6}, - ), + named_row, + expected_named_row, ), ( "positional list and tuple", @@ -458,22 +576,63 @@ class Result: attributes={"count": 13}, ), ), + ( + "property attribute", + Result( + id=14, + details=PropertyDetails(label="property", scores=[15]), + attributes={"count": 16}, + ), + Row( + id=14, + details=Row(label="PROPERTY", scores=[15]), + attributes={"count": 16}, + ), + ), + ( + "missing object attribute", + Result( + id=18, + details=MissingLabelDetails(scores=[19]), + attributes={"count": 20}, + ), + Row( + id=18, + details=Row(label=None, scores=[19]), + attributes={"count": 20}, + ), + ), ] for case_name, value, expected in cases: with self.subTest(case=case_name): self.assertEqual( - _normalize_user_value(value, table_type), expected + result_normalizer(value), expected ) with self.assertRaisesRegex(ValueError, "Expected 3 value"): - _normalize_user_value((1, 2), table_type) + result_normalizer((1, 2)) with self.assertRaisesRegex(TypeError, "Expected a Mapping"): - _normalize_user_value(object(), table_type) + result_normalizer(object()) + with self.assertRaisesRegex(AttributeError, "label lookup failed"): + result_normalizer( + { + "id": 21, + "details": FailingPropertyDetails(), + "attributes": {}, + }, + ) def test_invalid_declarations_fail_eagerly(self): def missing_return(value): return value + def unresolved_return(value): + return value + + unresolved_return.__annotations__ = { + "return": "UnavailableReturn" + } + def pandas_identity(values: pd.Series) -> pd.Series: return values @@ -535,6 +694,12 @@ def eval(self, value): TypeError, "Cannot infer return_dtype", ), + ( + "unresolved return", + lambda: pf.udf(unresolved_return), + TypeError, + "Cannot infer return_dtype", + ), ( "Table return type", lambda: pf.udf( @@ -615,6 +780,60 @@ def eval(self, value): class DataFrameUDFAdapterTests(unittest.TestCase): + def test_general_result_normalizers_are_bound_by_return_type(self): + from pyflink.dataframe.udf import ( + _DataFrameAsyncScalarFunctionAdapter, + _DataFrameScalarFunctionAdapter, + _UDFUsage, + _resolve_udf_source, + ) + + return_dtype = pf.DataType.struct( + { + "value": pf.DataType.int64(), + "labels": pf.DataType.list(pf.DataType.string()), + } + ) + + def describe(value): + return {"value": value, "labels": (str(value),)} + + async def describe_async(value): + return {"value": value, "labels": (str(value),)} + + sync_adapter = _DataFrameScalarFunctionAdapter( + _resolve_udf_source(describe), + return_dtype, + True, + _UDFUsage.EXPRESSION, + "general", + ) + async_adapter = _DataFrameAsyncScalarFunctionAdapter( + _resolve_udf_source(describe_async), + return_dtype, + True, + _UDFUsage.EXPRESSION, + "general", + ) + sync_adapter.open(object()) + async_adapter.open(object()) + expected = Row(value=3, labels=["3"]) + self.assertEqual(sync_adapter.eval(3), expected) + self.assertEqual(asyncio.run(async_adapter.eval(3)), expected) + + def identity(value): + return value + + leaf_adapter = _DataFrameScalarFunctionAdapter( + _resolve_udf_source(identity), + pf.DataType.int64(), + True, + _UDFUsage.EXPRESSION, + "general", + ) + leaf_adapter.open(object()) + self.assertIs(leaf_adapter._invocation(), identity) + def test_scalar_function_lifecycle_and_cleanup(self): from pyflink.dataframe.udf import ( _DataFrameAsyncScalarFunctionAdapter, diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index cda1f420785bc..21a7c3993efe9 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -28,6 +28,7 @@ Callable, Dict, FrozenSet, + List, Optional, Tuple, Type, @@ -55,6 +56,7 @@ _UDFInput = Union[Callable[..., Any], ScalarFunction, AsyncScalarFunction, Type] _DataTypeLike = Union[DataType, Type, str] +_UNRESOLVED_TYPE_HINT = object() class _UDFUsage(Enum): @@ -82,21 +84,28 @@ class _ResolvedUDFSource: is_async: bool ignored_hint_names: FrozenSet[str] = frozenset() - @property - def inspection_target(self) -> Callable[..., Any]: + def _inspection_target_and_skip_first( + self, + ) -> Tuple[Callable[..., Any], bool]: if self.kind is _UDFSourceKind.DIRECT_CALLABLE: - return _get_callable_inspection_target( - cast(Callable[..., Any], self.source) + return ( + _get_callable_inspection_target( + cast(Callable[..., Any], self.source) + ), + False, ) if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: - return cast( - Union[ScalarFunction, AsyncScalarFunction], self.source - ).eval + return ( + cast( + Union[ScalarFunction, AsyncScalarFunction], self.source + ).eval, + False, + ) if self.kind in ( _UDFSourceKind.CALLABLE_CLASS, _UDFSourceKind.SCALAR_FUNCTION_CLASS, ): - hint_method, _ = _get_callable_class_hint_method( + hint_method, skip_first_parameter = _get_callable_class_hint_method( cast(Type, self.source), "eval" if self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS @@ -104,8 +113,28 @@ def inspection_target(self) -> Callable[..., Any]: ) if hint_method is None: raise RuntimeError("Resolved UDF class has no inspection target.") - return hint_method - return cast(Callable[..., Any], getattr(self.source, "__call__")) + return hint_method, skip_first_parameter + return cast(Callable[..., Any], getattr(self.source, "__call__")), False + + @property + def inspection_target(self) -> Callable[..., Any]: + target, _ = self._inspection_target_and_skip_first() + return target + + @property + def invocation_signature(self) -> Optional[inspect.Signature]: + target, skip_first_parameter = self._inspection_target_and_skip_first() + if not self.is_scalar_function and not self.constructs_on_worker: + target = cast(Callable[..., Any], self.source) + + try: + signature = inspect.signature(target) + if skip_first_parameter: + parameters = tuple(signature.parameters.values()) + signature = signature.replace(parameters=parameters[1:]) + except Exception: + return None + return signature @property def default_name(self) -> str: @@ -232,6 +261,9 @@ def __init__( functools.update_wrapper(self, declaration_metadata, updated=()) object.__setattr__(self, "__name__", name) object.__setattr__(self, "__wrapped__", source.source) + invocation_signature = source.invocation_signature + if invocation_signature is not None: + object.__setattr__(self, "__signature__", invocation_signature) object.__setattr__(self, "_frozen", True) def __setattr__(self, name: str, value: Any) -> None: @@ -673,14 +705,14 @@ def _infer_return_dtype( if return_dtype is not None: return _convert_to_dtype(return_dtype) - hints = _get_callable_type_hints(func) - if "return" not in hints: + return_hint = _get_callable_return_type_hint(func) + if return_hint is _UNRESOLVED_TYPE_HINT: func_name = _default_udf_name(func) raise TypeError( f"Cannot infer return_dtype for '{func_name}': add a return annotation " "or specify return_dtype explicitly." ) - return _data_type_from_type_hint(hints["return"]) + return _data_type_from_type_hint(return_hint) def _convert_to_dtype(dtype_like: _DataTypeLike) -> DataType: @@ -766,19 +798,43 @@ def _get_callable_type_hints( try: if fallback_globals is None: return get_type_hints(func) - func_globals = getattr(func, "__globals__", None) - if func_globals is None: - func_globals = getattr( - getattr(func, "__func__", None), "__globals__", {} - ) return get_type_hints( func, - globalns={**fallback_globals, **func_globals}, + globalns={**fallback_globals, **_get_callable_globals(func)}, ) except (NameError, TypeError): return {} +def _get_callable_return_type_hint(func: Callable[..., Any]) -> Any: + annotations = getattr(func, "__annotations__", {}) + if "return" not in annotations: + return _UNRESOLVED_TYPE_HINT + + def return_annotation_holder() -> None: + pass + + return_annotation_holder.__annotations__ = { + "return": annotations["return"] + } + try: + return get_type_hints( + return_annotation_holder, + globalns=_get_callable_globals(func), + ).get("return", _UNRESOLVED_TYPE_HINT) + except (NameError, TypeError): + return _UNRESOLVED_TYPE_HINT + + +def _get_callable_globals(func: Callable[..., Any]) -> Dict[str, Any]: + func_globals = getattr(func, "__globals__", None) + if func_globals is None: + func_globals = getattr( + getattr(func, "__func__", None), "__globals__", {} + ) + return cast(Dict[str, Any], func_globals) + + def _resolve_deterministic( source: _ResolvedUDFSource, deterministic: bool ) -> bool: @@ -812,22 +868,22 @@ def _default_udf_name(func: _UDFInput) -> str: def _wrap_scalar_general_result( - func: Callable[..., Any], return_dtype: DataType, is_async: bool + func: Callable[..., Any], + result_normalizer: Callable[[Any], Any], + is_async: bool, ) -> Callable[..., Any]: - result_type = return_dtype._to_table_data_type() - if is_async: @functools.wraps(func) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: - return _normalize_user_value(await func(*args, **kwargs), result_type) + return result_normalizer(await func(*args, **kwargs)) wrapper = async_wrapper else: @functools.wraps(func) def sync_wrapper(*args: Any, **kwargs: Any) -> Any: - return _normalize_user_value(func(*args, **kwargs), result_type) + return result_normalizer(func(*args, **kwargs)) wrapper = sync_wrapper @@ -893,9 +949,14 @@ def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: f"DataFrame UDF usage {self._usage.value!r} is not supported yet." ) if self._func_type == "general": + result_normalizer = _create_result_normalizer( + cast(DataType, self._return_dtype)._to_table_data_type() + ) + if result_normalizer is None: + return invoke_func return _wrap_scalar_general_result( invoke_func, - cast(DataType, self._return_dtype), + result_normalizer, self._source.is_async, ) return invoke_func @@ -946,28 +1007,39 @@ async def eval(self, *args: Any) -> Any: # ======================== Result Normalization ======================== -def _row_value_by_type(value: Any, row_type: RowType, index: int) -> Any: - field_name = row_type.field_names()[index] +def _row_field_value( + value: Any, + field_name: str, + field_names: List[str], + field_count: int, + index: int, +) -> Any: if isinstance(value, Mapping): return value.get(field_name) if isinstance(value, Row) and hasattr(value, "_fields"): return _named_row_field_value(value, field_name) if isinstance(value, (Row, tuple, list)): - if len(value) != len(row_type.fields): + if len(value) != field_count: raise ValueError( - f"Expected {len(row_type.fields)} value(s) for RowType " - f"{row_type.field_names()}, got {len(value)}." + f"Expected {field_count} value(s) for RowType " + f"{field_names}, got {len(value)}." ) return value[index] - attributes = getattr(value, "__dict__", None) - if isinstance(attributes, Mapping): - return attributes.get(field_name) try: return getattr(value, field_name) except AttributeError: + try: + inspect.getattr_static(value, field_name) + except AttributeError: + attributes = getattr(value, "__dict__", None) + has_slots = any("__slots__" in cls.__dict__ for cls in type(value).__mro__) + if isinstance(attributes, Mapping) or has_slots: + return None + else: + raise raise TypeError( f"Expected a Mapping, Row, tuple, list, or object with fields for RowType " - f"{row_type.field_names()}, got {type(value).__name__}." + f"{field_names}, got {type(value).__name__}." ) from None @@ -985,52 +1057,89 @@ def _named_row_field_value(row: Row, field_name: str) -> Any: return row[field_name] -def _normalize_user_value(value: Any, data_type: Any) -> Any: - """Normalize nested user values to the Python shape expected by Table coders.""" - if value is None: - return None +def _create_result_normalizer( + data_type: Any, +) -> Optional[Callable[[Any], Any]]: if isinstance(data_type, RowType): - row = Row( - *[ - _normalize_user_value( - _row_value_by_type(value, data_type, index), field.data_type - ) - for index, field in enumerate(data_type) - ] + field_names = data_type.field_names() + field_count = len(data_type.fields) + field_normalizers = tuple( + _create_result_normalizer(field.data_type) for field in data_type ) - row.set_field_names(data_type.field_names()) - if isinstance(value, Row): - row.set_row_kind(value.get_row_kind()) - return row + + def normalize_row(value: Any) -> Any: + if value is None: + return None + normalized_fields = [] + for index, (field_name, field_normalizer) in enumerate( + zip(field_names, field_normalizers) + ): + field_value = _row_field_value( + value, field_name, field_names, field_count, index + ) + normalized_fields.append( + field_value + if field_normalizer is None + else field_normalizer(field_value) + ) + row = Row(*normalized_fields) + row.set_field_names(field_names) + if isinstance(value, Row): + row.set_row_kind(value.get_row_kind()) + return row + + return normalize_row if isinstance(data_type, ArrayType): - return [ - _normalize_user_value(item, data_type.element_type) for item in value - ] + element_normalizer = _create_result_normalizer(data_type.element_type) + if element_normalizer is None: + + def normalize_leaf_array(value: Any) -> Any: + return None if value is None else list(value) + + return normalize_leaf_array + + def normalize_array(value: Any) -> Any: + if value is None: + return None + return [element_normalizer(item) for item in value] + + return normalize_array if isinstance(data_type, MapType): - items_method = getattr(value, "items", None) - if callable(items_method): - items = list(items_method()) - else: - try: - items = list(value) - except TypeError as exc: + key_normalizer = _create_result_normalizer(data_type.key_type) + value_normalizer = _create_result_normalizer(data_type.value_type) + + def normalize_map(value: Any) -> Any: + if value is None: + return None + items_method = getattr(value, "items", None) + if callable(items_method): + items = list(items_method()) + else: + try: + items = list(value) + except TypeError as exc: + raise TypeError( + f"Expected a Mapping or iterable of key/value pairs for " + f"{data_type}, got {type(value).__name__}." + ) from exc + if any( + not isinstance(item, (tuple, list)) or len(item) != 2 + for item in items + ): raise TypeError( f"Expected a Mapping or iterable of key/value pairs for {data_type}, " f"got {type(value).__name__}." - ) from exc - if any( - not isinstance(item, (tuple, list)) or len(item) != 2 for item in items - ): - raise TypeError( - f"Expected a Mapping or iterable of key/value pairs for {data_type}, " - f"got {type(value).__name__}." - ) - if any(item[0] is None for item in items): - raise TypeError(f"MapType keys must not be null for {data_type}.") - return { - _normalize_user_value(key, data_type.key_type): _normalize_user_value( - item_value, data_type.value_type - ) - for key, item_value in items - } - return value + ) + if any(item[0] is None for item in items): + raise TypeError(f"MapType keys must not be null for {data_type}.") + return { + key if key_normalizer is None else key_normalizer(key): ( + item_value + if value_normalizer is None + else value_normalizer(item_value) + ) + for key, item_value in items + } + + return normalize_map + return None From 37776322c02ec300f784797670f670fdc8407353 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 20:52:33 +0800 Subject: [PATCH 06/18] [FLINK-40431][python] Hide DataFrame UDF wrapper implementation Generated-by: OpenAI Codex (GPT-5) --- .../docs/reference/pyflink.dataframe/udf.rst | 7 -- .../pyflink/dataframe/tests/test_udf.py | 35 ++++++---- flink-python/pyflink/dataframe/udf.py | 66 +++++-------------- 3 files changed, 37 insertions(+), 71 deletions(-) diff --git a/flink-python/docs/reference/pyflink.dataframe/udf.rst b/flink-python/docs/reference/pyflink.dataframe/udf.rst index 4d046bd61b2c8..73d1ff2abe76d 100644 --- a/flink-python/docs/reference/pyflink.dataframe/udf.rst +++ b/flink-python/docs/reference/pyflink.dataframe/udf.rst @@ -39,10 +39,3 @@ API Reference :toctree: api/ udf - -.. currentmodule:: pyflink.dataframe.udf - -.. autosummary:: - :toctree: api/ - - DataFrameUDFWrapper diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 098625df85d7b..b7f2f04045010 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -18,17 +18,19 @@ import asyncio import functools +import importlib import inspect import operator import unittest from dataclasses import dataclass -from typing import TypedDict +from typing import Any, Callable, TypedDict, cast import pandas as pd import pyarrow as pa import pyflink.dataframe as pf from pyflink.common import Row, RowKind from pyflink.table import DataTypes as TableDataTypes +from pyflink.table.expression import Expression from pyflink.table.types import RowType from pyflink.table.udf import AsyncScalarFunction, ScalarFunction, TableFunction from pyflink.testing.test_case_utils import ( @@ -37,6 +39,10 @@ ) +def _return_dtype(declaration: Callable[..., Expression]) -> pf.DataType: + return cast(Any, declaration).return_dtype + + class DataFrameUDFDeclarationTests(unittest.TestCase): def test_function_declarations_return_types_and_metadata(self): class Details(TypedDict): @@ -76,24 +82,25 @@ def postponed_return_with_unresolved_input(value): "return": "int", } - decorated = pf.udf(add_one) - - from pyflink.dataframe.udf import DataFrameUDFWrapper + decorated: Callable[..., Expression] = pf.udf(add_one) - self.assertIsInstance(decorated, DataFrameUDFWrapper) self.assertFalse(hasattr(pf, "DataFrameUDFWrapper")) - self.assertEqual(decorated.return_dtype, pf.DataType.int64()) + udf_module = importlib.import_module("pyflink.dataframe.udf") + self.assertFalse(hasattr(udf_module, "DataFrameUDFWrapper")) + self.assertEqual(_return_dtype(decorated), pf.DataType.int64()) self.assertEqual(decorated.__name__, "add_one") self.assertEqual(decorated.__doc__, "Add one to a value.") self.assertIs(decorated.__wrapped__, add_one) - configured = pf.udf(return_dtype=pf.DataType.string())( + configured: Callable[..., Expression] = pf.udf( + return_dtype=pf.DataType.string() + )( lambda value: str(value) ) direct = pf.udf(functools.partial(add_one), name="partial_add_one") - self.assertEqual(configured.return_dtype, pf.DataType.string()) - self.assertEqual(direct.return_dtype, pf.DataType.int64()) + self.assertEqual(_return_dtype(configured), pf.DataType.string()) + self.assertEqual(_return_dtype(direct), pf.DataType.int64()) self.assertEqual(direct.__name__, "partial_add_one") declarations = [ @@ -130,7 +137,7 @@ def postponed_return_with_unresolved_input(value): ] for case_name, declare, expected in declarations: with self.subTest(case=case_name): - self.assertEqual(declare().return_dtype, expected) + self.assertEqual(_return_dtype(declare()), expected) def test_callable_classes_and_instances_infer_from_invocation_method(self): plain_constructor_calls = [] @@ -197,7 +204,7 @@ def eval(self, *values: int) -> int: for source in callables: with self.subTest(source=source): decorated = pf.udf(source) - self.assertEqual(decorated.return_dtype, pf.DataType.int64()) + self.assertEqual(_return_dtype(decorated), pf.DataType.int64()) self.assertEqual(plain_constructor_calls, []) self.assertEqual(scalar_constructor_calls, []) @@ -267,9 +274,11 @@ def variadic_add(*values: int) -> int: ) uninspectable = pf.udf(operator.itemgetter(0), return_dtype=int) - self.assertEqual(uninspectable.return_dtype, pf.DataType.int64()) + self.assertEqual(_return_dtype(uninspectable), pf.DataType.int64()) exploding_signature = pf.udf(ExplodingSignature(), return_dtype=int) - self.assertEqual(exploding_signature.return_dtype, pf.DataType.int64()) + self.assertEqual( + _return_dtype(exploding_signature), pf.DataType.int64() + ) def test_func_type_resolution_and_async_detection(self): def pandas_add_one(values: pd.Series) -> pd.Series: diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 21a7c3993efe9..311e3e4deeccf 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -52,7 +52,7 @@ ) from pyflink.util.api_stability_decorators import PublicEvolving -__all__ = ["DataFrameUDFWrapper", "udf"] +__all__ = ["udf"] _UDFInput = Union[Callable[..., Any], ScalarFunction, AsyncScalarFunction, Type] _DataTypeLike = Union[DataType, Type, str] @@ -216,24 +216,8 @@ def close_worker_source(self, worker_source: Optional[_UDFInput]) -> None: ).close() -@PublicEvolving() -class DataFrameUDFWrapper: - """ - A callable DataFrame scalar UDF declaration. - - Instances are created with :func:`udf` and can be called with DataFrame - expressions or Python literals to produce an expression. - - Example:: - - >>> import pyflink.dataframe as pf - >>> @pf.udf - ... def add_one(value: int) -> int: - ... return value + 1 - >>> expression = add_one(pf.col("value")) - - .. versionadded:: 2.4.0 - """ +class _DataFrameUDFWrapper: + """Internal callable binding a DataFrame scalar UDF to Table expressions.""" _source: _ResolvedUDFSource _return_dtype: DataType @@ -268,22 +252,10 @@ def __init__( def __setattr__(self, name: str, value: Any) -> None: if getattr(self, "_frozen", False): - raise AttributeError("DataFrameUDFWrapper declarations are immutable.") + raise AttributeError("DataFrame UDF declarations are immutable.") object.__setattr__(self, name, value) - @PublicEvolving() def __call__(self, *args: Any) -> Expression: - """ - Create an expression that calls this UDF. - - Example:: - - >>> import pyflink.dataframe as pf - >>> @pf.udf - ... def add_one(value: int) -> int: - ... return value + 1 - >>> expression = add_one(pf.col("value")) - """ return table_call(self._table_udf_wrapper, *args) @property @@ -326,20 +298,7 @@ def _create_table_udf_wrapper( ) @property - @PublicEvolving() def return_dtype(self) -> DataType: - """ - The logical result type of this UDF. - - Example:: - - >>> import pyflink.dataframe as pf - >>> @pf.udf - ... def add_one(value: int) -> int: - ... return value + 1 - >>> add_one.return_dtype == pf.DataType.int64() - True - """ return self._return_dtype @@ -351,7 +310,7 @@ def udf( deterministic: bool = ..., name: Optional[str] = ..., func_type: Optional[str] = ..., -) -> DataFrameUDFWrapper: +) -> Callable[..., Expression]: ... @@ -363,7 +322,7 @@ def udf( deterministic: bool = ..., name: Optional[str] = ..., func_type: Optional[str] = ..., -) -> Callable[[_UDFInput], DataFrameUDFWrapper]: +) -> Callable[[_UDFInput], Callable[..., Expression]]: ... @@ -375,7 +334,10 @@ def udf( deterministic: bool = True, name: Optional[str] = None, func_type: Optional[str] = None, -) -> Union[DataFrameUDFWrapper, Callable[[_UDFInput], DataFrameUDFWrapper]]: +) -> Union[ + Callable[..., Expression], + Callable[[_UDFInput], Callable[..., Expression]], +]: """ Create a scalar UDF for DataFrame expressions. @@ -495,12 +457,14 @@ def udf( :param name: Non-empty function identity used by the Table planner. :param func_type: ``"general"`` or ``"pandas"``. If omitted, any unbound pandas container annotation selects pandas mode. - :return: A :class:`DataFrameUDFWrapper`, or a decorator when ``func`` is omitted. + :return: A callable that accepts DataFrame expressions or Python literals and + returns an :class:`~pyflink.table.expression.Expression`, or a decorator + producing such a callable when ``func`` is omitted. .. versionadded:: 2.4.0 """ - def decorator(f: _UDFInput) -> DataFrameUDFWrapper: + def decorator(f: _UDFInput) -> Callable[..., Expression]: source = _resolve_udf_source(f) actual_func_type = ( func_type @@ -516,7 +480,7 @@ def decorator(f: _UDFInput) -> DataFrameUDFWrapper: actual_deterministic = _resolve_deterministic(source, deterministic) actual_name = _resolve_name(source, name) - return DataFrameUDFWrapper( + return _DataFrameUDFWrapper( source, actual_return_dtype, actual_deterministic, From 332c1846818734feaf72e9548ddf08a876ec0996 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 22:01:20 +0800 Subject: [PATCH 07/18] [FLINK-40431][python] Optimize DataFrame UDF Row normalization Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 1 - flink-python/pyflink/dataframe/udf.py | 91 ++++++++++--------- 2 files changed, 46 insertions(+), 46 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index b7f2f04045010..7ab149b747367 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -617,7 +617,6 @@ class Result: self.assertEqual( result_normalizer(value), expected ) - with self.assertRaisesRegex(ValueError, "Expected 3 value"): result_normalizer((1, 2)) with self.assertRaisesRegex(TypeError, "Expected a Mapping"): diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 311e3e4deeccf..5849f931ed62a 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -341,11 +341,6 @@ def udf( """ Create a scalar UDF for DataFrame expressions. - The function may be synchronous or asynchronous. Pandas UDFs operate on - ``pandas.Series`` or ``pandas.DataFrame`` batches and must declare - ``return_dtype``. Callable and scalar-function class objects must have a - zero-argument constructor and are initialized on the TaskManager. - A UDF can be declared with a bare decorator, a configured decorator, or a direct call. General UDFs may infer ``return_dtype`` from the return annotation of the function, ``__call__``, or ``eval``. A ``TypedDict`` @@ -422,10 +417,13 @@ def udf( ... async def async_add_one(value: int) -> int: ... return value + 1 - Pandas UDFs receive and return ``pandas.Series`` or ``pandas.DataFrame`` - batches and always require an explicit logical ``return_dtype``. Pandas - mode can be selected explicitly, or inferred from a pandas container - annotation on any unbound parameter or the return value:: + Pandas UDFs always require an explicit logical ``return_dtype``. Each + ``ROW``-typed argument is received as a ``pandas.DataFrame`` with one column + per field; other arguments are received as ``pandas.Series``. A ``ROW``-typed + result should be returned as a ``pandas.DataFrame``, while other results + should be returned as ``pandas.Series``. Pandas mode can be selected + explicitly, or inferred from a pandas container annotation on any unbound + parameter or the return value:: >>> import pandas as pd @@ -971,24 +969,45 @@ async def eval(self, *args: Any) -> Any: # ======================== Result Normalization ======================== -def _row_field_value( - value: Any, - field_name: str, - field_names: List[str], - field_count: int, - index: int, -) -> Any: +def _row_field_values(value: Any, field_names: List[str]) -> List[Any]: if isinstance(value, Mapping): - return value.get(field_name) + return [value.get(field_name) for field_name in field_names] if isinstance(value, Row) and hasattr(value, "_fields"): - return _named_row_field_value(value, field_name) + field_indices: Dict[str, int] = {} + for index, field_name in enumerate(value._fields): + field_indices.setdefault(field_name, index) + field_values: List[Any] = [] + for field_name in field_names: + if field_name not in field_indices: + raise ValueError( + f"Field name {field_name!r} does not exist in Row fields " + f"{value._fields}." + ) + field_index = field_indices[field_name] + if field_index >= len(value): + raise ValueError( + f"Field name {field_name!r} is declared in Row fields " + f"{value._fields} but has no value." + ) + field_values.append(value[field_index]) + return field_values if isinstance(value, (Row, tuple, list)): + field_count = len(field_names) if len(value) != field_count: raise ValueError( f"Expected {field_count} value(s) for RowType " f"{field_names}, got {len(value)}." ) - return value[index] + return list(value) + return [ + _object_row_field_value(value, field_name, field_names) + for field_name in field_names + ] + + +def _object_row_field_value( + value: Any, field_name: str, field_names: List[str] +) -> Any: try: return getattr(value, field_name) except AttributeError: @@ -1007,26 +1026,11 @@ def _row_field_value( ) from None -def _named_row_field_value(row: Row, field_name: str) -> Any: - if field_name not in row._fields: - raise ValueError( - f"Field name {field_name!r} does not exist in Row fields {row._fields}." - ) - field_index = row._fields.index(field_name) - if field_index >= len(row): - raise ValueError( - f"Field name {field_name!r} is declared in Row fields {row._fields} " - "but has no value." - ) - return row[field_name] - - def _create_result_normalizer( data_type: Any, ) -> Optional[Callable[[Any], Any]]: if isinstance(data_type, RowType): field_names = data_type.field_names() - field_count = len(data_type.fields) field_normalizers = tuple( _create_result_normalizer(field.data_type) for field in data_type ) @@ -1034,18 +1038,15 @@ def _create_result_normalizer( def normalize_row(value: Any) -> Any: if value is None: return None - normalized_fields = [] - for index, (field_name, field_normalizer) in enumerate( - zip(field_names, field_normalizers) - ): - field_value = _row_field_value( - value, field_name, field_names, field_count, index - ) - normalized_fields.append( - field_value - if field_normalizer is None - else field_normalizer(field_value) + field_values = _row_field_values(value, field_names) + normalized_fields = [ + field_value + if field_normalizer is None + else field_normalizer(field_value) + for field_value, field_normalizer in zip( + field_values, field_normalizers ) + ] row = Row(*normalized_fields) row.set_field_names(field_names) if isinstance(value, Row): From f7eedd4e6d489680ff74fa58c61430d53b9c72f3 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 22:31:50 +0800 Subject: [PATCH 08/18] [FLINK-40431][python] Clarify DataFrame UDF type handling Generated-by: OpenAI Codex (GPT-5) --- flink-python/pyflink/dataframe/dataframe.py | 2 +- flink-python/pyflink/dataframe/udf.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 6cfe09f662d9c..27185987d13cf 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -200,7 +200,7 @@ def with_column( Add a column, or replace an existing column with the same name. ``expr`` may be an expression or a callable that receives this DataFrame and returns an - expression. See :func:`~pyflink.dataframe.udf.udf` for supported UDF declaration forms. + expression. :param name: Name of the added or replaced column. :param expr: Expression or callable used to compute the column value. diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 5849f931ed62a..3247a727e8100 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -28,6 +28,7 @@ Callable, Dict, FrozenSet, + Iterable, List, Optional, Tuple, @@ -773,6 +774,8 @@ def _get_callable_return_type_hint(func: Callable[..., Any]) -> Any: if "return" not in annotations: return _UNRESOLVED_TYPE_HINT + # Resolve the return annotation in isolation so an unresolvable parameter + # annotation does not prevent return-type inference. def return_annotation_holder() -> None: pass @@ -1078,7 +1081,7 @@ def normalize_map(value: Any) -> Any: return None items_method = getattr(value, "items", None) if callable(items_method): - items = list(items_method()) + items = list(cast(Iterable[Any], items_method())) else: try: items = list(value) From 3ee2651595f91af4e6216d354113c4f8ca8feaf0 Mon Sep 17 00:00:00 2001 From: auroflow Date: Sun, 30 Aug 2026 12:49:22 +0800 Subject: [PATCH 09/18] [FLINK-40431][python] Fix DataFrame UDF type hint resolution Resolve callable annotations independently so unrelated unresolved hints do not hide pandas annotations. Reuse recursive type-hint conversion for explicit TypedDict return types. Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 46 ++++++++++---- flink-python/pyflink/dataframe/udf.py | 63 +++++++++---------- 2 files changed, 65 insertions(+), 44 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 7ab149b747367..ee03c3ae4605c 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -103,6 +103,17 @@ def postponed_return_with_unresolved_input(value): self.assertEqual(_return_dtype(direct), pf.DataType.int64()) self.assertEqual(direct.__name__, "partial_add_one") + expected_result_dtype = pf.DataType.struct( + { + "id": pf.DataType.int64(), + "details": pf.DataType.struct( + { + "label": pf.DataType.string(), + "scores": pf.DataType.list(pf.DataType.int64()), + } + ), + } + ) declarations = [ ( "Python type", @@ -112,17 +123,12 @@ def postponed_return_with_unresolved_input(value): ( "nested TypedDict annotation", lambda: pf.udf(describe), - pf.DataType.struct( - { - "id": pf.DataType.int64(), - "details": pf.DataType.struct( - { - "label": pf.DataType.string(), - "scores": pf.DataType.list(pf.DataType.int64()), - } - ), - } - ), + expected_result_dtype, + ), + ( + "explicit nested TypedDict", + lambda: pf.udf(identity, return_dtype=Result), + expected_result_dtype, ), ( "concrete return with unresolved input", @@ -292,6 +298,15 @@ def pandas_forward_reference(values): pandas_forward_reference.__annotations__["values"] = "pandas.Series" + def pandas_with_unresolved_annotation( + values: pd.Series, context + ) -> pd.Series: + return values + + pandas_with_unresolved_annotation.__annotations__[ + "context" + ] = "UnavailableContext" + def mixed(values: pd.Series, offset: int): return values + offset @@ -342,6 +357,15 @@ async def eval(self, *values: int) -> int: "pandas", False, ), + ( + "unresolved annotation does not hide pandas annotation", + lambda: pf.udf( + pandas_with_unresolved_annotation, + return_dtype=pf.DataType.int64(), + ), + "pandas", + False, + ), ( "any pandas annotation selects pandas", lambda: pf.udf(mixed, return_dtype=pf.DataType.int64()), diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 3247a727e8100..02a4d148637b0 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -684,7 +684,7 @@ def _convert_to_dtype(dtype_like: _DataTypeLike) -> DataType: if isinstance(dtype_like, str): return DataType._from_sql(dtype_like) try: - return DataType._from_type_hint(dtype_like) + return _data_type_from_type_hint(dtype_like) except TypeError as exc: raise TypeError( "return_dtype must be a DataFrame DataType, Python type, or SQL " @@ -725,19 +725,19 @@ def _detect_func_type(source: _ResolvedUDFSource) -> str: import pandas as pd except ImportError: return "general" - hints = _get_callable_type_hints( - hint_func, fallback_globals={"pandas": pd, "pd": pd} - ) pandas_types = (pd.Series, pd.DataFrame) - return ( - "pandas" - if any( - name not in source.ignored_hint_names and hint in pandas_types - for name, hint in hints.items() + for name in getattr(hint_func, "__annotations__", {}): + if name in source.ignored_hint_names: + continue + hint = _resolve_callable_annotation( + hint_func, + name, + fallback_globals={"pandas": pd, "pd": pd}, ) - else "general" - ) + if hint in pandas_types: + return "pandas" + return "general" def _unwrap_partial(func: Any) -> Any: @@ -755,38 +755,35 @@ def _get_callable_inspection_target( return cast(Callable[..., Any], target) -def _get_callable_type_hints( - func: Callable[..., Any], fallback_globals: Optional[Dict[str, Any]] = None -) -> Dict[str, Any]: - try: - if fallback_globals is None: - return get_type_hints(func) - return get_type_hints( - func, - globalns={**fallback_globals, **_get_callable_globals(func)}, - ) - except (NameError, TypeError): - return {} +def _get_callable_return_type_hint(func: Callable[..., Any]) -> Any: + # Resolve the return annotation in isolation so an unresolvable parameter + # annotation does not prevent return-type inference. + return _resolve_callable_annotation(func, "return") -def _get_callable_return_type_hint(func: Callable[..., Any]) -> Any: +def _resolve_callable_annotation( + func: Callable[..., Any], + annotation_name: str, + fallback_globals: Optional[Dict[str, Any]] = None, +) -> Any: annotations = getattr(func, "__annotations__", {}) - if "return" not in annotations: + if annotation_name not in annotations: return _UNRESOLVED_TYPE_HINT - # Resolve the return annotation in isolation so an unresolvable parameter - # annotation does not prevent return-type inference. - def return_annotation_holder() -> None: + def annotation_holder() -> None: pass - return_annotation_holder.__annotations__ = { - "return": annotations["return"] + annotation_holder.__annotations__ = { + annotation_name: annotations[annotation_name] } try: return get_type_hints( - return_annotation_holder, - globalns=_get_callable_globals(func), - ).get("return", _UNRESOLVED_TYPE_HINT) + annotation_holder, + globalns={ + **(fallback_globals or {}), + **_get_callable_globals(func), + }, + ).get(annotation_name, _UNRESOLVED_TYPE_HINT) except (NameError, TypeError): return _UNRESOLVED_TYPE_HINT From 508a3074d9810ebae13874e4368cdcb0b7166051 Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 1 Sep 2026 11:49:03 +0800 Subject: [PATCH 10/18] [FLINK-40431][python] Fix callable UDF annotation resolution Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 289 +++++++++- flink-python/pyflink/dataframe/udf.py | 541 +++++++++++------- 2 files changed, 629 insertions(+), 201 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index ee03c3ae4605c..bcca450d492cc 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -21,6 +21,7 @@ import importlib import inspect import operator +import types import unittest from dataclasses import dataclass from typing import Any, Callable, TypedDict, cast @@ -43,6 +44,21 @@ def _return_dtype(declaration: Callable[..., Expression]) -> pf.DataType: return cast(Any, declaration).return_dtype +_UDF_TEST_ALIAS = int + + +def _module_alias_method(self, value: int) -> "_UDF_TEST_ALIAS": + return value + + +def _module_alias_function(value: int) -> "_UDF_TEST_ALIAS": + return value + + +def _call_module_alias_function(value): + return _module_alias_function(value) + + class DataFrameUDFDeclarationTests(unittest.TestCase): def test_function_declarations_return_types_and_metadata(self): class Details(TypedDict): @@ -220,6 +236,84 @@ def eval(self, *values: int) -> int: self.assertIs(decorated_class.__wrapped__, Double) self.assertEqual(decorated_class.__qualname__, Double.__qualname__) + def test_callable_class_resolves_class_local_return_annotation(self): + class Describe: + class Output(TypedDict): + value: int + + def __call__(self, value: int) -> "Output": + return {"value": value} + + expected = pf.DataType.struct({"value": pf.DataType.int64()}) + for source in (Describe, Describe()): + with self.subTest(source=source): + self.assertEqual(_return_dtype(pf.udf(source)), expected) + + def test_callable_annotations_use_lexical_defining_class(self): + class BoundMethodOwner: + class Output(TypedDict): + value: int + + def describe(self, value: int) -> "Output": + return {"value": value} + + class InheritedMethodOwner: + class Output(TypedDict): + value: int + + def __call__(self, value: int) -> "Output": + return {"value": value} + + class InheritedCallable(InheritedMethodOwner): + pass + + class SelfQualified: + class Output(TypedDict): + value: int + + def __call__(self, value: int) -> "SelfQualified.Output": + return {"value": value} + + expected = pf.DataType.struct({"value": pf.DataType.int64()}) + bound_method = BoundMethodOwner().describe + for source in ( + bound_method, + functools.partial(bound_method), + InheritedCallable, + SelfQualified, + ): + with self.subTest(source=source): + self.assertEqual(_return_dtype(pf.udf(source)), expected) + + class PandasCallable: + Batch = pd.Series + + def __call__(self, values: "Batch") -> int: + return len(values) + + pandas_declaration = pf.udf(PandasCallable, return_dtype=int) + self.assertEqual(pandas_declaration._func_type, "pandas") + + class ReceivingCallable: + _UDF_TEST_ALIAS = str + __call__ = _module_alias_method + + self.assertEqual( + _return_dtype(pf.udf(ReceivingCallable)), pf.DataType.int64() + ) + + class OverriddenScalarFunction(ScalarFunction): + _UDF_TEST_ALIAS = str + + def eval(self, value: int) -> str: + return str(value) + + overridden = OverriddenScalarFunction() + overridden.eval = types.MethodType(_module_alias_method, overridden) + self.assertEqual( + _return_dtype(pf.udf(overridden)), pf.DataType.int64() + ) + def test_wrapped_signature_describes_udf_invocation(self): def add(value: int, amount: int = 1) -> int: return value + amount @@ -238,6 +332,20 @@ class ClassMethodCallableClass: def __call__(cls, value: int, amount: int = 1) -> int: return value + amount + def pandas_identity(values: pd.Series) -> pd.Series: + return values + + class WrappedCallableClass: + @functools.wraps(pandas_identity) + def __call__(self, *args, **kwargs): + return pandas_identity(*args, **kwargs) + + class WrappedClassMethodCallableClass: + @classmethod + @functools.wraps(pandas_identity) + def __call__(cls, *args, **kwargs): + return pandas_identity(*args, **kwargs) + class AddFunction(ScalarFunction): def eval(self, *values: int) -> int: return sum(values) @@ -285,6 +393,37 @@ def variadic_add(*values: int) -> int: self.assertEqual( _return_dtype(exploding_signature), pf.DataType.int64() ) + exploding_partial_signature = pf.udf( + functools.partial(ExplodingSignature()), return_dtype=int + ) + self.assertEqual( + _return_dtype(exploding_partial_signature), pf.DataType.int64() + ) + self.assertNotIn("__signature__", vars(exploding_partial_signature)) + + for source in (WrappedCallableClass, WrappedClassMethodCallableClass): + with self.subTest(wrapped_source=source): + declaration = pf.udf( + source, return_dtype=pf.DataType.int64() + ) + self.assertEqual( + inspect.signature(declaration), + inspect.signature(pandas_identity), + ) + self.assertEqual(declaration._func_type, "pandas") + + cross_namespace_wrapper = types.FunctionType( + _call_module_alias_function.__code__, + { + "_module_alias_function": _module_alias_function, + "_UDF_TEST_ALIAS": str, + }, + ) + functools.update_wrapper(cross_namespace_wrapper, _module_alias_function) + self.assertEqual( + _return_dtype(pf.udf(cross_namespace_wrapper)), + pf.DataType.int64(), + ) def test_func_type_resolution_and_async_detection(self): def pandas_add_one(values: pd.Series) -> pd.Series: @@ -307,6 +446,20 @@ def pandas_with_unresolved_annotation( "context" ] = "UnavailableContext" + def pandas_after_missing_attribute( + context: Any, values: pd.Series + ) -> int: + return len(values) + + pandas_after_missing_attribute.__annotations__[ + "context" + ] = "pd.Missing" + + def only_missing_attribute(context: Any) -> int: + return 1 + + only_missing_attribute.__annotations__["context"] = "pd.Missing" + def mixed(values: pd.Series, offset: int): return values + offset @@ -366,6 +519,24 @@ async def eval(self, *values: int) -> int: "pandas", False, ), + ( + "missing annotation attribute does not hide pandas annotation", + lambda: pf.udf( + pandas_after_missing_attribute, + return_dtype=pf.DataType.int64(), + ), + "pandas", + False, + ), + ( + "missing annotation attribute falls back to general", + lambda: pf.udf( + only_missing_attribute, + return_dtype=pf.DataType.int64(), + ), + "general", + False, + ), ( "any pandas annotation selects pandas", lambda: pf.udf(mixed, return_dtype=pf.DataType.int64()), @@ -423,7 +594,9 @@ async def eval(self, *values: int) -> int: with self.subTest(case=case_name): wrapped = declare() self.assertEqual(wrapped._func_type, expected_type) - self.assertEqual(wrapped._source.is_async, expected_async) + self.assertEqual( + wrapped._runtime_source.is_async, expected_async + ) invalid_declarations = [ ( @@ -448,6 +621,104 @@ async def eval(self, *values: int) -> int: with self.assertRaisesRegex(error_type, message): declare() + def test_sync_wrapper_around_async_target_is_rejected(self): + async def async_add_one(value: int) -> int: + return value + 1 + + @functools.wraps(async_add_one) + def sync_wrapper(*args, **kwargs): + return async_add_one(*args, **kwargs) + + with self.assertRaisesRegex(TypeError, "async def"): + pf.udf(sync_wrapper) + + def test_invalid_class_invocation_descriptors_fail_eagerly(self): + class CallableBase: + def __call__(self, value: int) -> int: + return value + + class HiddenCallable(CallableBase): + __call__ = None + + class ScalarBase(ScalarFunction): + def eval(self, value: int) -> int: + return value + + class HiddenScalarFunction(ScalarBase): + eval = None + + class InvalidStaticCallable: + __call__ = staticmethod(None) + + class InvalidClassMethodCallable: + __call__ = classmethod(None) + + invalid_classes = ( + HiddenCallable, + HiddenScalarFunction, + InvalidStaticCallable, + InvalidClassMethodCallable, + ) + for source in invalid_classes: + with self.subTest(source=source): + with self.assertRaisesRegex(TypeError, "must be defined as a method"): + pf.udf(source, return_dtype=int) + + def test_descriptor_based_callable_classes_require_instances(self): + class PartialMethodCallable: + def invoke(self, offset: int, value: int) -> int: + return offset + value + + __call__ = functools.partialmethod(invoke, 1) + + class PartialDescriptorCallable: + __call__ = functools.partial(lambda: 1) + + for source in (PartialMethodCallable, PartialDescriptorCallable): + with self.subTest(class_source=source): + with self.assertRaisesRegex( + TypeError, "must be defined as a method" + ): + pf.udf(source, return_dtype=int) + + with self.subTest(instance_source=source): + declaration = pf.udf(source(), return_dtype=int) + self.assertEqual( + _return_dtype(declaration), pf.DataType.int64() + ) + + def test_unresolved_typed_dict_fields_have_actionable_errors(self): + class Describe: + OuterAlias = int + + class Output(TypedDict): + value: Any + + Output.__annotations__["value"] = "OuterAlias" + + def __call__(self, value: int) -> "Output": + return {"value": value} + + with self.assertRaisesRegex(TypeError, "Cannot infer return_dtype"): + pf.udf(Describe) + + with self.assertRaisesRegex(TypeError, "DataType or SQL"): + pf.udf(lambda value: value, return_dtype=Describe.Output) + + class InvalidOutput(TypedDict): + value: Any + + InvalidOutput.__annotations__["value"] = "list[" + + def invalid_output(value: int) -> InvalidOutput: + return {"value": value} + + with self.assertRaisesRegex(TypeError, "Cannot infer return_dtype"): + pf.udf(invalid_output) + + with self.assertRaisesRegex(TypeError, "DataType or SQL"): + pf.udf(lambda value: value, return_dtype=InvalidOutput) + def test_determinism_and_name_metadata(self): class NonDeterministic(ScalarFunction): def eval(self, *values: int) -> int: @@ -817,7 +1088,7 @@ def test_general_result_normalizers_are_bound_by_return_type(self): _DataFrameAsyncScalarFunctionAdapter, _DataFrameScalarFunctionAdapter, _UDFUsage, - _resolve_udf_source, + _resolve_udf, ) return_dtype = pf.DataType.struct( @@ -834,14 +1105,14 @@ async def describe_async(value): return {"value": value, "labels": (str(value),)} sync_adapter = _DataFrameScalarFunctionAdapter( - _resolve_udf_source(describe), + _resolve_udf(describe).runtime_source, return_dtype, True, _UDFUsage.EXPRESSION, "general", ) async_adapter = _DataFrameAsyncScalarFunctionAdapter( - _resolve_udf_source(describe_async), + _resolve_udf(describe_async).runtime_source, return_dtype, True, _UDFUsage.EXPRESSION, @@ -857,7 +1128,7 @@ def identity(value): return value leaf_adapter = _DataFrameScalarFunctionAdapter( - _resolve_udf_source(identity), + _resolve_udf(identity).runtime_source, pf.DataType.int64(), True, _UDFUsage.EXPRESSION, @@ -871,7 +1142,7 @@ def test_scalar_function_lifecycle_and_cleanup(self): _DataFrameAsyncScalarFunctionAdapter, _DataFrameScalarFunctionAdapter, _UDFUsage, - _resolve_udf_source, + _resolve_udf, ) events = [] @@ -883,7 +1154,7 @@ def create_adapter(source, deterministic=True, async_mode=False): else _DataFrameScalarFunctionAdapter ) return adapter_type( - _resolve_udf_source(source), + _resolve_udf(source).runtime_source, pf.DataType.int64(), deterministic, _UDFUsage.EXPRESSION, @@ -1053,7 +1324,7 @@ def test_binding_failure_closes_and_resets_deferred_scalar_class(self): from pyflink.dataframe.udf import ( _DataFrameScalarFunctionAdapter, _UDFUsage, - _resolve_udf_source, + _resolve_udf, ) events = [] @@ -1073,7 +1344,7 @@ def close(self): raise RuntimeError("close failed") adapter = _DataFrameScalarFunctionAdapter( - _resolve_udf_source(BindingFailureFunction), + _resolve_udf(BindingFailureFunction).runtime_source, pf.DataType.int64(), True, _UDFUsage.MAP, diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 02a4d148637b0..db6cf3d5cd878 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -77,69 +77,28 @@ class _UDFSourceKind(Enum): @dataclass(frozen=True) -class _ResolvedUDFSource: - """Callable metadata resolved once on the client and reused on workers.""" +class _UDFDeclarationContext: + """Client-only metadata used while declaring a UDF.""" - source: _UDFInput - kind: _UDFSourceKind - is_async: bool - ignored_hint_names: FrozenSet[str] = frozenset() + annotation_target: Callable[..., Any] + defining_class: Optional[Type] + globalns: Dict[str, Any] + localns: Optional[Dict[str, Any]] + invocation_signature: Optional[inspect.Signature] + ignored_hint_names: FrozenSet[str] - def _inspection_target_and_skip_first( - self, - ) -> Tuple[Callable[..., Any], bool]: - if self.kind is _UDFSourceKind.DIRECT_CALLABLE: - return ( - _get_callable_inspection_target( - cast(Callable[..., Any], self.source) - ), - False, - ) - if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: - return ( - cast( - Union[ScalarFunction, AsyncScalarFunction], self.source - ).eval, - False, - ) - if self.kind in ( - _UDFSourceKind.CALLABLE_CLASS, - _UDFSourceKind.SCALAR_FUNCTION_CLASS, - ): - hint_method, skip_first_parameter = _get_callable_class_hint_method( - cast(Type, self.source), - "eval" - if self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS - else "__call__", - ) - if hint_method is None: - raise RuntimeError("Resolved UDF class has no inspection target.") - return hint_method, skip_first_parameter - return cast(Callable[..., Any], getattr(self.source, "__call__")), False - @property - def inspection_target(self) -> Callable[..., Any]: - target, _ = self._inspection_target_and_skip_first() - return target - - @property - def invocation_signature(self) -> Optional[inspect.Signature]: - target, skip_first_parameter = self._inspection_target_and_skip_first() - if not self.is_scalar_function and not self.constructs_on_worker: - target = cast(Callable[..., Any], self.source) +@dataclass(frozen=True) +class _UDFRuntimeSource: + """Worker-facing metadata used to initialize and invoke a UDF.""" - try: - signature = inspect.signature(target) - if skip_first_parameter: - parameters = tuple(signature.parameters.values()) - signature = signature.replace(parameters=parameters[1:]) - except Exception: - return None - return signature + callable_source: _UDFInput + kind: _UDFSourceKind + is_async: bool @property def default_name(self) -> str: - return _default_udf_name(self.source) + return _default_udf_name(self.callable_source) @property def is_scalar_function(self) -> bool: @@ -157,8 +116,8 @@ def constructs_on_worker(self) -> bool: def create_worker_source(self) -> _UDFInput: if not self.constructs_on_worker: - return self.source - source_class = cast(Type, self.source) + return self.callable_source + source_class = cast(Type, self.callable_source) source = source_class() if self.is_scalar_function: if not isinstance(source, (ScalarFunction, AsyncScalarFunction)): @@ -178,7 +137,7 @@ def validate_deterministic( ) -> None: source: Optional[_UDFInput] if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: - source = self.source + source = self.callable_source elif self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS: source = worker_source else: @@ -217,10 +176,18 @@ def close_worker_source(self, worker_source: Optional[_UDFInput]) -> None: ).close() +@dataclass(frozen=True) +class _ResolvedUDF: + """A resolved declaration split into client and worker metadata.""" + + runtime_source: _UDFRuntimeSource + declaration_context: _UDFDeclarationContext + + class _DataFrameUDFWrapper: """Internal callable binding a DataFrame scalar UDF to Table expressions.""" - _source: _ResolvedUDFSource + _runtime_source: _UDFRuntimeSource _return_dtype: DataType _deterministic: bool _func_type: str @@ -230,23 +197,23 @@ class _DataFrameUDFWrapper: def __init__( self, - source: _ResolvedUDFSource, + runtime_source: _UDFRuntimeSource, return_dtype: DataType, deterministic: bool, name: str, func_type: str, + invocation_signature: Optional[inspect.Signature], ) -> None: - object.__setattr__(self, "_source", source) + object.__setattr__(self, "_runtime_source", runtime_source) object.__setattr__(self, "_return_dtype", return_dtype) object.__setattr__(self, "_deterministic", deterministic) object.__setattr__(self, "_func_type", func_type) object.__setattr__(self, "_cached_table_udf_wrapper", None) - declaration_metadata = _unwrap_partial(source.source) + declaration_metadata = _unwrap_partial(runtime_source.callable_source) functools.update_wrapper(self, declaration_metadata, updated=()) object.__setattr__(self, "__name__", name) - object.__setattr__(self, "__wrapped__", source.source) - invocation_signature = source.invocation_signature + object.__setattr__(self, "__wrapped__", runtime_source.callable_source) if invocation_signature is not None: object.__setattr__(self, "__signature__", invocation_signature) object.__setattr__(self, "_frozen", True) @@ -274,13 +241,13 @@ def _create_table_udf_wrapper( ) -> UserDefinedFunctionWrapper: adapter_type = ( _DataFrameAsyncScalarFunctionAdapter - if self._source.is_async + if self._runtime_source.is_async else _DataFrameScalarFunctionAdapter ) actual_func = cast( Union[ScalarFunction, AsyncScalarFunction], adapter_type( - self._source, + self._runtime_source, self._return_dtype, self._deterministic, usage, @@ -464,27 +431,32 @@ def udf( """ def decorator(f: _UDFInput) -> Callable[..., Expression]: - source = _resolve_udf_source(f) + resolved_udf = _resolve_udf(f) + runtime_source = resolved_udf.runtime_source + declaration_context = resolved_udf.declaration_context actual_func_type = ( func_type if func_type is not None - else _detect_func_type(source) + else _detect_func_type(declaration_context) ) _validate_scalar_udf_options( - actual_func_type, return_dtype, source.is_async + actual_func_type, return_dtype, runtime_source.is_async ) actual_return_dtype = _infer_return_dtype( - source.inspection_target, return_dtype + declaration_context, return_dtype + ) + actual_deterministic = _resolve_deterministic( + runtime_source, deterministic ) - actual_deterministic = _resolve_deterministic(source, deterministic) - actual_name = _resolve_name(source, name) + actual_name = _resolve_name(runtime_source, name) return _DataFrameUDFWrapper( - source, + runtime_source, actual_return_dtype, actual_deterministic, actual_name, actual_func_type, + declaration_context.invocation_signature, ) return decorator if func is None else decorator(func) @@ -517,38 +489,185 @@ def _validate_scalar_udf_options( # ======================== Callable Inspection and Resolution ======================== -def _has_custom_call(cls: Type) -> bool: - """Check whether a class defines ``__call__`` in its MRO.""" - return any("__call__" in base.__dict__ for base in cls.__mro__ if base is not object) +def _first_parameter_name(func: Callable[..., Any]) -> Optional[str]: + try: + parameters = tuple( + inspect.signature(func, follow_wrapped=False).parameters.values() + ) + except (TypeError, ValueError): + return None + return parameters[0].name if parameters else None + + +def _resolve_class_invocation_target( + func_class: Type, method_name: str +) -> Tuple[Optional[Callable[..., Any]], Optional[Type], Optional[str]]: + """Resolve the nearest supported invocation method without constructing a class.""" + descriptor_owner = None + descriptor = None + for candidate in func_class.__mro__: + if method_name in candidate.__dict__: + descriptor_owner = candidate + descriptor = candidate.__dict__[method_name] + break + + if descriptor_owner is None: + return None, None, None + if isinstance(descriptor, staticmethod): + target = descriptor.__func__ + implicit_parameter_name = None + elif isinstance(descriptor, classmethod): + target = descriptor.__func__ + implicit_parameter_name = ( + _first_parameter_name(target) if inspect.isroutine(target) else None + ) + elif inspect.isroutine(descriptor): + target = descriptor + implicit_parameter_name = _first_parameter_name(target) + else: + return None, descriptor_owner, None + if not callable(target) or not inspect.isroutine(target): + return None, descriptor_owner, None + return cast(Callable[..., Any], target), descriptor_owner, implicit_parameter_name -def _get_callable_class_hint_method( - func_class: Type, method_name: str = "__call__" -) -> Tuple[Optional[Callable[..., Any]], bool]: - """Return a class method that can be inspected without constructing the class.""" - descriptor = inspect.getattr_static(func_class, method_name, None) - if isinstance(descriptor, staticmethod): - return cast(Callable[..., Any], descriptor.__func__), False - if isinstance(descriptor, classmethod): - return cast(Callable[..., Any], descriptor.__func__), True - if inspect.isroutine(descriptor): - return cast(Callable[..., Any], descriptor), True - return None, False + +def _function_qualname(func: Callable[..., Any]) -> Optional[str]: + target = _unwrap_partial(func) + target = getattr(target, "__func__", target) + qualname = getattr(target, "__qualname__", None) + return qualname if isinstance(qualname, str) else None + + +def _lexical_defining_class( + target: Callable[..., Any], candidate: Optional[Type] = None +) -> Optional[Type]: + qualname = _function_qualname(target) + if qualname is None: + return None + owner_qualname, separator, _ = qualname.rpartition(".") + if not separator: + return None + if candidate is not None: + return candidate if candidate.__qualname__ == owner_qualname else None + + bound_target = _unwrap_partial(target) + receiver = getattr(bound_target, "__self__", None) + if receiver is None: + return None + receiver_class = receiver if inspect.isclass(receiver) else type(receiver) + return next( + ( + owner + for owner in receiver_class.__mro__ + if owner.__qualname__ == owner_qualname + ), + None, + ) + + +def _partial_bound_hint_names(func: Any) -> FrozenSet[str]: + if not isinstance(func, functools.partial): + return frozenset() + target = func.func + try: + return frozenset( + inspect.signature(target) + .bind_partial(*func.args, **(func.keywords or {})) + .arguments + ) + except Exception: + return frozenset() + + +def _create_declaration_context( + annotation_target: Callable[..., Any], + signature_target: Callable[..., Any], + *, + descriptor_owner: Optional[Type] = None, + implicit_parameter_name: Optional[str] = None, + partial_source: Any = None, +) -> _UDFDeclarationContext: + annotation_target = cast( + Callable[..., Any], _get_callable_inspection_target(annotation_target) + ) + if implicit_parameter_name is None: + bound_target = _unwrap_partial(annotation_target) + bound_function = getattr(bound_target, "__func__", None) + if bound_function is not None: + implicit_parameter_name = _first_parameter_name(bound_function) + + defining_class = _lexical_defining_class( + annotation_target, descriptor_owner + ) + localns = None + if defining_class is not None: + localns = dict(vars(defining_class)) + localns[defining_class.__name__] = defining_class + + try: + invocation_signature = inspect.signature(signature_target) + parameters = tuple(invocation_signature.parameters.values()) + if ( + implicit_parameter_name is not None + and parameters + and parameters[0].name == implicit_parameter_name + ): + invocation_signature = invocation_signature.replace( + parameters=parameters[1:] + ) + except Exception: + invocation_signature = None + + ignored_hint_names = set(_partial_bound_hint_names(partial_source)) + if implicit_parameter_name is not None: + ignored_hint_names.add(implicit_parameter_name) + return _UDFDeclarationContext( + annotation_target=annotation_target, + defining_class=defining_class, + globalns=_get_annotation_globals(annotation_target), + localns=localns, + invocation_signature=invocation_signature, + ignored_hint_names=frozenset(ignored_hint_names), + ) -def _resolve_udf_source( +def _create_resolved_udf( func: _UDFInput, -) -> _ResolvedUDFSource: - """Validate and classify one callable declaration.""" + kind: _UDFSourceKind, + declaration_context: _UDFDeclarationContext, + *, + async_marker: bool = False, +) -> _ResolvedUDF: + target = declaration_context.annotation_target + target_is_async = inspect.iscoroutinefunction(target) + try: + unwrapped_target = inspect.unwrap(target) + except ValueError as exc: + raise TypeError( + "Cannot inspect a UDF with a wrapper cycle." + ) from exc + if not target_is_async and inspect.iscoroutinefunction(unwrapped_target): + raise TypeError( + "A synchronous UDF wrapper cannot wrap an async target; define the " + "wrapper with async def." + ) + is_async = async_marker or target_is_async + return _ResolvedUDF( + _UDFRuntimeSource(func, kind, is_async), declaration_context + ) + + +def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: + """Validate a UDF and resolve its declaration and runtime metadata.""" if isinstance(func, functools.partial) or inspect.isroutine(func): - inspection_target = _get_callable_inspection_target( - cast(Callable[..., Any], func) + declaration_context = _create_declaration_context( + cast(Callable[..., Any], func), + cast(Callable[..., Any], func), + partial_source=func, ) - return _ResolvedUDFSource( - func, - _UDFSourceKind.DIRECT_CALLABLE, - inspect.iscoroutinefunction(inspection_target), - _ignored_hint_names(func, inspection_target), + return _create_resolved_udf( + func, _UDFSourceKind.DIRECT_CALLABLE, declaration_context ) if inspect.isclass(func): @@ -558,41 +677,45 @@ def _resolve_udf_source( raise TypeError(f"func must be a scalar UDF, got {func.__name__}.") if issubclass(func, (ScalarFunction, AsyncScalarFunction)): _validate_zero_argument_class(func) - hint_method, skip_first_parameter = _get_callable_class_hint_method( - func, "eval" + target, descriptor_owner, implicit_parameter_name = ( + _resolve_class_invocation_target(func, "eval") ) - if hint_method is None: + if target is None: raise TypeError( f"Scalar UDF class '{func.__name__}': eval must be defined as a " "method." ) - return _ResolvedUDFSource( + declaration_context = _create_declaration_context( + target, + target, + descriptor_owner=descriptor_owner, + implicit_parameter_name=implicit_parameter_name, + ) + return _create_resolved_udf( func, _UDFSourceKind.SCALAR_FUNCTION_CLASS, - issubclass(func, AsyncScalarFunction) - or inspect.iscoroutinefunction(hint_method), - _ignored_hint_names( - func, hint_method, skip_first=skip_first_parameter - ), + declaration_context, + async_marker=issubclass(func, AsyncScalarFunction), ) - if not _has_custom_call(func): - raise TypeError(f"func must be callable, got {func.__name__}.") - _validate_zero_argument_class(func) - class_hint_method, skip_first_parameter = ( - _get_callable_class_hint_method(func) + target, descriptor_owner, implicit_parameter_name = ( + _resolve_class_invocation_target(func, "__call__") ) - if class_hint_method is None: + if target is None: + if descriptor_owner is None: + raise TypeError(f"func must be callable, got {func.__name__}.") raise TypeError( f"Callable class '{func.__name__}': __call__ must be defined as a method." ) - return _ResolvedUDFSource( - func, - _UDFSourceKind.CALLABLE_CLASS, - inspect.iscoroutinefunction(class_hint_method), - _ignored_hint_names( - func, class_hint_method, skip_first=skip_first_parameter - ), + _validate_zero_argument_class(func) + declaration_context = _create_declaration_context( + target, + target, + descriptor_owner=descriptor_owner, + implicit_parameter_name=implicit_parameter_name, + ) + return _create_resolved_udf( + func, _UDFSourceKind.CALLABLE_CLASS, declaration_context ) if isinstance(func, UserDefinedFunction) and not isinstance( @@ -600,46 +723,38 @@ def _resolve_udf_source( ): raise TypeError(f"func must be a scalar UDF, got {type(func).__name__}.") if isinstance(func, (ScalarFunction, AsyncScalarFunction)): - hint_method = func.eval - is_async = isinstance(func, AsyncScalarFunction) or inspect.iscoroutinefunction( - hint_method + target = func.eval + if not callable(target): + raise TypeError( + f"Scalar UDF instance '{type(func).__name__}': eval must be callable." + ) + declaration_context = _create_declaration_context( + cast(Callable[..., Any], target), + cast(Callable[..., Any], target), + partial_source=target, ) - return _ResolvedUDFSource( + return _create_resolved_udf( func, _UDFSourceKind.SCALAR_FUNCTION_INSTANCE, - is_async, + declaration_context, + async_marker=isinstance(func, AsyncScalarFunction), ) if not callable(func): raise TypeError(f"func must be callable, got {type(func).__name__}.") - hint_method = cast(Callable[..., Any], getattr(func, "__call__")) - return _ResolvedUDFSource( - func, - _UDFSourceKind.CALLABLE_INSTANCE, - inspect.iscoroutinefunction(hint_method), + target = getattr(func, "__call__") + if not callable(target): + raise TypeError( + f"Callable instance '{type(func).__name__}': __call__ must be callable." + ) + declaration_context = _create_declaration_context( + cast(Callable[..., Any], target), + cast(Callable[..., Any], func), + partial_source=target, + ) + return _create_resolved_udf( + func, _UDFSourceKind.CALLABLE_INSTANCE, declaration_context ) - - -def _ignored_hint_names( - func: _UDFInput, - inspection_target: Callable[..., Any], - skip_first: bool = False, -) -> FrozenSet[str]: - try: - parameters = list(inspect.signature(inspection_target).parameters) - except (TypeError, ValueError): - parameters = [] - ignored_names = set(parameters[:1]) if skip_first else set() - if isinstance(func, functools.partial): - try: - ignored_names.update( - inspect.signature(inspection_target) - .bind_partial(*func.args, **(func.keywords or {})) - .arguments - ) - except (TypeError, ValueError): - pass - return frozenset(ignored_names) def _validate_zero_argument_class(func_class: Type) -> None: @@ -662,20 +777,29 @@ def _validate_zero_argument_class(func_class: Type) -> None: def _infer_return_dtype( - func: Callable[..., Any], return_dtype: Optional[_DataTypeLike] + declaration_context: _UDFDeclarationContext, + return_dtype: Optional[_DataTypeLike], ) -> DataType: """Infer the DataFrame return type or validate its explicit declaration.""" if return_dtype is not None: return _convert_to_dtype(return_dtype) - return_hint = _get_callable_return_type_hint(func) + return_hint = _get_callable_return_type_hint(declaration_context) + func_name = _default_udf_name( + declaration_context.annotation_target + ) if return_hint is _UNRESOLVED_TYPE_HINT: - func_name = _default_udf_name(func) raise TypeError( f"Cannot infer return_dtype for '{func_name}': add a return annotation " "or specify return_dtype explicitly." ) - return _data_type_from_type_hint(return_hint) + try: + return _data_type_from_type_hint(return_hint) + except (NameError, AttributeError, SyntaxError, TypeError) as exc: + raise TypeError( + f"Cannot infer return_dtype for '{func_name}': add a return annotation " + "or specify return_dtype explicitly." + ) from exc def _convert_to_dtype(dtype_like: _DataTypeLike) -> DataType: @@ -685,7 +809,12 @@ def _convert_to_dtype(dtype_like: _DataTypeLike) -> DataType: return DataType._from_sql(dtype_like) try: return _data_type_from_type_hint(dtype_like) - except TypeError as exc: + except (NameError, AttributeError, SyntaxError, TypeError) as exc: + if _is_typed_dict(dtype_like): + raise TypeError( + "Cannot resolve return_dtype from the supplied TypedDict; use a " + "concrete DataFrame DataType or SQL type string." + ) from exc raise TypeError( "return_dtype must be a DataFrame DataType, Python type, or SQL " f"type string, got {type(dtype_like).__name__}." @@ -718,22 +847,27 @@ def _data_type_from_type_hint(type_hint: Any) -> DataType: return DataType._from_type_hint(type_hint) -def _detect_func_type(source: _ResolvedUDFSource) -> str: +def _detect_func_type(declaration_context: _UDFDeclarationContext) -> str: """Detect pandas mode from an unbound pandas container annotation.""" - hint_func = source.inspection_target + hint_func = declaration_context.annotation_target try: import pandas as pd except ImportError: return "general" pandas_types = (pd.Series, pd.DataFrame) + pandas_globalns = { + "pandas": pd, + "pd": pd, + **declaration_context.globalns, + } for name in getattr(hint_func, "__annotations__", {}): - if name in source.ignored_hint_names: + if name in declaration_context.ignored_hint_names: continue hint = _resolve_callable_annotation( - hint_func, + declaration_context, name, - fallback_globals={"pandas": pd, "pd": pd}, + globalns=pandas_globalns, ) if hint in pandas_types: return "pandas" @@ -755,17 +889,34 @@ def _get_callable_inspection_target( return cast(Callable[..., Any], target) -def _get_callable_return_type_hint(func: Callable[..., Any]) -> Any: +def _get_annotation_globals(func: Callable[..., Any]) -> Dict[str, Any]: + try: + unwrapped = inspect.unwrap(func) + except ValueError: + return _get_callable_globals(func) + + annotations = getattr(func, "__annotations__", None) + if annotations is not None and annotations is getattr( + unwrapped, "__annotations__", None + ): + return _get_callable_globals(unwrapped) + return _get_callable_globals(func) + + +def _get_callable_return_type_hint( + declaration_context: _UDFDeclarationContext, +) -> Any: # Resolve the return annotation in isolation so an unresolvable parameter # annotation does not prevent return-type inference. - return _resolve_callable_annotation(func, "return") + return _resolve_callable_annotation(declaration_context, "return") def _resolve_callable_annotation( - func: Callable[..., Any], + declaration_context: _UDFDeclarationContext, annotation_name: str, - fallback_globals: Optional[Dict[str, Any]] = None, + globalns: Optional[Dict[str, Any]] = None, ) -> Any: + func = declaration_context.annotation_target annotations = getattr(func, "__annotations__", {}) if annotation_name not in annotations: return _UNRESOLVED_TYPE_HINT @@ -779,12 +930,12 @@ def annotation_holder() -> None: try: return get_type_hints( annotation_holder, - globalns={ - **(fallback_globals or {}), - **_get_callable_globals(func), - }, + globalns=( + declaration_context.globalns if globalns is None else globalns + ), + localns=declaration_context.localns, ).get(annotation_name, _UNRESOLVED_TYPE_HINT) - except (NameError, TypeError): + except (NameError, AttributeError, TypeError): return _UNRESOLVED_TYPE_HINT @@ -798,11 +949,11 @@ def _get_callable_globals(func: Callable[..., Any]) -> Dict[str, Any]: def _resolve_deterministic( - source: _ResolvedUDFSource, deterministic: bool + runtime_source: _UDFRuntimeSource, deterministic: bool ) -> bool: if not isinstance(deterministic, bool): raise TypeError("deterministic must be a bool.") - source.validate_deterministic(deterministic) + runtime_source.validate_deterministic(deterministic) return deterministic @@ -811,8 +962,10 @@ def _validate_deterministic(declared: bool, actual: bool) -> None: raise ValueError(f"Inconsistent deterministic: {declared} and {actual}.") -def _resolve_name(source: _ResolvedUDFSource, name: Optional[str]) -> str: - actual_name = source.default_name if name is None else name +def _resolve_name( + runtime_source: _UDFRuntimeSource, name: Optional[str] +) -> str: + actual_name = runtime_source.default_name if name is None else name if not isinstance(actual_name, str): raise TypeError("name must be a str or None.") if not actual_name: @@ -859,15 +1012,17 @@ class _DataFrameUDFAdapterBase: def __init__( self, - source: _ResolvedUDFSource, + runtime_source: _UDFRuntimeSource, return_dtype: DataType, deterministic: bool, usage: _UDFUsage, func_type: str, ) -> None: - self._source = source + self._runtime_source = runtime_source self._func: Optional[_UDFInput] = ( - None if source.constructs_on_worker else source.source + None + if runtime_source.constructs_on_worker + else runtime_source.callable_source ) self._return_dtype = return_dtype if func_type == "general" else None self._deterministic = deterministic @@ -875,33 +1030,35 @@ def __init__( self._func_type = func_type self._bound_func: Optional[Callable[..., Any]] = None self._lifecycle_opened = False - self.__name__ = source.default_name - self.__doc__ = getattr(source.source, "__doc__", None) + self.__name__ = runtime_source.default_name + self.__doc__ = getattr(runtime_source.callable_source, "__doc__", None) def open(self, function_context: Any) -> None: lifecycle_opened = False try: - if self._source.constructs_on_worker: - self._func = self._source.create_worker_source() + if self._runtime_source.constructs_on_worker: + self._func = self._runtime_source.create_worker_source() func = self._func if func is None: raise RuntimeError("DataFrame UDF source was not initialized.") - self._source.validate_deterministic(self._deterministic, func) - self._source.open_worker_source(func, function_context) - lifecycle_opened = self._source.is_scalar_function - invoke_func = self._source.worker_invocation(func) + self._runtime_source.validate_deterministic( + self._deterministic, func + ) + self._runtime_source.open_worker_source(func, function_context) + lifecycle_opened = self._runtime_source.is_scalar_function + invoke_func = self._runtime_source.worker_invocation(func) self._bound_func = self._bind_func(invoke_func) self._lifecycle_opened = lifecycle_opened except Exception: if lifecycle_opened: try: - self._source.close_worker_source(self._func) + self._runtime_source.close_worker_source(self._func) except Exception: pass self._bound_func = None self._lifecycle_opened = False - if self._source.constructs_on_worker: + if self._runtime_source.constructs_on_worker: self._func = None raise @@ -919,7 +1076,7 @@ def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: return _wrap_scalar_general_result( invoke_func, result_normalizer, - self._source.is_async, + self._runtime_source.is_async, ) return invoke_func @@ -927,11 +1084,11 @@ def close(self) -> None: func = self._func try: if self._lifecycle_opened: - self._source.close_worker_source(func) + self._runtime_source.close_worker_source(func) finally: self._bound_func = None self._lifecycle_opened = False - if self._source.constructs_on_worker: + if self._runtime_source.constructs_on_worker: self._func = None def is_deterministic(self) -> bool: From 1adc1a8783e93b02ee353d7d9491be5f79a155cf Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 1 Sep 2026 12:41:47 +0800 Subject: [PATCH 11/18] [FLINK-40431][python] Harden callable UDF inference Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 134 +++++++++++++++++- flink-python/pyflink/dataframe/udf.py | 101 +++++++++---- 2 files changed, 206 insertions(+), 29 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index bcca450d492cc..2a6df3d4a126a 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -335,17 +335,32 @@ def __call__(cls, value: int, amount: int = 1) -> int: def pandas_identity(values: pd.Series) -> pd.Series: return values + def self_named_pandas_identity( + self: pd.Series, amount: int = 1 + ) -> pd.Series: + return self + amount + class WrappedCallableClass: @functools.wraps(pandas_identity) def __call__(self, *args, **kwargs): return pandas_identity(*args, **kwargs) + class SelfNamedWrappedCallableClass: + @functools.wraps(self_named_pandas_identity) + def __call__(self, *args, **kwargs): + return self_named_pandas_identity(*args, **kwargs) + class WrappedClassMethodCallableClass: @classmethod @functools.wraps(pandas_identity) def __call__(cls, *args, **kwargs): return pandas_identity(*args, **kwargs) + class WrappedScalarFunction(ScalarFunction): + @functools.wraps(pandas_identity) + def eval(self, *args, **kwargs): + return pandas_identity(*args, **kwargs) + class AddFunction(ScalarFunction): def eval(self, *values: int) -> int: return sum(values) @@ -401,7 +416,16 @@ def variadic_add(*values: int) -> int: ) self.assertNotIn("__signature__", vars(exploding_partial_signature)) - for source in (WrappedCallableClass, WrappedClassMethodCallableClass): + wrapped_callable_instance = WrappedCallableClass() + for source in ( + WrappedCallableClass, + wrapped_callable_instance, + wrapped_callable_instance.__call__, + WrappedClassMethodCallableClass, + WrappedClassMethodCallableClass(), + WrappedScalarFunction, + WrappedScalarFunction(), + ): with self.subTest(wrapped_source=source): declaration = pf.udf( source, return_dtype=pf.DataType.int64() @@ -412,6 +436,47 @@ def variadic_add(*values: int) -> int: ) self.assertEqual(declaration._func_type, "pandas") + partial_bound_wrapper = functools.partial( + wrapped_callable_instance.__call__, pd.Series([1]) + ) + partial_declaration = pf.udf( + partial_bound_wrapper, + return_dtype=pf.DataType.int64(), + func_type="general", + ) + self.assertEqual( + inspect.signature(partial_declaration), + inspect.signature(functools.partial(pandas_identity, pd.Series([1]))), + ) + + self_named_instance = SelfNamedWrappedCallableClass() + for source in ( + SelfNamedWrappedCallableClass, + self_named_instance, + self_named_instance.__call__, + ): + with self.subTest(self_named_source=source): + self.assertEqual( + inspect.signature( + pf.udf(source, return_dtype=pf.DataType.int64()) + ), + inspect.signature(self_named_pandas_identity), + ) + self.assertEqual( + inspect.signature( + pf.udf( + functools.partial( + self_named_instance.__call__, pd.Series([1]) + ), + return_dtype=pf.DataType.int64(), + func_type="general", + ) + ), + inspect.signature( + functools.partial(self_named_pandas_identity, pd.Series([1])) + ), + ) + cross_namespace_wrapper = types.FunctionType( _call_module_alias_function.__code__, { @@ -632,6 +697,38 @@ def sync_wrapper(*args, **kwargs): with self.assertRaisesRegex(TypeError, "async def"): pf.udf(sync_wrapper) + def test_unrelated_methodtype_owner_requires_explicit_metadata(self): + class MethodOwner: + Batch = pd.Series + + class Output(TypedDict): + value: int + + def eval(self, values: "Batch") -> "Output": + return {"value": len(values)} + + class ReplacedScalarFunction(ScalarFunction): + def eval(self, value: int) -> int: + return value + + replaced = ReplacedScalarFunction() + replaced.eval = types.MethodType(MethodOwner.eval, replaced) + + with self.assertRaisesRegex( + TypeError, "Cannot infer return_dtype.*return_dtype" + ): + pf.udf(replaced) + + return_dtype = pf.DataType.struct({"value": pf.DataType.int64()}) + inferred_mode = pf.udf(replaced, return_dtype=return_dtype) + self.assertEqual(inferred_mode._func_type, "general") + explicit_mode = pf.udf( + replaced, + return_dtype=return_dtype, + func_type="pandas", + ) + self.assertEqual(explicit_mode._func_type, "pandas") + def test_invalid_class_invocation_descriptors_fail_eagerly(self): class CallableBase: def __call__(self, value: int) -> int: @@ -719,6 +816,41 @@ def invalid_output(value: int) -> InvalidOutput: with self.assertRaisesRegex(TypeError, "DataType or SQL"): pf.udf(lambda value: value, return_dtype=InvalidOutput) + def test_malformed_forward_references_have_clean_inference_behavior(self): + def pandas_after_malformed( + context: Any, values: pd.Series + ) -> int: + return len(values) + + pandas_after_malformed.__annotations__["context"] = "list[" + self.assertEqual( + pf.udf( + pandas_after_malformed, + return_dtype=pf.DataType.int64(), + )._func_type, + "pandas", + ) + + def only_malformed(context: Any) -> int: + return 1 + + only_malformed.__annotations__["context"] = "list[" + self.assertEqual( + pf.udf( + only_malformed, return_dtype=pf.DataType.int64() + )._func_type, + "general", + ) + + def malformed_return(value: int) -> int: + return value + + malformed_return.__annotations__["return"] = "list[" + with self.assertRaisesRegex( + TypeError, "Cannot infer return_dtype.*return_dtype" + ): + pf.udf(malformed_return) + def test_determinism_and_name_metadata(self): class NonDeterministic(ScalarFunction): def eval(self, *values: int) -> int: diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index db6cf3d5cd878..9f0e6420714df 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -549,6 +549,8 @@ def _lexical_defining_class( if not separator: return None if candidate is not None: + # Same-qualified-name method transplantation is indistinguishable because + # Python functions do not retain an exact defining-class identity. return candidate if candidate.__qualname__ == owner_qualname else None bound_target = _unwrap_partial(target) @@ -566,18 +568,72 @@ def _lexical_defining_class( ) -def _partial_bound_hint_names(func: Any) -> FrozenSet[str]: - if not isinstance(func, functools.partial): - return frozenset() - target = func.func +def _apply_partial_to_signature( + signature: inspect.Signature, partial_source: functools.partial +) -> inspect.Signature: + def signature_proxy(*args: Any, **kwargs: Any) -> None: + pass + + # Delegate partial's signature transformation to inspect after supplying the + # normalized invocation signature. + setattr(signature_proxy, "__signature__", signature) + return inspect.signature( + functools.partial( + signature_proxy, + *partial_source.args, + **(partial_source.keywords or {}), + ) + ) + + +def _resolve_invocation_signature( + annotation_target: Callable[..., Any], + signature_target: Callable[..., Any], + implicit_parameter_name: Optional[str], + partial_source: Any, +) -> Tuple[Optional[inspect.Signature], FrozenSet[str]]: + ignored_hint_names = set() + if implicit_parameter_name is not None: + ignored_hint_names.add(implicit_parameter_name) + try: - return frozenset( - inspect.signature(target) - .bind_partial(*func.args, **(func.keywords or {})) - .arguments + signature_inspection_target = signature_target + bound_function = getattr(annotation_target, "__func__", None) + is_wrapped_bound_method = bound_function is not None and hasattr( + bound_function, "__wrapped__" ) + if is_wrapped_bound_method: + signature_inspection_target = bound_function + + invocation_signature = inspect.signature(signature_inspection_target) + parameters = tuple(invocation_signature.parameters.values()) + if ( + implicit_parameter_name is not None + and not hasattr(signature_inspection_target, "__wrapped__") + and parameters + and parameters[0].name == implicit_parameter_name + ): + invocation_signature = invocation_signature.replace( + parameters=parameters[1:] + ) + + if isinstance(partial_source, functools.partial): + partial_target_signature = ( + invocation_signature + if is_wrapped_bound_method + else inspect.signature(partial_source.func) + ) + bound_arguments = partial_target_signature.bind_partial( + *partial_source.args, **(partial_source.keywords or {}) + ) + ignored_hint_names.update(bound_arguments.arguments) + if is_wrapped_bound_method: + invocation_signature = _apply_partial_to_signature( + invocation_signature, partial_source + ) + return invocation_signature, frozenset(ignored_hint_names) except Exception: - return frozenset() + return None, frozenset(ignored_hint_names) def _create_declaration_context( @@ -605,30 +661,19 @@ def _create_declaration_context( localns = dict(vars(defining_class)) localns[defining_class.__name__] = defining_class - try: - invocation_signature = inspect.signature(signature_target) - parameters = tuple(invocation_signature.parameters.values()) - if ( - implicit_parameter_name is not None - and parameters - and parameters[0].name == implicit_parameter_name - ): - invocation_signature = invocation_signature.replace( - parameters=parameters[1:] - ) - except Exception: - invocation_signature = None - - ignored_hint_names = set(_partial_bound_hint_names(partial_source)) - if implicit_parameter_name is not None: - ignored_hint_names.add(implicit_parameter_name) + invocation_signature, ignored_hint_names = _resolve_invocation_signature( + annotation_target, + signature_target, + implicit_parameter_name, + partial_source, + ) return _UDFDeclarationContext( annotation_target=annotation_target, defining_class=defining_class, globalns=_get_annotation_globals(annotation_target), localns=localns, invocation_signature=invocation_signature, - ignored_hint_names=frozenset(ignored_hint_names), + ignored_hint_names=ignored_hint_names, ) @@ -935,7 +980,7 @@ def annotation_holder() -> None: ), localns=declaration_context.localns, ).get(annotation_name, _UNRESOLVED_TYPE_HINT) - except (NameError, AttributeError, TypeError): + except (NameError, AttributeError, SyntaxError, TypeError): return _UNRESOLVED_TYPE_HINT From aa356baf4a07936ec5237887034123e81edec107 Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 1 Sep 2026 13:25:56 +0800 Subject: [PATCH 12/18] [FLINK-40431][python] Clarify DataFrame UDF worker lifecycle Generated-by: OpenAI Codex (GPT-5) --- flink-python/pyflink/dataframe/udf.py | 385 +++++++++++++------------- 1 file changed, 195 insertions(+), 190 deletions(-) diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 9f0e6420714df..14911ee49c2d5 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -21,7 +21,7 @@ import functools import inspect from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum from typing import ( Any, @@ -56,6 +56,7 @@ __all__ = ["udf"] _UDFInput = Union[Callable[..., Any], ScalarFunction, AsyncScalarFunction, Type] +_ActiveUDFSource = Union[Callable[..., Any], ScalarFunction, AsyncScalarFunction] _DataTypeLike = Union[DataType, Type, str] _UNRESOLVED_TYPE_HINT = object() @@ -75,6 +76,13 @@ class _UDFSourceKind(Enum): SCALAR_FUNCTION_INSTANCE = "scalar_function_instance" SCALAR_FUNCTION_CLASS = "scalar_function_class" + @property + def is_scalar_function(self) -> bool: + return self in ( + _UDFSourceKind.SCALAR_FUNCTION_INSTANCE, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ) + @dataclass(frozen=True) class _UDFDeclarationContext: @@ -90,7 +98,7 @@ class _UDFDeclarationContext: @dataclass(frozen=True) class _UDFRuntimeSource: - """Worker-facing metadata used to initialize and invoke a UDF.""" + """Worker-facing recipe used to initialize a UDF.""" callable_source: _UDFInput kind: _UDFSourceKind @@ -100,13 +108,6 @@ class _UDFRuntimeSource: def default_name(self) -> str: return _default_udf_name(self.callable_source) - @property - def is_scalar_function(self) -> bool: - return self.kind in ( - _UDFSourceKind.SCALAR_FUNCTION_INSTANCE, - _UDFSourceKind.SCALAR_FUNCTION_CLASS, - ) - @property def constructs_on_worker(self) -> bool: return self.kind in ( @@ -114,74 +115,84 @@ def constructs_on_worker(self) -> bool: _UDFSourceKind.SCALAR_FUNCTION_CLASS, ) - def create_worker_source(self) -> _UDFInput: - if not self.constructs_on_worker: - return self.callable_source - source_class = cast(Type, self.callable_source) - source = source_class() - if self.is_scalar_function: - if not isinstance(source, (ScalarFunction, AsyncScalarFunction)): + def validate_declared_determinism(self, declared: bool) -> None: + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: + _validate_deterministic( + declared, + cast( + Union[ScalarFunction, AsyncScalarFunction], self.callable_source + ).is_deterministic(), + ) + + def create_worker_udf(self) -> "_WorkerUDF": + source = self.callable_source + if self.constructs_on_worker: + source_class = cast(Type, source) + source = source_class() + if self.kind.is_scalar_function: + if not isinstance(source, (ScalarFunction, AsyncScalarFunction)): + raise TypeError( + f"Scalar UDF class '{source_class.__name__}' constructed an " + f"unsupported object of type '{type(source).__name__}'." + ) + elif not callable(source): raise TypeError( - f"Scalar UDF class '{source_class.__name__}' constructed an " - f"unsupported object of type '{type(source).__name__}'." + f"Callable class '{source_class.__name__}' constructed a non-callable " + f"object of type '{type(source).__name__}'." ) - elif not callable(source): - raise TypeError( - f"Callable class '{source_class.__name__}' constructed a non-callable " - f"object of type '{type(source).__name__}'." - ) - return cast(_UDFInput, source) + return _WorkerUDF(cast(_ActiveUDFSource, source), self.kind) - def validate_deterministic( - self, declared: bool, worker_source: Optional[_UDFInput] = None - ) -> None: - source: Optional[_UDFInput] - if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: - source = self.callable_source - elif self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS: - source = worker_source - else: - source = None - if source is not None: + +@dataclass(frozen=True) +class _ResolvedUDF: + """A resolved declaration split into client and worker metadata.""" + + runtime_source: _UDFRuntimeSource + declaration_context: _UDFDeclarationContext + + +@dataclass +class _WorkerUDF: + """An initialized UDF owned by one worker adapter lifecycle.""" + + active_source: _ActiveUDFSource + kind: _UDFSourceKind + _lifecycle_opened: bool = field(default=False, init=False, repr=False) + + def validate_deterministic(self, declared: bool) -> None: + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS: _validate_deterministic( declared, cast( - Union[ScalarFunction, AsyncScalarFunction], source + Union[ScalarFunction, AsyncScalarFunction], self.active_source ).is_deterministic(), ) - def open_worker_source( - self, worker_source: _UDFInput, function_context: Any - ) -> None: - if self.is_scalar_function: + def open(self, function_context: Any) -> None: + if self.kind.is_scalar_function: cast( - Union[ScalarFunction, AsyncScalarFunction], worker_source + Union[ScalarFunction, AsyncScalarFunction], self.active_source ).open(function_context) + self._lifecycle_opened = True - def worker_invocation( - self, worker_source: _UDFInput - ) -> Callable[..., Any]: + @property + def invocation(self) -> Callable[..., Any]: if self.kind is _UDFSourceKind.DIRECT_CALLABLE: - return cast(Callable[..., Any], worker_source) - if self.is_scalar_function: + return cast(Callable[..., Any], self.active_source) + if self.kind.is_scalar_function: return cast( - Union[ScalarFunction, AsyncScalarFunction], worker_source + Union[ScalarFunction, AsyncScalarFunction], self.active_source ).eval - return cast(Callable[..., Any], getattr(worker_source, "__call__")) - - def close_worker_source(self, worker_source: Optional[_UDFInput]) -> None: - if self.is_scalar_function and worker_source is not None: - cast( - Union[ScalarFunction, AsyncScalarFunction], worker_source - ).close() + return cast(Callable[..., Any], getattr(self.active_source, "__call__")) - -@dataclass(frozen=True) -class _ResolvedUDF: - """A resolved declaration split into client and worker metadata.""" - - runtime_source: _UDFRuntimeSource - declaration_context: _UDFDeclarationContext + def close(self) -> None: + try: + if self._lifecycle_opened: + cast( + Union[ScalarFunction, AsyncScalarFunction], self.active_source + ).close() + finally: + self._lifecycle_opened = False class _DataFrameUDFWrapper: @@ -488,6 +499,23 @@ def _validate_scalar_udf_options( # ======================== Callable Inspection and Resolution ======================== +# ---- Invocation target resolution ---- + + +def _unwrap_partial(func: Any) -> Any: + while isinstance(func, functools.partial): + func = func.func + return func + + +def _get_callable_inspection_target( + func: Callable[..., Any], +) -> Callable[..., Any]: + target = _unwrap_partial(func) + if callable(target) and not inspect.isroutine(target) and not inspect.isclass(target): + return cast(Callable[..., Any], getattr(target, "__call__")) + return cast(Callable[..., Any], target) + def _first_parameter_name(func: Callable[..., Any]) -> Optional[str]: try: @@ -532,6 +560,28 @@ def _resolve_class_invocation_target( return cast(Callable[..., Any], target), descriptor_owner, implicit_parameter_name +def _validate_zero_argument_class(func_class: Type) -> None: + if inspect.isabstract(func_class): + raise TypeError(f"UDF class '{func_class.__name__}' must not be abstract.") + try: + constructor_signature = inspect.signature(func_class) + except (TypeError, ValueError) as exc: + raise TypeError( + f"Cannot verify that UDF class '{func_class.__name__}' has a zero-argument " + "constructor; pass a configured instance instead." + ) from exc + try: + constructor_signature.bind() + except TypeError as exc: + raise TypeError( + f"UDF class '{func_class.__name__}' must have a zero-argument constructor; " + "pass a configured instance instead." + ) from exc + + +# ---- Annotation namespace resolution ---- + + def _function_qualname(func: Callable[..., Any]) -> Optional[str]: target = _unwrap_partial(func) target = getattr(target, "__func__", target) @@ -568,6 +618,32 @@ def _lexical_defining_class( ) +def _get_callable_globals(func: Callable[..., Any]) -> Dict[str, Any]: + func_globals = getattr(func, "__globals__", None) + if func_globals is None: + func_globals = getattr( + getattr(func, "__func__", None), "__globals__", {} + ) + return cast(Dict[str, Any], func_globals) + + +def _get_annotation_globals(func: Callable[..., Any]) -> Dict[str, Any]: + try: + unwrapped = inspect.unwrap(func) + except ValueError: + return _get_callable_globals(func) + + annotations = getattr(func, "__annotations__", None) + if annotations is not None and annotations is getattr( + unwrapped, "__annotations__", None + ): + return _get_callable_globals(unwrapped) + return _get_callable_globals(func) + + +# ---- Signature and declaration context assembly ---- + + def _apply_partial_to_signature( signature: inspect.Signature, partial_source: functools.partial ) -> inspect.Signature: @@ -802,23 +878,43 @@ def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: ) -def _validate_zero_argument_class(func_class: Type) -> None: - if inspect.isabstract(func_class): - raise TypeError(f"UDF class '{func_class.__name__}' must not be abstract.") - try: - constructor_signature = inspect.signature(func_class) - except (TypeError, ValueError) as exc: - raise TypeError( - f"Cannot verify that UDF class '{func_class.__name__}' has a zero-argument " - "constructor; pass a configured instance instead." - ) from exc +# ---- Annotation and type inference ---- + + +def _resolve_callable_annotation( + declaration_context: _UDFDeclarationContext, + annotation_name: str, + globalns: Optional[Dict[str, Any]] = None, +) -> Any: + func = declaration_context.annotation_target + annotations = getattr(func, "__annotations__", {}) + if annotation_name not in annotations: + return _UNRESOLVED_TYPE_HINT + + def annotation_holder() -> None: + pass + + annotation_holder.__annotations__ = { + annotation_name: annotations[annotation_name] + } try: - constructor_signature.bind() - except TypeError as exc: - raise TypeError( - f"UDF class '{func_class.__name__}' must have a zero-argument constructor; " - "pass a configured instance instead." - ) from exc + return get_type_hints( + annotation_holder, + globalns=( + declaration_context.globalns if globalns is None else globalns + ), + localns=declaration_context.localns, + ).get(annotation_name, _UNRESOLVED_TYPE_HINT) + except (NameError, AttributeError, SyntaxError, TypeError): + return _UNRESOLVED_TYPE_HINT + + +def _get_callable_return_type_hint( + declaration_context: _UDFDeclarationContext, +) -> Any: + # Resolve the return annotation in isolation so an unresolvable parameter + # annotation does not prevent return-type inference. + return _resolve_callable_annotation(declaration_context, "return") def _infer_return_dtype( @@ -919,78 +1015,7 @@ def _detect_func_type(declaration_context: _UDFDeclarationContext) -> str: return "general" -def _unwrap_partial(func: Any) -> Any: - while isinstance(func, functools.partial): - func = func.func - return func - - -def _get_callable_inspection_target( - func: Callable[..., Any], -) -> Callable[..., Any]: - target = _unwrap_partial(func) - if callable(target) and not inspect.isroutine(target) and not inspect.isclass(target): - return cast(Callable[..., Any], getattr(target, "__call__")) - return cast(Callable[..., Any], target) - - -def _get_annotation_globals(func: Callable[..., Any]) -> Dict[str, Any]: - try: - unwrapped = inspect.unwrap(func) - except ValueError: - return _get_callable_globals(func) - - annotations = getattr(func, "__annotations__", None) - if annotations is not None and annotations is getattr( - unwrapped, "__annotations__", None - ): - return _get_callable_globals(unwrapped) - return _get_callable_globals(func) - - -def _get_callable_return_type_hint( - declaration_context: _UDFDeclarationContext, -) -> Any: - # Resolve the return annotation in isolation so an unresolvable parameter - # annotation does not prevent return-type inference. - return _resolve_callable_annotation(declaration_context, "return") - - -def _resolve_callable_annotation( - declaration_context: _UDFDeclarationContext, - annotation_name: str, - globalns: Optional[Dict[str, Any]] = None, -) -> Any: - func = declaration_context.annotation_target - annotations = getattr(func, "__annotations__", {}) - if annotation_name not in annotations: - return _UNRESOLVED_TYPE_HINT - - def annotation_holder() -> None: - pass - - annotation_holder.__annotations__ = { - annotation_name: annotations[annotation_name] - } - try: - return get_type_hints( - annotation_holder, - globalns=( - declaration_context.globalns if globalns is None else globalns - ), - localns=declaration_context.localns, - ).get(annotation_name, _UNRESOLVED_TYPE_HINT) - except (NameError, AttributeError, SyntaxError, TypeError): - return _UNRESOLVED_TYPE_HINT - - -def _get_callable_globals(func: Callable[..., Any]) -> Dict[str, Any]: - func_globals = getattr(func, "__globals__", None) - if func_globals is None: - func_globals = getattr( - getattr(func, "__func__", None), "__globals__", {} - ) - return cast(Dict[str, Any], func_globals) +# ---- Declaration options ---- def _resolve_deterministic( @@ -998,7 +1023,7 @@ def _resolve_deterministic( ) -> bool: if not isinstance(deterministic, bool): raise TypeError("deterministic must be a bool.") - runtime_source.validate_deterministic(deterministic) + runtime_source.validate_declared_determinism(deterministic) return deterministic @@ -1064,47 +1089,29 @@ def __init__( func_type: str, ) -> None: self._runtime_source = runtime_source - self._func: Optional[_UDFInput] = ( - None - if runtime_source.constructs_on_worker - else runtime_source.callable_source - ) + self._worker_udf: Optional[_WorkerUDF] = None self._return_dtype = return_dtype if func_type == "general" else None self._deterministic = deterministic self._usage = usage self._func_type = func_type - self._bound_func: Optional[Callable[..., Any]] = None - self._lifecycle_opened = False + self._bound_invocation: Optional[Callable[..., Any]] = None self.__name__ = runtime_source.default_name self.__doc__ = getattr(runtime_source.callable_source, "__doc__", None) def open(self, function_context: Any) -> None: - lifecycle_opened = False + worker_udf = self._runtime_source.create_worker_udf() try: - if self._runtime_source.constructs_on_worker: - self._func = self._runtime_source.create_worker_source() - - func = self._func - if func is None: - raise RuntimeError("DataFrame UDF source was not initialized.") - self._runtime_source.validate_deterministic( - self._deterministic, func - ) - self._runtime_source.open_worker_source(func, function_context) - lifecycle_opened = self._runtime_source.is_scalar_function - invoke_func = self._runtime_source.worker_invocation(func) - self._bound_func = self._bind_func(invoke_func) - self._lifecycle_opened = lifecycle_opened + worker_udf.validate_deterministic(self._deterministic) + worker_udf.open(function_context) + self._bound_invocation = self._bind_func(worker_udf.invocation) + self._worker_udf = worker_udf except Exception: - if lifecycle_opened: - try: - self._runtime_source.close_worker_source(self._func) - except Exception: - pass - self._bound_func = None - self._lifecycle_opened = False - if self._runtime_source.constructs_on_worker: - self._func = None + try: + worker_udf.close() + except Exception: + pass + self._bound_invocation = None + self._worker_udf = None raise def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: @@ -1126,23 +1133,21 @@ def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: return invoke_func def close(self) -> None: - func = self._func + worker_udf = self._worker_udf try: - if self._lifecycle_opened: - self._runtime_source.close_worker_source(func) + if worker_udf is not None: + worker_udf.close() finally: - self._bound_func = None - self._lifecycle_opened = False - if self._runtime_source.constructs_on_worker: - self._func = None + self._bound_invocation = None + self._worker_udf = None def is_deterministic(self) -> bool: return self._deterministic def _invocation(self) -> Callable[..., Any]: - if self._bound_func is None: + if self._bound_invocation is None: raise RuntimeError("DataFrame UDF was invoked before open().") - return self._bound_func + return self._bound_invocation class _DataFrameScalarFunctionAdapter(_DataFrameUDFAdapterBase, ScalarFunction): From 9c575b5711de89ab7fd64a03fcedb6d7d1c805fc Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 1 Sep 2026 13:56:46 +0800 Subject: [PATCH 13/18] [FLINK-40431][python] Validate partial UDF declarations Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 24 +++++++++ flink-python/pyflink/dataframe/udf.py | 52 ++++++++++++------- 2 files changed, 56 insertions(+), 20 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 2a6df3d4a126a..9de39770a2dae 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -401,6 +401,10 @@ def variadic_add(*values: int) -> int: self.assertEqual( inspect.signature(pf.udf(partial_add)), inspect.signature(partial_add) ) + with self.assertRaisesRegex( + TypeError, "Invalid functools.partial UDF 'add'.*unexpected keyword" + ): + pf.udf(functools.partial(add, missing=1)) uninspectable = pf.udf(operator.itemgetter(0), return_dtype=int) self.assertEqual(_return_dtype(uninspectable), pf.DataType.int64()) @@ -1101,6 +1105,14 @@ class NonScalarFunction(TableFunction): def eval(self, value): return value + class MissingCallableReturn: + def __call__(self, value): + return value + + class MissingScalarReturn(ScalarFunction): + def eval(self, value): + return value + invalid_declarations = [ ( "not callable", @@ -1135,6 +1147,18 @@ def eval(self, value): TypeError, "Cannot infer return_dtype", ), + ( + "callable class missing return", + lambda: pf.udf(MissingCallableReturn), + TypeError, + "Cannot infer return_dtype for 'MissingCallableReturn'", + ), + ( + "scalar function class missing return", + lambda: pf.udf(MissingScalarReturn), + TypeError, + "Cannot infer return_dtype for 'MissingScalarReturn'", + ), ( "Table return type", lambda: pf.udf( diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 14911ee49c2d5..6b73d51d6f6c8 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -454,7 +454,7 @@ def decorator(f: _UDFInput) -> Callable[..., Expression]: actual_func_type, return_dtype, runtime_source.is_async ) actual_return_dtype = _infer_return_dtype( - declaration_context, return_dtype + declaration_context, return_dtype, runtime_source.default_name ) actual_deterministic = _resolve_deterministic( runtime_source, deterministic @@ -672,15 +672,15 @@ def _resolve_invocation_signature( if implicit_parameter_name is not None: ignored_hint_names.add(implicit_parameter_name) - try: - signature_inspection_target = signature_target - bound_function = getattr(annotation_target, "__func__", None) - is_wrapped_bound_method = bound_function is not None and hasattr( - bound_function, "__wrapped__" - ) - if is_wrapped_bound_method: - signature_inspection_target = bound_function + signature_inspection_target = signature_target + bound_function = getattr(annotation_target, "__func__", None) + is_wrapped_bound_method = bound_function is not None and hasattr( + bound_function, "__wrapped__" + ) + if is_wrapped_bound_method: + signature_inspection_target = bound_function + try: invocation_signature = inspect.signature(signature_inspection_target) parameters = tuple(invocation_signature.parameters.values()) if ( @@ -692,24 +692,38 @@ def _resolve_invocation_signature( invocation_signature = invocation_signature.replace( parameters=parameters[1:] ) + except Exception: + invocation_signature = None - if isinstance(partial_source, functools.partial): + if isinstance(partial_source, functools.partial): + try: partial_target_signature = ( invocation_signature if is_wrapped_bound_method else inspect.signature(partial_source.func) ) + except Exception: + return None, frozenset(ignored_hint_names) + if partial_target_signature is None: + return None, frozenset(ignored_hint_names) + try: bound_arguments = partial_target_signature.bind_partial( *partial_source.args, **(partial_source.keywords or {}) ) - ignored_hint_names.update(bound_arguments.arguments) - if is_wrapped_bound_method: + except TypeError as exc: + raise TypeError( + f"Invalid functools.partial UDF " + f"'{_default_udf_name(partial_source)}': {exc}." + ) from exc + ignored_hint_names.update(bound_arguments.arguments) + if is_wrapped_bound_method: + try: invocation_signature = _apply_partial_to_signature( invocation_signature, partial_source ) - return invocation_signature, frozenset(ignored_hint_names) - except Exception: - return None, frozenset(ignored_hint_names) + except Exception: + return None, frozenset(ignored_hint_names) + return invocation_signature, frozenset(ignored_hint_names) def _create_declaration_context( @@ -920,25 +934,23 @@ def _get_callable_return_type_hint( def _infer_return_dtype( declaration_context: _UDFDeclarationContext, return_dtype: Optional[_DataTypeLike], + udf_name: str, ) -> DataType: """Infer the DataFrame return type or validate its explicit declaration.""" if return_dtype is not None: return _convert_to_dtype(return_dtype) return_hint = _get_callable_return_type_hint(declaration_context) - func_name = _default_udf_name( - declaration_context.annotation_target - ) if return_hint is _UNRESOLVED_TYPE_HINT: raise TypeError( - f"Cannot infer return_dtype for '{func_name}': add a return annotation " + f"Cannot infer return_dtype for '{udf_name}': add a return annotation " "or specify return_dtype explicitly." ) try: return _data_type_from_type_hint(return_hint) except (NameError, AttributeError, SyntaxError, TypeError) as exc: raise TypeError( - f"Cannot infer return_dtype for '{func_name}': add a return annotation " + f"Cannot infer return_dtype for '{udf_name}': add a return annotation " "or specify return_dtype explicitly." ) from exc From 47098f9566b65d0e50f6773920edca472910da50 Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 1 Sep 2026 15:46:41 +0800 Subject: [PATCH 14/18] [FLINK-40431][python] Simplify UDF declaration helpers Generated-by: OpenAI Codex (GPT-5) --- flink-python/pyflink/dataframe/udf.py | 80 ++++++++++----------------- 1 file changed, 29 insertions(+), 51 deletions(-) diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 6b73d51d6f6c8..b7bbc9e99f187 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -117,12 +117,13 @@ def constructs_on_worker(self) -> bool: def validate_declared_determinism(self, declared: bool) -> None: if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: - _validate_deterministic( - declared, - cast( - Union[ScalarFunction, AsyncScalarFunction], self.callable_source - ).is_deterministic(), - ) + actual = cast( + Union[ScalarFunction, AsyncScalarFunction], self.callable_source + ).is_deterministic() + if declared != actual: + raise ValueError( + f"Inconsistent deterministic: {declared} and {actual}." + ) def create_worker_udf(self) -> "_WorkerUDF": source = self.callable_source @@ -161,12 +162,13 @@ class _WorkerUDF: def validate_deterministic(self, declared: bool) -> None: if self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS: - _validate_deterministic( - declared, - cast( - Union[ScalarFunction, AsyncScalarFunction], self.active_source - ).is_deterministic(), - ) + actual = cast( + Union[ScalarFunction, AsyncScalarFunction], self.active_source + ).is_deterministic() + if declared != actual: + raise ValueError( + f"Inconsistent deterministic: {declared} and {actual}." + ) def open(self, function_context: Any) -> None: if self.kind.is_scalar_function: @@ -456,15 +458,19 @@ def decorator(f: _UDFInput) -> Callable[..., Expression]: actual_return_dtype = _infer_return_dtype( declaration_context, return_dtype, runtime_source.default_name ) - actual_deterministic = _resolve_deterministic( - runtime_source, deterministic - ) - actual_name = _resolve_name(runtime_source, name) + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool.") + runtime_source.validate_declared_determinism(deterministic) + actual_name = runtime_source.default_name if name is None else name + if not isinstance(actual_name, str): + raise TypeError("name must be a str or None.") + if not actual_name: + raise ValueError("name must not be empty.") return _DataFrameUDFWrapper( runtime_source, actual_return_dtype, - actual_deterministic, + deterministic, actual_name, actual_func_type, declaration_context.invocation_signature, @@ -508,6 +514,12 @@ def _unwrap_partial(func: Any) -> Any: return func +def _default_udf_name(func: _UDFInput) -> str: + target = _unwrap_partial(func) + name = getattr(target, "__name__", None) + return name if isinstance(name, str) else type(target).__name__ + + def _get_callable_inspection_target( func: Callable[..., Any], ) -> Callable[..., Any]: @@ -1027,40 +1039,6 @@ def _detect_func_type(declaration_context: _UDFDeclarationContext) -> str: return "general" -# ---- Declaration options ---- - - -def _resolve_deterministic( - runtime_source: _UDFRuntimeSource, deterministic: bool -) -> bool: - if not isinstance(deterministic, bool): - raise TypeError("deterministic must be a bool.") - runtime_source.validate_declared_determinism(deterministic) - return deterministic - - -def _validate_deterministic(declared: bool, actual: bool) -> None: - if declared != actual: - raise ValueError(f"Inconsistent deterministic: {declared} and {actual}.") - - -def _resolve_name( - runtime_source: _UDFRuntimeSource, name: Optional[str] -) -> str: - actual_name = runtime_source.default_name if name is None else name - if not isinstance(actual_name, str): - raise TypeError("name must be a str or None.") - if not actual_name: - raise ValueError("name must not be empty.") - return actual_name - - -def _default_udf_name(func: _UDFInput) -> str: - target = _unwrap_partial(func) - name = getattr(target, "__name__", None) - return name if isinstance(name, str) else type(target).__name__ - - # ======================== Worker Adapters ======================== From ed503a5a56dc9ad5efdba9564a56eb3b719a820c Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 1 Sep 2026 15:52:27 +0800 Subject: [PATCH 15/18] [FLINK-40431][python] Fix callable signature type narrowing Generated-by: OpenAI Codex (GPT-5) --- flink-python/pyflink/dataframe/udf.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index b7bbc9e99f187..dc1eca7490587 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -684,13 +684,13 @@ def _resolve_invocation_signature( if implicit_parameter_name is not None: ignored_hint_names.add(implicit_parameter_name) - signature_inspection_target = signature_target + signature_inspection_target: Callable[..., Any] = signature_target bound_function = getattr(annotation_target, "__func__", None) is_wrapped_bound_method = bound_function is not None and hasattr( bound_function, "__wrapped__" ) if is_wrapped_bound_method: - signature_inspection_target = bound_function + signature_inspection_target = cast(Callable[..., Any], bound_function) try: invocation_signature = inspect.signature(signature_inspection_target) @@ -729,6 +729,8 @@ def _resolve_invocation_signature( ) from exc ignored_hint_names.update(bound_arguments.arguments) if is_wrapped_bound_method: + if invocation_signature is None: + return None, frozenset(ignored_hint_names) try: invocation_signature = _apply_partial_to_signature( invocation_signature, partial_source From 58e32614ac15bae20ed49b47fa98305e540c15f5 Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 1 Sep 2026 16:38:35 +0800 Subject: [PATCH 16/18] [FLINK-40431][python] Harden wrapped and async UDF declarations Validate the AsyncScalarFunction eval protocol, preserve method binding through functools.wraps, and share determinism agreement validation. Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 65 ++++++++++++++++- flink-python/pyflink/dataframe/udf.py | 72 ++++++++++++++----- 2 files changed, 118 insertions(+), 19 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 9de39770a2dae..a929614940dcc 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -305,7 +305,8 @@ class ReceivingCallable: class OverriddenScalarFunction(ScalarFunction): _UDF_TEST_ALIAS = str - def eval(self, value: int) -> str: + def eval(self, *values: int) -> str: + value, = values return str(value) overridden = OverriddenScalarFunction() @@ -494,6 +495,48 @@ def variadic_add(*values: int) -> int: pf.DataType.int64(), ) + def test_wrapped_methods_preserve_bound_signatures(self): + def method_decorator(method): + @functools.wraps(method) + def wrapper(*args, **kwargs): + return method(*args, **kwargs) + + return wrapper + + def add(value: int, amount: int = 1) -> int: + return value + amount + + class WrappedCallable: + @method_decorator + def __call__(self, value: int, amount: int = 1) -> int: + return value + amount + + class WrappedClassMethodCallable: + @classmethod + @method_decorator + def __call__(cls, value: int, amount: int = 1) -> int: + return value + amount + + expected_signature = inspect.signature(add) + instance = WrappedCallable() + for source in ( + WrappedCallable, + instance, + instance.__call__, + WrappedClassMethodCallable, + WrappedClassMethodCallable(), + ): + with self.subTest(source=source): + self.assertEqual( + inspect.signature(pf.udf(source)), expected_signature + ) + + partial_source = functools.partial(instance.__call__, 1) + self.assertEqual( + inspect.signature(pf.udf(partial_source)), + inspect.signature(functools.partial(add, 1)), + ) + def test_func_type_resolution_and_async_detection(self): def pandas_add_one(values: pd.Series) -> pd.Series: return values + 1 @@ -701,6 +744,20 @@ def sync_wrapper(*args, **kwargs): with self.assertRaisesRegex(TypeError, "async def"): pf.udf(sync_wrapper) + def test_sync_async_scalar_eval_is_rejected(self): + class SyncAsyncScalarFunction(AsyncScalarFunction): + def eval(self, *values: int) -> int: + value, = values + return value + 1 + + for source in (SyncAsyncScalarFunction, SyncAsyncScalarFunction()): + with self.subTest(source=source): + with self.assertRaisesRegex( + TypeError, + "AsyncScalarFunction 'SyncAsyncScalarFunction'.*async def", + ): + pf.udf(source, return_dtype=pf.DataType.int64()) + def test_unrelated_methodtype_owner_requires_explicit_metadata(self): class MethodOwner: Batch = pd.Series @@ -712,7 +769,8 @@ def eval(self, values: "Batch") -> "Output": return {"value": len(values)} class ReplacedScalarFunction(ScalarFunction): - def eval(self, value: int) -> int: + def eval(self, *values: int) -> int: + value, = values return value replaced = ReplacedScalarFunction() @@ -742,7 +800,8 @@ class HiddenCallable(CallableBase): __call__ = None class ScalarBase(ScalarFunction): - def eval(self, value: int) -> int: + def eval(self, *values: int) -> int: + value, = values return value class HiddenScalarFunction(ScalarBase): diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index dc1eca7490587..0a8fa448b15b3 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -120,10 +120,7 @@ def validate_declared_determinism(self, declared: bool) -> None: actual = cast( Union[ScalarFunction, AsyncScalarFunction], self.callable_source ).is_deterministic() - if declared != actual: - raise ValueError( - f"Inconsistent deterministic: {declared} and {actual}." - ) + _validate_determinism_agreement(declared, actual) def create_worker_udf(self) -> "_WorkerUDF": source = self.callable_source @@ -165,10 +162,7 @@ def validate_deterministic(self, declared: bool) -> None: actual = cast( Union[ScalarFunction, AsyncScalarFunction], self.active_source ).is_deterministic() - if declared != actual: - raise ValueError( - f"Inconsistent deterministic: {declared} and {actual}." - ) + _validate_determinism_agreement(declared, actual) def open(self, function_context: Any) -> None: if self.kind.is_scalar_function: @@ -482,6 +476,13 @@ def decorator(f: _UDFInput) -> Callable[..., Expression]: # ======================== Declaration Validation ======================== +def _validate_determinism_agreement(declared: bool, actual: bool) -> None: + if declared != actual: + raise ValueError( + f"Inconsistent deterministic: {declared} and {actual}." + ) + + def _validate_scalar_udf_options( func_type: str, return_dtype: Optional[_DataTypeLike], @@ -674,11 +675,28 @@ def signature_proxy(*args: Any, **kwargs: Any) -> None: ) +def _preserves_method_binding( + target: Callable[..., Any], defining_class: Optional[Type] +) -> bool: + target = _unwrap_partial(target) + target = getattr(target, "__func__", target) + if defining_class is None or not hasattr(target, "__wrapped__"): + return False + try: + unwrapped_target = inspect.unwrap(target) + except ValueError: + return False + return _lexical_defining_class( + cast(Callable[..., Any], unwrapped_target), defining_class + ) is defining_class + + def _resolve_invocation_signature( annotation_target: Callable[..., Any], signature_target: Callable[..., Any], implicit_parameter_name: Optional[str], partial_source: Any, + preserves_method_binding: bool, ) -> Tuple[Optional[inspect.Signature], FrozenSet[str]]: ignored_hint_names = set() if implicit_parameter_name is not None: @@ -689,7 +707,10 @@ def _resolve_invocation_signature( is_wrapped_bound_method = bound_function is not None and hasattr( bound_function, "__wrapped__" ) - if is_wrapped_bound_method: + uses_unbound_wrapped_signature = ( + is_wrapped_bound_method and not preserves_method_binding + ) + if uses_unbound_wrapped_signature: signature_inspection_target = cast(Callable[..., Any], bound_function) try: @@ -697,7 +718,13 @@ def _resolve_invocation_signature( parameters = tuple(invocation_signature.parameters.values()) if ( implicit_parameter_name is not None - and not hasattr(signature_inspection_target, "__wrapped__") + and ( + not hasattr( + getattr(annotation_target, "__func__", annotation_target), + "__wrapped__", + ) + or preserves_method_binding + ) and parameters and parameters[0].name == implicit_parameter_name ): @@ -711,7 +738,7 @@ def _resolve_invocation_signature( try: partial_target_signature = ( invocation_signature - if is_wrapped_bound_method + if uses_unbound_wrapped_signature else inspect.signature(partial_source.func) ) except Exception: @@ -728,7 +755,7 @@ def _resolve_invocation_signature( f"'{_default_udf_name(partial_source)}': {exc}." ) from exc ignored_hint_names.update(bound_arguments.arguments) - if is_wrapped_bound_method: + if uses_unbound_wrapped_signature: if invocation_signature is None: return None, frozenset(ignored_hint_names) try: @@ -751,15 +778,22 @@ def _create_declaration_context( annotation_target = cast( Callable[..., Any], _get_callable_inspection_target(annotation_target) ) + defining_class = _lexical_defining_class( + annotation_target, descriptor_owner + ) + preserves_method_binding = _preserves_method_binding( + annotation_target, defining_class + ) + if preserves_method_binding: + implicit_parameter_name = _first_parameter_name( + cast(Callable[..., Any], inspect.unwrap(annotation_target)) + ) if implicit_parameter_name is None: bound_target = _unwrap_partial(annotation_target) bound_function = getattr(bound_target, "__func__", None) if bound_function is not None: implicit_parameter_name = _first_parameter_name(bound_function) - defining_class = _lexical_defining_class( - annotation_target, descriptor_owner - ) localns = None if defining_class is not None: localns = dict(vars(defining_class)) @@ -770,6 +804,7 @@ def _create_declaration_context( signature_target, implicit_parameter_name, partial_source, + preserves_method_binding, ) return _UDFDeclarationContext( annotation_target=annotation_target, @@ -801,7 +836,12 @@ def _create_resolved_udf( "A synchronous UDF wrapper cannot wrap an async target; define the " "wrapper with async def." ) - is_async = async_marker or target_is_async + if async_marker and not target_is_async: + raise TypeError( + f"AsyncScalarFunction '{_default_udf_name(func)}': eval must be " + "defined with async def." + ) + is_async = target_is_async return _ResolvedUDF( _UDFRuntimeSource(func, kind, is_async), declaration_context ) From 6beb126bde0a19b34de34c56daaaa5c89224a265 Mon Sep 17 00:00:00 2001 From: auroflow Date: Tue, 1 Sep 2026 22:13:27 +0800 Subject: [PATCH 17/18] [FLINK-40431][python] Improve UDF declaration errors Clarify unsupported callable descriptors and distinguish missing return annotations from annotations that cannot be resolved or converted. Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 105 ++++++++---------- flink-python/pyflink/dataframe/udf.py | 24 ++-- 2 files changed, 66 insertions(+), 63 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index a929614940dcc..1dbc9e60ff9c6 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -25,6 +25,7 @@ import unittest from dataclasses import dataclass from typing import Any, Callable, TypedDict, cast +from unittest import mock import pandas as pd import pyarrow as pa @@ -118,6 +119,14 @@ def postponed_return_with_unresolved_input(value): self.assertEqual(_return_dtype(configured), pf.DataType.string()) self.assertEqual(_return_dtype(direct), pf.DataType.int64()) self.assertEqual(direct.__name__, "partial_add_one") + with mock.patch.object( + pf.DataType, + "_from_sql", + return_value=pf.DataType.int64(), + ) as from_sql: + sql_typed = pf.udf(identity, return_dtype="BIGINT") + self.assertEqual(_return_dtype(sql_typed), pf.DataType.int64()) + from_sql.assert_called_once_with("BIGINT") expected_result_dtype = pf.DataType.struct( { @@ -370,14 +379,6 @@ class AsyncAddFunction(AsyncScalarFunction): async def eval(self, *values: int) -> int: return sum(values) - class ExplodingSignature: - @property - def __signature__(self): - raise RuntimeError("signature lookup failed") - - def __call__(self, value): - return value - def variadic_add(*values: int) -> int: return sum(values) @@ -409,17 +410,6 @@ def variadic_add(*values: int) -> int: uninspectable = pf.udf(operator.itemgetter(0), return_dtype=int) self.assertEqual(_return_dtype(uninspectable), pf.DataType.int64()) - exploding_signature = pf.udf(ExplodingSignature(), return_dtype=int) - self.assertEqual( - _return_dtype(exploding_signature), pf.DataType.int64() - ) - exploding_partial_signature = pf.udf( - functools.partial(ExplodingSignature()), return_dtype=int - ) - self.assertEqual( - _return_dtype(exploding_partial_signature), pf.DataType.int64() - ) - self.assertNotIn("__signature__", vars(exploding_partial_signature)) wrapped_callable_instance = WrappedCallableClass() for source in ( @@ -777,7 +767,9 @@ def eval(self, *values: int) -> int: replaced.eval = types.MethodType(MethodOwner.eval, replaced) with self.assertRaisesRegex( - TypeError, "Cannot infer return_dtype.*return_dtype" + TypeError, + r"Cannot infer return_dtype for 'ReplacedScalarFunction' from its " + r"return annotation\.\nSpecify return_dtype explicitly\.", ): pf.udf(replaced) @@ -814,14 +806,19 @@ class InvalidClassMethodCallable: __call__ = classmethod(None) invalid_classes = ( - HiddenCallable, - HiddenScalarFunction, - InvalidStaticCallable, - InvalidClassMethodCallable, + (HiddenCallable, "Callable class", "__call__"), + (HiddenScalarFunction, "Scalar UDF class", "eval"), + (InvalidStaticCallable, "Callable class", "__call__"), + (InvalidClassMethodCallable, "Callable class", "__call__"), ) - for source in invalid_classes: + for source, source_kind, method_name in invalid_classes: with self.subTest(source=source): - with self.assertRaisesRegex(TypeError, "must be defined as a method"): + message = ( + rf"{source_kind} '{source.__name__}' has an unsupported " + rf"{method_name} definition\.\nDefine {method_name} as an " + r"instance, class, or static method\." + ) + with self.assertRaisesRegex(TypeError, message): pf.udf(source, return_dtype=int) def test_descriptor_based_callable_classes_require_instances(self): @@ -837,12 +834,17 @@ class PartialDescriptorCallable: for source in (PartialMethodCallable, PartialDescriptorCallable): with self.subTest(class_source=source): with self.assertRaisesRegex( - TypeError, "must be defined as a method" + TypeError, + rf"Callable class '{source.__name__}' has an unsupported " + r"__call__ definition\.\nDefine __call__ as an instance, " + r"class, or static method\.", ): pf.udf(source, return_dtype=int) with self.subTest(instance_source=source): - declaration = pf.udf(source(), return_dtype=int) + declaration = pf.udf( + source(), return_dtype=int, func_type="general" + ) self.assertEqual( _return_dtype(declaration), pf.DataType.int64() ) @@ -859,7 +861,11 @@ class Output(TypedDict): def __call__(self, value: int) -> "Output": return {"value": value} - with self.assertRaisesRegex(TypeError, "Cannot infer return_dtype"): + with self.assertRaisesRegex( + TypeError, + r"Cannot infer return_dtype for 'Describe' from its return annotation\.\n" + r"Specify return_dtype explicitly\.", + ): pf.udf(Describe) with self.assertRaisesRegex(TypeError, "DataType or SQL"): @@ -873,7 +879,11 @@ class InvalidOutput(TypedDict): def invalid_output(value: int) -> InvalidOutput: return {"value": value} - with self.assertRaisesRegex(TypeError, "Cannot infer return_dtype"): + with self.assertRaisesRegex( + TypeError, + r"Cannot infer return_dtype for 'invalid_output' from its return " + r"annotation\.\nSpecify return_dtype explicitly\.", + ): pf.udf(invalid_output) with self.assertRaisesRegex(TypeError, "DataType or SQL"): @@ -910,7 +920,9 @@ def malformed_return(value: int) -> int: malformed_return.__annotations__["return"] = "list[" with self.assertRaisesRegex( - TypeError, "Cannot infer return_dtype.*return_dtype" + TypeError, + r"Cannot infer return_dtype for 'malformed_return' from its return " + r"annotation\.\nSpecify return_dtype explicitly\.", ): pf.udf(malformed_return) @@ -1198,25 +1210,25 @@ def eval(self, value): "missing return", lambda: pf.udf(missing_return), TypeError, - "Cannot infer return_dtype", + "add a return annotation", ), ( "unresolved return", lambda: pf.udf(unresolved_return), TypeError, - "Cannot infer return_dtype", + r"from its return annotation\.\nSpecify return_dtype explicitly\.", ), ( "callable class missing return", lambda: pf.udf(MissingCallableReturn), TypeError, - "Cannot infer return_dtype for 'MissingCallableReturn'", + "add a return annotation", ), ( "scalar function class missing return", lambda: pf.udf(MissingScalarReturn), TypeError, - "Cannot infer return_dtype for 'MissingScalarReturn'", + "add a return annotation", ), ( "Table return type", @@ -1578,8 +1590,6 @@ def close(self): class DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase): def test_with_columns_binds_expressions_and_resolves_output_schema(self): - sql_typed = pf.udf(lambda value: value, return_dtype="BIGINT") - @pf.udf(name="render_value") def render(value: int, suffix: str) -> str: return f"{value}{suffix}" @@ -1598,12 +1608,11 @@ def describe(value): result = pf.from_records([(1,)], schema=["id"]).with_columns( rendered=render(pf.col("id"), "-literal"), description=describe(pf.col("id")), - sql_value=sql_typed(pf.col("id")), ) self.assert_dataframe_schema( result, - ["id", "rendered", "description", "sql_value"], + ["id", "rendered", "description"], [ TableDataTypes.BIGINT(), TableDataTypes.STRING(), @@ -1615,7 +1624,6 @@ def describe(value): ), ] ), - TableDataTypes.BIGINT(), ], ) @@ -1627,10 +1635,6 @@ class Details: doubled: int labels: list - @pf.udf - def add_one(value: int) -> int: - return value + 1 - @pf.udf async def add_two(value: int) -> int: return value + 2 @@ -1662,34 +1666,23 @@ def eval(self, *values: int) -> int: value, = values return value + self._increment - class ClassNonDeterministic(ScalarFunction): - def eval(self, *values: int) -> int: - value, = values - return value + 6 - - def is_deterministic(self): - return False - deferred = pf.udf(DeferredCallable) opened_scalar_class = pf.udf(OpenedScalarFunction) - scalar_class = pf.udf(ClassNonDeterministic, deterministic=False) result = ( pf.from_records([(1,)], schema=["id"]) .with_columns(async_value=add_two(pf.col("id"))) .with_columns( - sync_value=add_one(pf.col("id")), pandas_value=add_three(pf.col("id")), details=details(pf.col("id")), deferred_value=deferred(pf.col("id")), scalar_value=opened_scalar_class(pf.col("id")), - scalar_class_value=scalar_class(pf.col("id")), ) ) self.assertEqual( result.collect(), - [Row(1, 3, 2, 4, Row(2, ["1"]), 5, 6, 7)], + [Row(1, 3, 4, Row(2, ["1"]), 5, 6)], ) diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 0a8fa448b15b3..a9300b7680702 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -871,8 +871,8 @@ def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: ) if target is None: raise TypeError( - f"Scalar UDF class '{func.__name__}': eval must be defined as a " - "method." + f"Scalar UDF class '{func.__name__}' has an unsupported eval " + "definition.\nDefine eval as an instance, class, or static method." ) declaration_context = _create_declaration_context( target, @@ -894,7 +894,8 @@ def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: if descriptor_owner is None: raise TypeError(f"func must be callable, got {func.__name__}.") raise TypeError( - f"Callable class '{func.__name__}': __call__ must be defined as a method." + f"Callable class '{func.__name__}' has an unsupported __call__ " + "definition.\nDefine __call__ as an instance, class, or static method." ) _validate_zero_argument_class(func) declaration_context = _create_declaration_context( @@ -994,18 +995,27 @@ def _infer_return_dtype( if return_dtype is not None: return _convert_to_dtype(return_dtype) - return_hint = _get_callable_return_type_hint(declaration_context) - if return_hint is _UNRESOLVED_TYPE_HINT: + annotations = getattr( + declaration_context.annotation_target, "__annotations__", {} + ) or {} + if "return" not in annotations: raise TypeError( f"Cannot infer return_dtype for '{udf_name}': add a return annotation " "or specify return_dtype explicitly." ) + + return_hint = _get_callable_return_type_hint(declaration_context) + if return_hint is _UNRESOLVED_TYPE_HINT: + raise TypeError( + f"Cannot infer return_dtype for '{udf_name}' from its return annotation.\n" + "Specify return_dtype explicitly." + ) try: return _data_type_from_type_hint(return_hint) except (NameError, AttributeError, SyntaxError, TypeError) as exc: raise TypeError( - f"Cannot infer return_dtype for '{udf_name}': add a return annotation " - "or specify return_dtype explicitly." + f"Cannot infer return_dtype for '{udf_name}' from its return annotation.\n" + "Specify return_dtype explicitly." ) from exc From ec0d673a9968359e96e2160f82c5604a33873ecd Mon Sep 17 00:00:00 2001 From: auroflow Date: Wed, 2 Sep 2026 19:05:20 +0800 Subject: [PATCH 18/18] [FLINK-40431][python] Align UDF wrapper reflection Expose the DataFrame declaration object's actual expression-building call protocol instead of the worker callable signature. Keep signature inspection internal to partial validation and pandas inference. Generated-by: OpenAI Codex (GPT-5) --- .../pyflink/dataframe/tests/test_udf.py | 171 ++++-------------- flink-python/pyflink/dataframe/udf.py | 120 +++--------- 2 files changed, 60 insertions(+), 231 deletions(-) diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index 1dbc9e60ff9c6..12be3306aebeb 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -106,8 +106,19 @@ def postponed_return_with_unresolved_input(value): self.assertFalse(hasattr(udf_module, "DataFrameUDFWrapper")) self.assertEqual(_return_dtype(decorated), pf.DataType.int64()) self.assertEqual(decorated.__name__, "add_one") + self.assertEqual(decorated.__qualname__, add_one.__qualname__) + self.assertEqual(decorated.__module__, add_one.__module__) self.assertEqual(decorated.__doc__, "Add one to a value.") - self.assertIs(decorated.__wrapped__, add_one) + self.assertNotIn("__wrapped__", vars(decorated)) + self.assertNotIn("__signature__", vars(decorated)) + self.assertNotIn("__annotations__", vars(decorated)) + wrapper_signature = inspect.signature(decorated) + parameters = tuple(wrapper_signature.parameters.values()) + self.assertEqual(len(parameters), 1) + self.assertEqual(parameters[0].name, "args") + self.assertIs(parameters[0].kind, inspect.Parameter.VAR_POSITIONAL) + self.assertIs(parameters[0].annotation, Any) + self.assertIs(wrapper_signature.return_annotation, Expression) configured: Callable[..., Expression] = pf.udf( return_dtype=pf.DataType.string() @@ -242,7 +253,6 @@ def eval(self, *values: int) -> int: self.assertEqual(pf.udf(named_callable).__name__, "configured_add") decorated_class = pf.udf(Double) - self.assertIs(decorated_class.__wrapped__, Double) self.assertEqual(decorated_class.__qualname__, Double.__qualname__) def test_callable_class_resolves_class_local_return_annotation(self): @@ -324,42 +334,18 @@ def eval(self, *values: int) -> str: _return_dtype(pf.udf(overridden)), pf.DataType.int64() ) - def test_wrapped_signature_describes_udf_invocation(self): + def test_wrapped_callable_annotations_and_partial_validation(self): def add(value: int, amount: int = 1) -> int: return value + amount - class CallableClass: - def __call__(self, value: int, amount: int = 1) -> int: - return value + amount - - class StaticCallableClass: - @staticmethod - def __call__(value: int, amount: int = 1) -> int: - return value + amount - - class ClassMethodCallableClass: - @classmethod - def __call__(cls, value: int, amount: int = 1) -> int: - return value + amount - def pandas_identity(values: pd.Series) -> pd.Series: return values - def self_named_pandas_identity( - self: pd.Series, amount: int = 1 - ) -> pd.Series: - return self + amount - class WrappedCallableClass: @functools.wraps(pandas_identity) def __call__(self, *args, **kwargs): return pandas_identity(*args, **kwargs) - class SelfNamedWrappedCallableClass: - @functools.wraps(self_named_pandas_identity) - def __call__(self, *args, **kwargs): - return self_named_pandas_identity(*args, **kwargs) - class WrappedClassMethodCallableClass: @classmethod @functools.wraps(pandas_identity) @@ -371,38 +357,6 @@ class WrappedScalarFunction(ScalarFunction): def eval(self, *args, **kwargs): return pandas_identity(*args, **kwargs) - class AddFunction(ScalarFunction): - def eval(self, *values: int) -> int: - return sum(values) - - class AsyncAddFunction(AsyncScalarFunction): - async def eval(self, *values: int) -> int: - return sum(values) - - def variadic_add(*values: int) -> int: - return sum(values) - - expected_signature = inspect.signature(add) - variadic_signature = inspect.signature(variadic_add) - declarations = [ - (add, expected_signature), - (CallableClass, expected_signature), - (CallableClass(), expected_signature), - (StaticCallableClass, expected_signature), - (ClassMethodCallableClass, expected_signature), - (AddFunction, variadic_signature), - (AddFunction(), variadic_signature), - (AsyncAddFunction, variadic_signature), - (AsyncAddFunction(), variadic_signature), - ] - for source, expected in declarations: - with self.subTest(source=source): - self.assertEqual(inspect.signature(pf.udf(source)), expected) - - partial_add = functools.partial(add, 1) - self.assertEqual( - inspect.signature(pf.udf(partial_add)), inspect.signature(partial_add) - ) with self.assertRaisesRegex( TypeError, "Invalid functools.partial UDF 'add'.*unexpected keyword" ): @@ -425,53 +379,8 @@ def variadic_add(*values: int) -> int: declaration = pf.udf( source, return_dtype=pf.DataType.int64() ) - self.assertEqual( - inspect.signature(declaration), - inspect.signature(pandas_identity), - ) self.assertEqual(declaration._func_type, "pandas") - partial_bound_wrapper = functools.partial( - wrapped_callable_instance.__call__, pd.Series([1]) - ) - partial_declaration = pf.udf( - partial_bound_wrapper, - return_dtype=pf.DataType.int64(), - func_type="general", - ) - self.assertEqual( - inspect.signature(partial_declaration), - inspect.signature(functools.partial(pandas_identity, pd.Series([1]))), - ) - - self_named_instance = SelfNamedWrappedCallableClass() - for source in ( - SelfNamedWrappedCallableClass, - self_named_instance, - self_named_instance.__call__, - ): - with self.subTest(self_named_source=source): - self.assertEqual( - inspect.signature( - pf.udf(source, return_dtype=pf.DataType.int64()) - ), - inspect.signature(self_named_pandas_identity), - ) - self.assertEqual( - inspect.signature( - pf.udf( - functools.partial( - self_named_instance.__call__, pd.Series([1]) - ), - return_dtype=pf.DataType.int64(), - func_type="general", - ) - ), - inspect.signature( - functools.partial(self_named_pandas_identity, pd.Series([1])) - ), - ) - cross_namespace_wrapper = types.FunctionType( _call_module_alias_function.__code__, { @@ -485,7 +394,7 @@ def variadic_add(*values: int) -> int: pf.DataType.int64(), ) - def test_wrapped_methods_preserve_bound_signatures(self): + def test_func_type_resolution_and_async_detection(self): def method_decorator(method): @functools.wraps(method) def wrapper(*args, **kwargs): @@ -493,41 +402,6 @@ def wrapper(*args, **kwargs): return wrapper - def add(value: int, amount: int = 1) -> int: - return value + amount - - class WrappedCallable: - @method_decorator - def __call__(self, value: int, amount: int = 1) -> int: - return value + amount - - class WrappedClassMethodCallable: - @classmethod - @method_decorator - def __call__(cls, value: int, amount: int = 1) -> int: - return value + amount - - expected_signature = inspect.signature(add) - instance = WrappedCallable() - for source in ( - WrappedCallable, - instance, - instance.__call__, - WrappedClassMethodCallable, - WrappedClassMethodCallable(), - ): - with self.subTest(source=source): - self.assertEqual( - inspect.signature(pf.udf(source)), expected_signature - ) - - partial_source = functools.partial(instance.__call__, 1) - self.assertEqual( - inspect.signature(pf.udf(partial_source)), - inspect.signature(functools.partial(add, 1)), - ) - - def test_func_type_resolution_and_async_detection(self): def pandas_add_one(values: pd.Series) -> pd.Series: return values + 1 @@ -588,6 +462,13 @@ async def eval(self, *values: int) -> int: value, = values return value + 1 + class WrappedPandasContext: + @method_decorator + def __call__(self, context: pd.Series, value: int) -> int: + return value + + wrapped_pandas_context = WrappedPandasContext() + declarations = [ ( "inferred pandas", @@ -603,6 +484,16 @@ async def eval(self, *values: int) -> int: "general", False, ), + ( + "bound wrapped pandas annotation is ignored", + lambda: pf.udf( + functools.partial( + wrapped_pandas_context.__call__, pd.Series([1]) + ), + ), + "general", + False, + ), ( "pandas forward reference", lambda: pf.udf( diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index a9300b7680702..b5e55c3c14cfc 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -92,7 +92,6 @@ class _UDFDeclarationContext: defining_class: Optional[Type] globalns: Dict[str, Any] localns: Optional[Dict[str, Any]] - invocation_signature: Optional[inspect.Signature] ignored_hint_names: FrozenSet[str] @@ -209,7 +208,6 @@ def __init__( deterministic: bool, name: str, func_type: str, - invocation_signature: Optional[inspect.Signature], ) -> None: object.__setattr__(self, "_runtime_source", runtime_source) object.__setattr__(self, "_return_dtype", return_dtype) @@ -218,11 +216,13 @@ def __init__( object.__setattr__(self, "_cached_table_udf_wrapper", None) declaration_metadata = _unwrap_partial(runtime_source.callable_source) - functools.update_wrapper(self, declaration_metadata, updated=()) + for attribute_name in ("__module__", "__qualname__", "__doc__"): + try: + attribute_value = getattr(declaration_metadata, attribute_name) + except AttributeError: + continue + object.__setattr__(self, attribute_name, attribute_value) object.__setattr__(self, "__name__", name) - object.__setattr__(self, "__wrapped__", runtime_source.callable_source) - if invocation_signature is not None: - object.__setattr__(self, "__signature__", invocation_signature) object.__setattr__(self, "_frozen", True) def __setattr__(self, name: str, value: Any) -> None: @@ -467,7 +467,6 @@ def decorator(f: _UDFInput) -> Callable[..., Expression]: deterministic, actual_name, actual_func_type, - declaration_context.invocation_signature, ) return decorator if func is None else decorator(func) @@ -657,24 +656,6 @@ def _get_annotation_globals(func: Callable[..., Any]) -> Dict[str, Any]: # ---- Signature and declaration context assembly ---- -def _apply_partial_to_signature( - signature: inspect.Signature, partial_source: functools.partial -) -> inspect.Signature: - def signature_proxy(*args: Any, **kwargs: Any) -> None: - pass - - # Delegate partial's signature transformation to inspect after supplying the - # normalized invocation signature. - setattr(signature_proxy, "__signature__", signature) - return inspect.signature( - functools.partial( - signature_proxy, - *partial_source.args, - **(partial_source.keywords or {}), - ) - ) - - def _preserves_method_binding( target: Callable[..., Any], defining_class: Optional[Type] ) -> bool: @@ -691,18 +672,19 @@ def _preserves_method_binding( ) is defining_class -def _resolve_invocation_signature( +def _resolve_ignored_hint_names( annotation_target: Callable[..., Any], - signature_target: Callable[..., Any], implicit_parameter_name: Optional[str], partial_source: Any, preserves_method_binding: bool, -) -> Tuple[Optional[inspect.Signature], FrozenSet[str]]: +) -> FrozenSet[str]: ignored_hint_names = set() if implicit_parameter_name is not None: ignored_hint_names.add(implicit_parameter_name) - signature_inspection_target: Callable[..., Any] = signature_target + if not isinstance(partial_source, functools.partial): + return frozenset(ignored_hint_names) + bound_function = getattr(annotation_target, "__func__", None) is_wrapped_bound_method = bound_function is not None and hasattr( bound_function, "__wrapped__" @@ -710,66 +692,29 @@ def _resolve_invocation_signature( uses_unbound_wrapped_signature = ( is_wrapped_bound_method and not preserves_method_binding ) - if uses_unbound_wrapped_signature: - signature_inspection_target = cast(Callable[..., Any], bound_function) - try: - invocation_signature = inspect.signature(signature_inspection_target) - parameters = tuple(invocation_signature.parameters.values()) - if ( - implicit_parameter_name is not None - and ( - not hasattr( - getattr(annotation_target, "__func__", annotation_target), - "__wrapped__", - ) - or preserves_method_binding - ) - and parameters - and parameters[0].name == implicit_parameter_name - ): - invocation_signature = invocation_signature.replace( - parameters=parameters[1:] - ) + partial_target_signature = inspect.signature( + cast(Callable[..., Any], bound_function) + if uses_unbound_wrapped_signature + else partial_source.func + ) except Exception: - invocation_signature = None - - if isinstance(partial_source, functools.partial): - try: - partial_target_signature = ( - invocation_signature - if uses_unbound_wrapped_signature - else inspect.signature(partial_source.func) - ) - except Exception: - return None, frozenset(ignored_hint_names) - if partial_target_signature is None: - return None, frozenset(ignored_hint_names) - try: - bound_arguments = partial_target_signature.bind_partial( - *partial_source.args, **(partial_source.keywords or {}) - ) - except TypeError as exc: - raise TypeError( - f"Invalid functools.partial UDF " - f"'{_default_udf_name(partial_source)}': {exc}." - ) from exc - ignored_hint_names.update(bound_arguments.arguments) - if uses_unbound_wrapped_signature: - if invocation_signature is None: - return None, frozenset(ignored_hint_names) - try: - invocation_signature = _apply_partial_to_signature( - invocation_signature, partial_source - ) - except Exception: - return None, frozenset(ignored_hint_names) - return invocation_signature, frozenset(ignored_hint_names) + return frozenset(ignored_hint_names) + try: + bound_arguments = partial_target_signature.bind_partial( + *partial_source.args, **(partial_source.keywords or {}) + ) + except TypeError as exc: + raise TypeError( + f"Invalid functools.partial UDF " + f"'{_default_udf_name(partial_source)}': {exc}." + ) from exc + ignored_hint_names.update(bound_arguments.arguments) + return frozenset(ignored_hint_names) def _create_declaration_context( annotation_target: Callable[..., Any], - signature_target: Callable[..., Any], *, descriptor_owner: Optional[Type] = None, implicit_parameter_name: Optional[str] = None, @@ -799,9 +744,8 @@ def _create_declaration_context( localns = dict(vars(defining_class)) localns[defining_class.__name__] = defining_class - invocation_signature, ignored_hint_names = _resolve_invocation_signature( + ignored_hint_names = _resolve_ignored_hint_names( annotation_target, - signature_target, implicit_parameter_name, partial_source, preserves_method_binding, @@ -811,7 +755,6 @@ def _create_declaration_context( defining_class=defining_class, globalns=_get_annotation_globals(annotation_target), localns=localns, - invocation_signature=invocation_signature, ignored_hint_names=ignored_hint_names, ) @@ -851,7 +794,6 @@ def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: """Validate a UDF and resolve its declaration and runtime metadata.""" if isinstance(func, functools.partial) or inspect.isroutine(func): declaration_context = _create_declaration_context( - cast(Callable[..., Any], func), cast(Callable[..., Any], func), partial_source=func, ) @@ -875,7 +817,6 @@ def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: "definition.\nDefine eval as an instance, class, or static method." ) declaration_context = _create_declaration_context( - target, target, descriptor_owner=descriptor_owner, implicit_parameter_name=implicit_parameter_name, @@ -899,7 +840,6 @@ def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: ) _validate_zero_argument_class(func) declaration_context = _create_declaration_context( - target, target, descriptor_owner=descriptor_owner, implicit_parameter_name=implicit_parameter_name, @@ -919,7 +859,6 @@ def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: f"Scalar UDF instance '{type(func).__name__}': eval must be callable." ) declaration_context = _create_declaration_context( - cast(Callable[..., Any], target), cast(Callable[..., Any], target), partial_source=target, ) @@ -939,7 +878,6 @@ def _resolve_udf(func: _UDFInput) -> _ResolvedUDF: ) declaration_context = _create_declaration_context( cast(Callable[..., Any], target), - cast(Callable[..., Any], func), partial_source=target, ) return _create_resolved_udf(