adriangb commented on code in PR #25255:
URL: https://github.com/apache/datafusion/pull/25255#discussion_r4000258378
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -979,6 +979,58 @@ impl HashJoinExec {
Arc::new(DynamicFilterPhysicalExpr::new(right_keys, lit(true)))
}
+ /// Join types whose output rows all carry a matching key on both sides.
+ ///
+ /// For these a parent filter over one side's join keys can be transferred
+ /// to the other side's input: an input row that fails the transferred
+ /// filter can only pair with rows that fail the original, so pruning it
+ /// changes nothing, and once the transferred filter is applied exactly on
+ /// one side every output row satisfies the original. Outer, anti and mark
+ /// joins also emit unmatched rows, whose key on the other side is absent,
+ /// so the transferred filter is not exact for them.
+ fn supports_key_transfer(join_type: JoinType) -> bool {
+ matches!(
+ join_type,
+ JoinType::Inner | JoinType::LeftSemi | JoinType::RightSemi
+ )
+ }
+
+ /// Maps each output column that is a plain `Column` join key on one side
+ /// to the key expression on the other side, as `(to_right, to_left)`.
+ ///
+ /// `column_indices` are the (projected) output columns of this join. A key
+ /// column that appears in several `on` pairs maps to the first of them.
+ fn key_transfer_maps(
+ &self,
+ column_indices: &[ColumnIndex],
+ ) -> (KeyTransferMap, KeyTransferMap) {
+ let mut to_right = HashMap::new();
+ let mut to_left = HashMap::new();
+ for (output_idx, ci) in column_indices.iter().enumerate() {
+ let (map, other_key) = match ci.side {
+ JoinSide::Left => (
+ &mut to_right,
+ self.on
+ .iter()
+ .find(|(left_key, _)| is_column_at(left_key, ci.index))
Review Comment:
The comment says that the first `on` pair wins. No test checks this. A
mutant that selects the last pair passes all tests.
Both choices are correct, because all pairs are equal for a matched row. A
test with `ON a.k = b.x AND a.k = b.y` can pin the documented choice.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -979,6 +979,58 @@ impl HashJoinExec {
Arc::new(DynamicFilterPhysicalExpr::new(right_keys, lit(true)))
}
+ /// Join types whose output rows all carry a matching key on both sides.
+ ///
+ /// For these a parent filter over one side's join keys can be transferred
+ /// to the other side's input: an input row that fails the transferred
+ /// filter can only pair with rows that fail the original, so pruning it
+ /// changes nothing, and once the transferred filter is applied exactly on
+ /// one side every output row satisfies the original. Outer, anti and mark
+ /// joins also emit unmatched rows, whose key on the other side is absent,
+ /// so the transferred filter is not exact for them.
+ fn supports_key_transfer(join_type: JoinType) -> bool {
+ matches!(
+ join_type,
+ JoinType::Inner | JoinType::LeftSemi | JoinType::RightSemi
Review Comment:
No test covers `RightSemi`. A mutant that removes `RightSemi` from this
match passes all tests.
Please add a `RightSemi` variant of
`test_hashjoin_parent_filter_transfer_semi_join_different_key_names`. On
`main`, that plan pushes nothing to the left scan. With this PR, it pushes `k@0
= x` to the left scan.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -979,6 +979,58 @@ impl HashJoinExec {
Arc::new(DynamicFilterPhysicalExpr::new(right_keys, lit(true)))
}
+ /// Join types whose output rows all carry a matching key on both sides.
+ ///
+ /// For these a parent filter over one side's join keys can be transferred
+ /// to the other side's input: an input row that fails the transferred
+ /// filter can only pair with rows that fail the original, so pruning it
+ /// changes nothing, and once the transferred filter is applied exactly on
+ /// one side every output row satisfies the original. Outer, anti and mark
+ /// joins also emit unmatched rows, whose key on the other side is absent,
+ /// so the transferred filter is not exact for them.
+ fn supports_key_transfer(join_type: JoinType) -> bool {
+ matches!(
+ join_type,
+ JoinType::Inner | JoinType::LeftSemi | JoinType::RightSemi
+ )
+ }
+
+ /// Maps each output column that is a plain `Column` join key on one side
+ /// to the key expression on the other side, as `(to_right, to_left)`.
+ ///
+ /// `column_indices` are the (projected) output columns of this join. A key
+ /// column that appears in several `on` pairs maps to the first of them.
+ fn key_transfer_maps(
+ &self,
+ column_indices: &[ColumnIndex],
+ ) -> (KeyTransferMap, KeyTransferMap) {
+ let mut to_right = HashMap::new();
+ let mut to_left = HashMap::new();
+ for (output_idx, ci) in column_indices.iter().enumerate() {
+ let (map, other_key) = match ci.side {
+ JoinSide::Left => (
+ &mut to_right,
+ self.on
+ .iter()
+ .find(|(left_key, _)| is_column_at(left_key, ci.index))
+ .map(|(_, right_key)| right_key),
+ ),
+ JoinSide::Right => (
+ &mut to_left,
+ self.on
+ .iter()
+ .find(|(_, right_key)| is_column_at(right_key,
ci.index))
+ .map(|(left_key, _)| left_key),
+ ),
+ JoinSide::None => continue,
Review Comment:
This arm cannot run. Only mark joins produce `JoinSide::None`, and
`supports_key_transfer` excludes mark joins.
If you keep the arm as a safety check, please add a comment that says so.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2512,6 +2540,74 @@ mod proto_tests {
}
}
+/// Output column index of a join, mapped to the equivalent join-key expression
+/// on the other side of the join (in that side's input schema).
+type KeyTransferMap = HashMap<usize, PhysicalExprRef>;
+
+fn is_column_at(expr: &PhysicalExprRef, index: usize) -> bool {
+ expr.downcast_ref::<Column>()
+ .is_some_and(|column| column.index() == index)
+}
+
+/// Marks every parent filter whose columns are all join keys in `key_map` as
+/// supported for `child`, rewritten over the other side's key expressions.
+///
+/// Filters `child` already accepts directly are left alone, as are filters
+/// that reference a non-key column or no column at all: the former cannot be
+/// expressed on the other side, the latter were already routed by the plain
+/// column analysis.
+fn transfer_key_filters(
+ parent_filters: &[Arc<dyn PhysicalExpr>],
+ key_map: &KeyTransferMap,
+ child: &mut ChildFilterDescription,
+) -> Result<()> {
+ if key_map.is_empty() {
+ return Ok(());
+ }
+ for (filter, pushed) in
parent_filters.iter().zip(child.parent_filters.iter_mut()) {
+ if matches!(pushed.discriminant, PushedDown::Yes) {
Review Comment:
This guard cannot trigger. `to_right` contains only left-side output
indices, and `to_left` contains only right-side output indices. A filter is
directly pushable to a child only when all its columns are on that side. Thus a
filter cannot be both directly pushable and transferable to the same child.
A mutant that removes the guard passes all tests. Please remove the guard,
or add a comment that says it is a safety check.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2512,6 +2540,74 @@ mod proto_tests {
}
}
+/// Output column index of a join, mapped to the equivalent join-key expression
+/// on the other side of the join (in that side's input schema).
+type KeyTransferMap = HashMap<usize, PhysicalExprRef>;
+
+fn is_column_at(expr: &PhysicalExprRef, index: usize) -> bool {
+ expr.downcast_ref::<Column>()
+ .is_some_and(|column| column.index() == index)
+}
+
+/// Marks every parent filter whose columns are all join keys in `key_map` as
+/// supported for `child`, rewritten over the other side's key expressions.
+///
+/// Filters `child` already accepts directly are left alone, as are filters
+/// that reference a non-key column or no column at all: the former cannot be
+/// expressed on the other side, the latter were already routed by the plain
+/// column analysis.
+fn transfer_key_filters(
+ parent_filters: &[Arc<dyn PhysicalExpr>],
+ key_map: &KeyTransferMap,
+ child: &mut ChildFilterDescription,
+) -> Result<()> {
+ if key_map.is_empty() {
+ return Ok(());
+ }
+ for (filter, pushed) in
parent_filters.iter().zip(child.parent_filters.iter_mut()) {
+ if matches!(pushed.discriminant, PushedDown::Yes) {
+ continue;
+ }
+ if let Some(transferred) = transfer_filter_across_keys(filter,
key_map)? {
+ *pushed = PushedDownPredicate::supported(transferred);
+ }
+ }
+ Ok(())
+}
+
+/// Rewrites `filter` over the other side's join keys, or returns `None` when
+/// it references a column that is not a transferable key, or no column.
+///
+/// A [`DynamicFilterPhysicalExpr`] comes out as a view sharing the original's
+/// state with its key columns remapped, so it keeps tracking the build side.
+fn transfer_filter_across_keys(
+ filter: &Arc<dyn PhysicalExpr>,
+ key_map: &KeyTransferMap,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ let mut all_keys = true;
+ let mut any_column = false;
Review Comment:
The `any_column` condition is redundant. `try_remap` already marks a filter
with no columns as supported for both children, as the comment above says.
A mutant that removes this condition passes all tests.
##########
datafusion/core/tests/physical_optimizer/filter_pushdown.rs:
##########
@@ -1754,6 +1754,293 @@ fn
test_hashjoin_parent_filter_pushdown_semi_anti_join() {
assert_parent_filter_remains(plan);
}
+/// A parent filter over one side's join keys is transferred to the other side,
+/// rewritten over that side's key expressions, even when the key names differ.
+/// Filters over non-key columns stay on their own side.
+#[test]
+fn test_hashjoin_parent_filter_transferred_across_join_keys() {
+ let build_side_schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Utf8, false),
+ Field::new("build_val", DataType::Utf8, false),
+ ]));
+ let build_scan = TestScanBuilder::new(Arc::clone(&build_side_schema))
+ .with_support(true)
+ .build();
+
+ let probe_side_schema = Arc::new(Schema::new(vec![
+ Field::new("pid", DataType::Utf8, false),
+ Field::new("probe_val", DataType::Utf8, false),
+ ]));
+ let probe_scan = TestScanBuilder::new(Arc::clone(&probe_side_schema))
+ .with_support(true)
+ .build();
+
+ let on = vec![(
+ col("id", &build_side_schema).unwrap(),
+ col("pid", &probe_side_schema).unwrap(),
+ )];
+ let join = Arc::new(
+ HashJoinExec::try_new(
+ build_scan,
+ probe_scan,
+ on,
+ None,
+ &JoinType::Inner,
+ None,
+ PartitionMode::Partitioned,
+ datafusion_common::NullEquality::NullEqualsNothing,
+ false,
+ )
+ .unwrap(),
+ );
+
+ let join_schema = join.schema();
+
+ let build_key_filter = col_lit_predicate("id", "aa", &join_schema);
+ let probe_key_filter = col_lit_predicate("pid", "ab", &join_schema);
+ let build_val_filter = col_lit_predicate("build_val", "x", &join_schema);
+
+ let filter =
+ Arc::new(FilterExec::try_new(build_key_filter, Arc::clone(&join) as
_).unwrap());
+ let filter = Arc::new(FilterExec::try_new(probe_key_filter,
filter).unwrap());
+ let plan = Arc::new(FilterExec::try_new(build_val_filter, filter).unwrap())
+ as Arc<dyn ExecutionPlan>;
+
+ insta::assert_snapshot!(
+ OptimizationTest::new(Arc::clone(&plan), FilterPushdown::new(), true),
+ @r"
+ OptimizationTest:
+ input:
+ - FilterExec: build_val@1 = x
+ - FilterExec: pid@2 = ab
+ - FilterExec: id@0 = aa
+ - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(id@0,
pid@0)]
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[id, build_val], file_type=test, pushdown_supported=true
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[pid, probe_val], file_type=test, pushdown_supported=true
+ output:
+ Ok:
+ - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(id@0, pid@0)]
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[id, build_val], file_type=test, pushdown_supported=true,
predicate=id@0 = aa AND id@0 = ab AND build_val@1 = x
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[pid, probe_val], file_type=test, pushdown_supported=true,
predicate=pid@0 = aa AND pid@0 = ab
+ "
+ );
+}
+
+/// The non-output side of a semi join receives key filters through the same
+/// transfer, so differently named keys work too.
+#[test]
+fn test_hashjoin_parent_filter_transfer_semi_join_different_key_names() {
Review Comment:
This test uses different key names, but it does not cover the bug that the
removed code had.
On `main` (a407990b44), the removed code inserted the key's output index
into the other side's allowed set, and then `FilterRemapper::try_remap` mapped
the column by name. Test shape: `LeftSemi`, `on = [(k@0, j@0)]`, left schema
`[k, v]`, right schema `[j, w, k]`, where the right `k` is not a key. Parent
filter: `k@0 = 'x'`.
```text
main: right scan gets predicate=k@2 = x (wrong column), and the
FilterExec is removed
PR: right scan gets predicate=j@0 = x
```
That is a wrong-result bug on `main`, and this PR corrects it. The
description says that the old code "silently did nothing" when the key names
differ. That is not the full story. The existing test
`test_hashjoin_parent_filter_pushdown_semi_anti_join` did not find the bug,
because both keys in that test have the name `k`. The `RightSemi` branch had
the mirror-image bug.
Please add this shape as a regression test, and add the fix to the PR
description.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2523,7 +2619,8 @@ fn lr_is_preserved(join_type: JoinType) -> (bool, bool) {
JoinType::Left => (true, false),
JoinType::Right => (false, true),
JoinType::Full => (false, false),
- // Callers restrict the non-output side of semi joins to join-key
columns.
+ // The non-output side of a semi join only receives filters transferred
Review Comment:
This comment describes a rule that the code does not enforce. See my comment
on the `transfer_key_filters` calls above.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2512,6 +2540,74 @@ mod proto_tests {
}
}
+/// Output column index of a join, mapped to the equivalent join-key expression
+/// on the other side of the join (in that side's input schema).
+type KeyTransferMap = HashMap<usize, PhysicalExprRef>;
+
+fn is_column_at(expr: &PhysicalExprRef, index: usize) -> bool {
+ expr.downcast_ref::<Column>()
+ .is_some_and(|column| column.index() == index)
+}
+
+/// Marks every parent filter whose columns are all join keys in `key_map` as
+/// supported for `child`, rewritten over the other side's key expressions.
+///
+/// Filters `child` already accepts directly are left alone, as are filters
+/// that reference a non-key column or no column at all: the former cannot be
+/// expressed on the other side, the latter were already routed by the plain
+/// column analysis.
+fn transfer_key_filters(
+ parent_filters: &[Arc<dyn PhysicalExpr>],
+ key_map: &KeyTransferMap,
+ child: &mut ChildFilterDescription,
+) -> Result<()> {
+ if key_map.is_empty() {
+ return Ok(());
+ }
+ for (filter, pushed) in
parent_filters.iter().zip(child.parent_filters.iter_mut()) {
+ if matches!(pushed.discriminant, PushedDown::Yes) {
+ continue;
+ }
+ if let Some(transferred) = transfer_filter_across_keys(filter,
key_map)? {
+ *pushed = PushedDownPredicate::supported(transferred);
+ }
+ }
+ Ok(())
+}
+
+/// Rewrites `filter` over the other side's join keys, or returns `None` when
+/// it references a column that is not a transferable key, or no column.
+///
+/// A [`DynamicFilterPhysicalExpr`] comes out as a view sharing the original's
+/// state with its key columns remapped, so it keeps tracking the build side.
+fn transfer_filter_across_keys(
+ filter: &Arc<dyn PhysicalExpr>,
+ key_map: &KeyTransferMap,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ let mut all_keys = true;
+ let mut any_column = false;
+ let transformed = Arc::clone(filter).transform_down(|expr| {
+ let Some(column) = expr.downcast_ref::<Column>() else {
+ return Ok(Transformed::no(expr));
+ };
+ any_column = true;
+ match key_map.get(&column.index()) {
+ // The replacement is already in the other side's schema: do not
+ // descend into it, its columns are not indices of this join.
+ Some(other_key) => Ok(Transformed::new(
+ Arc::clone(other_key),
+ true,
+ TreeNodeRecursion::Jump,
Review Comment:
This `Jump` is necessary for termination, and no test covers it. A mutant
that uses `Continue` passes all tests.
Example: `on = [(CAST(k@0 AS Int64), j@0)]`, `JoinType::Inner`, `projection
= Some(vec![2])`. Then `to_left = {0: CAST(k@0 AS Int64)}`, and the inner
column of the substituted `CAST` also has index 0. With `Continue`, the
traversal substitutes the `CAST` again without end, and the test process runs
out of memory.
Please add a test with a `CastExpr` key and a projection that puts the other
side's key at output index 0.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -1903,6 +1919,18 @@ impl ExecutionPlan for HashJoinExec {
ChildFilterDescription::all_unsupported(&parent_filters)
};
+ // Transfer filters across the equi-join keys: a parent filter over one
+ // side's join-key columns holds for every matching row of the other
+ // side too, so it is also pushed there, rewritten over that side's key
+ // expressions. This is how a dynamic filter from a join above reaches
+ // the scans on both sides of this join, and how a semi join prunes its
+ // non-output side.
+ if Self::supports_key_transfer(self.join_type) {
+ let (to_right, to_left) = self.key_transfer_maps(&column_indices);
+ transfer_key_filters(&parent_filters, &to_right, &mut
right_child)?;
Review Comment:
`transfer_key_filters` writes into the child descriptions after the
`lr_is_preserved` gate. It also overwrites `all_unsupported` entries.
For the three permitted join types, both sides are preserved. Thus the gate
has no effect today. If a future change adds a join type to
`supports_key_transfer` that is not preserved on one side, the transfer pushes
filters to that side without a warning. A mutant that returns `(true, false)`
for semi joins passes all tests.
Please transfer only into a side that `lr_is_preserved` permits, or add an
assertion that couples the two functions.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -2512,6 +2540,74 @@ mod proto_tests {
}
}
+/// Output column index of a join, mapped to the equivalent join-key expression
+/// on the other side of the join (in that side's input schema).
+type KeyTransferMap = HashMap<usize, PhysicalExprRef>;
+
+fn is_column_at(expr: &PhysicalExprRef, index: usize) -> bool {
+ expr.downcast_ref::<Column>()
+ .is_some_and(|column| column.index() == index)
+}
+
+/// Marks every parent filter whose columns are all join keys in `key_map` as
+/// supported for `child`, rewritten over the other side's key expressions.
+///
+/// Filters `child` already accepts directly are left alone, as are filters
+/// that reference a non-key column or no column at all: the former cannot be
+/// expressed on the other side, the latter were already routed by the plain
+/// column analysis.
+fn transfer_key_filters(
+ parent_filters: &[Arc<dyn PhysicalExpr>],
+ key_map: &KeyTransferMap,
+ child: &mut ChildFilterDescription,
+) -> Result<()> {
+ if key_map.is_empty() {
+ return Ok(());
+ }
+ for (filter, pushed) in
parent_filters.iter().zip(child.parent_filters.iter_mut()) {
+ if matches!(pushed.discriminant, PushedDown::Yes) {
+ continue;
+ }
+ if let Some(transferred) = transfer_filter_across_keys(filter,
key_map)? {
+ *pushed = PushedDownPredicate::supported(transferred);
+ }
+ }
+ Ok(())
+}
+
+/// Rewrites `filter` over the other side's join keys, or returns `None` when
+/// it references a column that is not a transferable key, or no column.
+///
+/// A [`DynamicFilterPhysicalExpr`] comes out as a view sharing the original's
+/// state with its key columns remapped, so it keeps tracking the build side.
+fn transfer_filter_across_keys(
Review Comment:
Minor. `HashJoinExec::try_new` does not check that the two key expressions
have the same data type. With a `Utf8` left key and an `Int32` right key, this
function pushes `j@0 = x`, which compares an `Int32` column with a `Utf8`
literal. On `main`, nothing was pushed there.
The planner always coerces the keys, so this needs a hand-built plan. A
`debug_assert` on `data_type` equality makes the assumption explicit.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
