From 10538fdb074fea3b113ca4e4432bc716ebd27426 Mon Sep 17 00:00:00 2001 From: auroflow Date: Wed, 26 Aug 2026 20:32:30 +0800 Subject: [PATCH 01/10] [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 e07a631d1fe4f1..b4a00f08ce34be 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 datatype 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 00000000000000..4d046bd61b2c8d --- /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 88e7ca2aad55a1..969987bc264fe7 100644 --- a/flink-python/pyflink/dataframe/__init__.py +++ b/flink-python/pyflink/dataframe/__init__.py @@ -54,6 +54,7 @@ from pyflink.dataframe.dataframe import DataFrame, GroupedDataFrame, col, lit from pyflink.dataframe.datatype import DataType from pyflink.dataframe.io import read_generic +from pyflink.dataframe.udf import udf __all__ = [ "DataFrame", @@ -61,6 +62,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 00000000000000..78e0cb30cde43d --- /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 00000000000000..525e2b3f3f1ffb --- /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 4bb5b9ff06c20602e8c822aa5e31eb7eef40d723 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 27 Aug 2026 10:49:40 +0800 Subject: [PATCH 02/10] [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 78e0cb30cde43d..732da8d22fc922 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 525e2b3f3f1ffb..71451722aa151d 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 7722acbede9670af24bb37f278f5977a52c7df70 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 27 Aug 2026 14:16:47 +0800 Subject: [PATCH 03/10] [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 732da8d22fc922..88dea6ed8cdede 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 71451722aa151d..cda1f420785bcc 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 9deab133a93a7ac62e3028b237c1db015aac7b5b Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 27 Aug 2026 14:43:49 +0800 Subject: [PATCH 04/10] [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 88dea6ed8cdede..6af60d8abe536b 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 b5a0e5d5848dedf5ac7d1e15d654c6b575f70f8e Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 15:26:34 +0800 Subject: [PATCH 05/10] [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 ac4c12d24d533e..b4d0da623ada60 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -199,7 +199,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. @@ -210,10 +210,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 6af60d8abe536b..098625df85d7bb 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 cda1f420785bcc..21a7c3993efe90 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 13b4378275444eb1fcc1c25e485cd72f21c17f75 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 20:52:33 +0800 Subject: [PATCH 06/10] [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 4d046bd61b2c8d..73d1ff2abe76d5 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 098625df85d7bb..b7f2f040450104 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 21a7c3993efe90..311e3e4deeccfd 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 c2571c74657c76e145552547e73d4beef5a8ddf9 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 22:01:20 +0800 Subject: [PATCH 07/10] [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 b7f2f040450104..7ab149b747367f 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 311e3e4deeccfd..5849f931ed62a7 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 135d38f677076d6a64b4b395fc07f5aa9e4ba073 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 22:31:50 +0800 Subject: [PATCH 08/10] [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 b4d0da623ada60..4e199b0e34beed 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -199,7 +199,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 5849f931ed62a7..3247a727e8100a 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 20404b305eefc325a871495f3c702efbd1130039 Mon Sep 17 00:00:00 2001 From: auroflow Date: Sun, 30 Aug 2026 12:49:22 +0800 Subject: [PATCH 09/10] [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 7ab149b747367f..ee03c3ae4605cd 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 3247a727e8100a..02a4d148637b01 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 a98fcffcae3b5279706d095b3ae9391ed25a9ef6 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 28 Aug 2026 01:16:09 +0800 Subject: [PATCH 10/10] [FLINK-40431][python] Support per-UDF concurrency and batch size Propagate execution options through Python UDF metadata, split incompatible planner operators, and apply the minimum fused Arrow batch size. Generated-by: Codex (GPT-5) --- .../docs/reference/pyflink.dataframe/udf.rst | 6 + .../pyflink/dataframe/tests/test_udf.py | 151 +++++++++++++- flink-python/pyflink/dataframe/udf.py | 36 +++- flink-python/pyflink/table/udf.py | 59 ++++-- .../batch/PythonUdfExecutionOptionsTest.java | 197 ++++++++++++++++++ .../python/PythonAsyncScalarFunction.java | 90 +++++++- .../functions/python/PythonFunction.java | 10 + .../python/PythonScalarFunction.java | 90 +++++++- .../common/CommonExecPythonAsyncCalc.java | 64 +++--- .../exec/common/CommonExecPythonCalc.java | 17 +- .../nodes/exec/utils/CommonPythonUtil.java | 59 ++++++ .../PythonCalcSplitConcurrencyRule.java | 74 +++++++ .../rules/logical/PythonMapMergeRule.java | 34 ++- .../table/planner/plan/utils/PythonUtil.java | 78 +++++++ .../plan/rules/FlinkBatchRuleSets.scala | 1 + .../plan/rules/FlinkStreamRuleSets.scala | 2 + .../rules/logical/PythonCalcSplitRule.scala | 5 +- .../exec/utils/CommonPythonUtilTest.java | 111 ++++++++++ .../utils/JavaUserDefinedScalarFunctions.java | 15 ++ .../rules/logical/PythonCalcSplitRuleTest.xml | 54 +++++ .../rules/logical/PythonMapMergeRuleTest.xml | 37 ++++ .../logical/PythonCalcSplitRuleTest.scala | 27 +++ .../logical/PythonMapMergeRuleTest.scala | 45 ++++ 23 files changed, 1207 insertions(+), 55 deletions(-) create mode 100644 flink-python/src/test/java/org/apache/flink/table/planner/runtime/batch/PythonUdfExecutionOptionsTest.java create mode 100644 flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitConcurrencyRule.java create mode 100644 flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/utils/CommonPythonUtilTest.java diff --git a/flink-python/docs/reference/pyflink.dataframe/udf.rst b/flink-python/docs/reference/pyflink.dataframe/udf.rst index 73d1ff2abe76d5..fbc279620b1766 100644 --- a/flink-python/docs/reference/pyflink.dataframe/udf.rst +++ b/flink-python/docs/reference/pyflink.dataframe/udf.rst @@ -30,6 +30,12 @@ DataFrame scalar UDFs support synchronous, asynchronous, and pandas-vectorized callables. See :func:`pyflink.dataframe.udf` for declaration forms, type inference, execution modes, and examples. +Use the ``concurrency`` argument to set the parallelism of the Python operator +that executes a UDF. UDFs with different explicit concurrency values are placed +in separate operators. For pandas UDFs, ``batch_size`` sets the maximum Arrow +batch size. If compatible pandas UDFs share an operator, the smallest explicit +batch size is used; otherwise the configured default applies. + API Reference ============= diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py b/flink-python/pyflink/dataframe/tests/test_udf.py index ee03c3ae4605cd..80bd955e6f7b3e 100644 --- a/flink-python/pyflink/dataframe/tests/test_udf.py +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -32,7 +32,12 @@ 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.table.udf import ( + AsyncScalarFunction, + ScalarFunction, + TableFunction, + udf as table_udf, +) from pyflink.testing.test_case_utils import ( PyFlinkDataFrameUTTestCase, PyFlinkStreamDataFrameTestCase, @@ -448,6 +453,104 @@ async def eval(self, *values: int) -> int: with self.assertRaisesRegex(error_type, message): declare() + def test_execution_resource_options_are_propagated(self): + def add_one(value: int) -> int: + return value + 1 + + def pandas_add_one(values: pd.Series) -> pd.Series: + return values + 1 + + general = pf.udf(add_one, concurrency=2) + pandas = pf.udf( + pandas_add_one, + return_dtype=pf.DataType.int64(), + concurrency=3, + batch_size=64, + ) + table_pandas = table_udf( + pandas_add_one, + result_type=TableDataTypes.BIGINT(), + func_type="pandas", + concurrency=4, + batch_size=128, + ) + + declarations = [ + (general, 2, None), + (pandas, 3, 64), + ] + for declaration, expected_concurrency, expected_batch_size in declarations: + with self.subTest(declaration=declaration.__name__): + self.assertEqual(declaration._concurrency, expected_concurrency) + self.assertEqual(declaration._batch_size, expected_batch_size) + self.assertEqual( + declaration._table_udf_wrapper._concurrency, + expected_concurrency, + ) + self.assertEqual( + declaration._table_udf_wrapper._batch_size, + expected_batch_size, + ) + + self.assertEqual(table_pandas._concurrency, 4) + self.assertEqual(table_pandas._batch_size, 128) + + def test_execution_resource_options_are_validated(self): + def add_one(value: int) -> int: + return value + 1 + + def pandas_add_one(values: pd.Series) -> pd.Series: + return values + 1 + + invalid_declarations = [ + ( + "zero concurrency", + lambda: pf.udf(add_one, concurrency=0), + "concurrency must be a positive integer", + ), + ( + "boolean concurrency", + lambda: pf.udf(add_one, concurrency=True), + "concurrency must be a positive integer", + ), + ( + "batch size on general UDF", + lambda: pf.udf(add_one, batch_size=32), + "batch_size is only supported for pandas UDFs", + ), + ( + "negative pandas batch size", + lambda: pf.udf( + pandas_add_one, + return_dtype=pf.DataType.int64(), + batch_size=-1, + ), + "batch_size must be a positive integer", + ), + ( + "Table API boolean concurrency", + lambda: table_udf( + add_one, + result_type=TableDataTypes.BIGINT(), + concurrency=True, + ), + "concurrency must be a positive integer", + ), + ( + "Table API general batch size", + lambda: table_udf( + add_one, + result_type=TableDataTypes.BIGINT(), + batch_size=32, + ), + "batch_size is only supported for pandas UDFs", + ), + ] + for case_name, declare, message in invalid_declarations: + with self.subTest(case=case_name): + with self.assertRaisesRegex(ValueError, message): + declare() + def test_determinism_and_name_metadata(self): class NonDeterministic(ScalarFunction): def eval(self, *values: int) -> int: @@ -1091,6 +1194,41 @@ def close(self): class DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase): + def test_execution_resource_options_reach_java_functions(self): + @pf.udf(concurrency=2) + def add_one(value: int) -> int: + return value + 1 + + @pf.udf(concurrency=3) + async def async_add_one(value: int) -> int: + return value + 1 + + @pf.udf( + return_dtype=pf.DataType.int64(), + func_type="pandas", + concurrency=4, + batch_size=64, + ) + def pandas_add_one(values: pd.Series) -> pd.Series: + return values + 1 + + cases = [ + (add_one, 2, -1), + (async_add_one, 3, -1), + (pandas_add_one, 4, 64), + ] + for declaration, expected_concurrency, expected_batch_size in cases: + with self.subTest(declaration=declaration.__name__): + java_function = ( + declaration._table_udf_wrapper._java_user_defined_function() + ) + self.assertEqual( + java_function.getParallelism(), expected_concurrency + ) + self.assertEqual( + java_function.getMaxArrowBatchSize(), expected_batch_size + ) + def test_with_columns_binds_expressions_and_resolves_output_schema(self): sql_typed = pf.udf(lambda value: value, return_dtype="BIGINT") @@ -1141,15 +1279,20 @@ class Details: doubled: int labels: list - @pf.udf + @pf.udf(concurrency=1) def add_one(value: int) -> int: return value + 1 - @pf.udf + @pf.udf(concurrency=1) async def add_two(value: int) -> int: return value + 2 - @pf.udf(return_dtype=pf.DataType.int64(), func_type="pandas") + @pf.udf( + return_dtype=pf.DataType.int64(), + func_type="pandas", + concurrency=1, + batch_size=1, + ) def add_three(values: pd.Series) -> pd.Series: return values + 3 diff --git a/flink-python/pyflink/dataframe/udf.py b/flink-python/pyflink/dataframe/udf.py index 02a4d148637b01..389380539ded53 100644 --- a/flink-python/pyflink/dataframe/udf.py +++ b/flink-python/pyflink/dataframe/udf.py @@ -49,6 +49,7 @@ ScalarFunction, UserDefinedFunction, UserDefinedFunctionWrapper, + _validate_udf_execution_options, udf as table_udf, ) from pyflink.util.api_stability_decorators import PublicEvolving @@ -224,6 +225,8 @@ class _DataFrameUDFWrapper: _return_dtype: DataType _deterministic: bool _func_type: str + _concurrency: Optional[int] + _batch_size: Optional[int] _cached_table_udf_wrapper: Optional[UserDefinedFunctionWrapper] _frozen: bool __name__: str @@ -235,11 +238,15 @@ def __init__( deterministic: bool, name: str, func_type: str, + concurrency: Optional[int], + batch_size: Optional[int], ) -> None: 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, "_concurrency", concurrency) + object.__setattr__(self, "_batch_size", batch_size) object.__setattr__(self, "_cached_table_udf_wrapper", None) declaration_metadata = _unwrap_partial(source.source) @@ -295,6 +302,8 @@ def _create_table_udf_wrapper( deterministic=self._deterministic, name=self.__name__, func_type=self._func_type, + concurrency=self._concurrency, + batch_size=self._batch_size, ), ) @@ -311,6 +320,8 @@ def udf( deterministic: bool = ..., name: Optional[str] = ..., func_type: Optional[str] = ..., + concurrency: Optional[int] = ..., + batch_size: Optional[int] = ..., ) -> Callable[..., Expression]: ... @@ -323,6 +334,8 @@ def udf( deterministic: bool = ..., name: Optional[str] = ..., func_type: Optional[str] = ..., + concurrency: Optional[int] = ..., + batch_size: Optional[int] = ..., ) -> Callable[[_UDFInput], Callable[..., Expression]]: ... @@ -335,6 +348,8 @@ def udf( deterministic: bool = True, name: Optional[str] = None, func_type: Optional[str] = None, + concurrency: Optional[int] = None, + batch_size: Optional[int] = None, ) -> Union[ Callable[..., Expression], Callable[[_UDFInput], Callable[..., Expression]], @@ -428,7 +443,11 @@ def udf( >>> import pandas as pd - >>> @pf.udf(return_dtype=pf.DataType.int64(), func_type="pandas") + >>> @pf.udf( + ... return_dtype=pf.DataType.int64(), + ... func_type="pandas", + ... batch_size=256, + ... ) ... def pandas_add_one(values): ... return values + 1 @@ -456,6 +475,10 @@ 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. + :param concurrency: Optional parallelism for the Python operator that executes this UDF. + UDFs with different explicit values run in separate operators. + :param batch_size: Optional maximum Arrow batch size for a pandas UDF. If compatible pandas + UDFs are fused, the smallest explicit value is used. :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. @@ -471,7 +494,11 @@ def decorator(f: _UDFInput) -> Callable[..., Expression]: else _detect_func_type(source) ) _validate_scalar_udf_options( - actual_func_type, return_dtype, source.is_async + actual_func_type, + return_dtype, + source.is_async, + concurrency, + batch_size, ) actual_return_dtype = _infer_return_dtype( source.inspection_target, return_dtype @@ -485,6 +512,8 @@ def decorator(f: _UDFInput) -> Callable[..., Expression]: actual_deterministic, actual_name, actual_func_type, + concurrency, + batch_size, ) return decorator if func is None else decorator(func) @@ -497,6 +526,8 @@ def _validate_scalar_udf_options( func_type: str, return_dtype: Optional[_DataTypeLike], is_async: bool, + concurrency: Optional[int], + batch_size: Optional[int], ) -> None: if func_type not in ("general", "pandas"): raise ValueError( @@ -512,6 +543,7 @@ def _validate_scalar_udf_options( "Async scalar functions do not support pandas func_type. " "Use func_type='general'." ) + _validate_udf_execution_options(func_type, concurrency, batch_size) # ======================== Callable Inspection and Resolution ======================== diff --git a/flink-python/pyflink/table/udf.py b/flink-python/pyflink/table/udf.py index 92bd180d1899e1..dec24457f4dcce 100644 --- a/flink-python/pyflink/table/udf.py +++ b/flink-python/pyflink/table/udf.py @@ -18,7 +18,7 @@ import abc import functools import inspect -from typing import Union, List, Type, Callable, TypeVar, Generic, Iterable +from typing import Union, List, Type, Callable, TypeVar, Generic, Iterable, Optional from pyflink.java_gateway import get_gateway from pyflink.metrics import MetricGroup @@ -445,7 +445,8 @@ class UserDefinedFunctionWrapper(object): etc. It's for internal use only. """ - def __init__(self, func, input_types, func_type, deterministic=None, name=None): + def __init__(self, func, input_types, func_type, deterministic=None, name=None, + concurrency=None, batch_size=None): if inspect.isclass(func) or ( not isinstance(func, UserDefinedFunction) and not callable(func)): raise TypeError( @@ -483,6 +484,8 @@ def __init__(self, func, input_types, func_type, deterministic=None, name=None): self._func_type = func_type self._judf_placeholder = None self._takes_row_as_input = False + self._concurrency = concurrency + self._batch_size = batch_size def __call__(self, *args) -> Expression: from pyflink.table import expressions as expr @@ -541,9 +544,10 @@ class UserDefinedScalarFunctionWrapper(UserDefinedFunctionWrapper): Wrapper for Python user-defined scalar function. """ - def __init__(self, func, input_types, result_type, func_type, deterministic, name): + def __init__(self, func, input_types, result_type, func_type, deterministic, name, + concurrency=None, batch_size=None): super(UserDefinedScalarFunctionWrapper, self).__init__( - func, input_types, func_type, deterministic, name) + func, input_types, func_type, deterministic, name, concurrency, batch_size) if not isinstance(result_type, (DataType, str)): raise TypeError( @@ -568,7 +572,9 @@ def _create_judf(self, serialized_func, j_input_types, j_function_kind): j_function_kind, self._deterministic, self._takes_row_as_input, - _get_python_env()) + _get_python_env(), + self._concurrency if self._concurrency is not None else -1, + self._batch_size if self._batch_size is not None else -1) return j_scalar_function def _create_delegate_function(self) -> UserDefinedFunction: @@ -580,9 +586,10 @@ class UserDefinedAsyncScalarFunctionWrapper(UserDefinedFunctionWrapper): Wrapper for Python user-defined async scalar function. """ - def __init__(self, func, input_types, result_type, func_type, deterministic, name): + def __init__(self, func, input_types, result_type, func_type, deterministic, name, + concurrency=None, batch_size=None): super(UserDefinedAsyncScalarFunctionWrapper, self).__init__( - func, input_types, func_type, deterministic, name) + func, input_types, func_type, deterministic, name, concurrency, batch_size) if not isinstance(result_type, (DataType, str)): raise TypeError( @@ -607,7 +614,9 @@ def _create_judf(self, serialized_func, j_input_types, j_function_kind): j_function_kind, self._deterministic, self._takes_row_as_input, - _get_python_env()) + _get_python_env(), + self._concurrency if self._concurrency is not None else -1, + self._batch_size if self._batch_size is not None else -1) return j_async_scalar_function def _create_delegate_function(self) -> UserDefinedFunction: @@ -759,17 +768,20 @@ def _get_python_env(): return gateway.jvm.org.apache.flink.table.functions.python.PythonEnv(exec_type) -def _create_udf(f, input_types, result_type, func_type, deterministic, name): +def _create_udf(f, input_types, result_type, func_type, deterministic, name, + concurrency=None, batch_size=None): if isinstance(f, AsyncScalarFunction) or inspect.iscoroutinefunction(f): if func_type == 'pandas': raise ValueError( "Async scalar functions do not support pandas func_type. " "Please use func_type='general' (default) for async functions.") return UserDefinedAsyncScalarFunctionWrapper( - f, input_types, result_type, func_type, deterministic, name) + f, input_types, result_type, func_type, deterministic, name, + concurrency, batch_size) else: return UserDefinedScalarFunctionWrapper( - f, input_types, result_type, func_type, deterministic, name) + f, input_types, result_type, func_type, deterministic, name, + concurrency, batch_size) def _create_udtf(f, input_types, result_types, deterministic, name): @@ -789,7 +801,8 @@ def _create_udtaf(f, input_types, result_type, accumulator_type, func_type, dete def udf(f: Union[Callable, ScalarFunction, AsyncScalarFunction, Type] = None, input_types: Union[List[DataType], DataType, str, List[str]] = None, result_type: Union[DataType, str] = None, - deterministic: bool = None, name: str = None, func_type: str = "general" + deterministic: bool = None, name: str = None, func_type: str = "general", + concurrency: Optional[int] = None, batch_size: Optional[int] = None ) -> Union[ UserDefinedScalarFunctionWrapper, UserDefinedAsyncScalarFunctionWrapper, Callable]: """ @@ -840,6 +853,8 @@ def udf(f: Union[Callable, ScalarFunction, AsyncScalarFunction, Type] = None, :param name: the function name. :param func_type: the type of the python function, available value: general, pandas, (default: general) + :param concurrency: optional parallelism for the Python operator that executes this UDF. + :param batch_size: optional maximum Arrow batch size. Only supported for pandas UDFs. :return: UserDefinedScalarFunctionWrapper, UserDefinedAsyncScalarFunctionWrapper, or function. .. versionadded:: 1.10.0 @@ -848,14 +863,30 @@ def udf(f: Union[Callable, ScalarFunction, AsyncScalarFunction, Type] = None, if func_type not in ('general', 'pandas'): raise ValueError("The func_type must be one of 'general, pandas', got %s." % func_type) + _validate_udf_execution_options(func_type, concurrency, batch_size) # decorator if f is None: return functools.partial(_create_udf, input_types=input_types, result_type=result_type, func_type=func_type, deterministic=deterministic, - name=name) + name=name, concurrency=concurrency, batch_size=batch_size) else: - return _create_udf(f, input_types, result_type, func_type, deterministic, name) + return _create_udf(f, input_types, result_type, func_type, deterministic, name, + concurrency, batch_size) + + +def _validate_udf_execution_options(func_type, concurrency, batch_size): + for option_name, option_value in ( + ('concurrency', concurrency), ('batch_size', batch_size)): + if option_value is not None and ( + isinstance(option_value, bool) + or not isinstance(option_value, int) + or option_value <= 0): + raise ValueError( + f"{option_name} must be a positive integer, got: {option_value}") + + if batch_size is not None and func_type != 'pandas': + raise ValueError("batch_size is only supported for pandas UDFs.") def udtf(f: Union[Callable, TableFunction, Type] = None, diff --git a/flink-python/src/test/java/org/apache/flink/table/planner/runtime/batch/PythonUdfExecutionOptionsTest.java b/flink-python/src/test/java/org/apache/flink/table/planner/runtime/batch/PythonUdfExecutionOptionsTest.java new file mode 100644 index 00000000000000..86269d8e4777c6 --- /dev/null +++ b/flink-python/src/test/java/org/apache/flink/table/planner/runtime/batch/PythonUdfExecutionOptionsTest.java @@ -0,0 +1,197 @@ +/* + * 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. + */ + +package org.apache.flink.table.planner.runtime.batch; + +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.python.PythonOptions; +import org.apache.flink.streaming.api.operators.SimpleOperatorFactory; +import org.apache.flink.streaming.api.transformations.OneInputTransformation; +import org.apache.flink.table.api.TableConfig; +import org.apache.flink.table.api.config.ExecutionConfigOptions; +import org.apache.flink.table.operations.ModifyOperation; +import org.apache.flink.table.operations.Operation; +import org.apache.flink.table.planner.runtime.utils.JavaUserDefinedScalarFunctions.AsyncPythonScalarFunction; +import org.apache.flink.table.planner.runtime.utils.JavaUserDefinedScalarFunctions.PandasScalarFunction; +import org.apache.flink.table.planner.utils.BatchTableTestUtil; +import org.apache.flink.table.planner.utils.TableTestBase; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for applying per-UDF execution options to translated Python calc operators. */ +class PythonUdfExecutionOptionsTest extends TableTestBase { + + private BatchTableTestUtil util; + + @BeforeEach + void setUp() { + util = batchTestUtil(TableConfig.getDefault()); + util.getTableEnv() + .getConfig() + .set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 4); + util.tableEnv() + .executeSql( + "CREATE TABLE source_table (a INT, b INT) WITH (" + + "'connector' = 'filesystem', " + + "'format' = 'testcsv', " + + "'path' = '/tmp')"); + util.getTableEnv() + .executeSql( + "CREATE TABLE sink_table (a INT, b INT) WITH (" + + "'connector' = 'values')"); + } + + @Test + void testPandasOptionsSetParallelismAndMinimumBatchSize() throws Exception { + util.addTemporarySystemFunction( + "pandas64", new ResourcePandasScalarFunction("pandas64", 2, 64)); + util.addTemporarySystemFunction( + "pandas32", new ResourcePandasScalarFunction("pandas32", 2, 32)); + + final OneInputTransformation pythonCalc = + getPythonTransformation( + "INSERT INTO sink_table " + + "SELECT pandas64(a, b), pandas32(a, b) FROM source_table"); + + assertThat(pythonCalc.getParallelism()).isEqualTo(2); + assertThat(pythonCalc.isParallelismConfigured()).isTrue(); + assertThat(getPythonOperatorConfig(pythonCalc)) + .returns(32, config -> config.get(PythonOptions.MAX_ARROW_BATCH_SIZE)); + } + + @Test + void testUnspecifiedParallelismInheritsInput() throws Exception { + util.getTableEnv().getConfig().set(PythonOptions.MAX_ARROW_BATCH_SIZE, 257); + util.addTemporarySystemFunction( + "pandas_default", new ResourcePandasScalarFunction("pandas_default", -1, -1)); + + final OneInputTransformation pythonCalc = + getPythonTransformation( + "INSERT INTO sink_table " + + "SELECT pandas_default(a, b), b FROM source_table"); + + assertThat(pythonCalc.getParallelism()).isEqualTo(4); + assertThat(pythonCalc.isParallelismConfigured()).isFalse(); + assertThat(getPythonOperatorConfig(pythonCalc)) + .returns(257, config -> config.get(PythonOptions.MAX_ARROW_BATCH_SIZE)); + } + + @Test + void testAsyncConcurrencySetsConfiguredParallelism() { + util.addTemporarySystemFunction( + "async_func", new ResourceAsyncPythonScalarFunction("async_func", 3)); + + final OneInputTransformation pythonCalc = + getPythonTransformation( + "INSERT INTO sink_table SELECT async_func(a, b), b FROM source_table"); + + assertThat(pythonCalc.getParallelism()).isEqualTo(3); + assertThat(pythonCalc.isParallelismConfigured()).isTrue(); + } + + private OneInputTransformation getPythonTransformation(String statement) { + final List operations = util.getPlanner().getParser().parse(statement); + final List> transformations = + util.getPlanner() + .translate(Collections.singletonList((ModifyOperation) operations.get(0))); + assertThat(transformations).hasSize(1); + return findPythonTransformation(transformations.get(0)); + } + + private OneInputTransformation findPythonTransformation( + Transformation transformation) { + if (transformation instanceof OneInputTransformation) { + final OneInputTransformation oneInput = + (OneInputTransformation) transformation; + if (oneInput.getOperatorFactory() instanceof SimpleOperatorFactory) { + final Object operator = + ((SimpleOperatorFactory) oneInput.getOperatorFactory()).getOperator(); + if (operator.getClass().getSimpleName().contains("Python")) { + return oneInput; + } + } + } + for (Transformation input : transformation.getInputs()) { + final OneInputTransformation result = findPythonTransformation(input); + if (result != null) { + return result; + } + } + return null; + } + + private Configuration getPythonOperatorConfig(OneInputTransformation transformation) + throws Exception { + final Object operator = + ((SimpleOperatorFactory) transformation.getOperatorFactory()).getOperator(); + Class operatorClass = operator.getClass(); + while (operatorClass != null) { + try { + final Field configField = operatorClass.getDeclaredField("config"); + configField.setAccessible(true); + return (Configuration) configField.get(operator); + } catch (NoSuchFieldException ignored) { + operatorClass = operatorClass.getSuperclass(); + } + } + throw new AssertionError("Python operator configuration field was not found."); + } + + public static class ResourcePandasScalarFunction extends PandasScalarFunction { + private final int parallelism; + private final int maxArrowBatchSize; + + public ResourcePandasScalarFunction(String name, int parallelism, int maxArrowBatchSize) { + super(name); + this.parallelism = parallelism; + this.maxArrowBatchSize = maxArrowBatchSize; + } + + @Override + public int getParallelism() { + return parallelism; + } + + @Override + public int getMaxArrowBatchSize() { + return maxArrowBatchSize; + } + } + + public static class ResourceAsyncPythonScalarFunction extends AsyncPythonScalarFunction { + private final int parallelism; + + public ResourceAsyncPythonScalarFunction(String name, int parallelism) { + super(name); + this.parallelism = parallelism; + } + + @Override + public int getParallelism() { + return parallelism; + } + } +} diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonAsyncScalarFunction.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonAsyncScalarFunction.java index 85c166424c0245..48dcc6ecafe566 100644 --- a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonAsyncScalarFunction.java +++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonAsyncScalarFunction.java @@ -43,6 +43,8 @@ public class PythonAsyncScalarFunction extends AsyncScalarFunction implements Py private final boolean deterministic; private final PythonEnv pythonEnv; private final boolean takesRowAsInput; + private final int parallelism; + private final int maxArrowBatchSize; private DataType[] inputTypes; private String[] inputTypesString; @@ -61,10 +63,36 @@ public PythonAsyncScalarFunction( this( name, serializedAsyncScalarFunction, + inputTypes, + resultType, pythonFunctionKind, deterministic, takesRowAsInput, - pythonEnv); + pythonEnv, + -1, + -1); + } + + public PythonAsyncScalarFunction( + String name, + byte[] serializedAsyncScalarFunction, + DataType[] inputTypes, + DataType resultType, + PythonFunctionKind pythonFunctionKind, + boolean deterministic, + boolean takesRowAsInput, + PythonEnv pythonEnv, + int parallelism, + int maxArrowBatchSize) { + this( + name, + serializedAsyncScalarFunction, + pythonFunctionKind, + deterministic, + takesRowAsInput, + pythonEnv, + parallelism, + maxArrowBatchSize); this.inputTypes = inputTypes; this.resultType = resultType; } @@ -78,13 +106,39 @@ public PythonAsyncScalarFunction( boolean deterministic, boolean takesRowAsInput, PythonEnv pythonEnv) { + this( + name, + serializedAsyncScalarFunction, + inputTypesString, + resultTypeString, + pythonFunctionKind, + deterministic, + takesRowAsInput, + pythonEnv, + -1, + -1); + } + + public PythonAsyncScalarFunction( + String name, + byte[] serializedAsyncScalarFunction, + String[] inputTypesString, + String resultTypeString, + PythonFunctionKind pythonFunctionKind, + boolean deterministic, + boolean takesRowAsInput, + PythonEnv pythonEnv, + int parallelism, + int maxArrowBatchSize) { this( name, serializedAsyncScalarFunction, pythonFunctionKind, deterministic, takesRowAsInput, - pythonEnv); + pythonEnv, + parallelism, + maxArrowBatchSize); this.inputTypesString = inputTypesString; this.resultTypeString = resultTypeString; } @@ -96,12 +150,34 @@ public PythonAsyncScalarFunction( boolean deterministic, boolean takesRowAsInput, PythonEnv pythonEnv) { + this( + name, + serializedAsyncScalarFunction, + pythonFunctionKind, + deterministic, + takesRowAsInput, + pythonEnv, + -1, + -1); + } + + public PythonAsyncScalarFunction( + String name, + byte[] serializedAsyncScalarFunction, + PythonFunctionKind pythonFunctionKind, + boolean deterministic, + boolean takesRowAsInput, + PythonEnv pythonEnv, + int parallelism, + int maxArrowBatchSize) { this.name = name; this.serializedAsyncScalarFunction = serializedAsyncScalarFunction; this.pythonFunctionKind = pythonFunctionKind; this.deterministic = deterministic; this.pythonEnv = pythonEnv; this.takesRowAsInput = takesRowAsInput; + this.parallelism = parallelism; + this.maxArrowBatchSize = maxArrowBatchSize; } public void eval(CompletableFuture future, Object... args) { @@ -129,6 +205,16 @@ public boolean takesRowAsInput() { return takesRowAsInput; } + @Override + public int getParallelism() { + return parallelism; + } + + @Override + public int getMaxArrowBatchSize() { + return maxArrowBatchSize; + } + @Override public boolean isDeterministic() { return deterministic; diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunction.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunction.java index a1a77764cdb47c..a6ba85eef209bc 100644 --- a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunction.java +++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunction.java @@ -44,4 +44,14 @@ default PythonFunctionKind getPythonFunctionKind() { default boolean takesRowAsInput() { return false; } + + /** Returns the configured parallelism, or a non-positive value if it is not configured. */ + default int getParallelism() { + return -1; + } + + /** Returns the maximum Arrow batch size, or a non-positive value if it is not configured. */ + default int getMaxArrowBatchSize() { + return -1; + } } diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunction.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunction.java index 9d14db265f9917..211bca823f6cc4 100644 --- a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunction.java +++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunction.java @@ -44,6 +44,8 @@ public class PythonScalarFunction extends ScalarFunction implements PythonFuncti private final boolean deterministic; private final PythonEnv pythonEnv; private final boolean takesRowAsInput; + private final int parallelism; + private final int maxArrowBatchSize; private DataType[] inputTypes; private String[] inputTypesString; @@ -62,10 +64,36 @@ public PythonScalarFunction( this( name, serializedScalarFunction, + inputTypes, + resultType, pythonFunctionKind, deterministic, takesRowAsInput, - pythonEnv); + pythonEnv, + -1, + -1); + } + + public PythonScalarFunction( + String name, + byte[] serializedScalarFunction, + DataType[] inputTypes, + DataType resultType, + PythonFunctionKind pythonFunctionKind, + boolean deterministic, + boolean takesRowAsInput, + PythonEnv pythonEnv, + int parallelism, + int maxArrowBatchSize) { + this( + name, + serializedScalarFunction, + pythonFunctionKind, + deterministic, + takesRowAsInput, + pythonEnv, + parallelism, + maxArrowBatchSize); this.inputTypes = inputTypes; this.resultType = resultType; } @@ -79,13 +107,39 @@ public PythonScalarFunction( boolean deterministic, boolean takesRowAsInput, PythonEnv pythonEnv) { + this( + name, + serializedScalarFunction, + inputTypesString, + resultTypeString, + pythonFunctionKind, + deterministic, + takesRowAsInput, + pythonEnv, + -1, + -1); + } + + public PythonScalarFunction( + String name, + byte[] serializedScalarFunction, + String[] inputTypesString, + String resultTypeString, + PythonFunctionKind pythonFunctionKind, + boolean deterministic, + boolean takesRowAsInput, + PythonEnv pythonEnv, + int parallelism, + int maxArrowBatchSize) { this( name, serializedScalarFunction, pythonFunctionKind, deterministic, takesRowAsInput, - pythonEnv); + pythonEnv, + parallelism, + maxArrowBatchSize); this.inputTypesString = inputTypesString; this.resultTypeString = resultTypeString; } @@ -97,12 +151,34 @@ public PythonScalarFunction( boolean deterministic, boolean takesRowAsInput, PythonEnv pythonEnv) { + this( + name, + serializedScalarFunction, + pythonFunctionKind, + deterministic, + takesRowAsInput, + pythonEnv, + -1, + -1); + } + + public PythonScalarFunction( + String name, + byte[] serializedScalarFunction, + PythonFunctionKind pythonFunctionKind, + boolean deterministic, + boolean takesRowAsInput, + PythonEnv pythonEnv, + int parallelism, + int maxArrowBatchSize) { this.name = name; this.serializedScalarFunction = serializedScalarFunction; this.pythonFunctionKind = pythonFunctionKind; this.deterministic = deterministic; this.pythonEnv = pythonEnv; this.takesRowAsInput = takesRowAsInput; + this.parallelism = parallelism; + this.maxArrowBatchSize = maxArrowBatchSize; } public Object eval(Object... args) { @@ -130,6 +206,16 @@ public boolean takesRowAsInput() { return takesRowAsInput; } + @Override + public int getParallelism() { + return parallelism; + } + + @Override + public int getMaxArrowBatchSize() { + return maxArrowBatchSize; + } + @Override public boolean isDeterministic() { return deterministic; diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonAsyncCalc.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonAsyncCalc.java index ea20e465a98732..631eef3e0422e0 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonAsyncCalc.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonAsyncCalc.java @@ -58,9 +58,9 @@ import org.apache.calcite.rex.RexNode; import java.lang.reflect.Constructor; -import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; +import java.util.Optional; import java.util.stream.Collectors; import static org.apache.flink.util.Preconditions.checkArgument; @@ -121,21 +121,32 @@ private OneInputTransformation createPythonAsyncOneInputTransf ClassLoader classLoader, Configuration pythonConfig) { + final Tuple2 extractResult = + extractPythonFunctionInfos(classLoader); + final Optional explicitParallelism = + CommonPythonUtil.deriveExplicitPythonFunctionParallelism( + extractResult.f1, "Python async functions"); + InternalTypeInfo inputTypeInfo = (InternalTypeInfo) inputTransform.getOutputType(); InternalTypeInfo outputTypeInfo = InternalTypeInfo.of((RowType) getOutputType()); OneInputStreamOperator pythonAsyncOperator = getPythonAsyncOperator( - config, classLoader, pythonConfig, inputTypeInfo, outputTypeInfo); + config, + classLoader, + pythonConfig, + inputTypeInfo, + outputTypeInfo, + extractResult); return ExecNodeUtil.createOneInputTransformation( inputTransform, createTransformationMeta(PYTHON_ASYNC_CALC_TRANSFORMATION, config), pythonAsyncOperator, outputTypeInfo, - inputTransform.getParallelism(), - false); + explicitParallelism.orElse(inputTransform.getParallelism()), + explicitParallelism.isPresent()); } /** Gets the async operator for executing Python async scalar functions. */ @@ -144,7 +155,8 @@ private OneInputStreamOperator getPythonAsyncOperator( ClassLoader classLoader, Configuration pythonConfig, InternalTypeInfo inputTypeInfo, - InternalTypeInfo outputTypeInfo) { + InternalTypeInfo outputTypeInfo, + Tuple2 extractResult) { boolean isInProcessMode = CommonPythonUtil.isPythonWorkerInProcessMode(pythonConfig, classLoader); @@ -153,28 +165,12 @@ private OneInputStreamOperator getPythonAsyncOperator( "Python async scalar function is still not supported for 'thread' mode."); } - // Separate async function calls from forwarded fields - List asyncRexCalls = new ArrayList<>(); - List forwardedFields = new ArrayList<>(); - - for (RexNode rexNode : projection) { - if (rexNode instanceof RexCall) { - RexCall rexCall = (RexCall) rexNode; - if (isPythonAsyncCall(rexCall)) { - asyncRexCalls.add(rexCall); - } - } else if (rexNode instanceof RexInputRef) { - forwardedFields.add(((RexInputRef) rexNode).getIndex()); - } - } - - if (asyncRexCalls.isEmpty()) { - throw new IllegalStateException("No Python async scalar function found in projection"); - } - - // Extract Python function information - Tuple2 extractResult = - extractPythonAsyncScalarFunctionInfos(asyncRexCalls, classLoader); + final List forwardedFields = + projection.stream() + .filter(RexInputRef.class::isInstance) + .map(RexInputRef.class::cast) + .map(RexInputRef::getIndex) + .collect(Collectors.toList()); int[] udfInputOffsets = extractResult.f0; PythonFunctionInfo[] pythonFunctionInfos = extractResult.f1; @@ -281,6 +277,20 @@ private Tuple2 extractPythonAsyncScalarFunctionInfo return Tuple2.of(udfInputOffsets, pythonFunctionInfos); } + private Tuple2 extractPythonFunctionInfos( + ClassLoader classLoader) { + List asyncRexCalls = + projection.stream() + .filter(RexCall.class::isInstance) + .map(RexCall.class::cast) + .filter(this::isPythonAsyncCall) + .collect(Collectors.toList()); + if (asyncRexCalls.isEmpty()) { + throw new IllegalStateException("No Python async scalar function found in projection"); + } + return extractPythonAsyncScalarFunctionInfos(asyncRexCalls, classLoader); + } + private boolean isPythonAsyncCall(RexCall rexCall) { if (rexCall.getOperator() instanceof BridgingSqlFunction) { BridgingSqlFunction function = (BridgingSqlFunction) rexCall.getOperator(); diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonCalc.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonCalc.java index 9c6013055cab6e..74be2c43b7034e 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonCalc.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonCalc.java @@ -20,6 +20,7 @@ import org.apache.flink.api.dag.Transformation; import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.configuration.ConfigOptions; import org.apache.flink.configuration.Configuration; import org.apache.flink.configuration.ReadableConfig; import org.apache.flink.core.memory.ManagedMemoryUseCase; @@ -61,6 +62,7 @@ import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; +import java.util.Optional; import java.util.stream.Collectors; import static org.apache.flink.util.Preconditions.checkArgument; @@ -166,6 +168,17 @@ private OneInputTransformation createPythonOneInputTransformat fieldsLogicalTypes.addAll(pythonCallLogicalTypes); InternalTypeInfo pythonOperatorResultTyeInfo = InternalTypeInfo.ofFields(fieldsLogicalTypes.toArray(new LogicalType[0])); + final int maxArrowBatchSize = CommonPythonUtil.deriveArrowBatchSize(pythonFunctionInfos); + if (maxArrowBatchSize > 0) { + pythonConfig.set( + ConfigOptions.key("python.fn-execution.arrow.batch.size") + .intType() + .noDefaultValue(), + maxArrowBatchSize); + } + final Optional explicitParallelism = + CommonPythonUtil.deriveExplicitPythonFunctionParallelism( + pythonFunctionInfos, "Python functions"); OneInputStreamOperator pythonOperator = getPythonScalarFunctionOperator( config, @@ -187,8 +200,8 @@ private OneInputTransformation createPythonOneInputTransformat createTransformationMeta(PYTHON_CALC_TRANSFORMATION, config), pythonOperator, pythonOperatorResultTyeInfo, - inputTransform.getParallelism(), - false); + explicitParallelism.orElse(inputTransform.getParallelism()), + explicitParallelism.isPresent()); } private Tuple2 extractPythonScalarFunctionInfos( diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/utils/CommonPythonUtil.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/utils/CommonPythonUtil.java index d82a37f5e7c8a0..0505a015732950 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/utils/CommonPythonUtil.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/utils/CommonPythonUtil.java @@ -83,10 +83,14 @@ import java.math.BigDecimal; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; +import java.util.IdentityHashMap; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; +import java.util.Set; import java.util.stream.IntStream; /** A utility class used in PyFlink. */ @@ -120,6 +124,61 @@ public static Configuration extractPythonConfiguration( } } + public static Optional deriveExplicitPythonFunctionParallelism( + PythonFunctionInfo[] pythonFunctionInfos, String functionsDescription) { + int explicitParallelism = -1; + for (PythonFunction function : collectPythonFunctions(pythonFunctionInfos)) { + final int parallelism = function.getParallelism(); + if (parallelism > 0) { + if (explicitParallelism > 0 && explicitParallelism != parallelism) { + throw new TableException( + functionsDescription + + " with different concurrency values should have been split " + + "into different operators."); + } + explicitParallelism = parallelism; + } + } + return explicitParallelism > 0 ? Optional.of(explicitParallelism) : Optional.empty(); + } + + public static int deriveArrowBatchSize(PythonFunctionInfo[] pythonFunctionInfos) { + int minBatchSize = -1; + for (PythonFunction function : collectPythonFunctions(pythonFunctionInfos)) { + final int batchSize = function.getMaxArrowBatchSize(); + if (batchSize > 0) { + minBatchSize = minBatchSize > 0 ? Math.min(minBatchSize, batchSize) : batchSize; + } + } + return minBatchSize; + } + + private static List collectPythonFunctions( + PythonFunctionInfo[] pythonFunctionInfos) { + final List functions = new ArrayList<>(); + final Set visited = Collections.newSetFromMap(new IdentityHashMap<>()); + for (PythonFunctionInfo functionInfo : pythonFunctionInfos) { + collectPythonFunctions(functionInfo, functions, visited); + } + return functions; + } + + private static void collectPythonFunctions( + PythonFunctionInfo functionInfo, + List functions, + Set visited) { + if (!visited.add(functionInfo)) { + return; + } + + functions.add(functionInfo.getPythonFunction()); + for (Object input : functionInfo.getInputs()) { + if (input instanceof PythonFunctionInfo) { + collectPythonFunctions((PythonFunctionInfo) input, functions, visited); + } + } + } + public static PythonFunctionInfo createPythonFunctionInfo( RexCall pythonRexCall, Map inputNodes, ClassLoader classLoader) { SqlOperator operator = pythonRexCall.getOperator(); diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitConcurrencyRule.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitConcurrencyRule.java new file mode 100644 index 00000000000000..aaa2dfd3d85b10 --- /dev/null +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitConcurrencyRule.java @@ -0,0 +1,74 @@ +/* + * 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. + */ + +package org.apache.flink.table.planner.plan.rules.logical; + +import org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc; +import org.apache.flink.table.planner.plan.utils.PythonUtil; + +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexProgram; + +import java.util.Optional; + +import scala.Option; + +/** + * Splits projections containing Python functions with different explicit concurrency values. + * Functions without explicit concurrency can remain with any explicit group. + */ +public class PythonCalcSplitConcurrencyRule extends RemoteCalcSplitProjectionRuleBase { + + public PythonCalcSplitConcurrencyRule(RemoteCallFinder callFinder) { + super("PythonCalcSplitConcurrencyRule", callFinder); + } + + @Override + public boolean matches(RelOptRuleCall call) { + final FlinkLogicalCalc calc = call.rel(0); + return calc.getProgram().getProjectList().stream() + .map(calc.getProgram()::expandLocalRef) + .flatMap(node -> PythonUtil.extractDistinctParallelisms(node).stream()) + .distinct() + .count() + > 1; + } + + @Override + public boolean needConvert(RexProgram program, RexNode node, Option matchState) { + if (!callFinder().isRemoteCall(node)) { + return false; + } + + final Optional ownParallelism = PythonUtil.getOwnParallelism(node); + if (!ownParallelism.isPresent()) { + return false; + } + + final Optional targetParallelism = + program.getProjectList().stream() + .map(program::expandLocalRef) + .map(PythonUtil::firstExplicitParallelism) + .filter(Optional::isPresent) + .map(Optional::get) + .findFirst(); + return targetParallelism.isPresent() + && !targetParallelism.get().equals(ownParallelism.get()); + } +} diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRule.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRule.java index 57d2f2cea68815..9cfc17ed6742d0 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRule.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRule.java @@ -18,6 +18,7 @@ package org.apache.flink.table.planner.plan.rules.logical; +import org.apache.flink.table.functions.python.PythonFunction; import org.apache.flink.table.functions.python.PythonFunctionKind; import org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc; import org.apache.flink.table.planner.plan.utils.PythonUtil; @@ -33,7 +34,9 @@ import org.apache.calcite.rex.RexProgramBuilder; import java.util.Collections; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.stream.Collectors; /** @@ -94,7 +97,36 @@ public boolean matches(RelOptRuleCall call) { return PythonUtil.isFlattenCalc(middleCalc) && isTopCalcTakesWholeMiddleCalcAsInputs( - (RexCall) topProjects.get(0), middleCalc.getRowType().getFieldCount()); + (RexCall) topProjects.get(0), middleCalc.getRowType().getFieldCount()) + && hasCompatibleExecutionOptions( + (RexCall) topProjects.get(0), (RexCall) bottomProjects.get(0)); + } + + private boolean hasCompatibleExecutionOptions(RexCall topCall, RexCall bottomCall) { + final Set parallelisms = new HashSet<>(); + final Set batchSizes = new HashSet<>(); + collectExecutionOptions(topCall, parallelisms, batchSizes); + collectExecutionOptions(bottomCall, parallelisms, batchSizes); + return parallelisms.size() <= 1 && batchSizes.size() <= 1; + } + + private void collectExecutionOptions( + RexNode node, Set parallelisms, Set batchSizes) { + if (!(node instanceof RexCall)) { + return; + } + final RexCall call = (RexCall) node; + final PythonFunction function = PythonUtil.extractPythonFunction(call); + if (function != null) { + if (function.getParallelism() > 0) { + parallelisms.add(function.getParallelism()); + } + if (function.getMaxArrowBatchSize() > 0) { + batchSizes.add(function.getMaxArrowBatchSize()); + } + } + call.getOperands() + .forEach(operand -> collectExecutionOptions(operand, parallelisms, batchSizes)); } private boolean isTopCalcTakesWholeMiddleCalcAsInputs( diff --git a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/utils/PythonUtil.java b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/utils/PythonUtil.java index 360129870cf291..2f986ce5261b27 100644 --- a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/utils/PythonUtil.java +++ b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/utils/PythonUtil.java @@ -44,8 +44,10 @@ import org.apache.calcite.sql.SqlAggFunction; import org.apache.calcite.sql.SqlKind; +import java.util.LinkedHashSet; import java.util.List; import java.util.Optional; +import java.util.Set; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -165,6 +167,82 @@ public static boolean takesRowAsInput(RexCall call) { return false; } + /** Returns the Python function represented by the call, or {@code null} for other calls. */ + public static PythonFunction extractPythonFunction(RexCall call) { + final FunctionDefinition definition; + if (call.getOperator() instanceof ScalarSqlFunction) { + definition = ((ScalarSqlFunction) call.getOperator()).scalarFunction(); + } else if (call.getOperator() instanceof TableSqlFunction) { + definition = ((TableSqlFunction) call.getOperator()).udtf(); + } else if (call.getOperator() instanceof BridgingSqlFunction) { + definition = ((BridgingSqlFunction) call.getOperator()).getDefinition(); + } else { + return null; + } + return definition instanceof PythonFunction ? (PythonFunction) definition : null; + } + + /** Returns the call's explicit parallelism without inspecting nested calls. */ + public static Optional getOwnParallelism(RexNode node) { + if (!(node instanceof RexCall)) { + return Optional.empty(); + } + final PythonFunction function = extractPythonFunction((RexCall) node); + if (function == null || function.getParallelism() <= 0) { + return Optional.empty(); + } + return Optional.of(function.getParallelism()); + } + + /** Returns the first explicit Python parallelism in a pre-order traversal. */ + public static Optional firstExplicitParallelism(RexNode node) { + final Optional ownParallelism = getOwnParallelism(node); + if (ownParallelism.isPresent()) { + return ownParallelism; + } + if (node instanceof RexCall) { + for (RexNode operand : ((RexCall) node).getOperands()) { + final Optional nestedParallelism = firstExplicitParallelism(operand); + if (nestedParallelism.isPresent()) { + return nestedParallelism; + } + } + } else if (node instanceof RexFieldAccess) { + return firstExplicitParallelism(((RexFieldAccess) node).getReferenceExpr()); + } + return Optional.empty(); + } + + /** Returns all explicit Python parallelisms in the expression tree. */ + public static Set extractDistinctParallelisms(RexNode node) { + final Set parallelisms = new LinkedHashSet<>(); + node.accept( + new RexDefaultVisitor() { + @Override + public Void visitCall(RexCall call) { + getOwnParallelism(call).ifPresent(parallelisms::add); + call.getOperands().forEach(operand -> operand.accept(this)); + return null; + } + + @Override + public Void visitFieldAccess(RexFieldAccess fieldAccess) { + return fieldAccess.getReferenceExpr().accept(this); + } + + @Override + public Void visitNode(RexNode rexNode) { + return null; + } + + @Override + public Void visitNodeAndFieldIndex(RexNodeAndFieldIndex nodeAndFieldIndex) { + throw new UnsupportedOperationException("not supported yet"); + } + }); + return parallelisms; + } + private static boolean isPythonFunction( FunctionDefinition function, PythonFunctionKind pythonFunctionKind) { if (function instanceof PythonFunction) { diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkBatchRuleSets.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkBatchRuleSets.scala index 26553f24ddcd04..eb3b8dde01b71d 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkBatchRuleSets.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkBatchRuleSets.scala @@ -393,6 +393,7 @@ object FlinkBatchRuleSets { PythonCalcSplitRule.SPLIT_CONDITION, PythonCalcSplitRule.SPLIT_PROJECT, PythonCalcSplitRule.SPLIT_PANDAS_IN_PROJECT, + PythonCalcSplitRule.SPLIT_CONCURRENCY_IN_PROJECT, PythonCalcSplitRule.EXPAND_PROJECT, PythonCalcSplitRule.PUSH_CONDITION, PythonCalcSplitRule.REWRITE_PROJECT, diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkStreamRuleSets.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkStreamRuleSets.scala index de637d07a244bd..04b2a3ebc4b47b 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkStreamRuleSets.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/FlinkStreamRuleSets.scala @@ -419,6 +419,8 @@ object FlinkStreamRuleSets { PythonCalcSplitRule.SPLIT_PROJECT, // Splits calcs which contain both general Python functions and pandas Python functions PythonCalcSplitRule.SPLIT_PANDAS_IN_PROJECT, + // Splits calcs which contain Python functions with different concurrency values + PythonCalcSplitRule.SPLIT_CONCURRENCY_IN_PROJECT, // Avoid accessing a field as input to an async call with a single calc. PythonCalcSplitRule.EXPAND_PROJECT, // Avoid having any condition in a python calc by pushing it first. diff --git a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRule.scala b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRule.scala index 7e54ffbca03c9c..e542e36a4e43f0 100644 --- a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRule.scala +++ b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRule.scala @@ -97,12 +97,15 @@ object PythonCalcSplitRule { /** * These rules should be applied sequentially in the order of SPLIT_CONDITION, SPLIT_PROJECT, - * SPLIT_PANDAS_IN_PROJECT, EXPAND_PROJECT, PUSH_CONDITION and REWRITE_PROJECT. + * SPLIT_PANDAS_IN_PROJECT, SPLIT_CONCURRENCY_IN_PROJECT, EXPAND_PROJECT, PUSH_CONDITION and + * REWRITE_PROJECT. */ private val callFinder = new PythonRemoteCallFinder() val SPLIT_CONDITION: RelOptRule = new RemoteCalcSplitConditionRule(callFinder) val SPLIT_PROJECT: RelOptRule = new RemoteCalcSplitProjectionRule(callFinder) val SPLIT_PANDAS_IN_PROJECT: RelOptRule = new PythonCalcSplitPandasInProjectionRule(callFinder) + val SPLIT_CONCURRENCY_IN_PROJECT: RelOptRule = + new PythonCalcSplitConcurrencyRule(callFinder) val SPLIT_PROJECTION_REX_FIELD: RelOptRule = new RemoteCalcSplitProjectionRexFieldRule(callFinder) val SPLIT_CONDITION_REX_FIELD: RelOptRule = new RemoteCalcSplitConditionRexFieldRule(callFinder) val EXPAND_PROJECT: RelOptRule = new RemoteCalcExpandProjectRule(callFinder) diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/utils/CommonPythonUtilTest.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/utils/CommonPythonUtilTest.java new file mode 100644 index 00000000000000..8e0400e0759974 --- /dev/null +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/nodes/exec/utils/CommonPythonUtilTest.java @@ -0,0 +1,111 @@ +/* + * 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. + */ + +package org.apache.flink.table.planner.plan.nodes.exec.utils; + +import org.apache.flink.table.api.TableException; +import org.apache.flink.table.functions.python.PythonEnv; +import org.apache.flink.table.functions.python.PythonFunction; +import org.apache.flink.table.functions.python.PythonFunctionInfo; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for deriving execution resources from Python function trees. */ +class CommonPythonUtilTest { + + @Test + void testDeriveParallelismFromNestedFunctions() { + final PythonFunctionInfo nested = + functionInfo(-1, -1, functionInfo(3, -1), functionInfo(-1, -1)); + + assertThat( + CommonPythonUtil.deriveExplicitPythonFunctionParallelism( + new PythonFunctionInfo[] {nested}, "Python functions")) + .hasValue(3); + assertThat( + CommonPythonUtil.deriveExplicitPythonFunctionParallelism( + new PythonFunctionInfo[] {functionInfo(-1, -1)}, + "Python functions")) + .isEmpty(); + } + + @Test + void testRejectConflictingNestedParallelism() { + final PythonFunctionInfo nested = + functionInfo(2, -1, functionInfo(-1, -1, functionInfo(3, -1))); + + assertThatThrownBy( + () -> + CommonPythonUtil.deriveExplicitPythonFunctionParallelism( + new PythonFunctionInfo[] {nested}, "Python functions")) + .isInstanceOf(TableException.class) + .hasMessageContaining("different concurrency values"); + } + + @Test + void testDeriveMinimumBatchSizeFromNestedFunctions() { + final PythonFunctionInfo nested = + functionInfo(2, 64, functionInfo(-1, 32), functionInfo(-1, -1)); + + assertThat(CommonPythonUtil.deriveArrowBatchSize(new PythonFunctionInfo[] {nested})) + .isEqualTo(32); + assertThat( + CommonPythonUtil.deriveArrowBatchSize( + new PythonFunctionInfo[] {functionInfo(-1, -1)})) + .isEqualTo(-1); + } + + private static PythonFunctionInfo functionInfo( + int parallelism, int maxArrowBatchSize, PythonFunctionInfo... inputs) { + return new PythonFunctionInfo( + new TestPythonFunction(parallelism, maxArrowBatchSize), inputs); + } + + private static final class TestPythonFunction implements PythonFunction { + private final int parallelism; + private final int maxArrowBatchSize; + + private TestPythonFunction(int parallelism, int maxArrowBatchSize) { + this.parallelism = parallelism; + this.maxArrowBatchSize = maxArrowBatchSize; + } + + @Override + public byte[] getSerializedPythonFunction() { + return new byte[0]; + } + + @Override + public PythonEnv getPythonEnv() { + return null; + } + + @Override + public int getParallelism() { + return parallelism; + } + + @Override + public int getMaxArrowBatchSize() { + return maxArrowBatchSize; + } + } +} diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/runtime/utils/JavaUserDefinedScalarFunctions.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/runtime/utils/JavaUserDefinedScalarFunctions.java index dfd11550fb3824..ecf151a9f9007b 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/runtime/utils/JavaUserDefinedScalarFunctions.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/runtime/utils/JavaUserDefinedScalarFunctions.java @@ -175,6 +175,21 @@ public PythonEnv getPythonEnv() { } } + /** Test Python scalar function with explicitly configured parallelism. */ + public static class ParallelPythonScalarFunction extends PythonScalarFunction { + private final int parallelism; + + public ParallelPythonScalarFunction(String name, int parallelism) { + super(name); + this.parallelism = parallelism; + } + + @Override + public int getParallelism() { + return parallelism; + } + } + /** Test for Python Scalar Function. */ public static class BooleanPythonScalarFunction extends ScalarFunction implements PythonFunction { diff --git a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.xml b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.xml index 14a1873bf5184f..82b9838ad39d6f 100644 --- a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.xml +++ b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.xml @@ -69,6 +69,16 @@ LogicalProject(a=[$0], EXPR$1=[pyFunc1($1, pyFunc1($2, $3._1))]) FlinkLogicalCalc(select=[a, pyFunc1(b, pyFunc1(c, f0)) AS EXPR$1]) +- FlinkLogicalCalc(select=[a, b, c, d._1 AS f0]) +- FlinkLogicalTableSourceScan(table=[[default_catalog, default_database, MyTable]], fields=[a, b, c, d]) +]]> + + + + + @@ -105,6 +115,15 @@ LogicalProject(a=[$0], b=[$1], EXPR$2=[pyFunc1($0, $2)], EXPR$3=[1]) FlinkLogicalCalc(select=[a, b, f0 AS EXPR$2, 1 AS EXPR$3]) +- FlinkLogicalCalc(select=[a, b, pyFunc1(a, c) AS f0]) +- FlinkLogicalTableSourceScan(table=[[default_catalog, default_database, MyTable]], fields=[a, b, c, d]) +]]> + + + + + @@ -520,6 +539,23 @@ LogicalProject(a=[$0], EXPR$1=[pyFunc1($0, $2)], b=[$1]) FlinkLogicalCalc(select=[a, f0 AS EXPR$1, b]) +- FlinkLogicalCalc(select=[a, b, pyFunc1(a, c) AS f0]) +- FlinkLogicalTableSourceScan(table=[[default_catalog, default_database, MyTable]], fields=[a, b, c, d]) +]]> + + + + + + + + + + + @@ -539,6 +575,24 @@ LogicalProject(a=[$0], EXPR$1=[pyFunc1($0, $2)]) FlinkLogicalCalc(select=[a, f0 AS EXPR$1], where=[>(f0, 0)]) +- FlinkLogicalCalc(select=[a, pyFunc1(a, c) AS f0]) +- FlinkLogicalTableSourceScan(table=[[default_catalog, default_database, MyTable]], fields=[a, b, c, d]) +]]> + + + + + + + + + + + diff --git a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRuleTest.xml b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRuleTest.xml index 3b2b2f472a60fc..77acc3634c858d 100644 --- a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRuleTest.xml +++ b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRuleTest.xml @@ -16,6 +16,26 @@ See the License for the specific language governing permissions and limitations under the License. --> + + + + + + + + + + + + + + + + + + diff --git a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.scala b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.scala index c4f989d244316b..83f9dcc777c30a 100644 --- a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.scala +++ b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRuleTest.scala @@ -58,6 +58,9 @@ class PythonCalcSplitRuleTest extends TableTestBase { util.addTemporarySystemFunction("pyFunc3", new PythonScalarFunction("pyFunc3")) util.addTemporarySystemFunction("pyFunc4", new BooleanPythonScalarFunction("pyFunc4")) util.addTemporarySystemFunction("pyFunc5", new RowPythonScalarFunction("pyFunc5")) + util.addTemporarySystemFunction("pyFuncP1", new ParallelPythonScalarFunction("pyFuncP1", 1)) + util.addTemporarySystemFunction("pyFuncP2", new ParallelPythonScalarFunction("pyFuncP2", 2)) + util.addTemporarySystemFunction("pyFuncP4", new ParallelPythonScalarFunction("pyFuncP4", 4)) util.addTemporarySystemFunction("RowJavaFunc", new RowJavaScalarFunction("RowJavaFunc")) util.addTemporarySystemFunction("pandasFunc1", new PandasScalarFunction("pandasFunc1")) util.addTemporarySystemFunction("pandasFunc2", new PandasScalarFunction("pandasFunc2")) @@ -198,6 +201,30 @@ class PythonCalcSplitRuleTest extends TableTestBase { util.verifyRelPlan(sqlQuery) } + @Test + def testDifferentConcurrencySplitsCalc(): Unit = { + val sqlQuery = "SELECT pyFuncP2(a, b), pyFuncP4(a, c) FROM MyTable" + util.verifyRelPlan(sqlQuery) + } + + @Test + def testSameAndUnspecifiedConcurrencyStayTogether(): Unit = { + val sqlQuery = "SELECT pyFuncP2(a, b), pyFuncP2(a, c), pyFunc1(a, b) FROM MyTable" + util.verifyRelPlan(sqlQuery) + } + + @Test + def testNestedDifferentConcurrencySplitsDeterministically(): Unit = { + val sqlQuery = "SELECT pyFuncP2(a, pyFuncP1(a, c)) FROM MyTable" + util.verifyRelPlan(sqlQuery) + } + + @Test + def testUnspecifiedOuterWithDifferentInnerConcurrencies(): Unit = { + val sqlQuery = "SELECT pyFunc1(pyFuncP1(a, b), pyFuncP2(a, c)) FROM MyTable" + util.verifyRelPlan(sqlQuery) + } + @Test def testPandasFunctionWithCompositeInputs(): Unit = { val sqlQuery = "SELECT a, pandasFunc1(b, d._1) FROM MyTable" diff --git a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRuleTest.scala b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRuleTest.scala index 5212d4a487f332..447981fdba41d8 100644 --- a/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRuleTest.scala +++ b/flink-table/flink-table-planner/src/test/scala/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRuleTest.scala @@ -27,6 +27,15 @@ import org.apache.flink.table.planner.utils.TableTestBase import org.apache.calcite.plan.hep.HepMatchOrder import org.junit.jupiter.api.{BeforeEach, Test} +class ResourceRowPythonScalarFunction( + name: String, + parallelism: Int = -1, + maxArrowBatchSize: Int = -1) + extends RowPythonScalarFunction(name) { + override def getParallelism: Int = parallelism + override def getMaxArrowBatchSize: Int = maxArrowBatchSize +} + /** Test for [[PythonMapMergeRule]]. */ class PythonMapMergeRuleTest extends TableTestBase { private val util = batchTestUtil() @@ -75,6 +84,42 @@ class PythonMapMergeRuleTest extends TableTestBase { util.verifyRelPlan(result) } + @Test + def testCompatibleExecutionOptionsAreMerged(): Unit = { + val sourceTable = util.addTableSource[(Int, Int, Int)]("source", 'a, 'b, 'c) + val explicit = new ResourceRowPythonScalarFunction("explicit", 2, 64) + val unspecified = new ResourceRowPythonScalarFunction("unspecified") + val result = sourceTable + .map(explicit(withColumns('*))) + .map(unspecified(withColumns('*))) + .map(explicit(withColumns('*))) + util.verifyRelPlan(result) + } + + @Test + def testNestedConflictingConcurrencyIsNotMerged(): Unit = { + val sourceTable = util.addTableSource[(Int, Int, Int)]("source", 'a, 'b, 'c) + val parallelism2 = new ResourceRowPythonScalarFunction("parallelism2", parallelism = 2) + val unspecified = new ResourceRowPythonScalarFunction("unspecified") + val parallelism4 = new ResourceRowPythonScalarFunction("parallelism4", parallelism = 4) + val result = sourceTable + .map(parallelism2(withColumns('*))) + .map(unspecified(withColumns('*))) + .map(parallelism4(withColumns('*))) + util.verifyRelPlan(result) + } + + @Test + def testConflictingBatchSizesAreNotMerged(): Unit = { + val sourceTable = util.addTableSource[(Int, Int, Int)]("source", 'a, 'b, 'c) + val batch64 = new ResourceRowPythonScalarFunction("batch64", maxArrowBatchSize = 64) + val batch128 = new ResourceRowPythonScalarFunction("batch128", maxArrowBatchSize = 128) + val result = sourceTable + .map(batch64(withColumns('*))) + .map(batch128(withColumns('*))) + util.verifyRelPlan(result) + } + @Test def testProjectWithOneField(): Unit = { val sourceTable = util.addTableSource[(Int, Int, Int)]("source", 'a, 'b, 'c)