dianfu commented on code in PR #29513:
URL: https://github.com/apache/flink/pull/29513#discussion_r4310684525
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -3456,6 +3586,90 @@ def _prepare_join(
)
+def _asof_field_index(expression) -> Optional[int]:
+ if expression.getClass().getSimpleName() == "FieldReferenceExpression":
+ return expression.getFieldIndex()
+ if expression.getClass().getSimpleName() == "RexNodeExpression":
+ node = expression.getRexNode()
+ 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, -1, 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, time_index: int, left_arity: int) ->
Tuple[Set[int], bool]:
+ class_name = expression.getClass().getSimpleName()
+ if class_name == "RexNodeExpression":
+ return _asof_expression_info(expression.getRexNode(), time_index,
left_arity)
+ if class_name == "FieldReferenceExpression":
Review Comment:
This check also rejects the time column in residual predicates, even when it
is not an equality key:
```python
left.join_asof(
right,
on="ts",
how="left",
by=(pf.col("id") == pf.col("d_id")) & (pf.col("ts") > pf.col("rt")),
)
```
This raises `ValueError`, while the equivalent SQL works. Moving `ts > rt`
to a filter before or after the join changes left-join semantics. Could we
restrict this validation to equality keys?
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -3389,13 +3494,38 @@ 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)
+ time_index = left_columns.index(asof_time) if table is
left_table else -1
+ _asof_expression_info(resolved, time_index, len(columns))
+ field_index = _asof_field_index(resolved)
Review Comment:
When `d_id` is already an INT primary key, these equivalent forms behave
differently:
```python
key = pf.col("d_id").cast(DataTypes.INT())
# Plans successfully.
left.join_asof(right, on="ts", by=pf.col("id") == key).to_table().explain()
# Fails: the right primary key is not included in the equality condition.
left.join_asof(right, on="ts", left_by="id",
right_by=key).to_table().explain()
```
The paired-key form projects the no-op cast into a new column, losing the
primary-key reference.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
