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..fbc279620b1766 --- /dev/null +++ b/flink-python/docs/reference/pyflink.dataframe/udf.rst @@ -0,0 +1,47 @@ +.. ################################################################################ + 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. + +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 +============= + +.. currentmodule:: pyflink.dataframe + +.. autosummary:: + :toctree: api/ + + udf 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/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index ac4c12d24d533e..4e199b0e34beed 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -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 new file mode 100644 index 00000000000000..80bd955e6f7b3e --- /dev/null +++ b/flink-python/pyflink/dataframe/tests/test_udf.py @@ -0,0 +1,1354 @@ +################################################################################ +# 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 asyncio +import functools +import importlib +import inspect +import operator +import unittest +from dataclasses import dataclass +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, + udf as table_udf, +) +from pyflink.testing.test_case_utils import ( + PyFlinkDataFrameUTTestCase, + PyFlinkStreamDataFrameTestCase, +) + + +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): + 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]}, + } + + 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: Callable[..., Expression] = pf.udf(add_one) + + self.assertFalse(hasattr(pf, "DataFrameUDFWrapper")) + 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: 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(_return_dtype(configured), pf.DataType.string()) + 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", + lambda: pf.udf(identity, return_dtype=int), + pf.DataType.int64(), + ), + ( + "nested TypedDict annotation", + lambda: pf.udf(describe), + expected_result_dtype, + ), + ( + "explicit nested TypedDict", + lambda: pf.udf(identity, return_dtype=Result), + expected_result_dtype, + ), + ( + "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): + self.assertEqual(_return_dtype(declare()), 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, *values: int) -> int: + value, = values + return value * 2 + + class AsyncDouble(AsyncScalarFunction): + def __init__(self): + scalar_constructor_calls.append("AsyncDouble") + + 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, *values: int) -> int: + value, = values + return value + self.offset + + 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, + AddScalarOffset(2), + ] + for source in callables: + with self.subTest(source=source): + decorated = pf.udf(source) + self.assertEqual(_return_dtype(decorated), pf.DataType.int64()) + + self.assertEqual(plain_constructor_calls, []) + self.assertEqual(scalar_constructor_calls, []) + 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_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(_return_dtype(uninspectable), pf.DataType.int64()) + exploding_signature = pf.udf(ExplodingSignature(), return_dtype=int) + 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: + 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 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 + + 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 + + class PandasCallable: + def __call__(self, values: pd.Series) -> pd.Series: + return values + 1 + + class PandasScalarFunction(ScalarFunction): + def eval(self, *values: pd.Series) -> pd.Series: + value, = values + return value + 1 + + class AsyncScalarClass(AsyncScalarFunction): + async def eval(self, *values: int) -> int: + value, = values + return value + 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, + ), + ( + "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()), + "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, + ), + ( + "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): + wrapped = declare() + self.assertEqual(wrapped._func_type, expected_type) + self.assertEqual(wrapped._source.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_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: + value, = values + return value + + def is_deterministic(self): + return False + + class DefaultDeterministic(ScalarFunction): + def eval(self, *values: int) -> int: + value, = values + 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) + self.assertTrue(pf.udf(NonDeterministic)._deterministic) + + 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 _create_result_normalizer + + 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 + + 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 + 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) + 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 = [ + ( + "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", + named_row, + expected_named_row, + ), + ( + "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}, + ), + ), + ( + "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( + result_normalizer(value), expected + ) + with self.assertRaisesRegex(ValueError, "Expected 3 value"): + result_normalizer((1, 2)) + with self.assertRaisesRegex(TypeError, "Expected a Mapping"): + 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 + + class RequiresArgument: + def __init__(self, value): + self.value = value + + def __call__(self, other: int) -> int: + return other + self.value + + class RequiresScalarArgument(ScalarFunction): + def __init__(self, value): + self.value = 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, *values: int) -> int: + value, = values + return value + self.value + + class NotCallable: + pass + + class NonScalarFunction(TableFunction): + def eval(self, value): + return value + + 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", + ), + ( + "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), + TypeError, + "Cannot infer return_dtype", + ), + ( + "unresolved return", + lambda: pf.udf(unresolved_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", + ), + ( + "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( + 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_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, + _DataFrameScalarFunctionAdapter, + _UDFUsage, + _resolve_udf_source, + ) + + 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") + + 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 = create_adapter(LifecycleFunction) + self.assertEqual(events, []) + + 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", + "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: + def __init__(self): + deferred_constructor_calls.append("init") + + def __call__(self, value): + return value + 1 + + deferred_adapter = create_adapter(DeferredCallable) + 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 = 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_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") + + @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(concurrency=1) + def add_one(value: int) -> int: + return value + 1 + + @pf.udf(concurrency=1) + async def add_two(value: int) -> int: + return value + 2 + + @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 + + @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, *values: int) -> int: + value, = values + return value + self._increment + + class ClassNonDeterministic(ScalarFunction): + def eval(self, *values: int) -> int: + value, = values + return value + 6 + + def is_deterministic(self): + return False + + deferred = pf.udf(DeferredCallable) + opened_scalar_class = pf.udf(OpenedScalarFunction) + scalar_class = pf.udf(ClassNonDeterministic, deterministic=False) + + result = ( + pf.from_records([(1,)], schema=["id"]) + .with_columns(async_value=add_two(pf.col("id"))) + .with_columns( + sync_value=add_one(pf.col("id")), + pandas_value=add_three(pf.col("id")), + details=details(pf.col("id")), + deferred_value=deferred(pf.col("id")), + scalar_value=opened_scalar_class(pf.col("id")), + scalar_class_value=scalar_class(pf.col("id")), + ) + ) + + self.assertEqual( + result.collect(), + [Row(1, 3, 2, 4, Row(2, ["1"]), 5, 6, 7)], + ) + + +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..389380539ded53 --- /dev/null +++ b/flink-python/pyflink/dataframe/udf.py @@ -0,0 +1,1142 @@ +################################################################################ +# 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 dataclasses import dataclass +from enum import Enum +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Iterable, + List, + Optional, + 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, + _validate_udf_execution_options, + udf as table_udf, +) +from pyflink.util.api_stability_decorators import PublicEvolving + +__all__ = ["udf"] + +_UDFInput = Union[Callable[..., Any], ScalarFunction, AsyncScalarFunction, Type] +_DataTypeLike = Union[DataType, Type, str] +_UNRESOLVED_TYPE_HINT = object() + + +class _UDFUsage(Enum): + EXPRESSION = "expression" + MAP = "map" + 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_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.""" + + source: _UDFInput + kind: _UDFSourceKind + is_async: bool + ignored_hint_names: FrozenSet[str] = frozenset() + + def _inspection_target_and_skip_first( + self, + ) -> Tuple[Callable[..., Any], bool]: + if self.kind is _UDFSourceKind.DIRECT_CALLABLE: + return ( + _get_callable_inspection_target( + cast(Callable[..., Any], self.source) + ), + False, + ) + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: + return ( + cast( + Union[ScalarFunction, AsyncScalarFunction], self.source + ).eval, + False, + ) + if self.kind in ( + _UDFSourceKind.CALLABLE_CLASS, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ): + hint_method, skip_first_parameter = _get_callable_class_hint_method( + cast(Type, self.source), + "eval" + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS + else "__call__", + ) + if hint_method is None: + raise RuntimeError("Resolved UDF class has no inspection target.") + return hint_method, skip_first_parameter + return cast(Callable[..., Any], getattr(self.source, "__call__")), False + + @property + def inspection_target(self) -> Callable[..., Any]: + target, _ = self._inspection_target_and_skip_first() + return target + + @property + def invocation_signature(self) -> Optional[inspect.Signature]: + target, skip_first_parameter = self._inspection_target_and_skip_first() + if not self.is_scalar_function and not self.constructs_on_worker: + target = cast(Callable[..., Any], self.source) + + 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: + 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 in ( + _UDFSourceKind.CALLABLE_CLASS, + _UDFSourceKind.SCALAR_FUNCTION_CLASS, + ) + + def create_worker_source(self) -> _UDFInput: + if not self.constructs_on_worker: + return self.source + source_class = cast(Type, self.source) + source = source_class() + 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__}'." + ) + return cast(_UDFInput, source) + + def validate_deterministic( + self, declared: bool, worker_source: Optional[_UDFInput] = None + ) -> None: + source: Optional[_UDFInput] + if self.kind is _UDFSourceKind.SCALAR_FUNCTION_INSTANCE: + source = self.source + elif self.kind is _UDFSourceKind.SCALAR_FUNCTION_CLASS: + source = worker_source + else: + source = None + if source is not None: + _validate_deterministic( + declared, + cast( + Union[ScalarFunction, AsyncScalarFunction], source + ).is_deterministic(), + ) + + def open_worker_source( + self, worker_source: _UDFInput, function_context: Any + ) -> None: + if self.is_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.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.is_scalar_function and worker_source is not None: + cast( + Union[ScalarFunction, AsyncScalarFunction], worker_source + ).close() + + +class _DataFrameUDFWrapper: + """Internal callable binding a DataFrame scalar UDF to Table expressions.""" + + _source: _ResolvedUDFSource + _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 + + def __init__( + self, + source: _ResolvedUDFSource, + return_dtype: DataType, + 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) + 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: + if getattr(self, "_frozen", False): + raise AttributeError("DataFrame UDF declarations are immutable.") + object.__setattr__(self, name, value) + + def __call__(self, *args: Any) -> Expression: + 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._source.is_async + else _DataFrameScalarFunctionAdapter + ) + actual_func = cast( + Union[ScalarFunction, AsyncScalarFunction], + adapter_type( + self._source, + 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, + concurrency=self._concurrency, + batch_size=self._batch_size, + ), + ) + + @property + def return_dtype(self) -> DataType: + return self._return_dtype + + +@overload +def udf( + func: _UDFInput, + *, + return_dtype: Optional[_DataTypeLike] = ..., + deterministic: bool = ..., + name: Optional[str] = ..., + func_type: Optional[str] = ..., + concurrency: Optional[int] = ..., + batch_size: Optional[int] = ..., +) -> Callable[..., Expression]: + ... + + +@overload +def udf( + func: None = ..., + *, + return_dtype: Optional[_DataTypeLike] = ..., + deterministic: bool = ..., + name: Optional[str] = ..., + func_type: Optional[str] = ..., + concurrency: Optional[int] = ..., + batch_size: Optional[int] = ..., +) -> Callable[[_UDFInput], Callable[..., Expression]]: + ... + + +@PublicEvolving() +def udf( + func: Optional[_UDFInput] = None, + *, + return_dtype: Optional[_DataTypeLike] = None, + deterministic: bool = True, + name: Optional[str] = None, + func_type: Optional[str] = None, + concurrency: Optional[int] = None, + batch_size: Optional[int] = None, +) -> Union[ + Callable[..., Expression], + Callable[[_UDFInput], Callable[..., Expression]], +]: + """ + Create a scalar UDF for DataFrame expressions. + + 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 + initialized on the TaskManager, so expensive initialization is deferred:: + + >>> 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 initialized on the + TaskManager, where their ``open`` and ``close`` methods also run:: + + >>> 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 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 + + >>> @pf.udf( + ... return_dtype=pf.DataType.int64(), + ... func_type="pandas", + ... batch_size=256, + ... ) + ... 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. + :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. + + .. versionadded:: 2.4.0 + """ + + def decorator(f: _UDFInput) -> Callable[..., Expression]: + source = _resolve_udf_source(f) + actual_func_type = ( + func_type + if func_type is not None + else _detect_func_type(source) + ) + _validate_scalar_udf_options( + actual_func_type, + return_dtype, + source.is_async, + concurrency, + batch_size, + ) + 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( + source, + actual_return_dtype, + actual_deterministic, + actual_name, + actual_func_type, + concurrency, + batch_size, + ) + + return decorator if func is None else decorator(func) + + +# ======================== Declaration Validation ======================== + + +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( + 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'." + ) + _validate_udf_execution_options(func_type, concurrency, batch_size) + + +# ======================== 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 _resolve_udf_source( + func: _UDFInput, +) -> _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, + _UDFSourceKind.DIRECT_CALLABLE, + inspect.iscoroutinefunction(inspection_target), + _ignored_hint_names(func, inspection_target), + ) + + 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 issubclass(func, (ScalarFunction, AsyncScalarFunction)): + _validate_zero_argument_class(func) + 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, + _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): + 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." + ) + return _ResolvedUDFSource( + 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 + ) + return _ResolvedUDFSource( + func, + _UDFSourceKind.SCALAR_FUNCTION_INSTANCE, + 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, + _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: + 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) + + 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(return_hint) + + +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 _data_type_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(source: _ResolvedUDFSource) -> str: + """Detect pandas mode from an unbound pandas container annotation.""" + hint_func = source.inspection_target + try: + import pandas as pd + except ImportError: + return "general" + + pandas_types = (pd.Series, pd.DataFrame) + 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}, + ) + if hint in pandas_types: + return "pandas" + return "general" + + +def _unwrap_partial(func: Any) -> Any: + while isinstance(func, functools.partial): + func = func.func + return func + + +def _get_callable_inspection_target( + func: Callable[..., Any], +) -> Callable[..., Any]: + target = _unwrap_partial(func) + if callable(target) and not inspect.isroutine(target) and not inspect.isclass(target): + return cast(Callable[..., Any], getattr(target, "__call__")) + return cast(Callable[..., Any], target) + + +def _get_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 _resolve_callable_annotation( + func: Callable[..., Any], + annotation_name: str, + fallback_globals: Optional[Dict[str, Any]] = None, +) -> Any: + annotations = getattr(func, "__annotations__", {}) + if annotation_name not in annotations: + return _UNRESOLVED_TYPE_HINT + + def annotation_holder() -> None: + pass + + annotation_holder.__annotations__ = { + annotation_name: annotations[annotation_name] + } + try: + return get_type_hints( + annotation_holder, + globalns={ + **(fallback_globals or {}), + **_get_callable_globals(func), + }, + ).get(annotation_name, _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: + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool.") + source.validate_deterministic(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(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: + 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], + result_normalizer: Callable[[Any], Any], + is_async: bool, +) -> Callable[..., Any]: + if is_async: + + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + return result_normalizer(await func(*args, **kwargs)) + + wrapper = async_wrapper + else: + + @functools.wraps(func) + def sync_wrapper(*args: Any, **kwargs: Any) -> Any: + return result_normalizer(func(*args, **kwargs)) + + 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, + source: _ResolvedUDFSource, + return_dtype: DataType, + deterministic: bool, + usage: _UDFUsage, + func_type: str, + ) -> None: + self._source = source + self._func: Optional[_UDFInput] = ( + 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.source, "__doc__", None) + + def open(self, function_context: Any) -> None: + lifecycle_opened = False + try: + if self._source.constructs_on_worker: + self._func = self._source.create_worker_source() + + 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: + raise NotImplementedError( + 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, + result_normalizer, + self._source.is_async, + ) + return invoke_func + + def close(self) -> None: + func = self._func + try: + 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 + + 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_field_values(value: Any, field_names: List[str]) -> List[Any]: + if isinstance(value, Mapping): + return [value.get(field_name) for field_name in field_names] + if isinstance(value, Row) and hasattr(value, "_fields"): + 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 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: + 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"{field_names}, got {type(value).__name__}." + ) from None + + +def _create_result_normalizer( + data_type: Any, +) -> Optional[Callable[[Any], Any]]: + if isinstance(data_type, RowType): + field_names = data_type.field_names() + field_normalizers = tuple( + _create_result_normalizer(field.data_type) for field in data_type + ) + + def normalize_row(value: Any) -> Any: + if value is None: + return None + 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): + row.set_row_kind(value.get_row_kind()) + return row + + return normalize_row + if isinstance(data_type, ArrayType): + 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): + 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(cast(Iterable[Any], 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__}." + ) + 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 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)