Skip to content

Commit b73054a

Browse files
authored
feat: allow scaling RangePartitioning (#24766)
## Which issue does this PR close? - Closes #24712. ## 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? - Add validated sample-backed construction, sample/current/maximum partition-count accessors, and `scale(usize) -> Option<Self>`. Zero and above-capacity requests return `None`. - Retain samples across scaling, projection, repartition, join rewrites, protobuf, and FFI round trips. Keep structural equality for interleave safety and use `has_same_layout` for effective-layout comparisons. - Integrate scaling into `EnsureRequirements`, recheck key requirements after singleton scaling, and fall back to key partitioning at the requested parallelism when scaling is unsupported. During co-partitioning, a source already providing the required layout remains native; a newly inserted range exchange does not gain that preference. - Let `RepartitionExec::repartitioned` scale supported ranges with fresh runtime state and metrics, retaining ordering and batch-size settings. Remove the optimizer's redundant outer `Option`. - Deprecate unchecked `new` for 56.0.0, retain validated exact construction and `split_points()`, and document sample sizing and upgrade requirements. - Preserve effective boundaries in protobuf field 2 for older readers, and validate the added sample/count representation against those boundaries. ## 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. A separate join regression fails before the native-reference fix and passes afterward. It checks both input orders and verifies that a larger, newly scaled range exchange does not displace a source already providing the required hash layout. The end-to-end range join test now checks the resulting hash layout and exact output rows when scaling requires exchanges on both sides. Validation with the pinned Rust 1.98.1 toolchain: - Formatting, all-target/all-feature Clippy with warnings denied, and the complete `./dev/rust_lint.sh` suite pass. - The extended workspace suite passes 12,022 Rust tests, with eight ignored, and all 522 SQL logic-test files. - All 155 all-feature FFI tests pass. Performance caveat: six existing range benchmarks were compared with base `925d7f8ffd` using 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 head `1d2d59e699`, 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. `RangePartitioning` gains sample-backed construction and scaling APIs. `scale` returns `Option<Self>` rather than a public error type. Unchecked `new` is deprecated; validated exact construction remains supported. Unsupported optimizer scaling falls back to key partitioning rather than retaining lower parallelism. `FFI_RangePartitioning` changes layout and the generated `PhysicalRangePartitioning` protobuf struct gains fields. These API/ABI changes require the `api change` label. This targets main, not a patch release. FFI consumers must rebuild against the compatible release. Existing protobuf payloads remain readable. --------- Signed-off-by: goutamadwant <workwithgoutam@gmail.com>
1 parent f50c17d commit b73054a

20 files changed

Lines changed: 1614 additions & 135 deletions

File tree

‎datafusion/core/tests/physical_optimizer/enforce_distribution.rs‎

Lines changed: 497 additions & 0 deletions
Large diffs are not rendered by default.

‎datafusion/ffi/src/execution_plan.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -600,6 +600,12 @@ pub mod tests {
600600
self
601601
}
602602

603+
pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self {
604+
self.props =
605+
Arc::new(self.props.as_ref().clone().with_partitioning(partitioning));
606+
self
607+
}
608+
603609
pub fn with_expressions(
604610
mut self,
605611
expressions: Vec<Arc<dyn PhysicalExpr>>,

‎datafusion/ffi/src/physical_expr/partitioning.rs‎

Lines changed: 46 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,9 @@ use crate::physical_expr::sort::FFI_PhysicalSortExpr;
3333
#[repr(C)]
3434
#[derive(Debug)]
3535
pub struct FFI_RangePartitioning {
36-
split_points: SVec<SVec<WrappedArray>>,
36+
samples: SVec<SVec<WrappedArray>>,
3737
ordering: SVec<FFI_PhysicalSortExpr>,
38+
partition_count: usize,
3839
}
3940

4041
/// A stable struct for sharing [`Partitioning`] across FFI boundaries.
@@ -62,8 +63,8 @@ impl From<&Partitioning> for FFI_Partitioning {
6263
}
6364
Partitioning::Range(range) => {
6465
// Producer-side conversion should be infallible at ABI boundary
65-
let split_points = range
66-
.split_points()
66+
let samples = range
67+
.samples()
6768
.iter()
6869
.map(|split_point| {
6970
split_point
@@ -83,8 +84,9 @@ impl From<&Partitioning> for FFI_Partitioning {
8384
.map(FFI_PhysicalSortExpr::from)
8485
.collect();
8586
Self::Range(FFI_RangePartitioning {
86-
split_points,
87+
samples,
8788
ordering,
89+
partition_count: range.partition_count(),
8890
})
8991
}
9092
Partitioning::UnknownPartitioning(size) => Self::UnknownPartitioning(*size),
@@ -105,8 +107,8 @@ impl TryFrom<FFI_Partitioning> for Partitioning {
105107
Self::Hash(exprs, size)
106108
}
107109
FFI_Partitioning::Range(range) => {
108-
let split_points = range
109-
.split_points
110+
let samples = range
111+
.samples
110112
.into_iter()
111113
.map(|split_point| {
112114
split_point
@@ -126,7 +128,11 @@ impl TryFrom<FFI_Partitioning> for Partitioning {
126128
)
127129
})?;
128130

129-
Self::Range(RangePartitioning::try_new(ordering, split_points)?)
131+
Self::Range(RangePartitioning::try_new_with_samples(
132+
ordering,
133+
samples,
134+
range.partition_count,
135+
)?)
130136
}
131137
FFI_Partitioning::UnknownPartitioning(size) => {
132138
Self::UnknownPartitioning(size)
@@ -210,11 +216,43 @@ mod tests {
210216
Ok(())
211217
}
212218

219+
#[test]
220+
fn round_trip_ffi_sampled_range_partitioning() -> Result<()> {
221+
let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new(
222+
Column::new("a", 0),
223+
))])
224+
.expect("non-empty ordering");
225+
let samples = [10, 20, 30, 40, 50]
226+
.into_iter()
227+
.map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))]))
228+
.collect();
229+
let partitioning = Partitioning::Range(RangePartitioning::try_new_with_samples(
230+
ordering, samples, 3,
231+
)?);
232+
233+
let ffi_partitioning: FFI_Partitioning = (&partitioning).into();
234+
let returned: Partitioning = ffi_partitioning.try_into()?;
235+
let Partitioning::Range(returned) = returned else {
236+
panic!("expected range partitioning");
237+
};
238+
let Partitioning::Range(original) = partitioning else {
239+
panic!("expected range partitioning");
240+
};
241+
242+
assert_eq!(returned, original);
243+
assert_eq!(returned.samples(), original.samples());
244+
assert_eq!(returned.max_partition_count(), 6);
245+
assert_eq!(returned.partition_count(), 3);
246+
247+
Ok(())
248+
}
249+
213250
#[test]
214251
fn ffi_range_partitioning_rejects_empty_ordering() {
215252
let ffi_partitioning = FFI_Partitioning::Range(FFI_RangePartitioning {
216-
split_points: SVec::new(),
253+
samples: SVec::new(),
217254
ordering: SVec::new(),
255+
partition_count: 1,
218256
});
219257

220258
let err = Partitioning::try_from(ffi_partitioning).unwrap_err();

‎datafusion/ffi/src/plan_properties.rs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -290,11 +290,14 @@ mod tests {
290290
let col = datafusion::physical_plan::expressions::col("a", &schema)?;
291291
let ordering = LexOrdering::new([PhysicalSortExpr::new_default(col)])
292292
.expect("non-empty ordering");
293-
let split_points = vec![
293+
let samples = vec![
294294
SplitPoint::new(vec![ScalarValue::Int64(Some(10))]),
295295
SplitPoint::new(vec![ScalarValue::Int64(Some(20))]),
296+
SplitPoint::new(vec![ScalarValue::Int64(Some(30))]),
297+
SplitPoint::new(vec![ScalarValue::Int64(Some(40))]),
298+
SplitPoint::new(vec![ScalarValue::Int64(Some(50))]),
296299
];
297-
let range = RangePartitioning::try_new(ordering, split_points)?;
300+
let range = RangePartitioning::try_new_with_samples(ordering, samples, 3)?;
298301

299302
Ok(PlanProperties::new(
300303
EquivalenceProperties::new(schema),
@@ -314,7 +317,6 @@ mod tests {
314317
let foreign_props: PlanProperties = local_props_ptr.try_into()?;
315318

316319
assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}"));
317-
318320
Ok(())
319321
}
320322

@@ -351,6 +353,10 @@ mod tests {
351353
format!("{:?}", original_props.output_partitioning())
352354
);
353355
assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}"));
356+
let Partitioning::Range(range) = foreign_props.output_partitioning() else {
357+
panic!("expected range partitioning");
358+
};
359+
assert_eq!(range.max_partition_count(), 6);
354360

355361
Ok(())
356362
}

‎datafusion/ffi/src/tests/mod.rs‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,13 @@ use datafusion_catalog::MemTable;
2727
use datafusion_catalog::{Session, TableProvider};
2828
use datafusion_common::stats::Precision;
2929
use datafusion_common::{ColumnStatistics, Statistics};
30-
use datafusion_common::{Result, ScalarValue, exec_err};
30+
use datafusion_common::{Result, ScalarValue, SplitPoint, exec_err};
3131
use datafusion_expr::{Expr, TableType, col, lit};
32-
use datafusion_physical_expr::PhysicalExpr;
33-
use datafusion_physical_plan::ExecutionPlan;
32+
use datafusion_physical_expr::expressions::Column;
33+
use datafusion_physical_expr::{
34+
LexOrdering, PhysicalExpr, PhysicalSortExpr, RangePartitioning,
35+
};
36+
use datafusion_physical_plan::{ExecutionPlan, Partitioning};
3437
use sync_provider::create_sync_table_provider;
3538
use udf_udaf_udwf::{
3639
create_ffi_abs_func, create_ffi_first_value_func, create_ffi_random_func,
@@ -228,7 +231,22 @@ pub fn make_test_statistics() -> Statistics {
228231

229232
pub(crate) extern "C" fn create_exec_with_statistics() -> FFI_ExecutionPlan {
230233
let schema = create_test_schema();
231-
let plan = Arc::new(EmptyExec::new(schema).with_statistics(make_test_statistics()));
234+
let ordering =
235+
LexOrdering::new([PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))])
236+
.expect("non-empty ordering");
237+
let samples = [10, 20, 30, 40, 50]
238+
.into_iter()
239+
.map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))]))
240+
.collect();
241+
let partitioning = Partitioning::Range(
242+
RangePartitioning::try_new_with_samples(ordering, samples, 3)
243+
.expect("valid sampled range partitioning"),
244+
);
245+
let plan = Arc::new(
246+
EmptyExec::new(schema)
247+
.with_statistics(make_test_statistics())
248+
.with_partitioning(partitioning),
249+
);
232250
FFI_ExecutionPlan::new(plan, None)
233251
}
234252

‎datafusion/ffi/tests/ffi_execution_plan.rs‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ mod tests {
3030
use datafusion_physical_expr_common::metrics::{MetricCategory, MetricValue};
3131
use datafusion_physical_plan::execution_plan::InvariantLevel;
3232
use datafusion_physical_plan::{
33-
ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions,
33+
ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions,
3434
};
3535
use std::sync::Arc;
3636

@@ -142,6 +142,30 @@ mod tests {
142142
Ok(())
143143
}
144144

145+
#[test]
146+
fn test_ffi_range_partitioning_cross_library() -> Result<(), DataFusionError> {
147+
let module = get_module()?;
148+
let plan = (module.create_exec_with_statistics)();
149+
let plan: Arc<dyn ExecutionPlan> = (&plan).try_into()?;
150+
let Partitioning::Range(range) = plan.properties().output_partitioning() else {
151+
panic!("expected range partitioning");
152+
};
153+
154+
assert_eq!(range.partition_count(), 3);
155+
assert_eq!(range.max_partition_count(), 6);
156+
assert_eq!(range.samples().len(), 5);
157+
assert_eq!(
158+
range
159+
.split_points()
160+
.iter()
161+
.map(|point| point.to_string())
162+
.collect::<Vec<_>>(),
163+
vec!["(20)", "(40)"]
164+
);
165+
166+
Ok(())
167+
}
168+
145169
#[test]
146170
fn test_ffi_execution_plan_expressions_cross_library() -> Result<(), DataFusionError>
147171
{

0 commit comments

Comments
 (0)