diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 7c9d4984a6a68..b02c52b483874 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1123,6 +1123,433 @@ fn range_hash_join_repartitions_unpartitioned_side_to_match_range() -> Result<() Ok(()) } +#[test] +fn range_requirement_scales_within_sample_resolution() -> Result<()> { + let ordering = [PhysicalSortExpr::new_default(col("a", &schema())?)].into(); + let samples = (10..=50) + .step_by(10) + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect(); + let range = RangePartitioning::try_new_with_samples(ordering, samples, 3)?; + for target in [2, 3, 5, 6, 7] { + let input = + parquet_exec_with_output_partitioning(Partitioning::Range(range.clone())); + let requirement = RequirementsTestExec::new(input) + .with_required_input_distribution(Distribution::KeyPartitioned(vec![col( + "a", + &schema(), + )?])) + .into_arc(); + let plan = TestConfig::default() + .with_query_execution_partitions(target) + .to_plan(requirement, &DISTRIB_DISTRIB_SORT); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + let child = Arc::clone(plan.children()[0]); + if target <= range.max_partition_count() { + let Partitioning::Range(actual) = child.output_partitioning() else { + panic!("range partitioning lost for target {target}: {rendered}"); + }; + let expected_count = if (4..=6).contains(&target) { target } else { 3 }; + assert_eq!(actual.partition_count(), expected_count, "{rendered}"); + assert_eq!(actual.samples(), range.samples()); + assert_eq!( + actual.split_points(), + range.scale(expected_count).unwrap().split_points() + ); + } else { + let Partitioning::Hash(_, actual_count) = child.output_partitioning() else { + panic!("insufficient samples must fall back to hash: {rendered}"); + }; + assert_eq!(*actual_count, target, "{rendered}"); + } + assert_eq!( + rendered.matches("RepartitionExec:").count(), + usize::from(target > 3) + ); + } + Ok(()) +} + +#[test] +fn range_singleton_scaling_respects_required_keys() -> Result<()> { + let ordering = [PhysicalSortExpr::new_default(col("a", &schema())?)].into(); + let samples = (10..=50) + .step_by(10) + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect(); + let range = RangePartitioning::try_new_with_samples(ordering, samples, 1)?; + for key in ["a", "b"] { + let input = + parquet_exec_with_output_partitioning(Partitioning::Range(range.clone())); + let required = Distribution::KeyPartitioned(vec![col(key, &schema())?]); + let requirement = RequirementsTestExec::new(input) + .with_required_input_distribution(required.clone()) + .into_arc(); + let plan = TestConfig::default() + .with_query_execution_partitions(5) + .to_plan(requirement, &DISTRIB_DISTRIB_SORT); + let child = Arc::clone(plan.children()[0]); + assert_eq!(child.output_partitioning().partition_count(), 5); + assert!( + child + .output_partitioning() + .satisfaction(&required, child.equivalence_properties(), false) + .is_satisfied() + ); + assert_eq!( + matches!(child.output_partitioning(), Partitioning::Range(_)), + key == "a" + ); + } + Ok(()) +} + +#[test] +fn range_join_accepts_matching_layouts_with_different_sample_capacity() -> Result<()> { + let ordering = [PhysicalSortExpr::new_default(col("a", &schema())?)].into(); + let sampled = RangePartitioning::try_new_with_samples( + ordering, + (10..=50) + .step_by(10) + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect(), + 3, + )?; + let exact = RangePartitioning::try_new_with_samples( + sampled.ordering().clone(), + sampled.split_points().to_vec(), + 3, + )?; + + for swap in [false, true] { + let (left_range, right_range) = if swap { + (exact.clone(), sampled.clone()) + } else { + (sampled.clone(), exact.clone()) + }; + let left = parquet_exec_with_output_partitioning(Partitioning::Range(left_range)); + let right = + parquet_exec_with_output_partitioning(Partitioning::Range(right_range)); + let join_on = vec![(col("a", &schema())?, col("a", &schema())?)]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + let plan = TestConfig::default() + .with_query_execution_partitions(3) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + assert_eq!( + rendered.matches("RepartitionExec:").count(), + 0, + "{rendered}" + ); + let capacities = plan + .children() + .into_iter() + .map(|child| { + let Partitioning::Range(range) = child.output_partitioning() else { + panic!("expected range partitioning: {rendered}"); + }; + range.max_partition_count() + }) + .collect::>(); + assert_eq!(capacities, if swap { vec![3, 6] } else { vec![6, 3] }); + } + Ok(()) +} + +#[test] +fn range_preservation_keeps_existing_native_reference() -> Result<()> { + for (first_count, second_count, target, expected_range_count) in [ + (4, 3, 4, Some(4)), + (5, 3, 4, Some(5)), + (4, 4, 4, None), + (3, 2, 4, None), + ] { + for swap in [false, true] { + let first = parquet_exec_with_output_partitioning(range_partitioning( + "a", + (1..first_count).map(|i| i * 10), + SortOptions::default(), + )?); + let second = parquet_exec_with_output_partitioning(range_partitioning( + "a", + (1..second_count).map(|i| i * 10 + 5), + SortOptions::default(), + )?); + let (left, right) = if swap { + (second, first) + } else { + (first, second) + }; + let join_on = vec![(col("a", &schema())?, col("a", &schema())?)]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + let plan = TestConfig::default() + .with_query_execution_partitions(target) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + for child in plan.children() { + match (expected_range_count, child.output_partitioning()) { + (Some(count), Partitioning::Range(range)) => { + assert_eq!(range.partition_count(), count, "{rendered}") + } + (None, Partitioning::Hash(_, count)) => { + assert_eq!(*count, target, "{rendered}") + } + _ => panic!("unexpected reference choice: {rendered}"), + } + } + } + } + Ok(()) +} + +#[test] +fn range_fallback_preserves_target_parallelism() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let input = + |count: usize, rows_per_partition: usize| -> Result> { + let partitions = (0..count) + .map(|partition| { + Ok(vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![ + partition as i64 * 10; + rows_per_partition + ]))], + )?]) + }) + .collect::>>()?; + let source = + MemorySourceConfig::try_new_exec(&partitions, Arc::clone(&schema), None)?; + Ok(Arc::new(source.as_ref().clone().with_partitioning( + range_partitioning( + "a", + (1..count).map(|i| i as i64 * 10), + SortOptions::default(), + )?, + ))) + }; + for swap in [false, true] { + let first = input(4, 1)?; + let second = input(3, 100)?; + let (left, right) = if swap { + (second, first) + } else { + (first, second) + }; + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + for child in plan.children() { + let Partitioning::Range(range) = child.output_partitioning() else { + panic!("{rendered}") + }; + assert_eq!(range.partition_count(), 4, "{rendered}"); + } + } + Ok(()) +} + +#[test] +fn range_hash_join_falls_back_above_max_partitions() -> Result<()> { + let left = parquet_exec(); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![(col("a", &left.schema())?, col("a", &right.schema())?)]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + let plan = TestConfig::default() + .with_query_execution_partitions(8) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + assert_eq!( + rendered.matches("partitioning=Hash([a@0], 8)").count(), + 2, + "{rendered}" + ); + assert_eq!( + rendered.matches("RepartitionExec:").count(), + 2, + "{rendered}" + ); + Ok(()) +} + +#[tokio::test] +async fn range_singleton_scaling_preserves_aggregate_groups() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![0, 10, 20, 30, 40, 50])), + Arc::new(Int64Array::from(vec![1, 2, 1, 2, 1, 2])), + ], + )?; + let range = RangePartitioning::try_new_with_samples( + [PhysicalSortExpr::new_default(col("a", &schema)?)].into(), + (10..=50) + .step_by(10) + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect(), + 1, + )?; + let source = + MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + let source = Arc::new( + source + .as_ref() + .clone() + .with_partitioning(Partitioning::Range(range)), + ); + // The input represents partial groups. Repeated b values must be combined, + // even though they lie on opposite sides of the retained a boundaries. + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + PhysicalGroupBy::new_single(vec![(col("b", &schema)?, "b".to_string())]), + vec![], + vec![], + source, + schema, + )?); + let plan = TestConfig::default() + .with_query_execution_partitions(5) + .to_plan(aggregate, &DISTRIB_DISTRIB_SORT); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + assert!( + rendered.contains("partitioning=Hash([b@1], 5)"), + "{rendered}" + ); + let batches = collect(plan, SessionContext::new().task_ctx()).await?; + let mut groups = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + groups.sort_unstable(); + assert_eq!(groups, vec![1, 2]); + Ok(()) +} + +#[tokio::test] +async fn range_hash_join_scaling_preserves_rows() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])); + let batch = |values: Vec>| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + ) + }; + for descending in [false, true] { + // Physical rows must match the declared boundaries in each sort direction. + let partitions = if descending { + vec![ + vec![batch(vec![Some(99), Some(50)])?], + vec![batch(vec![Some(40), Some(39), Some(30)])?], + vec![batch(vec![Some(20), Some(19), Some(10), Some(0), None])?], + ] + } else { + vec![ + vec![batch(vec![None, Some(0), Some(10), Some(19)])?], + vec![batch(vec![Some(20), Some(30), Some(39)])?], + vec![batch(vec![Some(40), Some(50), Some(99)])?], + ] + }; + let ordering = [PhysicalSortExpr::new( + col("a", &schema)?, + SortOptions::new(descending, !descending), + )] + .into(); + let samples = (1..=5) + .map(|index| { + if descending { + 60 - index * 10 + } else { + index * 10 + } + }) + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect(); + let range = RangePartitioning::try_new_with_samples(ordering, samples, 3)?; + for target in [3, 5, 6, 8] { + let source = + MemorySourceConfig::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let right: Arc = Arc::new( + source + .as_ref() + .clone() + .with_partitioning(Partitioning::Range(range.clone())), + ); + let left = MemorySourceConfig::try_new_exec( + &[partitions.iter().flatten().cloned().collect()], + Arc::clone(&schema), + None, + )?; + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + let plan = TestConfig::default() + .with_query_execution_partitions(target) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + let rendered = displayable(plan.as_ref()).indent(true).to_string(); + let expected_count = target; + for child in plan.children() { + if target <= range.max_partition_count() { + let Partitioning::Range(actual) = child.output_partitioning() else { + panic!("target {target}: {rendered}"); + }; + assert_eq!(actual.partition_count(), expected_count, "{rendered}"); + assert_eq!( + actual.split_points(), + range.scale(expected_count).unwrap().split_points() + ); + } else { + let Partitioning::Hash(_, actual_count) = child.output_partitioning() + else { + panic!("target {target}: {rendered}"); + }; + assert_eq!(*actual_count, expected_count, "{rendered}"); + } + } + let batches = collect(plan, SessionContext::new().task_ctx()).await?; + let mut rows = vec![]; + for batch in &batches { + let left = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let right = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend(left.iter().zip(right.iter())); + } + rows.sort(); + let expected = [0, 10, 19, 20, 30, 39, 40, 50, 99] + .into_iter() + .map(|value| (Some(value), Some(value))) + .collect::>(); + assert_eq!(rows, expected, "target {target}: {rendered}"); + } + } + Ok(()) +} + #[test] fn range_hash_join_rehashes_incompatible_data_type() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 056d29d16aecf..b242d1b2757b2 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -600,6 +600,12 @@ pub mod tests { self } + pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self { + self.props = + Arc::new(self.props.as_ref().clone().with_partitioning(partitioning)); + self + } + pub fn with_expressions( mut self, expressions: Vec>, diff --git a/datafusion/ffi/src/physical_expr/partitioning.rs b/datafusion/ffi/src/physical_expr/partitioning.rs index 2a9a8528c6c3e..05535caefaef3 100644 --- a/datafusion/ffi/src/physical_expr/partitioning.rs +++ b/datafusion/ffi/src/physical_expr/partitioning.rs @@ -33,8 +33,9 @@ use crate::physical_expr::sort::FFI_PhysicalSortExpr; #[repr(C)] #[derive(Debug)] pub struct FFI_RangePartitioning { - split_points: SVec>, + samples: SVec>, ordering: SVec, + partition_count: usize, } /// A stable struct for sharing [`Partitioning`] across FFI boundaries. @@ -62,8 +63,8 @@ impl From<&Partitioning> for FFI_Partitioning { } Partitioning::Range(range) => { // Producer-side conversion should be infallible at ABI boundary - let split_points = range - .split_points() + let samples = range + .samples() .iter() .map(|split_point| { split_point @@ -83,8 +84,9 @@ impl From<&Partitioning> for FFI_Partitioning { .map(FFI_PhysicalSortExpr::from) .collect(); Self::Range(FFI_RangePartitioning { - split_points, + samples, ordering, + partition_count: range.partition_count(), }) } Partitioning::UnknownPartitioning(size) => Self::UnknownPartitioning(*size), @@ -105,8 +107,8 @@ impl TryFrom for Partitioning { Self::Hash(exprs, size) } FFI_Partitioning::Range(range) => { - let split_points = range - .split_points + let samples = range + .samples .into_iter() .map(|split_point| { split_point @@ -126,7 +128,11 @@ impl TryFrom for Partitioning { ) })?; - Self::Range(RangePartitioning::try_new(ordering, split_points)?) + Self::Range(RangePartitioning::try_new_with_samples( + ordering, + samples, + range.partition_count, + )?) } FFI_Partitioning::UnknownPartitioning(size) => { Self::UnknownPartitioning(size) @@ -174,6 +180,20 @@ mod tests { )?)) } + fn sampled_range_partitioning() -> Result { + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("a", 0), + ))]) + .expect("non-empty ordering"); + let samples = [10, 20, 30, 40, 50] + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect(); + Ok(Partitioning::Range( + RangePartitioning::try_new_with_samples(ordering, samples, 3)?, + )) + } + #[test] fn round_trip_ffi_partitioning() -> Result<()> { for partitioning in [ @@ -181,6 +201,7 @@ mod tests { Partitioning::Hash(vec![lit(1)], 10), Partitioning::UnknownPartitioning(10), range_partitioning()?, + sampled_range_partitioning()?, ] { let ffi_partitioning: FFI_Partitioning = (&partitioning).into(); let returned: Partitioning = ffi_partitioning.try_into()?; @@ -210,11 +231,33 @@ mod tests { Ok(()) } + #[test] + fn round_trip_ffi_sampled_range_partitioning() -> Result<()> { + let partitioning = sampled_range_partitioning()?; + + let ffi_partitioning: FFI_Partitioning = (&partitioning).into(); + let returned: Partitioning = ffi_partitioning.try_into()?; + let Partitioning::Range(returned) = returned else { + panic!("expected range partitioning"); + }; + let Partitioning::Range(original) = partitioning else { + panic!("expected range partitioning"); + }; + + assert_eq!(returned, original); + assert_eq!(returned.samples(), original.samples()); + assert_eq!(returned.max_partition_count(), 6); + assert_eq!(returned.partition_count(), 3); + + Ok(()) + } + #[test] fn ffi_range_partitioning_rejects_empty_ordering() { let ffi_partitioning = FFI_Partitioning::Range(FFI_RangePartitioning { - split_points: SVec::new(), + samples: SVec::new(), ordering: SVec::new(), + partition_count: 1, }); let err = Partitioning::try_from(ffi_partitioning).unwrap_err(); diff --git a/datafusion/ffi/src/plan_properties.rs b/datafusion/ffi/src/plan_properties.rs index dcbdbe59d14ab..6913bc0a56435 100644 --- a/datafusion/ffi/src/plan_properties.rs +++ b/datafusion/ffi/src/plan_properties.rs @@ -290,11 +290,14 @@ mod tests { let col = datafusion::physical_plan::expressions::col("a", &schema)?; let ordering = LexOrdering::new([PhysicalSortExpr::new_default(col)]) .expect("non-empty ordering"); - let split_points = vec![ + let samples = vec![ SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(30))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(40))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(50))]), ]; - let range = RangePartitioning::try_new(ordering, split_points)?; + let range = RangePartitioning::try_new_with_samples(ordering, samples, 3)?; Ok(PlanProperties::new( EquivalenceProperties::new(schema), @@ -314,7 +317,6 @@ mod tests { let foreign_props: PlanProperties = local_props_ptr.try_into()?; assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}")); - Ok(()) } @@ -351,6 +353,10 @@ mod tests { format!("{:?}", original_props.output_partitioning()) ); assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}")); + let Partitioning::Range(range) = foreign_props.output_partitioning() else { + panic!("expected range partitioning"); + }; + assert_eq!(range.max_partition_count(), 6); Ok(()) } diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 7e583b6c5d5bf..a334bd6d21f43 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -27,10 +27,13 @@ use datafusion_catalog::MemTable; use datafusion_catalog::{Session, TableProvider}; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, Statistics}; -use datafusion_common::{Result, ScalarValue, exec_err}; +use datafusion_common::{Result, ScalarValue, SplitPoint, exec_err}; use datafusion_expr::{Expr, TableType, col, lit}; -use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::{ + LexOrdering, PhysicalExpr, PhysicalSortExpr, RangePartitioning, +}; +use datafusion_physical_plan::{ExecutionPlan, Partitioning}; use sync_provider::create_sync_table_provider; use udf_udaf_udwf::{ create_ffi_abs_func, create_ffi_first_value_func, create_ffi_random_func, @@ -228,7 +231,22 @@ pub fn make_test_statistics() -> Statistics { pub(crate) extern "C" fn create_exec_with_statistics() -> FFI_ExecutionPlan { let schema = create_test_schema(); - let plan = Arc::new(EmptyExec::new(schema).with_statistics(make_test_statistics())); + let ordering = + LexOrdering::new([PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))]) + .expect("non-empty ordering"); + let samples = [10, 20, 30, 40, 50] + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(); + let partitioning = Partitioning::Range( + RangePartitioning::try_new_with_samples(ordering, samples, 3) + .expect("valid sampled range partitioning"), + ); + let plan = Arc::new( + EmptyExec::new(schema) + .with_statistics(make_test_statistics()) + .with_partitioning(partitioning), + ); FFI_ExecutionPlan::new(plan, None) } diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index 28828695d01af..43a50f1286e4f 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -30,7 +30,7 @@ mod tests { use datafusion_physical_expr_common::metrics::{MetricCategory, MetricValue}; use datafusion_physical_plan::execution_plan::InvariantLevel; use datafusion_physical_plan::{ - ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, + ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, }; use std::sync::Arc; @@ -142,6 +142,30 @@ mod tests { Ok(()) } + #[test] + fn test_ffi_range_partitioning_cross_library() -> Result<(), DataFusionError> { + let module = get_module()?; + let plan = (module.create_exec_with_statistics)(); + let plan: Arc = (&plan).try_into()?; + let Partitioning::Range(range) = plan.properties().output_partitioning() else { + panic!("expected range partitioning"); + }; + + assert_eq!(range.partition_count(), 3); + assert_eq!(range.max_partition_count(), 6); + assert_eq!(range.samples().len(), 5); + assert_eq!( + range + .split_points() + .iter() + .map(|point| point.to_string()) + .collect::>(), + vec!["(20)", "(40)"] + ); + + Ok(()) + } + #[test] fn test_ffi_execution_plan_expressions_cross_library() -> Result<(), DataFusionError> { diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 052fec6ebe97e..53e9d88112691 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -23,7 +23,7 @@ use crate::{ }; use arrow::datatypes::Schema; pub use datafusion_common::SplitPoint; -use datafusion_common::{Result, validate_range_split_points}; +use datafusion_common::{Result, plan_err, validate_range_split_points}; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; #[cfg(feature = "proto")] @@ -153,15 +153,23 @@ impl Display for Partitioning { /// Physical range partitioning. /// -/// [`RangePartitioning`] describes an ordered key space with split points. +/// [`RangePartitioning`] describes an ordered key space with sampled split points. /// /// - `ordering` defines the partitioning key and ordering. -/// - `split_points` define the boundaries between adjacent partitions. +/// - `samples` are the maximum-resolution split points supplied by the caller. +/// - The effective `split_points` are derived from those samples for the selected +/// partition count; the count itself is `split_points.len() + 1` and is not +/// stored separately. /// /// Comparisons use the lexicographic order defined by `ordering`, including -/// `ASC`/`DESC` and null ordering. Split points must be strictly ordered -/// according to that ordering, and each split point must have one value per -/// ordering expression. See [`SplitPoint`] for the shared boundary convention. +/// `ASC`/`DESC` and null ordering. Samples must be strictly ordered according +/// to that ordering, and each sample must have one value per ordering +/// expression. See [`SplitPoint`] for the shared boundary convention. +/// +/// When `partition_count` is smaller than [`Self::max_partition_count`], the +/// samples are evenly down-sampled to derive the effective split points. This +/// allows planners to reduce or later restore the number of partitions without +/// losing the original distribution sample. /// /// Like other user-specified data properties such as sortedness, if a source /// declares range partitioning, it is responsible for placing each row in the @@ -198,40 +206,84 @@ impl Display for Partitioning { /// partition 2: keys at/after (2023, Allston) /// ``` /// -/// NOTE: Optimizer and execution behavior for this partitioning is intentionally -/// not implemented and will be introduced incrementally. See -/// . +/// Equality includes retained samples, since they determine which future scales +/// are possible. Use [`Self::has_same_layout`] to compare only the current layout. #[derive(Debug, Clone, PartialEq)] pub struct RangePartitioning { /// Ordered partitioning key. ordering: LexOrdering, - /// Boundaries between adjacent partitions. - split_points: Vec, + /// Caller-supplied maximum-resolution split points used to derive the + /// effective boundaries. + samples: Arc<[SplitPoint]>, + /// Effective boundaries for the current partition count. + split_points: Arc<[SplitPoint]>, } impl RangePartitioning { /// Creates range partitioning metadata without validating split points. /// - /// Use [`Self::try_new`] to validate the contract documented on - /// [`RangePartitioning`]. + /// Prefer [`Self::try_new_with_samples`] to validate the boundaries and retain + /// additional samples for scaling up. [`Self::try_new`] remains available for + /// validated exact boundaries. + #[deprecated( + since = "56.0.0", + note = "Use RangePartitioning::try_new_with_samples instead" + )] pub fn new(ordering: LexOrdering, split_points: Vec) -> Self { + let split_points: Arc<[SplitPoint]> = Arc::from(split_points); Self { ordering, + samples: Arc::clone(&split_points), split_points, } } /// Creates range partitioning metadata and validates split point shape and /// ordering. + /// + /// The exact boundaries are also the retained samples. This allows scaling + /// down and back up to the original count, but not beyond it. Prefer + /// [`Self::try_new_with_samples`] when additional sample points are available. pub fn try_new(ordering: LexOrdering, split_points: Vec) -> Result { + let partition_count = split_points.len() + 1; + Self::try_new_with_samples(ordering, split_points, partition_count) + } + + /// 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. + /// + /// Retain at least `maximum_expected_partitions - 1` samples to support that + /// many partitions later. For example, samples at `[10, 20, ..., 90]` with + /// `partition_count = 4` produce effective split points `[30, 50, 70]` and + /// can later scale to any count from 1 through 10. Supplying approximately + /// `K * target_partitions` samples therefore leaves room to scale by about + /// `K`; choose `K` for the workload. Small inputs may not have enough distinct + /// values, in which case [`Self::scale`] returns `None` above the supported + /// maximum. + pub fn try_new_with_samples( + ordering: LexOrdering, + samples: Vec, + partition_count: usize, + ) -> Result { validate_range_split_points( - &split_points, + &samples, &ordering .iter() .map(|sort_expr| sort_expr.options) .collect::>(), )?; - Ok(Self::new(ordering, split_points)) + validate_range_partition_count(partition_count, samples.len() + 1)?; + let samples: Arc<[SplitPoint]> = Arc::from(samples); + let split_points = downsample_split_points(&samples, partition_count); + Ok(Self { + ordering, + samples, + split_points, + }) } /// Returns the ordering that defines the range key. @@ -239,7 +291,12 @@ impl RangePartitioning { &self.ordering } - /// Returns the ordered split points between partitions. + /// Returns the maximum-resolution sample points. + pub fn samples(&self) -> &[SplitPoint] { + &self.samples + } + + /// Returns the effective split points between partitions. pub fn split_points(&self) -> &[SplitPoint] { &self.split_points } @@ -249,6 +306,44 @@ impl RangePartitioning { self.split_points.len() + 1 } + /// Returns the largest partition count supported by the stored samples. + pub fn max_partition_count(&self) -> usize { + self.samples.len() + 1 + } + + /// Whether two range partitionings have the same current key ordering and + /// effective boundaries, irrespective of their retained samples. + /// + /// This does not imply equal scaling capacity. In particular, a plan that + /// combines inputs must not use this comparison to inherit one input's samples + /// for all inputs. Structural equality is required for that use case. + pub fn has_same_layout(&self, other: &Self) -> bool { + self.ordering == other.ordering && self.split_points == other.split_points + } + + /// 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`]. + /// Returns `None` when `target_partitions` is zero or the retained samples do + /// not support that many partitions. Insufficient samples are an expected + /// condition; callers should retain the layout or choose another partitioning. + /// This method changes metadata only; the caller must ensure that the actual + /// row distribution matches the resulting boundaries. + pub fn scale(&self, target_partitions: usize) -> Option { + if target_partitions == 0 || target_partitions > self.max_partition_count() { + return None; + } + if target_partitions == self.partition_count() { + return Some(self.clone()); + } + Some(Self { + ordering: self.ordering.clone(), + samples: Arc::clone(&self.samples), + split_points: downsample_split_points(&self.samples, target_partitions), + }) + } + /// Calculates the range partitioning after applying the given projection. /// /// Returns `None` if any range key cannot be projected or if projection @@ -279,7 +374,8 @@ impl RangePartitioning { Some(Self { ordering, - split_points: self.split_points.clone(), + samples: Arc::clone(&self.samples), + split_points: Arc::clone(&self.split_points), }) } @@ -292,7 +388,7 @@ impl RangePartitioning { if self.ordering.len() != exprs.len() { return false; } - if let Some(first_split) = self.split_points.first() { + if let Some(first_split) = self.samples.first() { exprs.iter().zip(first_split.values()).all(|(expr, val)| { expr.data_type(schema) .map(|dt| dt == val.data_type()) @@ -323,7 +419,14 @@ impl RangePartitioning { }) .collect::>(), )?; - Self::try_new(new_ordering, self.split_points.clone()).ok() + if new_ordering.len() != self.ordering.len() { + return None; + } + Some(Self { + ordering: new_ordering, + samples: Arc::clone(&self.samples), + split_points: Arc::clone(&self.split_points), + }) } } @@ -332,12 +435,52 @@ impl Display for RangePartitioning { let split_points = format_range_split_points(&self.split_points); write!( f, - "Range([{}], [{}], {})", + "Range([{}], [{}], {}", self.ordering, split_points, self.partition_count() - ) + )?; + if self.max_partition_count() != self.partition_count() { + write!(f, ", max {}", self.max_partition_count())?; + } + write!(f, ")") + } +} + +fn downsample_split_points( + samples: &Arc<[SplitPoint]>, + partition_count: usize, +) -> Arc<[SplitPoint]> { + if partition_count == samples.len() + 1 { + return Arc::clone(samples); + } + + let sample_count = samples.len(); + (1..partition_count) + .map(|partition| { + // Use a wider intermediate so valid slice lengths cannot overflow + // when calculating the evenly spaced sample index. + let sample_index = ((partition as u128 * sample_count as u128) + / partition_count as u128) as usize; + samples[sample_index].clone() + }) + .collect::>() + .into() +} + +fn validate_range_partition_count( + partition_count: usize, + max_partition_count: usize, +) -> Result<()> { + if partition_count == 0 { + return plan_err!("Range partitioning partition count must be at least 1"); + } + if partition_count > max_partition_count { + return plan_err!( + "Range partitioning partition count {partition_count} exceeds maximum {max_partition_count}" + ); } + Ok(()) } fn format_range_split_points(split_points: &[SplitPoint]) -> String { @@ -629,26 +772,28 @@ impl Partitioning { ) } Partitioning::Range(range) => { - let RangePartitioning { - ordering, - split_points, - } = range; - let sort_expr = sort_exprs_try_to_proto(ordering.iter(), ctx)?; - let split_point = split_points - .iter() - .map(|split_point| { - let value = split_point - .values() - .iter() - .map(|value| value.try_into().map_err(Into::into)) - .collect::>>()?; - Ok(protobuf::PhysicalRangeSplitPoint { value }) - }) - .collect::>>()?; + let sort_expr = sort_exprs_try_to_proto(range.ordering().iter(), ctx)?; + let encode_split_points = |split_points: &[SplitPoint]| { + split_points + .iter() + .map(|split_point| { + let value = split_point + .values() + .iter() + .map(|value| value.try_into().map_err(Into::into)) + .collect::>>()?; + Ok(protobuf::PhysicalRangeSplitPoint { value }) + }) + .collect::>>() + }; + let split_point = encode_split_points(range.split_points())?; + let sample_point = encode_split_points(range.samples())?; protobuf::partitioning::PartitionMethod::Range( protobuf::PhysicalRangePartitioning { sort_expr, split_point, + sample_point, + partition_count: partition_count(range.partition_count())?, }, ) } @@ -699,11 +844,7 @@ impl Partitioning { Partitioning::UnknownPartitioning(partition_count(*n)?) } protobuf::partitioning::PartitionMethod::Range(range) => { - let protobuf::PhysicalRangePartitioning { - sort_expr, - split_point, - } = range; - let sort_exprs = sort_exprs_try_from_proto(sort_expr, ctx)?; + let sort_exprs = sort_exprs_try_from_proto(&range.sort_expr, ctx)?; let sort_expr_count = sort_exprs.len(); let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { internal_datafusion_err!( @@ -715,18 +856,44 @@ impl Partitioning { "Range partitioning ordering must not contain duplicate expressions" ); } - let split_points = split_point - .iter() - .map(|split_point| { - let protobuf::PhysicalRangeSplitPoint { value } = split_point; - let values = value + let decode_split_points = + |split_points: &[protobuf::PhysicalRangeSplitPoint]| { + split_points .iter() - .map(|value| ScalarValue::try_from(value).map_err(Into::into)) - .collect::>>()?; - Ok(SplitPoint::new(values)) - }) - .collect::>>()?; - Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) + .map(|split_point| { + let values = split_point + .value + .iter() + .map(|value| { + ScalarValue::try_from(value).map_err(Into::into) + }) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect::>>() + }; + let split_points = decode_split_points(&range.split_point)?; + if range.partition_count == 0 { + // Older payloads derive their partition count from the exact + // split points and do not carry this field. + Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?) + } else { + let samples = decode_split_points(&range.sample_point)?; + let range_partitioning = RangePartitioning::try_new_with_samples( + ordering, + samples, + partition_count(range.partition_count)?, + )?; + if range_partitioning.split_points() != split_points { + return internal_err!( + "Range partitioning effective split points do not match its samples and partition count" + ); + } + Partitioning::Range(range_partitioning) + } } }; Ok(Some(partitioning)) @@ -913,17 +1080,6 @@ mod tests { ) -> Partitioning { Partitioning::Range(self.range(indices, split_points)) } - - fn range_partitioning_with_ordering( - &self, - ordering: LexOrdering, - split_points: Vec, - ) -> Partitioning { - Partitioning::Range( - RangePartitioning::try_new(ordering, split_points) - .expect("test range partitioning should be valid"), - ) - } } fn assert_satisfaction( @@ -1192,6 +1348,160 @@ mod tests { Ok(()) } + #[test] + fn test_range_partitioning_scales_from_samples() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a"])?; + let samples = (10..=90) + .step_by(10) + .map(|value| int_split_point([value])) + .collect::>(); + let range = RangePartitioning::try_new_with_samples( + fixture.range_ordering([0]), + samples.clone(), + 4, + )?; + + assert_eq!(range.partition_count(), 4); + assert_eq!(range.max_partition_count(), 10); + assert_eq!(range.samples(), samples); + assert_eq!( + range.split_points(), + vec![ + int_split_point([30]), + int_split_point([50]), + int_split_point([70]), + ] + ); + assert_eq!( + range.to_string(), + "Range([a@0 ASC], [(30), (50), (70)], 4, max 10)" + ); + + let single = range.scale(1).expect("one partition is supported"); + assert_eq!(single.partition_count(), 1); + assert!(single.split_points().is_empty()); + assert_eq!(single.max_partition_count(), 10); + assert_eq!(single.to_string(), "Range([a@0 ASC], [], 1, max 10)"); + + let restored = single + .scale(single.max_partition_count()) + .expect("retained samples support restoration"); + assert_eq!(restored.split_points(), samples); + assert_eq!(restored.samples(), samples); + + Ok(()) + } + + #[test] + fn test_range_partitioning_rejects_invalid_partition_count() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a"])?; + let ordering = fixture.range_ordering([0]); + let samples = vec![int_split_point([10]), int_split_point([20])]; + + let error = + RangePartitioning::try_new_with_samples(ordering.clone(), samples.clone(), 0) + .unwrap_err() + .to_string(); + assert!(error.contains("must be at least 1"), "{error}"); + + let error = + RangePartitioning::try_new_with_samples(ordering.clone(), samples.clone(), 4) + .unwrap_err() + .to_string(); + assert!(error.contains("exceeds maximum 3"), "{error}"); + + let range = RangePartitioning::try_new(ordering, samples)?; + assert!(range.scale(4).is_none()); + + Ok(()) + } + + #[test] + fn test_range_partitioning_equality_includes_scaling_capacity() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a"])?; + let ordering = fixture.range_ordering([0]); + let sampled = RangePartitioning::try_new_with_samples( + ordering.clone(), + (10..=90) + .step_by(10) + .map(|value| int_split_point([value])) + .collect(), + 4, + )?; + let exact = RangePartitioning::try_new( + ordering, + vec![ + int_split_point([30]), + int_split_point([50]), + int_split_point([70]), + ], + )?; + + assert_eq!(sampled.split_points(), exact.split_points()); + assert!(sampled.has_same_layout(&exact)); + assert!(exact.has_same_layout(&sampled)); + assert!(sampled.scale(10).is_some()); + assert!(exact.scale(10).is_none()); + assert_ne!(sampled, exact); + assert_ne!(Partitioning::Range(sampled), Partitioning::Range(exact)); + + Ok(()) + } + + #[test] + fn test_range_partitioning_layout_requires_keys_options_and_boundaries() -> Result<()> + { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let range = RangePartitioning::try_new_with_samples( + fixture.range_ordering([0]), + vec![int_split_point([10])], + 2, + )?; + assert!(range.has_same_layout(&range.clone())); + let different_boundary = RangePartitioning::try_new_with_samples( + fixture.range_ordering([0]), + vec![int_split_point([20])], + 2, + )?; + assert!(!range.has_same_layout(&different_boundary)); + assert!(!range.has_same_layout(&range.scale(1).unwrap())); + let singleton = range.scale(1).unwrap(); + for ordering in [ + fixture.range_ordering([1]), + [fixture.range_sort_expr(0, SortOptions::new(true, false))].into(), + [fixture.range_sort_expr(0, SortOptions::new(false, false))].into(), + ] { + let other = RangePartitioning::try_new_with_samples(ordering, vec![], 1)?; + assert!(!singleton.has_same_layout(&other)); + } + Ok(()) + } + + #[test] + fn test_range_partitioning_scaling_limits() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a"])?; + let range = RangePartitioning::try_new_with_samples( + fixture.range_ordering([0]), + vec![int_split_point([10])], + 1, + )?; + assert_eq!(range.scale(0), None); + for target_partitions in [3, usize::MAX] { + assert_eq!(range.scale(target_partitions), None); + } + assert_eq!(range.scale(1).unwrap(), range); + assert_eq!(range.scale(2).unwrap().scale(1).unwrap(), range); + assert_eq!(range.scale(2).unwrap().partition_count(), 2); + let empty = RangePartitioning::try_new_with_samples( + fixture.range_ordering([0]), + vec![], + 1, + )?; + assert_eq!(empty.scale(1).unwrap(), empty); + assert_eq!(empty.scale(2), None); + Ok(()) + } + #[test] fn test_range_partitioning_try_new_validates_split_points() -> Result<()> { let fixture = PartitioningTestFixture::int64(&["a", "b"])?; @@ -1241,18 +1551,29 @@ mod tests { #[test] fn test_range_partitioning_project_preserves_or_degrades() -> Result<()> { let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let range_partitioning = fixture.range_partitioning_with_ordering( - [fixture.range_sort_expr(1, SortOptions::new(true, false))].into(), - vec![int_split_point([10])], - ); + let range_partitioning = + Partitioning::Range(RangePartitioning::try_new_with_samples( + [fixture.range_sort_expr(1, SortOptions::new(true, false))].into(), + vec![ + int_split_point([30]), + int_split_point([20]), + int_split_point([10]), + ], + 2, + )?); let keep_b_mapping = ProjectionMapping::from_indices(&[1], &fixture.schema)?; let projected = range_partitioning.project(&keep_b_mapping, &fixture.eq_properties); assert_eq!( projected.to_string(), - "Range([b@0 DESC NULLS LAST], [(10)], 2)" + "Range([b@0 DESC NULLS LAST], [(20)], 2, max 4)" ); + let Partitioning::Range(projected_range) = &projected else { + panic!("expected range partitioning, got {projected:?}"); + }; + assert_eq!(projected_range.max_partition_count(), 4); + assert_eq!(projected_range.scale(4).unwrap().split_points().len(), 3); let drop_b_mapping = ProjectionMapping::from_indices(&[0], &fixture.schema)?; let projected = @@ -1364,22 +1685,40 @@ mod tests { ("c", DataType::Int32), ])?; - let range = fixture.range( - [0], + let range = RangePartitioning::try_new_with_samples( + fixture.range_ordering([0]), vec![ SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(15))]), SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(25))]), ], - ); + 3, + )?; // Adapting to col_c (same type Int32) succeeds let adapted = range.adapt(&[fixture.col(2)], &fixture.schema).unwrap(); assert_eq!(adapted.ordering().len(), 1); assert!(adapted.ordering()[0].expr.eq(&fixture.col(2))); assert_eq!(adapted.partition_count(), 3); + assert_eq!(adapted.max_partition_count(), 5); + assert_eq!(adapted.scale(5).unwrap().partition_count(), 5); // Adapting to col_b (different type Int64) fails assert!(range.adapt(&[fixture.col(1)], &fixture.schema).is_none()); + // Scaling to one partition removes effective boundaries, not sample types. + let singleton = range.scale(1).unwrap(); + assert!( + singleton + .adapt(&[fixture.col(1)], &fixture.schema) + .is_none() + ); + let adapted_singleton = + singleton.adapt(&[fixture.col(2)], &fixture.schema).unwrap(); + assert_eq!( + adapted_singleton.scale(5).unwrap().samples(), + range.samples() + ); // Adapting to empty or mismatch count fails assert!(range.adapt(&[], &fixture.schema).is_none()); @@ -1421,6 +1760,33 @@ mod tests { Ok(()) } + #[test] + fn test_range_partitioning_adapt_rejects_duplicate_keys() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![ + ("a", DataType::Int32), + ("b", DataType::Int32), + ("c", DataType::Int32), + ])?; + let range = RangePartitioning::try_new_with_samples( + fixture.range_ordering([0, 1]), + vec![SplitPoint::new(vec![ + ScalarValue::Int32(Some(10)), + ScalarValue::Int32(Some(20)), + ])], + 1, + )?; + for count in [1, 2] { + assert!( + range + .scale(count) + .unwrap() + .adapt(&fixture.cols([2, 2]), &fixture.schema) + .is_none() + ); + } + Ok(()) + } + #[test] fn test_range_partitioning_adapt_multi_key() -> Result<()> { let fixture = PartitioningTestFixture::new(vec![ @@ -1578,6 +1944,127 @@ mod ordering_proto_tests { } } +#[cfg(all(test, feature = "proto"))] +mod range_partitioning_proto_tests { + use std::sync::Arc; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{Result, ScalarValue, SplitPoint}; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion_proto_models::protobuf; + + use super::{Partitioning, RangePartitioning}; + use crate::expressions::Column; + use crate::proto_test_util::{StubDecoder, StubEncoder}; + + fn sampled_partitioning() -> Result { + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("a", 0), + ))]) + .expect("non-empty ordering"); + let samples = [10, 20, 30, 40, 50] + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(); + Ok(Partitioning::Range( + RangePartitioning::try_new_with_samples(ordering, samples, 3)?, + )) + } + + fn decode(partitioning: &protobuf::Partitioning) -> Result { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let decoder = StubDecoder::ok(); + let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + Ok(Partitioning::try_from_proto(partitioning, &decode_ctx)? + .expect("partitioning method is present")) + } + + #[test] + fn sampled_range_partitioning_round_trip_preserves_resolution() -> Result<()> { + let partitioning = sampled_partitioning()?; + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let encoded = partitioning.try_to_proto(&encode_ctx)?; + let Some(protobuf::partitioning::PartitionMethod::Range(encoded_range)) = + encoded.partition_method.as_ref() + else { + panic!("expected range partitioning"); + }; + + // Field 2 remains the effective boundary list for older readers. + assert_eq!(encoded_range.split_point.len(), 2); + assert_eq!(encoded_range.sample_point.len(), 5); + assert_eq!(encoded_range.partition_count, 3); + + let decoded = decode(&encoded)?; + let Partitioning::Range(decoded) = decoded else { + panic!("expected range partitioning"); + }; + let Partitioning::Range(original) = partitioning else { + panic!("expected range partitioning"); + }; + assert_eq!(decoded.partition_count(), original.partition_count()); + assert_eq!(decoded.split_points(), original.split_points()); + assert_eq!( + decoded.ordering()[0].options, + original.ordering()[0].options + ); + assert_eq!(decoded.samples(), original.samples()); + assert_eq!(decoded.max_partition_count(), 6); + + Ok(()) + } + + #[test] + fn legacy_range_partitioning_payload_remains_exact() -> Result<()> { + let partitioning = sampled_partitioning()?; + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let mut encoded = partitioning.try_to_proto(&encode_ctx)?; + let Some(protobuf::partitioning::PartitionMethod::Range(encoded_range)) = + encoded.partition_method.as_mut() + else { + panic!("expected range partitioning"); + }; + encoded_range.sample_point.clear(); + encoded_range.partition_count = 0; + + let decoded = decode(&encoded)?; + let Partitioning::Range(decoded) = decoded else { + panic!("expected range partitioning"); + }; + assert_eq!(decoded.partition_count(), 3); + assert_eq!(decoded.max_partition_count(), 3); + assert_eq!(decoded.split_points().len(), 2); + + Ok(()) + } + + #[test] + fn sampled_range_partitioning_rejects_inconsistent_effective_points() -> Result<()> { + let partitioning = sampled_partitioning()?; + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let mut encoded = partitioning.try_to_proto(&encode_ctx)?; + let Some(protobuf::partitioning::PartitionMethod::Range(encoded_range)) = + encoded.partition_method.as_mut() + else { + panic!("expected range partitioning"); + }; + encoded_range.split_point.pop(); + + let error = decode(&encoded).unwrap_err().to_string(); + assert!( + error.contains("effective split points do not match"), + "{error}" + ); + + Ok(()) + } +} + /// Partition counts are `usize` in memory and `u64` on the wire, so every /// counted [`Partitioning`] variant crosses a width boundary in both /// directions. These pin that neither crossing wraps or panics. diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 776e3c3bf14bd..04045db3fa0a4 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -952,6 +952,9 @@ struct RepartitionRequirementStatus { /// Per-child state while enforcing a parent's distribution requirements. struct DistributionChildState { + /// Scaling a native range layout must not lose its preference over a newly + /// introduced hash exchange when choosing a co-partitioning reference. + scaled_native_range: bool, context: DistributionContext, required_input_ordering: Option, maintains_input_order: bool, @@ -1154,7 +1157,8 @@ fn enforce_distribution_relationships( let partitioning = plan.output_partitioning(); match partitioning { Partitioning::Range(_) | Partitioning::Hash(_, _) => { - let is_native = !plan.is::(); + let is_native = + !plan.is::() || child.scaled_native_range; satisfied_children.push((i, partitioning.clone(), is_native)); } _ => {} @@ -1194,11 +1198,9 @@ fn enforce_distribution_relationships( }) .collect(); - // Only select a reference candidate if there is a unique, strictly - // larger winner (`size_a > size_b`). If candidates have equal or - // incomparable sizes (e.g. non-overlapping metrics), return None - // so the optimizer avoids arbitrary tie-breaking and falls back to - // standard distribution. + // Prefer a unique, strictly larger winner (`size_a > size_b`). + // Otherwise avoid arbitrary tie-breaking, except for the new + // preserved-range case below. candidates .iter() .find(|(size_a, idx_a, _)| { @@ -1282,7 +1284,9 @@ fn enforce_distribution_relationships( .map(|s| s.is_satisfied()) .unwrap_or(false) } - (Partitioning::Range(r1), Partitioning::Range(r2)) if r1 == r2 => true, + (Partitioning::Range(r1), Partitioning::Range(r2)) => { + r1.has_same_layout(r2) + } _ => false, }; @@ -1312,6 +1316,7 @@ fn enforce_distribution_relationships( let plan = Arc::new(repartition) as _; children[child_idx].context = DistributionContext::new(plan, true, vec![original_child]); + children[child_idx].scaled_native_range = false; repartitioned_for_relationship[child_idx] = true; changed = true; } @@ -1554,13 +1559,13 @@ pub fn ensure_distribution_with_stats( child.plan = new_child; } + let mut scaled_native_range = false; // Satisfy the distribution requirement if it is unmet. match &requirement { Distribution::SinglePartition => { child = add_merge_on_top(child); } - Distribution::HashPartitioned(exprs) - | Distribution::KeyPartitioned(exprs) => { + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) => { let child_partitions = child.plan.output_partitioning().partition_count(); let partitioning_satisfied = input_distributions @@ -1578,21 +1583,56 @@ pub fn ensure_distribution_with_stats( && target_partitions > child_partitions; // When subset satisfaction is enabled, preserve an - // already-satisfying partitioning. Otherwise, hash + // already-satisfying partitioning. Otherwise, // repartition may also increase parallelism. - let needs_hash_repartition = if allow_subset_satisfy_partitioning { + let needs_repartition = if allow_subset_satisfy_partitioning { !partitioning_satisfied } else { !partitioning_satisfied || (target_partitions > child_partitions && !preserve_satisfying_file_partitioning) }; - let should_add_hash_repartition = - hash_necessary && needs_hash_repartition; + let should_add_repartition = hash_necessary && needs_repartition; // See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background - // When inserting hash is necessary to satisfy hash requirement, insert hash repartition. - if should_add_hash_repartition { + // Enforce unmet requirements, or increase parallelism when beneficial. + if should_add_repartition { + let partitioning = match child.plan.output_partitioning() { + Partitioning::Range(range) if partitioning_satisfied => { + match range.scale(target_partitions) { + Some(range) => { + let scaled = Partitioning::Range(range); + // A single partition satisfies any key requirement, + // but scaling it must still use compatible keys. + if scaled + .satisfaction( + &requirement, + child.plan.equivalence_properties(), + false, + ) + .is_satisfied() + { + scaled_native_range = + !child.plan.is::(); + scaled + } else { + requirement + .clone() + .create_partitioning(target_partitions) + } + } + // Insufficient samples are expected. Preserve + // main's policy by falling back to key + // repartitioning at the requested parallelism. + None => requirement + .clone() + .create_partitioning(target_partitions), + } + } + _ => { + requirement.clone().create_partitioning(target_partitions) + } + }; // When there is an existing ordering, we preserve ordering during // repartition. This will be rolled back in the future if any of the // following conditions is true: @@ -1600,8 +1640,6 @@ pub fn ensure_distribution_with_stats( // requirements. // - Usage of order preserving variants is not desirable (per the flag // `config.optimizer.prefer_existing_sort`). - let partitioning = Distribution::KeyPartitioned(exprs.to_vec()) - .create_partitioning(target_partitions); let repartition = RepartitionExec::try_new( Arc::clone(&child.plan), partitioning, @@ -1621,6 +1659,7 @@ pub fn ensure_distribution_with_stats( } Ok(DistributionChildState { + scaled_native_range, context: child, required_input_ordering, maintains_input_order: maintains, @@ -1644,6 +1683,7 @@ pub fn ensure_distribution_with_stats( .into_iter() .map( |DistributionChildState { + scaled_native_range: _, mut context, required_input_ordering, maintains_input_order, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 30b2972312f8c..16f61d740985f 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -152,10 +152,11 @@ pub fn adjust_right_output_partitioning( "Offsetting range partitioning produced an empty ordering" ) })?; - Partitioning::Range(RangePartitioning::new( + Partitioning::Range(RangePartitioning::try_new_with_samples( ordering, - range.split_points().to_vec(), - )) + range.samples().to_vec(), + range.partition_count(), + )?) } result => result.clone(), }; @@ -4409,8 +4410,20 @@ mod tests { ScalarValue::Int32(Some(20)), ScalarValue::Int32(Some(50)), ]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(30)), + ScalarValue::Int32(Some(40)), + ]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(40)), + ScalarValue::Int32(Some(30)), + ]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(50)), + ScalarValue::Int32(Some(20)), + ]), ]; - let range = RangePartitioning::try_new( + let range = RangePartitioning::try_new_with_samples( LexOrdering::new([ PhysicalSortExpr::new( Arc::new(Column::new("a", 0)), @@ -4423,10 +4436,16 @@ mod tests { ]) .unwrap(), split_points.clone(), + 3, )?; let adjusted = adjust_right_output_partitioning(&Partitioning::Range(range), 3)?; - let expected = Partitioning::Range(RangePartitioning::new( + let Partitioning::Range(adjusted_range) = &adjusted else { + panic!("expected range partitioning"); + }; + assert_eq!(adjusted_range.max_partition_count(), 6); + assert_eq!(adjusted_range.samples(), split_points); + let expected = Partitioning::Range(RangePartitioning::try_new_with_samples( LexOrdering::new([ PhysicalSortExpr::new( Arc::new(Column::new("a", 3)), @@ -4439,7 +4458,8 @@ mod tests { ]) .unwrap(), split_points, - )); + 3, + )?); assert_eq!(adjusted, expected); Ok(()) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 794888b4ea827..d2d8ec38dacc4 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1989,9 +1989,10 @@ impl ExecutionPlan for RepartitionExec { ); }; - Partitioning::Range(RangePartitioning::try_new( + Partitioning::Range(RangePartitioning::try_new_with_samples( ordering, - range_partitioning.split_points().to_vec(), + range_partitioning.samples().to_vec(), + range_partitioning.partition_count(), )?) } others => others.clone(), @@ -2052,9 +2053,18 @@ impl ExecutionPlan for RepartitionExec { new_properties.partitioning = match new_properties.partitioning { RoundRobinBatch(_) => RoundRobinBatch(target_partitions), Hash(hash, _) => Hash(hash, target_partitions), - Range(_) => { - // Number of partitions is constrained by the split points and cannot be changed - return Ok(None); + Range(range) => { + let Some(range) = range.scale(target_partitions) else { + return Ok(None); + }; + // A different layout needs its own channels, router, and metrics. + let mut repartition = + Self::try_new(Arc::clone(&self.input), Range(range))?; + if self.preserve_order { + repartition = repartition.with_preserve_order(); + } + repartition.batch_size = self.batch_size; + return Ok(Some(Arc::new(repartition))); } UnknownPartitioning(_) => UnknownPartitioning(target_partitions), }; @@ -3083,6 +3093,112 @@ mod tests { Ok(()) } + #[tokio::test] + async fn range_repartitioned_scales_with_fresh_execution_state() -> Result<()> { + let schema = test_schema(false); + let ordering = + LexOrdering::new([PhysicalSortExpr::new_default(col("c0", &schema)?)]) + .unwrap(); + let samples = [10, 20, 30] + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::UInt32(Some(value))])) + .collect::>(); + let partitions = [vec![5, 15, 25, 35], vec![6, 16, 26, 36]] + .into_iter() + .map(|values| -> Result<_> { + Ok(vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(values))], + )?]) + }) + .collect::>>()?; + + for preserve_order in [false, true] { + let source = TestMemoryExec::try_new(&partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering.clone()])?; + let source = Arc::new(TestMemoryExec::update_cache(&Arc::new(source))); + let mut exec = Arc::new( + RepartitionExec::try_new( + source, + Partitioning::Range(RangePartitioning::try_new_with_samples( + ordering.clone(), + samples.clone(), + 2, + )?), + )? + .with_batch_size(2)?, + ); + if preserve_order { + exec = Arc::new(Arc::unwrap_or_clone(exec).with_preserve_order()); + } + let context = Arc::new(TaskContext::default()); + + // Initialize the old exchange before resizing: its channels and router + // must not be reused for a different set of output boundaries. + let initial = crate::collect_partitioned( + Arc::::clone(&exec), + Arc::clone(&context), + ) + .await?; + assert_eq!( + initial + .iter() + .map(|p| partition_row_count(p)) + .sum::(), + 8 + ); + + for target in [4, 1, 4] { + let scaled = exec + .repartitioned(target, &ConfigOptions::default())? + .expect("retained samples support the requested partition count"); + let repartition = scaled.downcast_ref::().unwrap(); + let range = expect_range_partitioning(repartition.partitioning()); + assert_eq!(range.partition_count(), target); + assert_eq!(range.samples(), samples); + assert_eq!(range.ordering(), &ordering); + assert_eq!(repartition.preserve_order, preserve_order); + assert_eq!(repartition.batch_size, Some(2)); + assert_eq!( + repartition.properties().output_ordering(), + exec.properties().output_ordering() + ); + assert!(!Arc::ptr_eq(&repartition.state, &exec.state)); + assert!(Arc::ptr_eq(&repartition.input, &exec.input)); + assert_eq!(repartition.metrics().unwrap().output_rows(), None); + + let output = + crate::collect_partitioned(Arc::clone(&scaled), Arc::clone(&context)) + .await?; + assert_eq!(output.len(), target); + for (index, batches) in output.iter().enumerate() { + let mut values = collect_partition_u32_values(batches); + if !preserve_order { + values.sort_unstable(); + } + let expected = if target == 1 { + vec![5, 6, 15, 16, 25, 26, 35, 36] + } else { + vec![index as u32 * 10 + 5, index as u32 * 10 + 6] + }; + assert_eq!( + values, + expected.into_iter().map(Some).collect::>() + ); + } + assert_eq!(repartition.metrics().unwrap().output_rows(), Some(8)); + exec = Arc::new(repartition.clone()); + } + for unsupported in [0, 5] { + assert!( + exec.repartitioned(unsupported, &ConfigOptions::default())? + .is_none() + ); + } + } + Ok(()) + } + #[tokio::test] async fn range_repartition_routes_rows_desc() -> Result<()> { let schema = test_schema(false); @@ -3448,9 +3564,18 @@ mod tests { Field::new("region", DataType::Utf8, false), Field::new("payload", DataType::UInt32, false), ])); + let ordering = + LexOrdering::new([PhysicalSortExpr::new_default(col("id", &schema)?)]) + .expect("non-empty ordering"); + let samples = [10, 20, 30, 40, 50] + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::UInt32(Some(value))])) + .collect(); let repartition = Arc::new(RepartitionExec::try_new( Arc::new(EmptyExec::new(Arc::clone(&schema))), - range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + Partitioning::Range(RangePartitioning::try_new_with_samples( + ordering, samples, 2, + )?), )?); let projection = @@ -3466,9 +3591,10 @@ mod tests { assert!(swapped_repartition.input().is::()); let range = expect_range_partitioning(swapped_repartition.partitioning()); assert_eq!(range.ordering()[0].to_string(), "id@1 ASC"); + assert_eq!(range.max_partition_count(), 6); assert_eq!( range.split_points(), - &[SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])] + &[SplitPoint::new(vec![ScalarValue::UInt32(Some(30))])] ); Ok(()) @@ -5003,12 +5129,11 @@ mod test { })?; assert_eq!(expressions, ["c0@0"]); - // Range partition count is fixed by split points, so repartitioned() - // cannot change it to an arbitrary target. + // Scaling cannot exceed the retained sample capacity. let result = exec.repartitioned(10, &Default::default())?; assert!( result.is_none(), - "range repartitioning should not support changing partition count" + "range repartitioning should reject counts above sample capacity" ); Ok(()) } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 1f5985fe3f7d2..b764098e7cbb8 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -1877,6 +1877,32 @@ mod tests { Ok(()) } + #[test] + fn test_interleave_rejects_different_range_scaling_capacity() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let ordering = [PhysicalSortExpr::new_default(col("a", &schema)?)].into(); + let samples = (10..=90) + .step_by(10) + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(); + let sampled = RangePartitioning::try_new_with_samples(ordering, samples, 4)?; + let source = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?); + let sampled: Arc = Arc::new(RepartitionExec::try_new( + source, + Partitioning::Range(sampled), + )?); + let exact = make_range_exec(&schema, vec![30, 50, 70], SortOptions::default())?; + assert!(can_interleave([&sampled, &sampled].into_iter())); + for inputs in [ + vec![Arc::clone(&sampled), Arc::clone(&exact)], + vec![exact, sampled], + ] { + assert!(!can_interleave(inputs.iter())); + assert!(InterleaveExec::try_new(inputs).is_err()); + } + Ok(()) + } + #[test] fn test_can_interleave_matrix() -> Result<()> { let name_column = "name"; diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 15dc272fabca3..dd551ead1f56f 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1620,7 +1620,12 @@ message PhysicalHashRepartition { message PhysicalRangePartitioning { repeated PhysicalSortExprNode sort_expr = 1; + // Effective split points. Kept for compatibility with older readers. repeated PhysicalRangeSplitPoint split_point = 2; + // Maximum-resolution sample points used to derive effective split points. + repeated PhysicalRangeSplitPoint sample_point = 3; + // Zero in legacy payloads means split_point.len() + 1. + uint64 partition_count = 4; } message PhysicalRangeSplitPoint { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index cada03c9a6d0f..744429ae0aea3 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -22092,6 +22092,12 @@ impl serde::Serialize for PhysicalRangePartitioning { if !self.split_point.is_empty() { len += 1; } + if !self.sample_point.is_empty() { + len += 1; + } + if self.partition_count != 0 { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalRangePartitioning", len)?; if !self.sort_expr.is_empty() { struct_ser.serialize_field("sortExpr", &self.sort_expr)?; @@ -22099,6 +22105,14 @@ impl serde::Serialize for PhysicalRangePartitioning { if !self.split_point.is_empty() { struct_ser.serialize_field("splitPoint", &self.split_point)?; } + if !self.sample_point.is_empty() { + struct_ser.serialize_field("samplePoint", &self.sample_point)?; + } + if self.partition_count != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("partitionCount", ToString::to_string(&self.partition_count).as_str())?; + } struct_ser.end() } } @@ -22113,12 +22127,18 @@ impl<'de> serde::Deserialize<'de> for PhysicalRangePartitioning { "sortExpr", "split_point", "splitPoint", + "sample_point", + "samplePoint", + "partition_count", + "partitionCount", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { SortExpr, SplitPoint, + SamplePoint, + PartitionCount, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -22142,6 +22162,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalRangePartitioning { match value { "sortExpr" | "sort_expr" => Ok(GeneratedField::SortExpr), "splitPoint" | "split_point" => Ok(GeneratedField::SplitPoint), + "samplePoint" | "sample_point" => Ok(GeneratedField::SamplePoint), + "partitionCount" | "partition_count" => Ok(GeneratedField::PartitionCount), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -22163,6 +22185,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalRangePartitioning { { let mut sort_expr__ = None; let mut split_point__ = None; + let mut sample_point__ = None; + let mut partition_count__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::SortExpr => { @@ -22177,11 +22201,27 @@ impl<'de> serde::Deserialize<'de> for PhysicalRangePartitioning { } split_point__ = Some(map_.next_value()?); } + GeneratedField::SamplePoint => { + if sample_point__.is_some() { + return Err(serde::de::Error::duplicate_field("samplePoint")); + } + sample_point__ = Some(map_.next_value()?); + } + GeneratedField::PartitionCount => { + if partition_count__.is_some() { + return Err(serde::de::Error::duplicate_field("partitionCount")); + } + partition_count__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } } } Ok(PhysicalRangePartitioning { sort_expr: sort_expr__.unwrap_or_default(), split_point: split_point__.unwrap_or_default(), + sample_point: sample_point__.unwrap_or_default(), + partition_count: partition_count__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 4bb4af1e8532a..a1272dc242487 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2432,8 +2432,15 @@ pub struct PhysicalHashRepartition { pub struct PhysicalRangePartitioning { #[prost(message, repeated, tag = "1")] pub sort_expr: ::prost::alloc::vec::Vec, + /// Effective split points. Kept for compatibility with older readers. #[prost(message, repeated, tag = "2")] pub split_point: ::prost::alloc::vec::Vec, + /// Maximum-resolution sample points used to derive effective split points. + #[prost(message, repeated, tag = "3")] + pub sample_point: ::prost::alloc::vec::Vec, + /// Zero in legacy payloads means split_point.len() + 1. + #[prost(uint64, tag = "4")] + pub partition_count: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalRangeSplitPoint { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 940c732f81318..2b349035b3292 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -301,10 +301,14 @@ mod file_scan_config_serde { Column::new("value", 0), ))]) .expect("single expression ordering"); - Partitioning::Range(RangePartitioning::new( - ordering, - vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], - )) + Partitioning::Range( + RangePartitioning::try_new_with_samples( + ordering, + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + 2, + ) + .unwrap(), + ) } fn decode_source(conf: &protobuf::FileScanExecConf) -> Result> { diff --git a/datafusion/proto/tests/cases/plans/misc.rs b/datafusion/proto/tests/cases/plans/misc.rs index a8b5102da3fdc..65c99a3fca7ba 100644 --- a/datafusion/proto/tests/cases/plans/misc.rs +++ b/datafusion/proto/tests/cases/plans/misc.rs @@ -449,10 +449,12 @@ fn roundtrip_repartition_preserve_order() -> Result<()> { fn roundtrip_range_partitioning() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let range_partitioning = Partitioning::Range(RangePartitioning::new( - [PhysicalSortExpr::new_default(col("a", &schema)?)].into(), - vec![SplitPoint::new(vec![ScalarValue::Int64(Some(10))])], - )); + let range_partitioning = + Partitioning::Range(RangePartitioning::try_new_with_samples( + [PhysicalSortExpr::new_default(col("a", &schema)?)].into(), + vec![SplitPoint::new(vec![ScalarValue::Int64(Some(10))])], + 2, + )?); // RepartitionExec is used only to carry the partitioning through proto. // Executing range repartitioning is intentionally unsupported. let repartition = RepartitionExec::try_new(input, range_partitioning)?; diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index c4897357ee00b..04d002e1e623c 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -1100,13 +1100,15 @@ fn roundtrip_parquet_exec_range_output_partitioning() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, false)])); let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); - let output_partitioning = Partitioning::Range(RangePartitioning::new( - LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new( - "col", 0, - )))]) - .unwrap(), - vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], - )); + let output_partitioning = + Partitioning::Range(RangePartitioning::try_new_with_samples( + LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new( + "col", 0, + )))]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + 2, + )?); let scan_config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) .with_file_groups(vec![ diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 01c6799e96453..f944daf322ff1 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -119,6 +119,25 @@ serialized physical plan reconstructed with local type identities. See the pattern. `ForeignSession::query_planner`, `optimize`, and `physical_optimizers` continue to forward to the owning session across the FFI boundary. +### Range partitioning supports retained samples and scaling + +`RangePartitioning::new` is deprecated. Prefer the validated +`RangePartitioning::try_new_with_samples(ordering, samples, partition_count)`. +For exact boundaries, pass `split_points.len() + 1` as the count, or keep using +the validated `try_new` constructor. Retain additional samples when future +scaling above that count is required; scaling does not create new sample values. + +`scale` returns `Option`. It returns `None` when the target +partition count is zero or exceeds the retained sample capacity; callers can +then retain the existing layout or choose another partitioning. Equality includes +retained samples; use `has_same_layout` to compare the current key ordering and +effective boundaries. + +`FFI_RangePartitioning` has a new layout; rebuild FFI consumers against the +compatible release. The generated `PhysicalRangePartitioning` protobuf struct +also gains fields, affecting exhaustive struct literals. Older protobuf payloads +remain readable, and effective split points remain available to older readers. + ### `GroupColumn` now requires `values_preserving` Custom implementations of the public `GroupColumn` trait must implement