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 |
| /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) | ||
| /// right: Range(right.x ASC, [15, 20]), required KeyPartitioned(right.x) | ||
| /// ``` | ||
| pub fn co_partitioned_with( |
There was a problem hiding this comment.
I realized the concept was a bit off. I think this is ok if we haven't had a release... 😅
I linked the mirroring concepts that are used in trino in spark in the PR description
Lesson learned to have consumer of the public API first
There was a problem hiding this comment.
👍 I don't think it's a big deal, this is anyways not the typical method external consumers rely one.
There was a problem hiding this comment.
Without taking a look at the usages of this method, I do notice that the signature seems a bit off: why would it need the required Distribution for checking if two Partitionings are co-partitioned? it seems unrelated.
I'll keep reading though
There was a problem hiding this comment.
Ya I thought similarly but this matters for multi-children operators. For joins, each side has a different may have a diff schema:
left: Partitioning::Range(left.a)
right: Partitioning::Range(right.x)
join: left.a = right.x
So the expressions left.a and right.x are not equal since they are from different children. Because of this we need to check each sides req Distribution to know what each partitioning has to satisfy:
left requirement: key partitioned on left.a
right requirement: key partitioned on right.x
So we need to check:
- left partitioning satisfies left requirement
- right partitioning satisfies right requirement
- the two partition mappings can be paired
I could do something like this:
struct RequiredPartitioning<'a> {
partitioning: &'a Partitioning,
requirement: &'a Distribution,
eq_properties: &'a EquivalenceProperties,
}
impl RequiredPartitioning<'_> {
fn co_partitioned_with(&self, other:
&RequiredPartitioning<'_>) -> bool
}or derive the EquivalenceProperties in the method if we pass the children directly
|
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.
| /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) | ||
| /// right: Range(right.x ASC, [15, 20]), required KeyPartitioned(right.x) | ||
| /// ``` | ||
| pub fn co_partitioned_with( |
There was a problem hiding this comment.
👍 I don't think it's a big deal, this is anyways not the typical method external consumers rely one.
| /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) | ||
| /// right: Range(right.x ASC, [15, 20]), required KeyPartitioned(right.x) | ||
| /// ``` | ||
| pub fn co_partitioned_with( |
There was a problem hiding this comment.
Without taking a look at the usages of this method, I do notice that the signature seems a bit off: why would it need the required Distribution for checking if two Partitionings are co-partitioned? it seems unrelated.
I'll keep reading though
|
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 |
stuhood
left a comment
There was a problem hiding this comment.
Thanks a lot!
I was able to incorporate this branch into datafusion-distributed and paradedb, and observe partitioned hash joins both with and without df-d. With df-d, stage boundaries were placed appropriately above the join.
I agree with collapsing Key and Hash in #23236.
| vec![ | ||
| Distribution::HashPartitioned(left_expr), | ||
| Distribution::HashPartitioned(right_expr), | ||
| ] |
There was a problem hiding this comment.
Worth a TODO pointing at followup?
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes apache#23191. - Related discussion: apache#23184, apache#23236. ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> Range partitioning can satisfy aggregate hash partitioning: equal group keys are already partitioned, even though the partitioning is not hash-based. This is the first unary-operator implementation from the range partitioning discussion before making broader public API changes around `HashPartitioned` / `KeyPartitioned`. ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> - Let compatible range partitioning satisfy aggregate hash distribution requirements in `EnforceDistribution` - Keep this private to aggregate planning for now to not make public API changes to `Distribution` enum variants yet until more operators are supported ## Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> Yes. ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> Yes. Range-partitioned aggregate plans can now avoid hash repartitioning. <!-- If there are any breaking changes to public APIs, please add the `api change` label. -->
2642d4d to
f20e0bf
Compare
|
@stuhood @2010YOUY01 @gabotechs I have explored this idea:
I implemented this idea through |
c3f4319 to
2dbbdf7
Compare
stuhood
left a comment
There was a problem hiding this comment.
I implemented this idea through
RequiredInputDistributionsindistribution_requirements.rsI am not done but wanted to flesh out this rough draft just to get feedback at a high level if you guys think this is worth purusing in this PR
I like this! I don't have much to say about the exact shape of struct RequiredInputDistributions, but the ability to expose requirements across multiple children makes a lot of sense.
| /// Designates whether hash repartitioning is required. | ||
| hash_repartition_required: bool, |
There was a problem hiding this comment.
Since you're renaming this, why isn't this repartition_required or aligned_repartition_required? Or is the idea that this flag would be set in cases where Hash in particular would be beneficial (to balance out skew, for example), and then you'd eventually have another flag for Range?
There was a problem hiding this comment.
the idea was that we are specifically having logic for hash repartitioning right now in the rule, then if / when we decide to add logic for range repartitioning there will be distinction
2dbbdf7 to
001b10c
Compare
It sounds a bit... complex. I wonder if there's a simpler way of modeling this. I'll try to make some suggestions there. |
gabotechs
left a comment
There was a problem hiding this comment.
Flushing some [not super useful] comments, I'll need to invest some time in understanding what's happening in enforce_distribution before actually making some better suggestions.
| children: &mut [( | ||
| DistributionContext, | ||
| Option<OrderingRequirements>, | ||
| bool, | ||
| Distribution, | ||
| )], | ||
| target_partitions: usize, |
There was a problem hiding this comment.
This type is pretty messed up, I'm surprised that clippy allows this. It would be nice to avoid these wide unnamed tuple arguments, as reading the signature is impossible to figure out what's any of that.
There was a problem hiding this comment.
I can pull it into some type of struct, its really trying to model children's dtribution state
There was a problem hiding this comment.
maybe a good chance for a refactor of this logic
| // There is an ordering requirement of the operator: | ||
| if let Some(required_input_ordering) = required_input_ordering { | ||
| // Either: | ||
| // - Ordering requirement cannot be satisfied by preserving ordering through repartitions, or | ||
| // - using order preserving variant is not desirable. | ||
| let sort_req = required_input_ordering.into_single(); | ||
| let ordering_satisfied = child | ||
| .plan | ||
| .equivalence_properties() | ||
| .ordering_satisfy_requirement(sort_req.clone())?; | ||
|
|
||
| if (!ordering_satisfied || !order_preserving_variants_desirable) | ||
| && !streaming_benefit | ||
| && child.data | ||
| { | ||
| child = replace_order_preserving_variants(child)?; | ||
| // If ordering requirements were satisfied before repartitioning, | ||
| // make sure ordering requirements are still satisfied after. | ||
| if ordering_satisfied { | ||
| // Make sure to satisfy ordering requirement: |
There was a problem hiding this comment.
The amount of if statement branching happening here is most likely going to take a huge hit in maintainability.
This was already a difficult optimizer rule to understand, but I fear now it's worst.
I'll likely need to invest some time in understanding this before I can contribute any meaningful suggestion. In the mean time, if you see opportunities for simplyfing this code, that's more than welcome.
There was a problem hiding this comment.
😂 damn, I'm realizing this is not even new code...
There was a problem hiding this comment.
Hahaha yes I'm trying my hardest to make it more readable jnremently 😭. But please if it's harder now let me know
| /// Specifies data distribution requirements for all children, including | ||
| /// non-default satisfaction policies or cross-child relationships. | ||
| fn required_input_distributions(&self) -> RequiredInputDistributions { | ||
| RequiredInputDistributions::new(self.required_input_distribution()) | ||
| } | ||
|
|
||
| /// Specifies the data distribution requirements for all the |
There was a problem hiding this comment.
Is there any chance we can not maintain these two methods?
required_input_distributionsrequired_input_distribution
Having the two at the same time does not seem like a nice public API design
There was a problem hiding this comment.
We would either deprecate the old one or just flat out change the existing. Didn't know if we wanted the churn
| AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned | ||
| if !self.group_by.has_grouping_set() => | ||
| { | ||
| requirements.with_range_key_partitioning(0) |
There was a problem hiding this comment.
Without looking at the specific implementation, this reads:
"The required input distribution must be range key partitioned"
Which I assume is wrong right?
| match ( | ||
| self.left.output_partitioning(), | ||
| self.right.output_partitioning(), | ||
| ) { | ||
| ( | ||
| Partitioning::Hash(_, left_partition_count), | ||
| Partitioning::Hash(_, right_partition_count), | ||
| ) => left_partition_count == right_partition_count, | ||
| (left_partitioning, right_partitioning) => { |
There was a problem hiding this comment.
I imagine that dynamic filters should work at some point with co-partitioned joins on a range key right? probably another issue worth adding to your batch?
| /// TODO: remove this field once [`Partitioning::Range`] generally | ||
| /// satisfies [`Distribution::KeyPartitioned`] through | ||
| /// [`Partitioning::satisfaction`]. See | ||
| /// <https://github.com/apache/datafusion/issues/23266>. | ||
| per_child_satisfaction: Vec<InputDistributionSatisfaction>, | ||
| /// Cross-child distribution relationships required by the plan. |
There was a problem hiding this comment.
There is a good amount, its really just tradeoff between if we want to test lots of operators upfront in a single PR or have this bridge and do one by one
This is centralizing the opt in so from this point on we just need to add the allow_range_partitioning()
| /// <https://github.com/apache/datafusion/issues/23266>. | ||
| per_child_satisfaction: Vec<InputDistributionSatisfaction>, | ||
| /// Cross-child distribution relationships required by the plan. | ||
| relationships: Vec<InputDistributionRelationship>, |
There was a problem hiding this comment.
Is there any scenario where this Vec would have a length greater than 1?
There was a problem hiding this comment.
unions or intervleave execs
There was a problem hiding this comment.
I don't fully understand how more than 1 would be necessary for unions or interleave execs.
InterleaveExec only works if all the inputs have the same hash partitioning:
https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/union.rs#L457-L457
So I don't think there's any scenario where there'll be more than 1 relationship between a subset of the InterleaveExec children
For UnionExec, how could this operator try to enforce a specific input distribution requirement to its children? unions don't really care what was the partitioning scheme of children, and I cannot think of a situation it'll want to enforce any specific distribution requirement
There was a problem hiding this comment.
wiat I think I stand corrected here, this can be Option
EDIT didnt refresh before I made this comment and saw the one above, yes that is right. This is saying that children might have different relatioships between them but this should not be true in the case of an Interleave, and ya think we should keep this fitting the current needs rather than speculative 👍
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.