From 33b6793a3ff71a0058fdfecf09cc125d016bd808 Mon Sep 17 00:00:00 2001 From: victor Date: Fri, 18 Sep 2026 12:32:24 +0800 Subject: [PATCH 1/2] [FLINK-40434][python] Add flat_map support to the DataFrame API Co-Authored-By: Codex AI-Model: gpt-6-astra AI-Contributed/Feature: 0/0 AI-Contributed/UT: 110/110 --- .../reference/pyflink.dataframe/dataframe.rst | 1 + .../docs/reference/pyflink.dataframe/udf.rst | 22 +- flink-python/pyflink/dataframe/__init__.py | 2 + flink-python/pyflink/dataframe/dataframe.py | 53 ++ .../pyflink/dataframe/tests/test_udtf.py | 568 ++++++++++++++++++ flink-python/pyflink/dataframe/udtf.py | 447 ++++++++++++++ 6 files changed, 1090 insertions(+), 3 deletions(-) create mode 100644 flink-python/pyflink/dataframe/tests/test_udtf.py create mode 100644 flink-python/pyflink/dataframe/udtf.py diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst index e78bbaf0d93c5..57850d6fa579a 100644 --- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst +++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst @@ -68,6 +68,7 @@ Transformations DataFrame.limit DataFrame.offset DataFrame.head + DataFrame.flat_map DataFrame.__getitem__ Set Operations diff --git a/flink-python/docs/reference/pyflink.dataframe/udf.rst b/flink-python/docs/reference/pyflink.dataframe/udf.rst index 73d1ff2abe76d..d2e803c347d42 100644 --- a/flink-python/docs/reference/pyflink.dataframe/udf.rst +++ b/flink-python/docs/reference/pyflink.dataframe/udf.rst @@ -16,9 +16,12 @@ limitations under the License. ################################################################################ -============================= -User-Defined Scalar Functions -============================= +====================== +User-Defined Functions +====================== + +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 @@ -30,6 +33,18 @@ DataFrame scalar UDFs support synchronous, asynchronous, and pandas-vectorized callables. See :func:`pyflink.dataframe.udf` for declaration forms, type inference, execution modes, and examples. +Table Functions +=============== + +Use :func:`pyflink.dataframe.udtf` to declare a Python function that emits zero or +more rows per invocation. A table UDF can be used with +:meth:`~pyflink.dataframe.DataFrame.flat_map`, which returns only the emitted columns. + +DataFrame table UDFs support synchronous functions, callable instances, and +``TableFunction`` instances or classes. See :func:`pyflink.dataframe.udtf` for +declaration forms, type inference, lifecycle, and examples, and +:meth:`~pyflink.dataframe.DataFrame.flat_map` for row input and output semantics. + API Reference ============= @@ -39,3 +54,4 @@ API Reference :toctree: api/ udf + udtf diff --git a/flink-python/pyflink/dataframe/__init__.py b/flink-python/pyflink/dataframe/__init__.py index c8ed1cf08ae5f..bce542f48c227 100644 --- a/flink-python/pyflink/dataframe/__init__.py +++ b/flink-python/pyflink/dataframe/__init__.py @@ -57,6 +57,7 @@ from pyflink.dataframe.io import read_generic from pyflink.dataframe.sql import sql from pyflink.dataframe.udf import udf +from pyflink.dataframe.udtf import udtf __all__ = [ "DataFrame", @@ -65,6 +66,7 @@ "col", "lit", "udf", + "udtf", "from_arrow", "from_dict", "from_pandas", diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 2e1b8fd8208e9..49ca24dd3e146 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -33,6 +33,8 @@ if TYPE_CHECKING: import pandas + from pyflink.dataframe.udf import _DataTypeLike + from pyflink.dataframe.udtf import _DataFrameUDTFWrapper from pyflink.table.table_schema import TableSchema from pyflink.common import Row @@ -134,6 +136,57 @@ class DataFrame: def __init__(self, table: Table): self._table = table + @PublicEvolving() + def flat_map( + self, + func: Union[Callable[[Dict[str, Any]], Any], "_DataFrameUDTFWrapper"], + *, + return_dtype: Optional["_DataTypeLike"] = None, + ) -> "DataFrame": + """ + Apply a function to each row, emitting zero or more output rows. + + A plain callable receives a dictionary keyed by column name. A declaration + created with :func:`pyflink.dataframe.udtf` receives a named Flink ``Row``. + Output column names come from a ``TypedDict`` or an explicit named struct; + scalar outputs use ``f0``. Multi-field outputs require named fields. + + :param func: Row-based callable or a declaration created with ``pf.udtf``. + :param return_dtype: Emitted row type, inferred from annotations when omitted. + Required if inference is not possible; must be omitted + for a UDTF declaration. + :return: A DataFrame containing only the emitted output columns. + + Example:: + + >>> from typing import Any, Dict, Iterator, TypedDict + >>> import pyflink.dataframe as pf + >>> class Token(TypedDict): + ... word: str + >>> def split(row: Dict[str, Any]) -> Iterator[Token]: + ... for word in row["text"].split(): + ... yield {"word": word} + >>> df = pf.from_dict({"text": ["hello world", "flink"]}) + >>> result = df.flat_map(split) + + A reusable UDTF declaration receives a named ``Row``:: + + >>> from pyflink.common import Row + >>> @pf.udtf + ... def split_row(row: Row) -> Iterator[Token]: + ... for word in row["text"].split(): + ... yield {"word": word} + >>> result = df.flat_map(split_row) + + .. versionadded:: 2.4.0 + """ + from pyflink.dataframe.udtf import _resolve_flat_map_udtf + + expression, output_columns = _resolve_flat_map_udtf(func, return_dtype, self.columns) + table = self._table.flat_map(expression) + # Table UDTFs expose positional field names, so restore the declared names. + return DataFrame(table.alias(output_columns[0], *output_columns[1:])) + # ======================== Core Operations ======================== @PublicEvolving() diff --git a/flink-python/pyflink/dataframe/tests/test_udtf.py b/flink-python/pyflink/dataframe/tests/test_udtf.py new file mode 100644 index 0000000000000..44219a9e77e9a --- /dev/null +++ b/flink-python/pyflink/dataframe/tests/test_udtf.py @@ -0,0 +1,568 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################ + +import functools +import os +import threading +import unittest +from dataclasses import replace +from typing import ( + Annotated, Any, Dict, Generator, Iterable, Iterator, List, Mapping, MutableMapping, + Optional, Tuple, TypedDict, Union, +) + +import cloudpickle +import pyflink.dataframe as pf +from pyflink.common import Row, RowKind +from pyflink.dataframe.udtf import _resolve_flat_map_udtf +from pyflink.table import DataTypes +from pyflink.table.udf import ScalarFunction, TableFunction +from pyflink.testing.test_case_utils import ( + PyFlinkBatchTableTestCase, + PyFlinkDataFrameUTTestCase, + PyFlinkStreamDataFrameTestCase, +) + + +class _Output(TypedDict): + value: int + label: str + + +def _eval_udtf(declaration, *args): + adapter = declaration._create_table_wrapper()._func + adapter.open(None) + try: + return list(adapter.eval(*args)) + finally: + adapter.close() + + +class DataFrameUDTFDeclarationTests(unittest.TestCase): + def test_decorator_forms_and_type_inference(self): + def emit(value) -> Iterator[_Output]: + yield {"value": value, "label": "ok"} + + expected = pf.DataType.struct({"value": pf.DataType.int64(), "label": pf.DataType.string()}) + for declaration in (pf.udtf(emit), pf.udtf()(emit), pf.udtf(return_dtype=expected)(emit), + pf.udtf(emit, return_dtype=_Output)): + self.assertEqual(declaration.return_dtype, expected) + self.assertEqual(declaration.__name__, "emit") + + def test_iterable_annotations_describe_one_emitted_item(self): + hints = (Iterator[int], Iterable[int], Generator[int, None, None], List[int], list[int]) + for hint in hints: + def emit(value): + return [value] + emit.__annotations__ = {"return": hint} + with self.subTest(hint=hint): + self.assertEqual(pf.udtf(emit).return_dtype, pf.DataType.int64()) + + def test_explicit_type_does_not_resolve_annotations(self): + def emit(value): + return [value] + emit.__annotations__ = {"value": "UnavailableInput", "return": "UnavailableOutput"} + self.assertEqual(pf.udtf(emit, return_dtype=int).return_dtype, pf.DataType.int64()) + emit.__annotations__["return"] = Iterator[int] + self.assertEqual(pf.udtf(emit).return_dtype, pf.DataType.int64()) + + def test_callable_instance_and_partial(self): + class Repeat: + def __init__(self, count): + self.count = count + + def __call__(self, value) -> Iterator[int]: + yield from [value] * self.count + + def repeat(count, value) -> Iterator[int]: + yield from [value] * count + + for func in (Repeat(2), functools.partial(repeat, 2)): + self.assertEqual(_eval_udtf(pf.udtf(func), 3), [Row(3), Row(3)]) + + def test_invalid_declarations(self): + class Scalar(ScalarFunction): + def eval(self, value): + return value + + class CallableClass: + def __call__(self, value): + return [value] + + class RequiresArgument(TableFunction): + def __init__(self, value): + self.value = value + + def eval(self, value): + return [value] + + for func in (123, Scalar(), Scalar, CallableClass, RequiresArgument): + with self.subTest(func=func), self.assertRaises(TypeError): + pf.udtf(func, return_dtype=int) + for hint in (Iterator, Iterator[Tuple[int, ...]], Iterator[Tuple]): + def emit(row): + return [] + emit.__annotations__ = {"return": hint} + with self.subTest(hint=hint), self.assertRaises(TypeError): + pf.udtf(emit) + with self.assertRaisesRegex(TypeError, "return_dtype"): + pf.udtf(lambda row: []) + with self.assertRaisesRegex(ValueError, "at least one"): + pf.udtf(lambda row: [], return_dtype=pf.DataType.struct({})) + + def test_async_functions_are_rejected(self): + async def coroutine(row): + return [row] + + async def generator(row): + yield row + + class AsyncCallable: + async def __call__(self, row): + yield row + + @functools.wraps(generator) + def hidden_generator(row): + return generator(row) + + for func in (coroutine, generator, AsyncCallable(), hidden_generator): + with self.subTest(func=func), self.assertRaisesRegex(TypeError, "async"): + pf.udtf(func, return_dtype=int) + + def test_names_follow_table_udtf_defaults(self): + def emit(row) -> Iterator[int]: + yield row[0] + + class Expand(TableFunction): + def eval(self, row) -> Iterator[int]: + yield row[0] + + class Repeat: + def __call__(self, row) -> Iterator[int]: + yield row[0] + + cases = ( + (emit, None, "emit"), (Expand, None, "Expand"), (Expand(), None, "Expand"), + (Repeat(), None, "Repeat"), (emit, "", "emit"), (emit, "custom_name", "custom_name"), + ) + for func, name, expected in cases: + with self.subTest(func=func, name=name): + declaration = pf.udtf(func, name=name) + self.assertEqual(declaration.__name__, expected) + self.assertEqual(declaration._create_table_wrapper()._name, expected) + + def test_determinism_validation(self): + with self.assertRaises(TypeError): + pf.udtf(lambda row: [], return_dtype=int, deterministic=1) + + class Random(TableFunction): + def eval(self, value) -> Iterator[int]: + yield value + + def is_deterministic(self): + return False + + with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"): + pf.udtf(Random()) + adapter = pf.udtf(Random, deterministic=False, name="random")._create_table_wrapper()._func + adapter.open(None) + try: + self.assertFalse(adapter.is_deterministic()) + self.assertEqual(list(adapter.eval(1)), [Row(1)]) + finally: + adapter.close() + with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"): + pf.udtf(Random) + + def test_table_function_class_constructs_on_client(self): + events = [] + + class Expand(TableFunction): + def __init__(self): + events.append("init") + + def open(self, context): + events.append(("open", context)) + + def eval(self, value) -> Iterator[int]: + yield value + + def close(self): + events.append("close") + + declaration = pf.udtf(Expand) + self.assertEqual(events, ["init"]) + adapter = declaration._create_table_wrapper()._func + self.assertEqual(events, ["init"]) + adapter.open("worker") + self.assertEqual(list(adapter.eval(7)), [Row(7)]) + adapter.close() + self.assertEqual(events, ["init", ("open", "worker"), "close"]) + with self.assertRaisesRegex(RuntimeError, "before open"): + adapter.eval(7) + + def test_adapter_serializes_only_worker_metadata(self): + def emit(value) -> Iterator[int]: + yield value + + declaration = pf.udtf(emit) + context = replace(declaration._declaration_context, + localns={"client_lock": threading.Lock()}) + declaration = replace(declaration, _declaration_context=context) + adapter = cloudpickle.loads(cloudpickle.dumps(declaration._create_table_wrapper()._func)) + adapter.open(None) + try: + self.assertEqual(list(adapter.eval(5)), [Row(5)]) + finally: + adapter.close() + + def test_table_function_constructor_failure_is_reported_on_client(self): + class FailingConstructor(TableFunction): + def __init__(self): + raise RuntimeError("constructor failed") + + def eval(self, value) -> Iterator[int]: + yield value + + with self.assertRaisesRegex(RuntimeError, "constructor failed"): + pf.udtf(FailingConstructor) + + def test_table_function_close_failure_resets_adapter(self): + class FailingClose(TableFunction): + def eval(self, value) -> Iterator[int]: + yield value + + def close(self): + raise RuntimeError("close failed") + + adapter = pf.udtf(FailingClose)._create_table_wrapper()._func + adapter.open(None) + self.assertEqual(list(adapter.eval(1)), [Row(1)]) + with self.assertRaisesRegex(RuntimeError, "close failed"): + adapter.close() + with self.assertRaisesRegex(RuntimeError, "before open"): + adapter.eval(1) + adapter.close() + + def test_output_cardinality_and_lazy_iteration(self): + events = [] + + def emit(value): + events.append(value) + yield value + events.append(value + 1) + yield value + 1 + + adapter = pf.udtf(emit, return_dtype=int)._create_table_wrapper()._func + adapter.open(None) + try: + output = adapter.eval(2) + self.assertEqual(events, []) + self.assertEqual(next(output), Row(2)) + self.assertEqual(events, [2]) + self.assertEqual(list(output), [Row(3)]) + finally: + adapter.close() + + cases = [(None, []), ([], []), (3, [Row(3)]), ((3,), [Row(3)]), + (Row(3), [Row(3)]), ([None, 3], [Row(None), Row(3)])] + for result, expected in cases: + self.assertEqual(_eval_udtf(pf.udtf(lambda: result, return_dtype=int)), expected) + + def test_struct_and_map_results(self): + dtype = pf.DataType.struct({ + "value": pf.DataType.int64(), + "nested": pf.DataType.list(pf.DataType.struct({"label": pf.DataType.string()})), + "attributes": pf.DataType.map(pf.DataType.string(), pf.DataType.int64()), + }) + result = {"attributes": {"x": 1}, "nested": [{"label": "ok"}], "value": 7} + self.assertEqual(_eval_udtf(pf.udtf(lambda: result, return_dtype=dtype)), [ + Row(value=7, nested=[Row(label="ok")], attributes={"x": 1})]) + map_type = pf.DataType.map(pf.DataType.string(), pf.DataType.int64()) + for result in ({"x": 1}, Row({"x": 1}), ({"x": 1},)): + self.assertEqual(_eval_udtf(pf.udtf(lambda: result, return_dtype=map_type)), + [Row({"x": 1})]) + + def test_runtime_errors_are_propagated(self): + dtype = pf.DataType.struct({"a": pf.DataType.int64(), "b": pf.DataType.int64()}) + for result, return_dtype in (((1,), dtype), ((1, 2), pf.DataType.int64())): + declaration = pf.udtf(lambda: result, return_dtype=return_dtype) + with self.assertRaises(ValueError): + _eval_udtf(declaration) + + def fail(row): + raise RuntimeError("user failure") + with self.assertRaisesRegex(RuntimeError, "user failure"): + _eval_udtf(pf.udtf(fail, return_dtype=int), Row(1)) + + def test_single_named_field_accepts_scalar_items(self): + cases = [ + (pf.DataType.string(), ["hello", "world"], [Row(word="hello"), Row(word="world")]), + (pf.DataType.list(pf.DataType.int64()), [[1, 2], []], [Row(word=[1, 2]), Row(word=[])]), + ] + for field_type, result, expected in cases: + dtype = pf.DataType.struct({"word": field_type}) + self.assertEqual(_eval_udtf(pf.udtf(lambda: result, return_dtype=dtype)), expected) + + def test_flat_map_validates_row_input_before_building_expression(self): + def named(row: Dict[str, Any], optional=0) -> Iterator[int]: + yield row["x"] + optional + + def column(value: int) -> Iterator[int]: + yield value + + def multiple(left, right) -> Iterator[int]: + yield left + right + + def keyword(row, *, required) -> Iterator[int]: + yield required + + for func in (column, multiple, keyword, pf.udtf(column), pf.udtf(multiple), pf.udtf(named)): + with self.subTest(func=func), self.assertRaisesRegex(ValueError, "row argument"): + _resolve_flat_map_udtf(func, None, ["x"]) + with self.assertRaisesRegex(ValueError, "return_dtype"): + _resolve_flat_map_udtf(pf.udtf(column), int, ["x"]) + with self.assertRaises(TypeError): + _resolve_flat_map_udtf(None, int, ["x"]) + + def test_incompatible_input_annotations(self): + for hint in (Optional[int], Union[int, str], Annotated[int, "column value"]): + def expand(row) -> Iterator[int]: + yield row + expand.__annotations__["row"] = hint + for func in (expand, pf.udtf(expand)): + with self.subTest(hint=hint, func=func): + with self.assertRaisesRegex(ValueError, "row argument"): + _resolve_flat_map_udtf(func, None, ["x"]) + + def test_flat_map_names_row_without_mutating_input(self): + received = [] + + @pf.udtf + def expand(row: Row) -> Iterator[int]: + received.append(row) + yield row["x"] + + adapter = expand._create_table_wrapper(("z", "x"))._func + adapter.open(None) + self.addCleanup(adapter.close) + original = Row.of_kind(RowKind.DELETE, original_z=7, original_x=3) + for value, kind in (((7, 3), RowKind.INSERT), (original, RowKind.DELETE)): + self.assertEqual(list(adapter.eval(value)), [Row(3)]) + self.assertIsInstance(received[-1], Row) + self.assertEqual(received[-1].as_dict(), {"z": 7, "x": 3}) + self.assertEqual(received[-1].get_row_kind(), kind) + self.assertEqual(original.as_dict(), {"original_z": 7, "original_x": 3}) + + +class DataFrameUDTFPlanningTests(PyFlinkDataFrameUTTestCase): + def test_compatible_input_annotations(self): + class Input(TypedDict): + x: int + + source = pf.from_dict({"x": [1]}) + for hint in (Input, Optional[Input], Optional[Dict[str, Any]], + Union[int, Dict[str, Any]], Mapping[str, Any], MutableMapping[str, Any], + Annotated[Dict[str, Any], "input row"]): + def expand(row) -> Iterator[int]: + yield row["x"] + expand.__annotations__["row"] = hint + with self.subTest(hint=hint): + self.assertEqual(source.flat_map(expand).columns, ["f0"]) + + @pf.udtf + def expand_row(row: Optional[Row]) -> Iterator[int]: + yield row["x"] + + self.assertEqual(source.flat_map(expand_row).columns, ["f0"]) + + def test_unnamed_multiple_fields_require_schema(self): + def emit(row) -> Iterator[Tuple[int, str]]: + yield row["x"], "ok" + + source = pf.from_dict({"x": [1]}) + with self.assertRaisesRegex(ValueError, "named output fields"): + source.flat_map(emit) + named = source.flat_map(emit, return_dtype="ROW") + self.assertEqual(named.columns, ["value", "label"]) + + def test_call_metadata_aliases_and_shared_declaration(self): + @pf.udtf + def emit(value) -> Iterator[_Output]: + yield {"value": value, "label": "ok"} + + source = pf.from_dict({"x": [1]}) + call = emit(pf.col("x")) + aliased = call.alias("v", "l") + self.assertIsNone(call.output_aliases) + self.assertEqual(aliased.output_aliases, ("v", "l")) + for names in (("one",), ("same", "same"), ("", "other")): + with self.subTest(names=names), self.assertRaises(ValueError): + call.alias(*names) + table = source.to_table().join_lateral(aliased.expression) + self.assertEqual(pf.from_table(table).columns, ["x", "v", "l"]) + + def test_constructing_plan_does_not_reconstruct_user_class(self): + events = [] + + class Expand(TableFunction): + def __init__(self): + events.append("init") + + def eval(self, row: Row) -> Iterator[int]: + yield row[0] + + declaration = pf.udtf(Expand) + self.assertEqual(events, ["init"]) + for name in ("x", "y"): + result = pf.from_dict({name: [1]}).flat_map(declaration) + self.assertEqual(result.columns, ["f0"]) + self.assertEqual(events, ["init"]) + + +class _DataFrameFlatMapTests: + def setUp(self): + super().setUp() + previous = pf.get_table_environment() + self.addCleanup(pf.set_table_environment, previous) + pf.set_table_environment(self.t_env) + + def test_typed_dict_input_and_named_output_pipeline(self): + class Input(TypedDict): + count: int + label: str + + def expand(row: Input) -> Iterator[_Output]: + for value in range(row["count"]): + yield {"label": row["label"], "value": value} + + source = pf.from_records([(0, "empty"), (1, "a"), (2, "b")], schema=["count", "label"]) + result = source.flat_map(expand) + self.assertEqual(result.columns, ["value", "label"]) + self.assertEqual(result.schema.get_field_data_types(), + [DataTypes.BIGINT(), DataTypes.STRING()]) + rows = result.filter(pf.col("value") >= 0).select("label", "value").collect() + self.assertCountEqual(rows, [Row("a", 0), Row("b", 0), Row("b", 1)]) + + def test_wrapper_receives_named_row(self): + @pf.udtf(return_dtype=pf.DataType.struct({ + "is_row": pf.DataType.bool(), + "first": pf.DataType.int64(), + "by_name": pf.DataType.int64(), + "as_dict_value": pf.DataType.int64(), + })) + def expand(row: Row): + yield isinstance(row, Row), row[0], row["x"], row.as_dict()["z"] + + source = pf.from_records([(7, 3)], schema=["z", "x"]) + self.assertEqual(source.flat_map(expand).collect(), [Row(True, 7, 3, 7)]) + + def test_wrapper_row_input_reuse_and_column_calls(self): + @pf.udtf + def expand(value) -> Iterator[int]: + value = value[0] if isinstance(value, Row) else value + yield value + yield value + 1 + + a = pf.from_dict({"x": [1]}) + b = pf.from_dict({"y": [3]}) + before = expand(pf.col("x")).alias("out") + first, second = a.flat_map(expand), b.flat_map(expand) + self.assertEqual(first.columns, ["f0"]) + after = expand(pf.col("x")).alias("out") + combined = first.union_all(second) + for call in (before, after): + lateral = pf.from_table(a.to_table().join_lateral(call.expression)).select("out") + combined = combined.union_all(lateral) + self.assertCountEqual(combined.collect(), [Row(1), Row(2)] * 3 + [Row(3), Row(4)]) + + def test_table_function_and_callable_instance(self): + client_pid = os.getpid() + + class Expand(TableFunction): + def __init__(self): + self.constructor_pid = os.getpid() + + def open(self, context): + self.offset = 10 + + def eval(self, row: Row) -> Iterator[int]: + if self.constructor_pid != client_pid: + raise AssertionError("TableFunction must be constructed on the client") + yield row[0] + self.offset + + class Repeat: + def __call__(self, row: Dict[str, Any]) -> Iterator[int]: + yield row["x"] + yield row["x"] + + source = pf.from_dict({"x": [1]}) + first = source.flat_map(pf.udtf(Expand)) + second = source.flat_map(pf.udtf(Expand())) + third = source.flat_map(Repeat()) + self.assertCountEqual(first.union_all(second).union_all(third).collect(), + [Row(11), Row(11), Row(1), Row(1)]) + + def test_nested_values_and_nulls(self): + dtype = pf.DataType.struct({ + "payload": pf.DataType.struct({"label": pf.DataType.string()}), + "values": pf.DataType.list(pf.DataType.int64()), + "attributes": pf.DataType.map(pf.DataType.string(), pf.DataType.int64()), + "missing": pf.DataType.string(), + }) + + def expand(row: Optional[Dict[str, Any]]): + if row["x"] == 0: + return None + return {"payload": {"label": "ok"}, "values": [1, None], "attributes": {"a": 2}} + + result = pf.from_dict({"x": [0, 1]}).flat_map(expand, return_dtype=dtype) + # Table.collect cannot decode NULL array elements, so inspect them in the JVM. + projected = result.select( + "payload", pf.col("values").cardinality, pf.col("values").at(1), + pf.col("values").at(2).is_null, "attributes", "missing") + self.assertEqual(projected.collect(), [Row(Row("ok"), 2, 1, True, {"a": 2}, None)]) + + def test_map_output_is_one_column(self): + def expand(row) -> Iterator[Dict[str, int]]: + yield {"value": row["x"]} + yield {"value": row["x"] + 1} + + result = pf.from_dict({"x": [1]}).flat_map(expand) + self.assertEqual(result.columns, ["f0"]) + self.assertEqual(result.schema.get_field_data_types(), + [DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT())]) + self.assertCountEqual(result.collect(), [Row({"value": 1}), Row({"value": 2})]) + + +class DataFrameFlatMapStreamTests(_DataFrameFlatMapTests, PyFlinkStreamDataFrameTestCase): + pass + + +class DataFrameFlatMapBatchTests(_DataFrameFlatMapTests, PyFlinkBatchTableTestCase): + pass + + +class DataFrameFlatMapThreadTests(_DataFrameFlatMapTests, PyFlinkStreamDataFrameTestCase): + def setUp(self): + super().setUp() + self.t_env.get_config().set("python.execution-mode", "thread") + + +if __name__ == "__main__": + unittest.main() diff --git a/flink-python/pyflink/dataframe/udtf.py b/flink-python/pyflink/dataframe/udtf.py new file mode 100644 index 0000000000000..93a678422f689 --- /dev/null +++ b/flink-python/pyflink/dataframe/udtf.py @@ -0,0 +1,447 @@ +################################################################################ +# 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 table functions for the DataFrame API.""" + +import collections.abc +import inspect +import types +from dataclasses import dataclass +from typing import ( + Annotated, Any, Callable, Iterator, List, Optional, Tuple, Type, Union, cast, get_args, + get_origin, overload, +) + +from pyflink.common import Row +from pyflink.dataframe.datatype import DataType +from pyflink.dataframe.udf import ( + _DataTypeLike, + _UDFDeclarationContext, + _UNRESOLVED_TYPE_HINT, + _convert_to_dtype, + _create_declaration_context, + _create_result_normalizer, + _data_type_from_type_hint, + _get_callable_return_type_hint, + _is_typed_dict, + _resolve_callable_annotation, + _resolve_udf, + _validate_determinism_agreement, + _validate_zero_argument_class, +) +from pyflink.table.expression import Expression +from pyflink.table.expressions import col, with_columns +from pyflink.table.types import RowType +from pyflink.table.udf import ( + TableFunction, + UserDefinedFunction, + UserDefinedTableFunctionWrapper, + udtf as table_udtf, +) +from pyflink.util.api_stability_decorators import PublicEvolving + +__all__ = ["udtf"] + +_UDTFInput = Union[Callable[..., Any], TableFunction, Type[TableFunction]] + + +@dataclass(frozen=True) +class _DataFrameUDTFCall: + expression: Expression + return_dtype: DataType + has_named_fields: bool + output_aliases: Optional[Tuple[str, ...]] = None + + def alias(self, name: str, *extra_names: str) -> "_DataFrameUDTFCall": + names = (name,) + extra_names + table_type = self.return_dtype._to_table_data_type() + arity = len(table_type.fields) if isinstance(table_type, RowType) else 1 + if any(not isinstance(value, str) for value in names): + raise TypeError("UDTF output aliases must be strings.") + if len(names) != arity or len(set(names)) != len(names) or not all(names): + raise ValueError( + "UDTF output aliases must be non-empty, unique, and match output arity.") + return _DataFrameUDTFCall( + self.expression.alias(*names), self.return_dtype, self.has_named_fields, names + ) + + +@dataclass(frozen=True) +class _DataFrameUDTFWrapper: + _func: Union[Callable[..., Any], TableFunction] + _declaration_context: _UDFDeclarationContext + return_dtype: DataType + _has_named_fields: bool + _deterministic: bool + _name: Optional[str] + + @property + def __name__(self) -> str: + return self._create_table_wrapper()._name + + def __call__(self, *args: Any) -> _DataFrameUDTFCall: + return _DataFrameUDTFCall( + self._create_table_wrapper()(*args), self.return_dtype, self._has_named_fields + ) + + def _create_table_wrapper( + self, input_columns: Optional[Tuple[str, ...]] = None, *, input_as_dict: bool = False + ) -> UserDefinedTableFunctionWrapper: + """Keep Table's mutable row-input flag local to each use of this declaration.""" + return cast(UserDefinedTableFunctionWrapper, table_udtf( + _DataFrameTableFunctionAdapter( + self._func, + self.return_dtype, + self._deterministic, + input_columns, + input_as_dict, + ), + result_types=self.return_dtype._to_table_data_type(), + deterministic=self._deterministic, + name=self._name, + )) + + +@overload +def udtf( + func: _UDTFInput, + *, + return_dtype: Optional[_DataTypeLike] = ..., + deterministic: bool = ..., + name: Optional[str] = ..., +) -> _DataFrameUDTFWrapper: + ... + + +@overload +def udtf( + func: None = ..., + *, + return_dtype: Optional[_DataTypeLike] = ..., + deterministic: bool = ..., + name: Optional[str] = ..., +) -> Callable[[_UDTFInput], _DataFrameUDTFWrapper]: + ... + + +@PublicEvolving() +def udtf( + func: Optional[_UDTFInput] = None, + *, + return_dtype: Optional[_DataTypeLike] = None, + deterministic: bool = True, + name: Optional[str] = None, +) -> Union[_DataFrameUDTFWrapper, Callable[[_UDTFInput], _DataFrameUDTFWrapper]]: + """ + Declare a synchronous function that emits zero or more rows per invocation. + + Supports functions, callable instances, and ``TableFunction`` instances or + classes with a zero-argument constructor. Classes are instantiated on the client + when declared. The resulting or explicitly provided instance is serialized with + the job. ``TableFunction.open`` and + ``close`` run on workers, making ``open`` suitable for resource initialization. + + The output type may be explicit or inferred from ``Iterator[T]``, ``Iterable[T]``, + ``Generator[T, ...]``, or ``list[T]``. A ``TypedDict`` supplies output field names. + Return ``None`` for no rows, a list/generator for multiple rows, or a scalar, + tuple, ``Row``, or dict for one row. Nested output values follow scalar UDF rules. + + A declaration passed to :meth:`DataFrame.flat_map` receives a named Flink ``Row``; + an undecorated callable passed to that method receives a column-name dictionary. + + Example:: + + >>> from typing import Iterator, TypedDict + >>> import pyflink.dataframe as pf + >>> from pyflink.common import Row + >>> class Output(TypedDict): + ... value: int + >>> @pf.udtf + ... def expand(row: Row) -> Iterator[Output]: + ... yield {"value": row["x"]} + ... yield {"value": row["x"] + 1} + >>> result = pf.from_dict({"x": [1, 2]}).flat_map(expand) + + :param func: Function, callable instance, or ``TableFunction`` instance/class. + :param return_dtype: Emitted row type, as a DataFrame DataType, Python type, + or SQL type string. Inferred from annotations when omitted. + :param deterministic: Whether equal inputs produce equal results; must agree + with ``TableFunction.is_deterministic``. + :param name: Optional function name. None or an empty string uses the callable's name. + :return: A reusable UDTF declaration, or a decorator when ``func`` is omitted. + + .. versionadded:: 2.4.0 + """ + def decorator(f: _UDTFInput) -> _DataFrameUDTFWrapper: + actual_func, context = _resolve_udtf(f) + dtype, has_named_fields = _infer_udtf_return_dtype(context, return_dtype) + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool.") + if isinstance(actual_func, TableFunction): + _validate_determinism_agreement(deterministic, actual_func.is_deterministic()) + return _DataFrameUDTFWrapper( + actual_func, context, dtype, has_named_fields, deterministic, name) + + return decorator if func is None else decorator(func) + + +def _resolve_udtf( + func: _UDTFInput, +) -> Tuple[Union[Callable[..., Any], TableFunction], _UDFDeclarationContext]: + if inspect.isclass(func): + if not issubclass(func, TableFunction): + raise TypeError("UDTF classes must extend TableFunction; pass a callable instance.") + _validate_zero_argument_class(func) + func = func() + if not isinstance(func, TableFunction): + raise TypeError("A TableFunction class must construct a TableFunction instance.") + if isinstance(func, TableFunction): + if not callable(func.eval): + raise TypeError("TableFunction.eval must be callable.") + context = _create_declaration_context(func.eval, partial_source=func.eval) + else: + if isinstance(func, UserDefinedFunction): + raise TypeError("func must be a table UDF or a Python callable.") + context = _resolve_udf(func).declaration_context + target = context.annotation_target + try: + unwrapped_target = inspect.unwrap(target) + except ValueError as exc: + raise TypeError("Cannot inspect a UDTF with a wrapper cycle.") from exc + if any( + inspect.iscoroutinefunction(candidate) or inspect.isasyncgenfunction(candidate) + for candidate in (target, unwrapped_target) + ): + raise TypeError("DataFrame UDTFs must be synchronous; async functions are not supported.") + return func, context + + +def _infer_udtf_return_dtype( + context: _UDFDeclarationContext, return_dtype: Optional[_DataTypeLike] +) -> Tuple[DataType, bool]: + if return_dtype is not None: + dtype = _convert_to_dtype(return_dtype) + has_named_fields = isinstance(dtype._to_table_data_type(), RowType) + else: + hint = _get_callable_return_type_hint(context) + origin, arguments = get_origin(hint), get_args(hint) + if origin in ( + collections.abc.Iterator, collections.abc.Iterable, collections.abc.Generator, list + ) and arguments: + hint = arguments[0] + elif hint is _UNRESOLVED_TYPE_HINT or origin in ( + collections.abc.Iterator, collections.abc.Iterable, collections.abc.Generator, list + ): + raise TypeError("Cannot infer UDTF return type; specify return_dtype explicitly.") + if get_origin(hint) is tuple: + fields = get_args(hint) + if not fields or Ellipsis in fields: + raise TypeError("UDTF tuple outputs must have a fixed number of fields.") + dtype = DataType.struct([ + (f"f{index}", _data_type_from_type_hint(field)) + for index, field in enumerate(fields) + ]) + has_named_fields = False + else: + dtype = _data_type_from_type_hint(hint) + has_named_fields = isinstance(dtype._to_table_data_type(), RowType) + table_type = dtype._to_table_data_type() + if isinstance(table_type, RowType) and not table_type.fields: + raise ValueError("A UDTF must declare at least one output field.") + return dtype, has_named_fields + + +def _validate_flat_map_input(declaration: _DataFrameUDTFWrapper, raw_callable: bool) -> None: + source = declaration._func + target: Callable[..., Any] + if isinstance(source, TableFunction): + target = source.eval + else: + target = source + try: + signature = inspect.signature(target) + except (TypeError, ValueError): + return + parameters = list(signature.parameters.values()) + try: + signature.bind(object()) + except TypeError as exc: + raise ValueError("flat_map requires a function accepting one row argument.") from exc + positional = [p for p in parameters if p.kind in ( + inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD + )] + if not positional: + return + hint = _resolve_callable_annotation(declaration._declaration_context, positional[0].name) + if not _is_flat_map_row_hint(hint, raw_callable): + row_type = "dict" if raw_callable else "Row" + raise ValueError(f"flat_map receives one {row_type} row argument, got annotation {hint}.") + + +def _is_flat_map_row_hint(hint: Any, raw_callable: bool) -> bool: + if hint in (_UNRESOLVED_TYPE_HINT, Any, object): + return True + origin = get_origin(hint) + if origin is Annotated: + return _is_flat_map_row_hint(get_args(hint)[0], raw_callable) + if origin in (Union, getattr(types, "UnionType", Union)): + return any(_is_flat_map_row_hint(member, raw_callable) for member in get_args(hint)) + if _is_typed_dict(hint): + return raw_callable + target = origin or hint + if not raw_callable and target is tuple: + return True + try: + return issubclass(dict if raw_callable else Row, target) + except TypeError: + # Leave annotations that cannot be checked at runtime to the user's type checker. + return True + + +def _resolve_flat_map_udtf( + func: Union[Callable[..., Any], _DataFrameUDTFWrapper], + return_dtype: Optional[_DataTypeLike], + input_columns: List[str], +) -> Tuple[Expression, List[str]]: + raw_callable = not isinstance(func, _DataFrameUDTFWrapper) + if isinstance(func, _DataFrameUDTFWrapper): + if return_dtype is not None: + raise ValueError("return_dtype must not be specified for a DataFrame UDTF declaration.") + declaration = func + else: + if func is None: + raise TypeError("flat_map requires a callable or a pf.udtf declaration.") + if inspect.isclass(func) or isinstance(func, UserDefinedFunction): + raise TypeError("flat_map accepts a callable instance or a pf.udtf declaration.") + declaration = udtf(func, return_dtype=return_dtype) + _validate_flat_map_input(declaration, raw_callable) + table_type = declaration.return_dtype._to_table_data_type() + if declaration._has_named_fields: + output_columns = cast(RowType, table_type).field_names() + elif isinstance(table_type, RowType) and len(table_type.fields) > 1: + raise ValueError("flat_map requires named output fields; use TypedDict or a named struct.") + else: + output_columns = ["f0"] + wrapper = declaration._create_table_wrapper(tuple(input_columns), input_as_dict=raw_callable) + wrapper._set_takes_row_as_input() + return wrapper(with_columns(col("*"))), output_columns + + +def _iter_user_results(result: Any) -> Iterator[Any]: + if result is None: + return + if isinstance(result, (Row, tuple, collections.abc.Mapping, str, bytes, bytearray)): + yield result + elif isinstance(result, collections.abc.Iterable): + yield from result + else: + yield result + + +class _DataFrameTableFunctionAdapter(TableFunction): + def __init__( + self, + func: Union[Callable[..., Any], TableFunction], + return_dtype: DataType, + deterministic: bool, + input_columns: Optional[Tuple[str, ...]], + input_as_dict: bool, + ) -> None: + self._func = func + self._return_dtype = return_dtype + self._deterministic = deterministic + self._input_columns = input_columns + self._input_as_dict = input_as_dict + self._bound_invocation: Optional[Callable[..., Any]] = None + self._lifecycle_opened = False + self.__name__ = getattr(func, "__name__", type(func).__name__) + + def open(self, function_context: Any) -> None: + if isinstance(self._func, TableFunction): + self._func.open(function_context) + self._lifecycle_opened = True + invoke_func = self._func.eval + else: + invoke_func = self._func + try: + self._bound_invocation = self._bind_func(invoke_func) + except Exception: + try: + self.close() + except Exception: + pass + raise + + def close(self) -> None: + try: + if self._lifecycle_opened: + cast(TableFunction, self._func).close() + finally: + self._bound_invocation = None + self._lifecycle_opened = False + + def is_deterministic(self) -> bool: + if isinstance(self._func, TableFunction): + return self._func.is_deterministic() + return self._deterministic + + def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: + table_type = self._return_dtype._to_table_data_type() + normalizer = _create_result_normalizer(table_type) + is_struct = isinstance(table_type, RowType) + field_count = len(table_type.fields) if isinstance(table_type, RowType) else 1 + input_columns = list(self._input_columns) if self._input_columns is not None else None + + def invoke(*args: Any) -> Iterator[Row]: + if input_columns is not None: + input_row = args[0] + row: Union[dict, Row] + if self._input_as_dict: + row = {name: input_row[i] for i, name in enumerate(input_columns)} + else: + # Thread mode supplies a tuple; expose the same named Row in both modes. + row = Row(*input_row) + row.set_field_names(input_columns) + if isinstance(input_row, Row): + row.set_row_kind(input_row.get_row_kind()) + result = invoke_func(row) + else: + result = invoke_func(*args) + for item in _iter_user_results(result): + if is_struct and field_count == 1 and item is not None: + if not isinstance(item, (Row, tuple, collections.abc.Mapping)): + item = Row(item) + if not is_struct and isinstance(item, (Row, tuple)): + if len(item) != 1: + raise ValueError("UDTF scalar output requires exactly one field.") + item = item[0] + normalized = normalizer(item) if normalizer is not None else item + if is_struct: + if normalized is None: + yield Row(*([None] * field_count)) + else: + yield normalized + else: + yield Row(normalized) + + return invoke + + def eval(self, *args: Any) -> Iterator[Row]: + if self._bound_invocation is None: + raise RuntimeError("DataFrame UDTF was invoked before open().") + return self._bound_invocation(*args) From e8e399e1aa0eec4b43f208c968746ba835236b99 Mon Sep 17 00:00:00 2001 From: victor Date: Sun, 20 Sep 2026 19:31:41 +0800 Subject: [PATCH 2/2] [FLINK-40434][python] Support callable classes and SQL binding for DataFrame UDTFs AI-Contributed/Feature: 0/312 AI-Contributed/UT: 0/318 --- .../docs/reference/pyflink.dataframe/udf.rst | 2 +- flink-python/pyflink/dataframe/dataframe.py | 118 +++++---- flink-python/pyflink/dataframe/sql.py | 44 ++-- .../pyflink/dataframe/tests/test_sql.py | 94 ++++++++ .../pyflink/dataframe/tests/test_udtf.py | 224 ++++++++++++++---- flink-python/pyflink/dataframe/udtf.py | 148 +++++++----- 6 files changed, 468 insertions(+), 162 deletions(-) diff --git a/flink-python/docs/reference/pyflink.dataframe/udf.rst b/flink-python/docs/reference/pyflink.dataframe/udf.rst index d2e803c347d42..c71e915b646af 100644 --- a/flink-python/docs/reference/pyflink.dataframe/udf.rst +++ b/flink-python/docs/reference/pyflink.dataframe/udf.rst @@ -40,7 +40,7 @@ Use :func:`pyflink.dataframe.udtf` to declare a Python function that emits zero more rows per invocation. A table UDF can be used with :meth:`~pyflink.dataframe.DataFrame.flat_map`, which returns only the emitted columns. -DataFrame table UDFs support synchronous functions, callable instances, and +DataFrame table UDFs support synchronous functions, callable classes and instances, and ``TableFunction`` instances or classes. See :func:`pyflink.dataframe.udtf` for declaration forms, type inference, lifecycle, and examples, and :meth:`~pyflink.dataframe.DataFrame.flat_map` for row input and output semantics. diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 49ca24dd3e146..95d85a9a15f9a 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -26,6 +26,7 @@ Optional, Set, Tuple, + Type, TypeVar, Union, overload, @@ -136,57 +137,6 @@ class DataFrame: def __init__(self, table: Table): self._table = table - @PublicEvolving() - def flat_map( - self, - func: Union[Callable[[Dict[str, Any]], Any], "_DataFrameUDTFWrapper"], - *, - return_dtype: Optional["_DataTypeLike"] = None, - ) -> "DataFrame": - """ - Apply a function to each row, emitting zero or more output rows. - - A plain callable receives a dictionary keyed by column name. A declaration - created with :func:`pyflink.dataframe.udtf` receives a named Flink ``Row``. - Output column names come from a ``TypedDict`` or an explicit named struct; - scalar outputs use ``f0``. Multi-field outputs require named fields. - - :param func: Row-based callable or a declaration created with ``pf.udtf``. - :param return_dtype: Emitted row type, inferred from annotations when omitted. - Required if inference is not possible; must be omitted - for a UDTF declaration. - :return: A DataFrame containing only the emitted output columns. - - Example:: - - >>> from typing import Any, Dict, Iterator, TypedDict - >>> import pyflink.dataframe as pf - >>> class Token(TypedDict): - ... word: str - >>> def split(row: Dict[str, Any]) -> Iterator[Token]: - ... for word in row["text"].split(): - ... yield {"word": word} - >>> df = pf.from_dict({"text": ["hello world", "flink"]}) - >>> result = df.flat_map(split) - - A reusable UDTF declaration receives a named ``Row``:: - - >>> from pyflink.common import Row - >>> @pf.udtf - ... def split_row(row: Row) -> Iterator[Token]: - ... for word in row["text"].split(): - ... yield {"word": word} - >>> result = df.flat_map(split_row) - - .. versionadded:: 2.4.0 - """ - from pyflink.dataframe.udtf import _resolve_flat_map_udtf - - expression, output_columns = _resolve_flat_map_udtf(func, return_dtype, self.columns) - table = self._table.flat_map(expression) - # Table UDTFs expose positional field names, so restore the declared names. - return DataFrame(table.alias(output_columns[0], *output_columns[1:])) - # ======================== Core Operations ======================== @PublicEvolving() @@ -675,6 +625,72 @@ def top_n( distinct = drop_duplicates unique = drop_duplicates + @PublicEvolving() + def flat_map( + self, + func: Union[Callable[[Dict[str, Any]], Any], Type, "_DataFrameUDTFWrapper"], + *, + return_dtype: Optional["_DataTypeLike"] = None, + ) -> "DataFrame": + """ + Apply a function to each row, emitting zero or more output rows. + + The function receives a dictionary keyed by column name, including when + declared with :func:`pyflink.dataframe.udtf`. + Output column names come from a ``TypedDict`` or an explicit named struct; + scalar outputs use ``f0``. Multi-field outputs require named fields. + + :param func: Row-based callable, a callable class with a zero-argument constructor, + or a declaration created with ``pf.udtf``. Callable classes are + instantiated on workers. + :param return_dtype: Emitted row type, inferred from annotations when omitted. + Required if inference is not possible; must be omitted + for a UDTF declaration. + :return: A DataFrame containing only the emitted output columns. + + Example:: + + >>> from typing import Any, Dict, Iterator, TypedDict + >>> import pyflink.dataframe as pf + >>> class Token(TypedDict): + ... word: str + >>> def split(row: Dict[str, Any]) -> Iterator[Token]: + ... for word in row["text"].split(): + ... yield {"word": word} + >>> df = pf.from_dict({"text": ["hello world", "flink"]}) + >>> result = df.flat_map(split) + >>> result.columns + ['word'] + + A reusable UDTF declaration receives the same dictionary input:: + + >>> @pf.udtf + ... def split_row(row: Dict[str, Any]) -> Iterator[Token]: + ... for word in row["text"].split(): + ... yield {"word": word} + >>> result = df.flat_map(split_row) + + An explicit output type can be supplied for unannotated callables:: + + >>> words = df.flat_map(lambda row: row["text"].split(), return_dtype=str) + >>> words.columns + ['f0'] + >>> named = df.flat_map( + ... lambda row: row["text"].split(), return_dtype="ROW") + >>> named.columns + ['word'] + + See :func:`pyflink.dataframe.udtf` for callable class and ``TableFunction`` examples. + + .. versionadded:: 2.4.0 + """ + from pyflink.dataframe.udtf import _resolve_flat_map_udtf + + expression, output_columns = _resolve_flat_map_udtf(func, return_dtype, self.columns) + table = self._table.flat_map(expression) + # Table UDTFs expose positional field names, so restore the declared names. + return DataFrame(table.alias(output_columns[0], *output_columns[1:])) + # ======================== Filtering & Ordering ======================== @PublicEvolving() diff --git a/flink-python/pyflink/dataframe/sql.py b/flink-python/pyflink/dataframe/sql.py index 879a060b29c18..9a9e11e213444 100644 --- a/flink-python/pyflink/dataframe/sql.py +++ b/flink-python/pyflink/dataframe/sql.py @@ -26,9 +26,11 @@ from pyflink.dataframe.context import get_or_create_table_environment from pyflink.dataframe.dataframe import DataFrame from pyflink.dataframe.udf import _DataFrameUDFWrapper +from pyflink.dataframe.udtf import _DataFrameUDTFWrapper from pyflink.java_gateway import get_gateway from pyflink.table import Table, TableEnvironment from pyflink.table.expression import Expression +from pyflink.table.udf import UserDefinedFunctionWrapper from pyflink.util.api_stability_decorators import PublicEvolving from pyflink.util.java_utils import is_instance_of @@ -36,11 +38,13 @@ _LOG = logging.getLogger(__name__) -# UDFs are annotated with the declared return type of :func:`pyflink.dataframe.udf` +# Scalar UDFs use the declared return type of :func:`pyflink.dataframe.udf` # so that its result type-checks as a binding; at runtime a binding must be an actual # UDF object, which is what _BINDABLE_TYPES enforces. -_Binding = Union[DataFrame, Callable[..., Expression]] -_BINDABLE_TYPES = (DataFrame, _DataFrameUDFWrapper) +_FunctionBinding = Union[_DataFrameUDFWrapper, _DataFrameUDTFWrapper] +_Binding = Union[DataFrame, Callable[..., Expression], _DataFrameUDTFWrapper] +_FUNCTION_TYPES = (_DataFrameUDFWrapper, _DataFrameUDTFWrapper) +_BINDABLE_TYPES = (DataFrame, *_FUNCTION_TYPES) @PublicEvolving() @@ -51,9 +55,9 @@ def sql(query: str, *, auto_bind: bool = True, **bindings: _Binding) -> DataFram The query must be a single statement that returns a result, such as SELECT or VALUES (no INSERT / DDL; use :meth:`TableEnvironment.execute_sql` for those). The referenced DataFrames are registered as temporary views and the referenced - UDFs (created with :func:`~pyflink.dataframe.udf`) as temporary system functions - for the duration of the call, and both are dropped afterwards. The result can be - further transformed with the DataFrame API. + UDFs (created with :func:`~pyflink.dataframe.udf` or :func:`~pyflink.dataframe.udtf`) + as temporary system functions for the duration of the call, and both are dropped + afterwards. The result can be further transformed with the DataFrame API. When ``auto_bind`` is ``True`` (the default), the caller's local and global variables are scanned for :class:`DataFrame` and UDF objects and each is registered under its @@ -91,7 +95,8 @@ def sql(query: str, *, auto_bind: bool = True, **bindings: _Binding) -> DataFram auto-bound candidates belong to different TableEnvironments when there are no explicit bindings. :raises TypeError: If an explicit binding is neither a :class:`DataFrame` nor a - UDF created with :func:`~pyflink.dataframe.udf`. + UDF created with :func:`~pyflink.dataframe.udf` or + :func:`~pyflink.dataframe.udtf`. Example:: @@ -112,6 +117,11 @@ def sql(query: str, *, auto_bind: bool = True, **bindings: _Binding) -> DataFram ... return value + 1 >>> pf.sql("SELECT add_one(a) AS a1 FROM df1") >>> pf.sql("SELECT inc(a) FROM src", auto_bind=False, src=df1, inc=add_one) + >>> # Table UDFs can be used in lateral joins + >>> @pf.udtf(return_dtype=str) + ... def chars(text): + ... yield from text + >>> pf.sql("SELECT ch FROM df1, LATERAL TABLE(chars(b)) AS T(ch)") >>> # Mix SQL and the DataFrame API >>> pf.sql("SELECT a, b FROM df1").filter(pf.col("a") > 1).to_pandas() @@ -138,7 +148,7 @@ def sql(query: str, *, auto_bind: bool = True, **bindings: _Binding) -> DataFram if not isinstance(value, _BINDABLE_TYPES): raise TypeError( f"sql() binding '{name}' must be a DataFrame or a UDF created with " - f"pyflink.dataframe.udf, got {type(value).__name__}" + f"pyflink.dataframe.udf or pyflink.dataframe.udtf, got {type(value).__name__}" ) explicit_frames = _get_dataframes(bindings) explicit_udfs = _get_udfs(bindings) @@ -162,8 +172,8 @@ def _get_dataframes(namespace: Dict[str, Any]) -> Dict[str, DataFrame]: return {k: v for k, v in namespace.items() if isinstance(v, DataFrame)} -def _get_udfs(namespace: Dict[str, Any]) -> Dict[str, _DataFrameUDFWrapper]: - return {k: v for k, v in namespace.items() if isinstance(v, _DataFrameUDFWrapper)} +def _get_udfs(namespace: Dict[str, Any]) -> Dict[str, _FunctionBinding]: + return {k: v for k, v in namespace.items() if isinstance(v, _FUNCTION_TYPES)} def _drop_views(t_env: TableEnvironment, names: List[str]) -> None: @@ -322,10 +332,16 @@ def _register_views( return registered +def _get_table_udf_wrapper(value: _FunctionBinding) -> UserDefinedFunctionWrapper: + if isinstance(value, _DataFrameUDTFWrapper): + return value._create_table_wrapper(preserve_field_names=True) + return value._table_udf_wrapper + + def _register_functions( t_env: TableEnvironment, - explicit: Dict[str, _DataFrameUDFWrapper], - auto: Dict[str, _DataFrameUDFWrapper], + explicit: Dict[str, _FunctionBinding], + auto: Dict[str, _FunctionBinding], ) -> List[str]: """ Register explicit and auto-collected UDFs as temporary system functions and return @@ -364,7 +380,7 @@ def _register_functions( f"cannot bind '{name}': a temporary function with this name " "already exists" ) - t_env.create_temporary_system_function(name, value._table_udf_wrapper) + t_env.create_temporary_system_function(name, _get_table_udf_wrapper(value)) registered.append(name) # Auto-bound candidates are not expected to raise: problems are reported as @@ -377,7 +393,7 @@ def _register_functions( _warn_skipped(name, "a function with this name already exists") continue try: - t_env.create_temporary_system_function(name, value._table_udf_wrapper) + t_env.create_temporary_system_function(name, _get_table_udf_wrapper(value)) except Exception as e: _warn_skipped(name, f"registration failed: {e}") continue diff --git a/flink-python/pyflink/dataframe/tests/test_sql.py b/flink-python/pyflink/dataframe/tests/test_sql.py index bcc7b1190b7b1..be8d9033c940b 100644 --- a/flink-python/pyflink/dataframe/tests/test_sql.py +++ b/flink-python/pyflink/dataframe/tests/test_sql.py @@ -18,6 +18,7 @@ import unittest import warnings +from typing import Iterator, TypedDict from py4j.protocol import Py4JJavaError @@ -376,6 +377,99 @@ def test_udf_names_are_case_insensitive(self): self.assertEqual(pf.sql(query).collect(), [Row(2)]) self.assertNotIn("addone", self.t_env.list_user_defined_functions()) + def test_udtf_explicit_and_auto_bindings(self): + source = pf.from_dict({"text": ["ab", ""]}) + + @pf.udtf(return_dtype=str) + def chars(text): + yield from text + + query = "SELECT ch FROM source, LATERAL TABLE(chars(text)) AS T(ch)" + explicit = pf.sql(query, auto_bind=False, source=source, chars=chars) + automatic = pf.sql(query) + self.assertEqual(explicit.columns, ["ch"]) + self.assertEqual(explicit.schema.get_field_data_types(), [DataTypes.STRING()]) + self.assertCountEqual(explicit.union_all(automatic).collect(), [Row("a"), Row("b")] * 2) + self.assertNotIn("chars", self.t_env.list_user_defined_functions()) + self.assertNotIn("source", self.t_env.list_temporary_views()) + + def test_udtf_declared_field_names_and_sql_aliases(self): + class Token(TypedDict): + word: str + size: int + + def expand(text) -> Iterator[Token]: + yield {"word": text, "size": len(text)} + + source = pf.from_dict({"text": ["ab"]}) + for return_dtype in (None, "ROW", + pf.DataType.struct({"word": pf.DataType.string(), + "size": pf.DataType.int64()})): + with self.subTest(return_dtype=return_dtype): + declaration = pf.udtf(expand, return_dtype=return_dtype) + named = pf.sql( + "SELECT t.word, t.size FROM src, LATERAL TABLE(expand(text)) AS t", + auto_bind=False, src=source, expand=declaration) + aliased = pf.sql( + "SELECT t.w, t.n FROM src, LATERAL TABLE(expand(text)) AS t(w, n)", + auto_bind=False, src=source, expand=declaration) + self.assertEqual(named.columns, ["word", "size"]) + self.assertEqual(aliased.columns, ["w", "n"]) + self.assertEqual(named.schema.get_field_data_types(), + [DataTypes.STRING(), DataTypes.BIGINT()]) + self.assertEqual(named.union_all(aliased).collect(), [Row("ab", 2)] * 2) + self.assertNotIn("expand", self.t_env.list_user_defined_functions()) + + def test_udtf_sql_preserves_quoted_names_and_nested_types(self): + dtype = pf.DataType.struct({ + "word text": pf.DataType.string().not_null(), + "meta": pf.DataType.struct({"size": pf.DataType.int64()}), + }) + expand = pf.udtf(lambda text: {"word text": text, "meta": {"size": len(text)}}, + return_dtype=dtype) + result = pf.sql( + "SELECT t.`word text`, t.meta FROM src, LATERAL TABLE(expand(text)) AS t", + auto_bind=False, src=pf.from_dict({"text": ["ab"]}), expand=expand) + self.assertEqual(result.columns, ["word text", "meta"]) + self.assertEqual(result.schema.get_field_data_types(), [ + DataTypes.STRING().not_null(), + DataTypes.ROW([DataTypes.FIELD("size", DataTypes.BIGINT())])]) + self.assertEqual(result.collect(), [Row("ab", Row(2))]) + + def test_udtf_sql_binding_after_flat_map_with_callable_class(self): + class Expand: + def __call__(self, value): + text = value["text"] if isinstance(value, dict) else value + yield from text + + chars = pf.udtf(Expand, return_dtype=str) + source = pf.from_dict({"text": ["ab"]}) + mapped = source.flat_map(chars) + query = "SELECT ch FROM src, LATERAL TABLE(chars(text)) AS T(ch)" + result = pf.sql(query, auto_bind=False, src=source, chars=chars) + self.assertCountEqual(mapped.union_all(result).collect(), [Row("a"), Row("b")] * 2) + + def test_udtf_bindings_are_dropped_after_planning_failure(self): + source = pf.from_dict({"text": ["ab"]}) + chars = pf.udtf(lambda text: list(text), return_dtype=str) + with self.assertRaises(Py4JJavaError): + pf.sql("SELECT * FROM src, LATERAL TABLE(chars(missing)) AS T(ch)", + auto_bind=False, src=source, chars=chars) + self.assertNotIn("chars", self.t_env.list_user_defined_functions()) + self.assertNotIn("src", self.t_env.list_temporary_views()) + + def test_udtf_binding_conflict_rolls_back_earlier_registrations(self): + self.t_env.create_temporary_system_function( + "taken", table_udf(lambda i: i + 100, result_type=DataTypes.BIGINT())) + self.addCleanup(self.t_env.drop_temporary_system_function, "taken") + chars = pf.udtf(lambda text: list(text), return_dtype=str) + with self.assertRaisesRegex(ValueError, "'taken'.*already exists"): + pf.sql("SELECT * FROM src", auto_bind=False, + src=pf.from_dict({"text": ["ab"]}), first=chars, taken=chars) + self.assertNotIn("first", self.t_env.list_user_defined_functions()) + self.assertNotIn("src", self.t_env.list_temporary_views()) + self.assertIn("taken", self.t_env.list_user_defined_functions()) + def test_pandas_udfs_are_bindable(self): df = pf.from_dict({"a": [1, 2]}) # noqa: F841 add_one = pf.udf( # noqa: F841 diff --git a/flink-python/pyflink/dataframe/tests/test_udtf.py b/flink-python/pyflink/dataframe/tests/test_udtf.py index 44219a9e77e9a..80623aaaae123 100644 --- a/flink-python/pyflink/dataframe/tests/test_udtf.py +++ b/flink-python/pyflink/dataframe/tests/test_udtf.py @@ -28,7 +28,7 @@ import cloudpickle import pyflink.dataframe as pf -from pyflink.common import Row, RowKind +from pyflink.common import Row from pyflink.dataframe.udtf import _resolve_flat_map_udtf from pyflink.table import DataTypes from pyflink.table.udf import ScalarFunction, TableFunction @@ -100,9 +100,15 @@ class Scalar(ScalarFunction): def eval(self, value): return value - class CallableClass: + class NotCallable: + pass + + class RequiresCallableArgument: + def __init__(self, count): + self.count = count + def __call__(self, value): - return [value] + return [value] * self.count class RequiresArgument(TableFunction): def __init__(self, value): @@ -111,7 +117,8 @@ def __init__(self, value): def eval(self, value): return [value] - for func in (123, Scalar(), Scalar, CallableClass, RequiresArgument): + for func in (123, Scalar(), Scalar, NotCallable, + RequiresCallableArgument, RequiresArgument): with self.subTest(func=func), self.assertRaises(TypeError): pf.udtf(func, return_dtype=int) for hint in (Iterator, Iterator[Tuple[int, ...]], Iterator[Tuple]): @@ -140,7 +147,7 @@ async def __call__(self, row): def hidden_generator(row): return generator(row) - for func in (coroutine, generator, AsyncCallable(), hidden_generator): + for func in (coroutine, generator, AsyncCallable, AsyncCallable(), hidden_generator): with self.subTest(func=func), self.assertRaisesRegex(TypeError, "async"): pf.udtf(func, return_dtype=int) @@ -320,9 +327,6 @@ def test_single_named_field_accepts_scalar_items(self): self.assertEqual(_eval_udtf(pf.udtf(lambda: result, return_dtype=dtype)), expected) def test_flat_map_validates_row_input_before_building_expression(self): - def named(row: Dict[str, Any], optional=0) -> Iterator[int]: - yield row["x"] + optional - def column(value: int) -> Iterator[int]: yield value @@ -332,7 +336,8 @@ def multiple(left, right) -> Iterator[int]: def keyword(row, *, required) -> Iterator[int]: yield required - for func in (column, multiple, keyword, pf.udtf(column), pf.udtf(multiple), pf.udtf(named)): + for func in (column, multiple, keyword, + pf.udtf(column), pf.udtf(multiple), pf.udtf(keyword)): with self.subTest(func=func), self.assertRaisesRegex(ValueError, "row argument"): _resolve_flat_map_udtf(func, None, ["x"]) with self.assertRaisesRegex(ValueError, "return_dtype"): @@ -341,7 +346,8 @@ def keyword(row, *, required) -> Iterator[int]: _resolve_flat_map_udtf(None, int, ["x"]) def test_incompatible_input_annotations(self): - for hint in (Optional[int], Union[int, str], Annotated[int, "column value"]): + for hint in (Row, Tuple[int], Optional[int], Union[int, str], + Annotated[int, "column value"]): def expand(row) -> Iterator[int]: yield row expand.__annotations__["row"] = hint @@ -350,25 +356,61 @@ def expand(row) -> Iterator[int]: with self.assertRaisesRegex(ValueError, "row argument"): _resolve_flat_map_udtf(func, None, ["x"]) - def test_flat_map_names_row_without_mutating_input(self): + def test_flat_map_passes_dict_without_mutating_input(self): received = [] @pf.udtf - def expand(row: Row) -> Iterator[int]: - received.append(row) - yield row["x"] + def expand(row: Dict[str, Any]) -> Iterator[int]: + received.append(row.copy()) + yield row.pop("x") adapter = expand._create_table_wrapper(("z", "x"))._func adapter.open(None) self.addCleanup(adapter.close) - original = Row.of_kind(RowKind.DELETE, original_z=7, original_x=3) - for value, kind in (((7, 3), RowKind.INSERT), (original, RowKind.DELETE)): + original = Row(original_z=7, original_x=3) + for value in ((7, 3), original): self.assertEqual(list(adapter.eval(value)), [Row(3)]) - self.assertIsInstance(received[-1], Row) - self.assertEqual(received[-1].as_dict(), {"z": 7, "x": 3}) - self.assertEqual(received[-1].get_row_kind(), kind) + self.assertEqual(received[-1], {"z": 7, "x": 3}) self.assertEqual(original.as_dict(), {"original_z": 7, "original_x": 3}) + def test_callable_class_initializes_per_worker_adapter(self): + events = [] + + class Counter: + def __init__(self): + events.append("init") + self.count = 0 + + def __call__(self, value: int) -> Iterator[int]: + self.count += 1 + yield value + self.count + + declaration = pf.udtf(Counter) + adapters = [declaration._create_table_wrapper()._func for _ in range(2)] + self.assertEqual(events, []) + for adapter in adapters: + adapter.open(None) + self.addCleanup(adapter.close) + self.assertEqual(list(adapter.eval(10)), [Row(11)]) + self.assertEqual(list(adapter.eval(10)), [Row(12)]) + self.assertEqual(events, ["init", "init"]) + adapters[0].close() + with self.assertRaisesRegex(RuntimeError, "before open"): + adapters[0].eval(10) + + def test_callable_class_constructor_failure_is_reported_on_worker(self): + class FailingConstructor: + def __init__(self): + raise RuntimeError("constructor failed") + + def __call__(self, value) -> Iterator[int]: + yield value + + adapter = pf.udtf(FailingConstructor)._create_table_wrapper()._func + with self.assertRaisesRegex(RuntimeError, "constructor failed"): + adapter.open(None) + adapter.close() + class DataFrameUDTFPlanningTests(PyFlinkDataFrameUTTestCase): def test_compatible_input_annotations(self): @@ -382,14 +424,60 @@ class Input(TypedDict): def expand(row) -> Iterator[int]: yield row["x"] expand.__annotations__["row"] = hint - with self.subTest(hint=hint): - self.assertEqual(source.flat_map(expand).columns, ["f0"]) + for func in (expand, pf.udtf(expand)): + with self.subTest(hint=hint, func=func): + self.assertEqual(source.flat_map(func).columns, ["f0"]) - @pf.udtf - def expand_row(row: Optional[Row]) -> Iterator[int]: - yield row["x"] + def test_callable_class_signatures_and_annotations(self): + class Expand: + def __call__(self, row: Dict[str, Any], optional=0) -> Iterator[int]: + yield row["x"] + optional - self.assertEqual(source.flat_map(expand_row).columns, ["f0"]) + class ClassMethod: + @classmethod + def __call__(cls, row: Dict[str, Any]) -> Iterator[int]: + yield row["x"] + + class StaticMethod: + @staticmethod + def __call__(row: Dict[str, Any]) -> Iterator[int]: + yield row["x"] + + class MultipleArguments: + def __call__(self, left, right) -> Iterator[int]: + yield left + right + + class ColumnArgument: + def __call__(self, value: int) -> Iterator[int]: + yield value + + source = pf.from_dict({"x": [1]}) + for cls in (Expand, ClassMethod, StaticMethod): + for func in (cls, pf.udtf(cls)): + with self.subTest(func=func): + self.assertEqual(source.flat_map(func).columns, ["f0"]) + for cls in (MultipleArguments, ColumnArgument): + for func in (cls, pf.udtf(cls)): + with self.subTest(func=func), self.assertRaisesRegex(ValueError, "row argument"): + source.flat_map(func) + + def test_wrapped_static_method_still_validates_input_annotations(self): + def transparent(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + return wrapper + + class Expand: + @staticmethod + @transparent + def __call__(value: int) -> Iterator[int]: + yield value + + source = pf.from_dict({"x": [1]}) + for func in (Expand, pf.udtf(Expand), Expand(), pf.udtf(Expand())): + with self.subTest(func=func), self.assertRaisesRegex(ValueError, "dict row argument"): + source.flat_map(func) def test_unnamed_multiple_fields_require_schema(self): def emit(row) -> Iterator[Tuple[int, str]]: @@ -424,8 +512,8 @@ class Expand(TableFunction): def __init__(self): events.append("init") - def eval(self, row: Row) -> Iterator[int]: - yield row[0] + def eval(self, row: Dict[str, Any]) -> Iterator[int]: + yield next(iter(row.values())) declaration = pf.udtf(Expand) self.assertEqual(events, ["init"]) @@ -459,23 +547,27 @@ def expand(row: Input) -> Iterator[_Output]: rows = result.filter(pf.col("value") >= 0).select("label", "value").collect() self.assertCountEqual(rows, [Row("a", 0), Row("b", 0), Row("b", 1)]) - def test_wrapper_receives_named_row(self): - @pf.udtf(return_dtype=pf.DataType.struct({ - "is_row": pf.DataType.bool(), + def test_plain_and_decorated_functions_receive_dict(self): + dtype = pf.DataType.struct({ + "is_dict": pf.DataType.bool(), "first": pf.DataType.int64(), "by_name": pf.DataType.int64(), - "as_dict_value": pf.DataType.int64(), - })) - def expand(row: Row): - yield isinstance(row, Row), row[0], row["x"], row.as_dict()["z"] + }) + + def expand(row: Dict[str, Any]): + yield isinstance(row, dict), next(iter(row.values())), row["x"] source = pf.from_records([(7, 3)], schema=["z", "x"]) - self.assertEqual(source.flat_map(expand).collect(), [Row(True, 7, 3, 7)]) + plain = source.flat_map(expand, return_dtype=dtype) + decorated = source.flat_map(pf.udtf(expand, return_dtype=dtype)) + self.assertEqual(plain.columns, ["is_dict", "first", "by_name"]) + self.assertEqual(decorated.columns, plain.columns) + self.assertEqual(plain.union_all(decorated).collect(), [Row(True, 7, 3)] * 2) def test_wrapper_row_input_reuse_and_column_calls(self): @pf.udtf def expand(value) -> Iterator[int]: - value = value[0] if isinstance(value, Row) else value + value = next(iter(value.values())) if isinstance(value, dict) else value yield value yield value + 1 @@ -491,7 +583,7 @@ def expand(value) -> Iterator[int]: combined = combined.union_all(lateral) self.assertCountEqual(combined.collect(), [Row(1), Row(2)] * 3 + [Row(3), Row(4)]) - def test_table_function_and_callable_instance(self): + def test_table_function_and_callable_classes_and_instances(self): client_pid = os.getpid() class Expand(TableFunction): @@ -501,10 +593,10 @@ def __init__(self): def open(self, context): self.offset = 10 - def eval(self, row: Row) -> Iterator[int]: + def eval(self, row: Dict[str, Any]) -> Iterator[int]: if self.constructor_pid != client_pid: raise AssertionError("TableFunction must be constructed on the client") - yield row[0] + self.offset + yield row["x"] + self.offset class Repeat: def __call__(self, row: Dict[str, Any]) -> Iterator[int]: @@ -514,9 +606,57 @@ def __call__(self, row: Dict[str, Any]) -> Iterator[int]: source = pf.from_dict({"x": [1]}) first = source.flat_map(pf.udtf(Expand)) second = source.flat_map(pf.udtf(Expand())) - third = source.flat_map(Repeat()) - self.assertCountEqual(first.union_all(second).union_all(third).collect(), - [Row(11), Row(11), Row(1), Row(1)]) + combined = first.union_all(second) + for func in (Repeat, Repeat(), pf.udtf(Repeat), pf.udtf(Repeat())): + combined = combined.union_all(source.flat_map(func)) + self.assertCountEqual(combined.collect(), [Row(11)] * 2 + [Row(1)] * 8) + + def test_wrapped_callable_class_methods(self): + def transparent(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + return wrapper + + class InstanceMethod: + @transparent + def __call__(self, row: Dict[str, Any]) -> Iterator[int]: + yield row["x"] + + class StaticMethod: + @staticmethod + @transparent + def __call__(row: Dict[str, Any]) -> Iterator[int]: + yield row["x"] + + class ClassMethod: + @classmethod + @transparent + def __call__(cls, row: Dict[str, Any]) -> Iterator[int]: + yield row["x"] + + class InheritedStaticMethod(StaticMethod): + pass + + def expand(row: Dict[str, Any]) -> Iterator[int]: + yield row["x"] + + class WrappedFunction: + @functools.wraps(expand) + def __call__(self, *args, **kwargs): + return expand(*args, **kwargs) + + source = pf.from_dict({"x": [1]}) + combined = None + for cls in (InstanceMethod, StaticMethod, ClassMethod, + InheritedStaticMethod, WrappedFunction): + self.assertEqual(list(cls()({"x": 1})), [1]) + for func in (cls(), pf.udtf(cls())): + self.assertEqual(source.flat_map(func).columns, ["f0"]) + for func in (cls, pf.udtf(cls)): + result = source.flat_map(func) + combined = result if combined is None else combined.union_all(result) + self.assertCountEqual(combined.collect(), [Row(1)] * 10) def test_nested_values_and_nulls(self): dtype = pf.DataType.struct({ diff --git a/flink-python/pyflink/dataframe/udtf.py b/flink-python/pyflink/dataframe/udtf.py index 93a678422f689..7fc7880e2bd4d 100644 --- a/flink-python/pyflink/dataframe/udtf.py +++ b/flink-python/pyflink/dataframe/udtf.py @@ -19,6 +19,7 @@ """User-defined table functions for the DataFrame API.""" import collections.abc +import functools import inspect import types from dataclasses import dataclass @@ -37,8 +38,10 @@ _create_declaration_context, _create_result_normalizer, _data_type_from_type_hint, + _get_callable_inspection_target, _get_callable_return_type_hint, _is_typed_dict, + _preserves_method_binding, _resolve_callable_annotation, _resolve_udf, _validate_determinism_agreement, @@ -46,7 +49,7 @@ ) from pyflink.table.expression import Expression from pyflink.table.expressions import col, with_columns -from pyflink.table.types import RowType +from pyflink.table.types import RowType, _to_java_data_type from pyflink.table.udf import ( TableFunction, UserDefinedFunction, @@ -57,7 +60,7 @@ __all__ = ["udtf"] -_UDTFInput = Union[Callable[..., Any], TableFunction, Type[TableFunction]] +_UDTFInput = Union[Callable[..., Any], TableFunction, Type] @dataclass(frozen=True) @@ -83,7 +86,7 @@ def alias(self, name: str, *extra_names: str) -> "_DataFrameUDTFCall": @dataclass(frozen=True) class _DataFrameUDTFWrapper: - _func: Union[Callable[..., Any], TableFunction] + _func: _UDTFInput _declaration_context: _UDFDeclarationContext return_dtype: DataType _has_named_fields: bool @@ -100,18 +103,22 @@ def __call__(self, *args: Any) -> _DataFrameUDTFCall: ) def _create_table_wrapper( - self, input_columns: Optional[Tuple[str, ...]] = None, *, input_as_dict: bool = False + self, input_columns: Optional[Tuple[str, ...]] = None, *, + preserve_field_names: bool = False, ) -> UserDefinedTableFunctionWrapper: """Keep Table's mutable row-input flag local to each use of this declaration.""" + result_types = self.return_dtype._to_table_data_type() + if preserve_field_names and isinstance(result_types, RowType): + # Table's RowType path drops field names; its SQL type declaration preserves them. + result_types = _to_java_data_type(result_types).getLogicalType().asSerializableString() return cast(UserDefinedTableFunctionWrapper, table_udtf( _DataFrameTableFunctionAdapter( self._func, self.return_dtype, self._deterministic, input_columns, - input_as_dict, ), - result_types=self.return_dtype._to_table_data_type(), + result_types=result_types, deterministic=self._deterministic, name=self._name, )) @@ -150,34 +157,58 @@ def udtf( """ Declare a synchronous function that emits zero or more rows per invocation. - Supports functions, callable instances, and ``TableFunction`` instances or - classes with a zero-argument constructor. Classes are instantiated on the client - when declared. The resulting or explicitly provided instance is serialized with - the job. ``TableFunction.open`` and - ``close`` run on workers, making ``open`` suitable for resource initialization. + Supports functions, callable classes and instances, and ``TableFunction`` + classes and instances. Classes must have a zero-argument constructor. Plain + callable classes are instantiated on workers, as with :func:`pyflink.dataframe.udf`. + ``TableFunction`` classes are instantiated on the client when declared; + their ``open`` and ``close`` methods run on workers. Explicitly provided + instances are serialized with the job. The output type may be explicit or inferred from ``Iterator[T]``, ``Iterable[T]``, ``Generator[T, ...]``, or ``list[T]``. A ``TypedDict`` supplies output field names. Return ``None`` for no rows, a list/generator for multiple rows, or a scalar, tuple, ``Row``, or dict for one row. Nested output values follow scalar UDF rules. - A declaration passed to :meth:`DataFrame.flat_map` receives a named Flink ``Row``; - an undecorated callable passed to that method receives a column-name dictionary. + When used with :meth:`DataFrame.flat_map`, the function receives a dictionary + keyed by input column name. Expression and SQL calls pass the specified arguments. Example:: - >>> from typing import Iterator, TypedDict + >>> from typing import Any, Dict, Iterator, TypedDict >>> import pyflink.dataframe as pf - >>> from pyflink.common import Row >>> class Output(TypedDict): ... value: int >>> @pf.udtf - ... def expand(row: Row) -> Iterator[Output]: + ... def expand(row: Dict[str, Any]) -> Iterator[Output]: ... yield {"value": row["x"]} ... yield {"value": row["x"] + 1} >>> result = pf.from_dict({"x": [1, 2]}).flat_map(expand) - :param func: Function, callable instance, or ``TableFunction`` instance/class. + Callable classes and instances:: + + >>> class Repeat: + ... def __call__(self, row: Dict[str, Any]) -> Iterator[int]: + ... yield row["x"] + ... yield row["x"] + >>> df = pf.from_dict({"x": [1, 2]}) + >>> result = df.flat_map(Repeat) + >>> result = df.flat_map(pf.udtf(Repeat)) + >>> result = df.flat_map(pf.udtf(Repeat())) + >>> result.columns + ['f0'] + + ``TableFunction`` classes can initialize worker resources in ``open``:: + + >>> from pyflink.table.udf import TableFunction + >>> class Expand(TableFunction): + ... def open(self, context): + ... self.offset = 1 + ... def eval(self, row: Dict[str, Any]) -> Iterator[int]: + ... yield row["x"] + self.offset + >>> result = df.flat_map(pf.udtf(Expand)) + >>> result = df.flat_map(pf.udtf(Expand())) + + :param func: Function, callable class/instance, or ``TableFunction`` class/instance. :param return_dtype: Emitted row type, as a DataFrame DataType, Python type, or SQL type string. Inferred from annotations when omitted. :param deterministic: Whether equal inputs produce equal results; must agree @@ -202,10 +233,8 @@ def decorator(f: _UDTFInput) -> _DataFrameUDTFWrapper: def _resolve_udtf( func: _UDTFInput, -) -> Tuple[Union[Callable[..., Any], TableFunction], _UDFDeclarationContext]: - if inspect.isclass(func): - if not issubclass(func, TableFunction): - raise TypeError("UDTF classes must extend TableFunction; pass a callable instance.") +) -> Tuple[_UDTFInput, _UDFDeclarationContext]: + if inspect.isclass(func) and issubclass(func, TableFunction): _validate_zero_argument_class(func) func = func() if not isinstance(func, TableFunction): @@ -215,7 +244,9 @@ def _resolve_udtf( raise TypeError("TableFunction.eval must be callable.") context = _create_declaration_context(func.eval, partial_source=func.eval) else: - if isinstance(func, UserDefinedFunction): + if isinstance(func, UserDefinedFunction) or ( + inspect.isclass(func) and issubclass(func, UserDefinedFunction) + ): raise TypeError("func must be a table UDF or a Python callable.") context = _resolve_udf(func).declaration_context target = context.annotation_target @@ -266,17 +297,36 @@ def _infer_udtf_return_dtype( return dtype, has_named_fields -def _validate_flat_map_input(declaration: _DataFrameUDTFWrapper, raw_callable: bool) -> None: +def _validate_flat_map_input(declaration: _DataFrameUDTFWrapper) -> None: source = declaration._func target: Callable[..., Any] - if isinstance(source, TableFunction): + if inspect.isclass(source): + target = declaration._declaration_context.annotation_target + elif isinstance(source, TableFunction): target = source.eval - else: + elif isinstance(source, functools.partial): target = source + else: + target = _get_callable_inspection_target(source) + if inspect.ismethod(target) and hasattr(target, "__wrapped__") and not ( + _preserves_method_binding(target, declaration._declaration_context.defining_class) + ): + target = target.__func__ try: signature = inspect.signature(target) except (TypeError, ValueError): return + if inspect.isclass(source): + descriptor = inspect.getattr_static(source, "__call__") + binds_receiver = not isinstance(descriptor, staticmethod) + if hasattr(target, "__wrapped__"): + binds_receiver = binds_receiver and _preserves_method_binding( + target, declaration._declaration_context.defining_class) + parameters = list(signature.parameters.values()) + if binds_receiver and parameters and parameters[0].kind in ( + inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD + ): + signature = signature.replace(parameters=parameters[1:]) parameters = list(signature.parameters.values()) try: signature.bind(object()) @@ -288,37 +338,33 @@ def _validate_flat_map_input(declaration: _DataFrameUDTFWrapper, raw_callable: b if not positional: return hint = _resolve_callable_annotation(declaration._declaration_context, positional[0].name) - if not _is_flat_map_row_hint(hint, raw_callable): - row_type = "dict" if raw_callable else "Row" - raise ValueError(f"flat_map receives one {row_type} row argument, got annotation {hint}.") + if not _is_flat_map_row_hint(hint): + raise ValueError(f"flat_map receives one dict row argument, got annotation {hint}.") -def _is_flat_map_row_hint(hint: Any, raw_callable: bool) -> bool: +def _is_flat_map_row_hint(hint: Any) -> bool: if hint in (_UNRESOLVED_TYPE_HINT, Any, object): return True origin = get_origin(hint) if origin is Annotated: - return _is_flat_map_row_hint(get_args(hint)[0], raw_callable) + return _is_flat_map_row_hint(get_args(hint)[0]) if origin in (Union, getattr(types, "UnionType", Union)): - return any(_is_flat_map_row_hint(member, raw_callable) for member in get_args(hint)) + return any(_is_flat_map_row_hint(member) for member in get_args(hint)) if _is_typed_dict(hint): - return raw_callable - target = origin or hint - if not raw_callable and target is tuple: return True + target = origin or hint try: - return issubclass(dict if raw_callable else Row, target) + return issubclass(dict, target) except TypeError: # Leave annotations that cannot be checked at runtime to the user's type checker. return True def _resolve_flat_map_udtf( - func: Union[Callable[..., Any], _DataFrameUDTFWrapper], + func: Union[Callable[..., Any], Type, _DataFrameUDTFWrapper], return_dtype: Optional[_DataTypeLike], input_columns: List[str], ) -> Tuple[Expression, List[str]]: - raw_callable = not isinstance(func, _DataFrameUDTFWrapper) if isinstance(func, _DataFrameUDTFWrapper): if return_dtype is not None: raise ValueError("return_dtype must not be specified for a DataFrame UDTF declaration.") @@ -326,10 +372,12 @@ def _resolve_flat_map_udtf( else: if func is None: raise TypeError("flat_map requires a callable or a pf.udtf declaration.") - if inspect.isclass(func) or isinstance(func, UserDefinedFunction): - raise TypeError("flat_map accepts a callable instance or a pf.udtf declaration.") + if isinstance(func, UserDefinedFunction) or ( + inspect.isclass(func) and issubclass(func, UserDefinedFunction) + ): + raise TypeError("flat_map accepts Python callables or a pf.udtf declaration.") declaration = udtf(func, return_dtype=return_dtype) - _validate_flat_map_input(declaration, raw_callable) + _validate_flat_map_input(declaration) table_type = declaration.return_dtype._to_table_data_type() if declaration._has_named_fields: output_columns = cast(RowType, table_type).field_names() @@ -337,7 +385,7 @@ def _resolve_flat_map_udtf( raise ValueError("flat_map requires named output fields; use TypedDict or a named struct.") else: output_columns = ["f0"] - wrapper = declaration._create_table_wrapper(tuple(input_columns), input_as_dict=raw_callable) + wrapper = declaration._create_table_wrapper(tuple(input_columns)) wrapper._set_takes_row_as_input() return wrapper(with_columns(col("*"))), output_columns @@ -356,17 +404,15 @@ def _iter_user_results(result: Any) -> Iterator[Any]: class _DataFrameTableFunctionAdapter(TableFunction): def __init__( self, - func: Union[Callable[..., Any], TableFunction], + func: _UDTFInput, return_dtype: DataType, deterministic: bool, input_columns: Optional[Tuple[str, ...]], - input_as_dict: bool, ) -> None: self._func = func self._return_dtype = return_dtype self._deterministic = deterministic self._input_columns = input_columns - self._input_as_dict = input_as_dict self._bound_invocation: Optional[Callable[..., Any]] = None self._lifecycle_opened = False self.__name__ = getattr(func, "__name__", type(func).__name__) @@ -377,7 +423,9 @@ def open(self, function_context: Any) -> None: self._lifecycle_opened = True invoke_func = self._func.eval else: - invoke_func = self._func + invoke_func = self._func() if inspect.isclass(self._func) else self._func + if not callable(invoke_func): + raise TypeError("UDTF class must construct a callable instance.") try: self._bound_invocation = self._bind_func(invoke_func) except Exception: @@ -410,15 +458,7 @@ def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[..., Any]: def invoke(*args: Any) -> Iterator[Row]: if input_columns is not None: input_row = args[0] - row: Union[dict, Row] - if self._input_as_dict: - row = {name: input_row[i] for i, name in enumerate(input_columns)} - else: - # Thread mode supplies a tuple; expose the same named Row in both modes. - row = Row(*input_row) - row.set_field_names(input_columns) - if isinstance(input_row, Row): - row.set_row_kind(input_row.get_row_kind()) + row = {name: input_row[i] for i, name in enumerate(input_columns)} result = invoke_func(row) else: result = invoke_func(*args)