[FLINK-40434][python] Add flat_map support to the DataFrame API

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

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 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

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",

flink-python/pyflink/dataframe/dataframe.py

@@ -26,13 +26,16 @@
     Optional,
     Set,
     Tuple,
+    Type,
     TypeVar,
     Union,
     overload,
 )
 
 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
@@ -622,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<word STRING>")
+            >>> 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()

flink-python/pyflink/dataframe/sql.py

@@ -26,21 +26,25 @@
 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
 
 __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

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<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