This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22207-786d56f0e4aedbf8f859f5e40e9b9619d0ac1bb1 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 7a6b0626da7a733d94821b8d9d6011eb85d96593 Author: Gene Bordegaray <[email protected]> AuthorDate: Wed May 27 13:48:01 2026 -0400 Add Physical `Partitioning::Range` enum variant (#22207) ## Which issue does this PR close? - First mechanical PR for `ExprPartitioning` as described in thread: #21992. ## Rationale for this change DataFusion currently cannot truthfully represent range-partitioned physical data. Some sources may be range partitioned, but have to advertise another partitioning shape or fall back to unknown partitioning. This PR introduces the metadata shape for range partitioning without implementing optimizer or execution behavior yet. The goal is to establish the public representation first, then implement planning, compatibility, and execution behavior incrementally in follow-up PRs. ## What changes are included in this PR? - Adds `Partitioning::Range(RangePartitioning)`. - Adds range metadata types: - `RangePartitioning` - `RangePartition` - `RangeInterval` - `RangeBound` - Adds proto serialization/deserialization. - Adds `not_impl_err!` handling for range partitioning at call sites. - Preserves range partitioning through projection only when all partition expressions can be projected, otherwise `UnknownPartitioning`. ## Are these changes tested? Yes. ## Are there any user-facing changes? Yes. This adds new public physical partitioning API and proto for range partitioning. --------- Co-authored-by: Andrew Lamb <[email protected]> --- datafusion/ffi/src/physical_expr/partitioning.rs | 5 + datafusion/physical-expr/src/lib.rs | 2 +- datafusion/physical-expr/src/partitioning.rs | 828 ++++++++++++++++----- datafusion/physical-plan/src/joins/utils.rs | 7 + datafusion/physical-plan/src/lib.rs | 2 +- datafusion/physical-plan/src/repartition/mod.rs | 76 ++ datafusion/physical-plan/src/sorts/sort.rs | 3 +- .../src/sorts/sort_preserving_merge.rs | 6 +- datafusion/proto-models/proto/datafusion.proto | 14 +- datafusion/proto-models/src/generated/pbjson.rs | 214 ++++++ datafusion/proto-models/src/generated/prost.rs | 20 +- datafusion/proto/src/physical_plan/from_proto.rs | 59 +- datafusion/proto/src/physical_plan/to_proto.rs | 43 +- .../proto/tests/cases/roundtrip_physical_plan.rs | 18 +- 14 files changed, 1084 insertions(+), 213 deletions(-) diff --git a/datafusion/ffi/src/physical_expr/partitioning.rs b/datafusion/ffi/src/physical_expr/partitioning.rs index 434b6a097e..eec437639e 100644 --- a/datafusion/ffi/src/physical_expr/partitioning.rs +++ b/datafusion/ffi/src/physical_expr/partitioning.rs @@ -45,6 +45,11 @@ impl From<&Partitioning> for FFI_Partitioning { .collect(); Self::Hash(exprs, *size) } + // FFI does not yet expose range partition metadata. + // See https://github.com/apache/datafusion/issues/22394 + Partitioning::Range(range) => { + Self::UnknownPartitioning(range.partition_count()) + } Partitioning::UnknownPartitioning(size) => Self::UnknownPartitioning(*size), } } diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index e67046987b..c82d1c64dd 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -62,7 +62,7 @@ pub use equivalence::{ AcrossPartitions, ConstExpr, EquivalenceProperties, calculate_union, }; pub use expressions::{DynamicFilterTracker, DynamicFilterTracking}; -pub use partitioning::{Distribution, Partitioning}; +pub use partitioning::{Distribution, Partitioning, RangePartitioning, SplitPoint}; pub use physical_expr::{ add_offset_to_expr, add_offset_to_physical_sort_exprs, create_lex_ordering, create_ordering, create_physical_sort_expr, create_physical_sort_exprs, diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index d24c60b63e..bb46b8a957 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -21,7 +21,10 @@ use crate::{ EquivalenceProperties, PhysicalExpr, equivalence::ProjectionMapping, expressions::UnKnownColumn, physical_exprs_equal, }; +use datafusion_common::{Result, ScalarValue, plan_err}; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; +use std::cmp::Ordering; use std::fmt; use std::fmt::Display; use std::sync::Arc; @@ -117,6 +120,8 @@ pub enum Partitioning { /// Allocate rows based on a hash of one of more expressions and the specified number of /// partitions Hash(Vec<Arc<dyn PhysicalExpr>>, usize), + /// Partition rows by source-declared ranges + Range(RangePartitioning), /// Unknown partitioning scheme with a known number of partitions UnknownPartitioning(usize), } @@ -133,6 +138,7 @@ impl Display for Partitioning { .join(", "); write!(f, "Hash([{phy_exprs_str}], {size})") } + Partitioning::Range(range) => write!(f, "{range}"), Partitioning::UnknownPartitioning(size) => { write!(f, "UnknownPartitioning({size})") } @@ -140,6 +146,271 @@ impl Display for Partitioning { } } +/// Physical range partitioning. +/// +/// [`RangePartitioning`] describes an ordered key space with split points. +/// +/// - `ordering` defines the partitioning key and ordering. +/// - `split_points` define the boundaries between adjacent partitions. +/// +/// 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. +/// +/// `N` split points define `N + 1` partitions: +/// +/// ```text +/// partition 0: key < split_points[0] +/// partition 1: split_points[0] <= key < split_points[1] +/// ... +/// partition N - 1: split_points[N - 2] <= key < split_points[N - 1] +/// partition N: split_points[N - 1] <= key +/// ``` +/// +/// Values equal to split point `i` belong to partition `i + 1`, so interior +/// partitions are lower-inclusive and upper-exclusive. +/// +/// For a single range key: +/// +/// ```text +/// ordering = [date ASC NULLS LAST] +/// split_points = [ +/// (2022-01-01), +/// (2023-01-01), +/// ] +/// +/// partition 0: date before 2022-01-01 +/// partition 1: date between 2022-01-01 (inclusive) and 2023-01-01 (exclusive) +/// partition 2: date at/after 2023-01-01 +/// ``` +/// +/// The same model extends to compound keys. +/// For `ordering = [time ASC, city ASC]`, split points are ordered +/// lexicographically by `(time, city)`: +/// +/// ```text +/// ordering = [time ASC NULLS LAST, city ASC NULLS LAST] +/// split_points = [ +/// (2022, Allston), +/// (2023, Allston), +/// ] +/// +/// partition 0: keys before (2022, Allston) +/// partition 1: keys between (2022, Allston) and (2023, Allston) +/// 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 +/// <https://github.com/apache/datafusion/issues/22395>. +#[derive(Debug, Clone, PartialEq)] +pub struct RangePartitioning { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Boundaries between adjacent partitions. + split_points: Vec<SplitPoint>, +} + +/// A boundary between adjacent range partitions. +/// +/// A split point is a tuple with one [`ScalarValue`] per sort expression in the +/// parent [`RangePartitioning`] ordering. +#[derive(Debug, Clone, PartialEq)] +pub struct SplitPoint { + values: Vec<ScalarValue>, +} + +impl SplitPoint { + /// Creates a new split point from its tuple values. + pub fn new(values: Vec<ScalarValue>) -> Self { + Self { values } + } + + /// Returns the tuple values for this split point. + pub fn values(&self) -> &[ScalarValue] { + &self.values + } +} + +impl Display for SplitPoint { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let values = self + .values + .iter() + .map(ToString::to_string) + .collect::<Vec<_>>() + .join(", "); + write!(f, "({values})") + } +} + +impl RangePartitioning { + /// Creates range partitioning metadata without validating split points. + /// + /// Use [`Self::try_new`] to validate the contract documented on + /// [`RangePartitioning`]. + pub fn new(ordering: LexOrdering, split_points: Vec<SplitPoint>) -> Self { + Self { + ordering, + split_points, + } + } + + /// Creates range partitioning metadata and validates split point shape and + /// ordering. + pub fn try_new(ordering: LexOrdering, split_points: Vec<SplitPoint>) -> Result<Self> { + validate_range_split_points(&ordering, &split_points)?; + Ok(Self::new(ordering, split_points)) + } + + /// Returns the ordering that defines the range key. + pub fn ordering(&self) -> &LexOrdering { + &self.ordering + } + + /// Returns the ordered 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 + } + + fn project( + &self, + mapping: &ProjectionMapping, + input_eq_properties: &EquivalenceProperties, + ) -> Option<Self> { + let exprs = self + .ordering + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::<Vec<_>>(); + let projected_exprs = input_eq_properties + .project_expressions(&exprs, mapping) + .collect::<Option<Vec<_>>>()?; + let sort_exprs = self + .ordering + .iter() + .zip(projected_exprs) + .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, sort_expr.options)) + .collect::<Vec<_>>(); + let ordering = LexOrdering::new(sort_exprs)?; + if ordering.len() != self.ordering.len() { + return None; + } + + Some(Self { + ordering, + split_points: self.split_points.clone(), + }) + } +} + +impl Display for RangePartitioning { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let split_points = format_range_split_points(&self.split_points); + write!( + f, + "Range([{}], [{}], {})", + self.ordering, + split_points, + self.partition_count() + ) + } +} + +fn format_range_split_points(split_points: &[SplitPoint]) -> String { + split_points + .iter() + .map(ToString::to_string) + .collect::<Vec<_>>() + .join(", ") +} + +fn validate_range_split_points( + ordering: &LexOrdering, + split_points: &[SplitPoint], +) -> Result<()> { + let width = ordering.len(); + for (idx, split_point) in split_points.iter().enumerate() { + let split_point_width = split_point.values.len(); + if split_point_width != width { + return plan_err!( + "Range partitioning split point {idx} has width {split_point_width}, but ordering has width {width}" + ); + } + } + + for (idx, split_points) in split_points.windows(2).enumerate() { + if compare_split_points(ordering, &split_points[0], &split_points[1])? + != Ordering::Less + { + return plan_err!( + "Range partitioning split points must be strictly ordered: split point {idx} ({}) must be less than split point {} ({})", + split_points[0], + idx + 1, + split_points[1] + ); + } + } + + Ok(()) +} + +fn compare_split_points( + ordering: &LexOrdering, + left: &SplitPoint, + right: &SplitPoint, +) -> Result<Ordering> { + for ((left_value, right_value), sort_expr) in + left.values.iter().zip(&right.values).zip(ordering.iter()) + { + let value_ordering = + compare_scalar_values_for_sort(left_value, right_value, sort_expr)?; + if value_ordering != Ordering::Equal { + return Ok(value_ordering); + } + } + + Ok(Ordering::Equal) +} + +fn compare_scalar_values_for_sort( + left: &ScalarValue, + right: &ScalarValue, + sort_expr: &PhysicalSortExpr, +) -> Result<Ordering> { + match (left.is_null(), right.is_null()) { + (true, true) => Ok(Ordering::Equal), + (true, false) => Ok(if sort_expr.options.nulls_first { + Ordering::Less + } else { + Ordering::Greater + }), + (false, true) => Ok(if sort_expr.options.nulls_first { + Ordering::Greater + } else { + Ordering::Less + }), + (false, false) => { + let Some(ordering) = left.partial_cmp(right) else { + return plan_err!( + "Range partitioning split point values are not comparable: {left:?} and {right:?}" + ); + }; + Ok(if sort_expr.options.descending { + ordering.reverse() + } else { + ordering + }) + } + } +} + /// Represents how a [`Partitioning`] satisfies a [`Distribution`] requirement. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PartitioningSatisfaction { @@ -167,6 +438,7 @@ impl Partitioning { use Partitioning::*; match self { RoundRobinBatch(n) | Hash(_, n) | UnknownPartitioning(n) => *n, + Range(range) => range.partition_count(), } } @@ -265,9 +537,13 @@ impl Partitioning { PartitioningSatisfaction::NotSatisfied } - _ => PartitioningSatisfaction::NotSatisfied, + Partitioning::RoundRobinBatch(_) + | Partitioning::Range(_) + | Partitioning::UnknownPartitioning(_) => { + PartitioningSatisfaction::NotSatisfied + } }, - _ => PartitioningSatisfaction::NotSatisfied, + Distribution::SinglePartition => PartitioningSatisfaction::NotSatisfied, } } @@ -277,19 +553,29 @@ impl Partitioning { mapping: &ProjectionMapping, input_eq_properties: &EquivalenceProperties, ) -> Self { - if let Partitioning::Hash(exprs, part) = self { - let normalized_exprs = input_eq_properties - .project_expressions(exprs, mapping) - .zip(exprs) - .map(|(proj_expr, expr)| { - proj_expr.unwrap_or_else(|| { - Arc::new(UnKnownColumn::new(&expr.to_string())) + match self { + Partitioning::Hash(exprs, part) => { + let normalized_exprs = input_eq_properties + .project_expressions(exprs, mapping) + .zip(exprs) + .map(|(proj_expr, expr)| { + proj_expr.unwrap_or_else(|| { + Arc::new(UnKnownColumn::new(&expr.to_string())) + }) }) - }) - .collect(); - Partitioning::Hash(normalized_exprs, *part) - } else { - self.clone() + .collect(); + Partitioning::Hash(normalized_exprs, *part) + } + Partitioning::Range(range) => { + if let Some(projected) = range.project(mapping, input_eq_properties) { + Partitioning::Range(projected) + } else { + Partitioning::UnknownPartitioning(range.partition_count()) + } + } + Partitioning::RoundRobinBatch(_) | Partitioning::UnknownPartitioning(_) => { + self.clone() + } } } } @@ -306,6 +592,7 @@ impl PartialEq for Partitioning { { true } + (Partitioning::Range(left), Partitioning::Range(right)) => left == right, _ => false, } } @@ -356,56 +643,156 @@ mod tests { use super::*; use crate::expressions::Column; + use crate::projection::ProjectionTargets; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_common::Result; + struct PartitioningTestFixture { + schema: SchemaRef, + cols: Vec<Arc<dyn PhysicalExpr>>, + eq_properties: EquivalenceProperties, + } + + impl PartitioningTestFixture { + fn new(fields: Vec<(&str, DataType)>) -> Result<Self> { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, data_type)| Field::new(*name, data_type.clone(), false)) + .collect::<Vec<_>>(), + )); + let cols = fields + .iter() + .map(|(name, _)| { + Ok(Arc::new(Column::new_with_schema(name, &schema)?) + as Arc<dyn PhysicalExpr>) + }) + .collect::<Result<_>>()?; + let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + Ok(Self { + schema, + cols, + eq_properties, + }) + } + + fn int64(names: &[&str]) -> Result<Self> { + Self::new(names.iter().map(|name| (*name, DataType::Int64)).collect()) + } + + fn col(&self, index: usize) -> Arc<dyn PhysicalExpr> { + Arc::clone(&self.cols[index]) + } + + fn cols( + &self, + indices: impl IntoIterator<Item = usize>, + ) -> Vec<Arc<dyn PhysicalExpr>> { + indices.into_iter().map(|index| self.col(index)).collect() + } + + fn hash_partitioning( + &self, + indices: impl IntoIterator<Item = usize>, + partition_count: usize, + ) -> Partitioning { + Partitioning::Hash(self.cols(indices), partition_count) + } + + fn hash_distribution( + &self, + indices: impl IntoIterator<Item = usize>, + ) -> Distribution { + Distribution::HashPartitioned(self.cols(indices)) + } + + fn range_sort_expr( + &self, + index: usize, + options: SortOptions, + ) -> PhysicalSortExpr { + PhysicalSortExpr::new(self.col(index), options) + } + + fn range_ordering( + &self, + indices: impl IntoIterator<Item = usize>, + ) -> LexOrdering { + LexOrdering::new( + indices + .into_iter() + .map(|index| PhysicalSortExpr::new_default(self.col(index))), + ) + .expect("ordering must not be empty") + } + + fn range( + &self, + indices: impl IntoIterator<Item = usize>, + split_points: Vec<SplitPoint>, + ) -> RangePartitioning { + RangePartitioning::try_new(self.range_ordering(indices), split_points) + .expect("test range partitioning should be valid") + } + + fn range_partitioning( + &self, + indices: impl IntoIterator<Item = usize>, + split_points: Vec<SplitPoint>, + ) -> Partitioning { + Partitioning::Range(self.range(indices, split_points)) + } + + fn range_partitioning_with_ordering( + &self, + ordering: LexOrdering, + split_points: Vec<SplitPoint>, + ) -> Partitioning { + Partitioning::Range( + RangePartitioning::try_new(ordering, split_points) + .expect("test range partitioning should be valid"), + ) + } + } + #[test] fn partitioning_satisfy_distribution() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("column_1", DataType::Int64, false), - Field::new("column_2", DataType::Utf8, false), - ])); - - let partition_exprs1: Vec<Arc<dyn PhysicalExpr>> = vec![ - Arc::new(Column::new_with_schema("column_1", &schema).unwrap()), - Arc::new(Column::new_with_schema("column_2", &schema).unwrap()), - ]; - - let partition_exprs2: Vec<Arc<dyn PhysicalExpr>> = vec![ - Arc::new(Column::new_with_schema("column_2", &schema).unwrap()), - Arc::new(Column::new_with_schema("column_1", &schema).unwrap()), - ]; + let fixture = PartitioningTestFixture::new(vec![ + ("column_1", DataType::Int64), + ("column_2", DataType::Utf8), + ])?; let distribution_types = vec![ Distribution::UnspecifiedDistribution, Distribution::SinglePartition, - Distribution::HashPartitioned(partition_exprs1.clone()), + fixture.hash_distribution([0, 1]), ]; let single_partition = Partitioning::UnknownPartitioning(1); let unspecified_partition = Partitioning::UnknownPartitioning(10); let round_robin_partition = Partitioning::RoundRobinBatch(10); - let hash_partition1 = Partitioning::Hash(partition_exprs1, 10); - let hash_partition2 = Partitioning::Hash(partition_exprs2, 10); - let eq_properties = EquivalenceProperties::new(schema); + let hash_partition1 = fixture.hash_partitioning([0, 1], 10); + let hash_partition2 = fixture.hash_partitioning([1, 0], 10); for distribution in distribution_types { let result = ( single_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), unspecified_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), round_robin_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), hash_partition1 - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), hash_partition2 - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), ); @@ -427,72 +814,41 @@ mod tests { #[test] fn test_partitioning_satisfy_by_subset() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; let test_cases = vec![ ( "Hash([a]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), + fixture.hash_partitioning([0], 4), + fixture.hash_distribution([0, 1]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([a]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + fixture.hash_partitioning([0], 4), + fixture.hash_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([a, b]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + fixture.hash_partitioning([0, 1], 4), + fixture.hash_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([b]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + fixture.hash_partitioning([1], 4), + fixture.hash_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([b, a]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + fixture.hash_partitioning([1, 0], 4), + fixture.hash_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), @@ -501,13 +857,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); + let result = partition.satisfaction(&required, &fixture.eq_properties, true); assert_eq!( result, expected_with_subset, "Failed for {desc} with subset enabled" ); - let result = partition.satisfaction(&required, &eq_properties, false); + let result = partition.satisfaction(&required, &fixture.eq_properties, false); assert_eq!( result, expected_without_subset, "Failed for {desc} with subset disabled" @@ -519,48 +875,27 @@ mod tests { #[test] fn test_partitioning_current_superset() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; let test_cases = vec![ ( "Hash([a, b]) vs Hash([a])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + fixture.hash_partitioning([0, 1], 4), + fixture.hash_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([a, b, c]) vs Hash([a])", - Partitioning::Hash( - vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_c)], - 4, - ), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + fixture.hash_partitioning([0, 1, 2], 4), + fixture.hash_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([a, b, c]) vs Hash([a, b])", - Partitioning::Hash( - vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_c)], - 4, - ), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), + fixture.hash_partitioning([0, 1, 2], 4), + fixture.hash_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -569,13 +904,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); + let result = partition.satisfaction(&required, &fixture.eq_properties, true); assert_eq!( result, expected_with_subset, "Failed for {desc} with subset enabled" ); - let result = partition.satisfaction(&required, &eq_properties, false); + let result = partition.satisfaction(&required, &fixture.eq_properties, false); assert_eq!( result, expected_without_subset, "Failed for {desc} with subset disabled" @@ -587,24 +922,12 @@ mod tests { #[test] fn test_partitioning_partial_overlap() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; let test_cases = vec![( "Partial overlap: Hash([a, c]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_c)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a), Arc::clone(&col_b)]), + fixture.hash_partitioning([0, 2], 4), + fixture.hash_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, )]; @@ -612,13 +935,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); + let result = partition.satisfaction(&required, &fixture.eq_properties, true); assert_eq!( result, expected_with_subset, "Failed for {desc} with subset enabled" ); - let result = partition.satisfaction(&required, &eq_properties, false); + let result = partition.satisfaction(&required, &fixture.eq_properties, false); assert_eq!( result, expected_without_subset, "Failed for {desc} with subset disabled" @@ -630,35 +953,20 @@ mod tests { #[test] fn test_partitioning_no_overlap() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; let test_cases = vec![ ( "Hash([a]) vs Hash([b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + fixture.hash_partitioning([0], 4), + fixture.hash_distribution([1, 2]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([a, b]) vs Hash([c])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_c)]), + fixture.hash_partitioning([0, 1], 4), + fixture.hash_distribution([2]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -667,13 +975,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); + let result = partition.satisfaction(&required, &fixture.eq_properties, true); assert_eq!( result, expected_with_subset, "Failed for {desc} with subset enabled" ); - let result = partition.satisfaction(&required, &eq_properties, false); + let result = partition.satisfaction(&required, &fixture.eq_properties, false); assert_eq!( result, expected_without_subset, "Failed for {desc} with subset disabled" @@ -685,32 +993,20 @@ mod tests { #[test] fn test_partitioning_exact_match() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - ])); - - let col_a: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("b", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; let test_cases = vec![ ( "Hash([a, b]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), + fixture.hash_partitioning([0, 1], 4), + fixture.hash_distribution([0, 1]), PartitioningSatisfaction::Exact, PartitioningSatisfaction::Exact, ), ( "Hash([a]) vs Hash([a])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + fixture.hash_partitioning([0], 4), + fixture.hash_distribution([0]), PartitioningSatisfaction::Exact, PartitioningSatisfaction::Exact, ), @@ -719,13 +1015,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); + let result = partition.satisfaction(&required, &fixture.eq_properties, true); assert_eq!( result, expected_with_subset, "Failed for {desc} with subset enabled" ); - let result = partition.satisfaction(&required, &eq_properties, false); + let result = partition.satisfaction(&required, &fixture.eq_properties, false); assert_eq!( result, expected_without_subset, "Failed for {desc} with subset disabled" @@ -737,32 +1033,20 @@ mod tests { #[test] fn test_partitioning_unknown() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - ])); - - let col_a: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("b", &schema)?); + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; let unknown: Arc<dyn PhysicalExpr> = Arc::new(UnKnownColumn::new("dropped")); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); let test_cases = vec![ ( "Hash([unknown]) vs Hash([a, b])", Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), + fixture.hash_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([a, b]) vs Hash([unknown])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), + fixture.hash_partitioning([0, 1], 4), Distribution::HashPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, @@ -779,13 +1063,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); + let result = partition.satisfaction(&required, &fixture.eq_properties, true); assert_eq!( result, expected_with_subset, "Failed for {desc} with subset enabled" ); - let result = partition.satisfaction(&required, &eq_properties, false); + let result = partition.satisfaction(&required, &fixture.eq_properties, false); assert_eq!( result, expected_without_subset, "Failed for {desc} with subset disabled" @@ -797,23 +1081,19 @@ mod tests { #[test] fn test_partitioning_empty_hash() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - - let col_a: Arc<dyn PhysicalExpr> = - Arc::new(Column::new_with_schema("a", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let fixture = PartitioningTestFixture::int64(&["a"])?; let test_cases = vec![ ( "Hash([]) vs Hash([a])", Partitioning::Hash(vec![], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + fixture.hash_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( "Hash([a]) vs Hash([])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), + fixture.hash_partitioning([0], 4), Distribution::HashPartitioned(vec![]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, @@ -830,13 +1110,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); + let result = partition.satisfaction(&required, &fixture.eq_properties, true); assert_eq!( result, expected_with_subset, "Failed for {desc} with subset enabled" ); - let result = partition.satisfaction(&required, &eq_properties, false); + let result = partition.satisfaction(&required, &fixture.eq_properties, false); assert_eq!( result, expected_without_subset, "Failed for {desc} with subset disabled" @@ -845,4 +1125,160 @@ mod tests { Ok(()) } + + fn int_split_point(values: impl IntoIterator<Item = i64>) -> SplitPoint { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::Int64(Some(value))) + .collect(), + ) + } + + fn assert_range_try_new_error( + ordering: LexOrdering, + split_points: Vec<SplitPoint>, + expected: &str, + ) { + let error = RangePartitioning::try_new(ordering, split_points) + .unwrap_err() + .to_string(); + assert!(error.contains(expected), "{error}"); + } + + #[test] + fn test_range_partitioning_metadata() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + + let range_partitioning = + fixture.range([0], vec![int_split_point([10]), int_split_point([20])]); + assert_eq!(range_partitioning.ordering()[0].to_string(), "a@0 ASC"); + assert_eq!( + range_partitioning.split_points(), + &[int_split_point([10]), int_split_point([20])] + ); + let partitioning = Partitioning::Range(range_partitioning); + + assert_eq!(partitioning.partition_count(), 3); + assert_eq!( + partitioning.to_string(), + "Range([a@0 ASC], [(10), (20)], 3)" + ); + + Ok(()) + } + + #[test] + fn test_range_partitioning_try_new_validates_split_points() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let asc_a = fixture.range_ordering([0]); + let ordering_ab = fixture.range_ordering([0, 1]); + + assert_range_try_new_error( + ordering_ab.clone(), + vec![int_split_point([10])], + "split point 0 has width 1, but ordering has width 2", + ); + + RangePartitioning::try_new( + [fixture.range_sort_expr(0, SortOptions::new(true, false))].into(), + vec![int_split_point([20]), int_split_point([10])], + )?; + + assert_range_try_new_error( + asc_a, + vec![int_split_point([20]), int_split_point([10])], + "split points must be strictly ordered", + ); + + assert_range_try_new_error( + [fixture.range_sort_expr(0, SortOptions::new(false, false))].into(), + vec![ + SplitPoint::new(vec![ScalarValue::Int64(None)]), + int_split_point([10]), + ], + "split points must be strictly ordered", + ); + + RangePartitioning::try_new( + ordering_ab.clone(), + vec![int_split_point([10, 20]), int_split_point([10, 30])], + )?; + + assert_range_try_new_error( + ordering_ab, + vec![int_split_point([10, 30]), int_split_point([10, 20])], + "split points must be strictly ordered", + ); + + Ok(()) + } + + #[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 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)" + ); + + let drop_b_mapping = ProjectionMapping::from_indices(&[0], &fixture.schema)?; + let projected = + range_partitioning.project(&drop_b_mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + + Ok(()) + } + + #[test] + fn test_range_partitioning_project_degrades_if_ordering_collapses() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let target: Arc<dyn PhysicalExpr> = Arc::new(Column::new("x", 0)); + let range_partitioning = + fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); + let mapping = ProjectionMapping::from_iter([ + ( + fixture.col(0), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), + ), + ( + fixture.col(1), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), + ), + ]); + + let projected = range_partitioning.project(&mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + + Ok(()) + } + + #[test] + fn test_multi_partition_range_does_not_satisfy_hash_distribution() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let range_partitioning = + fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); + let required = fixture.hash_distribution([0, 1]); + + assert_eq!( + range_partitioning.satisfaction(&required, &fixture.eq_properties, false), + PartitioningSatisfaction::NotSatisfied + ); + + Ok(()) + } } diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index b4aa295562..9a6d1e5545 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -144,6 +144,13 @@ pub fn adjust_right_output_partitioning( .collect::<Result<_>>()?; Partitioning::Hash(new_exprs, *size) } + Partitioning::Range(_) => { + // Range partitioning optimizer propagation is tracked in + // https://github.com/apache/datafusion/issues/22395 + return not_impl_err!( + "Join output partitioning with range partitioning is not implemented" + ); + } result => result.clone(), }; Ok(result) diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 3005e97542..c7b1d4729e 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -37,7 +37,7 @@ pub use datafusion_expr::{Accumulator, ColumnarValue}; use datafusion_physical_expr::PhysicalSortExpr; pub use datafusion_physical_expr::window::WindowExpr; pub use datafusion_physical_expr::{ - Distribution, Partitioning, PhysicalExpr, expressions, + Distribution, Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, expressions, }; pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index a6363378ed..3d30dd8276 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -737,6 +737,13 @@ impl BatchPartitioner { num_input_partitions, )) } + Partitioning::Range(_) => { + // Range repartition execution is tracked in + // https://github.com/apache/datafusion/issues/22397 + not_impl_err!( + "Range partitioning execution is not implemented by RepartitionExec" + ) + } other => { not_impl_err!("Unsupported repartitioning scheme {other:?}") } @@ -1430,6 +1437,13 @@ impl ExecutionPlan for RepartitionExec { } Partitioning::Hash(new_partitions, *size) } + Partitioning::Range(_) => { + // Range partitioning optimizer propagation is tracked in + // https://github.com/apache/datafusion/issues/22395 + return not_impl_err!( + "Projection pushdown through RepartitionExec with range partitioning is not implemented" + ); + } others => others.clone(), }; @@ -1466,6 +1480,18 @@ impl ExecutionPlan for RepartitionExec { if !self.maintains_input_order()[0] { return Ok(SortOrderPushdownResult::Unsupported); } + match self.partitioning() { + Partitioning::Range(_) => { + // Range partitioning optimizer propagation is tracked in + // https://github.com/apache/datafusion/issues/22395 + return not_impl_err!( + "Sort pushdown through RepartitionExec with range partitioning is not implemented" + ); + } + Partitioning::RoundRobinBatch(_) + | Partitioning::Hash(_, _) + | Partitioning::UnknownPartitioning(_) => {} + } // Delegate to the child and wrap with a new RepartitionExec self.input.try_pushdown_sort(order)?.try_map(|new_input| { @@ -1489,6 +1515,13 @@ impl ExecutionPlan for RepartitionExec { RoundRobinBatch(_) => RoundRobinBatch(target_partitions), Hash(hash, _) => Hash(hash, target_partitions), UnknownPartitioning(_) => UnknownPartitioning(target_partitions), + Range(_) => { + // Range repartition execution is tracked in + // https://github.com/apache/datafusion/issues/22397 + return not_impl_err!( + "Changing RepartitionExec partition counts with range partitioning is not implemented" + ); + } }; Ok(Some(Arc::new(Self { input: Arc::clone(&self.input), @@ -1617,6 +1650,13 @@ impl RepartitionExec { num_input_partitions, ) } + Partitioning::Range(_) => { + // Range repartition execution is tracked in + // https://github.com/apache/datafusion/issues/22397 + return not_impl_err!( + "Range partitioning execution is not implemented by RepartitionExec" + ); + } other => { return not_impl_err!("Unsupported repartitioning scheme {other:?}"); } @@ -1968,12 +2008,14 @@ mod tests { use arrow::array::{ArrayRef, StringArray, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::ScalarValue; use datafusion_common::cast::as_string_array; use datafusion_common::exec_err; use datafusion_common::test_util::batches_to_sort_string; use datafusion_common_runtime::JoinSet; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint}; use insta::assert_snapshot; #[test] @@ -2266,6 +2308,40 @@ mod tests { ); } + #[tokio::test] + async fn unsupported_range_partitioning() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef, + )])?; + + let schema = batch.schema(); + let expr = col("my_awesome_field", &schema)?; + let input = MockExec::new(vec![Ok(batch)], Arc::clone(&schema)); + let partitioning = Partitioning::Range(RangePartitioning::new( + [PhysicalSortExpr::new_default(expr)].into(), + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "foo".to_string(), + ))])], + )); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning)?; + let output_stream = exec.execute(0, task_ctx)?; + + let result_string = crate::common::collect(output_stream) + .await + .unwrap_err() + .to_string(); + assert!( + result_string.contains( + "Range partitioning execution is not implemented by RepartitionExec" + ), + "actual: {result_string}" + ); + + Ok(()) + } + #[tokio::test] async fn error_for_input_exec() { // This generates an error on a call to execute. The error diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index f715de0b59..929ff4f7df 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1141,7 +1141,8 @@ impl ExecutionPlan for SortExec { vec![Distribution::UnspecifiedDistribution] } else { // global sort - // TODO support RangePartition and OrderedDistribution + // TODO support range partitioning and OrderedDistribution. + // See https://github.com/apache/datafusion/issues/22395 vec![Distribution::SinglePartition] } } diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 09570f14ba..eb9b5f09aa 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -1486,11 +1486,7 @@ mod tests { let task_ctx = Arc::new(TaskContext::default()); let schema = Schema::new(vec![Field::new("c1", DataType::UInt64, false)]); let properties = CongestedExec::compute_properties(Arc::new(schema.clone())); - let &partition_count = match properties.output_partitioning() { - Partitioning::RoundRobinBatch(partitions) => partitions, - Partitioning::Hash(_, partitions) => partitions, - Partitioning::UnknownPartitioning(partitions) => partitions, - }; + let partition_count = properties.output_partitioning().partition_count(); let source = CongestedExec { schema: schema.clone(), cache: Arc::new(properties), diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index ea6d078366..2185748c70 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1381,13 +1381,22 @@ message PhysicalHashRepartition { uint64 partition_count = 2; } +message PhysicalRangePartitioning { + repeated PhysicalSortExprNode sort_expr = 1; + repeated PhysicalRangeSplitPoint split_point = 2; +} + +message PhysicalRangeSplitPoint { + repeated datafusion_common.ScalarValue value = 1; +} + message RepartitionExecNode{ PhysicalPlanNode input = 1; - // oneof partition_method { + // Legacy direct partitioning fields: // uint64 round_robin = 2; // PhysicalHashRepartition hash = 3; // uint64 unknown = 4; - // } + // New partitioning variants are stored in `partitioning`. Partitioning partitioning = 5; bool preserve_order = 6; } @@ -1397,6 +1406,7 @@ message Partitioning { uint64 round_robin = 1; PhysicalHashRepartition hash = 2; uint64 unknown = 3; + PhysicalRangePartitioning range = 4; } } diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 8e6997757f..4136cd2785 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -15768,6 +15768,9 @@ impl serde::Serialize for Partitioning { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("unknown", ToString::to_string(&v).as_str())?; } + partitioning::PartitionMethod::Range(v) => { + struct_ser.serialize_field("range", v)?; + } } } struct_ser.end() @@ -15784,6 +15787,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { "roundRobin", "hash", "unknown", + "range", ]; #[allow(clippy::enum_variant_names)] @@ -15791,6 +15795,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { RoundRobin, Hash, Unknown, + Range, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> @@ -15815,6 +15820,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { "roundRobin" | "round_robin" => Ok(GeneratedField::RoundRobin), "hash" => Ok(GeneratedField::Hash), "unknown" => Ok(GeneratedField::Unknown), + "range" => Ok(GeneratedField::Range), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -15856,6 +15862,13 @@ impl<'de> serde::Deserialize<'de> for Partitioning { } partition_method__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| partitioning::PartitionMethod::Unknown(x.0)); } + GeneratedField::Range => { + if partition_method__.is_some() { + return Err(serde::de::Error::duplicate_field("range")); + } + partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(partitioning::PartitionMethod::Range) +; + } } } Ok(Partitioning { @@ -19051,6 +19064,207 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { deserializer.deserialize_struct("datafusion.PhysicalPlanNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalRangePartitioning { + #[allow(deprecated)] + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.sort_expr.is_empty() { + len += 1; + } + if !self.split_point.is_empty() { + 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)?; + } + if !self.split_point.is_empty() { + struct_ser.serialize_field("splitPoint", &self.split_point)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalRangePartitioning { + #[allow(deprecated)] + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "sort_expr", + "sortExpr", + "split_point", + "splitPoint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + SortExpr, + SplitPoint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str<E>(self, value: &str) -> std::result::Result<GeneratedField, E> + where + E: serde::de::Error, + { + match value { + "sortExpr" | "sort_expr" => Ok(GeneratedField::SortExpr), + "splitPoint" | "split_point" => Ok(GeneratedField::SplitPoint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalRangePartitioning; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalRangePartitioning") + } + + fn visit_map<V>(self, mut map_: V) -> std::result::Result<PhysicalRangePartitioning, V::Error> + where + V: serde::de::MapAccess<'de>, + { + let mut sort_expr__ = None; + let mut split_point__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::SortExpr => { + if sort_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("sortExpr")); + } + sort_expr__ = Some(map_.next_value()?); + } + GeneratedField::SplitPoint => { + if split_point__.is_some() { + return Err(serde::de::Error::duplicate_field("splitPoint")); + } + split_point__ = Some(map_.next_value()?); + } + } + } + Ok(PhysicalRangePartitioning { + sort_expr: sort_expr__.unwrap_or_default(), + split_point: split_point__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalRangePartitioning", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalRangeSplitPoint { + #[allow(deprecated)] + fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.value.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalRangeSplitPoint", len)?; + if !self.value.is_empty() { + struct_ser.serialize_field("value", &self.value)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalRangeSplitPoint { + #[allow(deprecated)] + fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "value", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Value, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error> + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str<E>(self, value: &str) -> std::result::Result<GeneratedField, E> + where + E: serde::de::Error, + { + match value { + "value" => Ok(GeneratedField::Value), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalRangeSplitPoint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalRangeSplitPoint") + } + + fn visit_map<V>(self, mut map_: V) -> std::result::Result<PhysicalRangeSplitPoint, V::Error> + where + V: serde::de::MapAccess<'de>, + { + let mut value__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Value => { + if value__.is_some() { + return Err(serde::de::Error::duplicate_field("value")); + } + value__ = Some(map_.next_value()?); + } + } + } + Ok(PhysicalRangeSplitPoint { + value: value__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalRangeSplitPoint", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PhysicalScalarSubqueryExprNode { #[allow(deprecated)] fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index d8187e65a5..4e473668e8 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2044,14 +2044,26 @@ pub struct PhysicalHashRepartition { pub partition_count: u64, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalRangePartitioning { + #[prost(message, repeated, tag = "1")] + pub sort_expr: ::prost::alloc::vec::Vec<PhysicalSortExprNode>, + #[prost(message, repeated, tag = "2")] + pub split_point: ::prost::alloc::vec::Vec<PhysicalRangeSplitPoint>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalRangeSplitPoint { + #[prost(message, repeated, tag = "1")] + pub value: ::prost::alloc::vec::Vec<super::datafusion_common::ScalarValue>, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct RepartitionExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box<PhysicalPlanNode>>, - /// oneof partition_method { + /// Legacy direct partitioning fields: /// uint64 round_robin = 2; /// PhysicalHashRepartition hash = 3; /// uint64 unknown = 4; - /// } + /// New partitioning variants are stored in `partitioning`. #[prost(message, optional, tag = "5")] pub partitioning: ::core::option::Option<Partitioning>, #[prost(bool, tag = "6")] @@ -2059,7 +2071,7 @@ pub struct RepartitionExecNode { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct Partitioning { - #[prost(oneof = "partitioning::PartitionMethod", tags = "1, 2, 3")] + #[prost(oneof = "partitioning::PartitionMethod", tags = "1, 2, 3, 4")] pub partition_method: ::core::option::Option<partitioning::PartitionMethod>, } /// Nested message and enum types in `Partitioning`. @@ -2072,6 +2084,8 @@ pub mod partitioning { Hash(super::PhysicalHashRepartition), #[prost(uint64, tag = "3")] Unknown(u64), + #[prost(message, tag = "4")] + Range(super::PhysicalRangePartitioning), } } #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index d2a48aa457..7d2e68d810 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -24,7 +24,9 @@ use arrow::compute::SortOptions; use arrow::datatypes::{Field, Schema}; use arrow::ipc::reader::StreamReader; use chrono::{TimeZone, Utc}; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, not_impl_err}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, internal_datafusion_err, not_impl_err, +}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; @@ -48,7 +50,9 @@ use datafusion_physical_plan::expressions::{ }; use datafusion_physical_plan::joins::{HashExpr, SeededRandomState}; use datafusion_physical_plan::windows::{create_window_expr, schema_add_window_field}; -use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; +use datafusion_physical_plan::{ + Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, +}; use datafusion_proto_common::common::proto_error; use object_store::ObjectMeta; use object_store::path::Path; @@ -560,6 +564,14 @@ pub fn parse_protobuf_partitioning( proto_converter, ) } + Some(protobuf::partitioning::PartitionMethod::Range(range_partitioning)) => { + Ok(Some(parse_protobuf_range_partitioning( + range_partitioning, + ctx, + input_schema, + proto_converter, + )?)) + } Some(protobuf::partitioning::PartitionMethod::Unknown(partition_count)) => { Ok(Some(Partitioning::UnknownPartitioning( *partition_count as usize, @@ -571,6 +583,49 @@ pub fn parse_protobuf_partitioning( } } +fn parse_protobuf_range_partitioning( + range_partitioning: &protobuf::PhysicalRangePartitioning, + ctx: &PhysicalPlanDecodeContext<'_>, + input_schema: &Schema, + proto_converter: &dyn PhysicalProtoConverterExtension, +) -> Result<Partitioning> { + let sort_exprs = parse_physical_sort_exprs( + &range_partitioning.sort_expr, + ctx, + input_schema, + proto_converter, + )?; + let sort_expr_count = sort_exprs.len(); + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!("Range partitioning requires non-empty ordering") + })?; + if ordering.len() != sort_expr_count { + return Err(internal_datafusion_err!( + "Range partitioning ordering must not contain duplicate expressions" + )); + } + let split_points = range_partitioning + .split_point + .iter() + .map(parse_protobuf_range_split_point) + .collect::<Result<_>>()?; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) +} + +fn parse_protobuf_range_split_point( + split_point: &protobuf::PhysicalRangeSplitPoint, +) -> Result<SplitPoint> { + let values = split_point + .value + .iter() + .map(|value| ScalarValue::try_from(value).map_err(Into::into)) + .collect::<Result<_>>()?; + Ok(SplitPoint::new(values)) +} + pub fn parse_protobuf_file_scan_schema( proto: &protobuf::FileScanExecConf, ) -> Result<Arc<Schema>> { diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 28c0a57e94..9cb9e89760 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -42,7 +42,9 @@ use datafusion_physical_plan::expressions::{ use datafusion_physical_plan::joins::HashExpr; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; -use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; +use datafusion_physical_plan::{ + Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, +}; use super::{ DefaultPhysicalProtoConverter, PhysicalExtensionCodec, @@ -542,6 +544,11 @@ pub fn serialize_partitioning( )), } } + Partitioning::Range(range) => protobuf::Partitioning { + partition_method: Some(protobuf::partitioning::PartitionMethod::Range( + serialize_range_partitioning(range, codec, proto_converter)?, + )), + }, Partitioning::UnknownPartitioning(partition_count) => protobuf::Partitioning { partition_method: Some(protobuf::partitioning::PartitionMethod::Unknown( *partition_count as u64, @@ -551,6 +558,40 @@ pub fn serialize_partitioning( Ok(serialized_partitioning) } +fn serialize_range_partitioning( + range: &RangePartitioning, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, +) -> Result<protobuf::PhysicalRangePartitioning> { + Ok(protobuf::PhysicalRangePartitioning { + sort_expr: serialize_physical_sort_exprs( + range.ordering().iter().cloned(), + codec, + proto_converter, + )?, + split_point: range + .split_points() + .iter() + .map(serialize_range_split_point) + .collect::<Result<_>>()?, + }) +} + +fn serialize_range_split_point( + split_point: &SplitPoint, +) -> Result<protobuf::PhysicalRangeSplitPoint> { + Ok(protobuf::PhysicalRangeSplitPoint { + value: split_point + .values() + .iter() + .map(|value| { + TryInto::<datafusion_proto_common::ScalarValue>::try_into(value) + .map_err(Into::into) + }) + .collect::<Result<_>>()?, + }) +} + fn serialize_when_then_expr( when_expr: &Arc<dyn PhysicalExpr>, then_expr: &Arc<dyn PhysicalExpr>, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index d88a360422..bd996eb692 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -90,7 +90,8 @@ use datafusion::physical_plan::windows::{ }; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning, - PhysicalExpr, SendableRecordBatchStream, Statistics, displayable, + PhysicalExpr, RangePartitioning, SendableRecordBatchStream, SplitPoint, Statistics, + displayable, }; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion::scalar::ScalarValue; @@ -1806,6 +1807,21 @@ fn roundtrip_repartition_preserve_order() -> Result<()> { roundtrip_test(Arc::new(repartition)) } +#[test] +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))])], + )); + // RepartitionExec is used only to carry the partitioning through proto. + // Executing range repartitioning is intentionally unsupported. + let repartition = RepartitionExec::try_new(input, range_partitioning)?; + + roundtrip_test(Arc::new(repartition)) +} + #[test] fn roundtrip_interleave() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
