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 b3028978261 [FLINK-40419][python] Add regular join support to
DataFrame API
b3028978261 is described below
commit b3028978261ff2aa6630f41f7573402ab9c92a20
Author: SulphurFH <[email protected]>
AuthorDate: Wed Sep 16 10:05:38 2026 +0800
[FLINK-40419][python] Add regular join support to DataFrame API
This closes #29201.
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 10 +
flink-python/pyflink/dataframe/dataframe.py | 438 ++++++++++++-
.../pyflink/dataframe/tests/test_dataframe.py | 727 +++++++++++++++++++++
3 files changed, 1174 insertions(+), 1 deletion(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index 13eef577d3e..dbeadf80fe9 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -94,6 +94,16 @@ Set Operations
DataFrame.minus
DataFrame.minus_all
+Joins
+-----
+
+.. currentmodule:: pyflink.dataframe
+
+.. autosummary::
+ :toctree: api/
+
+ DataFrame.join
+
Aggregations
------------
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index c48f8af4397..a2694e562a0 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -18,6 +18,7 @@
import datetime
import keyword
+from contextlib import ExitStack
from typing import (
TYPE_CHECKING,
Any,
@@ -37,12 +38,13 @@ if TYPE_CHECKING:
import pandas
from pyflink.dataframe.udf import _DataTypeLike
from pyflink.dataframe.udtf import _DataFrameUDTFWrapper
+ from pyflink.table.table_environment import TableEnvironment
from pyflink.table.table_schema import TableSchema
from pyflink.common import Row
from pyflink.dataframe.datatype import _INT_MAX, DataType
from pyflink.java_gateway import get_gateway
-from pyflink.table.expression import Expression
+from pyflink.table.expression import Expression, _get_java_expression
from pyflink.table.expressions import (
and_,
call_sql,
@@ -53,6 +55,7 @@ from pyflink.table.table import Table
from pyflink.table.types import ArrayType, MapType, MultisetType, RowType
from pyflink.table.table_descriptor import TableDescriptor
from pyflink.util.api_stability_decorators import PublicEvolving
+from pyflink.util.java_utils import to_jarray
__all__ = ["DataFrame", "GroupedDataFrame", "col", "lit"]
@@ -887,6 +890,124 @@ class DataFrame:
)
return DataFrame(table._t_env.sql_query(query))
+ # ======================== Joins ========================
+
+ @PublicEvolving()
+ def join(
+ self,
+ other: "DataFrame",
+ *,
+ on=None,
+ how: str = "inner",
+ left_on=None,
+ right_on=None,
+ ) -> "DataFrame":
+ """
+ Join this DataFrame with another DataFrame.
+
+ Use ``on`` when both sides share the same named join keys, or pass a
boolean expression as
+ the complete join predicate. Use ``left_on`` and ``right_on`` together
when the key names
+ differ. In streaming mode, a join without equality keys may use
singleton distribution
+ (a single parallel instance) and can be expensive.
+ Shared named keys occur once in the result; other duplicate column
names must be renamed
+ before joining. ``semi`` and ``anti`` joins return only columns from
this DataFrame, while
+ ``cross`` performs a Cartesian product and accepts no join keys.
+
+ :param other: DataFrame on the right side of the join.
+ :param on: Shared column name, list of shared column names, or a
boolean join expression.
+ :param how: Join type: ``"inner"``, ``"left"``, ``"right"``,
``"full"``, ``"outer"``,
+ ``"semi"``, ``"anti"``, or ``"cross"``.
+ :param left_on: Column name, expression, or list of column names from
this DataFrame.
+ :param right_on: Column name, expression, or list of column names from
``other``.
+ :return: A new DataFrame containing the join result.
+ :raises TypeError: If an argument has an unsupported type.
+ :raises ValueError: If the join type, keys, schemas, or argument
combination is invalid.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> orders = pf.from_records([(1, 10)], schema=["customer_id",
"amount"])
+ >>> customers = pf.from_records([(1, "Alice")],
schema=["customer_id", "name"])
+ >>> matched = orders.join(customers, on="customer_id")
+ >>> matched = orders.join(
+ ... customers.rename_columns({"customer_id": "id"}),
+ ... left_on="customer_id",
+ ... right_on="id",
+ ... how="left",
+ ... )
+
+ Expression-based predicates support equality, compound conditions, and
non-equi joins::
+
+ >>> customers_by_id = customers.rename_columns({"customer_id":
"id"})
+ >>> matched = orders.join(
+ ... customers_by_id, on=pf.col("customer_id") == pf.col("id"))
+ >>> rules = pf.from_records([(1, 5)], schema=["rule_customer_id",
"min_amount"])
+ >>> matched = orders.join(
+ ... rules,
+ ... on=(pf.col("customer_id") == pf.col("rule_customer_id"))
+ ... & (pf.col("amount") >= pf.col("min_amount")),
+ ... )
+ >>> matched = orders.join(rules, on=pf.col("amount") >=
pf.col("min_amount"))
+ >>> unmatched = orders.join(
+ ... customers_by_id,
+ ... on=pf.col("customer_id") == pf.col("id"),
+ ... how="anti",
+ ... )
+
+ .. versionadded:: 2.4.0
+ """
+ if not isinstance(other, DataFrame):
+ raise TypeError("other must be a pyflink.dataframe.DataFrame")
+ if self._table._t_env._j_tenv != other._table._t_env._j_tenv:
+ raise ValueError("DataFrames must belong to the same
TableEnvironment")
+
+ join_type = _normalize_join_type(how)
+ if join_type == "cross":
+ if on is not None or left_on is not None or right_on is not None:
+ raise ValueError("cross join does not accept on, left_on, or
right_on")
+ _validate_join_column_conflicts(self.columns, other.columns, set())
+ return DataFrame(self._table.join(other._table))
+
+ (
+ left_table,
+ right_table,
+ predicate,
+ shared_keys,
+ ) = _prepare_join(
+ self._table,
+ other._table,
+ on,
+ left_on,
+ right_on,
+ validate_column_conflicts=join_type not in ("semi", "anti"),
+ )
+
+ with _JoinSqlFactory(self._table._t_env) as sql_factory:
+ if join_type in ("semi", "anti"):
+ return DataFrame(
+ _build_semi_anti_join_sql(
+ left_table,
+ right_table,
+ predicate,
+ self.columns,
+ join_type,
+ sql_factory,
+ )
+ )
+
+ return DataFrame(
+ _build_regular_join_sql(
+ left_table,
+ right_table,
+ predicate,
+ self.columns,
+ other.columns,
+ shared_keys,
+ join_type,
+ sql_factory,
+ )
+ )
+
# ======================== Filtering & Ordering ========================
@PublicEvolving()
@@ -2108,6 +2229,321 @@ class GroupedDataFrame:
# ======================== Internal Helpers ========================
+def _normalize_join_type(how: str) -> str:
+ if not isinstance(how, str):
+ raise TypeError("how must be a string")
+ aliases = {
+ "inner": "inner",
+ "left": "left",
+ "right": "right",
+ "full": "full",
+ "outer": "full",
+ "semi": "semi",
+ "anti": "anti",
+ "cross": "cross",
+ }
+ if how not in aliases:
+ raise ValueError(
+ 'how must be one of "inner", "left", "right", "full", "outer", '
+ '"semi", "anti", or "cross"'
+ )
+ return aliases[how]
+
+
+def _normalize_join_keys(value, parameter_name: str) -> List[Union[str,
Expression]]:
+ if isinstance(value, (str, Expression)):
+ return [value]
+ if isinstance(value, list):
+ if not value:
+ raise ValueError("%s must not be empty" % parameter_name)
+ if not all(isinstance(key, str) for key in value):
+ raise TypeError(
+ "%s must be a string, an expression, or a list of strings" %
parameter_name
+ )
+ if len(set(value)) != len(value):
+ raise ValueError("%s must not contain duplicate column names" %
parameter_name)
+ return value
+ raise TypeError(
+ "%s must be a string, an expression, or a list of strings" %
parameter_name
+ )
+
+
+def _validate_join_columns(
+ keys: List[Union[str, Expression]], columns: List[str], parameter_name: str
+) -> None:
+ for key in keys:
+ if isinstance(key, str) and key not in columns:
+ raise ValueError(
+ "%s column '%s' does not exist, available columns: %s"
+ % (parameter_name, key, columns)
+ )
+
+
+def _validate_join_column_conflicts(
+ left_columns: List[str], right_columns: List[str], shared_keys: Set[str]
+) -> None:
+ conflicts = sorted((set(left_columns) & set(right_columns)) - shared_keys)
+ if conflicts:
+ raise ValueError(
+ "join() found duplicate non-key columns %s; rename them with
rename_columns() "
+ "before joining" % conflicts
+ )
+
+
+def _prepare_join(
+ left_table: Table,
+ right_table: Table,
+ on,
+ left_on,
+ right_on,
+ *,
+ validate_column_conflicts: bool,
+) -> Tuple[Table, Table, Expression, Dict[str, str]]:
+ left_columns = list(left_table.get_resolved_schema().get_column_names())
+ right_columns = list(right_table.get_resolved_schema().get_column_names())
+
+ if on is not None:
+ if left_on is not None or right_on is not None:
+ raise ValueError("on cannot be combined with left_on or right_on")
+ if isinstance(on, Expression):
+ if validate_column_conflicts:
+ _validate_join_column_conflicts(left_columns, right_columns,
set())
+ return left_table, right_table, on, {}
+ left_keys = _normalize_join_keys(on, "on")
+ right_keys = list(left_keys)
+ _validate_join_columns(left_keys, left_columns, "on")
+ _validate_join_columns(right_keys, right_columns, "on")
+ else:
+ if left_on is None and right_on is None:
+ raise ValueError("join() requires on or both left_on and right_on")
+ if left_on is None or right_on is None:
+ raise ValueError("left_on and right_on must be provided together")
+ left_keys = _normalize_join_keys(left_on, "left_on")
+ right_keys = _normalize_join_keys(right_on, "right_on")
+ if len(left_keys) != len(right_keys):
+ raise ValueError("left_on and right_on must have the same number
of keys")
+ _validate_join_columns(left_keys, left_columns, "left_on")
+ _validate_join_columns(right_keys, right_columns, "right_on")
+
+ shared_names = {
+ left_key
+ for left_key, right_key in zip(left_keys, right_keys)
+ if isinstance(left_key, str)
+ and isinstance(right_key, str)
+ and left_key == right_key
+ }
+ if validate_column_conflicts:
+ _validate_join_column_conflicts(left_columns, right_columns,
shared_names)
+
+ taken = set(left_columns) | set(right_columns)
+ shared_keys: Dict[str, str] = {}
+ right_rename_expressions: List[Expression] = []
+ for name in left_columns:
+ if name in shared_names:
+ temporary_name = _unique_name("__pf_join_right_%s" % name, taken)
+ taken.add(temporary_name)
+ shared_keys[name] = temporary_name
+
right_rename_expressions.append(table_col(name).alias(temporary_name))
+ left_key_names: List[str] = []
+ right_key_names: List[str] = []
+ left_computed_keys: List[Expression] = []
+ right_computed_keys: List[Expression] = []
+ for index, (left_key, right_key) in enumerate(zip(left_keys, right_keys)):
+ if isinstance(left_key, str):
+ left_key_names.append(left_key)
+ else:
+ temporary_name = _unique_name("__pf_join_left_key_%d" % index,
taken)
+ taken.add(temporary_name)
+ left_key_names.append(temporary_name)
+ left_computed_keys.append(left_key.alias(temporary_name))
+
+ if isinstance(right_key, str):
+ right_key_names.append(shared_keys.get(right_key, right_key))
+ else:
+ temporary_name = _unique_name("__pf_join_right_key_%d" % index,
taken)
+ taken.add(temporary_name)
+ right_key_names.append(temporary_name)
+ right_computed_keys.append(right_key.alias(temporary_name))
+
+ if left_computed_keys:
+ left_table = left_table.add_columns(*left_computed_keys)
+ if right_computed_keys:
+ right_table = right_table.add_columns(*right_computed_keys)
+ if right_rename_expressions:
+ right_table = right_table.rename_columns(*right_rename_expressions)
+
+ conditions = [
+ table_col(left_name) == table_col(right_name)
+ for left_name, right_name in zip(left_key_names, right_key_names)
+ ]
+ predicate = conditions[0] if len(conditions) == 1 else and_(*conditions)
+ return (
+ left_table,
+ right_table,
+ predicate,
+ shared_keys,
+ )
+
+
+class _JoinSqlFactory:
+ """Keep inline UDFs registered until the join SQL has been resolved."""
+
+ def __init__(self, t_env: "TableEnvironment"):
+ self._t_env = t_env
+ self._functions: Dict[Any, str] = {}
+ self._taken_names: Optional[Set[str]] = None
+ self._cleanup = ExitStack()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc_info):
+ return self._cleanup.__exit__(*exc_info)
+
+ def serializeInlineFunction(self, definition):
+ if definition not in self._functions:
+ if self._taken_names is None:
+ self._taken_names = {name.lower() for name in
self._t_env.list_functions()}
+ name = _unique_name("__pf_join_udf", self._taken_names)
+ self._t_env._j_tenv.createTemporarySystemFunction(name, definition)
+ self._cleanup.callback(self._t_env.drop_temporary_system_function,
name)
+ self._taken_names.add(name)
+ self._functions[definition] = name
+ return _quote_identifier(self._functions[definition])
+
+ class Java:
+ implements = ["org.apache.flink.table.expressions.SqlFactory"]
+
+
+def _serialize_join_predicate(
+ left_table: Table,
+ right_table: Table,
+ predicate: Expression,
+ sql_factory: _JoinSqlFactory,
+) -> Tuple[str, str, str]:
+ left_alias, right_alias = "__pf_join_left", "__pf_join_right"
+ operation_tree_builder = (
+ left_table._j_table.getTableEnvironment().getOperationTreeBuilder()
+ )
+ gateway = get_gateway()
+ query_operations = to_jarray(
+ gateway.jvm.org.apache.flink.table.operations.QueryOperation,
+ [
+ left_table._j_table.getQueryOperation(),
+ right_table._j_table.getQueryOperation(),
+ ],
+ )
+ resolved_predicate = operation_tree_builder.resolveExpression(
+ _get_java_expression(predicate), query_operations
+ )
+
+ aliases = gateway.jvm.java.util.HashMap()
+ aliases.put(0, left_alias)
+ aliases.put(1, right_alias)
+ operation_expression_utils = (
+
gateway.jvm.org.apache.flink.table.operations.utils.OperationExpressionsUtils
+ )
+ predicate_sql = operation_expression_utils.scopeReferencesWithAlias(
+ aliases, resolved_predicate
+ ).asSerializableString(sql_factory)
+ return left_alias, right_alias, predicate_sql
+
+
+def _build_regular_join_sql(
+ left_table: Table,
+ right_table: Table,
+ predicate: Expression,
+ left_output_columns: List[str],
+ right_output_columns: List[str],
+ shared_keys: Dict[str, str],
+ join_type: str,
+ sql_factory: _JoinSqlFactory,
+) -> Table:
+ left_alias, right_alias, predicate_sql = _serialize_join_predicate(
+ left_table, right_table, predicate, sql_factory
+ )
+ left_alias_sql = _quote_identifier(left_alias)
+ right_alias_sql = _quote_identifier(right_alias)
+
+ projections = []
+ for name in left_output_columns:
+ left_field = "%s.%s" % (left_alias_sql, _quote_identifier(name))
+ if name in shared_keys and join_type in ("right", "full"):
+ right_field = "%s.%s" % (
+ right_alias_sql,
+ _quote_identifier(shared_keys[name]),
+ )
+ expression = (
+ right_field
+ if join_type == "right"
+ else "COALESCE(%s, %s)" % (left_field, right_field)
+ )
+ else:
+ expression = left_field
+ projections.append("%s AS %s" % (expression, _quote_identifier(name)))
+ projections.extend(
+ "%s.%s AS %s"
+ % (right_alias_sql, _quote_identifier(name), _quote_identifier(name))
+ for name in right_output_columns
+ if name not in shared_keys
+ )
+
+ join_keyword = {
+ "inner": "INNER JOIN",
+ "left": "LEFT OUTER JOIN",
+ "right": "RIGHT OUTER JOIN",
+ "full": "FULL OUTER JOIN",
+ }[join_type]
+ query = (
+ "SELECT %s FROM %s AS %s %s %s AS %s ON %s"
+ % (
+ ", ".join(projections),
+ _quote_identifier(str(left_table)),
+ left_alias_sql,
+ join_keyword,
+ _quote_identifier(str(right_table)),
+ right_alias_sql,
+ predicate_sql,
+ )
+ )
+ return left_table._t_env.sql_query(query)
+
+
+def _build_semi_anti_join_sql(
+ left_table: Table,
+ right_table: Table,
+ predicate: Expression,
+ output_columns: List[str],
+ join_type: str,
+ sql_factory: _JoinSqlFactory,
+) -> Table:
+ left_alias, right_alias, predicate_sql = _serialize_join_predicate(
+ left_table, right_table, predicate, sql_factory
+ )
+ left_alias_sql = _quote_identifier(left_alias)
+ right_alias_sql = _quote_identifier(right_alias)
+ select_list = ", ".join(
+ "%s.%s" % (left_alias_sql, _quote_identifier(name)) for name in
output_columns
+ )
+ left_source_sql = _quote_identifier(str(left_table))
+ right_source_sql = _quote_identifier(str(right_table))
+ existence_predicate = {"semi": "EXISTS", "anti": "NOT EXISTS"}[join_type]
+ query = (
+ "SELECT %s FROM %s AS %s WHERE %s ("
+ "SELECT 1 FROM %s AS %s WHERE %s)"
+ % (
+ select_list,
+ left_source_sql,
+ left_alias_sql,
+ existence_predicate,
+ right_source_sql,
+ right_alias_sql,
+ predicate_sql,
+ )
+ )
+ return left_table._t_env.sql_query(query)
+
+
def _normalize_subset(
subset: Union[str, List[str], None], parameter_name: str = "subset"
) -> Optional[List[str]]:
diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py
b/flink-python/pyflink/dataframe/tests/test_dataframe.py
index f7be076ca5a..0716d2cc16c 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -49,6 +49,7 @@ from pyflink.testing.test_case_utils import (
PyFlinkITTestCase,
PyFlinkStreamDataFrameTestCase,
)
+from pyflink.util.exceptions import TableException
class _Point(NamedTuple):
@@ -1133,6 +1134,413 @@ class
DataFrameRenameColumnsTests(PyFlinkDataFrameUTTestCase):
invalid_call()
+class DataFrameJoinTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ self.left = pf.from_records(
+ [(1, "left")],
+ schema=["id", "left_value"],
+ )
+ self.right = pf.from_records(
+ [(1, "right")],
+ schema=["id", "right_value"],
+ )
+
+ def test_join_on_shared_key_keeps_key_once(self):
+ result = self.left.join(self.right, on="id")
+
+ self.assert_dataframe_schema(
+ result,
+ ["id", "left_value", "right_value"],
+ [
+ TableDataTypes.BIGINT(),
+ TableDataTypes.STRING(),
+ TableDataTypes.STRING(),
+ ],
+ )
+
+ left = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (CAST(1 AS BIGINT))) AS T(id)"
+ )
+ )
+ right = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (CAST(1 AS INT))) AS T(id)"
+ )
+ )
+ self.assert_dataframe_schema(
+ left.join(right, on="id", how="right"),
+ ["id"],
+ [TableDataTypes.INT().not_null()],
+ )
+
+ def test_join_supports_multiple_shared_keys(self):
+ left = pf.from_records(
+ [(1, "A", "left")],
+ schema=["id", "category", "left_value"],
+ )
+ right = pf.from_records(
+ [(1, "A", "right")],
+ schema=["id", "category", "right_value"],
+ )
+
+ result = left.join(right, on=["id", "category"])
+
+ self.assertEqual(
+ result.columns,
+ ["id", "category", "left_value", "right_value"],
+ )
+ self.assertEqual(
+ left.join(right, on=["id", "category"], how="semi").columns,
+ ["id", "category", "left_value"],
+ )
+
+ def test_join_supports_different_and_computed_keys(self):
+ right = self.right.rename_columns({"id": "right_id"})
+
+ named_result = self.left.join(
+ right,
+ left_on="id",
+ right_on="right_id",
+ how="left",
+ )
+ computed_result = self.left.join(
+ right,
+ left_on=pf.col("id") + 1,
+ right_on=pf.col("right_id"),
+ )
+
+ for result in (named_result, computed_result):
+ self.assertEqual(
+ result.columns,
+ ["id", "left_value", "right_id", "right_value"],
+ )
+ self.assertFalse(any(name.startswith("__pf_join") for name in
result.columns))
+
+ def test_join_supports_expression_predicate(self):
+ right = pf.from_records(
+ [(1, 0, 2, "right")],
+ schema=["right_id", "min_id", "max_id", "right_value"],
+ )
+
+ result = self.left.join(
+ right,
+ on=(pf.col("id") == pf.col("right_id"))
+ & (pf.col("id") >= pf.col("min_id"))
+ & (pf.col("id") < pf.col("max_id")),
+ )
+
+ self.assertEqual(
+ result.columns,
+ ["id", "left_value", "right_id", "min_id", "max_id",
"right_value"],
+ )
+
+ def test_join_supports_non_equi_only_predicate(self):
+ for settings in (
+ EnvironmentSettings.in_batch_mode(),
+ EnvironmentSettings.in_streaming_mode(),
+ ):
+ with self.subTest(streaming=settings.is_streaming_mode()):
+ t_env = TableEnvironment.create(settings)
+ left = pf.from_table(t_env.from_elements([(1,), (2,)], ["id"]))
+ right = pf.from_table(t_env.from_elements([(2,), (3,)],
["right_id"]))
+
+ result = left.join(right, on=pf.col("id") < pf.col("right_id"))
+
+ self.assertEqual(result.columns, ["id", "right_id"])
+ plan = result.to_table().explain()
+ self.assertIn("joinType=[InnerJoin]", plan)
+ if settings.is_streaming_mode():
+ self.assertIn("distribution=[single]", plan)
+
+ def test_anti_join_uses_native_anti_join(self):
+ for settings in (
+ EnvironmentSettings.in_batch_mode(),
+ EnvironmentSettings.in_streaming_mode(),
+ ):
+ with self.subTest(streaming=settings.is_streaming_mode()):
+ t_env = TableEnvironment.create(settings)
+ left = pf.from_table(t_env.from_elements([(1, "A")], ["id",
"left_value"]))
+ right = pf.from_table(t_env.from_elements([(1, "X")], ["id",
"right_value"]))
+
+ with patch("pyflink.table.table.Table.execute") as execute:
+ with patch("pyflink.table.table.Table.explain") as explain:
+ result = left.join(right, on="id", how="anti")
+ execute.assert_not_called()
+ explain.assert_not_called()
+ plan = result.to_table().explain()
+
+ self.assertEqual(result.columns, left.columns)
+ self.assertIn("joinType=[LeftAntiJoin]", plan)
+ self.assertNotIn("joinType=[LeftOuterJoin]", plan)
+
+ def test_anti_join_uses_native_anti_join_for_values(self):
+ for settings in (
+ EnvironmentSettings.in_batch_mode(),
+ EnvironmentSettings.in_streaming_mode(),
+ ):
+ with self.subTest(streaming=settings.is_streaming_mode()):
+ t_env = TableEnvironment.create(settings)
+ left = pf.from_table(t_env.sql_query("SELECT * FROM (VALUES
(1), (2)) T(id)"))
+ right = pf.from_table(t_env.sql_query("SELECT * FROM (VALUES
(1), (2)) T(id)"))
+
+ with patch.object(t_env, "sql_query", wraps=t_env.sql_query)
as sql_query:
+ result = left.join(right, on="id", how="anti")
+
+ sql_query.assert_called_once()
+ self.assertIn("WHERE NOT EXISTS", sql_query.call_args[0][0])
+ self.assertEqual(result.columns, left.columns)
+ plan = result.to_table().explain()
+ self.assertIn("joinType=[LeftAntiJoin]", plan)
+ self.assertNotIn("joinType=[LeftOuterJoin]", plan)
+
+ def test_anti_join_propagates_query_errors_and_cleans_up_udfs(self):
+ right = self.right.rename_columns({"id": "right_id"})
+ is_allowed = pf.udf(lambda value: value == "right", return_dtype=bool)
+ predicate = (pf.col("id") == pf.col("right_id")) &
is_allowed(pf.col("right_value"))
+ functions_before = set(self.t_env.list_user_defined_functions())
+ errors = (
+ TableException("injected query failure"),
+ RuntimeError("injected unexpected failure"),
+ )
+
+ for error in errors:
+ with self.subTest(error=str(error)):
+ with patch.object(self.t_env, "sql_query", side_effect=error)
as sql_query:
+ with self.assertRaises(type(error)) as raised:
+ self.left.join(right, on=predicate, how="anti")
+
+ self.assertIs(raised.exception, error)
+ sql_query.assert_called_once()
+
self.assertEqual(set(self.t_env.list_user_defined_functions()),
functions_before)
+
+ def test_join_with_nested_inline_udfs_preserves_registered_functions(self):
+ right = self.right.rename_columns({"id": "right_id"})
+ normalize = pf.udf(lambda value: value.lower(), return_dtype=str)
+ same_value = pf.udf(lambda left, right: left == right,
return_dtype=bool)
+ self.t_env.create_temporary_system_function(
+ "__pf_join_udf", normalize._table_udf_wrapper
+ )
+ functions_before = set(self.t_env.list_user_defined_functions())
+ functions_during_query = []
+ original_sql_query = self.t_env.sql_query
+
+ def capture_functions(query):
+
functions_during_query.append(set(self.t_env.list_user_defined_functions()))
+ return original_sql_query(query)
+
+ match = same_value(normalize(pf.col("left_value")),
normalize(pf.col("right_value")))
+ with patch.object(self.t_env, "sql_query",
side_effect=capture_functions):
+ result = self.left.join(
+ right, on=(pf.col("id") == pf.col("right_id")) & match & match
+ )
+
+ self.assertEqual(result.columns, ["id", "left_value", "right_id",
"right_value"])
+ self.assertEqual(len(functions_during_query), 1)
+ self.assertEqual(len(functions_during_query[0] - functions_before), 2)
+ self.assertEqual(set(self.t_env.list_user_defined_functions()),
functions_before)
+ registered_result = self.t_env.sql_query("SELECT
__pf_join_udf('HELLO')")
+
self.assertEqual(len(registered_result.get_resolved_schema().get_column_names()),
1)
+
+ def test_join_cleans_up_inline_udfs_after_invalid_sql(self):
+ right = self.right.rename_columns({"id": "right_id"})
+ identity = pf.udf(lambda value: value, return_dtype=int)
+ functions_before = set(self.t_env.list_user_defined_functions())
+
+ with patch.object(self.t_env, "sql_query", wraps=self.t_env.sql_query)
as sql_query:
+ with self.assertRaises(Py4JJavaError):
+ self.left.join(right, on=identity(pf.col("id")))
+
+ sql_query.assert_called_once()
+ self.assertEqual(set(self.t_env.list_user_defined_functions()),
functions_before)
+
+ def test_join_cleans_up_partially_registered_inline_udfs(self):
+ right = self.right.rename_columns({"id": "right_id"})
+ normalize = pf.udf(lambda value: value.lower(), return_dtype=str)
+ same_value = pf.udf(lambda left, right: left == right,
return_dtype=bool)
+ functions_before = set(self.t_env.list_user_defined_functions())
+ original_register = self.t_env._j_tenv.createTemporarySystemFunction
+ registration_attempts = []
+
+ def fail_second_registration(name, definition):
+ registration_attempts.append(name)
+ if len(registration_attempts) == 2:
+ raise RuntimeError("injected UDF registration failure")
+ return original_register(name, definition)
+
+ with patch.object(
+ self.t_env._j_tenv,
+ "createTemporarySystemFunction",
+ side_effect=fail_second_registration,
+ ):
+ with self.assertRaisesRegex(Py4JJavaError, "injected UDF
registration failure"):
+ self.left.join(
+ right,
+ on=(pf.col("id") == pf.col("right_id"))
+ & same_value(normalize(pf.col("left_value")),
pf.col("right_value")),
+ )
+
+ self.assertEqual(len(registration_attempts), 2)
+ self.assertEqual(set(self.t_env.list_user_defined_functions()),
functions_before)
+
+ def test_join_supports_semi_anti_and_cross(self):
+ for how in ("semi", "anti"):
+ with self.subTest(how=how):
+ self.assertEqual(
+ self.left.join(self.right, on="id", how=how).columns,
+ ["id", "left_value"],
+ )
+
+ right = self.right.rename_columns({"id": "right_id"})
+ semi = self.left.join(
+ right,
+ on=pf.col("id") < pf.col("right_id"),
+ how="semi",
+ )
+ self.assertEqual(semi.columns, ["id", "left_value"])
+ self.assertIn("joinType=[LeftSemiJoin]", semi.to_table().explain())
+
+ is_right = pf.udf(lambda value: value == "right", return_dtype=bool)
+ functions_before = set(self.t_env.list_user_defined_functions())
+ semi_with_udf = self.left.join(
+ right,
+ on=(pf.col("id") == pf.col("right_id"))
+ & is_right(pf.col("right_value")),
+ how="semi",
+ )
+ self.assertEqual(set(self.t_env.list_user_defined_functions()),
functions_before)
+ self.assertIn("joinType=[LeftSemiJoin]",
semi_with_udf.to_table().explain())
+
+ self.assertEqual(
+ self.left.join(right, how="cross").columns,
+ ["id", "left_value", "right_id", "right_value"],
+ )
+
+ def test_join_outer_alias_matches_full_schema(self):
+ full = self.left.join(self.right, on="id", how="full")
+ outer = self.left.join(self.right, on="id", how="outer")
+
+ self.assertEqual(full.columns, outer.columns)
+ self.assertEqual(
+ full._table.get_resolved_schema(),
+ outer._table.get_resolved_schema(),
+ )
+
+ def test_join_rejects_invalid_argument_combinations(self):
+ invalid_calls = [
+ (
+ "other",
+ lambda: self.left.join(object(), on="id"),
+ TypeError,
+ "other must be",
+ ),
+ (
+ "how_type",
+ lambda: self.left.join(self.right, on="id", how=1),
+ TypeError,
+ "how must be a string",
+ ),
+ (
+ "how_value",
+ lambda: self.left.join(self.right, on="id", how="sideways"),
+ ValueError,
+ "how must be one of",
+ ),
+ (
+ "missing_keys",
+ lambda: self.left.join(self.right),
+ ValueError,
+ "requires on or both",
+ ),
+ (
+ "mixed_keys",
+ lambda: self.left.join(
+ self.right,
+ on="id",
+ left_on="id",
+ right_on="id",
+ ),
+ ValueError,
+ "on cannot be combined",
+ ),
+ (
+ "missing_right_on",
+ lambda: self.left.join(self.right, left_on="id"),
+ ValueError,
+ "must be provided together",
+ ),
+ (
+ "different_key_counts",
+ lambda: self.left.join(
+ self.right,
+ left_on=["id", "left_value"],
+ right_on=["id"],
+ ),
+ ValueError,
+ "same number of keys",
+ ),
+ (
+ "empty_keys",
+ lambda: self.left.join(self.right, on=[]),
+ ValueError,
+ "on must not be empty",
+ ),
+ (
+ "invalid_key_type",
+ lambda: self.left.join(self.right, on=1),
+ TypeError,
+ "on must be a string",
+ ),
+ (
+ "missing_column",
+ lambda: self.left.join(self.right, on="missing"),
+ ValueError,
+ "on column 'missing' does not exist",
+ ),
+ (
+ "cross_keys",
+ lambda: self.left.join(self.right, on="id", how="cross"),
+ ValueError,
+ "cross join does not accept",
+ ),
+ ]
+ for name, invalid_call, error, message in invalid_calls:
+ with self.subTest(name=name):
+ with self.assertRaisesRegex(error, message):
+ invalid_call()
+
+ def test_join_rejects_duplicate_non_key_columns(self):
+ right = pf.from_records(
+ [(1, "duplicate")],
+ schema=["right_id", "left_value"],
+ )
+
+ with self.assertRaisesRegex(ValueError, "duplicate non-key
columns.*rename_columns"):
+ self.left.join(right, left_on="id", right_on="right_id")
+
+ for how in ("semi", "anti"):
+ with self.subTest(how=how):
+ self.assertEqual(
+ self.left.join(
+ right,
+ left_on="id",
+ right_on="right_id",
+ how=how,
+ ).columns,
+ ["id", "left_value"],
+ )
+
+ def test_join_rejects_different_table_environments(self):
+ other_environment =
TableEnvironment.create(EnvironmentSettings.in_batch_mode())
+ other = pf.from_table(
+ other_environment.sql_query("SELECT 1 AS id, 'right' AS
right_value")
+ )
+
+ with self.assertRaisesRegex(ValueError, "same TableEnvironment"):
+ self.left.join(other, on="id")
+
+
class DataFramePropertyTests(PyFlinkDataFrameUTTestCase):
def test_schema_exposes_ordered_metadata(self):
dataframe = pf.from_records(
@@ -2596,6 +3004,19 @@ class DataFrameBatchITTests(PyFlinkITTestCase):
)
return pf.from_table(table)
+ def _join_dataframes(self):
+ left = self.t_env.sql_query(
+ "SELECT * FROM (VALUES "
+ "(CAST(1 AS INT), 'L1'), (2, 'L2'), (CAST(NULL AS INT), 'LN')) "
+ "AS T(id, left_value)"
+ )
+ right = self.t_env.sql_query(
+ "SELECT * FROM (VALUES "
+ "(CAST(2 AS INT), 'R2'), (3, 'R3'), (CAST(NULL AS INT), 'RN')) "
+ "AS T(id, right_value)"
+ )
+ return pf.from_table(left), pf.from_table(right)
+
def test_sort_returns_rows_in_ascending_order(self):
self.assertEqual(
self._unsorted_dataframe().sort("id").collect(),
@@ -2684,6 +3105,312 @@ class DataFrameBatchITTests(PyFlinkITTestCase):
[Row("engineering", 30, 2), Row("sales", 5, 1)],
)
+ def test_join_types_and_null_keys(self):
+ expected = {
+ "inner": [Row(2, "L2", "R2")],
+ "left": [
+ Row(1, "L1", None),
+ Row(2, "L2", "R2"),
+ Row(None, "LN", None),
+ ],
+ "right": [
+ Row(2, "L2", "R2"),
+ Row(3, None, "R3"),
+ Row(None, None, "RN"),
+ ],
+ "full": [
+ Row(1, "L1", None),
+ Row(2, "L2", "R2"),
+ Row(3, None, "R3"),
+ Row(None, "LN", None),
+ Row(None, None, "RN"),
+ ],
+ "outer": [
+ Row(1, "L1", None),
+ Row(2, "L2", "R2"),
+ Row(3, None, "R3"),
+ Row(None, "LN", None),
+ Row(None, None, "RN"),
+ ],
+ }
+
+ for how, expected_rows in expected.items():
+ with self.subTest(how=how):
+ left, right = self._join_dataframes()
+ self.assertCountEqual(
+ left.join(right, on="id", how=how).collect(),
+ expected_rows,
+ )
+
+ def test_semi_and_anti_join_preserve_left_multiplicity(self):
+ left = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES "
+ "(CAST(1 AS INT), 'A'), (1, 'A'), (1, 'B'), (2, 'C'), (2,
'C'), "
+ "(CAST(NULL AS INT), 'N')) AS T(id, left_value)"
+ )
+ )
+ right = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES "
+ "(CAST(1 AS INT), 'X'), (1, 'Y'), (CAST(NULL AS INT), 'Z')) "
+ "AS T(id, right_value)"
+ )
+ )
+
+ self.assertCountEqual(
+ left.join(right, on="id", how="semi").collect(),
+ [Row(1, "A"), Row(1, "A"), Row(1, "B")],
+ )
+ self.assertCountEqual(
+ left.join(right, on="id", how="anti").collect(),
+ [Row(2, "C"), Row(2, "C"), Row(None, "N")],
+ )
+
+ def test_anti_join_with_multiple_nullable_keys(self):
+ left = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (1, 'A'), (1, 'B'), (1, CAST(NULL AS
STRING)), "
+ "(CAST(NULL AS INT), 'A'), (CAST(NULL AS INT), CAST(NULL AS
STRING)), "
+ "(2, 'A'), (2, 'A')) "
+ "AS T(id, category)"
+ )
+ )
+ right = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (1, 'A'), (1, CAST(NULL AS STRING)), "
+ "(CAST(NULL AS INT), 'A'), (CAST(NULL AS INT), CAST(NULL AS
STRING))) "
+ "AS T(id, category)"
+ )
+ )
+
+ self.assertCountEqual(
+ left.join(right, on=["id", "category"], how="anti").collect(),
+ [
+ Row(1, "B"),
+ Row(1, None),
+ Row(None, "A"),
+ Row(None, None),
+ Row(2, "A"),
+ Row(2, "A"),
+ ],
+ )
+
+ def test_anti_join_with_empty_inputs(self):
+ left, right = self._join_dataframes()
+ empty_left = left.filter(pf.col("id") < 0)
+ empty_right = right.filter(pf.col("id") < 0)
+
+ self.assertCountEqual(
+ left.join(empty_right, on="id", how="anti").collect(),
+ [Row(1, "L1"), Row(2, "L2"), Row(None, "LN")],
+ )
+ self.assertEqual(empty_left.join(right, on="id",
how="anti").collect(), [])
+ self.assertEqual(empty_left.join(empty_right, on="id",
how="anti").collect(), [])
+
+ def test_anti_join_with_computed_keys_and_non_equi_predicate(self):
+ left = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (1), (2), (3), (CAST(NULL AS INT))) AS
T(id)"
+ )
+ )
+ right = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (2), (CAST(NULL AS INT))) AS
T(right_id)"
+ )
+ )
+
+ self.assertCountEqual(
+ left.join(
+ right, left_on=pf.col("id") + 1, right_on="right_id",
how="anti"
+ ).collect(),
+ [Row(2), Row(3), Row(None)],
+ )
+ self.assertCountEqual(
+ left.join(right, on=pf.col("id") < pf.col("right_id"),
how="anti").collect(),
+ [Row(2), Row(3), Row(None)],
+ )
+
+ def test_join_with_different_names_computed_keys_and_expression(self):
+ left = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (CAST(1 AS INT), 'L1'), (2, 'L2')) "
+ "AS T(id, left_value)"
+ )
+ )
+ right = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES "
+ "(CAST(2 AS INT), 1, 3, 'R2'), (3, 2, 4, 'R3')) "
+ "AS T(right_id, min_id, max_id, right_value)"
+ )
+ )
+
+ self.assertCountEqual(
+ left.join(right, left_on=pf.col("id") + 1,
right_on="right_id").collect(),
+ [
+ Row(1, "L1", 2, 1, 3, "R2"),
+ Row(2, "L2", 3, 2, 4, "R3"),
+ ],
+ )
+ predicate = (
+ (pf.col("id") + 1 == pf.col("right_id"))
+ & (pf.col("id") >= pf.col("min_id"))
+ & (pf.col("id") < pf.col("max_id"))
+ )
+ self.assertCountEqual(
+ left.join(right, on=predicate).collect(),
+ [
+ Row(1, "L1", 2, 1, 3, "R2"),
+ Row(2, "L2", 3, 2, 4, "R3"),
+ ],
+ )
+ self.assertEqual(left.join(right, on=predicate, how="anti").collect(),
[])
+ self.assertCountEqual(
+ left.join(
+ right,
+ left_on=pf.col("id") + 1,
+ right_on="right_id",
+ how="semi",
+ ).collect(),
+ [Row(1, "L1"), Row(2, "L2")],
+ )
+
+ def test_join_with_non_equi_only_predicate(self):
+ left = pf.from_table(
+ self.t_env.sql_query("SELECT * FROM (VALUES (1), (2), (3)) AS
T(id)")
+ )
+ right = pf.from_table(
+ self.t_env.sql_query("SELECT * FROM (VALUES (2), (3)) AS
T(right_id)")
+ )
+ predicate = pf.col("id") < pf.col("right_id")
+
+ self.assertCountEqual(
+ left.join(right, on=predicate).collect(),
+ [Row(1, 2), Row(1, 3), Row(2, 3)],
+ )
+ self.assertCountEqual(
+ left.join(right, on=predicate, how="left").collect(),
+ [Row(1, 2), Row(1, 3), Row(2, 3), Row(3, None)],
+ )
+ self.assertCountEqual(
+ left.join(right, on=predicate, how="semi").collect(),
+ [Row(1), Row(2)],
+ )
+
+ def test_join_with_inline_udf_predicate_after_function_cleanup(self):
+ left = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (1, 'A'), (1, 'A'), (2, 'B'), "
+ "(3, CAST(NULL AS STRING)), (CAST(NULL AS INT), 'N')) "
+ "AS T(id, left_value)"
+ )
+ )
+ right = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (1, 'A'), (1, 'X'), (2, 'X'), "
+ "(3, CAST(NULL AS STRING)), (CAST(NULL AS INT), 'N')) "
+ "AS T(right_id, right_value)"
+ )
+ )
+ same_value = pf.udf(lambda left, right: left == right,
return_dtype=bool)
+ functions_before = set(self.t_env.list_user_defined_functions())
+
+ result = left.join(
+ right,
+ on=(pf.col("id") == pf.col("right_id"))
+ & same_value(pf.col("left_value"), pf.col("right_value")),
+ )
+
+ self.assertEqual(set(self.t_env.list_user_defined_functions()),
functions_before)
+ self.assertCountEqual(
+ result.collect(),
+ [Row(1, "A", 1, "A"), Row(1, "A", 1, "A"), Row(3, None, 3, None)],
+ )
+
+ def test_left_semi_and_anti_join_with_inline_udf_on_right_input(self):
+ left = pf.from_table(
+ self.t_env.sql_query("SELECT * FROM (VALUES (1), (2)) AS T(id)")
+ )
+ right = pf.from_table(
+ self.t_env.sql_query(
+ "SELECT * FROM (VALUES (1, 'no'), (2, 'yes')) AS T(right_id,
flag)"
+ )
+ )
+ is_allowed = pf.udf(lambda value: value == "yes", return_dtype=bool)
+ predicate = (pf.col("id") == pf.col("right_id")) &
is_allowed(pf.col("flag"))
+ functions_before = set(self.t_env.list_user_defined_functions())
+
+ outer = left.join(right, on=predicate, how="left")
+ semi = left.join(right, on=predicate, how="semi")
+ anti = left.join(right, on=predicate, how="anti")
+
+ self.assertEqual(set(self.t_env.list_user_defined_functions()),
functions_before)
+
self.assertTrue(outer.to_table().get_resolved_schema().get_column_data_types()[1]._nullable)
+ self.assertCountEqual(outer.collect(), [Row(1, None, None), Row(2, 2,
"yes")])
+ self.assertEqual(semi.collect(), [Row(2)])
+ self.assertEqual(anti.collect(), [Row(1)])
+
+ def test_anti_join_on_table_sources(self):
+ left = pf.from_table(self.t_env.from_elements([(1,), (2,), (3,),
(None,)], ["id"]))
+ right = pf.from_table(self.t_env.from_elements([(2,), (None,)],
["right_id"]))
+ cases = (
+ ({"left_on": "id", "right_on": "right_id"}, [Row(1), Row(3),
Row(None)]),
+ (
+ {"left_on": pf.col("id") + 1, "right_on": "right_id"},
+ [Row(2), Row(3), Row(None)],
+ ),
+ ({"on": pf.col("id") < pf.col("right_id")}, [Row(2), Row(3),
Row(None)]),
+ )
+
+ for keys, expected in cases:
+ with self.subTest(keys=keys):
+ result = left.join(right, how="anti", **keys)
+
+ self.assertIn("joinType=[LeftAntiJoin]",
result.to_table().explain())
+ self.assertCountEqual(result.collect(), expected)
+
+ def test_anti_join_on_table_sources_with_inline_udf_after_cleanup(self):
+ left = pf.from_table(
+ self.t_env.from_elements(
+ [(1, "A"), (1, "A"), (2, "B"), (3, "C"), (None, "N")],
+ ["id", "left_value"],
+ )
+ )
+ right = pf.from_table(
+ self.t_env.from_elements(
+ [(1, "no"), (2, "yes"), (2, "yes"), (None, "yes")],
+ ["right_id", "flag"],
+ )
+ )
+ is_allowed = pf.udf(lambda value: value == "yes", return_dtype=bool)
+ functions_before = set(self.t_env.list_user_defined_functions())
+
+ result = left.join(
+ right,
+ on=(pf.col("id") == pf.col("right_id")) &
is_allowed(pf.col("flag")),
+ how="anti",
+ )
+
+ self.assertEqual(set(self.t_env.list_user_defined_functions()),
functions_before)
+ self.assertIn("joinType=[LeftAntiJoin]", result.to_table().explain())
+ self.assertCountEqual(
+ result.collect(), [Row(1, "A"), Row(1, "A"), Row(3, "C"),
Row(None, "N")]
+ )
+
+ def test_cross_join(self):
+ left = pf.from_table(
+ self.t_env.sql_query("SELECT * FROM (VALUES (1), (2)) AS T(id)")
+ )
+ right = pf.from_table(
+ self.t_env.sql_query("SELECT * FROM (VALUES ('S'), ('M')) AS
T(size_name)")
+ )
+
+ self.assertCountEqual(
+ left.join(right, how="cross").collect(),
+ [Row(1, "S"), Row(1, "M"), Row(2, "S"), Row(2, "M")],
+ )
+
class DataFrameSetOperationITTests(PyFlinkITTestCase):
def setUp(self):
