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 14fa2718313 [FLINK-40192][python] Add as-of join support to DataFrame
API
14fa2718313 is described below
commit 14fa27183139cb9ec0a5f2b34f5d72f7477115c1
Author: beetle0915 <[email protected]>
AuthorDate: Mon Oct 19 01:11:25 2026 +0800
[FLINK-40192][python] Add as-of join support to DataFrame API
Support processing-time lookup and event-time temporal joins through the
existing join SQL helpers. Add validation, execution coverage, and API
documentation.
This closes #29513.
Generated-by: OpenAI Codex (codex-cli 0.160.0)
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 16 +-
flink-python/pyflink/dataframe/__init__.py | 3 +-
flink-python/pyflink/dataframe/dataframe.py | 255 +++++++++++++++-
.../pyflink/dataframe/tests/test_dataframe.py | 338 +++++++++++++++++++++
.../table/utils/TestCollectionTableFactory.java | 2 +-
5 files changed, 603 insertions(+), 11 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index e31a2bf176d..d6d575b4db8 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -113,8 +113,21 @@ Joins
:toctree: api/
DataFrame.join
+ DataFrame.join_asof
DataFrame.join_lateral
+Use ``join_asof`` to enrich rows from a dimension table at a left time
attribute.
+Processing-time lookup joins require a lookup-capable right table; event-time
temporal
+joins require watermarks on both inputs and the right primary key in the
equality keys::
+
+ enriched = orders.join_asof(customers, on=pf.proctime(), by="customer_id",
how="left")
+ historical = orders.join_asof(
+ rates, on="event_time", left_by="currency", right_by="currency_id")
+
+Only inner and left joins are supported. Conditions need at least one equality
between
+the inputs and may include additional AND-combined predicates, but not OR
conditions.
+The time attribute may also appear in equality keys and residual conditions.
+
Missing-Value Handling
----------------------
@@ -197,7 +210,7 @@ Windowing
Expressions
-----------
-Functions for constructing column references and literal expressions.
+Functions for constructing column references, literals, and processing-time
attributes.
.. currentmodule:: pyflink.dataframe
@@ -206,6 +219,7 @@ Functions for constructing column references and literal
expressions.
col
lit
+ proctime
OVER windows
------------
diff --git a/flink-python/pyflink/dataframe/__init__.py
b/flink-python/pyflink/dataframe/__init__.py
index 561f2ea318f..6783c606225 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -62,7 +62,7 @@ from pyflink.dataframe.context import (
get_table_environment,
set_table_environment,
)
-from pyflink.dataframe.dataframe import DataFrame, GroupedDataFrame, col, lit
+from pyflink.dataframe.dataframe import DataFrame, GroupedDataFrame, col, lit,
proctime
from pyflink.dataframe.dataframe_config import config
from pyflink.dataframe.datatype import DataType
from pyflink.dataframe.io import (
@@ -102,6 +102,7 @@ __all__ = [
"list_model_providers",
"col",
"lit",
+ "proctime",
"udf",
"udtf",
"udaf",
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index d54f3538a5f..7a90e1dfeb5 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -161,6 +161,15 @@ def lit(value: Any, data_type: Optional[DataType] = None)
-> Expression:
return table_lit(value, table_data_type.not_null())
+@PublicEvolving()
+def proctime() -> Expression:
+ """Return a processing-time attribute for a lookup
:meth:`DataFrame.join_asof`.
+
+ .. versionadded:: 2.4.0
+ """
+ return call_sql("PROCTIME()")
+
+
@PublicEvolving()
class DataFrame:
"""
@@ -1320,6 +1329,95 @@ class DataFrame:
)
)
+ @PublicEvolving()
+ def join_asof(
+ self,
+ other: "DataFrame",
+ *,
+ how: str = "inner",
+ on: Optional[Union[str, Expression]] = None,
+ by: Optional[Union[str, Expression, List[Union[str, Expression]]]] =
None,
+ left_by: Optional[Union[str, Expression, List[Union[str,
Expression]]]] = None,
+ right_by: Optional[Union[str, Expression, List[Union[str,
Expression]]]] = None,
+ ) -> "DataFrame":
+ """
+ Join a dimension table as of each left row's time attribute.
+
+ Use :func:`proctime` or a left processing-time column for a lookup
join. The
+ right table must support lookups. Use a left event-time column for a
temporal
+ join; both inputs need watermarks, and the right table's primary key
must be
+ covered by the equality keys. These requirements follow the underlying
Table API.
+
+ Shared named keys occur once in the result. Other duplicate column
names must
+ be renamed before joining. The time attribute may also appear in
equality keys
+ and residual conditions.
+ Conditions require an equality between the inputs, optionally
AND-combined
+ with other conditions; OR conditions and purely non-equi conditions
are rejected.
+
+ :param other: Dimension DataFrame in the same TableEnvironment.
+ :param how: Join type, ``"inner"`` or ``"left"``.
+ :param on: Required left time-attribute column or :func:`proctime`
expression.
+ :param by: Shared key name, boolean join expression, or a list of
these.
+ :param left_by: Left key name, expression, or list of names or
expressions.
+ :param right_by: Right keys paired with ``left_by``, instead of ``by``.
+ :return: A new DataFrame containing the join result.
+ :raises TypeError: If an argument has an unsupported type.
+ :raises ValueError: If the time attribute, keys, or argument
combination is invalid.
+
+ Example::
+
+ >>> enriched = orders.join_asof(
+ ... customers, on=pf.proctime(), by="customer_id", how="left")
+ >>> historical = orders.join_asof(
+ ... rates, on="event_time", left_by="currency",
right_by="currency_id")
+
+ .. 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")
+ if not isinstance(how, str):
+ raise TypeError("how must be a string")
+ if how not in ("inner", "left"):
+ raise ValueError('how must be "inner" or "left"')
+ if on is None:
+ raise ValueError("on must specify a left time attribute")
+ if not isinstance(on, (str, Expression)):
+ raise TypeError("on must be a column name or expression")
+ if by is not None and (left_by is not None or right_by is not None):
+ raise ValueError("by cannot be combined with left_by or right_by")
+ if by is None and (left_by is None or right_by is None):
+ raise ValueError("join_asof requires by or left_by and right_by
together")
+
+ left_table, time_name = _prepare_asof_time(self._table, on,
other.columns)
+ join_keys: Union[List[str], Expression, None] = None
+ extra_conditions = []
+ named_keys = []
+ if by is not None:
+ for key in _normalize_join_keys(by, "by", allow_expressions=True):
+ if isinstance(key, str):
+ named_keys.append(key)
+ else:
+ extra_conditions.append(key)
+ if named_keys:
+ join_keys = named_keys
+ else:
+ join_keys = (and_(*extra_conditions)
+ if len(extra_conditions) > 1 else
extra_conditions[0])
+ extra_conditions = []
+ left_table, right_table, predicate, shared_keys = _prepare_join(
+ left_table, other._table, join_keys, left_by, right_by,
+ validate_column_conflicts=True, asof_time=time_name,
+ )
+ if extra_conditions:
+ predicate = and_(predicate, *extra_conditions)
+ with _JoinSqlFactory(self._table._t_env) as sql_factory:
+ return DataFrame(_build_regular_join_sql(
+ left_table, right_table, predicate, self.columns,
other.columns,
+ shared_keys, how, sql_factory, asof_time=time_name,
+ ))
+
@PublicEvolving()
def join_lateral(
self,
@@ -3321,18 +3419,25 @@ def _normalize_join_type(how: str) -> str:
return aliases[how]
-def _normalize_join_keys(value, parameter_name: str) -> List[Union[str,
Expression]]:
+def _normalize_join_keys(
+ value, parameter_name: str, *, allow_expressions: bool = False
+) -> 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):
+ key_types = (str, Expression) if allow_expressions else (str,)
+ if not all(isinstance(key, key_types) 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)
+ names = set()
+ for key in value:
+ if isinstance(key, str):
+ if key in names:
+ raise ValueError("%s must not contain duplicate column
names" % parameter_name)
+ names.add(key)
return value
raise TypeError(
"%s must be a string, an expression, or a list of strings" %
parameter_name
@@ -3369,6 +3474,7 @@ def _prepare_join(
right_on,
*,
validate_column_conflicts: bool,
+ asof_time: Optional[str] = None,
) -> 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())
@@ -3389,13 +3495,37 @@ def _prepare_join(
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")
+ left_keys = _normalize_join_keys(
+ left_on, "left_on", allow_expressions=asof_time is not None)
+ right_keys = _normalize_join_keys(
+ right_on, "right_on", allow_expressions=asof_time is not None)
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")
+ if asof_time is not None:
+ # Preserve direct field references so temporal joins retain the right
primary key.
+ left_keys, right_keys = list(left_keys), list(right_keys)
+ for table, keys, columns in (
+ (left_table, left_keys, left_columns), (right_table, right_keys,
right_columns)
+ ):
+ for index, key in enumerate(keys):
+ if isinstance(key, Expression):
+ selected = table.select(key)
+ if len(selected.get_resolved_schema().get_column_names())
!= 1:
+ raise ValueError("each join key must select a single
column")
+ projection = selected._j_table.getQueryOperation()
+ if not projection.getChildren().get(0).equals(
+ table._j_table.getQueryOperation()
+ ):
+ raise ValueError("join keys must be row-wise
expressions, not aggregations")
+ resolved = projection.getProjectList().get(0)
+ _asof_expression_info(resolved, len(columns))
+ field_index = _asof_field_index(resolved)
+ if field_index is not None:
+ keys[index] = columns[field_index]
+
shared_names = {
left_key
for left_key, right_key in zip(left_keys, right_keys)
@@ -3456,6 +3586,102 @@ def _prepare_join(
)
+def _asof_field_index(expression) -> Optional[int]:
+ definitions =
get_gateway().jvm.org.apache.flink.table.functions.BuiltInFunctionDefinitions
+ checks =
get_gateway().jvm.org.apache.flink.table.types.logical.utils.LogicalTypeChecks
+ while (expression.getClass().getSimpleName() == "CallExpression" and
+ expression.getFunctionDefinition().equals(definitions.CAST)):
+ child = expression.getResolvedChildren().get(0)
+ # DataType equality ignores time-attribute kinds; casting materializes
such attributes.
+ if (not
expression.getOutputDataType().equals(child.getOutputDataType()) or
+
checks.isTimeAttribute(child.getOutputDataType().getLogicalType())):
+ return None
+ expression = child
+ if expression.getClass().getSimpleName() == "FieldReferenceExpression":
+ return expression.getFieldIndex()
+ if expression.getClass().getSimpleName() == "RexNodeExpression":
+ node = expression.getRexNode()
+ while node.getKind().name() == "CAST":
+ child = node.getOperands().get(0)
+ if not node.getType().equals(child.getType()):
+ return None
+ node = child
+ if node.getKind().name() == "INPUT_REF":
+ return node.getIndex()
+ return None
+
+
+def _prepare_asof_time(table: Table, on, right_columns: List[str]) ->
Tuple[Table, str]:
+ expression = table_col(on) if isinstance(on, str) else on
+ selected = table.select(expression)
+ projection = selected._j_table.getQueryOperation()
+ if not
projection.getChildren().get(0).equals(table._j_table.getQueryOperation()):
+ raise ValueError("on must be a row-wise time attribute, not an
aggregation")
+ schema = selected._j_table.getResolvedSchema()
+ checks =
get_gateway().jvm.org.apache.flink.table.types.logical.utils.LogicalTypeChecks
+ if schema.getColumnCount() != 1 or not checks.isTimeAttribute(
+ schema.getColumnDataTypes().get(0).getLogicalType()
+ ):
+ raise ValueError("on must be a left processing-time or event-time
attribute")
+ resolved = projection.getProjectList().get(0)
+ columns = list(table.get_resolved_schema().get_column_names())
+ _asof_expression_info(resolved, len(columns))
+ field_index = _asof_field_index(resolved)
+ if field_index is not None:
+ return table, columns[field_index]
+ name = _unique_name("__pf_asof_time", set(columns) | set(right_columns))
+ return table.add_columns(expression.alias(name)), name
+
+
+def _asof_expression_info(expression, left_arity: int) -> Tuple[Set[int],
bool]:
+ """Return referenced inputs and whether an equi-join key is present."""
+ class_name = expression.getClass().getSimpleName()
+ if class_name == "RexNodeExpression":
+ return _asof_expression_info(expression.getRexNode(), left_arity)
+ if class_name == "FieldReferenceExpression":
+ return {expression.getInputIndex()}, False
+ operator = None
+ if expression.getClass().getName().startswith("org.apache.calcite.rex."):
+ # SQL expressions use flattened input indices rather than Table API
input/field pairs.
+ operator = expression.getKind().name()
+ if operator == "INPUT_REF":
+ index = expression.getIndex()
+ return {0 if index < left_arity else 1}, False
+ if operator == "LITERAL":
+ return set(), False
+ if operator == "FIELD_ACCESS":
+ children = [expression.getReferenceExpr()]
+ else:
+ if expression.getOperator().isAggregator():
+ raise ValueError("join expressions must be row-wise, not
aggregations")
+ children = expression.getOperands()
+ else:
+ children = expression.getResolvedChildren()
+ if class_name == "CallExpression":
+ definitions = (
+
get_gateway().jvm.org.apache.flink.table.functions.BuiltInFunctionDefinitions
+ )
+ for name in ("AND", "OR", "EQUALS"):
+ if
expression.getFunctionDefinition().equals(getattr(definitions, name)):
+ operator = name
+ break
+ if operator == "OR":
+ raise ValueError("join_asof does not support OR conditions")
+ child_info = []
+ references = set()
+ for child in children:
+ info = _asof_expression_info(child, left_arity)
+ references.update(info[0])
+ child_info.append(info)
+ has_equality = False
+ if operator == "AND":
+ has_equality = any(info[1] for info in child_info)
+ elif operator == "EQUALS":
+ left, right = child_info[0][0], child_info[1][0]
+ has_equality = (left == {0} and right == {1}) or (left == {1} and
right == {0})
+ return references, has_equality
+
+
class _JoinSqlFactory:
"""Keep inline UDFs registered until the join SQL has been resolved."""
@@ -3491,6 +3717,7 @@ def _serialize_join_predicate(
right_table: Table,
predicate: Expression,
sql_factory: _JoinSqlFactory,
+ asof_time: Optional[str] = None,
) -> Tuple[str, str, str]:
left_alias, right_alias = "__pf_join_left", "__pf_join_right"
operation_tree_builder = (
@@ -3507,6 +3734,11 @@ def _serialize_join_predicate(
resolved_predicate = operation_tree_builder.resolveExpression(
_get_java_expression(predicate), query_operations
)
+ if asof_time is not None:
+ columns = list(left_table.get_resolved_schema().get_column_names())
+ _, has_equality = _asof_expression_info(resolved_predicate,
len(columns))
+ if not has_equality:
+ raise ValueError("join_asof requires an AND-combined equality
between the inputs")
aliases = gateway.jvm.java.util.HashMap()
aliases.put(0, left_alias)
@@ -3529,9 +3761,11 @@ def _build_regular_join_sql(
shared_keys: Dict[str, str],
join_type: str,
sql_factory: _JoinSqlFactory,
+ *,
+ asof_time: Optional[str] = None,
) -> Table:
left_alias, right_alias, predicate_sql = _serialize_join_predicate(
- left_table, right_table, predicate, sql_factory
+ left_table, right_table, predicate, sql_factory, asof_time
)
left_alias_sql = _quote_identifier(left_alias)
right_alias_sql = _quote_identifier(right_alias)
@@ -3565,14 +3799,19 @@ def _build_regular_join_sql(
"right": "RIGHT OUTER JOIN",
"full": "FULL OUTER JOIN",
}[join_type]
+ temporal_clause = ""
+ if asof_time is not None:
+ temporal_clause = " FOR SYSTEM_TIME AS OF %s.%s" % (
+ left_alias_sql, _quote_identifier(asof_time))
query = (
- "SELECT %s FROM %s AS %s %s %s AS %s ON %s"
+ "SELECT %s FROM %s AS %s %s %s%s AS %s ON %s"
% (
", ".join(projections),
_quote_identifier(str(left_table)),
left_alias_sql,
join_keyword,
_quote_identifier(str(right_table)),
+ temporal_clause,
right_alias_sql,
predicate_sql,
)
diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py
b/flink-python/pyflink/dataframe/tests/test_dataframe.py
index ce906b405aa..35034b43c18 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -36,6 +36,7 @@ from pyflink.dataframe.dataframe import (
_to_interval_expression,
)
from pyflink.dataframe.datatype import DataType
+from pyflink.java_gateway import get_gateway
from pyflink.table import (
DataTypes as TableDataTypes,
EnvironmentSettings,
@@ -1210,6 +1211,234 @@ class
DataFrameRenameColumnsTests(PyFlinkDataFrameUTTestCase):
invalid_call()
+class DataFrameAsOfJoinTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ self.t_env.execute_sql(
+ "CREATE TABLE orders (id INT, amount INT, pt AS PROCTIME(), "
+ "ts TIMESTAMP(3), WATERMARK FOR ts AS ts) "
+ "WITH ('connector' = 'COLLECTION')"
+ )
+ self.t_env.execute_sql(
+ "CREATE TABLE rates (id INT, rate INT, rt TIMESTAMP(3), "
+ "WATERMARK FOR rt AS rt, PRIMARY KEY (id) NOT ENFORCED) "
+ "WITH ('connector' = 'COLLECTION')"
+ )
+ self.left = pf.from_table(self.t_env.from_path("orders"))
+ self.versions = pf.from_table(self.t_env.from_path("rates"))
+ self.t_env.execute_sql(
+ "CREATE TABLE dimension (id INT, rate INT, rt TIMESTAMP(3)) "
+ "WITH ('connector' = 'COLLECTION')"
+ )
+ self.right = pf.from_table(self.t_env.from_path("dimension"))
+
+ def test_time_attributes_and_lazy_construction(self):
+ original = self.left.to_table()
+ for on, operator in (
+ ("pt", "LookupJoin"), (pf.col("pt"), "LookupJoin"),
+ (pf.proctime(), "LookupJoin"), ("ts", "TemporalJoin"),
+ ):
+ with self.subTest(on=str(on)):
+ with patch.object(Table, "execute") as execute:
+ with patch.object(Table, "explain") as explain:
+ right = self.versions if operator == "TemporalJoin"
else self.right
+ result = self.left.join_asof(right, on=on, by="id",
how="left")
+ execute.assert_not_called()
+ explain.assert_not_called()
+ self.assert_dataframe_schema(result, ["id", "amount", "pt",
"ts", "rate", "rt"])
+ self.assertIs(self.left.to_table(), original)
+ plan = result.to_table().explain()
+ self.assertIn(operator, plan)
+
+ def test_key_forms_and_inline_functions(self):
+ right = self.right.rename_columns({"id": "d_id"})
+ identity = pf.udf(lambda value: value, return_dtype=int)
+ for keys in (
+ {"by": pf.col("id") == pf.col("d_id")},
+ {"by": [pf.col("id") == pf.col("d_id"), pf.col("amount") >
pf.col("rate")]},
+ {"left_by": "id", "right_by": "d_id"},
+ {"left_by": [pf.col("id") + 1], "right_by": [pf.col("d_id")]},
+ {"by": (pf.col("id") == pf.col("d_id")) &
(identity(pf.col("amount")) > 0)},
+ ):
+ with self.subTest(keys=str(keys)):
+ before = set(self.t_env.list_user_defined_functions())
+ result = self.left.join_asof(right, on="pt", **keys)
+ self.assertEqual(result.columns, self.left.columns +
right.columns)
+
self.assertEqual(set(self.t_env.list_user_defined_functions()), before)
+ self.assertIn("LookupJoin", result.to_table().explain())
+
+ result = self.left.join_asof(
+ self.right, on="pt", by=["id", pf.col("amount") > pf.col("rate")])
+ self.assertIn("LookupJoin", result.to_table().explain())
+ result = self.left.join_asof(
+ self.right.rename_columns({"rate": "amount"}), on="pt", by=["id",
"amount"])
+ self.assertEqual(result.columns, self.left.columns + ["rt"])
+ self.assertIn("LookupJoin", result.to_table().explain())
+
+ def test_invalid_arguments(self):
+ for kwargs, error, message in (
+ ({"on": None, "by": "id"}, ValueError, "on"),
+ ({"on": 1, "by": "id"}, TypeError, "on"),
+ ({"on": "pt", "by": "id", "how": "full"}, ValueError, "how"),
+ ({"on": "pt", "by": "id", "how": 1}, TypeError, "how"),
+ ({"on": "pt"}, ValueError, "by"),
+ ({"on": "pt", "by": []}, ValueError, "empty"),
+ ({"on": "pt", "by": [1]}, TypeError, "by"),
+ ({"on": "pt", "by": ["id", "id"]}, ValueError, "duplicate"),
+ ({"on": "pt", "by": "id", "left_by": "id"}, ValueError,
"combined"),
+ ({"on": "pt", "left_by": "id"}, ValueError, "together"),
+ ({"on": "pt", "left_by": ["id", "amount"], "right_by": "id"},
+ ValueError, "same number"),
+ ({"on": "amount", "by": "id"}, ValueError, "time attribute"),
+ ({"on": "pt", "left_by": pf.col("*"), "right_by": "id"},
+ ValueError, "single column"),
+ ({"on": "pt", "left_by": pf.col("id").sum, "right_by": "id"},
+ ValueError, "row-wise"),
+ ({"on": "pt", "left_by": "id", "right_by": pf.col("id").sum},
+ ValueError, "row-wise"),
+ ):
+ with self.subTest(kwargs=str(kwargs)):
+ with self.assertRaisesRegex(error, message):
+ self.left.join_asof(self.right, **kwargs)
+ with self.assertRaisesRegex(TypeError, "other"):
+ self.left.join_asof(object(), on="pt", by="id")
+ other_env =
TableEnvironment.create(EnvironmentSettings.in_streaming_mode())
+ with self.assertRaisesRegex(ValueError, "same TableEnvironment"):
+ self.left.join_asof(pf.from_table(other_env.from_elements([(1,)],
["id"])),
+ on="pt", by="id")
+
+ def test_predicate_requires_cross_input_equality_and_no_or(self):
+ right = self.right.rename_columns({"id": "d_id"})
+ for predicate in (
+ pf.col("id") > pf.col("d_id"),
+ pf.col("id") == pf.col("amount"),
+ (pf.col("id") == pf.col("d_id")) | (pf.col("amount") > 1),
+ (pf.col("id") == pf.col("d_id")) &
+ ((pf.col("amount") > 1) | (pf.col("rate") < 10)),
+ ):
+ with self.subTest(predicate=str(predicate)):
+ with self.assertRaises(ValueError):
+ self.left.join_asof(right, on="ts", by=predicate)
+
+ def test_time_attribute_in_equality_keys(self):
+ from pyflink.table.expressions import call_sql
+
+ right = self.versions.rename_columns({"id": "d_id"})
+ for keys in (
+ {"by": (pf.col("id") == pf.col("d_id")) & (pf.col("ts") ==
pf.col("rt"))},
+ {"by": call_sql("id = d_id AND ts = rt")},
+ {"left_by": ["id", "ts"], "right_by": ["d_id", "rt"]},
+ {"left_by": ["id", pf.col("ts").alias("time_key")],
+ "right_by": ["d_id", pf.col("rt").alias("dim_time_key")]},
+ {"left_by": ["id", call_sql("ts")], "right_by": ["d_id",
call_sql("rt")]},
+ {"left_by": ["id", pf.col("ts").cast(TableDataTypes.STRING())],
+ "right_by": ["d_id", pf.col("rt").cast(TableDataTypes.STRING())]},
+ ):
+ with self.subTest(keys=str(keys)):
+ result = self.left.join_asof(right, on="ts", **keys)
+ self.assertEqual(result.columns, self.left.columns +
right.columns)
+ self.assertIn("TemporalJoin", result.to_table().explain())
+
+ shared = self.versions.rename_columns({"rt": "ts"})
+ result = self.left.join_asof(shared, on="ts", by=["id", "ts"])
+ self.assertEqual(result.columns, self.left.columns + ["rate"])
+ self.assertIn("TemporalJoin", result.to_table().explain())
+
+ def test_quoted_keys_and_internal_name_conflicts(self):
+ left = self.left.rename_columns({"id":
"order`id"}).with_column("__pf_asof_time", pf.lit(1))
+ right = self.right.rename_columns({"id": "order`id"})
+ result = left.join_asof(right, on=pf.proctime(), by="order`id")
+ self.assertEqual(result.columns, left.columns + ["rate", "rt"])
+ self.assertIn("LookupJoin", result.to_table().explain())
+
+ def test_planner_validates_dimension_requirements(self):
+ with self.assertRaisesRegex(ValueError, "time attribute"):
+ self.left.join_asof(self.right,
on=pf.col("ts").cast(TableDataTypes.STRING()), by="id")
+ self.t_env.execute_sql(
+ "CREATE TABLE no_primary_key (id INT, rate INT, rt TIMESTAMP(3), "
+ "WATERMARK FOR rt AS rt) WITH ('connector' = 'COLLECTION')")
+ right = pf.from_table(self.t_env.from_path("no_primary_key"))
+ with self.assertRaisesRegex(Exception, "primary key"):
+ self.left.join_asof(right, on="ts", by="id").to_table().explain()
+ with self.assertRaisesRegex(Exception, "primary key|Processing-time
temporal join"):
+ self.left.join_asof(pf.from_records([(1,)], schema=["id"]),
+ on="pt", by="id").to_table().explain()
+
+ def test_event_time_expression_keys_preserve_primary_key(self):
+ from pyflink.table.expressions import call_sql
+
+ right = self.versions.rename_columns({"id": "d_id"})
+ for left_key, right_key in (
+ (pf.col("id"), pf.col("d_id")),
+ ([pf.col("id").alias("key")], [pf.col("d_id").alias("dim_key")]),
+ ("id", pf.col("d_id").cast(TableDataTypes.INT())),
+ ("id", pf.col("d_id").cast(TableDataTypes.INT()).alias("dim_key")),
+ ("id", call_sql("CAST(d_id AS INT)")),
+ ):
+ with self.subTest(right_key=str(right_key)):
+ result = self.left.join_asof(right, on="ts", left_by=left_key,
right_by=right_key)
+ self.assertIn("TemporalJoin", result.to_table().explain())
+
+ def test_time_attribute_in_residual_predicates(self):
+ from pyflink.table.expressions import call_sql
+
+ right = self.versions.rename_columns({"id": "d_id"})
+ for predicate in (
+ (pf.col("id") == pf.col("d_id")) & (pf.col("ts") > pf.col("rt")),
+ call_sql("id = d_id AND ts > rt"),
+ (pf.col("id") == pf.col("d_id")) & ~(pf.col("ts") == pf.col("rt")),
+ call_sql("id = d_id AND CASE WHEN ts = rt THEN FALSE ELSE TRUE
END"),
+ (pf.col("id") == pf.col("d_id")) &
+ (pf.col("ts").cast(TableDataTypes.STRING()) == "2024-01-01
00:00:02"),
+ ):
+ with self.subTest(predicate=str(predicate)):
+ result = self.left.join_asof(right, on="ts", by=predicate,
how="left")
+ self.assertIn("TemporalJoin", result.to_table().explain())
+
+ def test_type_changing_key_cast_is_preserved(self):
+ from pyflink.dataframe.dataframe import _asof_field_index
+ from pyflink.table.expressions import call_sql
+
+ for key in (
+ pf.col("id").cast(TableDataTypes.BIGINT()), call_sql("CAST(id AS
BIGINT)"),
+ pf.col("rt").cast(TableDataTypes.TIMESTAMP(3)), call_sql("CAST(rt
AS TIMESTAMP(3))"),
+ ):
+ with self.subTest(key=str(key)):
+ selected = self.versions.to_table().select(key)
+ resolved =
selected._j_table.getQueryOperation().getProjectList().get(0)
+ self.assertIsNone(_asof_field_index(resolved))
+
+ def test_sql_expressions_follow_the_same_validation(self):
+ from pyflink.table.expressions import call_sql
+
+ right = self.versions.rename_columns({"id": "d_id"})
+ result = self.left.join_asof(right, on=call_sql("ts"), by=call_sql("id
= d_id"))
+ self.assertIn("TemporalJoin", result.to_table().explain())
+ result = self.left.join_asof(
+ right, on="ts", left_by=call_sql("id"), right_by=call_sql("d_id"))
+ self.assertIn("TemporalJoin", result.to_table().explain())
+ for predicate in (
+ call_sql("id > d_id"),
+ call_sql("id = d_id OR amount > rate"),
+ (pf.col("id") == pf.col("d_id")) & call_sql("amount > rate OR
amount = 0"),
+ ):
+ with self.subTest(predicate=str(predicate)):
+ with self.assertRaises(ValueError):
+ self.left.join_asof(right, on=call_sql("ts"), by=predicate)
+ time_first = self.left.select("pt", "id", "ts")
+ for on in (pf.col("ts").max, call_sql("MAX(ts)")):
+ with self.subTest(on=str(on)):
+ with self.assertRaisesRegex(ValueError, "row-wise"):
+ time_first.join_asof(self.right, on=on, by="id")
+ for keys in (
+ {"left_by": call_sql("SUM(id) + 1"), "right_by": "d_id"},
+ {"left_by": "id", "right_by": call_sql("SUM(d_id)")},
+ ):
+ with self.subTest(keys=str(keys)):
+ with self.assertRaisesRegex(ValueError, "row-wise"):
+ self.left.join_asof(right, on="ts", **keys)
+
+
class DataFrameJoinTests(PyFlinkDataFrameUTTestCase):
def setUp(self):
super().setUp()
@@ -3967,6 +4196,115 @@ class
DataFrameExplodeStreamITTests(PyFlinkStreamDataFrameTestCase):
self.assertCountEqual(results[0].union_all(results[1]).collect(),
expected_rows)
+class DataFrameAsOfJoinITTests(PyFlinkStreamDataFrameTestCase):
+ def test_lookup_join(self):
+ gateway = get_gateway()
+ factory =
gateway.jvm.org.apache.flink.table.utils.TestCollectionTableFactory
+ factory.reset()
+ self.addCleanup(factory.reset)
+ rows = gateway.jvm.java.util.ArrayList()
+ for key, rate in [(1, 10), (2, 20)]:
+ row = gateway.jvm.org.apache.flink.types.Row(2)
+ row.setField(0, key)
+ row.setField(1, rate)
+ rows.add(row)
+ factory.initData(gateway.jvm.java.util.ArrayList(), rows, None)
+ self.t_env.execute_sql(
+ "CREATE TEMPORARY TABLE lookup_rates (id INT, rate INT) "
+ "WITH ('connector' = 'COLLECTION')")
+ left = pf.from_table(self.t_env.sql_query(
+ "SELECT *, PROCTIME() AS pt FROM (VALUES (1, 15), (2, 5), (3, 30))
T(id, amount)"))
+ right = pf.from_table(self.t_env.from_path("lookup_rates"))
+ renamed = right.rename_columns({"id": "d_id"})
+ results = [
+ left.join_asof(right, on=pf.proctime(), by="id"),
+ left.join_asof(renamed, on="pt", left_by="id", right_by="d_id",
how="left"),
+ left.join_asof(right, on="pt", by=["id", pf.col("amount") >
pf.col("rate")],
+ how="left"),
+ ]
+ combined = results[0].select(pf.lit(0).alias("variant"), "id", "rate")
+ for index, result in enumerate(results[1:], 1):
+ combined = combined.union_all(result.select(
+ pf.lit(index).alias("variant"), "id", "rate"))
+ self.assertCountEqual(combined.collect(), [
+ Row(0, 1, 10), Row(0, 2, 20),
+ Row(1, 1, 10), Row(1, 2, 20), Row(1, 3, None),
+ Row(2, 1, 10), Row(2, 2, None), Row(2, 3, None),
+ ])
+
+ def test_event_time_join_selects_historical_versions(self):
+ orders_path = os.path.join(self.tempdir, "asof_orders.csv")
+ rates_path = os.path.join(self.tempdir, "asof_rates.csv")
+ with open(orders_path, "w", encoding="utf-8") as output:
+ output.write("1,2024-01-01 00:00:02\n2,2024-01-01 00:00:02\n"
+ "1,2024-01-01 00:00:06\n3,2024-01-01 00:00:06\n")
+ with open(rates_path, "w", encoding="utf-8") as output:
+ output.write("1,10,2024-01-01 00:00:01\n2,30,2024-01-01 00:00:03\n"
+ "1,20,2024-01-01 00:00:04\n")
+ self.t_env.execute_sql(
+ "CREATE TEMPORARY TABLE asof_orders (id INT, ts TIMESTAMP(3), "
+ "WATERMARK FOR ts AS ts - INTERVAL '1' SECOND) "
+ "WITH ('connector' = 'filesystem', 'format' = 'csv', 'path' =
'%s')" % orders_path)
+ self.t_env.execute_sql(
+ "CREATE TEMPORARY TABLE asof_rates (id INT, rate INT, rt
TIMESTAMP(3), "
+ "WATERMARK FOR rt AS rt - INTERVAL '1' SECOND, PRIMARY KEY (id)
NOT ENFORCED) "
+ "WITH ('connector' = 'filesystem', 'format' = 'csv', 'path' =
'%s')" % rates_path)
+ left = pf.from_table(self.t_env.from_path("asof_orders"))
+ right = pf.from_table(self.t_env.from_path("asof_rates"))
+ inner = left.join_asof(right, on="ts", by="id").select(
+ pf.lit("inner").alias("kind"), "id", "rate")
+ outer = left.join_asof(
+ right.rename_columns({"id": "d_id"}), on=pf.col("ts"),
+ left_by=pf.col("id"),
right_by=pf.col("d_id").cast(TableDataTypes.INT()),
+ how="left").select(
+ pf.lit("left").alias("kind"), "id", "rate")
+ residual = left.join_asof(
+ right.rename_columns({"id": "d_id"}), on="ts", how="left",
+ by=(pf.col("id") == pf.col("d_id")) &
+ (pf.col("ts") > pf.col("rt") + pf.lit(1).seconds)).select(
+ pf.lit("residual").alias("kind"), "id", "rate")
+
self.assertCountEqual(inner.union_all(outer).union_all(residual).collect(), [
+ Row("inner", 1, 10), Row("inner", 1, 20),
+ Row("left", 1, 10), Row("left", 1, 20),
+ Row("left", 2, None), Row("left", 3, None),
+ Row("residual", 1, None), Row("residual", 1, 20),
+ Row("residual", 2, None), Row("residual", 3, None),
+ ])
+
+ def test_event_time_equality_matches_only_exact_versions(self):
+ orders_path = os.path.join(self.tempdir, "time_key_orders.csv")
+ rates_path = os.path.join(self.tempdir, "time_key_rates.csv")
+ with open(orders_path, "w", encoding="utf-8") as output:
+ output.write("1,2024-01-01 00:00:02\n1,2024-01-01 00:00:05\n"
+ "1,2024-01-01 00:00:06\n2,2024-01-01 00:00:06\n")
+ with open(rates_path, "w", encoding="utf-8") as output:
+ output.write("1,10,2024-01-01 00:00:01\n1,20,2024-01-01
00:00:05\n")
+ self.t_env.execute_sql(
+ "CREATE TEMPORARY TABLE time_key_orders (id INT, ts TIMESTAMP(3), "
+ "WATERMARK FOR ts AS ts - INTERVAL '1' SECOND) "
+ "WITH ('connector' = 'filesystem', 'format' = 'csv', 'path' =
'%s')" % orders_path)
+ self.t_env.execute_sql(
+ "CREATE TEMPORARY TABLE time_key_rates (d_id INT, rate INT, rt
TIMESTAMP(3), "
+ "WATERMARK FOR rt AS rt - INTERVAL '1' SECOND, PRIMARY KEY (d_id)
NOT ENFORCED) "
+ "WITH ('connector' = 'filesystem', 'format' = 'csv', 'path' =
'%s')" % rates_path)
+ left = pf.from_table(self.t_env.from_path("time_key_orders"))
+ right = pf.from_table(self.t_env.from_path("time_key_rates"))
+ inner = left.join_asof(
+ right, on="ts",
+ by=(pf.col("id") == pf.col("d_id")) & (pf.col("ts") ==
pf.col("rt"))).select(
+ pf.lit("inner").alias("kind"), "id", "ts", "rate")
+ outer = left.join_asof(
+ right, on="ts", left_by=["id", "ts"], right_by=["d_id", "rt"],
+ how="left").select(pf.lit("left").alias("kind"), "id", "ts",
"rate")
+ self.assertCountEqual(inner.union_all(outer).collect(), [
+ Row("inner", 1, datetime(2024, 1, 1, 0, 0, 5), 20),
+ Row("left", 1, datetime(2024, 1, 1, 0, 0, 2), None),
+ Row("left", 1, datetime(2024, 1, 1, 0, 0, 5), 20),
+ Row("left", 1, datetime(2024, 1, 1, 0, 0, 6), None),
+ Row("left", 2, datetime(2024, 1, 1, 0, 0, 6), None),
+ ])
+
+
class DataFrameWindowITTests(PyFlinkStreamDataFrameTestCase):
@classmethod
def setUpClass(cls):
diff --git
a/flink-python/src/test/java/org/apache/flink/table/utils/TestCollectionTableFactory.java
b/flink-python/src/test/java/org/apache/flink/table/utils/TestCollectionTableFactory.java
index d59ad693926..0b65b63b66e 100644
---
a/flink-python/src/test/java/org/apache/flink/table/utils/TestCollectionTableFactory.java
+++
b/flink-python/src/test/java/org/apache/flink/table/utils/TestCollectionTableFactory.java
@@ -260,7 +260,7 @@ public class TestCollectionTableFactory
}
}
- static class TemporalTableFetcher extends TableFunction<Row> {
+ public static class TemporalTableFetcher extends TableFunction<Row> {
private final LinkedList<Row> dimData;
private final int[] keyes;
