dianfu commented on code in PR #29201:
URL: https://github.com/apache/flink/pull/29201#discussion_r4055848382
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -622,6 +623,119 @@ def top_n(
distinct = drop_duplicates
unique = drop_duplicates
+ # ======================== Joins ========================
+
+ @PublicEvolving()
+ def join(
Review Comment:
`DataFrame.join` should be added to
`flink-python/docs/reference/pyflink.dataframe/dataframe.rst`.
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -1463,6 +1417,327 @@ def agg(self, *aggs: Expression, **named_aggs:
Expression) -> DataFrame:
# ======================== 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], List[str], List[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,
+ left_key_names,
+ right_key_names,
+ )
+
+
+def _serialize_join_predicate(
+ left_table: Table,
+ right_table: Table,
+ predicate: Expression,
+) -> 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()
Review Comment:
Serializing the predicate with `asSerializableString()` breaks inline UDF
predicates. For example:
same_value = pf.udf(
lambda left, right: left == right,
return_dtype=bool,
)
result = left.join(
right,
on=(pf.col("id") == pf.col("right_id"))
& same_value(pf.col("left_value"), pf.col("right_value")),
)
This fails with `Only functions that have been registered before are
serializable`, while the underlying Table API accepts the same predicate
directly.
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -622,6 +623,119 @@ def top_n(
distinct = drop_duplicates
unique = drop_duplicates
+ # ======================== 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. A predicate must contain at least one
equality condition
Review Comment:
The SQL planner supports non-equi-only join predicates, so I don't think we
need to require an equality predicate here. This restriction is not validated
actually.
Could we remove this restriction from the documentation? For streaming, it
may also be worth mentioning that a join without equi keys can use singleton
distribution and may be expensive.
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -1463,6 +1417,327 @@ def agg(self, *aggs: Expression, **named_aggs:
Expression) -> DataFrame:
# ======================== 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], List[str], List[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,
+ left_key_names,
+ right_key_names,
+ )
+
+
+def _serialize_join_predicate(
+ left_table: Table,
+ right_table: Table,
+ predicate: Expression,
+) -> 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()
+ 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,
+) -> Table:
+ left_alias, right_alias, predicate_sql = _serialize_join_predicate(
+ left_table, right_table, predicate
+ )
+ 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 = "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_join_sql(
+ left_table: Table,
+ right_table: Table,
+ output_columns: List[str],
+ left_key_names: List[str],
+ right_key_names: List[str],
+) -> Table:
+ left_alias, right_alias = "__pf_join_left", "__pf_join_right"
+ 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_keys = ", ".join(
+ "%s.%s" % (left_alias_sql, _quote_identifier(name))
+ for name in left_key_names
+ )
+ right_keys = ", ".join(
+ "%s.%s" % (right_alias_sql, _quote_identifier(name))
+ for name in right_key_names
+ )
+ if len(left_key_names) > 1:
+ left_keys = "(%s)" % left_keys
+ query = (
+ "SELECT %s FROM %s AS %s WHERE %s IN ("
+ "SELECT %s FROM %s AS %s)"
+ % (
+ select_list,
+ _quote_identifier(str(left_table)),
+ left_alias_sql,
+ left_keys,
+ right_keys,
+ _quote_identifier(str(right_table)),
+ right_alias_sql,
+ )
+ )
+ return left_table._t_env.sql_query(query)
+
+
+def _build_anti_join_sql(
+ left_table: Table,
+ right_table: Table,
+ predicate: Expression,
+ output_columns: List[str],
+) -> Table:
+ left_alias, right_alias, predicate_sql = _serialize_join_predicate(
+ left_table, right_table, predicate
+ )
+ 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
+ )
+ right_columns = list(right_table.get_resolved_schema().get_column_names())
+ match_marker = _unique_name("__pf_join_match", set(right_columns))
+ match_marker_sql = _quote_identifier(match_marker)
+ query = (
Review Comment:
This rewrite remains a `LeftOuterJoin` followed by `IS NULL` in the
optimized plan. Could you rewrite it using `NOT EXISTS`, which produces a
native `LeftAntiJoin`?
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -622,6 +623,119 @@ def top_n(
distinct = drop_duplicates
unique = drop_duplicates
+ # ======================== 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. A predicate must contain at least one
equality condition
+ between the inputs. Use ``left_on`` and ``right_on`` together when the
key names differ.
+ 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. A
``semi`` join requires
+ shared named keys through ``on`` or keys through
``left_on``/``right_on`` instead of a
+ boolean predicate.
+
+ :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"])
+ >>> orders.join(customers, on="customer_id")
+ >>> orders.join(
Review Comment:
Could we expand the examples for expression-based join conditions? Currently
the examples only cover named keys. It would be helpful to include:
1. an equality expression;
2. a compound equality and range predicate;
3. a non-equi-only predicate;
4. a semi or anti join example.
For example:
```
orders.join(
customers,
on=pf.col("customer_id") == pf.col("id"),
)
orders.join(
rules,
on=(pf.col("customer_id") == pf.col("rule_customer_id"))
& (pf.col("amount") >= pf.col("min_amount")),
)
orders.join(
rules,
on=pf.col("amount") >= pf.col("min_amount"),
)
```
--
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]
