gene-bordegaray commented on code in PR #23854:
URL: https://github.com/apache/datafusion/pull/23854#discussion_r3718957052
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -891,16 +891,20 @@ impl HashJoinExec {
// https://github.com/apache/datafusion/issues/20195
if config.optimizer.preserve_file_partitions > 0
&& self.mode == PartitionMode::Partitioned
+ && !matches!(
Review Comment:
I think this might read easier if we chcked if it is `Partitioning::Hash`
instead, what do you think?
Could we also update the comment explaining what we allow and what we do not
as of now and the reasons why (partition / filter alignment)
##########
datafusion/core/tests/physical_optimizer/filter_pushdown.rs:
##########
@@ -1191,6 +1192,256 @@ async fn
test_hashjoin_dynamic_filter_pushdown_partitioned() {
);
}
+#[tokio::test]
+async fn test_hashjoin_dynamic_filter_pushdown_range_partitioned() {
+ use datafusion_common::JoinType;
+ use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode};
+
+ // Rough sketch of the Range-partitioned MRE we're trying to recreate. The
+ // test hand-wires identical Range repartitioning:
+ //
+ // EXPLAIN
+ // SELECT *
+ // FROM build
+ // JOIN probe
+ // ON build.a = probe.a AND build.b = probe.b;
+ //
+ //
+---------------+------------------------------------------------------------+
+ // | plan_type | plan
|
+ //
+---------------+------------------------------------------------------------+
+ // | physical_plan | ┌───────────────────────────┐
|
+ // | | │ HashJoinExec │
|
+ // | | │ -------------------- ├──────────────┐
|
+ // | | │ on: (a = a), (b = b) │ │
|
+ // | | └─────────────┬─────────────┘ │
|
+ // | |
┌─────────────┴─────────────┐┌─────────────┴─────────────┐ |
+ // | | │ RepartitionExec ││ RepartitionExec
│ |
+ // | | │ -------------------- ││
-------------------- │ |
+ // | | │ partition_count(in->out): ││
partition_count(in->out): │ |
+ // | | │ 1 -> 2 ││ 1 -> 2
│ |
+ // | | │ ││
│ |
+ // | | │ partitioning_scheme: ││
partitioning_scheme: │ |
+ // | | │ Range([a ASC, b ASC], 2) ││ Range([a ASC, b ASC],
2) │ |
+ // | | │ split: (aa, bb) ││ split: (aa, bb)
│ |
+ // | |
└─────────────┬─────────────┘└─────────────┬─────────────┘ |
+ // | |
┌─────────────┴─────────────┐┌─────────────┴─────────────┐ |
+ // | | │ DataSourceExec (build) ││ DataSourceExec (probe)
│ |
+ // | | │ -------------------- ││
-------------------- │ |
+ // | | │ rows: (aa,ba), (ab,bb) ││ rows: (aa,ba) ...
(ad,bd) │ |
+ // | | │ ││ predicate:
DynamicFilter │ |
+ // | | │ ││ range CASE ->
filter_0/1 │ |
+ // | |
└───────────────────────────┘└───────────────────────────┘ |
+ // | |
|
+ //
+---------------+------------------------------------------------------------+
+
+ // Create build side with limited values
+ let build_batches = vec![
Review Comment:
seems like this is boilerplate that is repreated for multiple tests would it
e worth to extract / reduce some of this with a helper?
##########
datafusion/core/tests/physical_optimizer/filter_pushdown.rs:
##########
Review Comment:
similarly this is repeated every test with the same plans checks, could we
reduce this boilerplate without losing visibility as to what the test is doing
##########
datafusion/core/tests/physical_optimizer/filter_pushdown.rs:
##########
@@ -1191,6 +1192,260 @@ async fn
test_hashjoin_dynamic_filter_pushdown_partitioned() {
);
}
+// Not portable to sqllogictest: this test pins `PartitionMode::Partitioned`
+// by hand-wiring matching Range repartitioning on both join sides, which the
+// SQL planner does not currently produce.
+#[tokio::test]
Review Comment:
I see, ya it is using a csv backed table but we could either:
- update it to be parquet backed and reuse them for your testing (I would
prefer this is the code changes are small)
- or could introduce a `ParquetSource` in the PR and add coverage there
For range partitioning I would really prefer there to be coverage for all
features in this file 🙇 .
NOTE: I would still keep the existing tests you have, they are great
##########
datafusion/core/tests/physical_optimizer/filter_pushdown.rs:
##########
@@ -1191,6 +1192,256 @@ async fn
test_hashjoin_dynamic_filter_pushdown_partitioned() {
);
}
+#[tokio::test]
+async fn test_hashjoin_dynamic_filter_pushdown_range_partitioned() {
+ use datafusion_common::JoinType;
+ use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode};
+
+ // Rough sketch of the Range-partitioned MRE we're trying to recreate. The
+ // test hand-wires identical Range repartitioning:
+ //
+ // EXPLAIN
+ // SELECT *
+ // FROM build
+ // JOIN probe
+ // ON build.a = probe.a AND build.b = probe.b;
+ //
+ //
+---------------+------------------------------------------------------------+
+ // | plan_type | plan
|
+ //
+---------------+------------------------------------------------------------+
+ // | physical_plan | ┌───────────────────────────┐
|
+ // | | │ HashJoinExec │
|
+ // | | │ -------------------- ├──────────────┐
|
+ // | | │ on: (a = a), (b = b) │ │
|
+ // | | └─────────────┬─────────────┘ │
|
+ // | |
┌─────────────┴─────────────┐┌─────────────┴─────────────┐ |
+ // | | │ RepartitionExec ││ RepartitionExec
│ |
+ // | | │ -------------------- ││
-------------------- │ |
+ // | | │ partition_count(in->out): ││
partition_count(in->out): │ |
+ // | | │ 1 -> 2 ││ 1 -> 2
│ |
+ // | | │ ││
│ |
+ // | | │ partitioning_scheme: ││
partitioning_scheme: │ |
+ // | | │ Range([a ASC, b ASC], 2) ││ Range([a ASC, b ASC],
2) │ |
+ // | | │ split: (aa, bb) ││ split: (aa, bb)
│ |
+ // | |
└─────────────┬─────────────┘└─────────────┬─────────────┘ |
+ // | |
┌─────────────┴─────────────┐┌─────────────┴─────────────┐ |
+ // | | │ DataSourceExec (build) ││ DataSourceExec (probe)
│ |
+ // | | │ -------------------- ││
-------------------- │ |
+ // | | │ rows: (aa,ba), (ab,bb) ││ rows: (aa,ba) ...
(ad,bd) │ |
+ // | | │ ││ predicate:
DynamicFilter │ |
+ // | | │ ││ range CASE ->
filter_0/1 │ |
+ // | |
└───────────────────────────┘└───────────────────────────┘ |
+ // | |
|
+ //
+---------------+------------------------------------------------------------+
+
+ // Create build side with limited values
+ let build_batches = vec![
+ record_batch!(
+ ("a", Utf8, ["aa", "ab"]),
+ ("b", Utf8, ["ba", "bb"]),
+ ("c", Float64, [1.0, 2.0]) // Extra column not used in join
+ )
+ .unwrap(),
+ ];
+ let build_side_schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Utf8, false),
+ Field::new("b", DataType::Utf8, false),
+ Field::new("c", DataType::Float64, false),
+ ]));
+ let build_scan = TestScanBuilder::new(Arc::clone(&build_side_schema))
+ .with_support(true)
+ .with_batches(build_batches)
+ .build();
+
+ // Create probe side with more values
+ let probe_batches = vec![
+ record_batch!(
+ ("a", Utf8, ["aa", "ab", "ac", "ad"]),
+ ("b", Utf8, ["ba", "bb", "bc", "bd"]),
+ ("e", Float64, [1.0, 2.0, 3.0, 4.0]) // Extra column not used in
join
+ )
+ .unwrap(),
+ ];
+ let probe_side_schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Utf8, false),
+ Field::new("b", DataType::Utf8, false),
+ Field::new("e", DataType::Float64, false),
+ ]));
+ let probe_scan = TestScanBuilder::new(Arc::clone(&probe_side_schema))
+ .with_support(true)
+ .with_batches(probe_batches)
+ .build();
+
+ let split_points = vec![SplitPoint::new(vec![
+ ScalarValue::Utf8(Some("aa".to_string())),
+ ScalarValue::Utf8(Some("bb".to_string())),
+ ])];
+
+ // Build side: DataSource -> RepartitionExec (Range)
+ let build_range_ordering = LexOrdering::new(vec![
+ PhysicalSortExpr::new(
+ col("a", &build_side_schema).unwrap(),
+ SortOptions::default(),
+ ),
+ PhysicalSortExpr::new(
+ col("b", &build_side_schema).unwrap(),
+ SortOptions::default(),
+ ),
+ ])
+ .unwrap();
+ let build_repartition = Arc::new(
+ RepartitionExec::try_new(
+ build_scan,
+ Partitioning::Range(
+ RangePartitioning::try_new(build_range_ordering,
split_points.clone())
+ .unwrap(),
+ ),
+ )
+ .unwrap(),
+ );
+
+ // Probe side: DataSource -> RepartitionExec (Range)
+ let probe_range_ordering = LexOrdering::new(vec![
+ PhysicalSortExpr::new(
+ col("a", &probe_side_schema).unwrap(),
+ SortOptions::default(),
+ ),
+ PhysicalSortExpr::new(
+ col("b", &probe_side_schema).unwrap(),
+ SortOptions::default(),
+ ),
+ ])
+ .unwrap();
+ let probe_repartition = Arc::new(
+ RepartitionExec::try_new(
+ Arc::clone(&probe_scan),
+ Partitioning::Range(
+ RangePartitioning::try_new(probe_range_ordering,
split_points).unwrap(),
+ ),
+ )
+ .unwrap(),
+ );
+
+ // Create HashJoinExec with partitioned inputs
+ let on = vec![
+ (
+ col("a", &build_side_schema).unwrap(),
+ col("a", &probe_side_schema).unwrap(),
+ ),
+ (
+ col("b", &build_side_schema).unwrap(),
+ col("b", &probe_side_schema).unwrap(),
+ ),
+ ];
+ let hash_join = Arc::new(
+ HashJoinExec::try_new(
+ build_repartition,
+ probe_repartition,
+ on,
+ None,
+ &JoinType::Inner,
+ None,
+ PartitionMode::Partitioned,
+ datafusion_common::NullEquality::NullEqualsNothing,
+ false,
+ )
+ .unwrap(),
+ );
+
+ // Top-level CoalescePartitionsExec
+ let cp = Arc::new(CoalescePartitionsExec::new(hash_join)) as Arc<dyn
ExecutionPlan>;
+ // Add a sort for deterministic output
+ let plan = Arc::new(SortExec::new(
+ LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("a", &probe_side_schema).unwrap(),
+ SortOptions::new(true, false), // descending, nulls_first
+ )])
+ .unwrap(),
+ cp,
+ )) as Arc<dyn ExecutionPlan>;
+
+ // expect the predicate to be pushed down into the probe side DataSource
+ insta::assert_snapshot!(
+ OptimizationTest::new(Arc::clone(&plan),
FilterPushdown::new_post_optimization(), true),
+ @r"
+ OptimizationTest:
+ input:
+ - SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
+ - CoalescePartitionsExec
+ - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0),
(b@1, b@1)]
+ - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa,
bb)], 2), input_partitions=1
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[a, b, c], file_type=test, pushdown_supported=true
+ - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa,
bb)], 2), input_partitions=1
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[a, b, e], file_type=test, pushdown_supported=true
+ output:
+ Ok:
+ - SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
+ - CoalescePartitionsExec
+ - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0,
a@0), (b@1, b@1)]
+ - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC],
[(aa, bb)], 2), input_partitions=1
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[a, b, c], file_type=test, pushdown_supported=true
+ - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC],
[(aa, bb)], 2), input_partitions=1
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[a, b, e], file_type=test, pushdown_supported=true,
predicate=DynamicFilter [ empty ]
+ "
+ );
+
+ // Actually apply the optimization to the plan and execute to see the
filter in action
+ let mut config = ConfigOptions::default();
+ config.execution.parquet.pushdown_filters = true;
+ config.optimizer.enable_dynamic_filter_pushdown = true;
+ config.optimizer.preserve_file_partitions = 1;
+ let plan = FilterPushdown::new_post_optimization()
+ .optimize(plan, &config)
+ .unwrap();
+
+ let config = SessionConfig::from(config).with_batch_size(10);
+ let session_ctx = SessionContext::new_with_config(config);
+ session_ctx.register_object_store(
+ ObjectStoreUrl::parse("test://").unwrap().as_ref(),
+ Arc::new(InMemory::new()),
+ );
+ let state = session_ctx.state();
+ let task_ctx = state.task_ctx();
+ let batches = collect(Arc::clone(&plan), Arc::clone(&task_ctx))
+ .await
+ .unwrap();
+
+ // Now check what our filter looks like
+ insta::assert_snapshot!(
+ format!("{}", format_plan_for_test(&plan)),
+ @r"
+ - SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]
+ - CoalescePartitionsExec
+ - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0),
(b@1, b@1)]
+ - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa,
bb)], 2), input_partitions=1
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[a, b, c], file_type=test, pushdown_supported=true
+ - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa,
bb)], 2), input_partitions=1
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[a, b, e], file_type=test, pushdown_supported=true,
predicate=DynamicFilter [ CASE WHEN a@0 IS NULL OR a@0 < aa OR a@0 = aa AND
(b@1 IS NULL OR b@1 < bb) THEN a@0 >= aa AND a@0 <= aa AND b@1 >= ba AND b@1 <=
ba AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}]) ELSE a@0 >= ab AND a@0 <= ab
AND b@1 >= bb AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:ab,c1:bb}]) END
]
Review Comment:
🎊
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -656,38 +661,99 @@ impl SharedBuildAccumulator {
}
}
- let filter_expr = if has_canceled_unknown {
- let mut when_then_branches = empty_partition_ids
- .into_iter()
- .map(|partition_id| {
- (
- lit(ScalarValue::UInt64(Some(partition_id as
u64))),
- lit(false),
- )
- })
- .collect::<Vec<_>>();
- when_then_branches.extend(real_branches);
-
- if when_then_branches.is_empty() {
- lit(true)
- } else {
- Arc::new(CaseExpr::try_new(
- Some(modulo_expr),
- when_then_branches,
- Some(lit(true)),
- )?) as Arc<dyn PhysicalExpr>
- }
- } else if real_branches.is_empty() {
+ let filter_expr = if has_canceled_unknown
+ && real_partition_ids.is_empty()
+ && empty_partition_ids.is_empty()
+ {
+ lit(true)
+ } else if !has_canceled_unknown &&
real_partition_ids.is_empty() {
lit(false)
- } else if real_branches.len() == 1
+ } else if !has_canceled_unknown
+ && real_partition_ids.len() == 1
&& empty_partition_ids.len() + 1 == num_partitions
{
- Arc::clone(&real_branches[0].1)
+ Arc::clone(&partition_filters[real_partition_ids[0]])
+ } else if let Some(range_partitioning) =
&self.probe_range_partitioning {
+ // Range partitioning
+ assert_eq!(
+ partition_filters.len(),
+ range_partitioning.partition_count()
+ );
+ assert_eq!(self.on_right.len(),
range_partitioning.ordering().len());
+ let sort_exprs = self
+ .on_right
+ .iter()
+ .zip(range_partitioning.ordering())
+ .map(|(expr, sort_expr)| {
+ PhysicalSortExpr::new(Arc::clone(expr),
sort_expr.options)
+ })
+ .collect::<Vec<_>>();
+ let else_expr = partition_filters
+ .pop()
+ .expect("Range partitioning always has at least one
partition");
+ // CASE evaluates in order
+ //
+ // CASE
+ // WHEN key <range split[0] THEN F0
+ // WHEN key <range split[1] THEN F1
+ // ...
+ // ELSE Fn
+ // END
+ let when_then_expr = range_partitioning
+ .split_points()
+ .iter()
+ .zip(partition_filters)
+ .map(|(split_point, then_expr)| {
+ let when_expr = build_lexicographic_filter(
Review Comment:
we need to be careful here. I don't think that this is always producing the
same ordering as `RangePartitioning`. Codex caught a nasty edge case.
The repartitioned for range uses `compare_rows` which uses
`ScalarValue::try_cmp`, `build_lexicographic_filter` uses a binary expression
comparisions which do not have 1-1 semantics.
An edge case is comparing floating point 0.0 values. If a split point was
0.0 then the bounds of range partitioning would be:
```text
partition 0: key < +0.0
partition 1: +0.0 <= key
```
`ScalarValue::try_cmp` will evaluate -0.0 to route to partition 0 since it
thinks of -0.0 < +0.0 as true. The binary expression comparision will normalize
these values to 0.0 and have -0.0 < +0.0 as false, thus would evaluate against
the wrong filter.
Can we create a 1-1 mirroring of the compare_rows sematntics and add a
regression test here.
Here is a pointer to the compare_rows I am talking about:
https://github.com/peterxcli/datafusion/blob/18f9db399331ef1c2163c8d5a230b7169d073e50/datafusion/common/src/utils/mod.rs#L120
and here is where it is used for range repartitioner:
https://github.com/peterxcli/datafusion/blob/18f9db399331ef1c2163c8d5a230b7169d073e50/datafusion/physical-plan/src/repartition/mod.rs#L962
##########
datafusion/core/tests/physical_optimizer/filter_pushdown.rs:
##########
@@ -1191,6 +1192,256 @@ async fn
test_hashjoin_dynamic_filter_pushdown_partitioned() {
);
}
+#[tokio::test]
+async fn test_hashjoin_dynamic_filter_pushdown_range_partitioned() {
+ use datafusion_common::JoinType;
Review Comment:
nit: could we keep these at the module level?
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -7064,8 +7075,57 @@ mod tests {
.enable_join_dynamic_filter_pushdown = true;
let join = HashJoinExec::try_new(
- left,
+ Arc::clone(&left) as Arc<dyn ExecutionPlan>,
right,
+ on.clone(),
+ None,
+ &JoinType::Inner,
+ None,
+ PartitionMode::Partitioned,
+ NullEquality::NullEqualsNothing,
+ false,
+ )?;
+
+
assert!(join.allow_join_dynamic_filter_pushdown(session_config.options()));
Review Comment:
I would prefer to split this test into two separate postive and negative
test cases so it is easier to understand the failure if it pops up 👍
--
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]
