Allow for scaling RangePartitioning
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Is your feature request related to a problem or challenge?
[RangePartitioning](https://github.com/apache/datafusion/blob/cb8faecc78cef2c8bf6f8c04368d6615670e4cdf/datafusion/physical-expr/src/partitioning.rs#L204-L251) currently requires an exact list of split points at construction time, with `partition_count()` strictly derived as `split_points.len() + 1`:
```rust
pub struct RangePartitioning {
ordering: LexOrdering,
split_points: Vec,
}
```
Because `split_points` cannot be adjusted or down-sampled at planning time, `Partitioning::Range` cannot be scaled when physical plans are adapted or distributed. This creates friction in several areas:
1. Stage Partition Scaling (`datafusion-distributed`):
- In multi-stage distributed execution, network shuffle boundaries need to scale partition counts across stages (e.g. scaling from N producer partitions to M consumer tasks).
- Because `RangePartitioning` cannot be resized without out-of-band information, scaling is forced to fall back to `UnknownPartitioning(M)`, destroying range co-partitioning guarantees for downstream joins and merges.
2. Storage Segment Partitioning and Join Pushdown (ParadeDB via `datafusion-distributed`):
- Storage layers often maintain physically over-partitioned segment or files (N boundaries, where N >> worker count).
- At query planning time, an optimizer rule seeks to dynamically choose K partition boundaries (e.g. matching worker count) to co-partition join inputs and push the join down into single workers without a shuffle.
- Because `RangePartitioning` cannot hold a larger sample set and down-sample to K partitions, engines have to maintain shadow data structures outside DataFusion's type system before materializing `RangePartitioning`.
3. Adaptive Query Execution and Cut Discovery (Ballista):
- In distributed range shuffles ([OrderedRangeRepartitionExec](https://github.com/apache/datafusion-ballista/blob/96991336284126a2114f3f63f7f984240ff9e0e6/ballista/core/src/execution_plans/ordered_range_repartition.rs#L191-L197)), quantile sketches ([discover_cuts](https://github.com/apache/datafusion-ballista/blob/96991336284126a2114f3f63f7f984240ff9e0e6/ballista/core/src/execution_plans/range_repartition_common.rs#L66-L110)) produce sample points, and stage adaptation often needs to coalesce or adjust target partition counts. Without sample down-sampling, plans drop to `UnknownPartitioning(K)`.
### Describe the solution you'd like
Extend `RangePartitioning` to hold sorted sample points wrapped in an `Arc` alongside an explicit `partition_count`:
```rust
pub struct RangePartitioning {
ordering: LexOrdering,
samples: Arc<[SplitPoint]>,
partition_count: usize,
}
```
When `partition_count == samples.len() + 1`, `samples` acts as the exact split points (matching existing behavior). When `partition_count < samples.len() + 1`, `RangePartitioning` evenly down-samples `samples` to produce `partition_count - 1` effective split points.
Key capabilities:
1. Backwards Compatible Construction:
- `RangePartitioning::try_new(ordering, split_points)` initializes `samples = Arc::from(split_points)` and `partition_count = samples.len() + 1`.
- Existing callers and static partition declarations continue working without changes.
2. Sample-Backed Construction:
- `RangePartitioning::try_new_with_samples(ordering, samples, partition_count)` allows providing an over-sampled distribution with a target partition count where `partition_count <= samples.len() + 1`.
3. Fallible Scaling & Max Partition Resolution:
- `pub fn max_partition_count(&self) -> usize { self.samples.len() + 1 }`
- `pub fn scale(&self, target_partitions: usize) -> Result` adjusts `partition_count` by down-sampling `samples`.
- Scaling up past `max_partition_count()` returns an error rather than fabricating arbitrary partition values or creating empty partitions with duplicate split points.
4. Derived Split Points & Comparability:
- `pub fn split_points(&self) -> Vec` returns the effective split points derived by down-sampling `samples` to `partition_count - 1` boundaries.
- `RangePartitioning` with identical derived `SplitPoint`s remains compatible for co-partitioning.
### Describe alternatives you've considered
1. Infallible scale-up via interpolation or padding:
- Synthesizing new split points between discrete sample values is not possible without domain knowledge of the underlying data type and distribution.
- Therefore, `scale` should be fallible and capped at `samples.len() + 1`.
2. Maintaining shadow sample structs in downstream engines:
- Engines can maintain external data structures (e.g. [`RangePartitioningSample`](https://github.com/paradedb/paradedb/blob/0cc4b3cd1b329346523724d359ffde565388a9c2/pg_search/src/scan/range_partitioning.rs#L34-L226)), but that prevents interoperability between producers and consumers of `Partitioning`.
### Additional context
https://github.com/apache/datafusion/issues/22395#issuecomment-5414199534
Contributor guide
Assessment
This issue has not been assessed yet.