feat: allow scaling RangePartitioning - #24766
goutamadwant wants to merge 13 commits into
Conversation
github-actions
Bot
added
physical-expr
|
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 82.38% 82.42% +0.03%
==========================================
Files 1138 1138
Lines 434309 436009 +1700
Branches 434309 436009 +1700
==========================================
+ Hits 357803 359362 +1559
- Misses 54875 54903 +28
- Partials 21631 21744 +113 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
jayzhan211
left a comment
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.
| @@ -296,6 +374,44 @@ impl Display for RangePartitioning { | |||
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.
|
Thanks for working on this! Before we actually land this, there is a little bit more work that I should do to prove that #24712 is actually sufficient to solve the case that I'm examining... could we hold off for a few days, so that I can update the ticket if needed? |
Merge main while preserving satisfying ranges, validating scaled keys, and retaining sample-aware equality.
github-actions
Bot
added
optimizer
@jayzhan211 Switched RangePartitioning to structural equality and added an interleave regression for matching effective boundaries with different retained samples. Both input orders are covered, so interleave cannot inherit unsupported scaling capacity from either input. |
stuhood
left a comment
stuhood
left a comment
There was a problem hiding this comment.
Thanks for working on this!
I've integrated it in datafusion-distributed and ParadeDB, and it seems to be working well.
| /// Returns this range partitioning scaled to `target_partitions`. | ||
| /// | ||
| /// Scaling retains the original samples, so a range partitioning that was | ||
| /// scaled down can later be scaled back up to [`Self::max_partition_count`]. | ||
| pub fn scale(&self, target_partitions: usize) -> Result<Self> { | ||
| validate_range_partition_count(target_partitions, self.max_partition_count())?; | ||
| if target_partitions == self.partition_count { | ||
| return Ok(self.clone()); | ||
| } | ||
| Ok(Self { | ||
| ordering: self.ordering.clone(), | ||
| samples: Arc::clone(&self.samples), | ||
| split_points: downsample_split_points(&self.samples, target_partitions), | ||
| partition_count: target_partitions, | ||
| }) | ||
| } |
There was a problem hiding this comment.
So, the case where we don't have enough samples to scale up to a particular target_partition count should be treated more like an expected case rather than a failure. Ideally we would always have enough samples to do this... but in cases where we don't (small tables, etc), it is definitely a recoverable case, and callers should convert the Partitioning::Range into something else... e.g. Partitioning::Unknown.
So rather than returning Result<Self>, maybe this should return Option<Self>? Alternatively, it could return a Result<Self, ScalingError> (name TBD) that allows the caller to differentiate failure cases.
There was a problem hiding this comment.
@stuhood changed scale to return a typed RangePartitioningScaleError, distinguishing InsufficientSamples from ZeroPartitions. EnsureRequirements retains a satisfying range when capacity is insufficient; other callers can choose their own fallback without parsing error strings.
| /// Creates sample-backed range partitioning and validates the sample shape, | ||
| /// ordering, and target partition count. | ||
| /// | ||
| /// `partition_count` must be at least one and no larger than | ||
| /// `samples.len() + 1`. When it is smaller than that maximum, the samples | ||
| /// are evenly down-sampled to derive the effective split points. | ||
| pub fn try_new_with_samples( | ||
| ordering: LexOrdering, | ||
| samples: Vec<SplitPoint>, | ||
| partition_count: usize, | ||
| ) -> Result<Self> { |
There was a problem hiding this comment.
This looks good. But I think that we should go ahead and deprecate the old constructor, and have all existing callers switch to a method like RangePartitioning::try_new_with_samples. RangePartitioning::new will create something which cannot be scaled, and that's much less useful.
In terms of the number of samples for callers to provide to try_new_with_samples: we might want to recommend something like K * target_partitions, where K is some value that we expect will give us enough flexibility to actually scale things for the relevant number of partitions.
There was a problem hiding this comment.
@stuhood deprecated the unchecked new constructor for 56.0.0 and migrated its internal callers to try_new_with_samples. Validated try_new remains supported and delegates to the sample-backed constructor. Added capacity-sizing guidance and an oversampling example, including the small-input limitation.
There was a problem hiding this comment.
I don't know if deprecating is needed rather than it jsut being an invariant. It seems we can alsways derive the samples form split points.
@stuhood do you think there is a use case for having a contructor that does not return a Result? I could see it if a use case vlaidated themselves and is really trying to squeeze perf
There was a problem hiding this comment.
@stuhood do you think there is a use case for having a contructor that does not return a
Result? I could see it if a use case vlaidated themselves and is really trying to squeeze perf
Hm, maybe... but I can't think of any cases where you'd be creating RangePartitioning instances in a tight loop.
| /// NOTE: Optimizer and execution behavior for this partitioning is intentionally | ||
| /// not implemented and will be introduced incrementally. See | ||
| /// <https://github.com/apache/datafusion/issues/22395>. | ||
| #[derive(Debug, Clone, PartialEq)] |
There was a problem hiding this comment.
The fact that PartialEq is derived here is going to confuse callers: enforce_distribution_relationships uses equality to check whether two RangePartitioning instances are compatible. The samples should not actually be included in that equality check: only the split_points.
There are two schools of thought on that, and I don't know which DataFusion prefers:
- School 1:
PartialEqshould always represent structural equality, and so should basically always be derived. Other definitions of equality should be provided by other methods. - School 2:
PartialEqshould be the most useful definition of equality, even if that is not structural.
I expect that @gene-bordegaray has an opinion on this one.
There was a problem hiding this comment.
sure @stuhood cc @gene-bordegaray kept structural PartialEq for InterleaveExec safety and added has_same_layout to compare ordering and effective boundaries. The optimizer now uses that comparison. A three-input regression fails before the change and passes afterward in both sampled/exact orientations: compatible inputs retain their own samples without another exchange.
| /// Effective boundaries for the current partition count. | ||
| split_points: Arc<[SplitPoint]>, | ||
| /// Number of effective partitions. | ||
| partition_count: usize, |
There was a problem hiding this comment.
The partition_count can be derived from split_points.len() + 1, so it doesn't need to be stored.
There was a problem hiding this comment.
makes sense @stuhood removed the stored partition_count field; it is now derived from split_points.len() + 1. Protobuf and FFI still retain the effective count needed to reconstruct the selected layout from the stored samples.
github-actions
Bot
added
documentation
|
This changes planning policy beyond the feature. When We already have a knob for "keep a satisfying layout instead of increasing parallelism": It can also be strictly worse than main for joins. In Suggest keeping main's behaviour on - Err(RangePartitioningScaleError::InsufficientSamples { .. }) => {
- preserved_unscalable_range = true;
- None
- }
+ // Samples cannot support the target parallelism:
+ // fall back to key repartitioning as before.
+ Err(RangePartitioningScaleError::InsufficientSamples { .. }) => Some(
+ requirement
+ .clone()
+ .create_partitioning(target_partitions),
+ ),and restoring the matrix / TEST 13 expectations. If retaining an unscalable range for the co-partitioning pass is desirable, it should be gated on the existing preserve_file_partitions threshold rather than unconditional, and the reference selection must never let a preserved range beat a candidate that already provides target_partitions. |
Agreed. Thanks! |
gene-bordegaray
left a comment
gene-bordegaray
left a comment
There was a problem hiding this comment.
flushing comments
|
|
||
| /// Why a [`RangePartitioning`] cannot be scaled to a requested partition count. | ||
| #[derive(Debug, Clone, Copy, PartialEq, Eq)] | ||
| pub enum RangePartitioningScaleError { |
There was a problem hiding this comment.
I dont think eposing a public error type is the best path here. Most likely could be a datafusion internal error and have some private enum if needed
| /// - `ordering` defines the partitioning key and ordering. | ||
| /// - `split_points` define the boundaries between adjacent partitions. | ||
| /// - `samples` define the maximum-resolution boundaries. | ||
| /// - `partition_count` selects how many ranges to derive from those samples. |
There was a problem hiding this comment.
isnt split points still in the struct?
| /// - `ordering` defines the partitioning key and ordering. | ||
| /// - `split_points` define the boundaries between adjacent partitions. | ||
| /// - `samples` define the maximum-resolution boundaries. | ||
| /// - `partition_count` selects how many ranges to derive from those samples. |
There was a problem hiding this comment.
I also am not seeing partition_count but may be missing
| /// samples provides capacity for up to `4 * partition_count + 1` partitions. | ||
| /// Choose the sampling factor for the workload; small inputs may not have | ||
| /// enough distinct values, and callers must handle insufficient capacity. | ||
| pub fn try_new_with_samples( |
There was a problem hiding this comment.
I would love to see a concrete example of samles and split points relationship being shown. Like what an actual scaling would look like and the contract it enforces, like not exceeding the max samples.
This also makes me think, samples is not the most telling name for a public facing property. Maybe something like max_partition_bounds?
| /// Creates sample-backed range partitioning and validates the sample shape, | ||
| /// ordering, and target partition count. | ||
| /// | ||
| /// `partition_count` must be at least one and no larger than | ||
| /// `samples.len() + 1`. When it is smaller than that maximum, the samples | ||
| /// are evenly down-sampled to derive the effective split points. | ||
| pub fn try_new_with_samples( | ||
| ordering: LexOrdering, | ||
| samples: Vec<SplitPoint>, | ||
| partition_count: usize, | ||
| ) -> Result<Self> { |
There was a problem hiding this comment.
I don't know if deprecating is needed rather than it jsut being an invariant. It seems we can alsways derive the samples form split points.
@stuhood do you think there is a use case for having a contructor that does not return a Result? I could see it if a use case vlaidated themselves and is really trying to squeeze perf
| # Co-partitioning satisfaction does not prevent a repartition that increases | ||
| # parallelism. With target_partitions larger than the Range partition count, | ||
| # both sides are hash repartitioned. | ||
| # TEST 13: Compatible Range Join Preserves Unscalable Partitioning |
There was a problem hiding this comment.
is this the inteded behavior we want though? What if we would benefit more from a repartition to 5 for more parallelism. We should make this configurable to the user in some way. Maybe a threshold of some type
| } | ||
| Err(RangePartitioningScaleError::InsufficientSamples { .. }) => { | ||
| preserved_unscalable_range = true; | ||
| None |
There was a problem hiding this comment.
this returns None which I have tracked as meaning “keep the current range partitioning” and this happens even when preserve_file_partitions is disabled so I think this shouldnt happen
There was a problem hiding this comment.
so like a 3 partiitoned range merging with a 8 partitioned hash, I would think if the user doesn't opt in we should decide to use more parallelism
There was a problem hiding this comment.
On InsufficientSamples the Range input keeps its partition count even with preserve_file_partitions = 0, and even for single-input requirements (e.g. an aggregate) where nothing needs to line up. Range sources today are built with try_new (max == current count), so all existing Range users lose parallelism. range_satisfaction_config_matrix (NOT_MET, DISABLED, GREATER) flips Hash → Reuse, and sqllogictest TEST 13 now joins at 4 partitions with target_partitions = 5. Please keep the layout only when the inputs are co-partitioned (and confirm on #24712 whether joins should also require preserve_file_partitions):
- Err(RangePartitioningScaleError::InsufficientSamples { .. }) => {
+ Err(RangePartitioningScaleError::InsufficientSamples { .. })
+ if input_distributions.is_co_partitioned() =>
+ {
preserved_unscalable_range = true;
None
}
+ Err(RangePartitioningScaleError::InsufficientSamples { .. }) => Some(
+ requirement
+ .clone()
+ .create_partitioning(target_partitions),
+ ),
Err(error) => return Err(error.into()),There was a problem hiding this comment.
are we intentionally rejecting here?
There was a problem hiding this comment.
shouldn't we attempt to scale
| ---- | ||
| physical_plan | ||
| 01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] | ||
| 02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 |
There was a problem hiding this comment.
I am very confused here, we are showing that we are not repartitioning but we are not doing this due to scaling. This was intentional to show that if we can increase parallelism in the system we should. I think a better test here would be showing increasing parallelism by repartitioning on range and scaling the points correctly
| if should_add_hash_repartition { | ||
| // Enforce unmet requirements, or increase parallelism when beneficial. | ||
| if should_add_repartition { | ||
| let partitioning = match child.plan.output_partitioning() { |
There was a problem hiding this comment.
I think using an option here too is pretty confusing. Could we make it more explicit what the variants are like Keep or Repartition on Range
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
A small suggestion, other LGTM
| } | ||
| Err(RangePartitioningScaleError::InsufficientSamples { .. }) => { | ||
| preserved_unscalable_range = true; | ||
| None |
There was a problem hiding this comment.
On InsufficientSamples the Range input keeps its partition count even with preserve_file_partitions = 0, and even for single-input requirements (e.g. an aggregate) where nothing needs to line up. Range sources today are built with try_new (max == current count), so all existing Range users lose parallelism. range_satisfaction_config_matrix (NOT_MET, DISABLED, GREATER) flips Hash → Reuse, and sqllogictest TEST 13 now joins at 4 partitions with target_partitions = 5. Please keep the layout only when the inputs are co-partitioned (and confirm on #24712 whether joins should also require preserve_file_partitions):
- Err(RangePartitioningScaleError::InsufficientSamples { .. }) => {
+ Err(RangePartitioningScaleError::InsufficientSamples { .. })
+ if input_distributions.is_co_partitioned() =>
+ {
preserved_unscalable_range = true;
None
}
+ Err(RangePartitioningScaleError::InsufficientSamples { .. }) => Some(
+ requirement
+ .clone()
+ .create_partitioning(target_partitions),
+ ),
Err(error) => return Err(error.into()),
|
Thanks again for working on this @goutamadwant! Some more breadcrumbs here: I've integrated this in datafusion-contrib/datafusion-distributed#730 and paradedb/paradedb#6342, and did not find any problems with the API other than those already mentioned here. Let me know if we can help get this landed. |
…g-review # Conflicts: # datafusion/physical-expr/src/partitioning.rs
Signed-off-by: goutamadwant <workwithgoutam@gmail.com>
Signed-off-by: goutamadwant <workwithgoutam@gmail.com>
Signed-off-by: goutamadwant <workwithgoutam@gmail.com>
5 participants
Which issue does this PR close?
Rationale for this change
Distributed and adaptive planners need to change a range-partitioned plan's task count without losing its higher-resolution boundary sample. The optimizer should use a compatible range layout where possible while preserving the requested parallelism when sample capacity is insufficient.
What changes are included in this PR?
scale(usize) -> Option<Self>. Zero and above-capacity requests returnNone.has_same_layoutfor effective-layout comparisons.EnsureRequirements, recheck key requirements after singleton scaling, and fall back to key partitioning at the requested parallelism when scaling is unsupported. Preserve native reference selection and the larger-input preference.RepartitionExec::repartitionedscale supported ranges with fresh runtime state and metrics, retaining ordering and batch-size settings. Remove the optimizer's redundant outerOption.newfor 56.0.0, retain validated exact construction andsplit_points(), and document sample sizing and upgrade requirements.What is the testing strategy for this PR?
Coverage includes count limits, sample retention, key and sort-option compatibility, structural versus effective-layout equality, interleave safety, sample-preserving rewrites, protobuf/FFI round trips, and aggregate/join results with NULLs.
The new exchange regression fails before the follow-up fix and passes afterward. It initializes the original exchange, scales 2 -> 4 -> 1 -> 4, and checks exact row routing, ordering, retained samples, batch configuration, and independent execution state and metrics. The 23 range optimizer integration regressions also pass.
Validation with the pinned Rust 1.98.1 toolchain:
./dev/rust_lint.shsuite pass.Performance caveat: six existing range benchmarks were compared with base
925d7f8ffdusing separate Rust 1.98.1 release-nonlto builds and alternating runs. The eight-partition integer routing case was 13.5% slower; the other five differed by -1.1% to +2.2%. A separate before/after check of that case found this follow-up within 0.3% of pre-follow-up head1d2d59e699, with both about 18% slower than the base in that run. This indicates an existing full-PR performance concern, not a measured slowdown introduced by this follow-up. Its cause remains unverified; these local results do not establish performance parity with main.Are there any user-facing changes?
Yes.
RangePartitioninggains sample-backed construction and scaling APIs.scalereturnsOption<Self>rather than a public error type. Uncheckednewis deprecated; validated exact construction remains supported. Unsupported optimizer scaling falls back to key partitioning rather than retaining lower parallelism.FFI_RangePartitioningchanges layout and the generatedPhysicalRangePartitioningprotobuf struct gains fields. These API/ABI changes require theapi changelabel. This targets main, not a patch release. FFI consumers must rebuild against the compatible release. Existing protobuf payloads remain readable.