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 fbc3e83ba49fc..9f166ec0913b1 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 4067d7eb49b2a..10b33f83bea70 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -29,7 +29,7 @@ mod tests { use datafusion_ffi::tests::utils::get_module; use datafusion_physical_plan::execution_plan::InvariantLevel; use datafusion_physical_plan::{ - ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, + ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, }; use std::sync::Arc; @@ -68,6 +68,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 e1ca2df9224c3..0773c10668db3 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -153,15 +153,21 @@ 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` define the maximum-resolution boundaries. +/// - `partition_count` selects how many ranges to derive from those samples. /// /// 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 @@ -201,12 +207,16 @@ impl Display for Partitioning { /// NOTE: Optimizer and execution behavior for this partitioning is intentionally /// not implemented and will be introduced incrementally. See /// . -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub struct RangePartitioning { /// Ordered partitioning key. ordering: LexOrdering, - /// Boundaries between adjacent partitions. - split_points: Vec, + /// Maximum-resolution boundaries used to derive split points. + samples: Arc<[SplitPoint]>, + /// Effective boundaries for the current partition count. + split_points: Arc<[SplitPoint]>, + /// Number of effective partitions. + partition_count: usize, } impl RangePartitioning { @@ -215,9 +225,13 @@ impl RangePartitioning { /// Use [`Self::try_new`] to validate the contract documented on /// [`RangePartitioning`]. pub fn new(ordering: LexOrdering, split_points: Vec) -> Self { + let partition_count = split_points.len() + 1; + let split_points: Arc<[SplitPoint]> = Arc::from(split_points); Self { ordering, + samples: Arc::clone(&split_points), split_points, + partition_count, } } @@ -234,19 +248,75 @@ impl RangePartitioning { Ok(Self::new(ordering, split_points)) } + /// Creates sample-backed range partitioning and validates the sample shape, + /// ordering, and target partition count. + /// + /// `partition_count` must be at least one and no larger than + /// `samples.len() + 1`. When it is smaller than that maximum, the samples + /// are evenly down-sampled to derive the effective split points. + pub fn try_new_with_samples( + ordering: LexOrdering, + samples: Vec, + partition_count: usize, + ) -> Result { + validate_range_split_points( + &samples, + &ordering + .iter() + .map(|sort_expr| sort_expr.options) + .collect::>(), + )?; + 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, + partition_count, + }) + } + /// Returns the ordering that defines the range key. pub fn ordering(&self) -> &LexOrdering { &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 } /// Returns the number of partitions. pub fn partition_count(&self) -> usize { - self.split_points.len() + 1 + self.partition_count + } + + /// Returns the largest partition count supported by the stored samples. + pub fn max_partition_count(&self) -> usize { + self.samples.len() + 1 + } + + /// Returns this range partitioning scaled to `target_partitions`. + /// + /// Scaling retains the original samples, so a range partitioning that was + /// scaled down can later be scaled back up to [`Self::max_partition_count`]. + pub fn scale(&self, target_partitions: usize) -> Result { + validate_range_partition_count(target_partitions, self.max_partition_count())?; + if target_partitions == self.partition_count { + return Ok(self.clone()); + } + Ok(Self { + ordering: self.ordering.clone(), + samples: Arc::clone(&self.samples), + split_points: downsample_split_points(&self.samples, target_partitions), + partition_count: target_partitions, + }) } /// Calculates the range partitioning after applying the given projection. @@ -279,7 +349,9 @@ impl RangePartitioning { Some(Self { ordering, - split_points: self.split_points.clone(), + samples: Arc::clone(&self.samples), + split_points: Arc::clone(&self.split_points), + partition_count: self.partition_count, }) } @@ -323,7 +395,18 @@ impl RangePartitioning { }) .collect::>(), )?; - Self::try_new(new_ordering, self.split_points.clone()).ok() + Some(Self { + ordering: new_ordering, + samples: Arc::clone(&self.samples), + split_points: Arc::clone(&self.split_points), + partition_count: self.partition_count, + }) + } +} + +impl PartialEq for RangePartitioning { + fn eq(&self, other: &Self) -> bool { + self.ordering == other.ordering && self.split_points == other.split_points } } @@ -332,12 +415,54 @@ 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 datafusion_common::plan_err!( + "Range partitioning partition count must be at least 1" + ); + } + if partition_count > max_partition_count { + return datafusion_common::plan_err!( + "Range partitioning partition count {partition_count} exceeds maximum {max_partition_count}" + ); } + Ok(()) } fn format_range_split_points(split_points: &[SplitPoint]) -> String { @@ -630,22 +755,27 @@ impl Partitioning { } Partitioning::Range(range) => { let sort_expr = sort_exprs_try_to_proto(range.ordering().iter(), ctx)?; - let split_point = range - .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 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())?, }, ) } @@ -704,19 +834,44 @@ impl Partitioning { "Range partitioning ordering must not contain duplicate expressions" ); } - let split_points = range - .split_point - .iter() - .map(|split_point| { - let values = split_point - .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)) @@ -903,17 +1058,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( @@ -1182,6 +1326,100 @@ 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)?; + 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())?; + 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)?; + let error = range.scale(4).unwrap_err().to_string(); + assert!(error.contains("exceeds maximum 3"), "{error}"); + + Ok(()) + } + + #[test] + fn test_range_partitioning_equality_uses_effective_split_points() -> 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, exact); + assert_eq!(Partitioning::Range(sampled), Partitioning::Range(exact)); + + Ok(()) + } + #[test] fn test_range_partitioning_try_new_validates_split_points() -> Result<()> { let fixture = PartitioningTestFixture::int64(&["a", "b"])?; @@ -1231,18 +1469,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)?.split_points().len(), 3); let drop_b_mapping = ProjectionMapping::from_indices(&[0], &fixture.schema)?; let projected = @@ -1354,19 +1603,24 @@ 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)?.partition_count(), 5); // Adapting to col_b (different type Int64) fails assert!(range.adapt(&[fixture.col(1)], &fixture.schema).is_none()); @@ -1568,6 +1822,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-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 798f22b1d5770..1d75162737456 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -153,10 +153,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(), }; @@ -4430,8 +4431,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)), @@ -4444,9 +4457,15 @@ mod tests { ]) .unwrap(), split_points.clone(), + 3, )?; let adjusted = adjust_right_output_partitioning(&Partitioning::Range(range), 3)?; + 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::new( LexOrdering::new([ PhysicalSortExpr::new( @@ -4459,7 +4478,16 @@ mod tests { ), ]) .unwrap(), - split_points, + vec![ + SplitPoint::new(vec![ + ScalarValue::Int32(Some(20)), + ScalarValue::Int32(Some(50)), + ]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(40)), + ScalarValue::Int32(Some(30)), + ]), + ], )); assert_eq!(adjusted, expected); diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 7822c6facd7c9..62ef21fc2a567 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1877,9 +1877,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(), @@ -3021,9 +3022,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 = @@ -3039,9 +3049,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(()) diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 8486410d07a2b..0fd05a64997fb 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1574,7 +1574,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 f568a6d853d56..07875f01b0193 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -21267,6 +21267,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)?; @@ -21274,6 +21280,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() } } @@ -21288,12 +21302,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 @@ -21317,6 +21337,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)), } } @@ -21338,6 +21360,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 => { @@ -21352,11 +21376,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 f5c7c01e7192f..65411b964880c 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2373,8 +2373,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 {