This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 8161afa99bb [FLINK-40434][python] Add flat_map support to the
DataFrame API
8161afa99bb is described below
commit 8161afa99bb4ede160cf3efdbb4d9d38e7c9d666
Author: victor <[email protected]>
AuthorDate: Fri Sep 18 12:32:24 2026 +0800
[FLINK-40434][python] Add flat_map support to the DataFrame API
This closes #29234.
---
.../docs/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 | 88 +++
flink-python/pyflink/dataframe/sql.py | 44 +-
flink-python/pyflink/dataframe/tests/test_sql.py | 94 +++
flink-python/pyflink/dataframe/tests/test_udtf.py | 708 +++++++++++++++++++++
flink-python/pyflink/dataframe/udtf.py | 518 +++++++++++++++
8 files changed, 1460 insertions(+), 17 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index ed02294ce9c..9c65ebdef80 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -74,6 +74,7 @@ Transformations
DataFrame.limit
DataFrame.offset
DataFrame.head
+ DataFrame.flat_map
DataFrame.__getitem__
DataFrame.__getattr__
diff --git a/flink-python/docs/reference/pyflink.dataframe/udf.rst
b/flink-python/docs/reference/pyflink.dataframe/udf.rst
index c0084bee45d..798297eea33 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 general synchronous and
asynchronous callables,
and synchronous pandas or Arrow 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 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.
+
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 c8ed1cf08ae..bce542f48c2 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -57,6 +57,7 @@ from pyflink.dataframe.datatype import DataType
from pyflink.dataframe.io import read_generic
from pyflink.dataframe.sql import sql
from pyflink.dataframe.udf import udf
+from pyflink.dataframe.udtf import udtf
__all__ = [
"DataFrame",
@@ -65,6 +66,7 @@ __all__ = [
"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 bdfacc22480..0341b1c8f8c 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -27,6 +27,7 @@ from typing import (
Optional,
Set,
Tuple,
+ Type,
TypeVar,
Union,
overload,
@@ -34,6 +35,8 @@ from typing import (
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
@@ -623,6 +626,91 @@ class DataFrame:
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(record: Dict[str, Any]) -> Iterator[Token]:
+ ... for word in record["text"].split():
+ ... yield {"word": word}
+ >>> df = pf.from_dict({"text": ["hello world", "flink"]})
+ >>> result = df.flat_map(split)
+ >>> result.columns
+ ['word']
+
+ The decorator is optional for plain callables. Use it to attach
reusable
+ metadata, such as the output schema, instead of repeating it in each
+ ``flat_map`` call::
+
+ >>> @pf.udtf(return_dtype="ROW<word STRING>")
+ ... def tokenize(record: Dict[str, Any]):
+ ... yield from record["text"].split()
+ >>> result = df.flat_map(tokenize)
+ >>> result.columns
+ ['word']
+
+ An explicit output type can be supplied for unannotated callables::
+
+ >>> words = df.flat_map(lambda record: record["text"].split(),
return_dtype=str)
+ >>> words.columns
+ ['f0']
+ >>> named = df.flat_map(
+ ... lambda record: record["text"].split(),
+ ... return_dtype="ROW<word STRING>")
+ >>> named.columns
+ ['word']
+
+ Callable classes can be passed directly and are instantiated on
workers::
+
+ >>> class SplitWords:
+ ... def __call__(self, record: Dict[str, Any]) ->
Iterator[str]:
+ ... yield from record["text"].split()
+ >>> words = df.flat_map(SplitWords)
+
+ ``TableFunction`` classes are declared with
:func:`pyflink.dataframe.udtf`::
+
+ >>> from pyflink.table.udf import TableFunction
+ >>> class SplitWordsFunction(TableFunction):
+ ... def eval(self, record: Dict[str, Any]) -> Iterator[str]:
+ ... yield from record["text"].split()
+ >>> words = df.flat_map(pf.udtf(SplitWordsFunction))
+
+ See :func:`pyflink.dataframe.udtf` for more details.
+
+ .. 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 879a060b29c..9a9e11e2134 100644
--- a/flink-python/pyflink/dataframe/sql.py
+++ b/flink-python/pyflink/dataframe/sql.py
@@ -26,9 +26,11 @@ from py4j.protocol import Py4JJavaError
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 @@ __all__ = ["sql"]
_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 bcc7b1190b7..be8d9033c94 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 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
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<word STRING, size BIGINT>",
+ 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
new file mode 100644
index 00000000000..80623aaaae1
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_udtf.py
@@ -0,0 +1,708 @@
+################################################################################
+# 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
+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 NotCallable:
+ pass
+
+ class RequiresCallableArgument:
+ def __init__(self, count):
+ self.count = count
+
+ def __call__(self, value):
+ return [value] * self.count
+
+ class RequiresArgument(TableFunction):
+ def __init__(self, value):
+ self.value = value
+
+ def eval(self, value):
+ return [value]
+
+ 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]):
+ 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, 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 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(keyword)):
+ 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 (Row, Tuple[int], 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_passes_dict_without_mutating_input(self):
+ received = []
+
+ @pf.udtf
+ 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(original_z=7, original_x=3)
+ for value in ((7, 3), original):
+ self.assertEqual(list(adapter.eval(value)), [Row(3)])
+ 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):
+ 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
+ for func in (expand, pf.udtf(expand)):
+ with self.subTest(hint=hint, func=func):
+ self.assertEqual(source.flat_map(func).columns, ["f0"])
+
+ 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
+
+ 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]]:
+ 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<value BIGINT, label
STRING>")
+ 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: Dict[str, Any]) -> Iterator[int]:
+ yield next(iter(row.values()))
+
+ 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_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(),
+ })
+
+ 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"])
+ 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 = next(iter(value.values())) if isinstance(value, dict) 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_classes_and_instances(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: Dict[str, Any]) -> Iterator[int]:
+ if self.constructor_pid != client_pid:
+ raise AssertionError("TableFunction must be constructed on
the client")
+ yield row["x"] + 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()))
+ 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({
+ "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 00000000000..f1b92a1839e
--- /dev/null
+++ b/flink-python/pyflink/dataframe/udtf.py
@@ -0,0 +1,518 @@
+################################################################################
+# 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 functools
+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_inspection_target,
+ _get_callable_return_type_hint,
+ _is_typed_dict,
+ _preserves_method_binding,
+ _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, _to_java_data_type
+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]
+
+
+@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: _UDTFInput
+ _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, *,
+ 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,
+ ),
+ result_types=result_types,
+ 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 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.
+
+ 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 Any, Dict, Iterator, TypedDict
+ >>> import pyflink.dataframe as pf
+
+ >>> class Output(TypedDict):
+ ... value: int
+
+ >>> @pf.udtf
+ ... def expand(record: Dict[str, Any]) -> Iterator[Output]:
+ ... yield {"value": record["x"]}
+ ... yield {"value": record["x"] + 1}
+
+ >>> df = pf.from_dict({"x": [1, 2]})
+ >>> result = df.flat_map(expand)
+
+ Pass a callable class directly for a one-off ``flat_map`` operation. The
class
+ must support zero-argument construction and is instantiated on workers::
+
+ >>> class Repeat:
+ ... def __init__(self, times=2):
+ ... self.times = times
+ ...
+ ... def __call__(self, record: Dict[str, Any]):
+ ... for _ in range(self.times):
+ ... yield record["x"]
+
+ >>> repeated = df.flat_map(Repeat, return_dtype=int)
+ >>> repeated.columns
+ ['f0']
+
+ Create a UDTF declaration when the callable will be reused or its metadata
should
+ be declared once. Passing the class keeps worker-side construction::
+
+ >>> repeat = pf.udtf(Repeat, return_dtype=int)
+ >>> first = df.flat_map(repeat)
+ >>> second = df.flat_map(repeat)
+
+ When constructor arguments are needed, create an instance on the client.
Pass it
+ directly for one-off use, or wrap it with :func:`udtf` to reuse the
declaration
+ without repeating its metadata::
+
+ >>> repeat_three_times = Repeat(3)
+ >>> one_off = df.flat_map(repeat_three_times, return_dtype=int)
+ >>> reusable_repeat_three_times = pf.udtf(repeat_three_times,
return_dtype=int)
+ >>> reused = df.flat_map(reusable_repeat_three_times)
+
+ A ``TableFunction`` class is instantiated on the client when declared. Use
an
+ instance for constructor configuration; in both cases, ``open`` runs on
workers::
+
+ >>> from pyflink.table.udf import TableFunction
+
+ >>> class Expand(TableFunction):
+ ... def __init__(self, offset=1):
+ ... self.offset = offset
+ ...
+ ... def open(self, context):
+ ... self.worker_offset = self.offset
+ ...
+ ... def eval(self, record: Dict[str, Any]) -> Iterator[int]:
+ ... yield record["x"] + self.worker_offset
+
+ >>> default_expand = pf.udtf(Expand)
+ >>> configured_expand = pf.udtf(Expand(offset=10))
+ >>> default_result = df.flat_map(default_expand)
+ >>> configured_result = df.flat_map(configured_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
+ 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[_UDTFInput, _UDFDeclarationContext]:
+ if inspect.isclass(func) and issubclass(func, TableFunction):
+ _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) 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
+ 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) -> None:
+ source = declaration._func
+ target: Callable[..., Any]
+ if inspect.isclass(source):
+ target = declaration._declaration_context.annotation_target
+ elif isinstance(source, TableFunction):
+ target = source.eval
+ 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())
+ 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):
+ raise ValueError(f"flat_map receives one dict row argument, got
annotation {hint}.")
+
+
+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])
+ if origin in (Union, getattr(types, "UnionType", Union)):
+ return any(_is_flat_map_row_hint(member) for member in get_args(hint))
+ if _is_typed_dict(hint):
+ return True
+ target = origin or hint
+ try:
+ 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], Type, _DataFrameUDTFWrapper],
+ return_dtype: Optional[_DataTypeLike],
+ input_columns: List[str],
+) -> Tuple[Expression, List[str]]:
+ 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 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)
+ 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))
+ 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: _UDTFInput,
+ return_dtype: DataType,
+ deterministic: bool,
+ input_columns: Optional[Tuple[str, ...]],
+ ) -> None:
+ self._func = func
+ self._return_dtype = return_dtype
+ self._deterministic = deterministic
+ self._input_columns = input_columns
+ 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() 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:
+ 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 = {name: input_row[i] for i, name in
enumerate(input_columns)}
+ 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)
