Support co-partitioned range inner equi joins#23184
Conversation
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
|
cc: @gabotechs @stuhood |
8583303 to
45d598b
Compare
|
Amazing timing. Will get you some feedback on this by early next week! Thank you! |
@stuhood great, thank you! I also have a huge line of follow up issues to support more join types that we can split up 👍 I am trying to make smaller tickets as more of the plumbing gets in to get more people involved |
45d598b to
2fca2bb
Compare
|
I have added a stacked PR that shows how I imagine Aggregations will consume the |
gabotechs
left a comment
There was a problem hiding this comment.
This is looking good! left a first round of comments.
The main one is about trying to collapse Distribution::KeyPartitioned and Distribution::HashPartitioned into just a single Distribution variant, I do see that in pretty much all places the code that handles both is the same, and I get the feeling that Distribution::HashPartitioned is just an unfortunate name that was choosen a while back, but it's not really coupled to hashes.
|
I think Key idea
Hash partitioning and range partitioning are two concrete ways to satisfy this requirement. More precisely: Equivalently, different output partitions contain disjoint sets of key values. This seems to cover the logical requirement that Example 1: aggregation has one input
That means all rows with the same For example, both of the following satisfy this requirement:
In either case, we do not need to insert an extra hash repartition before the aggregation. Example 2: hash join has two inputsFor a partitioned hash join, it is not enough for each side to be independently The two sides must also be co-partitioned with respect to the join keys. That means equal join-key values from the left and right inputs must be assigned to the same partition id. The physical execution pattern is: So the required co-partitioning rule is: This implies some extra compatibility requirements across the two inputs:
So I think there are two related concepts: Implementation planHere is my initial thoughts on a implementation plan (note I haven't read the related code carefully, so it's just some rough ideas)
|
|
@2010YOUY01 thank you for the comments and suggestions. I do think that I also believe my current implementation of co-partitioning is satisfying the same contract as described here. Since I would be ok with supporting aggregations before multi-input operators. Something I would like to point out though, is if we are going to collapse If we take this approach we could add explicit checks of what operator we are checking distribuion for in After thinking about it for a bit I would prefer adding private helpers for each operator and then have general satisfaction for PS: also see my linked aggregation pR in my fork 👍 it will give a good idea of what this looks like using this brnach as a base: gene-bordegaray#6 cc: @gabotechs |
Thank you for the extra context! This approach makes sense if it's easier to implement. I think we should start with aggregation first, since single-input cases are simpler. Two-input/co-partitioned operators are trickier. For example, when a hash join requires its inputs to be [HashPartitioned(left_key), HashPartitioned(right_key)], it's really enforcing the co-partitioning property we're discussing here. We'll need to figure out how to incrementally migrate that to the new API to avoid duplicated implementations in the long term. |
Making this co-partition requirement part of the If we can instead make this an explicit requirement at the // Before
fn required_input_distribution(&self) -> Vec<Distribution> {
vec![Distribution::UnspecifiedDistribution; self.children().len()]
}// After
pub struct RequiredInputDistributions {
pub per_child: Vec<Distribution>,
pub cross_child: Vec<CrossChildDistribution>,
}
trait ExecutionPlan {
fn required_input_distributions(&self) -> RequiredInputDistributions {
...
}
}The tradeoff is that this would require a large refactor upfront. The should not be a hard requirement for now, and I'm also unsure how to carry this out incrementally, but I would still love to see us move towards this direction sooner. |
This is an intersting idea, thank you. Let me cherry pick my aggregation commit to be stacked on main and can open something up for a unary operator first. I think something like this would be worth exploring 👍 |
|
Also, I would prefer to support aggregations and joins (and maybe a few more operators) before deciding to rename or replace For something like the aggregations, I can just introduce a private helper in |
There was a problem hiding this comment.
Left a last batch of non blocking comments.
This looks really good @gene-bordegaray, nice work!
Before merging, it would be nice to get the blessing from @2010YOUY01, as this seems quite an important change that could use another committer's +1.
| } else { | ||
| let streaming_benefit = if child.data { | ||
| preserving_order_enables_streaming(&plan, &child.plan)? | ||
| enforce_distribution_relationships( |
There was a problem hiding this comment.
TBH, it's not super pretty to have this function call here in the middle, but I'd not now how to make it better without a major refactor of this file...
My impressions is that this file is worth refactoring. In the current state, it's super difficult to understand what's happening.
This PR is not making it worst though, so one thing we could do is:
- Make up for it with some good comments at the top of this function call and probably below it.
- File an issue for refactoring this file, I might take this one myself.
48e163f to
29ab3a2
Compare
|
I have addressed the last round @gabotechs, thank you for the review. There are some PRs stacked on this and I have some issues lined up that depend on this work being in. We can wait to see if we can get another +1 👍 |
|
It'd be nice to also get a +1 from @2010YOUY01, as he's made some very useful suggestions that have helped shape this PR. |
2010YOUY01
left a comment
There was a problem hiding this comment.
LGTM, Thank you!
Alternative Design
I’d like to share a design idea I have in mind. If we were implementing this from scratch, I suspect this alternative might be simpler. I understand this is not feasible in a single PR, and the current approach is practical and quite good already. So this idea is only to provide additional perspectives.
The current design seems to treat per-child distribution requirements and cross-child requirements as separate mechanisms, then handles them individually.
An alternative approach:
- Treat co-partitioning as the top-level requirement to satisfy.
- Make
enforce_distributiondo one thing: given the current source physical layout, config values, and statistics, decide the cheapest operator(s) to insert.
I think this maps more directly to the essence of the problem, so the implementation might become simpler. Separating these concepts may require keeping implicit consistency across different places, and the mechanisms can sometimes cancel each other out.
enum DistributionRequirement {
// Single-child case
Distribution(Distribution),
// Multi-child co-partitioning case
CoPartition(CoPartition),
}
For example:
-- t1 and t2 are both source tables backed by Parquet
SELECT *
FROM t1
INNER JOIN t2
ON t1.v1 = t2.v1;
Suppose the physical layout is:
# t1: 4 similarly sized range partitions
t1:
p1.parquet # v1 range [0, 100)
p2.parquet # v1 range [100, 200)
p3.parquet # v1 range [200, 300)
p4.parquet # v1 range [300, 400)
# t2: one partition, uniform distribution over [0, 400)
t2:
p1.parquet # v1 range [0, 400)
And the configuration is:
target_partitions = 2
In this model, enforce_distribution would look at the physical layout, configuration, and statistics, and could decide to:
- Optimize
t1’s scan operator into this shape:- partition 1:
[p1.parquet, p2.parquet] - partition 2:
[p3.parquet, p4.parquet]
- partition 1:
- add
RepartitionExec(range, split_points = [200])above thet2scan operator
If we first satisfy the per-child requirements independently, it seems harder to reach this optimal plan without adding many special cases, and we may end up with an extra hash repartition on v1.
Related Issue
One related issue is that we currently enforce sorting and distribution through separate mechanisms. It looks like they can cancel each other out, which makes the implementation harder to maintain.
The essence of the issue is the same: given the input plan and required plan properties, find the cheapest plan change that satisfies the requirement. Unifying these mechanisms would likely simplify the overall implementation too, but that looks like an even larger refactor.
| JOIN range_partitioned r ON l.range_key = r.range_key; | ||
| ---- | ||
| physical_plan | ||
| 01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] |
There was a problem hiding this comment.
nit and optional: since we're only checking if extra hash repartitioning is added for most EXPLAIN tests, we could do
HashJoinExec: <slt:ignore>
--DataSourceExec: <slt:ignore>
--DataSourceExec: <slt:ignore>
This way less future maintenance is needed, if unrelated explain plan changes break this test.
There was a problem hiding this comment.
I added ignore to the long file list but kept the partitioning metadata for sanity checks 👍
29ab3a2 to
f2e6d2c
Compare
@2010YOUY01 I see what you are saying here, and I think this is a great long-term model to have in mind. It seems like what you are saying is the current contract is split into two:
This is a good design with what we are given, it would be better to just explicitly say upfront that these children need to be co-partitioned and derive that immediately. With the current design I could see us moving toward this shape cleanly. We could introduce something like you suggested enum InputDistributionRequirement {
PerChild {...},
CoPartitioned {...},
}and then rather than operators asking for a Distribution they can directly ask for per child or co-partitioned requirements. I think this would, again like you suggested, be pretty big PR (or PRs) but I am up to keep looking into this 👍 |
|
It might be a good topic for the DF 55 blog @gene-bordegaray WDYT? |
@comphead yes would be more than happy to talk about! Still much to do 🚀 |
Distribution::KeyPartitioned requirements before general Range satisfaction
#23451
…rtitioned requirements (apache#23416) ## Which issue does this PR close? - Closes apache#23289. Rebased on main now that apache#23184 has merged. Overlaps with apache#23355, which also opts the window execs into range satisfaction as part of its PartitionedTopK work — happy to rebase whichever lands second. ## What changes are included in this PR? In one commit: * Opt `WindowAggExec` and `BoundedWindowAggExec` with partition keys into range satisfaction via `InputDistributionRequirements::allow_range_satisfaction_for_key_partitioning`, mirroring `AggregateExec`. Compatible range-partitioned inputs then satisfy the window key requirement without a hash repartition; subset satisfaction and the hash fallback for incompatible keys come from the existing satisfaction machinery. Windows without `PARTITION BY` keep requiring a single partition. ## Are these changes tested? Yes: new `slt` tests in `range_partitioning.slt` (exact and subset reuse, rehash on incompatible keys, `subset_repartition_threshold` / `preserve_file_partitions` / `target_partitions` behavior, `WindowAggExec` via an unbounded frame, and no-`PARTITION BY`), plus plan-shape tests in `enforce_distribution.rs` and acceptance/rejection tests in `sanity_checker.rs`. ## Are there any user-facing changes? No: only plan changes.
## Which issue does this PR close? - Closes apache#23450. ## Rationale for this change apache#23184 allowed compatible range-partitioned inputs to satisfy partitioned inner joins without hash repartitioning. However, join output still downgraded `Partitioning::Range` to `UnknownPartitioning`, so downstream joins and aggregates could not reuse the preserved range layout and inserted avoidable hash repartitions. ## What changes are included in this PR? - Preserve `Partitioning::Range` in `adjust_right_output_partitioning`. - Shift right-side range-ordering column indexes into the joined schema while retaining split points and sort options. - Add unit coverage for compound range keys and non-default sort options. - Update SQL logic test plans for nested range joins and a range join feeding an aggregate. ## Are these changes tested? Yes. - `./dev/rust_lint.sh` - `cargo clippy --all-targets --all-features -- -D warnings` - `cargo test -p datafusion-physical-plan test_adjust_right_output_partitioning_preserves_range` - `cargo test --profile=ci --test sqllogictests -- range_partitioning.slt` - `ulimit -n 10240 && RUST_BACKTRACE=1 cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption` ## Are there any user-facing changes? Yes. Physical plans can preserve proven range partitioning through joins, allowing compatible downstream joins and aggregates to avoid unnecessary hash repartitioning. There are no public API changes.
Which issue does this PR close?
Rationale for this change
DataFusion can represent source-declared range partitioning, but partitioned hash joins still required hash partitioned inputs. So an inner join on compatible range-partitioned keys would insert unnecessary hash repartitions, even when each left/right partition already covered the same key domain.
This PR adds a partitioning requirement that means "equal key values are co-located" . I was calling this "compatibility" but found we can satisfy the requirement with looser conditions. Other systems call this "co-location" or "co-partitioning" (trino, spark). Which they (and now I am proposing) define as when both sides of a join are already partitioned so matching key values appear in corresponding partitions, so we can join partition pairs directly without repartitioning the sides.
This lets "co-partitioned" range inputs satisfy inner partitioned hash joins. This will also be applicable to other join types and operators but kept the first PR thin to keep scope more reviewable.
What changes are included in this PR?
Adds
Distribution::KeyPartitioned(Vec<Arc<dyn PhysicalExpr>>)as a public distribution requirement.HashPartitioned([a])means rows must be partitioned by hash ona.KeyPartitioned([a])means rows with equalavalues must be co-located, but the partitioning algorithm may be hash, range, or another compatible scheme.Adds
Partitioning::co_partitioned_with(...)to validate that two independently satisfying partitionings also can be paired by partition index.Changes inner partitioned
HashJoinExecrequirements fromHashPartitionedtoKeyPartitioned.HashPartitionedfor now.Updates
EnforceDistributionso co-partitioned range inner joins avoid repartitioning.idoes not represent the same key domain on both sides.KeyPartitionedin this PR.Keeps partitioned dynamic filter pushdown restricted to hash-compatible routing.
Degrades range join output partitioning to
UnknownPartitioning(n)rather than erroring. Adding this behavior would need more tests and careful thought about, I think its safert o just degrade for first PR.Are these changes tested?
Yes.
KeyPartitionedsatisfaction for hash and range partitioning.co_partitioned_withfor compatible and incompatible range/hash partitioning.EnforceDistributionbehavior for:Are there any user-facing changes?
Yes.
This PR changes public physical planning APIs:
Distribution::KeyPartitioned.Partitioning::co_partitioned_with.Distribution.