feat: allow scaling RangePartitioning - #24766
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 |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24766 +/- ##
==========================================
+ Coverage 81.47% 81.51% +0.03%
==========================================
Files 1122 1123 +1
Lines 404140 405483 +1343
Branches 404140 405483 +1343
==========================================
+ Hits 329284 330516 +1232
- Misses 55546 55618 +72
- Partials 19310 19349 +39 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
jayzhan211
left a comment
There was a problem hiding this comment.
Thanks @goutamadwant, there is a suggestion:
RangePartitioning::PartialEq compares only ordering + effective split_points, but samples is not derivable from split_points and is exactly the state scale() reads. So a == b does not imply a.scale(k) == b.scale(k), and doesn't even imply both succeed — your own test asserts try_new_with_samples(ord, [10..90], 4) == try_new(ord, [30,50,70]), yet the first scales to 10 and the second errors at 4.
There's already a consumer that turns this into wrong results. can_interleave gates on partition == *reference:
// datafusion/physical-plan/src/union.rs:931
matches!(reference, Partitioning::Hash(_, _) | Partitioning::Range(_))
&& inputs
.map(|plan| plan.borrow().output_partitioning().clone())
.all(|partition| partition == *reference)and InterleaveExec::compute_properties then adopts one input's metadata for the whole output:
// datafusion/physical-plan/src/union.rs:698
let output_partitioning = inputs[0].output_partitioning().clone();Concretely: input A = try_new_with_samples(ord, [10,20,...,90], 4), input B = try_new(ord, [30,50,70]). They compare equal, interleaving is allowed (correct — the effective boundaries do match), and the InterleaveExec output now advertises max_partition_count() == 10. A distributed planner — the use case this PR is for — calls scale(10) on that output and gets [10,20,...,90], boundaries B's rows were never placed against. Because range partitioning is a declared, unvalidated property, that's silently wrong rows per partition, not an error.
Suggest making equality structural:
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
pub struct RangePartitioning {
ordering: LexOrdering,
samples: Arc<[SplitPoint]>,
split_points: Arc<[SplitPoint]>,
partition_count: usize,
}
@@
-impl PartialEq for RangePartitioning {
- fn eq(&self, other: &Self) -> bool {
- self.ordering == other.ordering && self.split_points == other.split_points
- }
-}The one place that genuinely wants effective-boundary comparison already spells it out and is unaffected:
// datafusion/physical-plan/src/distribution_requirements.rs:348
(Partitioning::Range(left), Partitioning::Range(right)) => {
left.split_points() == right.split_points()
&& ...
}so co-partitioned joins keep the permissive behavior; only can_interleave gets stricter, which is the conservative direction. test_range_partitioning_equality_uses_effective_split_points would then need to flip to assert_ne!.
If you'd rather keep permissive equality, the alternative is for InterleaveExec::compute_properties to reduce the output to the coarsest sample set common to all inputs rather than inheriting inputs[0]'s — but that's more machinery for the same guarantee.
| self.ordering, | ||
| split_points, | ||
| self.partition_count() | ||
| ) |
There was a problem hiding this comment.
nit: Range([a@0 ASC], [(30), (50), (70)], 4) is identical for a partitioning that can scale to 10 and one that can't scale past 4. Since EXPLAIN output is the primary debugging surface for partitioning bugs, appending the max when it differs (e.g. , max 10) would pay for itself.
Which issue does this PR close?
Rationale for this change
Distributed and adaptive planners need to lower a range-partitioned plan's task count without discarding its higher-resolution boundary sample. Today scaling loses the range layout and falls back to unknown partitioning.
What changes are included in this PR?
split_points()API.Are these changes tested?
Yes. Tests cover deterministic scaling and invalid counts, equality by effective boundaries, projection and plan rewrites, legacy/current/malformed protobuf payloads, and FFI unit and cross-library round trips. The full required workspace test suite and workspace-wide clippy with all targets and features passed. After refreshing the unchanged patch to current
main, the focusedRangePartitioningtests passed again.Are there any user-facing changes?
Yes.
RangePartitioninggainstry_new_with_samples,samples,max_partition_count, andscale. Existing exact construction andsplit_points()callers remain source-compatible.FFI_RangePartitioningnow carries samples and the current partition count, so this PR requires theapi changelabel.