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 1465b13a4d1 [FLINK-40193][python] Add join_lateral support to the
DataFrame API (#29433)
1465b13a4d1 is described below
commit 1465b13a4d16b381464d7cf4b0f6d50c7d2209ad
Author: Liu Liu <[email protected]>
AuthorDate: Fri Oct 9 17:53:48 2026 +0800
[FLINK-40193][python] Add join_lateral support to the DataFrame API (#29433)
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 1 +
flink-python/pyflink/dataframe/dataframe.py | 72 ++++++++-
flink-python/pyflink/dataframe/tests/test_udtf.py | 162 ++++++++++++++++++++-
flink-python/pyflink/dataframe/udtf.py | 47 +++++-
4 files changed, 274 insertions(+), 8 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index 9f0f7764dcc..ec7a33281e4 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -103,6 +103,7 @@ Joins
:toctree: api/
DataFrame.join
+ DataFrame.join_lateral
Missing-Value Handling
----------------------
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index 8418a90dad1..3404a4d1861 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -40,7 +40,7 @@ if TYPE_CHECKING:
import pandas
import pyarrow
from pyflink.dataframe.udf import _DataTypeLike
- from pyflink.dataframe.udtf import _DataFrameUDTFWrapper
+ from pyflink.dataframe.udtf import _DataFrameUDTFCall,
_DataFrameUDTFWrapper
from pyflink.table.table_environment import TableEnvironment
from pyflink.table.table_schema import TableSchema
@@ -1047,6 +1047,76 @@ class DataFrame:
)
)
+ @PublicEvolving()
+ def join_lateral(
+ self,
+ table_function_call: Union["_DataFrameUDTFCall", Expression],
+ *,
+ on: Optional[Expression[bool]] = None,
+ ignore_empty: bool = True,
+ ) -> "DataFrame":
+ """
+ Join each input row with the rows emitted by a table function call.
+
+ The function receives the explicit arguments supplied in the call. The
result
+ retains the input columns followed by the function's output columns.
Output
+ names come from ``alias`` or a named ``TypedDict`` or struct
``return_dtype``.
+ Raw Table API expressions must specify output names with ``alias``.
Output
+ names must be non-empty, unique, and distinct from the input column
names.
+
+ :param table_function_call: A call to a function declared with
+ :func:`pyflink.dataframe.udtf`, or a Table API table function
expression.
+ :param on: Optional boolean join predicate. Left outer lateral joins
currently
+ require an omitted or literal ``True`` predicate, as in the Table
API.
+ :param ignore_empty: If ``True`` (default), use an inner lateral join.
If
+ ``False``, use a left outer lateral join, preserving input rows
with no
+ emitted rows and filling the output columns with NULLs.
+ :return: A new DataFrame containing the lateral join result.
+
+ Example::
+
+ >>> from typing import Iterator, TypedDict
+ >>> import pyflink.dataframe as pf
+ >>> @pf.udtf
+ ... def chars(text: str) -> Iterator[str]:
+ ... yield from text
+ >>> df = pf.from_records([(1, "ab"), (2, "")], schema=["id",
"text"])
+ >>> expanded = df.join_lateral(chars(pf.col("text")).alias("ch"))
+ >>> preserved = df.join_lateral(
+ ... chars(pf.col("text")).alias("ch"), ignore_empty=False)
+ >>> filtered = df.join_lateral(
+ ... chars(pf.col("text")).alias("ch"), on=pf.col("id") > 0)
+
+ Named multi-field outputs do not require aliases::
+
+ >>> class Token(TypedDict):
+ ... word: str
+ ... length: int
+ >>> @pf.udtf
+ ... def tokenize(text: str) -> Iterator[Token]:
+ ... for word in text.split():
+ ... yield {"word": word, "length": len(word)}
+ >>> tokens = df.join_lateral(tokenize(pf.col("text")))
+
+ .. versionadded:: 2.4.0
+ """
+ from pyflink.dataframe.udtf import _resolve_lateral_expression
+
+ if not isinstance(ignore_empty, bool):
+ raise TypeError("ignore_empty must be a bool")
+ if on is not None and not isinstance(on, Expression):
+ raise TypeError("on must be a boolean expression")
+ expression = _resolve_lateral_expression(table_function_call,
self.columns)
+ if ignore_empty:
+ table = self._table.join_lateral(expression)
+ if on is not None:
+ # Filtering an inner lateral join is equivalent to its ON
predicate and
+ # avoids the Table API's equi-join validation for ordinary
joins.
+ table = table.filter(on)
+ else:
+ table = self._table.left_outer_join_lateral(expression, on)
+ return DataFrame(table)
+
# ======================== Filtering & Ordering ========================
@PublicEvolving()
diff --git a/flink-python/pyflink/dataframe/tests/test_udtf.py
b/flink-python/pyflink/dataframe/tests/test_udtf.py
index 711050be7c7..4ea27607be8 100644
--- a/flink-python/pyflink/dataframe/tests/test_udtf.py
+++ b/flink-python/pyflink/dataframe/tests/test_udtf.py
@@ -27,11 +27,15 @@ from typing import (
)
import cloudpickle
+from py4j.protocol import Py4JJavaError
+
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.table.expression import Expression
+from pyflink.table.expressions import call as table_call
+from pyflink.table.udf import ScalarFunction, TableFunction, udtf as table_udtf
from pyflink.testing.test_case_utils import (
PyFlinkBatchTableTestCase,
PyFlinkDataFrameUTTestCase,
@@ -527,13 +531,117 @@ class
DataFrameUDTFPlanningTests(PyFlinkDataFrameUTTestCase):
self.assertEqual(events, ["init"])
-class _DataFrameFlatMapTests:
+class DataFrameLateralJoinPlanningTests(PyFlinkDataFrameUTTestCase):
+ def test_output_names_and_call_reuse(self):
+ source = pf.from_dict({"x": [1]})
+
+ @pf.udtf
+ def named(value) -> Iterator[_Output]:
+ yield {"value": value, "label": "ok"}
+
+ call = named(pf.col("x"))
+ inferred_types = [DataTypes.BIGINT().not_null(),
DataTypes.STRING().not_null()]
+ for function_call, names, output_types in (
+ (call, ["value", "label"], inferred_types),
+ (call.alias("v", "l"), ["v", "l"], inferred_types),
+ (pf.udtf(lambda x: [(x, "ok")], return_dtype="ROW<v BIGINT, l
STRING>")(
+ pf.col("x")), ["v", "l"], [DataTypes.BIGINT(),
DataTypes.STRING()]),
+ ):
+ with self.subTest(names=names):
+
self.assert_dataframe_schema(source.join_lateral(function_call), ["x"] + names,
+ [DataTypes.BIGINT()] +
output_types)
+ self.assertIsNone(call.output_aliases)
+ self.assert_dataframe_schema(source.join_lateral(call), ["x", "value",
"label"])
+
+ @pf.udtf
+ def pair(value, label: str) -> Iterator[Tuple[int, str]]:
+ yield value, label
+
+ result = source.join_lateral(pair(pf.col("x"),
pf.lit("ok")).alias("v", "l"))
+ self.assertEqual(result.columns, ["x", "v", "l"])
+
+ def test_raw_expression_aliases(self):
+ source = pf.from_dict({"x": [1]})
+ emit = table_udtf(lambda x: [x], result_types=[DataTypes.BIGINT()])
+ self.t_env.create_temporary_system_function("emit", emit)
+ for expression in (emit(pf.col("x")).alias("out"),
+ table_call("emit", pf.col("x")).alias("out")):
+ with self.subTest(expression=str(expression)):
+ # Alias names belong to the expression tree, not a Python
wrapper.
+ rebuilt = Expression(expression._j_expr)
+ self.assert_dataframe_schema(source.join_lateral(rebuilt),
["x", "out"],
+ [DataTypes.BIGINT(),
DataTypes.BIGINT()])
+ with self.assertRaisesRegex(ValueError, "require alias"):
+ source.join_lateral(emit(pf.col("x")))
+
+ def test_invalid_calls_and_output_names(self):
+ source = pf.from_dict({"x": [1]})
+ scalar = pf.udtf(lambda x: [x], return_dtype=int)
+ unnamed = scalar(pf.col("x"))
+ named = pf.udtf(lambda x: [(x, "ok")], return_dtype="ROW<x BIGINT,
label STRING>")
+ raw = table_udtf(lambda x: [(x, "ok")],
+ result_types=[DataTypes.BIGINT(),
DataTypes.STRING()])(pf.col("x"))
+ cases = (
+ (unnamed, ValueError, "require output names"),
+ (unnamed.alias("x"), ValueError, "conflict"),
+ (named(pf.col("x")), ValueError, "conflict"),
+ (raw.alias("same", "same"), ValueError, "unique"),
+ (raw.alias("", "label"), ValueError, "non-empty"),
+ (raw.alias("x", "label"), ValueError, "conflict"),
+ (None, TypeError, "table_function_call"),
+ (scalar, TypeError, "table_function_call"),
+ )
+ for function_call, error, message in cases:
+ with self.subTest(message=message, function_call=function_call):
+ with self.assertRaisesRegex(error, message):
+ source.join_lateral(function_call)
+
+ def test_option_types_and_predicate_validation(self):
+ source = pf.from_dict({"x": [1]})
+ emit = pf.udtf(lambda x: [x],
return_dtype=int)(pf.col("x")).alias("out")
+ for kwargs in ({"ignore_empty": 0}, {"ignore_empty": None}, {"on":
True}, {"on": "x"}):
+ with self.subTest(kwargs=kwargs), self.assertRaises(TypeError):
+ source.join_lateral(emit, **kwargs)
+ for predicate in (pf.lit(1), pf.col("missing") > 0):
+ with self.subTest(predicate=str(predicate)),
self.assertRaisesRegex(
+ Py4JJavaError, "ValidationException"):
+ source.join_lateral(emit, on=predicate)
+
+ def test_predicates_plan_in_stream_and_batch(self):
+ from pyflink.datastream import RuntimeExecutionMode
+
+ emit = pf.udtf(lambda x: [x], return_dtype=int)
+ positive = pf.udf(lambda x: x > 0, return_dtype=bool)
+ for mode in (RuntimeExecutionMode.STREAMING,
RuntimeExecutionMode.BATCH):
+ self.env.set_runtime_mode(mode)
+ source = pf.from_dict({"x": [1]})
+ call = emit(pf.col("x") + 1).alias("out")
+ for predicate in (None, pf.col("x") > 0, pf.col("out") >
pf.col("x"),
+ positive(pf.col("out"))):
+ with self.subTest(mode=mode, predicate=str(predicate)):
+ self.assertIn("PythonCorrelate", source.join_lateral(call,
on=predicate)
+ .to_table().explain())
+ for predicate in (None, pf.lit(True)):
+ with self.subTest(mode=mode, predicate=str(predicate)):
+ self.assertIn("PythonCorrelate", source.join_lateral(
+ call, on=predicate,
ignore_empty=False).to_table().explain())
+
+ def test_left_predicate_retains_table_api_restriction(self):
+ source = pf.from_dict({"x": [1]})
+ emit = pf.udtf(lambda x: [x],
return_dtype=int)(pf.col("x")).alias("out")
+ with self.assertRaisesRegex(Py4JJavaError, "can only be empty or
literal true"):
+ source.join_lateral(emit, on=pf.col("out") > 0, ignore_empty=False)
+
+
+class _DataFrameUDTFTests:
def setUp(self):
super().setUp()
previous = pf.get_table_environment()
self.addCleanup(pf.set_table_environment, previous)
pf.set_table_environment(self.t_env)
+
+class _DataFrameFlatMapTests(_DataFrameUDTFTests):
def test_typed_dict_input_and_named_output_pipeline(self):
class Input(TypedDict):
count: int
@@ -583,7 +691,7 @@ class _DataFrameFlatMapTests:
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")
+ lateral = a.join_lateral(call).select("out")
combined = combined.union_all(lateral)
self.assertCountEqual(combined.collect(), [Row(1), Row(2)] * 3 +
[Row(3), Row(4)])
@@ -695,15 +803,57 @@ class _DataFrameFlatMapTests:
self.assertCountEqual(result.collect(), [Row({"value": 1}),
Row({"value": 2})])
-class DataFrameFlatMapStreamTests(_DataFrameFlatMapTests,
PyFlinkStreamDataFrameTestCase):
+class _DataFrameLateralJoinTests(_DataFrameUDTFTests):
+ def test_lateral_join_semantics_and_predicates(self):
+ @pf.udtf
+ def words(text: Optional[str]) -> Iterator[str]:
+ yield from (text or "").split()
+
+ class Token(TypedDict):
+ word: str
+ length: int
+
+ @pf.udtf
+ def tokenize(text: Optional[str]) -> Iterator[Token]:
+ for word in (text or "").split():
+ yield {"word": word, "length": len(word)}
+
+ source = pf.from_records([(0, ""), (1, "a bb"), (2, "z"), (-1,
"skip"), (3, None)],
+ schema=["id", "text"])
+ call = words(pf.col("text")).alias("word")
+ inner = source.join_lateral(call)
+ left = source.join_lateral(call, ignore_empty=False, on=pf.lit(True))
+ filtered = source.join_lateral(tokenize(pf.col("text")),
+ on=(pf.col("id") > 0) &
(pf.col("length") > 1))
+ raw = table_udtf(lambda text: (text or "").split(),
result_types=[DataTypes.STRING()])
+ legacy = source.join_lateral(raw(pf.col("text")).alias("word"))
+ frames = [("inner", inner), ("left", left), ("filtered", filtered),
("raw", legacy)]
+ combined = None
+ for name, frame in frames:
+ result = frame.select(pf.lit(name).alias("kind"), "id", "text",
"word")
+ combined = result if combined is None else
combined.union_all(result)
+ emitted = [(1, "a bb", "a"), (1, "a bb", "bb"), (2, "z", "z"), (-1,
"skip", "skip")]
+ expected = [Row(name, *row) for name in ("inner", "left", "raw") for
row in emitted]
+ expected += [Row("left", 0, "", None), Row("left", 3, None, None),
+ Row("filtered", 1, "a bb", "bb")]
+ self.assertCountEqual(combined.collect(), expected)
+
+
+class DataFrameUDTFStreamTests(
+ _DataFrameFlatMapTests, _DataFrameLateralJoinTests,
PyFlinkStreamDataFrameTestCase
+):
pass
-class DataFrameFlatMapBatchTests(_DataFrameFlatMapTests,
PyFlinkBatchTableTestCase):
+class DataFrameUDTFBatchTests(
+ _DataFrameFlatMapTests, _DataFrameLateralJoinTests,
PyFlinkBatchTableTestCase
+):
pass
-class DataFrameFlatMapThreadTests(_DataFrameFlatMapTests,
PyFlinkStreamDataFrameTestCase):
+class DataFrameUDTFThreadTests(
+ _DataFrameFlatMapTests, _DataFrameLateralJoinTests,
PyFlinkStreamDataFrameTestCase
+):
def setUp(self):
super().setUp()
self.t_env.get_config().set("python.execution-mode", "thread")
diff --git a/flink-python/pyflink/dataframe/udtf.py
b/flink-python/pyflink/dataframe/udtf.py
index 60c1ed4bf71..b63d504afe4 100644
--- a/flink-python/pyflink/dataframe/udtf.py
+++ b/flink-python/pyflink/dataframe/udtf.py
@@ -28,6 +28,8 @@ from typing import (
get_origin, overload,
)
+from py4j.java_gateway import get_java_class
+
from pyflink.common import Row
from pyflink.dataframe.datatype import DataType
from pyflink.dataframe.udf import (
@@ -46,6 +48,7 @@ from pyflink.dataframe.udf import (
_validate_determinism_agreement,
_validate_zero_argument_class,
)
+from pyflink.java_gateway import get_gateway
from pyflink.table.expression import Expression
from pyflink.table.expressions import col, with_columns
from pyflink.table.types import RowType, _to_java_data_type
@@ -169,7 +172,9 @@ def udtf(
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.
+ keyed by input column name. :meth:`DataFrame.join_lateral`, expression and
SQL
+ calls pass the specified arguments. Use ``alias`` to name lateral join
outputs,
+ or declare named fields with a ``TypedDict`` or struct ``return_dtype``.
Example::
@@ -420,6 +425,46 @@ def _resolve_flat_map_udtf(
return wrapper(with_columns(col("*"))), output_columns
+def _resolve_lateral_expression(
+ call: Union[_DataFrameUDTFCall, Expression], input_columns: List[str],
+) -> Expression:
+ if isinstance(call, _DataFrameUDTFCall):
+ expression = call.expression
+ if call.output_aliases is not None:
+ output_names = list(call.output_aliases)
+ elif call.has_named_fields:
+ output_names = cast(RowType,
call.return_dtype._to_table_data_type()).field_names()
+ expression = expression.alias(*output_names)
+ else:
+ raise ValueError(
+ "UDTF calls passed to join_lateral require output names; use
alias, "
+ "a TypedDict, or a named struct return_dtype.")
+ elif isinstance(call, Expression):
+ expression = call
+ gateway = get_gateway()
+ api_utils =
gateway.jvm.org.apache.flink.table.expressions.ApiExpressionUtils
+ definitions =
gateway.jvm.org.apache.flink.table.functions.BuiltInFunctionDefinitions
+ j_expression = api_utils.unwrapFromApi(expression._j_expr)
+ if not api_utils.isFunction(j_expression, definitions.AS):
+ raise ValueError("Expression arguments passed to join_lateral
require alias.")
+ string_class = get_java_class(gateway.jvm.String)
+ output_names = [
+ child.getValueAs(string_class).get() for child in
list(j_expression.getChildren())[1:]
+ ]
+ else:
+ raise TypeError(
+ "table_function_call must be a DataFrame UDTF call or Expression, "
+ f"got {type(call).__name__}")
+
+ if not output_names or not all(output_names) or len(set(output_names)) !=
len(output_names):
+ raise ValueError("join_lateral output names must be non-empty and
unique.")
+ conflicts = set(input_columns).intersection(output_names)
+ if conflicts:
+ raise ValueError(
+ f"join_lateral output names conflict with input columns:
{sorted(conflicts)}")
+ return expression
+
+
def _iter_user_results(result: Any) -> Iterator[Any]:
if result is None:
return
