From e2d9f846b73762193e3bcb69d5e0827d29fb7149 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Fri, 24 Jul 2026 05:17:30 -0400 Subject: [PATCH 01/22] Backport physical range partitioning to branch 54 Backport the physical range partitioning representation, partition-map compatibility, and co-partitioned range hash join support onto the DataDog branch-54 base. Adapt the optimizer tests to the branch-54 execution-plan APIs and keep range-partitioned dynamic-filter routing disabled until it is supported safely. Upstream-Commit: 7a6b0626da7a733d94821b8d9d6011eb85d96593 Upstream-Commit: 1465d6fa4e580638bb8e0987adc33d8892a31547 Upstream-Commit: 0aa76ecd1017ed9b69eabe90603cff290534f701 --- .../enforce_distribution.rs | 238 ++- .../physical_optimizer/sanity_checker.rs | 54 +- .../ffi/src/physical_expr/partitioning.rs | 5 + datafusion/physical-expr/src/lib.rs | 2 +- datafusion/physical-expr/src/partitioning.rs | 1336 ++++++++++++++--- .../src/enforce_distribution.rs | 147 +- .../src/output_requirements.rs | 18 +- .../physical-optimizer/src/sanity_checker.rs | 48 +- .../physical-plan/src/execution_plan.rs | 9 +- .../physical-plan/src/joins/hash_join/exec.rs | 334 ++++- datafusion/physical-plan/src/joins/utils.rs | 6 + datafusion/physical-plan/src/lib.rs | 2 +- .../physical-plan/src/repartition/mod.rs | 76 + datafusion/physical-plan/src/sorts/sort.rs | 3 +- .../src/sorts/sort_preserving_merge.rs | 6 +- datafusion/proto/proto/datafusion.proto | 14 +- datafusion/proto/src/generated/pbjson.rs | 214 +++ datafusion/proto/src/generated/prost.rs | 20 +- .../proto/src/physical_plan/from_proto.rs | 59 +- .../proto/src/physical_plan/to_proto.rs | 43 +- .../tests/cases/roundtrip_physical_plan.rs | 18 +- 21 files changed, 2332 insertions(+), 320 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 12abf79041091..4ff438fdddfca 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -59,9 +59,13 @@ use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; -use datafusion_physical_expr::Distribution; +use datafusion_physical_expr::{ + Distribution, EquivalenceProperties, RangePartitioning, SplitPoint, +}; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion_physical_plan::execution_plan::ExecutionPlan; +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, ExecutionPlan, +}; use datafusion_physical_plan::expressions::col; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::JoinOn; @@ -70,7 +74,8 @@ use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlanProperties, PlanProperties, displayable, + DisplayAs, DisplayFormatType, ExecutionPlanProperties, Partitioning, PlanProperties, + displayable, }; use insta::Settings; @@ -325,6 +330,104 @@ fn parquet_exec_multiple_sorted( DataSourceExec::from_data_source(config) } +#[derive(Debug)] +struct PartitionedTestExec { + cache: Arc, +} + +impl PartitionedTestExec { + fn new(output_partitioning: Partitioning) -> Self { + Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema()), + output_partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )), + } + } +} + +impl DisplayAs for PartitionedTestExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "PartitionedTestExec: output_partitioning={}", + self.cache.output_partitioning() + ), + DisplayFormatType::TreeRender => write!(f, ""), + } + } +} + +impl ExecutionPlan for PartitionedTestExec { + fn name(&self) -> &'static str { + "PartitionedTestExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert!(children.is_empty()); + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } +} + +fn parquet_exec_with_output_partitioning( + output_partitioning: Partitioning, +) -> Arc { + Arc::new(PartitionedTestExec::new(output_partitioning)) +} + +fn range_partitioning( + key: &str, + split_points: impl IntoIterator, + options: SortOptions, +) -> Result { + let ordering = [PhysicalSortExpr { + expr: col(key, &schema())?, + options, + }] + .into(); + let split_points = split_points + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect(); + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) +} + +fn hash_partitioning(key: &str, partition_count: usize) -> Result { + Ok(Partitioning::Hash( + vec![col(key, &schema())?], + partition_count, + )) +} + fn csv_exec() -> Arc { csv_exec_with_sort(vec![]) } @@ -701,6 +804,135 @@ impl TestConfig { } } +fn partitioned_join_plan( + left_partitioning: Partitioning, + right_partitioning: Partitioning, + join_type: JoinType, +) -> Arc { + let left = parquet_exec_with_output_partitioning(left_partitioning); + let right = parquet_exec_with_output_partitioning(right_partitioning); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &right.schema()).unwrap()) as _, + )]; + hash_join_exec(left, right, &join_on, &join_type) +} + +#[test] +fn inner_range_join_keeps_range_partitioning() -> Result<()> { + let join = partitioned_join_plan( + range_partitioning("a", [10, 20], SortOptions::default())?, + range_partitioning("b", [10, 20], SortOptions::default())?, + JoinType::Inner, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20)], 3) + PartitionedTestExec: output_partitioning=Range([b@1 ASC], [(10), (20)], 3) + " + ); + + Ok(()) +} + +#[test] +fn inner_range_join_rehashes_different_bounds() -> Result<()> { + let join = partitioned_join_plan( + range_partitioning("a", [10, 20], SortOptions::default())?, + range_partitioning("b", [15, 20], SortOptions::default())?, + JoinType::Inner, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20)], 3) + RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Range([b@1 ASC], [(15), (20)], 3) + " + ); + + Ok(()) +} + +#[test] +fn inner_hash_join_rehashes_mismatched_counts() -> Result<()> { + let join = partitioned_join_plan( + hash_partitioning("a", 11)?, + hash_partitioning("b", 12)?, + JoinType::Inner, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=11 + PartitionedTestExec: output_partitioning=Hash([a@0], 11) + RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=12 + PartitionedTestExec: output_partitioning=Hash([b@1], 12) + " + ); + Ok(()) +} + +#[test] +fn inner_hash_join_rehashes_to_target_count() -> Result<()> { + let join = partitioned_join_plan( + hash_partitioning("a", 3)?, + hash_partitioning("b", 3)?, + JoinType::Inner, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Hash([a@0], 3) + RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Hash([b@1], 3) + " + ); + Ok(()) +} + +#[test] +fn non_inner_range_join_rehashes() -> Result<()> { + let join = partitioned_join_plan( + range_partitioning("a", [10, 20], SortOptions::default())?, + range_partitioning("b", [10, 20], SortOptions::default())?, + JoinType::Left, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Left, on=[(a@0, b@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20)], 3) + RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Range([b@1 ASC], [(10), (20)], 3) + " + ); + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index 217570846d56e..f12e5d5f764b0 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -19,9 +19,9 @@ use insta::assert_snapshot; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - bounded_window_exec, global_limit_exec, local_limit_exec, memory_exec, - projection_exec, repartition_exec, sort_exec, sort_expr, sort_expr_options, - sort_merge_join_exec, sort_preserving_merge_exec, union_exec, + bounded_window_exec, global_limit_exec, hash_join_exec, local_limit_exec, + memory_exec, projection_exec, repartition_exec, sort_exec, sort_expr, + sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; use arrow::compute::SortOptions; @@ -30,8 +30,8 @@ use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTab use datafusion::prelude::{CsvReadOptions, SessionContext}; use datafusion_common::config::ConfigOptions; use datafusion_common::{JoinType, Result, ScalarValue}; -use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::expressions::{Literal, col}; +use datafusion_physical_expr::{Partitioning, RangePartitioning, SplitPoint}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; @@ -400,6 +400,52 @@ fn assert_sanity_check(plan: &Arc, is_sane: bool) { ); } +fn range_partitioned_exec( + schema: &SchemaRef, + key: &str, + split_points: impl IntoIterator, +) -> Result> { + let split_points = split_points + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [sort_expr(key, schema)].into(), + split_points, + )?); + let input = memory_exec(schema); + + RepartitionExec::try_new(input, partitioning) + .map(|exec| Arc::new(exec) as Arc) +} + +#[test] +fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("b", &schema)?)]; + let right = range_partitioned_exec(&schema, "b", [10])?; + + let valid_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + Arc::clone(&right), + join_on.clone(), + None, + &JoinType::Inner, + )?; + assert_sanity_check(&valid_join, true); + + let invalid_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [20])?, + right, + join_on, + None, + &JoinType::Inner, + )?; + assert_sanity_check(&invalid_join, false); + + Ok(()) +} + #[tokio::test] /// Tests that plan is valid when the sort requirements are satisfied. async fn test_bounded_window_agg_sort_requirement() -> Result<()> { diff --git a/datafusion/ffi/src/physical_expr/partitioning.rs b/datafusion/ffi/src/physical_expr/partitioning.rs index 434b6a097e645..eec437639e156 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 848bf81d15979..ad788c15d098e 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -58,7 +58,7 @@ pub use analysis::{AnalysisContext, ExprBoundaries, analyze}; pub use equivalence::{ AcrossPartitions, ConstExpr, EquivalenceProperties, calculate_union, }; -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 d24c60b63e6bd..77411ecb765a3 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>, 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,335 @@ 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 +/// . +#[derive(Debug, Clone, PartialEq)] +pub struct RangePartitioning { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Boundaries between adjacent partitions. + split_points: Vec, +} + +/// 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, +} + +impl SplitPoint { + /// Creates a new split point from its tuple values. + pub fn new(values: Vec) -> 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::>() + .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) -> 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) -> Result { + 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 + } + + /// Returns true when `self` and `other` have the same range boundaries. + /// + /// Single-partition range partitionings always have the same boundaries. Otherwise, + /// the two partitionings must have identical split points, ordering width, + /// and sort options. This does not compare ordering expressions, callers + /// should validate the range keys separately. + fn same_boundaries(&self, other: &Self) -> bool { + if self.partition_count() == 1 && other.partition_count() == 1 { + return true; + } + + if self.split_points != other.split_points + || self.ordering.len() != other.ordering.len() + { + return false; + } + + if !self + .ordering + .iter() + .zip(other.ordering.iter()) + .all(|(left, right)| left.options == right.options) + { + return false; + } + + true + } + + /// Calculates the range partitioning after applying the given projection. + /// + /// Returns `None` if any range key cannot be projected or if projection + /// collapses distinct range keys into duplicate output expressions. + fn project( + &self, + mapping: &ProjectionMapping, + input_eq_properties: &EquivalenceProperties, + ) -> Option { + let exprs = self + .ordering + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + let projected_exprs = input_eq_properties + .project_expressions(&exprs, mapping) + .collect::>>()?; + let sort_exprs = self + .ordering + .iter() + .zip(projected_exprs) + .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, sort_expr.options)) + .collect::>(); + 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::>() + .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 { + 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 { + 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 + }) + } + } +} + +fn equivalent_exprs( + left: &[Arc], + right: &[Arc], + eq_properties: &EquivalenceProperties, +) -> bool { + if physical_exprs_equal(left, right) { + return true; + } + + let eq_groups = eq_properties.eq_group(); + if eq_groups.is_empty() { + return false; + } + + let normalized_left = normalize_exprs(left, eq_properties); + let normalized_right = normalize_exprs(right, eq_properties); + + physical_exprs_equal(&normalized_left, &normalized_right) +} + +fn normalize_exprs( + exprs: &[Arc], + eq_properties: &EquivalenceProperties, +) -> Vec> { + let eq_groups = eq_properties.eq_group(); + exprs + .iter() + .map(|expr| eq_groups.normalize_expr(Arc::clone(expr))) + .collect() +} + /// Represents how a [`Partitioning`] satisfies a [`Distribution`] requirement. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PartitioningSatisfaction { @@ -167,6 +502,62 @@ impl Partitioning { use Partitioning::*; match self { RoundRobinBatch(n) | Hash(_, n) | UnknownPartitioning(n) => *n, + Range(range) => range.partition_count(), + } + } + + /// Returns true when two partitionings both satisfy their own distribution + /// requirements and can be paired by partition index. + /// + /// Use this for multi-input operators, such as partitioned joins, where + /// each child has a different schema, required [`Distribution`], and + /// expression-equivalence context. + /// + /// ```text + /// # co-partitioned: each side satisfies its own requirement, and boundaries match + /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) + /// right: Range(right.x ASC, [10, 20]), required KeyPartitioned(right.x) + /// + /// # not compatible: right side does not satisfy a hash-specific requirement + /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) + /// right: Range(right.x ASC, [10, 20]), required HashPartitioned(right.x) + /// + /// # not compatible: boundaries differ + /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) + /// right: Range(right.x ASC, [15, 20]), required KeyPartitioned(right.x) + /// ``` + pub fn co_partitioned_with( + &self, + required: &Distribution, + eq_properties: &EquivalenceProperties, + other: &Self, + other_required: &Distribution, + other_eq_properties: &EquivalenceProperties, + ) -> bool { + if !self + .satisfaction(required, eq_properties, false) + .is_satisfied() + || !other + .satisfaction(other_required, other_eq_properties, false) + .is_satisfied() + { + return false; + } + + if self.partition_count() == 1 && other.partition_count() == 1 { + return true; + } + + if self.partition_count() != other.partition_count() { + return false; + } + + match (self, other) { + (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, + (Partitioning::Range(left), Partitioning::Range(right)) => { + left.same_boundaries(right) + } + _ => false, } } @@ -189,6 +580,47 @@ impl Partitioning { }) } + fn key_expr_satisfaction( + partition_exprs: &[Arc], + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, + ) -> PartitioningSatisfaction { + if partition_exprs.is_empty() || required_exprs.is_empty() { + return PartitioningSatisfaction::NotSatisfied; + } + + if equivalent_exprs(required_exprs, partition_exprs, eq_properties) { + return PartitioningSatisfaction::Exact; + } + + if !allow_subset { + return PartitioningSatisfaction::NotSatisfied; + } + + let eq_groups = eq_properties.eq_group(); + if eq_groups.is_empty() { + if Self::is_subset_partitioning(partition_exprs, required_exprs) { + PartitioningSatisfaction::Subset + } else { + PartitioningSatisfaction::NotSatisfied + } + } else { + let normalized_partition_exprs = + normalize_exprs(partition_exprs, eq_properties); + let normalized_required_exprs = + normalize_exprs(required_exprs, eq_properties); + if Self::is_subset_partitioning( + &normalized_partition_exprs, + &normalized_required_exprs, + ) { + PartitioningSatisfaction::Subset + } else { + PartitioningSatisfaction::NotSatisfied + } + } + } + #[deprecated(since = "52.0.0", note = "Use satisfaction instead")] pub fn satisfy( &self, @@ -212,62 +644,54 @@ impl Partitioning { Distribution::SinglePartition if self.partition_count() == 1 => { PartitioningSatisfaction::Exact } - // When partition count is 1, hash requirement is satisfied. - Distribution::HashPartitioned(_) if self.partition_count() == 1 => { + // When partition count is 1, partitioned requirements are satisfied. + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) + if self.partition_count() == 1 => + { PartitioningSatisfaction::Exact } + Distribution::KeyPartitioned(required_exprs) => match self { + Partitioning::Hash(partition_exprs, _) => Self::key_expr_satisfaction( + partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ), + Partitioning::Range(range) => { + let partition_exprs = range + .ordering + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + Self::key_expr_satisfaction( + &partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ) + } + Partitioning::RoundRobinBatch(_) + | Partitioning::UnknownPartitioning(_) => { + PartitioningSatisfaction::NotSatisfied + } + }, Distribution::HashPartitioned(required_exprs) => match self { // Here we do not check the partition count for hash partitioning and assumes the partition count // and hash functions in the system are the same. In future if we plan to support storage partition-wise joins, // then we need to have the partition count and hash functions validation. - Partitioning::Hash(partition_exprs, _) => { - // Empty hash partitioning is invalid - if partition_exprs.is_empty() || required_exprs.is_empty() { - return PartitioningSatisfaction::NotSatisfied; - } - - // Fast path: exact match - if physical_exprs_equal(required_exprs, partition_exprs) { - return PartitioningSatisfaction::Exact; - } - - // Normalization path using equivalence groups - let eq_groups = eq_properties.eq_group(); - if !eq_groups.is_empty() { - let normalized_required_exprs = required_exprs - .iter() - .map(|e| eq_groups.normalize_expr(Arc::clone(e))) - .collect::>(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|e| eq_groups.normalize_expr(Arc::clone(e))) - .collect::>(); - if physical_exprs_equal( - &normalized_required_exprs, - &normalized_partition_exprs, - ) { - return PartitioningSatisfaction::Exact; - } - - if allow_subset - && Self::is_subset_partitioning( - &normalized_partition_exprs, - &normalized_required_exprs, - ) - { - return PartitioningSatisfaction::Subset; - } - } else if allow_subset - && Self::is_subset_partitioning(partition_exprs, required_exprs) - { - return PartitioningSatisfaction::Subset; - } - + Partitioning::Hash(partition_exprs, _) => Self::key_expr_satisfaction( + partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ), + Partitioning::RoundRobinBatch(_) + | Partitioning::UnknownPartitioning(_) => { PartitioningSatisfaction::NotSatisfied } - _ => PartitioningSatisfaction::NotSatisfied, + Partitioning::Range(_) => PartitioningSatisfaction::NotSatisfied, }, - _ => PartitioningSatisfaction::NotSatisfied, + Distribution::SinglePartition => PartitioningSatisfaction::NotSatisfied, } } @@ -277,19 +701,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 +740,7 @@ impl PartialEq for Partitioning { { true } + (Partitioning::Range(left), Partitioning::Range(right)) => left == right, _ => false, } } @@ -322,9 +757,31 @@ pub enum Distribution { /// Requires children to be distributed in such a way that the same /// values of the keys end up in the same partition HashPartitioned(Vec>), + /// Requires rows with equal values for the given keys to be colocated in + /// the same partition, without requiring a specific partitioning algorithm. + /// + /// Unlike [`Self::HashPartitioned`], this can be satisfied by non-hash + /// partitioning such as range partitioning. A partitioning on a subset of + /// these keys can also satisfy this requirement because rows equal on all + /// required keys are also equal on any subset. + /// + /// For multi-input operators, satisfaction alone is not enough: each input + /// may satisfy its own key requirement while using incompatible partition + /// boundaries. Use [`Partitioning::co_partitioned_with`] before pairing + /// partitions by index. + KeyPartitioned(Vec>), } impl Distribution { + /// Returns key expressions for distribution variants that require + /// co-locating equal key values. + pub fn key_exprs(&self) -> Option<&[Arc]> { + match self { + Self::HashPartitioned(exprs) | Self::KeyPartitioned(exprs) => Some(exprs), + Self::UnspecifiedDistribution | Self::SinglePartition => None, + } + } + /// Creates a `Partitioning` that satisfies this `Distribution` pub fn create_partitioning(self, partition_count: usize) -> Partitioning { match self { @@ -332,7 +789,7 @@ impl Distribution { Partitioning::UnknownPartitioning(partition_count) } Distribution::SinglePartition => Partitioning::UnknownPartitioning(1), - Distribution::HashPartitioned(expr) => { + Distribution::HashPartitioned(expr) | Distribution::KeyPartitioned(expr) => { Partitioning::Hash(expr, partition_count) } } @@ -347,6 +804,9 @@ impl Display for Distribution { Distribution::HashPartitioned(exprs) => { write!(f, "HashPartitioned[{}])", format_physical_expr_list(exprs)) } + Distribution::KeyPartitioned(exprs) => { + write!(f, "KeyPartitioned[{}])", format_physical_expr_list(exprs)) + } } } } @@ -356,56 +816,164 @@ 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>, + eq_properties: EquivalenceProperties, + } + + impl PartitioningTestFixture { + fn new(fields: Vec<(&str, DataType)>) -> Result { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, data_type)| Field::new(*name, data_type.clone(), false)) + .collect::>(), + )); + let cols = fields + .iter() + .map(|(name, _)| { + Ok(Arc::new(Column::new_with_schema(name, &schema)?) + as Arc) + }) + .collect::>()?; + let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + Ok(Self { + schema, + cols, + eq_properties, + }) + } + + fn int64(names: &[&str]) -> Result { + Self::new(names.iter().map(|name| (*name, DataType::Int64)).collect()) + } + + fn col(&self, index: usize) -> Arc { + Arc::clone(&self.cols[index]) + } + + fn cols( + &self, + indices: impl IntoIterator, + ) -> Vec> { + indices.into_iter().map(|index| self.col(index)).collect() + } + + fn hash_partitioning( + &self, + indices: impl IntoIterator, + partition_count: usize, + ) -> Partitioning { + Partitioning::Hash(self.cols(indices), partition_count) + } + + fn hash_distribution( + &self, + indices: impl IntoIterator, + ) -> Distribution { + Distribution::HashPartitioned(self.cols(indices)) + } + + fn key_partitioned_distribution( + &self, + indices: impl IntoIterator, + ) -> Distribution { + Distribution::KeyPartitioned(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, + ) -> 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, + split_points: Vec, + ) -> RangePartitioning { + RangePartitioning::try_new(self.range_ordering(indices), split_points) + .expect("test range partitioning should be valid") + } + + fn range_partitioning( + &self, + indices: impl IntoIterator, + split_points: Vec, + ) -> 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"), + ) + } + } + #[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> = 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> = 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]), + Distribution::KeyPartitioned(fixture.cols([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(), ); @@ -419,6 +987,9 @@ mod tests { Distribution::HashPartitioned(_) => { assert_eq!(result, (true, false, false, true, false)) } + Distribution::KeyPartitioned(_) => { + assert_eq!(result, (true, false, false, true, false)) + } } } @@ -427,72 +998,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 = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - 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 +1041,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 +1059,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 = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - 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 +1088,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 +1106,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 = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - 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 +1119,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 +1137,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 = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - 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 +1159,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 +1177,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 = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - 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 +1199,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 +1217,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 = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; let unknown: Arc = 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 +1247,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 +1265,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 = - 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 +1294,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 +1309,392 @@ mod tests { Ok(()) } + + #[test] + fn key_partitioned_satisfaction_is_exact() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let hash_a = fixture.hash_partitioning([0], 2); + let hash_ab = fixture.hash_partitioning([0, 1], 2); + let range_a = fixture + .range_partitioning([0], vec![int_split_point([10]), int_split_point([20])]); + let range_ab = + fixture.range_partitioning([0, 1], vec![int_split_point([10, 20])]); + + let exact_cases = [ + ( + "Hash([a]) vs KeyPartitioned([a])", + &hash_a, + fixture.key_partitioned_distribution([0]), + ), + ( + "Range([a]) vs KeyPartitioned([a])", + &range_a, + fixture.key_partitioned_distribution([0]), + ), + ]; + for (desc, partitioning, requirement) in exact_cases { + for allow_subset in [true, false] { + assert_eq!( + partitioning.satisfaction( + &requirement, + &fixture.eq_properties, + allow_subset, + ), + PartitioningSatisfaction::Exact, + "Failed for {desc} with allow_subset={allow_subset}" + ); + } + } + + let subset_cases = [ + ( + "Hash([a]) vs KeyPartitioned([a, b])", + &hash_a, + fixture.key_partitioned_distribution([0, 1]), + ), + ( + "Range([a]) vs KeyPartitioned([a, b])", + &range_a, + fixture.key_partitioned_distribution([0, 1]), + ), + ]; + for (desc, partitioning, requirement) in subset_cases { + assert_eq!( + partitioning.satisfaction(&requirement, &fixture.eq_properties, true), + PartitioningSatisfaction::Subset, + "Failed for {desc} with subset enabled" + ); + assert_eq!( + partitioning.satisfaction(&requirement, &fixture.eq_properties, false), + PartitioningSatisfaction::NotSatisfied, + "Failed for {desc} with subset disabled" + ); + } + + let not_satisfied_cases = [ + ( + "Range([a]) vs KeyPartitioned([b])", + &range_a, + fixture.key_partitioned_distribution([1]), + ), + ( + "Hash([a, b]) vs KeyPartitioned([a])", + &hash_ab, + fixture.key_partitioned_distribution([0]), + ), + ( + "Range([a, b]) vs KeyPartitioned([a])", + &range_ab, + fixture.key_partitioned_distribution([0]), + ), + ]; + for (desc, partitioning, requirement) in not_satisfied_cases { + for allow_subset in [true, false] { + assert_eq!( + partitioning.satisfaction( + &requirement, + &fixture.eq_properties, + allow_subset, + ), + PartitioningSatisfaction::NotSatisfied, + "Failed for {desc} with allow_subset={allow_subset}" + ); + } + } + + Ok(()) + } + + fn int_split_point(values: impl IntoIterator) -> SplitPoint { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::Int64(Some(value))) + .collect(), + ) + } + + fn assert_range_try_new_error( + ordering: LexOrdering, + split_points: Vec, + 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 = 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 range_partitionings_are_co_partitioned_by_boundaries() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let left = fixture + .range_partitioning([0], vec![int_split_point([10]), int_split_point([20])]); + let right_same_map = fixture + .range_partitioning([1], vec![int_split_point([10]), int_split_point([20])]); + let right_different_split = fixture + .range_partitioning([1], vec![int_split_point([15]), int_split_point([20])]); + let right_desc = fixture.range_partitioning_with_ordering( + [fixture.range_sort_expr(1, SortOptions::new(true, false))].into(), + vec![int_split_point([20]), int_split_point([10])], + ); + + let test_cases = [ + ( + "same boundaries with matching key requirements", + fixture.key_partitioned_distribution([0]), + right_same_map.clone(), + fixture.key_partitioned_distribution([1]), + true, + ), + ( + "different split points", + fixture.key_partitioned_distribution([0]), + right_different_split, + fixture.key_partitioned_distribution([1]), + false, + ), + ( + "different sort options", + fixture.key_partitioned_distribution([0]), + right_desc, + fixture.key_partitioned_distribution([1]), + false, + ), + ( + "range cannot satisfy hash requirement", + fixture.hash_distribution([0]), + right_same_map, + fixture.key_partitioned_distribution([1]), + false, + ), + ]; + for (desc, left_requirement, right, right_requirement, expected) in test_cases { + assert_eq!( + left.co_partitioned_with( + &left_requirement, + &fixture.eq_properties, + &right, + &right_requirement, + &fixture.eq_properties, + ), + expected, + "Failed for {desc}" + ); + } + + Ok(()) + } + + #[test] + fn co_partitioned_with_rejects_subset_key_satisfaction() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let left = fixture + .range_partitioning([0], vec![int_split_point([10]), int_split_point([20])]); + let right = fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); + + assert_eq!( + right.satisfaction( + &fixture.key_partitioned_distribution([0]), + &fixture.eq_properties, + false, + ), + PartitioningSatisfaction::NotSatisfied + ); + assert_eq!( + left.satisfaction( + &fixture.key_partitioned_distribution([0, 1]), + &fixture.eq_properties, + true, + ), + PartitioningSatisfaction::Subset + ); + assert!(!left.co_partitioned_with( + &fixture.key_partitioned_distribution([0, 1]), + &fixture.eq_properties, + &right, + &fixture.key_partitioned_distribution([0]), + &fixture.eq_properties, + )); + + Ok(()) + } + + #[test] + fn hash_partitionings_are_co_partitioned_by_count() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let left = fixture.hash_partitioning([0], 2); + + let test_cases = [ + ( + "same partition count", + fixture.hash_partitioning([1], 2), + fixture.key_partitioned_distribution([1]), + true, + ), + ( + "different partition count", + fixture.hash_partitioning([1], 3), + fixture.key_partitioned_distribution([1]), + false, + ), + ( + "mixed hash and range partitioning", + fixture.range_partitioning([1], vec![int_split_point([10])]), + fixture.key_partitioned_distribution([1]), + false, + ), + ]; + for (desc, right, right_requirement, expected) in test_cases { + assert_eq!( + left.co_partitioned_with( + &fixture.key_partitioned_distribution([0]), + &fixture.eq_properties, + &right, + &right_requirement, + &fixture.eq_properties, + ), + expected, + "Failed for {desc}" + ); + } + + 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-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index c522867c05196..9d50f07d6a3a6 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -867,6 +867,8 @@ fn add_roundrobin_on_top( /// current executor is less than this value. Partition number will be increased. /// * `allow_subset_satisfy_partitioning`: Whether to allow subset partitioning logic in satisfaction checks. /// Set to `false` for partitioned hash joins to ensure exact hash matching. +/// * `force_to_target`: Whether to repartition even when the hash expressions +/// are already satisfied but the partition count differs from `n_target`. /// /// # Returns /// @@ -877,6 +879,7 @@ fn add_hash_on_top( hash_exprs: Vec>, n_target: usize, allow_subset_satisfy_partitioning: bool, + force_to_target: bool, ) -> Result { // Early return if hash repartition is unnecessary // `RepartitionExec: partitioning=Hash([...], 1), input_partitions=1` is unnecessary. @@ -885,6 +888,7 @@ fn add_hash_on_top( } let dist = Distribution::HashPartitioned(hash_exprs); + let current_partitions = input.plan.output_partitioning().partition_count(); let satisfaction = input.plan.output_partitioning().satisfaction( &dist, input.plan.equivalence_properties(), @@ -894,11 +898,12 @@ fn add_hash_on_top( // Add hash repartitioning when: // - When subset satisfaction is enabled (current >= threshold): only repartition if not satisfied // - When below threshold (current < threshold): repartition if expressions don't match OR to increase parallelism - let needs_repartition = if allow_subset_satisfy_partitioning { + let needs_repartition = if force_to_target { + !satisfaction.is_satisfied() || n_target != current_partitions + } else if allow_subset_satisfy_partitioning { !satisfaction.is_satisfied() } else { - !satisfaction.is_satisfied() - || n_target > input.plan.output_partitioning().partition_count() + !satisfaction.is_satisfied() || n_target > current_partitions }; if needs_repartition { @@ -1092,6 +1097,32 @@ struct RepartitionRequirementStatus { roundrobin_beneficial_stats: bool, /// Designates whether hash partitioning is necessary. hash_necessary: bool, + /// Designates whether hash repartitioning should force the target + /// partition count even when the hash expressions are already satisfied. + force_hash_to_target: bool, +} + +#[derive(Debug, Clone, Copy, Default)] +struct PartitionedJoinDistribution { + /// Inner partitioned hash join children can be paired by existing Range + /// partitions, so hash repartitioning is not needed. + compatible_range: bool, + /// Partitioned join children have different partition counts. If hash + /// repartitioning is used, both sides must be forced to the target count. + needs_count_alignment: bool, +} + +fn requirement_includes_grouping_id(requirement: &Distribution) -> bool { + // Grouping set aggregates (ROLLUP, CUBE, GROUPING SETS) require exact hash + // partitioning on all group columns including __grouping_id to ensure + // partial aggregates from different partitions are correctly combined. + requirement.key_exprs().is_some_and(|exprs| { + exprs.iter().any(|expr| { + (expr.as_ref() as &dyn Any) + .downcast_ref::() + .is_some_and(|col| col.name() == Aggregate::INTERNAL_GROUPING_ID) + }) + }) } /// Calculates the `RepartitionRequirementStatus` for each children to generate @@ -1134,6 +1165,7 @@ fn get_repartition_requirement_status( let children = plan.children(); let rr_beneficial = plan.benefits_from_input_partitioning(); let requirements = plan.required_input_distribution(); + let join_distribution = partitioned_join_distribution(plan); let mut repartition_status_flags = vec![]; for (child, requirement, roundrobin_beneficial) in izip!(children.into_iter(), requirements, rr_beneficial) @@ -1146,30 +1178,39 @@ fn get_repartition_requirement_status( Precision::Inexact(n_rows) => !should_use_estimates || (n_rows > batch_size), Precision::Absent => true, }; - let is_hash = matches!(requirement, Distribution::HashPartitioned(_)); - // Hash re-partitioning is necessary when the input has more than one - // partitions: + let is_partitioned_requirement = requirement.key_exprs().is_some(); + // Hash repartitioning may be necessary when the input has more than one + // partition, or when repartitioning one sibling requires aligning all + // key-partitioned siblings. let multi_partitions = child.output_partitioning().partition_count() > 1; let roundrobin_sensible = roundrobin_beneficial && roundrobin_beneficial_stats; - needs_alignment |= is_hash && (multi_partitions || roundrobin_sensible); + needs_alignment |= is_partitioned_requirement + && !join_distribution.compatible_range + && (multi_partitions || roundrobin_sensible); repartition_status_flags.push(( - is_hash, + is_partitioned_requirement, RepartitionRequirementStatus { requirement, roundrobin_beneficial, roundrobin_beneficial_stats, - hash_necessary: is_hash && multi_partitions, + hash_necessary: is_partitioned_requirement + && multi_partitions + && !join_distribution.compatible_range, + // Hash satisfaction checks key expressions, not matching + // partition counts, so force repartition when join sides differ. + force_hash_to_target: is_partitioned_requirement + && join_distribution.needs_count_alignment, }, )); } - // Align hash necessary flags for hash partitions to generate consistent + // Align hash necessary flags for key partitions to generate consistent // hash partitions at each children: if needs_alignment { - // When there is at least one hash requirement that is necessary or - // beneficial according to statistics, make all children require hash - // repartitioning: - for (is_hash, status) in &mut repartition_status_flags { - if *is_hash { + // When there is at least one key-partitioned requirement that is necessary + // or beneficial according to statistics, make all key-partitioned + // children require hash repartitioning: + for (is_partitioned_requirement, status) in &mut repartition_status_flags { + if *is_partitioned_requirement { status.hash_necessary = true; } } @@ -1180,6 +1221,59 @@ fn get_repartition_requirement_status( .collect()) } +/// Returns distribution state for a partitioned join's children. +/// +/// This is optimizer policy: partitioned joins require children that can be +/// paired by partition index. Inner hash joins can reuse compatible range +/// partitioning; otherwise the existing hash repartitioning policy applies. +fn partitioned_join_distribution( + plan: &Arc, +) -> PartitionedJoinDistribution { + let Some(hash_join) = plan.downcast_ref::() else { + return Default::default(); + }; + + if hash_join.mode != PartitionMode::Partitioned { + return Default::default(); + } + + let children = plan.children(); + let [left, right] = children.as_slice() else { + return Default::default(); + }; + let needs_count_alignment = left.output_partitioning().partition_count() + != right.output_partitioning().partition_count(); + + let requirements = plan.required_input_distribution(); + let left_partitioning = left.output_partitioning(); + let right_partitioning = right.output_partitioning(); + let compatible_range = match requirements.as_slice() { + [ + left_requirement @ Distribution::KeyPartitioned(_), + right_requirement @ Distribution::KeyPartitioned(_), + ] if hash_join.join_type == JoinType::Inner + && matches!( + (left_partitioning, right_partitioning), + (Partitioning::Range(_), Partitioning::Range(_)) + ) => + { + left_partitioning.co_partitioned_with( + left_requirement, + left.equivalence_properties(), + right_partitioning, + right_requirement, + right.equivalence_properties(), + ) + } + _ => false, + }; + + PartitionedJoinDistribution { + compatible_range, + needs_count_alignment, + } +} + /// This function checks whether we need to add additional data exchange /// operators to satisfy distribution requirements. Since this function /// takes care of such requirements, we should avoid manually adding data @@ -1301,6 +1395,7 @@ pub fn ensure_distribution( roundrobin_beneficial, roundrobin_beneficial_stats, hash_necessary, + force_hash_to_target, }, )| { let increases_partition_count = @@ -1319,18 +1414,6 @@ pub fn ensure_distribution( // 3. Not a grouping set aggregate (requires exact hash including __grouping_id) let current_partitions = child.plan.output_partitioning().partition_count(); - // Check if the hash partitioning requirement includes __grouping_id column. - // Grouping set aggregates (ROLLUP, CUBE, GROUPING SETS) require exact hash - // partitioning on all group columns including __grouping_id to ensure partial - // aggregates from different partitions are correctly combined. - let requires_grouping_id = matches!(&requirement, Distribution::HashPartitioned(exprs) - if exprs.iter().any(|expr| { - (expr.as_ref() as &dyn Any) - .downcast_ref::() - .is_some_and(|col| col.name() == Aggregate::INTERNAL_GROUPING_ID) - }) - ); - let allow_subset_satisfy_partitioning = (current_partitions >= subset_satisfaction_threshold // `preserve_file_partitions` exposes existing file-group @@ -1340,7 +1423,7 @@ pub fn ensure_distribution( || (config.optimizer.preserve_file_partitions > 0 && current_partitions < target_partitions)) && !is_partitioned_join - && !requires_grouping_id; + && !requirement_includes_grouping_id(&requirement); // When `repartition_file_scans` is set, attempt to increase // parallelism at the source. @@ -1361,7 +1444,8 @@ pub fn ensure_distribution( Distribution::SinglePartition => { child = add_merge_on_top(child); } - Distribution::HashPartitioned(exprs) => { + Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) => { // See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background // When inserting hash is necessary to satisfy hash requirement, insert hash repartition. if hash_necessary { @@ -1370,6 +1454,7 @@ pub fn ensure_distribution( exprs.to_vec(), target_partitions, allow_subset_satisfy_partitioning, + force_hash_to_target, )?; } } @@ -1428,7 +1513,9 @@ pub fn ensure_distribution( // no ordering requirement match requirement { // Operator requires specific distribution. - Distribution::SinglePartition | Distribution::HashPartitioned(_) => { + Distribution::SinglePartition + | Distribution::HashPartitioned(_) + | Distribution::KeyPartitioned(_) => { // If the parent doesn't maintain input order, preserving // ordering is pointless. However, if it does maintain // input order, we keep order-preserving variants so diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 24eb3af5f564c..679d1bc7b2e36 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -268,8 +268,9 @@ impl ExecutionPlan for OutputRequirementExec { requirements = OrderingRequirements::new_alternatives(updated_reqs, soft); } - let dist_req = match &self.required_input_distribution()[0] { - Distribution::HashPartitioned(exprs) => { + let dist_req = { + let dist_req = &self.required_input_distribution()[0]; + if let Some(exprs) = dist_req.key_exprs() { let mut updated_exprs = vec![]; for expr in exprs { let Some(new_expr) = update_expr(expr, projection.expr(), false)? @@ -278,9 +279,18 @@ impl ExecutionPlan for OutputRequirementExec { }; updated_exprs.push(new_expr); } - Distribution::HashPartitioned(updated_exprs) + match dist_req { + Distribution::HashPartitioned(_) => { + Distribution::HashPartitioned(updated_exprs) + } + Distribution::KeyPartitioned(_) => { + Distribution::KeyPartitioned(updated_exprs) + } + _ => unreachable!(), + } + } else { + dist_req.clone() } - dist => dist.clone(), }; make_with_child(projection, &self.input()).map(|input| { diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index 40c6245d894d4..96900490e0921 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -31,7 +31,9 @@ use datafusion_common::plan_err; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported}; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion_physical_plan::joins::SymmetricHashJoinExec; +use datafusion_physical_plan::joins::{ + HashJoinExec, PartitionMode, SymmetricHashJoinExec, +}; use datafusion_physical_plan::{ExecutionPlanProperties, get_plan_string}; use crate::PhysicalOptimizerRule; @@ -178,6 +180,50 @@ pub fn check_plan_sanity( } } + check_partitioned_join_distribution(plan)?; + + Ok(()) +} + +fn check_partitioned_join_distribution(plan: &Arc) -> Result<()> { + let Some(hash_join) = plan.downcast_ref::() else { + return Ok(()); + }; + + if hash_join.mode != PartitionMode::Partitioned { + return Ok(()); + } + + let children = plan.children(); + let requirements = plan.required_input_distribution(); + let ([left, right], [left_req, right_req]) = + (children.as_slice(), requirements.as_slice()) + else { + return plan_err!( + "Invalid HashJoinExec: expected two children and two distribution requirements" + ); + }; + + if !left.output_partitioning().co_partitioned_with( + left_req, + left.equivalence_properties(), + right.output_partitioning(), + right_req, + right.equivalence_properties(), + ) { + let plan_str = get_plan_string(plan); + return plan_err!( + "Plan: {:?} does not satisfy partitioned join co-partitioning requirements: \ + left requirement: {}, left output partitioning: {}; \ + right requirement: {}, right output partitioning: {}", + plan_str, + left_req, + left.output_partitioning(), + right_req, + right.output_partitioning() + ); + } + Ok(()) } diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 50eac566d90ef..9fc3725fc6457 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -161,8 +161,13 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { check_default_invariants(self, check) } - /// Specifies the data distribution requirements for all the - /// children for this `ExecutionPlan`, By default it's [[Distribution::UnspecifiedDistribution]] for each child, + /// Specifies the data distribution requirements for all the children for + /// this `ExecutionPlan`. + /// + /// By default, each child has [`Distribution::UnspecifiedDistribution`]. + /// Multi-input operators that use [`Distribution::KeyPartitioned`] must + /// use [`Partitioning::co_partitioned_with`] to verify that satisfied + /// children can be paired by partition index. fn required_input_distribution(&self) -> Vec { vec![Distribution::UnspecifiedDistribution; self.children().len()] } diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index f7391feb29cc3..17d72a856f93b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -24,8 +24,8 @@ use std::vec; use crate::ExecutionPlanProperties; use crate::execution_plan::{ - EmissionType, boundedness_from_children, has_same_children_properties, - stub_properties, + EmissionType, InvariantLevel, boundedness_from_children, check_default_invariants, + has_same_children_properties, stub_properties, }; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, @@ -845,21 +845,71 @@ impl HashJoinExec { return false; } - // `preserve_file_partitions` can report Hash partitioning for Hive-style - // file groups, but those partitions are not actually hash-distributed. - // Partitioned dynamic filters rely on hash routing, so disable them in - // this mode to avoid incorrect results. Follow-up work: enable dynamic - // filtering for preserve_file_partitioned scans (issue #20195). - // https://github.com/apache/datafusion/issues/20195 - if config.optimizer.preserve_file_partitions > 0 - && self.mode == PartitionMode::Partitioned - { + // Bounds and membership filters derived from the build side do not + // account for null-equal matching: a probe-side NULL key evaluates + // such predicates to NULL and would be pruned, even though it can + // match a build-side NULL when nulls compare equal. + if self.null_equality == NullEquality::NullEqualsNull { return false; } + if self.mode == PartitionMode::Partitioned { + // `preserve_file_partitions` can report Hash partitioning for + // Hive-style file groups, but those partitions are not actually + // hash-distributed. Partitioned dynamic filters rely on hash + // routing, so disable them in this mode to avoid incorrect + // results. Follow-up work: enable dynamic filtering for + // preserve_file_partitioned scans (issue #20195). + // https://github.com/apache/datafusion/issues/20195 + if config.optimizer.preserve_file_partitions > 0 { + return false; + } + + // Partitioned dynamic filters route probe rows with + // `hash(join_key) % partition_count`. That is only valid when + // partition ids are hash buckets with the same bucket count, or + // when there is only one partition and no routing choice exists. + // This also rejects non-hash partitioning such as + // `Partitioning::Range`. + if !self.has_partitioned_dynamic_filter_routing() { + return false; + } + } + true } + fn has_partitioned_dynamic_filter_routing(&self) -> bool { + match ( + self.left.output_partitioning(), + self.right.output_partitioning(), + ) { + ( + Partitioning::Hash(_, left_partition_count), + Partitioning::Hash(_, right_partition_count), + ) => left_partition_count == right_partition_count, + (left_partitioning, right_partitioning) => { + left_partitioning.partition_count() == 1 + && right_partitioning.partition_count() == 1 + } + } + } + + fn partitioned_children_co_partitioned(&self) -> bool { + let requirements = self.required_input_distribution(); + let [left_requirement, right_requirement] = requirements.as_slice() else { + return false; + }; + + self.left.output_partitioning().co_partitioned_with( + left_requirement, + self.left.equivalence_properties(), + self.right.output_partitioning(), + right_requirement, + self.right.equivalence_properties(), + ) + } + /// left (build) side which gets hashed pub fn left(&self) -> &Arc { &self.left @@ -1198,6 +1248,21 @@ impl ExecutionPlan for HashJoinExec { "HashJoinExec" } + fn check_invariants(&self, check: InvariantLevel) -> Result<()> { + check_default_invariants(self, check)?; + + if matches!(check, InvariantLevel::Executable) + && self.mode == PartitionMode::Partitioned + && !self.partitioned_children_co_partitioned() + { + return plan_err!( + "Invalid HashJoinExec, partitioned children are not co-partitioned, consider using RepartitionExec" + ); + } + + Ok(()) + } + fn properties(&self) -> &Arc { &self.cache } @@ -1214,10 +1279,17 @@ impl ExecutionPlan for HashJoinExec { .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), - ] + if self.join_type == JoinType::Inner { + vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ] + } else { + vec![ + Distribution::HashPartitioned(left_expr), + Distribution::HashPartitioned(right_expr), + ] + } } PartitionMode::Auto => vec![ Distribution::UnspecifiedDistribution, @@ -1286,6 +1358,12 @@ impl ExecutionPlan for HashJoinExec { consider using RepartitionExec" ); + assert_or_internal_err!( + self.mode != PartitionMode::Partitioned + || self.partitioned_children_co_partitioned(), + "Invalid HashJoinExec, partitioned children are not co-partitioned, consider using RepartitionExec" + ); + assert_or_internal_err!( self.mode != PartitionMode::CollectLeft || left_partitions == 1, "Invalid HashJoinExec, the output partition count of the left child must be 1 in CollectLeft mode,\ @@ -2119,6 +2197,88 @@ mod tests { Ok((left_schema, right_schema, on)) } + #[derive(Debug)] + struct PartitionedTestInput { + input: Arc, + cache: Arc, + } + + impl PartitionedTestInput { + fn new(input: Arc, partitioning: Partitioning) -> Self { + let cache = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + partitioning, + input.pipeline_behavior(), + input.boundedness(), + )); + Self { input, cache } + } + } + + impl DisplayAs for PartitionedTestInput { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "PartitionedTestInput") + } + DisplayFormatType::TreeRender => write!(f, ""), + } + } + } + + impl ExecutionPlan for PartitionedTestInput { + fn name(&self) -> &'static str { + "PartitionedTestInput" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!( + "PartitionedTestInput expected one child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new( + Arc::clone(&children[0]), + self.cache.output_partitioning().clone(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } + } + + fn declared_hash_test_input( + input: Arc, + key: &str, + key_index: usize, + partition_count: usize, + ) -> Result> { + Ok(Arc::new(PartitionedTestInput::new( + input, + Partitioning::Hash( + vec![Arc::new(Column::new(key, key_index))], + partition_count, + ), + ))) + } + use crate::coalesce_partitions::CoalescePartitionsExec; use crate::joins::hash_join::stream::lookup_join_hashmap; use crate::test::{TestMemoryExec, assert_join_metrics}; @@ -2142,6 +2302,7 @@ mod tests { use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint}; use hashbrown::HashTable; use insta::{allow_duplicates, assert_snapshot}; use rstest::*; @@ -5375,6 +5536,7 @@ mod tests { None, ) .unwrap(); + let left = declared_hash_test_input(left, "b1", 1, 2)?; let right_batch = build_table_i32( ("a2", &vec![10, 11]), ("b2", &vec![12, 13]), @@ -5386,6 +5548,7 @@ mod tests { None, ) .unwrap(); + let right = declared_hash_test_input(right, "b2", 1, 2)?; let on = vec![( Arc::new(Column::new_with_schema("b1", &left_batch.schema())?) as _, Arc::new(Column::new_with_schema("b2", &right_batch.schema())?) as _, @@ -5413,8 +5576,8 @@ mod tests { let task_ctx = Arc::new(task_ctx); let join = HashJoinExec::try_new( - Arc::clone(&left) as Arc, - Arc::clone(&right) as Arc, + Arc::clone(&left), + Arc::clone(&right), on.clone(), None, &join_type, @@ -5717,6 +5880,8 @@ mod tests { Arc::clone(&child_right_schema), None, )?; + let child_left = declared_hash_test_input(child_left, "child_key", 1, 4)?; + let child_right = declared_hash_test_input(child_right, "child_right_key", 1, 4)?; let parent_left: Arc = TestMemoryExec::try_new_exec( &[ vec![build_table_i32( @@ -5735,6 +5900,7 @@ mod tests { Arc::clone(&parent_left_schema), None, )?; + let parent_left = declared_hash_test_input(parent_left, "parent_key", 1, 4)?; let child_on = vec![( Arc::new(Column::new_with_schema("child_key", &child_left_schema)?) as _, @@ -6341,6 +6507,140 @@ mod tests { Ok(()) } + #[test] + fn test_dynamic_filter_pushdown_rejects_null_equal_join() -> Result<()> { + let (_, _, on) = build_schema_and_on()?; + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![1])); + let right = build_table(("a2", &vec![1]), ("b1", &vec![1]), ("c2", &vec![1])); + + let session_config = join_dynamic_filter_session_config(0); + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::RightSemi, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + )?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + fn range_partitioned_test_input( + schema: SchemaRef, + range_key: &str, + ) -> Result> { + let input = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let range_expr = Arc::new(Column::new_with_schema(range_key, &schema)?); + let range_partitioning = Partitioning::Range(RangePartitioning::new( + [PhysicalSortExpr::new_default(range_expr)].into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )); + RepartitionExec::try_new(input, range_partitioning) + .map(|exec| Arc::new(exec) as _) + } + + fn hash_partitioned_test_input( + schema: SchemaRef, + hash_key: &str, + partition_count: usize, + ) -> Result> { + let input = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let hash_expr = Arc::new(Column::new_with_schema(hash_key, &schema)?); + RepartitionExec::try_new( + input, + Partitioning::Hash(vec![hash_expr], partition_count), + ) + .map(|exec| Arc::new(exec) as _) + } + + fn join_dynamic_filter_session_config( + preserve_file_partitions: usize, + ) -> SessionConfig { + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + session_config + .options_mut() + .optimizer + .preserve_file_partitions = preserve_file_partitions; + session_config + } + + fn partitioned_inner_hash_join( + left: Arc, + right: Arc, + on: JoinOn, + ) -> Result { + HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + } + + #[test] + fn dynamic_filter_rejects_range_partitioning() -> Result<()> { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left = range_partitioned_test_input(left_schema, "b1")?; + let right = range_partitioned_test_input(right_schema, "b1")?; + + let session_config = join_dynamic_filter_session_config(0); + let join = partitioned_inner_hash_join(left, right, on)?; + + assert!(matches!( + join.required_input_distribution().as_slice(), + [ + Distribution::KeyPartitioned(_), + Distribution::KeyPartitioned(_) + ] + )); + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn dynamic_filter_rejects_preserve_file_partitions() -> Result<()> { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left = hash_partitioned_test_input(left_schema, "b1", 2)?; + let right = hash_partitioned_test_input(right_schema, "b1", 2)?; + + let session_config = join_dynamic_filter_session_config(1); + let join = partitioned_inner_hash_join(left, right, on)?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn dynamic_filter_rejects_mismatched_hash_counts() -> Result<()> { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left = hash_partitioned_test_input(left_schema, "b1", 2)?; + let right = hash_partitioned_test_input(right_schema, "b1", 3)?; + + let session_config = join_dynamic_filter_session_config(0); + let join = partitioned_inner_hash_join(left, right, on)?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + #[test] fn test_with_dynamic_filter_rejects_invalid_columns() -> Result<()> { let (_, _, on) = build_schema_and_on()?; diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index b4aa295562b67..3fb090a406286 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -144,6 +144,12 @@ pub fn adjust_right_output_partitioning( .collect::>()?; Partitioning::Hash(new_exprs, *size) } + // Range partitioning can satisfy join input requirements, but range + // output propagation needs broader join semantics coverage. + // https://github.com/apache/datafusion/issues/22395 + Partitioning::Range(range) => { + Partitioning::UnknownPartitioning(range.partition_count()) + } result => result.clone(), }; Ok(result) diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 3005e975424b4..c7b1d4729e21d 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 465ca4a99e961..2a8005c10ae00 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 f715de0b5964b..929ff4f7dfc85 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 09570f14ba734..eb9b5f09aa3ed 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/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 15a744f5c3500..b727c670152b9 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -1420,13 +1420,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; } @@ -1436,6 +1445,7 @@ message Partitioning { uint64 round_robin = 1; PhysicalHashRepartition hash = 2; uint64 unknown = 3; + PhysicalRangePartitioning range = 4; } } diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index bd859f2c080b7..4ee2bd30b2c77 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -16138,6 +16138,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() @@ -16154,6 +16157,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { "roundRobin", "hash", "unknown", + "range", ]; #[allow(clippy::enum_variant_names)] @@ -16161,6 +16165,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { RoundRobin, Hash, Unknown, + Range, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -16185,6 +16190,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)), } } @@ -16226,6 +16232,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 { @@ -19810,6 +19823,207 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { deserializer.deserialize_struct("datafusion.PhysicalPlanNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalRangePartitioning { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + 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(deserializer: D) -> std::result::Result + 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(deserializer: D) -> std::result::Result + 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(self, value: &str) -> std::result::Result + 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(self, mut map_: V) -> std::result::Result + 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(&self, serializer: S) -> std::result::Result + 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(deserializer: D) -> std::result::Result + 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(deserializer: D) -> std::result::Result + 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(self, value: &str) -> std::result::Result + 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(self, mut map_: V) -> std::result::Result + 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(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index ea25dba4abd8c..e82582bd3c1aa 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -2100,14 +2100,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, + #[prost(message, repeated, tag = "2")] + pub split_point: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalRangeSplitPoint { + #[prost(message, repeated, tag = "1")] + pub value: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct RepartitionExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - /// 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, #[prost(bool, tag = "6")] @@ -2115,7 +2127,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, } /// Nested message and enum types in `Partitioning`. @@ -2128,6 +2140,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 62f989a111ed2..b5b2652659401 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}; @@ -51,7 +53,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; @@ -674,6 +678,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, @@ -685,6 +697,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 { + 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::>()?; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) +} + +fn parse_protobuf_range_split_point( + split_point: &protobuf::PhysicalRangeSplitPoint, +) -> Result { + let values = split_point + .value + .iter() + .map(|value| ScalarValue::try_from(value).map_err(Into::into)) + .collect::>()?; + Ok(SplitPoint::new(values)) +} + pub fn parse_protobuf_file_scan_schema( proto: &protobuf::FileScanExecConf, ) -> Result> { diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 73fca1bbe6070..785c4df95fbaf 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -44,7 +44,9 @@ use datafusion_physical_plan::expressions::{ use datafusion_physical_plan::joins::{HashExpr, HashTableLookupExpr}; 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, @@ -658,6 +660,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, @@ -667,6 +674,40 @@ pub fn serialize_partitioning( Ok(serialized_partitioning) } +fn serialize_range_partitioning( + range: &RangePartitioning, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, +) -> Result { + 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::>()?, + }) +} + +fn serialize_range_split_point( + split_point: &SplitPoint, +) -> Result { + Ok(protobuf::PhysicalRangeSplitPoint { + value: split_point + .values() + .iter() + .map(|value| { + TryInto::::try_into(value) + .map_err(Into::into) + }) + .collect::>()?, + }) +} + fn serialize_when_then_expr( when_expr: &Arc, then_expr: &Arc, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index a0b7180ba90bc..19bc50933d3f0 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -89,7 +89,8 @@ use datafusion::physical_plan::windows::{ }; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning, - PhysicalExpr, PlanProperties, SendableRecordBatchStream, Statistics, displayable, + PhysicalExpr, PlanProperties, RangePartitioning, SendableRecordBatchStream, + SplitPoint, Statistics, displayable, }; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion::scalar::ScalarValue; @@ -1995,6 +1996,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); From 9baa581fe3ee86f27626ca5f724a44dd9db20f8d Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Thu, 2 Jul 2026 14:09:34 -0400 Subject: [PATCH 02/22] Aggregations Support `Partitioning::Range` (#23239) - Closes #23191. - Related discussion: #23184, #23236. Range partitioning can satisfy aggregate hash partitioning: equal group keys are already partitioned, even though the partitioning is not hash-based. This is the first unary-operator implementation from the range partitioning discussion before making broader public API changes around `HashPartitioned` / `KeyPartitioned`. - Let compatible range partitioning satisfy aggregate hash distribution requirements in `EnforceDistribution` - Keep this private to aggregate planning for now to not make public API changes to `Distribution` enum variants yet until more operators are supported Yes. Yes. Range-partitioned aggregate plans can now avoid hash repartitioning. --- .../enforce_distribution.rs | 80 ++++++++++++++++ .../src/enforce_distribution.rs | 92 ++++++++++++++++--- .../physical-optimizer/src/sanity_checker.rs | 25 ++++- datafusion/physical-optimizer/src/utils.rs | 78 +++++++++++++++- 4 files changed, 257 insertions(+), 18 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 4ff438fdddfca..1f53e79485d27 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -933,6 +933,86 @@ fn non_inner_range_join_rehashes() -> Result<()> { Ok(()) } +#[test] +fn range_aggregate_reuses_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let aggregate = + aggregate_exec_with_alias(input, vec![("a".to_string(), "a".to_string())]); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(aggregate, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] + AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + " + ); + + Ok(()) +} + +#[test] +fn range_grouping_set_aggregate_rehashes_with_grouping_id() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20], + SortOptions::default(), + )?); + let input_schema = input.schema(); + let group_by = PhysicalGroupBy::new( + vec![ + (col("a", &input_schema)?, "a".to_string()), + (col("b", &input_schema)?, "b".to_string()), + ], + vec![ + (lit(ScalarValue::Int64(None)), "a".to_string()), + (lit(ScalarValue::Int64(None)), "b".to_string()), + ], + vec![vec![false, true], vec![false, false]], + true, + ); + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + input, + Arc::clone(&input_schema), + )?); + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + group_by.as_final(), + vec![], + vec![], + Arc::clone(&partial) as _, + partial.schema(), + )?); + + let plan = TestConfig::default() + .with_query_execution_partitions(3) + .to_plan(aggregate, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, __grouping_id@2 as __grouping_id], aggr=[] + RepartitionExec: partitioning=Hash([a@0, b@1, __grouping_id@2], 3), input_partitions=3 + AggregateExec: mode=Partial, gby=[(a@0 as a, NULL as b), (a@0 as a, b@1 as b)], aggr=[] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20)], 3) + " + ); + + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index 9d50f07d6a3a6..b501d6ef8537d 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -28,8 +28,9 @@ use std::sync::Arc; use crate::optimizer::PhysicalOptimizerRule; use crate::output_requirements::OutputRequirementExec; use crate::utils::{ - add_sort_above_with_check, is_coalesce_partitions, is_repartition, - is_sort_preserving_merge, + add_sort_above_with_check, aggregate_can_reuse_range_partitioning, + is_coalesce_partitions, is_repartition, is_sort_preserving_merge, + range_partitioning_satisfies_key_partitioning, }; use arrow::compute::SortOptions; @@ -926,6 +927,49 @@ fn add_hash_on_top( Ok(input) } +// TODO: remove this private helper once Range generally satisfies +// KeyPartitioned requirements through Partitioning::satisfaction. +// See . +// +// Partial aggregates do not require key partitioning, but they preserve their +// input partitioning for the final aggregate. Until Range satisfies +// KeyPartitioned generally, this check keeps preserve_file_partitions from +// inserting RoundRobin between a reusable Range input and the partial aggregate. +fn partial_aggregate_preserves_reusable_partitioning( + plan: &Arc, + child: &Arc, + allow_subset_satisfy_partitioning: bool, +) -> bool { + let Some(aggregate) = plan.downcast_ref::() else { + return false; + }; + if aggregate.mode() != &AggregateMode::Partial + || aggregate.group_expr().is_empty() + || aggregate.group_expr().has_grouping_set() + { + return false; + } + + let group_exprs = aggregate.group_expr().input_exprs(); + let output_partitioning = child.output_partitioning(); + let eq_properties = child.equivalence_properties(); + let key_distribution = Distribution::KeyPartitioned(group_exprs.clone()); + + output_partitioning + .satisfaction( + &key_distribution, + eq_properties, + allow_subset_satisfy_partitioning, + ) + .is_satisfied() + || range_partitioning_satisfies_key_partitioning( + output_partitioning, + &group_exprs, + eq_properties, + allow_subset_satisfy_partitioning, + ) +} + /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator /// on top of the given plan node to satisfy a single partition requirement /// while preserving ordering constraints. @@ -1398,21 +1442,14 @@ pub fn ensure_distribution( force_hash_to_target, }, )| { - let increases_partition_count = - child.plan.output_partitioning().partition_count() < target_partitions; - - let add_roundrobin = enable_round_robin - // Operator benefits from partitioning (e.g. filter): - && roundrobin_beneficial - && roundrobin_beneficial_stats - // Unless partitioning increases the partition count, it is not beneficial: - && increases_partition_count; - // Allow subset satisfaction when: // 1. Current partition count >= threshold // 2. Not a partitioned join since must use exact hash matching for joins // 3. Not a grouping set aggregate (requires exact hash including __grouping_id) let current_partitions = child.plan.output_partitioning().partition_count(); + let preserve_file_partition_threshold_met = + config.optimizer.preserve_file_partitions > 0 + && current_partitions >= config.optimizer.preserve_file_partitions; let allow_subset_satisfy_partitioning = (current_partitions >= subset_satisfaction_threshold @@ -1420,11 +1457,29 @@ pub fn ensure_distribution( // partitioning to the optimizer. Respect it when the only // reason to repartition would be to increase partition count // beyond the preserved file-group count. - || (config.optimizer.preserve_file_partitions > 0 + || (preserve_file_partition_threshold_met && current_partitions < target_partitions)) && !is_partitioned_join && !requirement_includes_grouping_id(&requirement); + let increases_partition_count = current_partitions < target_partitions; + + let preserve_partial_aggregate_partitioning = + preserve_file_partition_threshold_met + && partial_aggregate_preserves_reusable_partitioning( + &plan, + &child.plan, + allow_subset_satisfy_partitioning, + ); + + let add_roundrobin = enable_round_robin + // Operator benefits from partitioning (e.g. filter): + && roundrobin_beneficial + && roundrobin_beneficial_stats + // Unless partitioning increases the partition count, it is not beneficial: + && increases_partition_count + && !preserve_partial_aggregate_partitioning; + // When `repartition_file_scans` is set, attempt to increase // parallelism at the source. // @@ -1446,9 +1501,18 @@ pub fn ensure_distribution( } Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs) => { + let range_satisfied_for_aggregate = + aggregate_can_reuse_range_partitioning(&plan) + && range_partitioning_satisfies_key_partitioning( + child.plan.output_partitioning(), + exprs, + child.plan.equivalence_properties(), + allow_subset_satisfy_partitioning, + ); + // See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background // When inserting hash is necessary to satisfy hash requirement, insert hash repartition. - if hash_necessary { + if hash_necessary && !range_satisfied_for_aggregate { child = add_hash_on_top( child, exprs.to_vec(), diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index 96900490e0921..08c014593cdc2 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use datafusion_common::Result; +use datafusion_physical_expr::Distribution; use datafusion_physical_plan::ExecutionPlan; use datafusion_common::config::{ConfigOptions, OptimizerOptions}; @@ -37,6 +38,9 @@ use datafusion_physical_plan::joins::{ use datafusion_physical_plan::{ExecutionPlanProperties, get_plan_string}; use crate::PhysicalOptimizerRule; +use crate::utils::{ + aggregate_can_reuse_range_partitioning, range_partitioning_satisfies_key_partitioning, +}; use datafusion_physical_expr_common::sort_expr::format_physical_sort_requirement_list; use itertools::izip; @@ -164,11 +168,26 @@ pub fn check_plan_sanity( } } - if !child + let child_satisfies_distribution = child .output_partitioning() .satisfaction(&dist_req, child_eq_props, true) - .is_satisfied() - { + .is_satisfied(); + let range_satisfies_aggregate_distribution = + aggregate_can_reuse_range_partitioning(plan) + && match &dist_req { + Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) => { + range_partitioning_satisfies_key_partitioning( + child.output_partitioning(), + exprs, + child_eq_props, + true, + ) + } + _ => false, + }; + + if !(child_satisfies_distribution || range_satisfies_aggregate_distribution) { let plan_str = get_plan_string(plan); return plan_err!( "Plan: {:?} does not satisfy distribution requirements: {}. Child-{} output partitioning: {}", diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index a6b01637c970e..2f928224da28b 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -18,7 +18,11 @@ use std::sync::Arc; use datafusion_common::Result; -use datafusion_physical_expr::{LexOrdering, LexRequirement}; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, LexRequirement, Partitioning, PhysicalExpr, + physical_exprs_equal, +}; +use datafusion_physical_plan::aggregates::{AggregateExec, AggregateMode}; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -108,6 +112,78 @@ pub fn is_repartition(plan: &Arc) -> bool { plan.is::() } +/// TODO: remove once Range generally satisfies KeyPartitioned requirements +/// through Partitioning::satisfaction. +/// See . +/// +/// Checks whether range partitioning satisfies a key partitioning requirement. +/// This is intentionally separate from general partitioning satisfaction while +/// range reuse is rolled out operator by operator. +pub(crate) fn range_partitioning_satisfies_key_partitioning( + partitioning: &Partitioning, + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, +) -> bool { + match partitioning { + Partitioning::Range(range) => { + let partition_exprs = range + .ordering() + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + + if partition_exprs.is_empty() || required_exprs.is_empty() { + return false; + } + + let eq_group = eq_properties.eq_group(); + let normalized_partition_exprs = partition_exprs + .iter() + .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) + .collect::>(); + let normalized_required_exprs = required_exprs + .iter() + .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) + .collect::>(); + + if physical_exprs_equal( + &normalized_required_exprs, + &normalized_partition_exprs, + ) { + return true; + } + + allow_subset + && normalized_partition_exprs.len() < normalized_required_exprs.len() + && normalized_partition_exprs.iter().all(|partition_expr| { + normalized_required_exprs + .iter() + .any(|required_expr| partition_expr.eq(required_expr)) + }) + } + _ => false, + } +} + +/// TODO: remove once Range generally satisfies KeyPartitioned requirements +/// through Partitioning::satisfaction. +/// See . +/// +/// Checks whether an aggregate can reuse range partitioning to satisfy its key +/// partitioning requirement. +pub(crate) fn aggregate_can_reuse_range_partitioning( + plan: &Arc, +) -> bool { + plan.downcast_ref::() + .is_some_and(|aggregate| { + matches!( + aggregate.mode(), + AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned + ) && !aggregate.group_expr().has_grouping_set() + }) +} + /// Checks whether the given operator is a limit; /// i.e. either a [`LocalLimitExec`] or a [`GlobalLimitExec`]. pub fn is_limit(plan: &Arc) -> bool { From 5f1da9870c143ee531f0674f2ff50b0275db637a Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Fri, 29 May 2026 16:48:54 -0400 Subject: [PATCH 03/22] Add range partitioning sqllogictest fixture (#22607) ## Which issue does this PR close? - Part of #22397. - Discussion: #21992. ## Rationale for this change This adds a focused sqllogictest fixture for source-provided `Range` partitioning before changing optimizer behavior. It follows the direction discussed in #21992 and gives later planning PRs stable baselines for current behavior. ## What changes are included in this PR? - Registers a `range_partitioned` test table for `range_partitioning.slt`. - Adds a sqllogictest-only source wrapper that reports `Range` partitioning when `range_key` is projected, and `UnknownPartitioning` when it is not. - Adds baselines for grouping on the range key, grouping on a non-range key, joining on the range key, and `UNION ALL` over range-partitioned inputs. ## Are these changes tested? Yes. - `cargo fmt --all` - `cargo test -p datafusion-sqllogictest --test sqllogictests range_partitioning` - `cargo clippy --all-targets --all-features -- -D warnings` ## Are there any user-facing changes? No. This is test-only infrastructure and sqllogictest coverage. --- Cargo.lock | 1 + datafusion/physical-expr/src/partitioning.rs | 5 + datafusion/sqllogictest/Cargo.toml | 1 + datafusion/sqllogictest/src/test_context.rs | 8 + .../src/test_context/range_partitioning.rs | 250 ++++++++++++++++++ .../test_files/range_partitioning.slt | 134 ++++++++++ 6 files changed, 399 insertions(+) create mode 100644 datafusion/sqllogictest/src/test_context/range_partitioning.rs create mode 100644 datafusion/sqllogictest/test_files/range_partitioning.slt diff --git a/Cargo.lock b/Cargo.lock index c6e5e61a807dd..46b2145c3a63e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2591,6 +2591,7 @@ dependencies = [ "chrono", "clap", "datafusion", + "datafusion-datasource", "datafusion-spark", "datafusion-substrait", "env_logger", diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 77411ecb765a3..c152c16bc5a46 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -171,6 +171,11 @@ impl Display for Partitioning { /// Values equal to split point `i` belong to partition `i + 1`, so interior /// partitions are lower-inclusive and upper-exclusive. /// +/// Like other user-specified data properties such as sortedness, if a source +/// declares range partitioning, it is responsible for placing each row in the +/// partition described by the split points. DataFusion will not validate this is +/// upheld. +/// /// For a single range key: /// /// ```text diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index e2ffe1415a1fb..a642fbe22a6e3 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -47,6 +47,7 @@ bytes = { workspace = true, optional = true } chrono = { workspace = true, optional = true } clap = { version = "4.5.60", features = ["derive", "env"] } datafusion = { workspace = true, default-features = true, features = ["avro"] } +datafusion-datasource = { workspace = true } datafusion-spark = { workspace = true, features = ["core"] } datafusion-substrait = { workspace = true, default-features = true, optional = true } futures = { workspace = true } diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 0edde71b939f4..a83db2bfb947f 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -53,6 +53,8 @@ use datafusion::{ use datafusion_spark::SessionStateBuilderSpark; use crate::is_spark_path; +use range_partitioning::register_range_partitioned_table; + use async_trait::async_trait; use datafusion::common::cast::as_float64_array; use datafusion::execution::SessionStateBuilder; @@ -61,6 +63,8 @@ use log::info; use sqlparser::ast; use tempfile::TempDir; +mod range_partitioning; + /// Context for running tests pub struct TestContext { /// Context for running queries @@ -167,6 +171,10 @@ impl TestContext { info!("Registering table with many types"); register_table_with_many_types(test_ctx.session_ctx()).await; } + "range_partitioning.slt" => { + info!("Registering range partitioned table"); + register_range_partitioned_table(test_ctx.session_ctx()); + } "metadata.slt" | "arrow_field.slt" => { info!("Registering metadata table tables"); register_metadata_tables(test_ctx.session_ctx()).await; diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs new file mode 100644 index 0000000000000..88e49708baf60 --- /dev/null +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -0,0 +1,250 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt; +use std::sync::Arc; + +use arrow::array::Int32Array; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use async_trait::async_trait; +use datafusion::catalog::Session; +use datafusion::common::{Result, ScalarValue, project_schema}; +use datafusion::datasource::source::{DataSource, DataSourceExec}; +use datafusion::datasource::{TableProvider, TableType}; +use datafusion::execution::context::TaskContext; +use datafusion::logical_expr::Expr; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_expr::expressions::col as physical_col; +use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion::physical_plan::execution_plan::SchedulingType; +use datafusion::physical_plan::projection::ProjectionExprs; +use datafusion::physical_plan::{ + DisplayFormatType, ExecutionPlan, Partitioning, RangePartitioning, + SendableRecordBatchStream, SplitPoint, Statistics, +}; +use datafusion::prelude::SessionContext; +use datafusion_datasource::memory::MemorySourceConfig; + +// ============================================================================== +// Range Partitioned Table (sqllogictest-only) +// ============================================================================== + +/// Simple range-partitioned table for testing before declaring such tables is +/// supported via SQL. +#[derive(Debug)] +struct RangePartitionedTable { + schema: SchemaRef, + partitions: Vec>, + range_column_index: usize, + split_points: Vec, +} + +#[async_trait] +impl TableProvider for RangePartitionedTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + let projected_schema = project_schema(&self.schema, projection)?; + let mut source = MemorySourceConfig::try_new( + &self.partitions, + Arc::clone(&self.schema), + projection.cloned(), + )?; + source = source.with_show_sizes(state.config_options().explain.show_sizes); + + let output_partitioning = + self.output_partitioning(projection, &projected_schema)?; + let source = RangePartitionedSource { + inner: source, + output_partitioning, + }; + + Ok(DataSourceExec::from_data_source(source)) + } +} + +impl RangePartitionedTable { + fn output_partitioning( + &self, + projection: Option<&Vec>, + projected_schema: &SchemaRef, + ) -> Result { + let Some(projected_range_index) = + projected_index(self.range_column_index, projection) + else { + return Ok(Partitioning::UnknownPartitioning(self.partitions.len())); + }; + + let range_column = projected_schema.field(projected_range_index).name(); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + physical_col(range_column, projected_schema)?, + SortOptions::default(), + )]) + .expect("range ordering should not be empty"); + + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + self.split_points.clone(), + )?)) + } +} + +fn projected_index( + column_index: usize, + projection: Option<&Vec>, +) -> Option { + projection + .map(|projection| projection.iter().position(|idx| *idx == column_index)) + .unwrap_or(Some(column_index)) +} + +#[derive(Clone, Debug)] +struct RangePartitionedSource { + inner: MemorySourceConfig, + output_partitioning: Partitioning, +} + +impl DataSource for RangePartitionedSource { + fn open( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.open(partition, context) + } + + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + self.inner.fmt_as(t, f)?; + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, ", output_partitioning={}", self.output_partitioning) + } + DisplayFormatType::TreeRender => Ok(()), + } + } + + fn output_partitioning(&self) -> Partitioning { + self.output_partitioning.clone() + } + + fn eq_properties(&self) -> EquivalenceProperties { + self.inner.eq_properties() + } + + fn scheduling_type(&self) -> SchedulingType { + self.inner.scheduling_type() + } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.inner.partition_statistics(partition) + } + + fn with_fetch(&self, limit: Option) -> Option> { + Some(Arc::new(Self { + inner: self.inner.clone().with_limit(limit), + output_partitioning: self.output_partitioning.clone(), + })) + } + + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn try_swapping_with_projection( + &self, + _projection: &ProjectionExprs, + ) -> Result>> { + // Range partitioning metadata is projection-sensitive. This fixture + // computes it in TableProvider::scan, so do not rewrite later + // ProjectionExec nodes into the source. + Ok(None) + } +} + +pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { + let schema = Arc::new(Schema::new(vec![ + Field::new("range_key", DataType::Int32, false), + Field::new("non_range_key", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + let partitions = vec![ + vec![range_partition_batch(&schema, &[1, 5], &[1, 2], &[10, 50])], + vec![range_partition_batch( + &schema, + &[10, 15], + &[1, 2], + &[100, 150], + )], + vec![range_partition_batch( + &schema, + &[20, 25], + &[1, 2], + &[200, 250], + )], + vec![range_partition_batch( + &schema, + &[30, 35], + &[1, 2], + &[300, 350], + )], + ]; + let split_points = vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ]; + let table = RangePartitionedTable { + schema, + partitions, + range_column_index: 0, + split_points, + }; + + ctx.register_table("range_partitioned", Arc::new(table)) + .expect("range partitioned table registration should succeed"); +} + +fn range_partition_batch( + schema: &SchemaRef, + range_key: &[i32], + non_range_key: &[i32], + value: &[i32], +) -> RecordBatch { + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(range_key.to_vec())), + Arc::new(Int32Array::from(non_range_key.to_vec())), + Arc::new(Int32Array::from(value.to_vec())), + ], + ) + .expect("range partition batch should be valid") +} diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt new file mode 100644 index 0000000000000..a61f17a039eb8 --- /dev/null +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -0,0 +1,134 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# The sqllogictest harness registers range_partitioned(range_key, non_range_key, value) +# as an in-memory source with four physical source partitions: +# +# partition 0: range_key in [..., 10), rows (1, 1, 10), (5, 2, 50) +# partition 1: range_key in [10, 20), rows (10, 1, 100), (15, 2, 150) +# partition 2: range_key in [20, 30), rows (20, 1, 200), (25, 2, 250) +# partition 3: range_key in [30, ...), rows (30, 1, 300), (35, 2, 350) + +statement ok +set datafusion.explain.physical_plan_only = true; + +########## +# TEST 1: Aggregate on Range Partition Column +# Scanning range_key preserves source Range partitioning metadata. +# Planning still inserts Hash repartitioning today; later optimizer PRs can +# use this baseline to show when the repartition is removed. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) + +query II +SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 2: Aggregate on Non-Range Column +# Projecting away range_key means the scan output no longer contains the +# expression needed to describe range partitioning, so it reports +# UnknownPartitioning with the same partition count. +########## + +query TT +EXPLAIN SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1], output_partitioning=UnknownPartitioning(4) + +query II +SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key ORDER BY non_range_key; +---- +1 610 +2 800 + + +########## +# TEST 3: Join on Range Partition Column +# Both inputs expose Range partitioning on range_key. Join planning currently +# reaches the unsupported Range output-partitioning path; later optimizer PRs +# can replace this baseline with a successful plan and result test. +########## + +query error This feature is not implemented: Join output partitioning with range partitioning is not implemented +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; + +########## +# TEST 4: Union of Range Partitioned Inputs +# Each input exposes Range partitioning on range_key. This records current +# UNION ALL behavior before later PRs decide whether compatible range inputs can +# preserve Range partitioning across the union. +########## + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +03)--DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +ORDER BY range_key, value; +---- +1 10 +1 10 +5 50 +5 50 +10 100 +10 100 +15 150 +15 150 +20 200 +20 200 +25 250 +25 250 +30 300 +30 300 +35 350 +35 350 + +statement ok +reset datafusion.explain.physical_plan_only; From 3598943e7ae17019b147f6b9cd56802f1da89341 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Wed, 10 Jun 2026 01:35:19 -0400 Subject: [PATCH 04/22] Add logical range partitioning representation (#22777) - Closes #22778. - Related: #21992, #22395. - Needed by #22657. Declared scan output partitioning should use logical partitioning metadata, not physical partitioning types. This adds logical range partitioning so range-partitioned sources can declare their layout at the logical layer. - Add logical `Partitioning::Range` and `RangePartitioning`. - Move `SplitPoint` and shared split-point validation to `datafusion-common`. - Wire logical range partitioning through expression traversal, rewrites, and display. - Keep planning, logical proto, and Substrait support explicitly unsupported for now. Yes. Unit tests added Yes. This adds public logical range partitioning API. No breaking API changes. --- datafusion/common/src/lib.rs | 2 + datafusion/common/src/partitioning.rs | 104 +++++++ datafusion/core/src/physical_planner.rs | 66 +++- datafusion/expr/src/logical_plan/display.rs | 17 ++ datafusion/expr/src/logical_plan/mod.rs | 4 +- datafusion/expr/src/logical_plan/plan.rs | 283 +++++++++++++++++- datafusion/expr/src/logical_plan/tree_node.rs | 15 + datafusion/physical-expr/src/lib.rs | 3 +- datafusion/physical-expr/src/partitioning.rs | 142 +-------- datafusion/proto/src/logical_plan/mod.rs | 5 + .../logical_plan/producer/rel/exchange_rel.rs | 10 + 11 files changed, 507 insertions(+), 144 deletions(-) create mode 100644 datafusion/common/src/partitioning.rs diff --git a/datafusion/common/src/lib.rs b/datafusion/common/src/lib.rs index e865c548bb554..2f6d9848b6e55 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -30,6 +30,7 @@ mod dfschema; mod functional_dependencies; mod join_type; mod param_value; +mod partitioning; mod schema_reference; mod table_reference; mod unnest; @@ -92,6 +93,7 @@ pub use join_type::{JoinConstraint, JoinSide, JoinType}; pub use nested_struct::cast_column; pub use null_equality::NullEquality; pub use param_value::ParamValues; +pub use partitioning::{SplitPoint, validate_range_split_points}; pub use scalar::{ScalarType, ScalarValue}; pub use schema_reference::SchemaReference; pub use spans::{Location, Span, Spans}; diff --git a/datafusion/common/src/partitioning.rs b/datafusion/common/src/partitioning.rs new file mode 100644 index 0000000000000..8a7212c2e3089 --- /dev/null +++ b/datafusion/common/src/partitioning.rs @@ -0,0 +1,104 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::utils::compare_rows; +use crate::{Result, ScalarValue, error::_plan_err}; +use arrow::compute::SortOptions; +use std::cmp::Ordering; +use std::fmt::{self, Display}; + +/// A boundary between adjacent range partitions. +/// +/// A split point is a tuple with one [`ScalarValue`] per partitioning +/// expression. Split points are interpreted lexicographically according to the +/// ordering of the range partitioning that owns them. +/// +/// `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. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct SplitPoint { + values: Vec, +} + +impl SplitPoint { + /// Creates a new split point from its tuple values. + pub fn new(values: Vec) -> 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::>() + .join(", "); + write!(f, "({values})") + } +} + +/// Validates that split points match the ordering width and are strictly +/// ordered according to the provided sort options. +pub fn validate_range_split_points( + split_points: &[SplitPoint], + sort_options: &[SortOptions], +) -> Result<()> { + let width = sort_options.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_rows( + split_points[0].values(), + split_points[1].values(), + sort_options, + )? != 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(()) +} diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 43af743caa030..070c49907b606 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -98,7 +98,7 @@ use datafusion_physical_expr::aggregate::{ }; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::{ - LexOrdering, PhysicalSortExpr, create_physical_sort_exprs, + LexOrdering, PhysicalSortExpr, RangePartitioning, create_physical_sort_exprs, }; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::empty::EmptyExec; @@ -1260,6 +1260,22 @@ impl DefaultPhysicalPlanner { .collect::>>()?; Partitioning::Hash(runtime_expr, *n) } + LogicalPartitioning::Range(range) => { + let sort_exprs = create_physical_sort_exprs( + range.ordering(), + input_dfschema, + execution_props, + )?; + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!( + "Range repartitioning requires non-empty ordering" + ) + })?; + Partitioning::Range(RangePartitioning::try_new( + ordering, + range.split_points().to_vec(), + )?) + } LogicalPartitioning::DistributeBy(_) => { return not_impl_err!( "Physical plan does not support DistributeBy partitioning" @@ -3231,8 +3247,8 @@ mod tests { use arrow_schema::{FieldRef, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::{ - DFSchemaRef, ScalarValue, TableReference, ToDFSchema as _, assert_batches_eq, - assert_contains, + DFSchemaRef, ScalarValue, SplitPoint, TableReference, ToDFSchema as _, + assert_batches_eq, assert_contains, }; use datafusion_execution::TaskContext; use datafusion_execution::runtime_env::RuntimeEnv; @@ -3241,8 +3257,8 @@ mod tests { use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::{ Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder, - Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, - WindowFunctionDefinition, col, lit, + RangePartitioning, Signature, TableSource, UserDefinedLogicalNodeCore, + Volatility, WindowFunctionDefinition, col, lit, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; @@ -3290,6 +3306,46 @@ mod tests { Field::new(name, DataType::Int64, nullable) } + #[tokio::test] + async fn logical_range_repartition_plans_output_partitioning() -> Result<()> { + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )])?; + let table = Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]])?); + let source = Arc::new(DefaultTableSource::new(table)); + let logical_plan = LogicalPlanBuilder::scan("test", source, None)? + .repartition(LogicalPartitioning::Range(RangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?))? + .build()?; + + let planner = DefaultPhysicalPlanner::default(); + let physical_plan = planner + .create_initial_plan(&logical_plan, &make_session_state()) + .await?; + let repartition = physical_plan + .as_ref() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!( + "expected RepartitionExec, got {}", + physical_plan.name() + ) + })?; + let Partitioning::Range(range) = repartition.partitioning() else { + return internal_err!( + "expected Range target partitioning, got {:?}", + repartition.partitioning() + ); + }; + assert_eq!(range.partition_count(), 2); + assert_eq!(physical_plan.output_partitioning().partition_count(), 2); + + Ok(()) + } + #[test] fn test_create_window_expr_unwraps_alias_with_metadata() -> Result<()> { use std::collections::HashMap; diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 58c7feb616179..27b86a6d8cdd5 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -515,6 +515,23 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { "Partitioning Key": hash_expr }) } + Partitioning::Range(range) => { + let range_expr: Vec = + range.ordering().iter().map(|e| format!("{e}")).collect(); + let split_points: Vec = range + .split_points() + .iter() + .map(|e| format!("{e}")) + .collect(); + + json!({ + "Node Type": "Repartition", + "Partitioning Scheme": "Range", + "Partition Count": range.partition_count(), + "Partitioning Key": range_expr, + "Split Points": split_points + }) + } Partitioning::DistributeBy(expr) => { let dist_by_expr: Vec = expr.iter().map(|e| format!("{e}")).collect(); diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index c2b01868c97f3..5da6d97d64403 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -41,8 +41,8 @@ pub use plan::{ Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection, - RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, Subquery, - SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, + RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, + Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, projection_schema, }; pub use statement::{ diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index c572b202f03ce..a5530af914164 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -50,6 +50,7 @@ use crate::{ WindowFunctionDefinition, build_join_schema, expr_vec_fmt, requalify_sides_if_needed, }; +use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use datafusion_common::cse::{NormalizeEq, Normalizeable}; use datafusion_common::format::ExplainFormat; @@ -60,10 +61,12 @@ use datafusion_common::tree_node::{ use datafusion_common::{ Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, Dependency, FunctionalDependence, FunctionalDependencies, NullEquality, ParamValues, Result, - ScalarValue, Spans, TableReference, UnnestOptions, aggregate_functional_dependencies, - assert_eq_or_internal_err, assert_or_internal_err, internal_err, plan_err, + ScalarValue, Spans, SplitPoint, TableReference, UnnestOptions, + aggregate_functional_dependencies, assert_eq_or_internal_err, assert_or_internal_err, + internal_err, plan_err, validate_range_split_points, }; use indexmap::IndexSet; +use itertools::Itertools as _; // backwards compatibility use crate::display::PgJsonVisitor; @@ -864,6 +867,32 @@ impl LogicalPlan { input: Arc::new(input), })) } + Partitioning::Range(range) => { + if expr.len() != range.ordering().len() { + return internal_err!( + "Incorrect number of expressions for Range partitioning" + ); + } + let input = self.only_input(inputs)?; + let ordering = range + .ordering() + .iter() + .zip(expr) + .map(|(sort_expr, expr)| SortExpr { + expr, + asc: sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }) + .collect(); + let range = RangePartitioning::try_new( + ordering, + range.split_points().to_vec(), + )?; + Ok(LogicalPlan::Repartition(Repartition { + partitioning_scheme: Partitioning::Range(range), + input: Arc::new(input), + })) + } Partitioning::DistributeBy(_) => { let input = self.only_input(inputs)?; Ok(LogicalPlan::Repartition(Repartition { @@ -2091,6 +2120,9 @@ impl LogicalPlan { n ) } + Partitioning::Range(range) => { + write!(f, "Repartition: {range}") + } Partitioning::DistributeBy(expr) => { let dist_by_expr: Vec = expr.iter().map(|e| format!("{e}")).collect(); @@ -4137,11 +4169,16 @@ impl Debug for Subquery { } } -/// Logical partitioning schemes supported by [`LogicalPlan::Repartition`] +/// Logical partitioning schemes. /// -/// See [`Partitioning`] for more details on partitioning +/// A scheme can describe either requested repartitioning in +/// [`LogicalPlan::Repartition`] or a partitioning property declared by a source. +/// Some schemes are only valid as metadata until planner support is added. /// -/// [`Partitioning`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/enum.Partitioning.html# +/// For physical execution partitioning, see +/// [`datafusion_physical_expr::Partitioning`]. +/// +/// [`datafusion_physical_expr::Partitioning`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/enum.Partitioning.html# #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub enum Partitioning { /// Allocate batches using a round-robin algorithm and the specified number of partitions @@ -4149,10 +4186,118 @@ pub enum Partitioning { /// Allocate rows based on a hash of one of more expressions and the specified number /// of partitions. Hash(Vec, usize), + /// Partition rows by ranges. + /// See [`RangePartitioning`] for the logical contract. + Range(RangePartitioning), /// The DISTRIBUTE BY clause is used to repartition the data based on the input expressions DistributeBy(Vec), } +impl Partitioning { + /// Return the number of partitions, if known. + pub fn partition_count(&self) -> Option { + match self { + Self::RoundRobinBatch(partition_count) | Self::Hash(_, partition_count) => { + Some(*partition_count) + } + Self::Range(range) => Some(range.partition_count()), + Self::DistributeBy(_) => None, + } + } +} + +/// Logical range partitioning. +/// +/// [`RangePartitioning`] describes an ordered logical key space with split points. +/// +/// - `ordering` defines the partitioning key and ordering using logical +/// [`SortExpr`]s. +/// - `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 ordered +/// according to that ordering, and each split point must have one value per +/// ordering expression. See [`SplitPoint`] for the shared boundary contract. +/// +/// The expressions are resolved against the declaring plan's schema. This +/// constructor does not validate split point value types against the resolved +/// expression types. Like other user-specified data properties such as +/// sortedness, if a source declares range partitioning, it is responsible for +/// placing each row in the partition described by the split points. DataFusion +/// will not validate this is upheld. +/// +/// NOTE: Range-aware optimizer and execution behavior will be introduced +/// incrementally. See +/// . +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct RangePartitioning { + /// Ordered logical partitioning key. + ordering: Vec, + /// Boundaries between adjacent partitions. + split_points: Vec, +} + +impl RangePartitioning { + /// Creates logical range partitioning metadata and validates split point + /// shape and ordering. + pub fn try_new( + ordering: Vec, + split_points: Vec, + ) -> Result { + if ordering.is_empty() { + return plan_err!("Range partitioning requires non-empty ordering"); + } + + validate_range_split_points(&split_points, &logical_sort_options(&ordering))?; + + Ok(Self { + ordering, + split_points, + }) + } + + /// Return the number of partitions. + pub fn partition_count(&self) -> usize { + self.split_points.len() + 1 + } + + /// Returns the ordering that defines the range key. + pub fn ordering(&self) -> &[SortExpr] { + &self.ordering + } + + /// Returns the ordered split points between partitions. + pub fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } +} + +fn logical_sort_options(ordering: &[SortExpr]) -> Vec { + ordering + .iter() + .map(|sort_expr| SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }) + .collect() +} + +impl Display for RangePartitioning { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let ordering = self.ordering().iter().map(ToString::to_string).join(", "); + let split_points = self + .split_points() + .iter() + .map(ToString::to_string) + .join(", "); + write!( + f, + "Range([{ordering}], [{split_points}], {})", + self.partition_count() + ) + } +} + /// Represent the unnesting operation on a list column, such as the recursion depth and /// the output column name after unnesting /// @@ -4490,6 +4635,134 @@ mod tests { ]) } + fn i32_split_point(value: i32) -> SplitPoint { + SplitPoint::new(vec![ScalarValue::Int32(Some(value))]) + } + + fn null_i32_split_point() -> SplitPoint { + SplitPoint::new(vec![ScalarValue::Int32(None)]) + } + + #[test] + fn logical_range_partitioning_validates_shape() { + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10), i32_split_point(20)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + + let range = RangePartitioning::try_new( + vec![col("id").sort(false, true)], + vec![i32_split_point(20), i32_split_point(10)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + + let err = RangePartitioning::try_new(vec![], vec![]).unwrap_err(); + assert!(err.to_string().contains("non-empty ordering")); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true), col("salary").sort(true, true)], + vec![i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split point 0 has width 1, but ordering has width 2") + ); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(20), i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split points must be strictly ordered") + ); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10), i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split points must be strictly ordered") + ); + + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![null_i32_split_point(), i32_split_point(10)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + } + + #[test] + fn logical_partitioning_reports_known_partition_count() -> Result<()> { + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10)], + )?; + + assert_eq!(Partitioning::RoundRobinBatch(4).partition_count(), Some(4)); + assert_eq!( + Partitioning::Hash(vec![col("id")], 8).partition_count(), + Some(8) + ); + assert_eq!(Partitioning::Range(range).partition_count(), Some(2)); + assert_eq!( + Partitioning::DistributeBy(vec![col("id")]).partition_count(), + None + ); + + Ok(()) + } + + #[test] + fn logical_range_partitioning_participates_in_expression_rewrite() -> Result<()> { + let input = + table_scan(Some("employee_csv"), &employee_schema(), None)?.build()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(input), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10)], + )?), + }); + + let mut visited_exprs = vec![]; + plan.apply_expressions(|expr| { + visited_exprs.push(expr.to_string()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(visited_exprs, vec!["id"]); + + let plan = plan + .map_expressions(|expr| { + if expr == col("id") { + Ok(Transformed::yes(col("salary"))) + } else { + Ok(Transformed::no(expr)) + } + })? + .data; + + let LogicalPlan::Repartition(Repartition { + partitioning_scheme: Partitioning::Range(range), + .. + }) = plan + else { + unreachable!("expected range repartition"); + }; + assert_eq!(range.ordering()[0].expr, col("salary")); + assert_eq!(range.partition_count(), 2); + + Ok(()) + } + fn display_plan() -> Result { let plan1 = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3]))? .build()?; diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index ef9382a57209a..32396d2cd26f2 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -37,6 +37,7 @@ //! * [`LogicalPlan::with_new_exprs`]: Create a new plan with different expressions //! * [`LogicalPlan::expressions`]: Return a copy of the plan's expressions +use crate::logical_plan::plan::RangePartitioning; use crate::{ Aggregate, Analyze, CreateMemoryTable, CreateView, DdlStatement, Distinct, DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, Limit, @@ -414,6 +415,7 @@ impl LogicalPlan { Partitioning::Hash(expr, _) | Partitioning::DistributeBy(expr) => { expr.apply_elements(f) } + Partitioning::Range(range) => range.ordering().to_vec().apply_elements(f), Partitioning::RoundRobinBatch(_) => Ok(TreeNodeRecursion::Continue), }, LogicalPlan::Window(Window { window_expr, .. }) => { @@ -519,6 +521,19 @@ impl LogicalPlan { Partitioning::DistributeBy(expr) => expr .map_elements(f)? .update_data(Partitioning::DistributeBy), + Partitioning::Range(range) => { + let split_points = range.split_points().to_vec(); + range + .ordering() + .to_vec() + .map_elements(f)? + .map_data(|ordering| { + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) + })? + } Partitioning::RoundRobinBatch(_) => Transformed::no(partitioning_scheme), } .update_data(|partitioning_scheme| { diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index ad788c15d098e..2f5d5f0fcb460 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -55,10 +55,11 @@ pub mod execution_props { pub use aggregate::groups_accumulator::{GroupsAccumulatorAdapter, NullState}; pub use analysis::{AnalysisContext, ExprBoundaries, analyze}; +pub use datafusion_common::SplitPoint; pub use equivalence::{ AcrossPartitions, ConstExpr, EquivalenceProperties, calculate_union, }; -pub use partitioning::{Distribution, Partitioning, RangePartitioning, SplitPoint}; +pub use partitioning::{Distribution, Partitioning, RangePartitioning}; 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 c152c16bc5a46..8ac5de2a87b9d 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -21,10 +21,10 @@ use crate::{ EquivalenceProperties, PhysicalExpr, equivalence::ProjectionMapping, expressions::UnKnownColumn, physical_exprs_equal, }; -use datafusion_common::{Result, ScalarValue, plan_err}; +pub use datafusion_common::SplitPoint; +use datafusion_common::{Result, validate_range_split_points}; 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; @@ -156,20 +156,7 @@ impl Display for Partitioning { /// 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. +/// ordering expression. See [`SplitPoint`] for the shared boundary convention. /// /// Like other user-specified data properties such as sortedness, if a source /// declares range partitioning, it is responsible for placing each row in the @@ -217,39 +204,6 @@ pub struct RangePartitioning { split_points: Vec, } -/// 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, -} - -impl SplitPoint { - /// Creates a new split point from its tuple values. - pub fn new(values: Vec) -> 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::>() - .join(", "); - write!(f, "({values})") - } -} - impl RangePartitioning { /// Creates range partitioning metadata without validating split points. /// @@ -265,7 +219,13 @@ impl RangePartitioning { /// Creates range partitioning metadata and validates split point shape and /// ordering. pub fn try_new(ordering: LexOrdering, split_points: Vec) -> Result { - validate_range_split_points(&ordering, &split_points)?; + validate_range_split_points( + &split_points, + &ordering + .iter() + .map(|sort_expr| sort_expr.options) + .collect::>(), + )?; Ok(Self::new(ordering, split_points)) } @@ -369,86 +329,6 @@ fn format_range_split_points(split_points: &[SplitPoint]) -> String { .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 { - 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 { - 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 - }) - } - } -} - fn equivalent_exprs( left: &[Arc], right: &[Arc], @@ -825,7 +705,7 @@ mod tests { use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; - use datafusion_common::Result; + use datafusion_common::{Result, ScalarValue}; struct PartitioningTestFixture { schema: SchemaRef, diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 4e01ebabb5f69..41449112df8d0 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -1631,6 +1631,11 @@ impl AsLogicalPlan for LogicalPlanNode { Partitioning::RoundRobinBatch(partition_count) => { PartitionMethod::RoundRobin(*partition_count as u64) } + Partitioning::Range(_) => { + // TODO: Support range repartition protobuf serialization. + // Tracked by https://github.com/apache/datafusion/issues/22787 + return not_impl_err!("Range repartition"); + } Partitioning::DistributeBy(_) => { return not_impl_err!("DistributeBy"); } diff --git a/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs index 50c4b3da86cbe..1b9e91c7c475a 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs @@ -32,6 +32,11 @@ pub fn from_repartition( let partition_count = match repartition.partitioning_scheme { Partitioning::RoundRobinBatch(num) => num, Partitioning::Hash(_, num) => num, + Partitioning::Range(_) => { + // TODO: Support range repartitioning in Substrait exchange output. + // Tracked by https://github.com/apache/datafusion/issues/22788 + return not_impl_err!("Substrait does not support Range repartitioning"); + } Partitioning::DistributeBy(_) => { return not_impl_err!( "Physical plan does not support DistributeBy partitioning" @@ -50,6 +55,11 @@ pub fn from_repartition( .collect::>>()?; ExchangeKind::ScatterByFields(ScatterFields { fields }) } + Partitioning::Range(_) => { + // TODO: Support range repartitioning in Substrait exchange output. + // Tracked by https://github.com/apache/datafusion/issues/22788 + return not_impl_err!("Substrait does not support Range repartitioning"); + } Partitioning::DistributeBy(_) => { return not_impl_err!( "Physical plan does not support DistributeBy partitioning" From 51c8780fdc5befa30e2ec9ccc4385037a2603137 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Tue, 23 Jun 2026 15:54:31 -0400 Subject: [PATCH 05/22] Add `ListingOptions::output_partitioning` and `FileScanConfig::output_partitioning` for pre-defined file partitioning (#22657) - Closes #22645. This follows up on #22607 by replacing range-partitioning sqllogictest boilerplate with a general file/listing scan API for declared output partitioning. Related: #21992, #22607, https://github.com/apache/datafusion/pull/22607#discussion_r3323904683 - Add declared `output_partitioning` to file scan and listing table configuration. - Preserve declared partition counts during listing-table file grouping. - Serialize scan `output_partitioning` through physical plan proto. - Refactor `range_partitioning.slt` to use a CSV `ListingTable` instead of a custom test-only `TableProvider` / `DataSource`. Contract: - Declared partitioning expressions are written against the full table schema before scan projection. For example, `Range([range_key@0], [(10), (20)], 3)` remains valid if the scan projects `range_key` and falls back to `UnknownPartitioning(3)` if `range_key` is not projected. - Listing tables create one file group per declared output partition (which can exceed `target_partitions`). It is up to the user to plan their partitioning. For example, a 4-partition range declaration creates four scan file groups, adding empty trailing groups when fewer files are present. - File group index is part of the contract: file group `i` must contain rows for declared output partition `i`. DataFusion does not validate row placement, matching other user-declared properties such as sortedness. Yes. Yes. This adds public API for declaring file/listing scan output partitioning. No breaking API changes. --------- Co-authored-by: Gabriel <45515538+gabotechs@users.noreply.github.com> --- Cargo.lock | 1 - datafusion/catalog-listing/src/helpers.rs | 4 +- datafusion/catalog-listing/src/options.rs | 54 +++- datafusion/catalog-listing/src/table.rs | 288 ++++++++++++++---- .../core/src/datasource/listing/table.rs | 259 +++++++++++++++- datafusion/core/src/physical_planner.rs | 61 +--- .../datasource/src/file_scan_config/mod.rs | 224 +++++++++++--- datafusion/physical-expr/src/lib.rs | 5 +- datafusion/physical-expr/src/physical_expr.rs | 101 +++++- datafusion/proto/proto/datafusion.proto | 2 + datafusion/proto/src/generated/pbjson.rs | 36 +++ datafusion/proto/src/generated/prost.rs | 4 + .../proto/src/physical_plan/from_proto.rs | 16 +- .../proto/src/physical_plan/to_proto.rs | 7 + .../tests/cases/roundtrip_physical_plan.rs | 288 ++++++++++++++++++ datafusion/sqllogictest/Cargo.toml | 1 - .../src/test_context/range_partitioning.rs | 286 +++++------------ .../test_files/range_partitioning.slt | 10 +- 18 files changed, 1268 insertions(+), 379 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 46b2145c3a63e..c6e5e61a807dd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2591,7 +2591,6 @@ dependencies = [ "chrono", "clap", "datafusion", - "datafusion-datasource", "datafusion-spark", "datafusion-substrait", "env_logger", diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index 0389b3cb17fe9..edb3294051cde 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -325,7 +325,7 @@ pub fn evaluate_partition_prefix<'a>( } } -fn filter_partitions( +pub fn filter_partitioned_file( pf: PartitionedFile, filters: &[Expr], df_schema: &DFSchema, @@ -447,7 +447,7 @@ pub async fn pruned_partition_list<'a>( )) }) .try_filter_map(move |pf| { - futures::future::ready(filter_partitions(pf, filters, &df_schema)) + futures::future::ready(filter_partitioned_file(pf, filters, &df_schema)) }) .boxed()) } diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 0ab15e05abba1..65444706dd02f 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -21,7 +21,7 @@ use datafusion_common::plan_err; use datafusion_datasource::ListingTableUrl; use datafusion_datasource::file_format::FileFormat; use datafusion_execution::config::SessionConfig; -use datafusion_expr::SortExpr; +use datafusion_expr::{Partitioning, SortExpr}; use futures::StreamExt; use futures::TryStreamExt; use itertools::Itertools; @@ -61,6 +61,46 @@ pub struct ListingOptions { /// multiple equivalent orderings, the outer `Vec` will have a /// single element. pub file_sort_order: Vec>, + /// Declared output partitioning for scans from this table. + /// + /// Expressions are logical expressions over the full table schema. When set, + /// [`ListingTable`](crate::ListingTable) creates one file group per + /// declared output partition. When unset, file grouping uses the scan-time + /// [`SessionConfig::target_partitions`](datafusion_execution::config::SessionConfig::target_partitions). + /// + /// Files are listed in path order, split into whole-file groups across the + /// declared partition count, and then padded with trailing empty groups when + /// needed. DataFusion does not route files by partition values or validate + /// row placement, so callers must ensure file group `i` contains rows for + /// partition `i`. Layouts that require explicit file-to-partition assignment + /// are not supported. + /// + /// For example, range partitioning on column `a` with split points + /// `[10, 20, 30]` declares four output partitions. With three path-ordered + /// files, the trailing partition is preserved as empty: + /// + /// ```text + /// files in path order: f0, f1, f2 + /// + /// file groups: + /// partition 0: [f0] + /// partition 1: [f1] + /// partition 2: [f2] + /// partition 3: [] + /// ``` + /// + /// With five path-ordered files, a partition can contain multiple files: + /// + /// ```text + /// files in path order: f0, f1, f2, f3, f4 + /// + /// file groups: + /// partition 0: [f0, f1] + /// partition 1: [f2, f3] + /// partition 2: [f4] + /// partition 3: [] + /// ``` + pub output_partitioning: Option, } impl ListingOptions { @@ -78,6 +118,7 @@ impl ListingOptions { collect_stat: false, target_partitions: 1, file_sort_order: vec![], + output_partitioning: None, } } @@ -136,6 +177,17 @@ impl ListingOptions { self } + /// Set declared output partitioning. + /// + /// See [`Self::output_partitioning`] for the contract. + pub fn with_output_partitioning( + mut self, + output_partitioning: Option, + ) -> Self { + self.output_partitioning = output_partitioning; + self + } + /// Set `table partition columns` on [`ListingOptions`] and returns self. /// /// "partition columns," used to support [Hive Partitioning], are diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 7ee743a6abe71..7f4d3e1965fea 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -16,14 +16,17 @@ // under the License. use crate::config::SchemaSource; -use crate::helpers::{expr_applicable_for_cols, pruned_partition_list}; +use crate::helpers::{ + expr_applicable_for_cols, filter_partitioned_file, pruned_partition_list, +}; use crate::{ListingOptions, ListingTableConfig}; use arrow::datatypes::{Field, Schema, SchemaBuilder, SchemaRef}; use async_trait::async_trait; use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider}; use datafusion_common::stats::Precision; use datafusion_common::{ - Constraints, SchemaExt, Statistics, internal_datafusion_err, plan_err, project_schema, + Constraints, DFSchema, SchemaExt, Statistics, internal_datafusion_err, plan_err, + project_schema, }; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; @@ -38,8 +41,10 @@ use datafusion_execution::cache::TableScopedPath; use datafusion_execution::cache::cache_manager::FileStatisticsCache; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::ExecutionProps; -use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; -use datafusion_physical_expr::create_lex_ordering; +use datafusion_expr::{ + Expr, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType, +}; +use datafusion_physical_expr::{create_lex_ordering, create_physical_partitioning}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::ExecutionPlan; @@ -433,6 +438,20 @@ fn derive_common_ordering_from_files(file_groups: &[FileGroup]) -> Option datafusion_common::Result { + let files = file_group + .into_inner() + .into_iter() + .map(|file| filter_partitioned_file(file, filters, df_schema)) + .filter_map(Result::transpose) + .collect::>>()?; + Ok(FileGroup::new(files)) +} + // Expressions can be used for partition pruning if they can be evaluated using // only the partition columns and there are partition columns. fn can_be_evaluated_for_partition_pruning( @@ -500,9 +519,19 @@ impl TableProvider for ListingTable { can_be_evaluated_for_partition_pruning(&table_partition_col_names, filter) }); - // We should not limit the number of partitioned files to scan if there are filters and limit - // at the same time. This is because the limit should be applied after the filters are applied. - let statistic_file_limit = if filters.is_empty() { limit } else { None }; + let declared_output_partitioning = self.options.output_partitioning.as_ref(); + + // We should not limit files before assigning declared output partitions + // or before applying non-partition filters. + let statistic_file_limit = + if filters.is_empty() && declared_output_partitioning.is_none() { + limit + } else { + None + }; + let file_group_count = declared_output_partitioning + .and_then(LogicalPartitioning::partition_count) + .unwrap_or_else(|| state.config().target_partitions()); let ListFilesResult { file_groups: mut partitioned_file_lists, @@ -522,17 +551,19 @@ impl TableProvider for ListingTable { state.execution_props(), &partitioned_file_lists, )?; - match state - .config_options() - .execution - .split_file_groups_by_statistics + let split_file_groups_by_statistics = declared_output_partitioning.is_none() + && state + .config_options() + .execution + .split_file_groups_by_statistics; + match split_file_groups_by_statistics .then(|| { output_ordering.first().map(|output_ordering| { FileScanConfig::split_groups_by_statistics_with_target_partitions( &self.table_schema, &partitioned_file_lists, output_ordering, - self.options.target_partitions, + file_group_count, ) }) }) @@ -540,7 +571,7 @@ impl TableProvider for ListingTable { { Some(Err(e)) => log::debug!("failed to split file groups by statistics: {e}"), Some(Ok(new_groups)) => { - if new_groups.len() <= self.options.target_partitions { + if new_groups.len() <= file_group_count { partitioned_file_lists = new_groups; } else { log::debug!( @@ -551,6 +582,41 @@ impl TableProvider for ListingTable { None => {} // no ordering required }; + let output_partitioning = if let Some(output_partitioning) = + declared_output_partitioning + { + let output_partitioning = match output_partitioning { + LogicalPartitioning::RoundRobinBatch(_) => { + return datafusion_common::not_impl_err!( + "RoundRobinBatch output partitioning is not supported for ListingTable" + ); + } + LogicalPartitioning::DistributeBy(_) => { + return datafusion_common::not_impl_err!( + "DistributeBy output partitioning is not supported for ListingTable" + ); + } + LogicalPartitioning::Hash(_, _) | LogicalPartitioning::Range(_) => { + let df_schema = DFSchema::try_from(Arc::clone(&self.table_schema))?; + create_physical_partitioning( + output_partitioning, + &df_schema, + state.execution_props(), + )? + } + }; + let partition_count = output_partitioning.partition_count(); + if partitioned_file_lists.len() != partition_count { + return plan_err!( + "ListingTable output_partitioning has {partition_count} partitions, but the scan has {} file groups", + partitioned_file_lists.len() + ); + } + Some(output_partitioning) + } else { + None + }; + let Some(object_store_url) = self.table_paths.first().map(ListingTableUrl::object_store) else { @@ -560,24 +626,23 @@ impl TableProvider for ListingTable { }; let file_source = self.create_file_source(); + let scan_config = FileScanConfigBuilder::new(object_store_url, file_source) + .with_file_groups(partitioned_file_lists) + .with_constraints(self.constraints.clone()) + .with_statistics(statistics) + .with_projection_indices(projection)? + .with_limit(limit) + .with_output_ordering(output_ordering) + .with_output_partitioning(output_partitioning) + .with_expr_adapter(self.expr_adapter_factory.clone()) + .with_partitioned_by_file_group(partitioned_by_file_group) + .build(); // create the execution plan let plan = self .options .format - .create_physical_plan( - state, - FileScanConfigBuilder::new(object_store_url, file_source) - .with_file_groups(partitioned_file_lists) - .with_constraints(self.constraints.clone()) - .with_statistics(statistics) - .with_projection_indices(projection)? - .with_limit(limit) - .with_output_ordering(output_ordering) - .with_expr_adapter(self.expr_adapter_factory.clone()) - .with_partitioned_by_file_group(partitioned_by_file_group) - .build(), - ) + .create_physical_plan(state, scan_config) .await?; Ok(ScanResult::new(plan)) @@ -689,28 +754,41 @@ impl ListingTable { /// Get the list of files for a scan as well as the file level statistics. /// The list is grouped to let the execution plan know how the files should /// be distributed to different threads / executors. + /// + /// If [`ListingOptions::output_partitioning`] is set, returns one file + /// group per declared partition, including empty trailing groups. pub async fn list_files_for_scan<'a>( &'a self, ctx: &'a dyn Session, filters: &'a [Expr], limit: Option, ) -> datafusion_common::Result { - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? + if let Some(output_partitioning) = self.options.output_partitioning.as_ref() { + self.list_files_for_declared_output_partitioning( + ctx, + output_partitioning, + filters, + ) + .await } else { - return Ok(ListFilesResult { - file_groups: vec![], - statistics: Statistics::new_unknown(&self.file_schema), - grouped_by_partition: false, - }); - }; + self.list_files_for_regular_scan(ctx, filters, limit).await + } + } + + async fn collect_files_for_scan<'a>( + &'a self, + ctx: &'a dyn Session, + store: &'a Arc, + listing_time_filters: &'a [Expr], + file_limit: Option, + ) -> datafusion_common::Result<(FileGroup, bool)> { // list files (with partitions) let file_list = future::try_join_all(self.table_paths.iter().map(|table_path| { pruned_partition_list( ctx, store.as_ref(), table_path, - filters, + listing_time_filters, &self.options.file_extension, &self.options.table_partition_cols, ) @@ -736,8 +814,34 @@ impl ListingTable { .boxed() .buffer_unordered(ctx.config_options().execution.meta_fetch_concurrency); - let (file_group, inexact_stats) = - get_files_with_limit(files, limit, self.options.collect_stat).await?; + get_files_with_limit(files, file_limit, self.options.collect_stat).await + } + + async fn list_files_for_regular_scan<'a>( + &'a self, + ctx: &'a dyn Session, + filters: &'a [Expr], + limit: Option, + ) -> datafusion_common::Result { + let file_group_count = self.options.target_partitions; + if file_group_count == 0 { + return plan_err!( + "ListingTable requires target_partitions to be greater than zero" + ); + } + + let store = if let Some(url) = self.table_paths.first() { + ctx.runtime_env().object_store(url)? + } else { + return Ok(ListFilesResult { + file_groups: vec![], + statistics: Statistics::new_unknown(&self.file_schema), + grouped_by_partition: false, + }); + }; + let (file_group, inexact_stats) = self + .collect_files_for_scan(ctx, &store, filters, limit) + .await?; // Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N // @@ -746,28 +850,102 @@ impl ListingTable { // hash repartitioning for aggregates and joins on partition columns. let threshold = ctx.config_options().optimizer.preserve_file_partitions; - let (file_groups, grouped_by_partition) = if threshold > 0 - && !self.options.table_partition_cols.is_empty() - { - let grouped = - file_group.group_by_partition_values(self.options.target_partitions); - if grouped.len() >= threshold { - (grouped, true) + let (file_groups, grouped_by_partition) = + if threshold > 0 && !self.options.table_partition_cols.is_empty() { + let grouped = file_group.group_by_partition_values(file_group_count); + if grouped.len() >= threshold { + (grouped, true) + } else { + let all_files: Vec<_> = + grouped.into_iter().flat_map(|g| g.into_inner()).collect(); + ( + FileGroup::new(all_files).split_files(file_group_count), + false, + ) + } } else { - let all_files: Vec<_> = - grouped.into_iter().flat_map(|g| g.into_inner()).collect(); - ( - FileGroup::new(all_files).split_files(self.options.target_partitions), - false, - ) - } + (file_group.split_files(file_group_count), false) + }; + + self.list_files_result_from_groups( + ctx, + file_groups, + inexact_stats, + grouped_by_partition, + ) + } + + async fn list_files_for_declared_output_partitioning<'a>( + &'a self, + ctx: &'a dyn Session, + output_partitioning: &LogicalPartitioning, + filters: &'a [Expr], + ) -> datafusion_common::Result { + let Some(file_group_count) = output_partitioning.partition_count() else { + return datafusion_common::not_impl_err!( + "DistributeBy output partitioning is not supported for ListingTable" + ); + }; + if file_group_count == 0 { + return plan_err!( + "ListingTable output_partitioning requires at least one partition" + ); + } + + let store = if let Some(url) = self.table_paths.first() { + ctx.runtime_env().object_store(url)? } else { - ( - file_group.split_files(self.options.target_partitions), - false, - ) + return Ok(ListFilesResult { + file_groups: vec![], + statistics: Statistics::new_unknown(&self.file_schema), + grouped_by_partition: false, + }); }; + let (file_group, inexact_stats) = + self.collect_files_for_scan(ctx, &store, &[], None).await?; + let mut file_groups = file_group.split_files(file_group_count); + if !file_groups.is_empty() { + file_groups.resize_with(file_group_count, || FileGroup::new(vec![])); + } + let file_groups = + self.filter_declared_file_groups_by_partition_filters(file_groups, filters)?; + self.list_files_result_from_groups(ctx, file_groups, inexact_stats, false) + } + + fn filter_declared_file_groups_by_partition_filters( + &self, + file_groups: Vec, + filters: &[Expr], + ) -> datafusion_common::Result> { + if filters.is_empty() { + return Ok(file_groups); + } + + let df_schema = DFSchema::from_unqualified_fields( + self.options + .table_partition_cols + .iter() + .map(|(name, data_type)| Field::new(name, data_type.clone(), true)) + .collect(), + Default::default(), + )?; + + file_groups + .into_iter() + .map(|file_group| { + filter_file_group_by_partition_filters(file_group, filters, &df_schema) + }) + .collect::>>() + } + + fn list_files_result_from_groups( + &self, + _ctx: &dyn Session, + file_groups: Vec, + inexact_stats: bool, + grouped_by_partition: bool, + ) -> datafusion_common::Result { let (file_groups, stats) = compute_all_files_statistics( file_groups, self.schema(), diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index d14ec1f56dce2..067cd380ba21f 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -125,7 +125,7 @@ mod tests { }, }; use arrow::{compute::SortOptions, record_batch::RecordBatch}; - use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_catalog::TableProvider; use datafusion_catalog_listing::{ ListingOptions, ListingTable, ListingTableConfig, SchemaSource, @@ -139,12 +139,18 @@ mod tests { use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_format::FileFormat; use datafusion_expr::dml::InsertOp; - use datafusion_expr::{BinaryExpr, LogicalPlanBuilder, Operator}; + use datafusion_expr::{ + BinaryExpr, LogicalPlanBuilder, Operator, Partitioning as LogicalPartitioning, + RangePartitioning as LogicalRangePartitioning, + }; use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr::expressions::binary; + use datafusion_physical_expr::expressions::{Column, binary}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::empty::EmptyExec; - use datafusion_physical_plan::{ExecutionPlanProperties, collect}; + use datafusion_physical_plan::statistics::StatisticsArgs; + use datafusion_physical_plan::{ + ExecutionPlanProperties, Partitioning, RangePartitioning, SplitPoint, collect, + }; use std::collections::HashMap; use std::io::Write; use std::sync::Arc; @@ -177,6 +183,21 @@ mod tests { .collect() } + fn listing_table_with_files( + ctx: &SessionContext, + files: &[&str], + table_path: &str, + options: ListingOptions, + schema: Schema, + ) -> Result { + register_test_store(ctx, &files.iter().map(|f| (*f, 10)).collect::>()); + + let config = ListingTableConfig::new(ListingTableUrl::parse(table_path)?) + .with_listing_options(options) + .with_schema(Arc::new(schema)); + ListingTable::try_new(config) + } + #[tokio::test] async fn test_schema_source_tracking_comprehensive() -> Result<()> { let ctx = SessionContext::new(); @@ -1286,6 +1307,236 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_list_files_uses_declared_output_partitioning_count() -> Result<()> { + let files = ["bucket/key-prefix/file0", "bucket/key-prefix/file1"]; + + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(1), + ); + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(LogicalPartitioning::Range( + LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::from(10i32)]), + SplitPoint::new(vec![ScalarValue::from(20i32)]), + SplitPoint::new(vec![ScalarValue::from(30i32)]), + ], + )?, + ))); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new("a", DataType::Int32, false)]), + )?; + + let result = table.list_files_for_scan(&ctx.state(), &[], None).await?; + let group_sizes = result + .file_groups + .iter() + .map(|group| group.len()) + .collect::>(); + + assert_eq!(group_sizes, vec![1, 1, 0, 0]); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_normalizes_split_point_types() -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(123_000_000_000), + None, + )])], + )?); + let expected_output_partitioning = + Partitioning::Range(RangePartitioning::try_new( + LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions::default(), + )]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::TimestampSecond( + Some(123), + None, + )])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new( + "a", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]), + )?; + + let scan = table.scan(&ctx.state(), None, &[], None).await?; + assert_eq!(scan.output_partitioning(), &expected_output_partitioning); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_rejects_invalid_split_point_type() + -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "not-an-int".to_string(), + ))])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new("a", DataType::Int32, false)]), + )?; + + let err = table.scan(&ctx.state(), None, &[], None).await.unwrap_err(); + assert_contains!( + err.to_string(), + "Range output partitioning split point 0 value 0 with type Utf8 cannot be represented exactly as ordering expression type Int32" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_rejects_lossy_timestamp_split_point() + -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(123_456), + None, + )])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new( + "a", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]), + )?; + + let err = table.scan(&ctx.state(), None, &[], None).await.unwrap_err(); + assert_contains!( + err.to_string(), + "Range output partitioning split point 0 value 0 with type Timestamp(ns) cannot be represented exactly as ordering expression type Timestamp(s)" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_partition_filter_preserves_declared_output_partitioning() -> Result<()> + { + let files = ["bucket/test/pid=1/file1", "bucket/test/pid=2/file2"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("pid").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::from(2i32)])], + )?); + let expected_output_partitioning = + Partitioning::Range(RangePartitioning::try_new( + LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("pid", 1)), + SortOptions::default(), + )]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::from(2i32)])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_table_partition_cols(vec![("pid".to_string(), DataType::Int32)]) + .with_output_partitioning(Some(output_partitioning.clone())); + + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/test/", + opt, + Schema::new(vec![Field::new("a", DataType::Boolean, false)]), + )?; + + let unfiltered = table.scan(&ctx.state(), None, &[], None).await?; + assert_eq!( + unfiltered.output_partitioning(), + &expected_output_partitioning + ); + + let filter = Expr::eq(col("pid"), lit(2_i32)); + let file_groups = table + .list_files_for_scan(&ctx.state(), std::slice::from_ref(&filter), None) + .await? + .file_groups + .into_iter() + .map(|group| { + group + .into_inner() + .into_iter() + .map(|file| file.path().to_string()) + .collect::>() + }) + .collect::>(); + assert_eq!( + file_groups, + vec![ + Vec::::new(), + vec!["bucket/test/pid=2/file2".to_string()] + ] + ); + + let filtered = table.scan(&ctx.state(), None, &[filter], None).await?; + assert_eq!( + filtered.output_partitioning(), + &expected_output_partitioning + ); + + Ok(()) + } + #[tokio::test] async fn test_listing_table_prunes_extra_files_in_hive() -> Result<()> { let files = [ diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 070c49907b606..2a41d83aac957 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -32,10 +32,11 @@ use crate::logical_expr::{ Aggregate, EmptyRelation, Join, Projection, Sort, TableScan, Unnest, Values, Window, }; use crate::logical_expr::{ - Expr, LogicalPlan, Partitioning as LogicalPartitioning, PlanType, Repartition, - UserDefinedLogicalNode, + Expr, LogicalPlan, PlanType, Repartition, UserDefinedLogicalNode, +}; +use crate::physical_expr::{ + create_physical_expr, create_physical_exprs, create_physical_partitioning, }; -use crate::physical_expr::{create_physical_expr, create_physical_exprs}; use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use crate::physical_plan::analyze::AnalyzeExec; use crate::physical_plan::explain::ExplainExec; @@ -52,8 +53,8 @@ use crate::physical_plan::union::UnionExec; use crate::physical_plan::unnest::UnnestExec; use crate::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use crate::physical_plan::{ - ExecutionPlan, ExecutionPlanProperties, InputOrderMode, Partitioning, PhysicalExpr, - WindowExpr, displayable, windows, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PhysicalExpr, WindowExpr, + displayable, windows, }; use crate::schema_equivalence::schema_satisfied_by; @@ -98,7 +99,7 @@ use datafusion_physical_expr::aggregate::{ }; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::{ - LexOrdering, PhysicalSortExpr, RangePartitioning, create_physical_sort_exprs, + LexOrdering, PhysicalSortExpr, create_physical_sort_exprs, }; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::empty::EmptyExec; @@ -1247,41 +1248,11 @@ impl DefaultPhysicalPlanner { }) => { let physical_input = children.one()?; let input_dfschema = input.as_ref().schema(); - let physical_partitioning = match partitioning_scheme { - LogicalPartitioning::RoundRobinBatch(n) => { - Partitioning::RoundRobinBatch(*n) - } - LogicalPartitioning::Hash(expr, n) => { - let runtime_expr = expr - .iter() - .map(|e| { - create_physical_expr(e, input_dfschema, execution_props) - }) - .collect::>>()?; - Partitioning::Hash(runtime_expr, *n) - } - LogicalPartitioning::Range(range) => { - let sort_exprs = create_physical_sort_exprs( - range.ordering(), - input_dfschema, - execution_props, - )?; - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - internal_datafusion_err!( - "Range repartitioning requires non-empty ordering" - ) - })?; - Partitioning::Range(RangePartitioning::try_new( - ordering, - range.split_points().to_vec(), - )?) - } - LogicalPartitioning::DistributeBy(_) => { - return not_impl_err!( - "Physical plan does not support DistributeBy partitioning" - ); - } - }; + let physical_partitioning = create_physical_partitioning( + partitioning_scheme, + input_dfschema, + execution_props, + )?; Arc::new(RepartitionExec::try_new( physical_input, physical_partitioning, @@ -3235,8 +3206,8 @@ mod tests { use crate::datasource::MemTable; use crate::datasource::file_format::options::CsvReadOptions; use crate::physical_plan::{ - DisplayAs, DisplayFormatType, PlanProperties, SendableRecordBatchStream, - expressions, + DisplayAs, DisplayFormatType, Partitioning, PlanProperties, + SendableRecordBatchStream, expressions, }; use crate::prelude::{SessionConfig, SessionContext}; use crate::test_util::{scan_empty, scan_empty_with_partitions}; @@ -3257,8 +3228,8 @@ mod tests { use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::{ Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder, - RangePartitioning, Signature, TableSource, UserDefinedLogicalNodeCore, - Volatility, WindowFunctionDefinition, col, lit, + Partitioning as LogicalPartitioning, RangePartitioning, Signature, TableSource, + UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, lit, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 4bf86e17d387d..5e4126793cb1f 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -27,8 +27,7 @@ use crate::{ file_stream::work_source::SharedWorkSource, source::DataSource, statistics::MinMaxStatistics, }; -use arrow::datatypes::FieldRef; -use arrow::datatypes::{DataType, Schema, SchemaRef}; +use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::{ Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err, @@ -40,7 +39,7 @@ use datafusion_expr::Operator; use crate::source::OpenArgs; use datafusion_physical_expr::expressions::{BinaryExpr, Column}; -use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::projection::{ProjectionExprs, ProjectionMapping}; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr::{EquivalenceProperties, Partitioning, split_conjunction}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; @@ -205,7 +204,17 @@ pub struct FileScanConfig { /// /// If the number of file partitions > target_partitions, the file partitions will be grouped /// in a round-robin fashion such that number of file partitions = target_partitions. + /// + /// Follow-up: remove this redundant field in favor of + /// `output_partitioning`, see . pub partitioned_by_file_group: bool, + /// Declared physical output partitioning for this scan. + /// + /// Expressions are against the full table schema, before scan projection or + /// filtering. `ListingTable` validates partition count before building the + /// scan, and direct builders with mismatched counts fall back to + /// `UnknownPartitioning`. + pub output_partitioning: Option, } /// A builder for [`FileScanConfig`]'s. @@ -272,6 +281,7 @@ pub struct FileScanConfigBuilder { file_groups: Vec, statistics: Option, output_ordering: Vec, + output_partitioning: Option, file_compression_type: Option, batch_size: Option, expr_adapter_factory: Option>, @@ -295,6 +305,7 @@ impl FileScanConfigBuilder { file_groups: vec![], statistics: None, output_ordering: vec![], + output_partitioning: None, file_compression_type: None, limit: None, preserve_order: false, @@ -461,6 +472,15 @@ impl FileScanConfigBuilder { self } + /// Set declared physical output partitioning for this scan. + pub fn with_output_partitioning( + mut self, + output_partitioning: Option, + ) -> Self { + self.output_partitioning = output_partitioning; + self + } + /// Set the file compression type pub fn with_file_compression_type( mut self, @@ -519,6 +539,7 @@ impl FileScanConfigBuilder { file_groups, statistics, output_ordering, + output_partitioning, file_compression_type, batch_size, expr_adapter_factory: expr_adapter, @@ -548,6 +569,7 @@ impl FileScanConfigBuilder { expr_adapter_factory: expr_adapter, statistics, partitioned_by_file_group, + output_partitioning, } } } @@ -560,6 +582,7 @@ impl From for FileScanConfigBuilder { file_groups: config.file_groups, statistics: Some(config.statistics), output_ordering: config.output_ordering, + output_partitioning: config.output_partitioning, file_compression_type: Some(config.file_compression_type), limit: config.limit, preserve_order: config.preserve_order, @@ -571,6 +594,52 @@ impl From for FileScanConfigBuilder { } } +fn hash_partitioning_from_partition_fields( + schema: &Schema, + partition_cols: &[FieldRef], + partition_count: usize, +) -> Option { + if partition_cols.is_empty() { + return None; + } + + let mut exprs: Vec> = Vec::with_capacity(partition_cols.len()); + for partition_col in partition_cols { + let name = partition_col.name(); + let idx = schema + .fields() + .iter() + .position(|field| field.name() == name)?; + exprs.push(Arc::new(Column::new(name, idx))); + } + + Some(Partitioning::Hash(exprs, partition_count)) +} + +fn project_output_partitioning( + partitioning: &Partitioning, + mapping: &ProjectionMapping, + input_schema: &SchemaRef, + partition_count: usize, +) -> Partitioning { + let input_eq_properties = EquivalenceProperties::new(Arc::clone(input_schema)); + match partitioning { + Partitioning::Hash(exprs, _) => { + let projected_exprs = input_eq_properties + .project_expressions(exprs, mapping) + .collect::>>(); + projected_exprs + .map(|exprs| Partitioning::Hash(exprs, partition_count)) + .unwrap_or_else(|| Partitioning::UnknownPartitioning(partition_count)) + } + Partitioning::Range(_) + | Partitioning::RoundRobinBatch(_) + | Partitioning::UnknownPartitioning(_) => { + partitioning.project(mapping, &input_eq_properties) + } + } +} + impl DataSource for FileScanConfig { fn open( &self, @@ -658,6 +727,10 @@ impl DataSource for FileScanConfig { display_orderings(f, &orderings)?; + if self.output_partitioning.is_some() { + write!(f, ", output_partitioning={}", self.output_partitioning())?; + } + if !self.constraints.is_empty() { write!(f, ", {}", self.constraints)?; } @@ -681,10 +754,9 @@ impl DataSource for FileScanConfig { repartition_file_min_size: usize, output_ordering: Option, ) -> Result>> { - // When files are grouped by partition values, we cannot allow byte-range - // splitting. It would mix rows from different partition values across - // file groups, breaking the Hash partitioning. - if self.partitioned_by_file_group { + // When file groups define output partitioning, repartitioning files + // would invalidate the partition-to-file-group mapping. + if self.output_partitioning.is_some() || self.partitioned_by_file_group { return Ok(None); } @@ -700,13 +772,18 @@ impl DataSource for FileScanConfig { /// Returns the output partitioning for this file scan. /// - /// When `partitioned_by_file_group` is true, this returns `Partitioning::Hash` on - /// the Hive partition columns, allowing the optimizer to skip hash repartitioning - /// for aggregates and joins on those columns. + /// When `output_partitioning` is set, this returns the declared partitioning + /// after applying scan projection. When `partitioned_by_file_group` is true, + /// this returns `Partitioning::Hash` on the Hive partition columns, allowing + /// the optimizer to skip hash repartitioning for aggregates and joins on + /// those columns. + /// + /// If projection or partition count validation fails, this returns + /// `UnknownPartitioning`. /// /// Tradeoffs - /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries with - /// `GROUP BY` or `ORDER BY` on partition columns. + /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries whose + /// required distribution is satisfied by the scan's output partitioning. /// - Cost: Files are grouped by partition values rather than split by byte /// ranges, which may reduce I/O parallelism when partition sizes are uneven. /// For simple aggregations without `ORDER BY`, this cost may outweigh the benefit. @@ -715,39 +792,45 @@ impl DataSource for FileScanConfig { /// - Idea: Could allow byte-range splitting within partition-aware groups, /// preserving I/O parallelism while maintaining partition semantics. fn output_partitioning(&self) -> Partitioning { - if self.partitioned_by_file_group { - let partition_cols = self.table_partition_cols(); - if !partition_cols.is_empty() { - let projected_schema = match self.projected_schema() { - Ok(schema) => schema, - Err(_) => { - debug!( - "Could not get projected schema, falling back to UnknownPartitioning." - ); - return Partitioning::UnknownPartitioning(self.file_groups.len()); - } - }; - - // Build Column expressions for partition columns based on their - // position in the projected schema - let mut exprs: Vec> = Vec::new(); - for partition_col in partition_cols { - if let Some((idx, _)) = projected_schema - .fields() - .iter() - .enumerate() - .find(|(_, f)| f.name() == partition_col.name()) - { - exprs.push(Arc::new(Column::new(partition_col.name(), idx))); - } - } + let Some(output_partitioning) = self.output_partitioning.clone().or_else(|| { + self.partitioned_by_file_group.then(|| { + hash_partitioning_from_partition_fields( + self.file_source.table_schema().table_schema(), + self.table_partition_cols(), + self.file_groups.len(), + ) + })? + }) else { + return Partitioning::UnknownPartitioning(self.file_groups.len()); + }; + if output_partitioning.partition_count() != self.file_groups.len() { + warn!( + "Declared output partitioning has {} partitions, but file scan has {} file groups. Falling back to UnknownPartitioning.", + output_partitioning.partition_count(), + self.file_groups.len() + ); + return Partitioning::UnknownPartitioning(self.file_groups.len()); + } - if exprs.len() == partition_cols.len() { - return Partitioning::Hash(exprs, self.file_groups.len()); + if let Some(projection) = self.file_source.projection() { + let schema = self.file_source.table_schema().table_schema(); + return match projection.projection_mapping(schema) { + Ok(mapping) => project_output_partitioning( + &output_partitioning, + &mapping, + schema, + self.file_groups.len(), + ), + Err(e) => { + debug!( + "Could not project output partitioning, falling back to UnknownPartitioning: {e}" + ); + Partitioning::UnknownPartitioning(self.file_groups.len()) } - } + }; } - Partitioning::UnknownPartitioning(self.file_groups.len()) + + output_partitioning } /// Computes the effective equivalence properties of this file scan, taking @@ -1041,7 +1124,10 @@ impl DataSource for FileScanConfig { /// when file order must be preserved or the file groups define the output /// partitioning needed for the rest of the plan fn create_sibling_state(&self) -> Option> { - if self.preserve_order || self.partitioned_by_file_group { + if self.preserve_order + || self.output_partitioning.is_some() + || self.partitioned_by_file_group + { return None; } @@ -2445,6 +2531,58 @@ mod tests { assert!(matches!(partitioning, Partitioning::UnknownPartitioning(_))); } + #[test] + fn test_declared_output_partitioning_projects_with_scan() { + let file_schema = aggr_test_schema(); + let output_partitioning = + Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4); + + let mut config = config_for_projection( + Arc::clone(&file_schema), + Some(vec![1, 2]), + Statistics::new_unknown(&file_schema), + vec![], + ); + config.file_groups = vec![ + FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]), + ]; + config.output_partitioning = Some(output_partitioning); + + match config.output_partitioning() { + Partitioning::Hash(exprs, num_partitions) => { + assert_eq!(num_partitions, 4); + assert_eq!(exprs.len(), 1); + let column = exprs[0].downcast_ref::().unwrap(); + assert_eq!(column.name(), "c2"); + assert_eq!(column.index(), 0); + } + _ => panic!("Expected Hash partitioning"), + } + + let mut config = config_for_projection( + Arc::clone(&file_schema), + Some(vec![2]), + Statistics::new_unknown(&file_schema), + vec![], + ); + config.file_groups = vec![ + FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]), + ]; + config.output_partitioning = + Some(Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4)); + + assert!(matches!( + config.output_partitioning(), + Partitioning::UnknownPartitioning(4) + )); + } + #[test] fn test_output_partitioning_no_partition_columns() { let file_schema = aggr_test_schema(); diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index 2f5d5f0fcb460..a0e9f8ee05363 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -62,8 +62,9 @@ pub use equivalence::{ pub use partitioning::{Distribution, Partitioning, RangePartitioning}; 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, - physical_exprs_bag_equal, physical_exprs_contains, physical_exprs_equal, + create_ordering, create_physical_partitioning, create_physical_sort_expr, + create_physical_sort_exprs, physical_exprs_bag_equal, physical_exprs_contains, + physical_exprs_equal, }; pub use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef}; diff --git a/datafusion/physical-expr/src/physical_expr.rs b/datafusion/physical-expr/src/physical_expr.rs index 77ede76e1daa8..6ff5be4e38229 100644 --- a/datafusion/physical-expr/src/physical_expr.rs +++ b/datafusion/physical-expr/src/physical_expr.rs @@ -21,15 +21,17 @@ use crate::expressions::{self, Column}; use crate::{LexOrdering, PhysicalSortExpr, create_physical_expr}; use arrow::compute::SortOptions; -use arrow::datatypes::{Schema, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{DFSchema, HashMap}; +use datafusion_common::{DFSchema, HashMap, ScalarValue, SplitPoint}; use datafusion_common::{Result, plan_err}; use datafusion_expr::execution_props::ExecutionProps; -use datafusion_expr::{Expr, SortExpr}; +use datafusion_expr::{Expr, Partitioning as LogicalPartitioning, SortExpr}; +use datafusion_expr_common::casts::try_cast_literal_to_type; use itertools::izip; // Exports: +use crate::{Partitioning, RangePartitioning}; pub(crate) use datafusion_physical_expr_common::physical_expr::PhysicalExpr; /// Adds the `offset` value to `Column` indices inside `expr`. This function is @@ -216,6 +218,99 @@ pub fn create_physical_sort_exprs( .collect() } +/// Create physical partitioning from logical partitioning. +pub fn create_physical_partitioning( + partitioning: &LogicalPartitioning, + input_dfschema: &DFSchema, + execution_props: &ExecutionProps, +) -> Result { + match partitioning { + LogicalPartitioning::RoundRobinBatch(n) => Ok(Partitioning::RoundRobinBatch(*n)), + LogicalPartitioning::Hash(exprs, partition_count) => { + let exprs = exprs + .iter() + .map(|expr| create_physical_expr(expr, input_dfschema, execution_props)) + .collect::>>()?; + Ok(Partitioning::Hash(exprs, *partition_count)) + } + LogicalPartitioning::Range(range) => { + let ordering = create_physical_sort_exprs( + range.ordering(), + input_dfschema, + execution_props, + )?; + let Some(ordering) = LexOrdering::new(ordering) else { + return plan_err!("Range partitioning requires non-empty ordering"); + }; + let split_points = normalize_range_split_points( + &ordering, + range.split_points(), + input_dfschema.as_arrow(), + )?; + let range = RangePartitioning::try_new(ordering, split_points)?; + Ok(Partitioning::Range(range)) + } + LogicalPartitioning::DistributeBy(_) => { + datafusion_common::not_impl_err!( + "Physical plan does not support DistributeBy partitioning" + ) + } + } +} + +fn normalize_range_split_points( + ordering: &LexOrdering, + split_points: &[SplitPoint], + schema: &Schema, +) -> Result> { + split_points + .iter() + .enumerate() + .map(|(split_idx, split_point)| { + let values = split_point + .values() + .iter() + .zip(ordering.iter()) + .enumerate() + .map(|(value_idx, (value, sort_expr))| { + let target_type = sort_expr.expr.data_type(schema)?; + normalize_range_split_point_value( + value, + &target_type, + split_idx, + value_idx, + ) + }) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect() +} + +fn normalize_range_split_point_value( + value: &ScalarValue, + target_type: &DataType, + split_idx: usize, + value_idx: usize, +) -> Result { + let value_type = value.data_type(); + if &value_type == target_type { + return Ok(value.clone()); + } + + if let Some(casted) = try_cast_literal_to_type(value, target_type) { + // Split points define physical partition boundaries, so normalization + // must reject casts that would change the advertised boundary. + if try_cast_literal_to_type(&casted, &value_type).as_ref() == Some(value) { + return Ok(casted); + } + } + + plan_err!( + "Range output partitioning split point {split_idx} value {value_idx} with type {value_type} cannot be represented exactly as ordering expression type {target_type}" + ) +} + pub fn add_offset_to_physical_sort_exprs( sort_exprs: impl IntoIterator, offset: isize, diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index b727c670152b9..23a0edc9ef23c 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -1160,6 +1160,8 @@ message FileScanExecConf { optional uint64 batch_size = 12; optional ProjectionExprs projection_exprs = 13; + optional bool partitioned_by_file_group = 14; + optional Partitioning output_partitioning = 15; } message ParquetScanExecNode { diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index 4ee2bd30b2c77..1e85d14f8cae1 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -6848,6 +6848,12 @@ impl serde::Serialize for FileScanExecConf { if self.projection_exprs.is_some() { len += 1; } + if self.partitioned_by_file_group.is_some() { + len += 1; + } + if self.output_partitioning.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.FileScanExecConf", len)?; if !self.file_groups.is_empty() { struct_ser.serialize_field("fileGroups", &self.file_groups)?; @@ -6884,6 +6890,12 @@ impl serde::Serialize for FileScanExecConf { if let Some(v) = self.projection_exprs.as_ref() { struct_ser.serialize_field("projectionExprs", v)?; } + if let Some(v) = self.partitioned_by_file_group.as_ref() { + struct_ser.serialize_field("partitionedByFileGroup", v)?; + } + if let Some(v) = self.output_partitioning.as_ref() { + struct_ser.serialize_field("outputPartitioning", v)?; + } struct_ser.end() } } @@ -6911,6 +6923,10 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "batchSize", "projection_exprs", "projectionExprs", + "partitioned_by_file_group", + "partitionedByFileGroup", + "output_partitioning", + "outputPartitioning", ]; #[allow(clippy::enum_variant_names)] @@ -6926,6 +6942,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { Constraints, BatchSize, ProjectionExprs, + PartitionedByFileGroup, + OutputPartitioning, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -6958,6 +6976,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "constraints" => Ok(GeneratedField::Constraints), "batchSize" | "batch_size" => Ok(GeneratedField::BatchSize), "projectionExprs" | "projection_exprs" => Ok(GeneratedField::ProjectionExprs), + "partitionedByFileGroup" | "partitioned_by_file_group" => Ok(GeneratedField::PartitionedByFileGroup), + "outputPartitioning" | "output_partitioning" => Ok(GeneratedField::OutputPartitioning), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -6988,6 +7008,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { let mut constraints__ = None; let mut batch_size__ = None; let mut projection_exprs__ = None; + let mut partitioned_by_file_group__ = None; + let mut output_partitioning__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::FileGroups => { @@ -7061,6 +7083,18 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { } projection_exprs__ = map_.next_value()?; } + GeneratedField::PartitionedByFileGroup => { + if partitioned_by_file_group__.is_some() { + return Err(serde::de::Error::duplicate_field("partitionedByFileGroup")); + } + partitioned_by_file_group__ = map_.next_value()?; + } + GeneratedField::OutputPartitioning => { + if output_partitioning__.is_some() { + return Err(serde::de::Error::duplicate_field("outputPartitioning")); + } + output_partitioning__ = map_.next_value()?; + } } } Ok(FileScanExecConf { @@ -7075,6 +7109,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { constraints: constraints__, batch_size: batch_size__, projection_exprs: projection_exprs__, + partitioned_by_file_group: partitioned_by_file_group__, + output_partitioning: output_partitioning__, }) } } diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index e82582bd3c1aa..e61189f69a222 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -1735,6 +1735,10 @@ pub struct FileScanExecConf { pub batch_size: ::core::option::Option, #[prost(message, optional, tag = "13")] pub projection_exprs: ::core::option::Option, + #[prost(bool, optional, tag = "14")] + pub partitioned_by_file_group: ::core::option::Option, + #[prost(message, optional, tag = "15")] + pub output_partitioning: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ParquetScanExecNode { diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index b5b2652659401..37e7d6b5d84c2 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -809,6 +809,12 @@ pub fn parse_protobuf_file_scan_config( )?; output_ordering.extend(LexOrdering::new(sort_exprs)); } + let output_partitioning = parse_protobuf_partitioning( + proto.output_partitioning.as_ref(), + ctx, + &schema, + proto_converter, + )?; // Parse projection expressions if present and apply to file source let file_source = if let Some(proto_projection_exprs) = &proto.projection_exprs { @@ -837,14 +843,18 @@ pub fn parse_protobuf_file_scan_config( file_source }; - let config = FileScanConfigBuilder::new(object_store_url, file_source) + let mut config_builder = FileScanConfigBuilder::new(object_store_url, file_source) .with_file_groups(file_groups) .with_constraints(constraints) .with_statistics(statistics) .with_limit(proto.limit.as_ref().map(|sl| sl.limit as usize)) .with_output_ordering(output_ordering) - .with_batch_size(proto.batch_size.map(|s| s as usize)) - .build(); + .with_output_partitioning(output_partitioning) + .with_batch_size(proto.batch_size.map(|s| s as usize)); + if proto.partitioned_by_file_group.unwrap_or(false) { + config_builder = config_builder.with_partitioned_by_file_group(true); + } + let config = config_builder.build(); Ok(config) } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 785c4df95fbaf..5d3fed8737ecf 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -786,6 +786,11 @@ pub fn serialize_file_scan_config( serialize_physical_sort_exprs(order.to_vec(), codec, proto_converter)?; output_orderings.push(ordering) } + let output_partitioning = conf + .output_partitioning + .as_ref() + .map(|partitioning| serialize_partitioning(partitioning, codec, proto_converter)) + .transpose()?; // Fields must be added to the schema so that they can persist in the protobuf, // and then they are to be removed from the schema in `parse_protobuf_file_scan_config` @@ -845,6 +850,8 @@ pub fn serialize_file_scan_config( constraints: Some(conf.constraints.clone().into()), batch_size: conf.batch_size.map(|s| s as u64), projection_exprs, + partitioned_by_file_group: Some(conf.partitioned_by_file_group), + output_partitioning, }) } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 19bc50933d3f0..1cf55a9afee40 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -4081,3 +4081,291 @@ fn test_sort_topk_with_dynamic_filter_roundtrip() -> Result<()> { Ok(()) } + +/// A custom [`ExecutionPlan`] which stores [`PhysicalExpr`]s. +struct CustomExecWithExprs { + exprs: Vec>, + child: Arc, +} + +#[derive(Clone, PartialEq, Message)] +struct CustomExecWithExprsProto { + #[prost(message, repeated, tag = "1")] + exprs: Vec, +} + +impl std::fmt::Debug for CustomExecWithExprs { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CustomExecWithExprs") + .field("exprs", &self.exprs) + .field("child", &self.child) + .finish() + } +} + +impl CustomExecWithExprs { + fn new(exprs: Vec>, child: Arc) -> Self { + Self { exprs, child } + } +} + +impl DisplayAs for CustomExecWithExprs { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CustomExecWithExprs") + } +} + +impl ExecutionPlan for CustomExecWithExprs { + fn name(&self) -> &str { + "CustomExecWithExprs" + } + + fn schema(&self) -> SchemaRef { + self.child.schema() + } + + fn properties(&self) -> &Arc { + self.child.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + unreachable!() + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } +} + +/// A [`PhysicalExtensionCodec`] for [`CustomExecWithExprs`]. +#[derive(Debug)] +struct CustomExecWithExprsCodec {} + +impl PhysicalExtensionCodec for CustomExecWithExprsCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self); + let input_schema = inputs[0].schema(); + let proto = CustomExecWithExprsProto::decode(buf) + .map_err(|e| internal_datafusion_err!("Failed to decode custom exec: {e}"))?; + let exprs = proto + .exprs + .iter() + .map(|expr_proto| { + proto_converter.proto_to_physical_expr( + expr_proto, + input_schema.as_ref(), + &decode_ctx, + ) + }) + .collect::>>()?; + + Ok(Arc::new(CustomExecWithExprs::new(exprs, inputs[0].clone()))) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + let custom = node + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected CustomExecWithExprs"))?; + let proto = CustomExecWithExprsProto { + exprs: custom + .exprs + .iter() + .map(|expr| proto_converter.physical_expr_to_proto(expr, self)) + .collect::>>()?, + }; + proto + .encode(buf) + .map_err(|e| internal_datafusion_err!("Failed to encode custom exec: {e}"))?; + + Ok(()) + } +} + +/// Tests that a custom [`ExecutionPlan`] with [`PhysicalExpr`] can +/// dedupe dynamic filters by using the proto converter in its +/// [`PhysicalExtensionCodec`] implementation. +#[test] +fn test_custom_node_with_dynamic_filter_dedup_roundtrip() -> Result<()> { + // Create the plan: + // + // FilterExec(dynamic_filter) + // -> CustomExecWithExprs(exprs: [dynamic_filter]) + // -> EmptyExec + // + // The same dynamic filter expression is saved in both the FilterExec and CustomExecWithExprs. + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0)) as Arc], + lit(true), + )); + let dynamic_filter_expr: Arc = dynamic_filter; + + let empty = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let custom_exec = Arc::new(CustomExecWithExprs::new( + vec![Arc::clone(&dynamic_filter_expr)], + empty, + )); + let filter_exec = Arc::new(FilterExec::try_new( + Arc::clone(&dynamic_filter_expr), + custom_exec, + )?) as Arc; + + // Roundtrip with DeduplicatingProtoConverter + let codec = CustomExecWithExprsCodec {}; + let converter = DeduplicatingProtoConverter {}; + + let bytes = physical_plan_to_bytes_with_proto_converter( + Arc::clone(&filter_exec), + &codec, + &converter, + )?; + + let ctx = SessionContext::new(); + let deser_converter = DeduplicatingProtoConverter {}; + let deserialized = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + &codec, + &deser_converter, + )?; + + // Extract the deserialized FilterExec's dynamic filter + let deser_filter = deserialized + .downcast_ref::() + .expect("Top-level should be FilterExec"); + let deser_filter_df = deser_filter.predicate(); + + // Extract the deserialized custom node's dynamic filter + let deser_custom = deser_filter + .input() + .downcast_ref::() + .expect("FilterExec child should be CustomExecWithExprs"); + assert_eq!(deser_custom.exprs.len(), 1, "Should have one expression"); + let [deser_custom_df] = deser_custom.exprs.as_slice() else { + return internal_err!("Custom node should have one expression"); + }; + + // Pass the un-remapped filter first so the helper's `with_new_children` + // rewrite can reconstruct the remapped form on the other side. + assert_dynamic_filters_equal(deser_custom_df, deser_filter_df); + assert_dynamic_filter_update_is_visible(deser_custom_df, deser_filter_df)?; + + Ok(()) +} + +fn roundtrip_file_scan_config(scan_config: FileScanConfig) -> Result { + let exec_plan: Arc = DataSourceExec::from_data_source(scan_config); + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let result_plan = + roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + + let data_source_exec = result_plan + .downcast_ref::() + .expect("Expected DataSourceExec"); + let file_scan_config = data_source_exec + .data_source() + .downcast_ref::() + .expect("Expected FileScanConfig"); + Ok(file_scan_config.clone()) +} + +#[test] +fn roundtrip_parquet_exec_partitioned_by_file_group() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_partitioned_by_file_group(true) + .build(); + + assert!(roundtrip_file_scan_config(scan_config)?.partitioned_by_file_group); + Ok(()) +} + +#[test] +fn roundtrip_parquet_exec_output_partitioning() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let output_partitioning = + Partitioning::Hash(vec![Arc::new(Column::new("col", 0))], 1); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_output_partitioning(Some(output_partitioning.clone())) + .build(); + + assert_eq!( + roundtrip_file_scan_config(scan_config)?.output_partitioning, + Some(output_partitioning) + ); + + Ok(()) +} + +#[test] +fn roundtrip_parquet_exec_range_output_partitioning() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let output_partitioning = Partitioning::Range(RangePartitioning::new( + LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new( + "col", 0, + )))]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![ + FileGroup::new(vec![PartitionedFile::new( + "/path/to/file-1.parquet".to_string(), + 1024, + )]), + FileGroup::new(vec![PartitionedFile::new( + "/path/to/file-2.parquet".to_string(), + 1024, + )]), + ]) + .with_output_partitioning(Some(output_partitioning.clone())) + .build(); + + assert_eq!( + roundtrip_file_scan_config(scan_config)?.output_partitioning, + Some(output_partitioning) + ); + + Ok(()) +} diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index a642fbe22a6e3..e2ffe1415a1fb 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -47,7 +47,6 @@ bytes = { workspace = true, optional = true } chrono = { workspace = true, optional = true } clap = { version = "4.5.60", features = ["derive", "env"] } datafusion = { workspace = true, default-features = true, features = ["avro"] } -datafusion-datasource = { workspace = true } datafusion-spark = { workspace = true, features = ["core"] } datafusion-substrait = { workspace = true, default-features = true, optional = true } futures = { workspace = true } diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index 88e49708baf60..a3e16eefd881a 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -15,236 +15,94 @@ // specific language governing permissions and limitations // under the License. -use std::fmt; +use std::fs::{create_dir_all, remove_dir_all, write}; +use std::path::Path; use std::sync::Arc; -use arrow::array::Int32Array; -use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use arrow::record_batch::RecordBatch; -use async_trait::async_trait; -use datafusion::catalog::Session; -use datafusion::common::{Result, ScalarValue, project_schema}; -use datafusion::datasource::source::{DataSource, DataSourceExec}; -use datafusion::datasource::{TableProvider, TableType}; -use datafusion::execution::context::TaskContext; -use datafusion::logical_expr::Expr; -use datafusion::physical_expr::EquivalenceProperties; -use datafusion::physical_expr::expressions::col as physical_col; -use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; -use datafusion::physical_plan::execution_plan::SchedulingType; -use datafusion::physical_plan::projection::ProjectionExprs; -use datafusion::physical_plan::{ - DisplayFormatType, ExecutionPlan, Partitioning, RangePartitioning, - SendableRecordBatchStream, SplitPoint, Statistics, +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion::common::{ScalarValue, SplitPoint}; +use datafusion::datasource::file_format::csv::CsvFormat; +use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; +use datafusion::logical_expr::{Partitioning, RangePartitioning, col}; use datafusion::prelude::SessionContext; -use datafusion_datasource::memory::MemorySourceConfig; // ============================================================================== // Range Partitioned Table (sqllogictest-only) // ============================================================================== -/// Simple range-partitioned table for testing before declaring such tables is -/// supported via SQL. -#[derive(Debug)] -struct RangePartitionedTable { - schema: SchemaRef, - partitions: Vec>, - range_column_index: usize, - split_points: Vec, -} - -#[async_trait] -impl TableProvider for RangePartitionedTable { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } - - fn table_type(&self) -> TableType { - TableType::Base - } - - async fn scan( - &self, - state: &dyn Session, - projection: Option<&Vec>, - _filters: &[Expr], - _limit: Option, - ) -> Result> { - let projected_schema = project_schema(&self.schema, projection)?; - let mut source = MemorySourceConfig::try_new( - &self.partitions, - Arc::clone(&self.schema), - projection.cloned(), - )?; - source = source.with_show_sizes(state.config_options().explain.show_sizes); - - let output_partitioning = - self.output_partitioning(projection, &projected_schema)?; - let source = RangePartitionedSource { - inner: source, - output_partitioning, - }; - - Ok(DataSourceExec::from_data_source(source)) - } -} - -impl RangePartitionedTable { - fn output_partitioning( - &self, - projection: Option<&Vec>, - projected_schema: &SchemaRef, - ) -> Result { - let Some(projected_range_index) = - projected_index(self.range_column_index, projection) - else { - return Ok(Partitioning::UnknownPartitioning(self.partitions.len())); - }; - - let range_column = projected_schema.field(projected_range_index).name(); - let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( - physical_col(range_column, projected_schema)?, - SortOptions::default(), - )]) - .expect("range ordering should not be empty"); - - Ok(Partitioning::Range(RangePartitioning::try_new( - ordering, - self.split_points.clone(), - )?)) - } -} - -fn projected_index( - column_index: usize, - projection: Option<&Vec>, -) -> Option { - projection - .map(|projection| projection.iter().position(|idx| *idx == column_index)) - .unwrap_or(Some(column_index)) -} - -#[derive(Clone, Debug)] -struct RangePartitionedSource { - inner: MemorySourceConfig, - output_partitioning: Partitioning, -} - -impl DataSource for RangePartitionedSource { - fn open( - &self, - partition: usize, - context: Arc, - ) -> Result { - self.inner.open(partition, context) - } - - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - self.inner.fmt_as(t, f)?; - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, ", output_partitioning={}", self.output_partitioning) - } - DisplayFormatType::TreeRender => Ok(()), - } - } - - fn output_partitioning(&self) -> Partitioning { - self.output_partitioning.clone() - } - - fn eq_properties(&self) -> EquivalenceProperties { - self.inner.eq_properties() - } - - fn scheduling_type(&self) -> SchedulingType { - self.inner.scheduling_type() - } - - fn partition_statistics(&self, partition: Option) -> Result> { - self.inner.partition_statistics(partition) - } - - fn with_fetch(&self, limit: Option) -> Option> { - Some(Arc::new(Self { - inner: self.inner.clone().with_limit(limit), - output_partitioning: self.output_partitioning.clone(), - })) - } - - fn fetch(&self) -> Option { - self.inner.fetch() - } - - fn try_swapping_with_projection( - &self, - _projection: &ProjectionExprs, - ) -> Result>> { - // Range partitioning metadata is projection-sensitive. This fixture - // computes it in TableProvider::scan, so do not rewrite later - // ProjectionExec nodes into the source. - Ok(None) - } -} - +/// Registers a simple range-partitioned listing table for testing before +/// declaring such tables is supported via SQL. pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { let schema = Arc::new(Schema::new(vec![ Field::new("range_key", DataType::Int32, false), Field::new("non_range_key", DataType::Int32, false), Field::new("value", DataType::Int32, false), ])); - let partitions = vec![ - vec![range_partition_batch(&schema, &[1, 5], &[1, 2], &[10, 50])], - vec![range_partition_batch( - &schema, - &[10, 15], - &[1, 2], - &[100, 150], - )], - vec![range_partition_batch( - &schema, - &[20, 25], - &[1, 2], - &[200, 250], - )], - vec![range_partition_batch( - &schema, - &[30, 35], - &[1, 2], - &[300, 350], - )], - ]; - let split_points = vec![ - SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), - SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), - SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), - ]; - let table = RangePartitionedTable { + let output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned"), schema, - partitions, - range_column_index: 0, - split_points, - }; - - ctx.register_table("range_partitioned", Arc::new(table)) - .expect("range partitioned table registration should succeed"); + [ + "1,1,10\n5,2,50\n", + "10,1,100\n15,2,150\n", + "20,1,200\n25,2,250\n", + "30,1,300\n35,2,350\n", + ], + Some(output_partitioning), + ); } -fn range_partition_batch( - schema: &SchemaRef, - range_key: &[i32], - non_range_key: &[i32], - value: &[i32], -) -> RecordBatch { - RecordBatch::try_new( - Arc::clone(schema), - vec![ - Arc::new(Int32Array::from(range_key.to_vec())), - Arc::new(Int32Array::from(non_range_key.to_vec())), - Arc::new(Int32Array::from(value.to_vec())), - ], - ) - .expect("range partition batch should be valid") +fn register_csv_listing_table( + ctx: &SessionContext, + name: &str, + table_dir: impl AsRef, + schema: Arc, + partitions: impl IntoIterator, + output_partitioning: Option, +) { + let table_dir = table_dir.as_ref(); + if table_dir.exists() { + remove_dir_all(table_dir).expect("test table dir should be removable"); + } + create_dir_all(table_dir).expect("test table dir should be created"); + for (idx, rows) in partitions.into_iter().enumerate() { + write(table_dir.join(format!("part-{idx}.csv")), rows) + .expect("test table csv partition should be written"); + } + + let table_path = format!( + "{}/", + table_dir + .to_str() + .expect("test table path should be valid utf8") + ); + let table_url = + ListingTableUrl::parse(&table_path).expect("test table url should parse"); + let options = + ListingOptions::new(Arc::new(CsvFormat::default().with_has_header(false))) + .with_output_partitioning(output_partitioning); + let config = ListingTableConfig::new(table_url) + .with_listing_options(options) + .with_schema(schema); + let table = + ListingTable::try_new(config).expect("test listing table should be valid"); + + ctx.register_table(name, Arc::new(table)) + .expect("test listing table registration should succeed"); } diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index a61f17a039eb8..2b7a2cfdf4083 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -16,7 +16,7 @@ # under the License. # The sqllogictest harness registers range_partitioned(range_key, non_range_key, value) -# as an in-memory source with four physical source partitions: +# as a CSV ListingTable with four declared range-partitioned file groups: # # partition 0: range_key in [..., 10), rows (1, 1, 10), (5, 2, 50) # partition 1: range_key in [10, 20), rows (10, 1, 100), (15, 2, 150) @@ -40,7 +40,7 @@ physical_plan 01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] 02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] -04)------DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query II SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key ORDER BY range_key; @@ -69,7 +69,7 @@ physical_plan 01)AggregateExec: mode=FinalPartitioned, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] 02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 03)----AggregateExec: mode=Partial, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] -04)------DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1], output_partitioning=UnknownPartitioning(4) +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false query II SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key ORDER BY non_range_key; @@ -104,8 +104,8 @@ SELECT range_key, value FROM range_partitioned; ---- physical_plan 01)UnionExec -02)--DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) -03)--DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +02)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query II SELECT range_key, value FROM range_partitioned From 7a056487d972908152453d54868029f10ea0fd12 Mon Sep 17 00:00:00 2001 From: Saad Tajwar <59696464+saadtajwar@users.noreply.github.com> Date: Mon, 22 Jun 2026 17:04:44 -0400 Subject: [PATCH 06/22] feat: logical plan protobuf representation for range repartitioning (#23030) - Closes #22787 The range repartitioning scheme for logical plans does not currently have a protobuf representation. A protobuf representation of the `RangeRepartition` struct was added to `datafusion.proto`, and the codegened Rust types were created. Added logic for serializing and deserializing to and from the protobuf representation, and a roundtrip test as well! Yes! Added a test in `roundtrip_logical_plan` No, adding internal protobuf serialization support for an existing logical plan variant --- datafusion/proto/proto/datafusion.proto | 10 + datafusion/proto/src/generated/pbjson.rs | 214 ++++++++++++++++++ datafusion/proto/src/generated/prost.rs | 16 +- .../proto/src/logical_plan/from_proto.rs | 13 +- datafusion/proto/src/logical_plan/mod.rs | 31 ++- datafusion/proto/src/logical_plan/to_proto.rs | 14 +- .../tests/cases/roundtrip_logical_plan.rs | 115 ++++++++-- 7 files changed, 386 insertions(+), 27 deletions(-) diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 23a0edc9ef23c..bfa40e61a8983 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -148,9 +148,19 @@ message RepartitionNode { oneof partition_method { uint64 round_robin = 2; HashRepartition hash = 3; + RangeRepartition range = 4; } } +message RangeSplitPoint { + repeated datafusion_common.ScalarValue value = 1; +} + +message RangeRepartition { + repeated SortExprNode sort_expr = 1; + repeated RangeSplitPoint split_point = 2; +} + message HashRepartition { repeated LogicalExprNode hash_expr = 1; uint64 partition_count = 2; diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index 1e85d14f8cae1..9f5785f4cc64a 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -22238,6 +22238,207 @@ impl<'de> serde::Deserialize<'de> for ProjectionNode { deserializer.deserialize_struct("datafusion.ProjectionNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for RangeRepartition { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + 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.RangeRepartition", 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 RangeRepartition { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + 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(deserializer: D) -> std::result::Result + 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(self, value: &str) -> std::result::Result + 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 = RangeRepartition; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.RangeRepartition") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + 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(RangeRepartition { + sort_expr: sort_expr__.unwrap_or_default(), + split_point: split_point__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.RangeRepartition", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for RangeSplitPoint { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + 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.RangeSplitPoint", len)?; + if !self.value.is_empty() { + struct_ser.serialize_field("value", &self.value)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for RangeSplitPoint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + 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(deserializer: D) -> std::result::Result + 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(self, value: &str) -> std::result::Result + 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 = RangeSplitPoint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.RangeSplitPoint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + 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(RangeSplitPoint { + value: value__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.RangeSplitPoint", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for RecursionUnnestOption { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -22666,6 +22867,9 @@ impl serde::Serialize for RepartitionNode { repartition_node::PartitionMethod::Hash(v) => { struct_ser.serialize_field("hash", v)?; } + repartition_node::PartitionMethod::Range(v) => { + struct_ser.serialize_field("range", v)?; + } } } struct_ser.end() @@ -22682,6 +22886,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { "round_robin", "roundRobin", "hash", + "range", ]; #[allow(clippy::enum_variant_names)] @@ -22689,6 +22894,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { Input, RoundRobin, Hash, + Range, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -22713,6 +22919,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { "input" => Ok(GeneratedField::Input), "roundRobin" | "round_robin" => Ok(GeneratedField::RoundRobin), "hash" => Ok(GeneratedField::Hash), + "range" => Ok(GeneratedField::Range), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -22753,6 +22960,13 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { return Err(serde::de::Error::duplicate_field("hash")); } partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(repartition_node::PartitionMethod::Hash) +; + } + GeneratedField::Range => { + if partition_method__.is_some() { + return Err(serde::de::Error::duplicate_field("range")); + } + partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(repartition_node::PartitionMethod::Range) ; } } diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index e61189f69a222..06db2b35bbf72 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -214,7 +214,7 @@ pub struct SortNode { pub struct RepartitionNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - #[prost(oneof = "repartition_node::PartitionMethod", tags = "2, 3")] + #[prost(oneof = "repartition_node::PartitionMethod", tags = "2, 3, 4")] pub partition_method: ::core::option::Option, } /// Nested message and enum types in `RepartitionNode`. @@ -225,9 +225,23 @@ pub mod repartition_node { RoundRobin(u64), #[prost(message, tag = "3")] Hash(super::HashRepartition), + #[prost(message, tag = "4")] + Range(super::RangeRepartition), } } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct RangeSplitPoint { + #[prost(message, repeated, tag = "1")] + pub value: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RangeRepartition { + #[prost(message, repeated, tag = "1")] + pub sort_expr: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "2")] + pub split_point: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct HashRepartition { #[prost(message, repeated, tag = "1")] pub hash_expr: ::prost::alloc::vec::Vec, diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index cff90ba7e31b7..bed64879f5a92 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field}; use datafusion_common::datatype::DataTypeExt; use datafusion_common::{ - NullEquality, RecursionUnnestOption, Result, ScalarValue, TableReference, + NullEquality, RecursionUnnestOption, Result, ScalarValue, SplitPoint, TableReference, UnnestOptions, exec_datafusion_err, internal_err, plan_datafusion_err, }; use datafusion_execution::TaskContext; @@ -834,3 +834,14 @@ fn parse_required_expr( fn proto_error>(message: S) -> Error { Error::General(message.into()) } + +pub(super) fn parse_protobuf_range_split_point( + split_point: &protobuf::RangeSplitPoint, +) -> Result { + let values = split_point + .value + .iter() + .map(ScalarValue::try_from) + .collect::>()?; + Ok(SplitPoint::new(values)) +} diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 41449112df8d0..5b2acfcac0e60 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -56,8 +56,8 @@ use datafusion_datasource_json::file_format::{ #[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RecursiveQuery, SkipType, - TableSource, Unnest, + AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning, + RecursiveQuery, SkipType, TableSource, Unnest, }; use datafusion_expr::{ DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, @@ -70,6 +70,7 @@ use datafusion_expr::{ }; use self::to_proto::{serialize_expr, serialize_exprs}; +use crate::logical_plan::to_proto::serialize_range_split_point; use crate::logical_plan::to_proto::serialize_sorts; use datafusion_catalog::TableProvider; use datafusion_catalog::default_table_source::{provider_as_source, source_as_provider}; @@ -675,6 +676,16 @@ impl AsLogicalPlan for LogicalPlanNode { PartitionMethod::RoundRobin(partition_count) => { Partitioning::RoundRobinBatch(*partition_count as usize) } + PartitionMethod::Range(protobuf::RangeRepartition { + sort_expr: pb_sort_expr, + split_point, + }) => Partitioning::Range(RangePartitioning::try_new( + from_proto::parse_sorts(pb_sort_expr, ctx, extension_codec)?, + split_point + .iter() + .map(from_proto::parse_protobuf_range_split_point) + .collect::, _>>()?, + )?), }; LogicalPlanBuilder::from(input) @@ -1631,10 +1642,18 @@ impl AsLogicalPlan for LogicalPlanNode { Partitioning::RoundRobinBatch(partition_count) => { PartitionMethod::RoundRobin(*partition_count as u64) } - Partitioning::Range(_) => { - // TODO: Support range repartition protobuf serialization. - // Tracked by https://github.com/apache/datafusion/issues/22787 - return not_impl_err!("Range repartition"); + Partitioning::Range(range_partitioning) => { + let ordering = range_partitioning.ordering(); + let split_point = range_partitioning + .split_points() + .iter() + .map(serialize_range_split_point) + .collect::, _>>()?; + + PartitionMethod::Range(protobuf::RangeRepartition { + sort_expr: serialize_sorts(ordering, extension_codec)?, + split_point, + }) } Partitioning::DistributeBy(_) => { return not_impl_err!("DistributeBy"); diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 0599e48d2795c..86b980a109b90 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -21,7 +21,7 @@ use std::collections::HashMap; -use datafusion_common::{NullEquality, TableReference, UnnestOptions}; +use datafusion_common::{NullEquality, SplitPoint, TableReference, UnnestOptions}; use datafusion_expr::WriteOp; use datafusion_expr::dml::InsertOp; use datafusion_expr::expr::{ @@ -706,6 +706,18 @@ where .collect::, Error>>() } +pub(super) fn serialize_range_split_point( + split_point: &SplitPoint, +) -> Result { + Ok(protobuf::RangeSplitPoint { + value: split_point + .values() + .iter() + .map(TryInto::::try_into) + .collect::>()?, + }) +} + impl From for protobuf::TableReference { fn from(t: TableReference) -> Self { use protobuf::table_reference::TableReferenceEnum; diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index fbe1af5617e42..3bb5b898677c7 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -45,6 +45,7 @@ use std::vec; use datafusion::catalog::{TableProvider, TableProviderFactory}; use datafusion::datasource::DefaultTableSource; +use datafusion::datasource::empty::EmptyTable; use datafusion::datasource::file_format::arrow::ArrowFormatFactory; use datafusion::datasource::file_format::csv::CsvFormatFactory; use datafusion::datasource::file_format::parquet::ParquetFormatFactory; @@ -71,8 +72,8 @@ use datafusion_common::config::TableOptions; use datafusion_common::format::ExplainFormat; use datafusion_common::scalar::ScalarStructBuilder; use datafusion_common::{ - DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference, - internal_datafusion_err, internal_err, not_impl_err, plan_err, + DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SplitPoint, + TableReference, internal_datafusion_err, internal_err, not_impl_err, plan_err, }; use datafusion_execution::TaskContext; use datafusion_expr::dml::CopyTo; @@ -84,8 +85,8 @@ use datafusion_expr::logical_plan::{Extension, UserDefinedLogicalNodeCore}; use datafusion_expr::{ Accumulator, AggregateUDF, ColumnarValue, ExprFunctionExt, ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, - PartitionEvaluator, ScalarUDF, Signature, TryCast, Volatility, WindowFrame, - WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, WindowUDF, + RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, Volatility, + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, WindowUDF, WindowUDFImpl, }; use datafusion_functions_aggregate::average::avg_udaf; @@ -3276,9 +3277,7 @@ async fn roundtrip_empty_table_scan() -> Result<()> { Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); - let table = Arc::new(datafusion::datasource::empty::EmptyTable::new(Arc::clone( - &schema, - ))); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); let ctx = SessionContext::new(); ctx.register_table("empty", table)?; @@ -3300,9 +3299,7 @@ async fn roundtrip_empty_table_scan_with_projection() -> Result<()> { Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); - let table = Arc::new(datafusion::datasource::empty::EmptyTable::new(Arc::clone( - &schema, - ))); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); let ctx = SessionContext::new(); ctx.register_table("empty", table)?; @@ -3378,14 +3375,8 @@ async fn roundtrip_join_null_equality() -> Result<()> { let left_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); let right_schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])); - ctx.register_table( - "t1", - Arc::new(datafusion::datasource::empty::EmptyTable::new(left_schema)), - )?; - ctx.register_table( - "t2", - Arc::new(datafusion::datasource::empty::EmptyTable::new(right_schema)), - )?; + ctx.register_table("t1", Arc::new(EmptyTable::new(left_schema)))?; + ctx.register_table("t2", Arc::new(EmptyTable::new(right_schema)))?; let left = ctx.table("t1").await?.into_optimized_plan()?; let right = ctx.table("t2").await?.into_optimized_plan()?; @@ -3406,3 +3397,91 @@ async fn roundtrip_join_null_equality() -> Result<()> { Ok(()) } + +// Single column, single split point range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_single_col() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} + +// Multi-column compound key with multiple split points for range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_multi_col() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("region", DataType::Utf8, false), + ])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("ts").sort(true, true), col("region").sort(true, true)], + vec![ + SplitPoint::new(vec![ + ScalarValue::Int64(Some(1000)), + ScalarValue::Utf8(Some("east".to_string())), + ]), + SplitPoint::new(vec![ + ScalarValue::Int64(Some(2000)), + ScalarValue::Utf8(Some("west".to_string())), + ]), + ], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} + +// Non-default sort options: descending with nulls last for range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_desc_nulls_last() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "score", + DataType::Float64, + true, + )])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("score").sort(false, false)], + vec![SplitPoint::new(vec![ScalarValue::Float64(Some(50.0))])], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} From 22159add9998ca0e72f33defdf414dcca7a97a11 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Wed, 1 Jul 2026 01:50:52 -0400 Subject: [PATCH 07/22] Add `Distribution::HashPartitioned` to `Distribution::KeyPartitioned` API bridge (#23259) - Closes #23236. `HashPartitioned` is historical naming for a key-partitioning requirement. This keeps the old variant as a deprecated compatibility bridge while moving DataFusion internals to `KeyPartitioned`. Adds `KeyPartitioned`, deprecates `HashPartitioned`, and treats both equivalently during the transition to avoid breaking changes for downstream consumers. A blast radius report on this was done here #23241 Yes. Yes, `Distribution::HashPartitioned` is deprecated and trainsitioned to `Distribution::KeyPartitioned`. --- .../physical_optimizer/projection_pushdown.rs | 10 +- datafusion/physical-expr/src/partitioning.rs | 128 ++++++++++++------ .../src/enforce_distribution.rs | 6 +- .../src/output_requirements.rs | 23 ++-- .../physical-optimizer/src/sanity_checker.rs | 4 + .../physical-plan/src/aggregates/mod.rs | 2 +- .../physical-plan/src/joins/hash_join/exec.rs | 15 +- .../src/joins/sort_merge_join/exec.rs | 4 +- .../src/joins/symmetric_hash_join.rs | 4 +- .../src/sorts/partitioned_topk.rs | 2 +- .../src/windows/bounded_window_agg_exec.rs | 2 +- .../src/windows/window_agg_exec.rs | 2 +- 12 files changed, 119 insertions(+), 83 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 6f88e01059fc9..21f03b14f58bd 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -724,7 +724,7 @@ fn test_output_req_after_projection() -> Result<()> { ] .into(), )), - Distribution::HashPartitioned(vec![ + Distribution::KeyPartitioned(vec![ Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1)), ]), @@ -746,7 +746,7 @@ fn test_output_req_after_projection() -> Result<()> { actual, @r" ProjectionExec: expr=[c@2 as c, a@0 as new_a, b@1 as b] - OutputRequirementExec: order_by=[(b@1, asc), (c@2 + a@0, asc)], dist_by=HashPartitioned[[a@0, b@1]]) + OutputRequirementExec: order_by=[(b@1, asc), (c@2 + a@0, asc)], dist_by=KeyPartitioned[[a@0, b@1]]) DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false " ); @@ -762,7 +762,7 @@ fn test_output_req_after_projection() -> Result<()> { assert_snapshot!( actual, @r" - OutputRequirementExec: order_by=[(b@2, asc), (c@0 + new_a@1, asc)], dist_by=HashPartitioned[[new_a@1, b@2]]) + OutputRequirementExec: order_by=[(b@2, asc), (c@0 + new_a@1, asc)], dist_by=KeyPartitioned[[new_a@1, b@2]]) DataSourceExec: file_groups={1 group: [[x]]}, projection=[c, a@0 as new_a, b], file_type=csv, has_header=false " ); @@ -797,7 +797,7 @@ fn test_output_req_after_projection() -> Result<()> { Arc::new(Column::new("new_a", 1)), Arc::new(Column::new("b", 2)), ]; - if let Distribution::HashPartitioned(vec) = after_optimize + if let Distribution::KeyPartitioned(vec) = after_optimize .downcast_ref::() .unwrap() .required_input_distribution()[0] @@ -809,7 +809,7 @@ fn test_output_req_after_projection() -> Result<()> { .all(|(actual, expected)| actual.eq(&expected)) ); } else { - panic!("Expected HashPartitioned distribution!"); + panic!("Expected KeyPartitioned distribution!"); }; Ok(()) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 8ac5de2a87b9d..a777f9df14a80 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -518,6 +518,10 @@ impl Partitioning { /// Returns how this [`Partitioning`] satisfies the partitioning scheme mandated /// by the `required` [`Distribution`]. + #[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" + )] pub fn satisfaction( &self, required: &Distribution, @@ -639,8 +643,9 @@ pub enum Distribution { UnspecifiedDistribution, /// A single partition is required SinglePartition, - /// Requires children to be distributed in such a way that the same - /// values of the keys end up in the same partition + /// Deprecated historical name for [`Distribution::KeyPartitioned`]. + /// See for details. + #[deprecated(since = "55.0.0", note = "Use Distribution::KeyPartitioned")] HashPartitioned(Vec>), /// Requires rows with equal values for the given keys to be colocated in /// the same partition, without requiring a specific partitioning algorithm. @@ -657,6 +662,10 @@ pub enum Distribution { KeyPartitioned(Vec>), } +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] impl Distribution { /// Returns key expressions for distribution variants that require /// co-locating equal key values. @@ -681,6 +690,10 @@ impl Distribution { } } +#[expect( + deprecated, + reason = "HashPartitioned display is preserved during the KeyPartitioned migration" +)] impl Display for Distribution { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { @@ -760,11 +773,11 @@ mod tests { Partitioning::Hash(self.cols(indices), partition_count) } - fn hash_distribution( + fn key_distribution( &self, indices: impl IntoIterator, ) -> Distribution { - Distribution::HashPartitioned(self.cols(indices)) + Distribution::KeyPartitioned(self.cols(indices)) } fn key_partitioned_distribution( @@ -824,6 +837,10 @@ mod tests { } #[test] + #[expect( + deprecated, + reason = "test intentionally covers deprecated HashPartitioned compatibility" + )] fn partitioning_satisfy_distribution() -> Result<()> { let fixture = PartitioningTestFixture::new(vec![ ("column_1", DataType::Int64), @@ -869,7 +886,7 @@ mod tests { Distribution::SinglePartition => { assert_eq!(result, (true, false, false, false, false)) } - Distribution::HashPartitioned(_) => { + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) => { assert_eq!(result, (true, false, false, true, false)) } Distribution::KeyPartitioned(_) => { @@ -881,43 +898,66 @@ mod tests { Ok(()) } + #[test] + #[expect( + deprecated, + reason = "test intentionally covers deprecated HashPartitioned compatibility" + )] + fn deprecated_hash_partitioned_matches_key_partitioned() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let partitioning = fixture.hash_partitioning([0, 1], 4); + let hash_distribution = Distribution::HashPartitioned(fixture.cols([0, 1])); + let key_distribution = fixture.key_distribution([0, 1]); + + assert_eq!( + partitioning.satisfaction(&hash_distribution, &fixture.eq_properties, false), + partitioning.satisfaction(&key_distribution, &fixture.eq_properties, false) + ); + assert_eq!( + hash_distribution.create_partitioning(4), + key_distribution.create_partitioning(4) + ); + + Ok(()) + } + #[test] fn test_partitioning_satisfy_by_subset() -> Result<()> { let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; let test_cases = vec![ ( - "Hash([a]) vs Hash([a, b])", + "KeyPartitioned([a, b]) satisfied by Hash([a])", fixture.hash_partitioning([0], 4), - fixture.hash_distribution([0, 1]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a]) vs Hash([a, b, c])", + "KeyPartitioned([a, b, c]) satisfied by Hash([a])", fixture.hash_partitioning([0], 4), - fixture.hash_distribution([0, 1, 2]), + fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b]) vs Hash([a, b, c])", + "KeyPartitioned([a, b, c]) satisfied by Hash([a, b])", fixture.hash_partitioning([0, 1], 4), - fixture.hash_distribution([0, 1, 2]), + fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([b]) vs Hash([a, b, c])", + "KeyPartitioned([a, b, c]) satisfied by Hash([b])", fixture.hash_partitioning([1], 4), - fixture.hash_distribution([0, 1, 2]), + fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([b, a]) vs Hash([a, b, c])", + "KeyPartitioned([a, b, c]) satisfied by Hash([b, a])", fixture.hash_partitioning([1, 0], 4), - fixture.hash_distribution([0, 1, 2]), + fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), @@ -948,23 +988,23 @@ mod tests { let test_cases = vec![ ( - "Hash([a, b]) vs Hash([a])", + "KeyPartitioned([a]) satisfied by Hash([a, b])", fixture.hash_partitioning([0, 1], 4), - fixture.hash_distribution([0]), + fixture.key_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b, c]) vs Hash([a])", + "KeyPartitioned([a]) satisfied by Hash([a, b, c])", fixture.hash_partitioning([0, 1, 2], 4), - fixture.hash_distribution([0]), + fixture.key_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b, c]) vs Hash([a, b])", + "KeyPartitioned([a, b]) satisfied by Hash([a, b, c])", fixture.hash_partitioning([0, 1, 2], 4), - fixture.hash_distribution([0, 1]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -994,9 +1034,9 @@ mod tests { let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; let test_cases = vec![( - "Partial overlap: Hash([a, c]) vs Hash([a, b])", + "Partial overlap: KeyPartitioned([a, b]) satisfied by Hash([a, c])", fixture.hash_partitioning([0, 2], 4), - fixture.hash_distribution([0, 1]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, )]; @@ -1026,16 +1066,16 @@ mod tests { let test_cases = vec![ ( - "Hash([a]) vs Hash([b, c])", + "KeyPartitioned([b, c]) satisfied by Hash([a])", fixture.hash_partitioning([0], 4), - fixture.hash_distribution([1, 2]), + fixture.key_distribution([1, 2]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b]) vs Hash([c])", + "KeyPartitioned([c]) satisfied by Hash([a, b])", fixture.hash_partitioning([0, 1], 4), - fixture.hash_distribution([2]), + fixture.key_distribution([2]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -1066,16 +1106,16 @@ mod tests { let test_cases = vec![ ( - "Hash([a, b]) vs Hash([a, b])", + "KeyPartitioned([a, b]) satisfied by Hash([a, b])", fixture.hash_partitioning([0, 1], 4), - fixture.hash_distribution([0, 1]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::Exact, PartitioningSatisfaction::Exact, ), ( - "Hash([a]) vs Hash([a])", + "KeyPartitioned([a]) satisfied by Hash([a])", fixture.hash_partitioning([0], 4), - fixture.hash_distribution([0]), + fixture.key_distribution([0]), PartitioningSatisfaction::Exact, PartitioningSatisfaction::Exact, ), @@ -1107,23 +1147,23 @@ mod tests { let test_cases = vec![ ( - "Hash([unknown]) vs Hash([a, b])", + "KeyPartitioned([a, b]) satisfied by Hash([unknown])", Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - fixture.hash_distribution([0, 1]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b]) vs Hash([unknown])", + "KeyPartitioned([unknown]) satisfied by Hash([a, b])", fixture.hash_partitioning([0, 1], 4), - Distribution::HashPartitioned(vec![Arc::clone(&unknown)]), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([unknown]) vs Hash([unknown])", + "KeyPartitioned([unknown]) satisfied by Hash([unknown])", Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&unknown)]), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -1154,23 +1194,23 @@ mod tests { let test_cases = vec![ ( - "Hash([]) vs Hash([a])", + "KeyPartitioned([a]) satisfied by Hash([])", Partitioning::Hash(vec![], 4), - fixture.hash_distribution([0]), + fixture.key_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a]) vs Hash([])", + "KeyPartitioned([]) satisfied by Hash([a])", fixture.hash_partitioning([0], 4), - Distribution::HashPartitioned(vec![]), + Distribution::KeyPartitioned(vec![]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([]) vs Hash([])", + "KeyPartitioned([]) satisfied by Hash([])", Partitioning::Hash(vec![], 4), - Distribution::HashPartitioned(vec![]), + Distribution::KeyPartitioned(vec![]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -1573,7 +1613,7 @@ mod tests { 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]); + let required = fixture.key_distribution([0, 1]); assert_eq!( range_partitioning.satisfaction(&required, &fixture.eq_properties, false), diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index b501d6ef8537d..b45f59bccf72e 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -888,7 +888,7 @@ fn add_hash_on_top( return Ok(input); } - let dist = Distribution::HashPartitioned(hash_exprs); + let dist = Distribution::KeyPartitioned(hash_exprs); let current_partitions = input.plan.output_partitioning().partition_count(); let satisfaction = input.plan.output_partitioning().satisfaction( &dist, @@ -1326,6 +1326,10 @@ fn partitioned_join_distribution( /// This function is intended to be used in a bottom up traversal, as it /// can first repartition (or newly partition) at the datasources -- these /// source partitions may be later repartitioned with additional data exchange operators. +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] pub fn ensure_distribution( dist_context: DistributionContext, config: &ConfigOptions, diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 679d1bc7b2e36..50c1ea8483d3e 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -244,6 +244,10 @@ impl ExecutionPlan for OutputRequirementExec { self.input.partition_statistics(partition) } + #[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" + )] fn try_swapping_with_projection( &self, projection: &ProjectionExec, @@ -268,9 +272,9 @@ impl ExecutionPlan for OutputRequirementExec { requirements = OrderingRequirements::new_alternatives(updated_reqs, soft); } - let dist_req = { - let dist_req = &self.required_input_distribution()[0]; - if let Some(exprs) = dist_req.key_exprs() { + let dist_req = match &self.required_input_distribution()[0] { + Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) => { let mut updated_exprs = vec![]; for expr in exprs { let Some(new_expr) = update_expr(expr, projection.expr(), false)? @@ -279,18 +283,9 @@ impl ExecutionPlan for OutputRequirementExec { }; updated_exprs.push(new_expr); } - match dist_req { - Distribution::HashPartitioned(_) => { - Distribution::HashPartitioned(updated_exprs) - } - Distribution::KeyPartitioned(_) => { - Distribution::KeyPartitioned(updated_exprs) - } - _ => unreachable!(), - } - } else { - dist_req.clone() + Distribution::KeyPartitioned(updated_exprs) } + dist => dist.clone(), }; make_with_child(projection, &self.input()).map(|input| { diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index 08c014593cdc2..4e74062268260 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -142,6 +142,10 @@ fn is_prunable(join: &SymmetricHashJoinExec) -> bool { /// Ensures that the plan is pipeline friendly and the order and /// distribution requirements from its children are satisfied. +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] pub fn check_plan_sanity( plan: &Arc, optimizer_options: &OptimizerOptions, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 5e6b8505764a2..4f2c2e58bf791 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1549,7 +1549,7 @@ impl ExecutionPlan for AggregateExec { vec![Distribution::UnspecifiedDistribution] } AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned => { - vec![Distribution::HashPartitioned(self.group_by.input_exprs())] + vec![Distribution::KeyPartitioned(self.group_by.input_exprs())] } AggregateMode::Final | AggregateMode::Single => { vec![Distribution::SinglePartition] diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 17d72a856f93b..f746ebcb6b9b1 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1279,17 +1279,10 @@ impl ExecutionPlan for HashJoinExec { .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - if self.join_type == JoinType::Inner { - vec![ - Distribution::KeyPartitioned(left_expr), - Distribution::KeyPartitioned(right_expr), - ] - } else { - vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), - ] - } + vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ] } PartitionMode::Auto => vec![ Distribution::UnspecifiedDistribution, diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 9e87b52696a57..1fe6f3bda922c 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -429,8 +429,8 @@ impl ExecutionPlan for SortMergeJoinExec { .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), ] } diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index ef92964fadf84..41caa8cba2d9a 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -434,8 +434,8 @@ impl ExecutionPlan for SymmetricHashJoinExec { .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) .unzip(); vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), ] } StreamJoinPartitionMode::SinglePartition => { diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index fe876eeddf7f2..5596783caf2e4 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -307,7 +307,7 @@ impl ExecutionPlan for PartitionedTopKExec { .iter() .map(|e| Arc::clone(&e.expr)) .collect(); - vec![Distribution::HashPartitioned(partition_exprs)] + vec![Distribution::KeyPartitioned(partition_exprs)] } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 6c6b26c9cf49f..4b31ce9b9245d 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -335,7 +335,7 @@ impl ExecutionPlan for BoundedWindowAggExec { debug!("No partition defined for BoundedWindowAggExec!!!"); vec![Distribution::SinglePartition] } else { - vec![Distribution::HashPartitioned(self.partition_keys().clone())] + vec![Distribution::KeyPartitioned(self.partition_keys().clone())] } } diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index ee3b071fc9167..0c77adc40dd92 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -243,7 +243,7 @@ impl ExecutionPlan for WindowAggExec { if self.partition_keys().is_empty() { vec![Distribution::SinglePartition] } else { - vec![Distribution::HashPartitioned(self.partition_keys())] + vec![Distribution::KeyPartitioned(self.partition_keys())] } } From d3a927141dd9af33bf6b36ca730c83b5f25ce31e Mon Sep 17 00:00:00 2001 From: Saad Tajwar <59696464+saadtajwar@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:48:16 -0700 Subject: [PATCH 08/22] feat: physical execution for range partitioning (#23231) ## Which issue does this PR close? - Closes #23137 ## Rationale for this change Range repartitioning was already planned and serialized into physical plans, but `RepartitionExec` could not execute it. This PR completes the core execution path so rows in an input batch are routed to the correct output partition based on range split points and the ordering defined on the partitioning scheme. ## What changes are included in this PR? This PR adds a `Range` variant to `BatchPartitioner` that evaluates the ordering expressions on each input batch, compares each row's key against split points using `compare_rows` (respecting ASC/DESC and null ordering), and assigns row indices to output partitions via binary search. The partitioned row indices are then materialized into sub-batches using the same `partition_grouped_take` path as hash repartitioning. `pull_from_input` is wired to construct a range partitioner for `Partitioning::Range`, replacing the previous `not_impl_err!` at execution time. Optimizer-related paths remain intentionally unimplemented and are tracked in [#23230](https://github.com/apache/datafusion/issues/23230): projection pushdown through `RepartitionExec` (`try_swapping_with_projection`), sort pushdown (`try_pushdown_sort`), and changing partition counts via `repartitioned()`. ## Are these changes tested? Yes! ## Are there any user-facing changes? No public API changes --- .../physical-plan/src/repartition/mod.rs | 590 ++++++++++++++---- 1 file changed, 483 insertions(+), 107 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 2a8005c10ae00..f1c06ccaff89d 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -19,10 +19,11 @@ //! partitions to M output partitions based on a partitioning scheme, optionally //! maintaining the order of the input rows in the output. +use std::cmp::Ordering; use std::fmt::{Debug, Formatter}; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::task::{Context, Poll}; use std::vec; @@ -45,20 +46,22 @@ use crate::{ check_if_same_properties, }; -use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions}; +use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions}; use arrow::compute::take_arrays; use arrow::datatypes::{SchemaRef, UInt32Type}; +use arrow_schema::SortOptions; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; -use datafusion_common::utils::transpose; +use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, assert_or_internal_err, internal_err, + ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, + assert_or_internal_err, internal_err, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryConsumer; -use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr, RangePartitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use crate::filter_pushdown::{ @@ -255,7 +258,7 @@ impl SharedCoalescer { /// sender, finalize the coalescer and return its residual batches; if /// other senders are still active, return `Ok(None)`. fn finalize(&self) -> Result> { - let was_last = self.active_senders.fetch_sub(1, Ordering::AcqRel) == 1; + let was_last = self.active_senders.fetch_sub(1, AtomicOrdering::AcqRel) == 1; if !was_last { return Ok(vec![]); } @@ -569,6 +572,18 @@ enum BatchPartitionerState { num_partitions: usize, next_idx: usize, }, + Range { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Sort options from the `LexOrdering` + sort_options: Vec, + /// Boundaries between adjacent partitions. + split_points: Vec, + /// Row indices grouped by output partition + indices: Vec>, + /// Buffer of `ScalarValue` used to represent the values for a row - based on the `LexOrdering` ordering - to compare against split points + partition_buffer: Vec, + }, } /// Fixed RandomState used for hash repartitioning to ensure consistent behavior across @@ -705,13 +720,40 @@ impl BatchPartitioner { timer, } } + + /// Create a new [`BatchPartitioner`] for range-based repartitioning. + /// + /// # Parameters + /// - `range_partitioning`: `RangePartitioning` struct used for ordering, split points, and number of partitions + /// - `timer`: Metric used to record time spent during repartitioning. + pub fn new_range_partitioner( + range_partitioning: &RangePartitioning, + timer: metrics::Time, + ) -> Self { + let ordering = range_partitioning.ordering().clone(); + let split_points = range_partitioning.split_points().to_vec(); + let num_partitions = range_partitioning.partition_count(); + let sort_options: Vec = ordering.iter().map(|e| e.options).collect(); + + Self { + state: BatchPartitionerState::Range { + partition_buffer: Vec::with_capacity(ordering.len()), + ordering, + sort_options, + split_points, + indices: vec![vec![]; num_partitions], + }, + timer, + } + } + /// Create a new [`BatchPartitioner`] based on the provided [`Partitioning`] scheme. /// /// This is a convenience constructor that delegates to the specialized - /// hash or round-robin constructors depending on the partitioning variant. + /// hash, round-robin, or range constructors depending on the partitioning variant. /// /// # Parameters - /// - `partitioning`: Partitioning scheme to apply (hash or round-robin). + /// - `partitioning`: Partitioning scheme to apply (hash, round-robin, or range). /// - `timer`: Metric used to record time spent during repartitioning. /// - `input_partition`: Index of the current input partition. /// - `num_input_partitions`: Total number of input partitions. @@ -737,12 +779,8 @@ 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" - ) + Partitioning::Range(range_repartitioning) => { + Ok(Self::new_range_partitioner(&range_repartitioning, timer)) } other => { not_impl_err!("Unsupported repartitioning scheme {other:?}") @@ -830,22 +868,95 @@ impl BatchPartitioner { Box::new(partitioned_batches.into_iter()) } + BatchPartitionerState::Range { + ordering, + sort_options, + split_points, + indices, + partition_buffer, + } => { + // Tracking time required for distributing indexes across output partitions + let timer = self.timer.timer(); + if split_points.is_empty() { + timer.done(); + Box::new(std::iter::once(Ok((0, batch)))) + } else { + let arrays = evaluate_expressions_to_arrays( + ordering.iter().map(|e| &e.expr), + &batch, + )?; + + indices.iter_mut().for_each(|v| v.clear()); + + Self::partition_range_indices( + &arrays, + split_points, + sort_options, + partition_buffer, + indices, + )?; + + // Finished building index-arrays for output partitions + timer.done(); + + let partitioned_batches = + Self::partition_grouped_take(&batch, indices, &self.timer)?; + + Box::new(partitioned_batches.into_iter()) + } + } }; Ok(it) } + /// Groups input row indices by range partition. This populates `indices[p]` with the + /// row indices from `arrays` that belong in output partition `p` according to `split_points` and `sort_options`. + fn partition_range_indices( + arrays: &[Arc], + split_points: &[SplitPoint], + sort_options: &[SortOptions], + row_key_buffer: &mut Vec, + indices: &mut [Vec], + ) -> Result<()> { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + // Note that `extract_row_at_idx_to_buf` clears the `row_key_buffer` on each invocation, creating a new row key for comparison for each row + extract_row_at_idx_to_buf(arrays, row_idx, row_key_buffer)?; + + let mut low = 0; + let mut high = split_points.len(); + while low < high { + let mid = low + (high - low) / 2; + let comparison = compare_rows( + row_key_buffer, + split_points[mid].values(), + sort_options, + )?; + match comparison { + Ordering::Less => high = mid, + Ordering::Equal | Ordering::Greater => low = mid + 1, + } + } + + indices[low].push(row_idx as u32) + } + + Ok(()) + } + // return the number of output partitions fn num_partitions(&self) -> usize { match &self.state { BatchPartitionerState::RoundRobin { num_partitions, .. } => *num_partitions, - BatchPartitionerState::Hash { indices, .. } => indices.len(), + BatchPartitionerState::Hash { indices, .. } + | BatchPartitionerState::Range { indices, .. } => indices.len(), } } - /// Build repartitioned hash output batches using one `take` per input batch. + /// Build repartitioned hash/range output batches using one `take` per input batch. /// - /// The hash router first fills one index vector per output partition. This method + /// The routers first fills one index vector per output partition. This method /// concatenates those index vectors, performs one grouped `take_arrays`, and /// then returns each output partition as a slice of the reordered batch. /// @@ -1439,10 +1550,8 @@ impl ExecutionPlan for RepartitionExec { } 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" - ); + // https://github.com/apache/datafusion/issues/23230 + return Ok(None); } others => others.clone(), }; @@ -1483,10 +1592,8 @@ impl ExecutionPlan for RepartitionExec { 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" - ); + // https://github.com/apache/datafusion/issues/23230 + return Ok(SortOrderPushdownResult::Unsupported); } Partitioning::RoundRobinBatch(_) | Partitioning::Hash(_, _) @@ -1516,11 +1623,9 @@ impl ExecutionPlan for RepartitionExec { 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" - ); + // Range repartition optimizations are tracked in + // https://github.com/apache/datafusion/issues/23230 + return Ok(None); } }; Ok(Some(Arc::new(Self { @@ -1634,33 +1739,12 @@ impl RepartitionExec { input_partition: usize, num_input_partitions: usize, ) -> Result<()> { - let mut partitioner = match &partitioning { - Partitioning::Hash(exprs, num_partitions) => { - BatchPartitioner::new_hash_partitioner( - exprs.clone(), - *num_partitions, - metrics.repartition_time.clone(), - )? - } - Partitioning::RoundRobinBatch(num_partitions) => { - BatchPartitioner::new_round_robin_partitioner( - *num_partitions, - metrics.repartition_time.clone(), - input_partition, - 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:?}"); - } - }; + let mut partitioner = BatchPartitioner::try_new( + partitioning, + metrics.repartition_time.clone(), + input_partition, + num_input_partitions, + )?; // While there are still outputs to send to, keep pulling inputs let mut batches_until_yield = partitioner.num_partitions(); @@ -2009,7 +2093,7 @@ 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::cast::{as_string_array, as_uint32_array}; use datafusion_common::exec_err; use datafusion_common::test_util::batches_to_sort_string; use datafusion_common_runtime::JoinSet; @@ -2113,7 +2197,7 @@ mod tests { #[tokio::test] async fn one_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition]; @@ -2136,7 +2220,7 @@ mod tests { #[tokio::test] async fn many_to_one_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2153,7 +2237,7 @@ mod tests { #[tokio::test] async fn many_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2174,7 +2258,7 @@ mod tests { #[tokio::test] async fn many_to_many_hash_partition() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2196,9 +2280,267 @@ mod tests { Ok(()) } + #[tokio::test] + async fn many_to_many_range_partition() -> Result<()> { + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; + + // create_batch values are [1, 2, 3, 4, 5, 6, 7, 8]; split at 3 and 6 yields + // 2, 3, and 3 rows per batch respectively + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![3, 6])?; + + let output_partitions = repartition(&schema, partitions, partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!(300, partition_row_count(&output_partitions[0])); + assert_eq!(450, partition_row_count(&output_partitions[1])); + assert_eq!(450, partition_row_count(&output_partitions[2])); + assert_eq!( + collect_partition_u32_values(&output_partitions[0]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([1, 2]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[1]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([3, 4, 5]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[2]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([6, 7, 8]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_compound_keys() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![5, 10, 10, 10, 10, 15])), + Arc::new(UInt32Array::from(vec![1, 1, 3, 5, 7, 0])), + ], + )?; + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, SortOptions::default()), + PhysicalSortExpr::new(col("b", &schema)?, SortOptions::default()), + ] + .into(), + vec![ + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(1)), + ]), + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(5)), + ]), + ], + )?); + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![(5, 1)], + collect_partition_u32_pairs(&output_partitions[0]) + ); + assert_eq!( + vec![(10, 1), (10, 3)], + collect_partition_u32_pairs(&output_partitions[1]) + ); + assert_eq!( + vec![(10, 5), (10, 7), (15, 0)], + collect_partition_u32_pairs(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_last() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, false), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![None, Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_first() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, true), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![None, Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_asc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![10, 20])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_desc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 20, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(true, false), vec![20, 10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(15), Some(20)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(5), Some(10)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_string_rows() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["bar", "baz", "foo", "qux"])) 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::try_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 mut partition_0 = Vec::new(); + let mut stream = exec.execute(0, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + partition_0.push(result?); + } + + let mut partition_1 = Vec::new(); + let mut stream = exec.execute(1, task_ctx)?; + while let Some(result) = stream.next().await { + partition_1.push(result?); + } + + assert_eq!( + vec!["bar", "baz"], + collect_partition_string_values(&partition_0) + ); + assert_eq!( + vec!["foo", "qux"], + collect_partition_string_values(&partition_1) + ); + + Ok(()) + } + #[tokio::test] async fn test_repartition_with_coalescing() -> Result<()> { - let schema = test_schema(); + let schema = test_schema(false); // create 50 batches, each having 8 rows let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone()]; @@ -2222,8 +2564,76 @@ mod tests { Ok(()) } - fn test_schema() -> Arc { - Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) + fn test_schema(nullable: bool) -> Arc { + Arc::new(Schema::new(vec![Field::new( + "c0", + DataType::UInt32, + nullable, + )])) + } + + fn u32_range_partitioning( + schema: &SchemaRef, + sort_options: SortOptions, + split_values: Vec, + ) -> Result { + let expr = col("c0", schema)?; + Ok(Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(expr, sort_options)].into(), + split_values + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::UInt32(Some(value))])) + .collect(), + )?)) + } + + fn partition_row_count(batches: &[RecordBatch]) -> usize { + batches.iter().map(|batch| batch.num_rows()).sum() + } + + fn collect_partition_u32_values(batches: &[RecordBatch]) -> Vec> { + batches + .iter() + .flat_map(|batch| { + let array = + as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + (0..array.len()) + .map(|idx| { + if array.is_null(idx) { + None + } else { + Some(array.value(idx)) + } + }) + .collect::>() + }) + .collect() + } + + fn collect_partition_u32_pairs(batches: &[RecordBatch]) -> Vec<(u32, u32)> { + batches + .iter() + .flat_map(|batch| { + let a = as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + let b = as_uint32_array(batch.column(1)).expect("expected UInt32 column"); + (0..a.len()) + .map(|idx| (a.value(idx), b.value(idx))) + .collect::>() + }) + .collect() + } + + fn collect_partition_string_values(batches: &[RecordBatch]) -> Vec<&str> { + batches + .iter() + .flat_map(|batch| { + let array = + as_string_array(batch.column(0)).expect("expected Utf8 column"); + (0..array.len()) + .map(|idx| array.value(idx)) + .collect::>() + }) + .collect() } async fn repartition( @@ -2256,7 +2666,7 @@ mod tests { let handle: SpawnedTask>>> = SpawnedTask::spawn(async move { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2308,40 +2718,6 @@ 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 @@ -2650,7 +3026,7 @@ mod tests { #[tokio::test] async fn repartition_with_spilling() -> Result<()> { // Test that repartition successfully spills to disk when memory is constrained - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2712,7 +3088,7 @@ mod tests { #[tokio::test] async fn repartition_with_partial_spilling() -> Result<()> { // Test that repartition can handle partial spilling (some batches in memory, some spilled) - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2782,7 +3158,7 @@ mod tests { #[tokio::test] async fn repartition_without_spilling() -> Result<()> { // Test that repartition does not spill when there's ample memory - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2844,7 +3220,7 @@ mod tests { use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; // Test that repartition fails with OOM when disk manager is disabled - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2887,7 +3263,7 @@ mod tests { /// Create batch fn create_batch() -> RecordBatch { - let schema = test_schema(); + let schema = test_schema(false); RecordBatch::try_new( schema, vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))], @@ -2897,7 +3273,7 @@ mod tests { /// Create batches with sequential values for ordering tests fn create_ordered_batches(num_batches: usize) -> Vec { - let schema = test_schema(); + let schema = test_schema(false); (0..num_batches) .map(|i| { let start = (i * 8) as u32; @@ -2918,7 +3294,7 @@ mod tests { // This tests the state machine fix where we must block on spill_stream // when a Spilled marker is received, rather than continuing to poll the channel - let schema = test_schema(); + let schema = test_schema(false); // Create batches with sequential values: batch 0 has [0,1,2,3,4,5,6,7], // batch 1 has [8,9,10,11,12,13,14,15], etc. let partition = create_ordered_batches(20); From 1d1ec1e162fc142b6e862b6b5318532636256872 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Fri, 10 Jul 2026 12:53:22 -0400 Subject: [PATCH 09/22] Support co-partitioned range inner equi joins (#23184) - Closes #23183. - Part of #22395. DataFusion can represent source-declared range partitioning, but partitioned hash joins still required hash partitioned inputs. So an inner join on compatible range-partitioned keys would insert unnecessary hash repartitions, even when each left/right partition already covered the same key domain. This PR adds a partitioning requirement that means "equal key values are co-located" . I was calling this "compatibility" but found we can satisfy the requirement with looser conditions. Other systems call this "co-location" or "co-partitioning" ([trino](https://trino.io/docs/current/admin/properties-optimizer.html#optimizer-colocated-joins-enabled), [spark](https://spark.apache.org/docs/latest/sql-performance-tuning.html#storage-partition-join)). Which they (and now I am proposing) define as when both sides of a join are already partitioned so matching key values appear in corresponding partitions, so we can join partition pairs directly without repartitioning the sides. This lets "co-partitioned" range inputs satisfy inner partitioned hash joins. This will also be applicable to other join types and operators but kept the first PR thin to keep scope more reviewable. - Adds `Distribution::KeyPartitioned(Vec>)` as a public distribution requirement. - `HashPartitioned([a])` means rows must be partitioned by hash on `a`. - `KeyPartitioned([a])` means rows with equal `a` values must be co-located, but the partitioning algorithm may be hash, range, or another compatible scheme. - Example: ```text Hash([left.a], 3) satisfies KeyPartitioned([left.a]) Range([right.b ASC], [(10), (20)], 3) satisfies KeyPartitioned([right.b]) ``` - Adds `Partitioning::co_partitioned_with(...)` to validate that two independently satisfying partitionings also can be paired by partition index. - Examples: - Accepted: both sides satisfy their own key requirement and have matching range boundaries. ```text left: Range([a ASC], [(10), (20)], 3), required KeyPartitioned([a]) right: Range([b ASC], [(10), (20)], 3), required KeyPartitioned([b]) ``` - Accepted: both sides satisfy their own key requirement and have matching hash partition counts. ```text left: Hash([a], 3), required KeyPartitioned([a]) right: Hash([b], 3), required KeyPartitioned([b]) ``` - Rejected: both sides satisfy their own key requirement, but range boundaries differ. ```text left: Range([a ASC], [(10), (20)], 3), required KeyPartitioned([a]) right: Range([b ASC], [(15), (20)], 3), required KeyPartitioned([b]) ``` - Rejected: both sides satisfy their own key requirement, but partition counts differ. ```text left: Hash([a], 3), required KeyPartitioned([a]) right: Hash([b], 4), required KeyPartitioned([b]) ``` - Changes inner partitioned `HashJoinExec` requirements from `HashPartitioned` to `KeyPartitioned`. - All other hash joins still require `HashPartitioned` for now. - Updates `EnforceDistribution` so co-partitioned range inner joins avoid repartitioning. - Examples: - Compatible range partitioning: no repartition is inserted because partitions can be joined by index. ```text HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a, b)] DataSourceExec: output_partitioning=Range([a ASC], [(10), (20)], 3) DataSourceExec: output_partitioning=Range([b ASC], [(10), (20)], 3) ``` - Incompatible range boundaries: both sides are repartitioned by hash because partition `i` does not represent the same key domain on both sides. ```text HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a, b)] RepartitionExec: partitioning=Hash([a], target_partitions) DataSourceExec: output_partitioning=Range([a ASC], [(10), (20)], 3) RepartitionExec: partitioning=Hash([b], target_partitions) DataSourceExec: output_partitioning=Range([b ASC], [(15), (20)], 3) ``` - Mismatched hash partition counts: both sides are forced to the target hash partition count so partition indexes line up. ```text HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a, b)] RepartitionExec: partitioning=Hash([a], target_partitions) DataSourceExec: output_partitioning=Hash([a], 11) RepartitionExec: partitioning=Hash([b], target_partitions) DataSourceExec: output_partitioning=Hash([b], 12) ``` - Non-inner joins: range inputs still get hash repartitioning because only inner partitioned hash joins use `KeyPartitioned` in this PR. ```text HashJoinExec: mode=Partitioned, join_type=Left, on=[(a, b)] RepartitionExec: partitioning=Hash([a], target_partitions) DataSourceExec: output_partitioning=Range([a ASC], [(10), (20)], 3) RepartitionExec: partitioning=Hash([b], target_partitions) DataSourceExec: output_partitioning=Range([b ASC], [(10), (20)], 3) ``` - Keeps partitioned dynamic filter pushdown restricted to hash-compatible routing. - Compatible range partitioning can satisfy the join, but dynamic filters still route by hash, so range/range partitioned joins disable dynamic filters. - Degrades range join output partitioning to `UnknownPartitioning(n)` rather than erroring. Adding this behavior would need more tests and careful thought about, I think its safert o just degrade for first PR. Yes. - `KeyPartitioned` satisfaction for hash and range partitioning. - `co_partitioned_with` for compatible and incompatible range/hash partitioning. - `EnforceDistribution` behavior for: - compatible range joins avoiding hash repartitioning - incompatible range bounds rehashing - mismatched hash partition counts rehashing - non-inner range joins rehashing - sanity checking for invalid partitioned hash joins. - dynamic filter rejection for range partitioning, preserved file partitions, and mismatched hash counts. - sqllogictest coverage for range-partitioned joins avoiding hash repartitioning and non-range joins still repartitioning. Yes. This PR changes public physical planning APIs: - Adds `Distribution::KeyPartitioned`. - Adds `Partitioning::co_partitioned_with`. - **NOTE**: This replaces the previous partition compatibility API with the new co-partitioning API. Since the compatibility API was never in a release I believe this is ok to do (lesson learned to not make API change until ew have definitive consumer). - Affects users matching exhaustively on `Distribution`. --- datafusion/catalog-listing/src/table.rs | 2 +- .../core/src/datasource/listing/table.rs | 1 - .../enforce_distribution.rs | 55 +- .../physical_optimizer/projection_pushdown.rs | 4 +- .../physical_optimizer/sanity_checker.rs | 19 +- .../tests/user_defined/user_defined_plan.rs | 8 +- datafusion/datasource/src/sink.rs | 14 +- datafusion/physical-expr/src/lib.rs | 4 +- datafusion/physical-expr/src/partitioning.rs | 232 +------ .../src/enforce_distribution.rs | 469 +++++++------- .../src/enforce_sorting/mod.rs | 22 +- .../src/output_requirements.rs | 33 +- .../physical-optimizer/src/sanity_checker.rs | 93 +-- datafusion/physical-optimizer/src/utils.rs | 19 - .../physical-plan/src/aggregates/mod.rs | 18 +- datafusion/physical-plan/src/analyze.rs | 8 +- .../src/distribution_requirements.rs | 510 ++++++++++++++++ .../physical-plan/src/execution_plan.rs | 33 +- .../physical-plan/src/joins/cross_join.rs | 8 +- .../physical-plan/src/joins/hash_join/exec.rs | 195 ++++-- .../src/joins/nested_loop_join.rs | 8 +- .../src/joins/piecewise_merge_join/exec.rs | 8 +- .../src/joins/sort_merge_join/exec.rs | 8 +- .../src/joins/symmetric_hash_join.rs | 8 +- datafusion/physical-plan/src/joins/utils.rs | 5 +- datafusion/physical-plan/src/lib.rs | 4 + datafusion/physical-plan/src/limit.rs | 6 +- .../physical-plan/src/recursive_query.rs | 8 +- .../physical-plan/src/sorts/partial_sort.rs | 8 +- .../src/sorts/partitioned_topk.rs | 8 +- datafusion/physical-plan/src/sorts/sort.rs | 8 +- .../src/sorts/sort_preserving_merge.rs | 8 +- datafusion/physical-plan/src/unnest.rs | 8 +- .../src/windows/bounded_window_agg_exec.rs | 8 +- .../src/windows/window_agg_exec.rs | 8 +- .../tests/cases/roundtrip_physical_plan.rs | 193 ------ .../src/test_context/range_partitioning.rs | 29 +- .../test_files/range_partitioning.slt | 574 +++++++++++++++++- 38 files changed, 1762 insertions(+), 892 deletions(-) create mode 100644 datafusion/physical-plan/src/distribution_requirements.rs diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 7f4d3e1965fea..3627e560a5096 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -802,7 +802,7 @@ impl ListingTable { .map(|part_file| async { let part_file = part_file?; let (statistics, ordering) = if self.options.collect_stat { - self.do_collect_statistics_and_ordering(ctx, &store, &part_file) + self.do_collect_statistics_and_ordering(ctx, store, &part_file) .await? } else { (Arc::new(Statistics::new_unknown(&self.file_schema)), None) diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 067cd380ba21f..0d6231ce27833 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -147,7 +147,6 @@ mod tests { use datafusion_physical_expr::expressions::{Column, binary}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::empty::EmptyExec; - use datafusion_physical_plan::statistics::StatisticsArgs; use datafusion_physical_plan::{ ExecutionPlanProperties, Partitioning, RangePartitioning, SplitPoint, collect, }; diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 1f53e79485d27..3a40a2cd2fc86 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -268,8 +268,12 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { vec![&self.input] } - fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + ]) } fn maintains_input_order(&self) -> Vec { @@ -826,7 +830,9 @@ fn inner_range_join_keeps_range_partitioning() -> Result<()> { JoinType::Inner, ); - let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + let plan = TestConfig::default() + .with_query_execution_partitions(3) + .to_plan(join, &DISTRIB_DISTRIB_SORT); assert_plan!( plan, @@ -1013,6 +1019,49 @@ fn range_grouping_set_aggregate_rehashes_with_grouping_id() -> Result<()> { Ok(()) } +#[test] +fn range_inner_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 30, 40], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + RepartitionExec: partitioning=Hash([a1@0], 4), input_partitions=4 + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4) + " + ); + + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 21f03b14f58bd..827a001b59894 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -800,7 +800,9 @@ fn test_output_req_after_projection() -> Result<()> { if let Distribution::KeyPartitioned(vec) = after_optimize .downcast_ref::() .unwrap() - .required_input_distribution()[0] + .input_distribution_requirements() + .child_distribution(0) + .unwrap() .clone() { assert!( diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index f12e5d5f764b0..e759156282306 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -413,35 +413,32 @@ fn range_partitioned_exec( [sort_expr(key, schema)].into(), split_points, )?); - let input = memory_exec(schema); - - RepartitionExec::try_new(input, partitioning) + RepartitionExec::try_new(memory_exec(schema), partitioning) .map(|exec| Arc::new(exec) as Arc) } #[test] fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> { let schema = create_test_schema2(); - let join_on = vec![(col("a", &schema)?, col("b", &schema)?)]; - let right = range_partitioned_exec(&schema, "b", [10])?; + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; - let valid_join = hash_join_exec( + let compatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, range_partitioned_exec(&schema, "a", [10])?, - Arc::clone(&right), join_on.clone(), None, &JoinType::Inner, )?; - assert_sanity_check(&valid_join, true); + assert_sanity_check(&compatible_join, true); - let invalid_join = hash_join_exec( + let incompatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, range_partitioned_exec(&schema, "a", [20])?, - right, join_on, None, &JoinType::Inner, )?; - assert_sanity_check(&invalid_join, false); + assert_sanity_check(&incompatible_join, false); Ok(()) } diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index e8ff6758ccdd4..b837373632f07 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -708,8 +708,12 @@ impl ExecutionPlan for TopKExec { &self.cache } - fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + ]) } fn children(&self) -> Vec<&Arc> { diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index e3df1ad6381f4..18ebe80773e8a 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -31,8 +31,9 @@ use datafusion_physical_expr_common::sort_expr::{LexRequirement, OrderingRequire use datafusion_physical_plan::metrics::MetricsSet; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - PlanProperties, SendableRecordBatchStream, execute_input_stream, + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + InputDistributionRequirements, Partitioning, PlanProperties, + SendableRecordBatchStream, execute_input_stream, }; use async_trait::async_trait; @@ -189,9 +190,16 @@ impl ExecutionPlan for DataSinkExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { // DataSink is responsible for dynamically partitioning its // own input at execution time, and so requires a single input partition. - vec![Distribution::SinglePartition; self.children().len()] + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition; + self.children().len() + ]) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index a0e9f8ee05363..ec7bf648e22ea 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -59,7 +59,9 @@ pub use datafusion_common::SplitPoint; pub use equivalence::{ AcrossPartitions, ConstExpr, EquivalenceProperties, calculate_union, }; -pub use partitioning::{Distribution, Partitioning, RangePartitioning}; +pub use partitioning::{ + Distribution, Partitioning, PartitioningSatisfaction, RangePartitioning, +}; pub use physical_expr::{ add_offset_to_expr, add_offset_to_physical_sort_exprs, create_lex_ordering, create_ordering, create_physical_partitioning, create_physical_sort_expr, diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index a777f9df14a80..75002a67e614e 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -244,35 +244,6 @@ impl RangePartitioning { self.split_points.len() + 1 } - /// Returns true when `self` and `other` have the same range boundaries. - /// - /// Single-partition range partitionings always have the same boundaries. Otherwise, - /// the two partitionings must have identical split points, ordering width, - /// and sort options. This does not compare ordering expressions, callers - /// should validate the range keys separately. - fn same_boundaries(&self, other: &Self) -> bool { - if self.partition_count() == 1 && other.partition_count() == 1 { - return true; - } - - if self.split_points != other.split_points - || self.ordering.len() != other.ordering.len() - { - return false; - } - - if !self - .ordering - .iter() - .zip(other.ordering.iter()) - .all(|(left, right)| left.options == right.options) - { - return false; - } - - true - } - /// Calculates the range partitioning after applying the given projection. /// /// Returns `None` if any range key cannot be projected or if projection @@ -391,61 +362,6 @@ impl Partitioning { } } - /// Returns true when two partitionings both satisfy their own distribution - /// requirements and can be paired by partition index. - /// - /// Use this for multi-input operators, such as partitioned joins, where - /// each child has a different schema, required [`Distribution`], and - /// expression-equivalence context. - /// - /// ```text - /// # co-partitioned: each side satisfies its own requirement, and boundaries match - /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) - /// right: Range(right.x ASC, [10, 20]), required KeyPartitioned(right.x) - /// - /// # not compatible: right side does not satisfy a hash-specific requirement - /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) - /// right: Range(right.x ASC, [10, 20]), required HashPartitioned(right.x) - /// - /// # not compatible: boundaries differ - /// left: Range(left.a ASC, [10, 20]), required KeyPartitioned(left.a) - /// right: Range(right.x ASC, [15, 20]), required KeyPartitioned(right.x) - /// ``` - pub fn co_partitioned_with( - &self, - required: &Distribution, - eq_properties: &EquivalenceProperties, - other: &Self, - other_required: &Distribution, - other_eq_properties: &EquivalenceProperties, - ) -> bool { - if !self - .satisfaction(required, eq_properties, false) - .is_satisfied() - || !other - .satisfaction(other_required, other_eq_properties, false) - .is_satisfied() - { - return false; - } - - if self.partition_count() == 1 && other.partition_count() == 1 { - return true; - } - - if self.partition_count() != other.partition_count() { - return false; - } - - match (self, other) { - (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, - (Partitioning::Range(left), Partitioning::Range(right)) => { - left.same_boundaries(right) - } - _ => false, - } - } - /// Returns true if `subset_exprs` is a subset of `exprs`. /// For example: Hash(a, b) is subset of Hash(a) since a partition with all occurrences of /// a distinct (a) must also contain all occurrences of a distinct (a, b) with the same (a). @@ -850,7 +766,7 @@ mod tests { let distribution_types = vec![ Distribution::UnspecifiedDistribution, Distribution::SinglePartition, - fixture.hash_distribution([0, 1]), + Distribution::HashPartitioned(fixture.cols([0, 1])), Distribution::KeyPartitioned(fixture.cols([0, 1])), ]; @@ -889,9 +805,6 @@ mod tests { Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) => { assert_eq!(result, (true, false, false, true, false)) } - Distribution::KeyPartitioned(_) => { - assert_eq!(result, (true, false, false, true, false)) - } } } @@ -1472,148 +1385,15 @@ mod tests { } #[test] - fn range_partitionings_are_co_partitioned_by_boundaries() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let left = fixture - .range_partitioning([0], vec![int_split_point([10]), int_split_point([20])]); - let right_same_map = fixture - .range_partitioning([1], vec![int_split_point([10]), int_split_point([20])]); - let right_different_split = fixture - .range_partitioning([1], vec![int_split_point([15]), int_split_point([20])]); - let right_desc = fixture.range_partitioning_with_ordering( - [fixture.range_sort_expr(1, SortOptions::new(true, false))].into(), - vec![int_split_point([20]), int_split_point([10])], - ); - - let test_cases = [ - ( - "same boundaries with matching key requirements", - fixture.key_partitioned_distribution([0]), - right_same_map.clone(), - fixture.key_partitioned_distribution([1]), - true, - ), - ( - "different split points", - fixture.key_partitioned_distribution([0]), - right_different_split, - fixture.key_partitioned_distribution([1]), - false, - ), - ( - "different sort options", - fixture.key_partitioned_distribution([0]), - right_desc, - fixture.key_partitioned_distribution([1]), - false, - ), - ( - "range cannot satisfy hash requirement", - fixture.hash_distribution([0]), - right_same_map, - fixture.key_partitioned_distribution([1]), - false, - ), - ]; - for (desc, left_requirement, right, right_requirement, expected) in test_cases { - assert_eq!( - left.co_partitioned_with( - &left_requirement, - &fixture.eq_properties, - &right, - &right_requirement, - &fixture.eq_properties, - ), - expected, - "Failed for {desc}" - ); - } - - Ok(()) - } - - #[test] - fn co_partitioned_with_rejects_subset_key_satisfaction() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let left = fixture - .range_partitioning([0], vec![int_split_point([10]), int_split_point([20])]); - let right = fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); - - assert_eq!( - right.satisfaction( - &fixture.key_partitioned_distribution([0]), - &fixture.eq_properties, - false, - ), - PartitioningSatisfaction::NotSatisfied - ); - assert_eq!( - left.satisfaction( - &fixture.key_partitioned_distribution([0, 1]), - &fixture.eq_properties, - true, - ), - PartitioningSatisfaction::Subset - ); - assert!(!left.co_partitioned_with( - &fixture.key_partitioned_distribution([0, 1]), - &fixture.eq_properties, - &right, - &fixture.key_partitioned_distribution([0]), - &fixture.eq_properties, - )); - - Ok(()) - } - - #[test] - fn hash_partitionings_are_co_partitioned_by_count() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let left = fixture.hash_partitioning([0], 2); - - let test_cases = [ - ( - "same partition count", - fixture.hash_partitioning([1], 2), - fixture.key_partitioned_distribution([1]), - true, - ), - ( - "different partition count", - fixture.hash_partitioning([1], 3), - fixture.key_partitioned_distribution([1]), - false, - ), - ( - "mixed hash and range partitioning", - fixture.range_partitioning([1], vec![int_split_point([10])]), - fixture.key_partitioned_distribution([1]), - false, - ), - ]; - for (desc, right, right_requirement, expected) in test_cases { - assert_eq!( - left.co_partitioned_with( - &fixture.key_partitioned_distribution([0]), - &fixture.eq_properties, - &right, - &right_requirement, - &fixture.eq_properties, - ), - expected, - "Failed for {desc}" - ); - } - - Ok(()) - } - - #[test] + #[expect( + deprecated, + reason = "test intentionally covers the hash-specific requirement" + )] 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.key_distribution([0, 1]); + let required = Distribution::HashPartitioned(fixture.cols([0, 1])); assert_eq!( range_partitioning.satisfaction(&required, &fixture.eq_properties, false), diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index b45f59bccf72e..d8f9fc880861a 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -28,9 +28,8 @@ use std::sync::Arc; use crate::optimizer::PhysicalOptimizerRule; use crate::output_requirements::OutputRequirementExec; use crate::utils::{ - add_sort_above_with_check, aggregate_can_reuse_range_partitioning, - is_coalesce_partitions, is_repartition, is_sort_preserving_merge, - range_partitioning_satisfies_key_partitioning, + add_sort_above_with_check, is_coalesce_partitions, is_repartition, + is_sort_preserving_merge, range_partitioning_satisfies_key_partitioning, }; use arrow::compute::SortOptions; @@ -42,7 +41,8 @@ use datafusion_expr::logical_plan::{Aggregate, JoinType}; use datafusion_physical_expr::expressions::{Column, NoOp}; use datafusion_physical_expr::utils::map_columns_before_projection; use datafusion_physical_expr::{ - EquivalenceProperties, PhysicalExpr, PhysicalExprRef, physical_exprs_equal, + EquivalenceProperties, OrderingRequirements, PhysicalExpr, PhysicalExprRef, + physical_exprs_equal, }; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::aggregates::{ @@ -60,7 +60,10 @@ use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::union::{InterleaveExec, UnionExec, can_interleave}; use datafusion_physical_plan::windows::WindowAggExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, get_best_fitting_window}; -use datafusion_physical_plan::{Distribution, ExecutionPlan, Partitioning}; +use datafusion_physical_plan::{ + ChildSatisfactionOptions, Distribution, ExecutionPlan, InputDistributionRequirements, + Partitioning, +}; use itertools::izip; @@ -854,82 +857,10 @@ fn add_roundrobin_on_top( } } -/// Adds a hash repartition operator: -/// - to increase parallelism, and/or -/// - to satisfy requirements of the subsequent operators. -/// -/// Repartition(Hash) is added on top of operator `input`. -/// -/// # Arguments -/// -/// * `input`: Current node. -/// * `hash_exprs`: Stores Physical Exprs that are used during hashing. -/// * `n_target`: desired target partition number, if partition number of the -/// current executor is less than this value. Partition number will be increased. -/// * `allow_subset_satisfy_partitioning`: Whether to allow subset partitioning logic in satisfaction checks. -/// Set to `false` for partitioned hash joins to ensure exact hash matching. -/// * `force_to_target`: Whether to repartition even when the hash expressions -/// are already satisfied but the partition count differs from `n_target`. -/// -/// # Returns -/// -/// A [`Result`] object that contains new execution plan where the desired -/// distribution is satisfied by adding a Hash repartition. -fn add_hash_on_top( - input: DistributionContext, - hash_exprs: Vec>, - n_target: usize, - allow_subset_satisfy_partitioning: bool, - force_to_target: bool, -) -> Result { - // Early return if hash repartition is unnecessary - // `RepartitionExec: partitioning=Hash([...], 1), input_partitions=1` is unnecessary. - if n_target == 1 && input.plan.output_partitioning().partition_count() == 1 { - return Ok(input); - } - - let dist = Distribution::KeyPartitioned(hash_exprs); - let current_partitions = input.plan.output_partitioning().partition_count(); - let satisfaction = input.plan.output_partitioning().satisfaction( - &dist, - input.plan.equivalence_properties(), - allow_subset_satisfy_partitioning, - ); - - // Add hash repartitioning when: - // - When subset satisfaction is enabled (current >= threshold): only repartition if not satisfied - // - When below threshold (current < threshold): repartition if expressions don't match OR to increase parallelism - let needs_repartition = if force_to_target { - !satisfaction.is_satisfied() || n_target != current_partitions - } else if allow_subset_satisfy_partitioning { - !satisfaction.is_satisfied() - } else { - !satisfaction.is_satisfied() || n_target > current_partitions - }; - - if needs_repartition { - // When there is an existing ordering, we preserve ordering during - // repartition. This will be rolled back in the future if any of the - // following conditions is true: - // - Preserving ordering is not helpful in terms of satisfying ordering - // requirements. - // - Usage of order preserving variants is not desirable (per the flag - // `config.optimizer.prefer_existing_sort`). - let partitioning = dist.create_partitioning(n_target); - let repartition = - RepartitionExec::try_new(Arc::clone(&input.plan), partitioning)? - .with_preserve_order(); - let plan = Arc::new(repartition) as _; - - return Ok(DistributionContext::new(plan, true, vec![input])); - } - - Ok(input) -} - -// TODO: remove this private helper once Range generally satisfies -// KeyPartitioned requirements through Partitioning::satisfaction. -// See . +// TODO: remove this temporary bridge once [`Partitioning::Range`] +// generally satisfies [`Distribution::KeyPartitioned`] through +// [`Partitioning::satisfaction`]. +// . // // Partial aggregates do not require key partitioning, but they preserve their // input partitioning for the final aggregate. Until Range satisfies @@ -1141,19 +1072,6 @@ struct RepartitionRequirementStatus { roundrobin_beneficial_stats: bool, /// Designates whether hash partitioning is necessary. hash_necessary: bool, - /// Designates whether hash repartitioning should force the target - /// partition count even when the hash expressions are already satisfied. - force_hash_to_target: bool, -} - -#[derive(Debug, Clone, Copy, Default)] -struct PartitionedJoinDistribution { - /// Inner partitioned hash join children can be paired by existing Range - /// partitions, so hash repartitioning is not needed. - compatible_range: bool, - /// Partitioned join children have different partition counts. If hash - /// repartitioning is used, both sides must be forced to the target count. - needs_count_alignment: bool, } fn requirement_includes_grouping_id(requirement: &Distribution) -> bool { @@ -1169,6 +1087,14 @@ fn requirement_includes_grouping_id(requirement: &Distribution) -> bool { }) } +/// Per-child state while enforcing a parent's distribution requirements. +struct DistributionChildState { + context: DistributionContext, + required_input_ordering: Option, + maintains_input_order: bool, + requirement: Distribution, +} + /// Calculates the `RepartitionRequirementStatus` for each children to generate /// consistent and sensible (in terms of performance) distribution requirements. /// As an example, a hash join's left (build) child might produce @@ -1208,8 +1134,7 @@ fn get_repartition_requirement_status( let mut needs_alignment = false; let children = plan.children(); let rr_beneficial = plan.benefits_from_input_partitioning(); - let requirements = plan.required_input_distribution(); - let join_distribution = partitioned_join_distribution(plan); + let requirements = plan.input_distribution_requirements().into_per_child(); let mut repartition_status_flags = vec![]; for (child, requirement, roundrobin_beneficial) in izip!(children.into_iter(), requirements, rr_beneficial) @@ -1223,27 +1148,19 @@ fn get_repartition_requirement_status( Precision::Absent => true, }; let is_partitioned_requirement = requirement.key_exprs().is_some(); - // Hash repartitioning may be necessary when the input has more than one - // partition, or when repartitioning one sibling requires aligning all - // key-partitioned siblings. + // Hash repartitioning is necessary when the input has more than one + // partition. let multi_partitions = child.output_partitioning().partition_count() > 1; let roundrobin_sensible = roundrobin_beneficial && roundrobin_beneficial_stats; - needs_alignment |= is_partitioned_requirement - && !join_distribution.compatible_range - && (multi_partitions || roundrobin_sensible); + needs_alignment |= + is_partitioned_requirement && (multi_partitions || roundrobin_sensible); repartition_status_flags.push(( is_partitioned_requirement, RepartitionRequirementStatus { requirement, roundrobin_beneficial, roundrobin_beneficial_stats, - hash_necessary: is_partitioned_requirement - && multi_partitions - && !join_distribution.compatible_range, - // Hash satisfaction checks key expressions, not matching - // partition counts, so force repartition when join sides differ. - force_hash_to_target: is_partitioned_requirement - && join_distribution.needs_count_alignment, + hash_necessary: is_partitioned_requirement && multi_partitions, }, )); } @@ -1265,56 +1182,87 @@ fn get_repartition_requirement_status( .collect()) } -/// Returns distribution state for a partitioned join's children. +/// Enforce cross-child distribution relationships after each child has already +/// satisfied its own distribution requirement. /// -/// This is optimizer policy: partitioned joins require children that can be -/// paired by partition index. Inner hash joins can reuse compatible range -/// partitioning; otherwise the existing hash repartitioning policy applies. -fn partitioned_join_distribution( - plan: &Arc, -) -> PartitionedJoinDistribution { - let Some(hash_join) = plan.downcast_ref::() else { - return Default::default(); - }; +/// See [`InputDistributionRequirements`] for the distinction between +/// independent per-child requirements and co-partitioned child relationships. +/// +/// Currently, unsatisfied co-partitioning is repaired by hash repartitioning +/// key-partitioned children and other relationship kinds are rejected. +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] +fn enforce_distribution_relationships( + plan_name: &str, + input_distributions: &InputDistributionRequirements, + children: &mut [DistributionChildState], + target_partitions: usize, +) -> Result<()> { + let mut repartitioned_for_relationship = vec![false; children.len()]; + + loop { + let child_plan_refs = children + .iter() + .map(|child| child.context.plan.as_ref()) + .collect::>(); + let unsatisfied_children = input_distributions + .unsatisfied_co_partitioned_children(plan_name, &child_plan_refs)?; - if hash_join.mode != PartitionMode::Partitioned { - return Default::default(); - } + if unsatisfied_children.is_empty() { + return Ok(()); + } - let children = plan.children(); - let [left, right] = children.as_slice() else { - return Default::default(); - }; - let needs_count_alignment = left.output_partitioning().partition_count() - != right.output_partitioning().partition_count(); - - let requirements = plan.required_input_distribution(); - let left_partitioning = left.output_partitioning(); - let right_partitioning = right.output_partitioning(); - let compatible_range = match requirements.as_slice() { - [ - left_requirement @ Distribution::KeyPartitioned(_), - right_requirement @ Distribution::KeyPartitioned(_), - ] if hash_join.join_type == JoinType::Inner - && matches!( - (left_partitioning, right_partitioning), - (Partitioning::Range(_), Partitioning::Range(_)) - ) => - { - left_partitioning.co_partitioned_with( - left_requirement, - left.equivalence_properties(), - right_partitioning, - right_requirement, - right.equivalence_properties(), - ) + let mut changed = false; + for child_idx in unsatisfied_children { + if repartitioned_for_relationship[child_idx] { + continue; + } + + let (Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs)) = &children[child_idx].requirement + else { + continue; + }; + + let already_target_hash = matches!( + children[child_idx].context.plan.output_partitioning(), + Partitioning::Hash(_, partition_count) if *partition_count == target_partitions + ) && input_distributions + .child_satisfaction( + child_idx, + children[child_idx].context.plan.as_ref(), + ChildSatisfactionOptions::new(), + )? + .is_satisfied(); + + if already_target_hash { + continue; + } + + let partitioning = Distribution::KeyPartitioned(exprs.to_vec()) + .create_partitioning(target_partitions); + let repartition = RepartitionExec::try_new( + Arc::clone(&children[child_idx].context.plan), + partitioning, + )? + .with_preserve_order(); + let plan = Arc::new(repartition) as _; + let original_child = std::mem::replace( + &mut children[child_idx].context, + DistributionContext::new(plan, true, vec![]), + ); + children[child_idx].context.children = vec![original_child]; + repartitioned_for_relationship[child_idx] = true; + changed = true; } - _ => false, - }; - PartitionedJoinDistribution { - compatible_range, - needs_count_alignment, + if !changed { + return datafusion_common::internal_err!( + "{plan_name} has distribution relationships that could not be enforced" + ); + } } } @@ -1420,6 +1368,7 @@ pub fn ensure_distribution( .is_some_and(|join| join.mode == PartitionMode::Partitioned) || plan.is::(); + let input_distributions = plan.input_distribution_requirements(); let repartition_status_flags = get_repartition_requirement_status(&plan, batch_size, should_use_estimates)?; // This loop iterates over all the children to: @@ -1427,7 +1376,8 @@ pub fn ensure_distribution( // - Satisfy the distribution requirements of every child, if it is not // already satisfied. // We store the updated children in `new_children`. - let children = izip!( + let mut children = izip!( + 0..children.len(), children.into_iter(), plan.required_input_ordering(), plan.maintains_input_order(), @@ -1435,6 +1385,7 @@ pub fn ensure_distribution( ) .map( |( + child_idx, mut child, required_input_ordering, maintains, @@ -1443,13 +1394,16 @@ pub fn ensure_distribution( roundrobin_beneficial, roundrobin_beneficial_stats, hash_necessary, - force_hash_to_target, }, )| { // Allow subset satisfaction when: // 1. Current partition count >= threshold // 2. Not a partitioned join since must use exact hash matching for joins // 3. Not a grouping set aggregate (requires exact hash including __grouping_id) + // + // Partitioned joins still require exact satisfaction. If that + // exact check already passes, preserve_file_partitions can skip + // repartitioning whose only purpose is increasing partition count. let current_partitions = child.plan.output_partitioning().partition_count(); let preserve_file_partition_threshold_met = config.optimizer.preserve_file_partitions > 0 @@ -1505,25 +1459,47 @@ pub fn ensure_distribution( } Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs) => { - let range_satisfied_for_aggregate = - aggregate_can_reuse_range_partitioning(&plan) - && range_partitioning_satisfies_key_partitioning( - child.plan.output_partitioning(), - exprs, - child.plan.equivalence_properties(), - allow_subset_satisfy_partitioning, - ); + let child_partitions = + child.plan.output_partitioning().partition_count(); + let partitioning_satisfied = input_distributions + .child_satisfaction( + child_idx, + child.plan.as_ref(), + ChildSatisfactionOptions::new() + .with_allow_subset(allow_subset_satisfy_partitioning), + )? + .is_satisfied(); + let preserve_satisfying_file_partitioning = + preserve_file_partition_threshold_met + && !requirement_includes_grouping_id(&requirement) + && partitioning_satisfied + && target_partitions > child_partitions; + + // When subset satisfaction is enabled, preserve an + // already-satisfying partitioning. Otherwise, hash + // repartition may also increase parallelism. + let needs_hash_repartition = if allow_subset_satisfy_partitioning { + !partitioning_satisfied + } else { + !partitioning_satisfied + || (target_partitions > child_partitions + && !preserve_satisfying_file_partitioning) + }; + let should_add_hash_repartition = + hash_necessary && needs_hash_repartition; // See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background // When inserting hash is necessary to satisfy hash requirement, insert hash repartition. - if hash_necessary && !range_satisfied_for_aggregate { - child = add_hash_on_top( - child, - exprs.to_vec(), - target_partitions, - allow_subset_satisfy_partitioning, - force_hash_to_target, - )?; + if should_add_hash_repartition { + let partitioning = Distribution::KeyPartitioned(exprs.to_vec()) + .create_partitioning(target_partitions); + let repartition = RepartitionExec::try_new( + Arc::clone(&child.plan), + partitioning, + )? + .with_preserve_order(); + let plan = Arc::new(repartition) as _; + child = DistributionContext::new(plan, true, vec![child]); } } Distribution::UnspecifiedDistribution => { @@ -1535,75 +1511,101 @@ pub fn ensure_distribution( } }; - let streaming_benefit = if child.data { - preserving_order_enables_streaming(&plan, &child.plan)? - } else { - false - }; + Ok(DistributionChildState { + context: child, + required_input_ordering, + maintains_input_order: maintains, + requirement, + }) + }, + ) + .collect::>>()?; - // There is an ordering requirement of the operator: - if let Some(required_input_ordering) = required_input_ordering { - // Either: - // - Ordering requirement cannot be satisfied by preserving ordering through repartitions, or - // - using order preserving variant is not desirable. - let sort_req = required_input_ordering.into_single(); - let ordering_satisfied = child - .plan - .equivalence_properties() - .ordering_satisfy_requirement(sort_req.clone())?; - - if (!ordering_satisfied || !order_preserving_variants_desirable) - && !streaming_benefit - && child.data - { - child = replace_order_preserving_variants(child)?; - // If ordering requirements were satisfied before repartitioning, - // make sure ordering requirements are still satisfied after. - if ordering_satisfied { - // Make sure to satisfy ordering requirement: - child = add_sort_above_with_check( - child, - sort_req, - plan.downcast_ref::() - .map(|output| output.fetch()) - .unwrap_or(None), - )?; - } - } - // Stop tracking distribution changing operators - child.data = false; - } else { - let streaming_benefit = if child.data { - preserving_order_enables_streaming(&plan, &child.plan)? + // This is called after each child satisfies its own distribution requirement. + // It enforces relationships between child partition layouts for multi-child + // operators that process matching partition indexes together. + enforce_distribution_relationships( + plan.name(), + &input_distributions, + &mut children, + target_partitions, + )?; + + let children = children + .into_iter() + .map( + |DistributionChildState { + mut context, + required_input_ordering, + maintains_input_order, + requirement, + }| { + let streaming_benefit = if context.data { + preserving_order_enables_streaming(&plan, &context.plan)? } else { false }; - // no ordering requirement - match requirement { - // Operator requires specific distribution. - Distribution::SinglePartition - | Distribution::HashPartitioned(_) - | Distribution::KeyPartitioned(_) => { - // If the parent doesn't maintain input order, preserving - // ordering is pointless. However, if it does maintain - // input order, we keep order-preserving variants so - // ordering can flow through to ancestors that need it. - if !maintains && !streaming_benefit { - child = replace_order_preserving_variants(child)?; + + // There is an ordering requirement of the operator: + if let Some(required_input_ordering) = required_input_ordering { + // Either: + // - Ordering requirement cannot be satisfied by preserving ordering through repartitions, or + // - using order preserving variant is not desirable. + let sort_req = required_input_ordering.into_single(); + let ordering_satisfied = context + .plan + .equivalence_properties() + .ordering_satisfy_requirement(sort_req.clone())?; + + if (!ordering_satisfied || !order_preserving_variants_desirable) + && !streaming_benefit + && context.data + { + context = replace_order_preserving_variants(context)?; + // If ordering requirements were satisfied before repartitioning, + // make sure ordering requirements are still satisfied after. + if ordering_satisfied { + // Make sure to satisfy ordering requirement: + context = add_sort_above_with_check( + context, + sort_req, + plan.downcast_ref::() + .map(|output| output.fetch()) + .unwrap_or(None), + )?; } } - Distribution::UnspecifiedDistribution => { - // Since ordering is lost, trying to preserve ordering is pointless - if !maintains || plan.is::() { - child = replace_order_preserving_variants(child)?; + // Stop tracking distribution changing operators + context.data = false; + } else { + // no ordering requirement + match requirement { + // Operator requires specific distribution. + Distribution::SinglePartition + | Distribution::HashPartitioned(_) + | Distribution::KeyPartitioned(_) => { + // If the parent doesn't maintain input order, preserving + // ordering is pointless. However, if it does maintain + // input order, we keep order-preserving variants so + // ordering can flow through to ancestors that need it. + if !maintains_input_order && !streaming_benefit { + context = replace_order_preserving_variants(context)?; + } + } + Distribution::UnspecifiedDistribution => { + // Since ordering is lost, trying to preserve ordering is pointless + if !maintains_input_order + || plan.is::() + { + context = replace_order_preserving_variants(context)?; + } } } } - } - Ok(child) - }, - ) - .collect::>>()?; + Ok(context) + }, + ) + .collect::>>()?; let children_plans = children .iter() @@ -1669,7 +1671,8 @@ fn update_children(mut dist_context: DistributionContext) -> Result Result { if node.data { let requires_single_partition = matches!( - parent.required_input_distribution()[child_idx], - Distribution::SinglePartition + parent + .input_distribution_requirements() + .child_distribution(child_idx), + Some(Distribution::SinglePartition) ); node = remove_corresponding_sort_from_sub_plan(node, requires_single_partition)?; } @@ -749,7 +757,7 @@ fn remove_corresponding_sort_from_sub_plan( } } else { let mut any_connection = false; - let required_dist = node.plan.required_input_distribution(); + let required_dist = node.plan.input_distribution_requirements().into_per_child(); node.children = node .children .into_iter() diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 50c1ea8483d3e..48cc41b8a21c4 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -28,7 +28,7 @@ use crate::PhysicalOptimizerRule; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{Result, Statistics}; +use datafusion_common::{Result, Statistics, internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::Distribution; use datafusion_physical_expr_common::sort_expr::OrderingRequirements; @@ -205,7 +205,15 @@ impl ExecutionPlan for OutputRequirementExec { } fn required_input_distribution(&self) -> Vec { - vec![self.dist_requirement.clone()] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + self.dist_requirement.clone(), + ]) } fn maintains_input_order(&self) -> Vec { @@ -272,9 +280,12 @@ impl ExecutionPlan for OutputRequirementExec { requirements = OrderingRequirements::new_alternatives(updated_reqs, soft); } - let dist_req = match &self.required_input_distribution()[0] { - Distribution::HashPartitioned(exprs) - | Distribution::KeyPartitioned(exprs) => { + let input_distributions = self.input_distribution_requirements(); + let dist_req = match input_distributions.child_distribution(0) { + Some( + Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs), + ) => { let mut updated_exprs = vec![]; for expr in exprs { let Some(new_expr) = update_expr(expr, projection.expr(), false)? @@ -285,7 +296,12 @@ impl ExecutionPlan for OutputRequirementExec { } Distribution::KeyPartitioned(updated_exprs) } - dist => dist.clone(), + Some(dist) => dist.clone(), + None => { + return internal_err!( + "OutputRequirementExec missing input distribution requirement" + ); + } }; make_with_child(projection, &self.input()).map(|input| { @@ -360,7 +376,10 @@ fn require_top_ordering_helper( // In case of constant columns, output ordering of the `SortExec` would // be an empty set. Therefore; we check the sort expression field to // assign the requirements. - let req_dist = sort_exec.required_input_distribution().swap_remove(0); + let req_dist = sort_exec + .input_distribution_requirements() + .into_per_child() + .swap_remove(0); let req_ordering = sort_exec.expr(); let reqs = OrderingRequirements::from(req_ordering.clone()); let fetch = sort_exec.fetch(); diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index 4e74062268260..713213b70612d 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -24,23 +24,21 @@ use std::sync::Arc; use datafusion_common::Result; -use datafusion_physical_expr::Distribution; use datafusion_physical_plan::ExecutionPlan; use datafusion_common::config::{ConfigOptions, OptimizerOptions}; use datafusion_common::plan_err; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported}; -use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion_physical_plan::joins::{ - HashJoinExec, PartitionMode, SymmetricHashJoinExec, +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, InvariantLevel, +}; +use datafusion_physical_plan::joins::SymmetricHashJoinExec; +use datafusion_physical_plan::{ + ChildSatisfactionOptions, ExecutionPlanProperties, get_plan_string, }; -use datafusion_physical_plan::{ExecutionPlanProperties, get_plan_string}; use crate::PhysicalOptimizerRule; -use crate::utils::{ - aggregate_can_reuse_range_partitioning, range_partitioning_satisfies_key_partitioning, -}; use datafusion_physical_expr_common::sort_expr::format_physical_sort_requirement_list; use itertools::izip; @@ -142,20 +140,17 @@ fn is_prunable(join: &SymmetricHashJoinExec) -> bool { /// Ensures that the plan is pipeline friendly and the order and /// distribution requirements from its children are satisfied. -#[expect( - deprecated, - reason = "HashPartitioned is accepted during the KeyPartitioned migration" -)] pub fn check_plan_sanity( plan: &Arc, optimizer_options: &OptimizerOptions, ) -> Result<()> { check_finiteness_requirements(plan.as_ref(), optimizer_options)?; + let input_distributions = plan.input_distribution_requirements(); for ((idx, child), sort_req, dist_req) in izip!( plan.children().into_iter().enumerate(), plan.required_input_ordering(), - plan.required_input_distribution(), + input_distributions.per_child_distributions(), ) { let child_eq_props = child.equivalence_properties(); if let Some(sort_req) = sort_req { @@ -172,26 +167,14 @@ pub fn check_plan_sanity( } } - let child_satisfies_distribution = child - .output_partitioning() - .satisfaction(&dist_req, child_eq_props, true) - .is_satisfied(); - let range_satisfies_aggregate_distribution = - aggregate_can_reuse_range_partitioning(plan) - && match &dist_req { - Distribution::HashPartitioned(exprs) - | Distribution::KeyPartitioned(exprs) => { - range_partitioning_satisfies_key_partitioning( - child.output_partitioning(), - exprs, - child_eq_props, - true, - ) - } - _ => false, - }; - - if !(child_satisfies_distribution || range_satisfies_aggregate_distribution) { + if !input_distributions + .child_satisfaction( + idx, + child.as_ref(), + ChildSatisfactionOptions::new().with_allow_subset(true), + )? + .is_satisfied() + { let plan_str = get_plan_string(plan); return plan_err!( "Plan: {:?} does not satisfy distribution requirements: {}. Child-{} output partitioning: {}", @@ -203,49 +186,7 @@ pub fn check_plan_sanity( } } - check_partitioned_join_distribution(plan)?; - - Ok(()) -} - -fn check_partitioned_join_distribution(plan: &Arc) -> Result<()> { - let Some(hash_join) = plan.downcast_ref::() else { - return Ok(()); - }; - - if hash_join.mode != PartitionMode::Partitioned { - return Ok(()); - } - - let children = plan.children(); - let requirements = plan.required_input_distribution(); - let ([left, right], [left_req, right_req]) = - (children.as_slice(), requirements.as_slice()) - else { - return plan_err!( - "Invalid HashJoinExec: expected two children and two distribution requirements" - ); - }; - - if !left.output_partitioning().co_partitioned_with( - left_req, - left.equivalence_properties(), - right.output_partitioning(), - right_req, - right.equivalence_properties(), - ) { - let plan_str = get_plan_string(plan); - return plan_err!( - "Plan: {:?} does not satisfy partitioned join co-partitioning requirements: \ - left requirement: {}, left output partitioning: {}; \ - right requirement: {}, right output partitioning: {}", - plan_str, - left_req, - left.output_partitioning(), - right_req, - right.output_partitioning() - ); - } + plan.check_invariants(InvariantLevel::Executable)?; Ok(()) } diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index 2f928224da28b..36e630c1f4ae3 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -22,7 +22,6 @@ use datafusion_physical_expr::{ EquivalenceProperties, LexOrdering, LexRequirement, Partitioning, PhysicalExpr, physical_exprs_equal, }; -use datafusion_physical_plan::aggregates::{AggregateExec, AggregateMode}; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -166,24 +165,6 @@ pub(crate) fn range_partitioning_satisfies_key_partitioning( } } -/// TODO: remove once Range generally satisfies KeyPartitioned requirements -/// through Partitioning::satisfaction. -/// See . -/// -/// Checks whether an aggregate can reuse range partitioning to satisfy its key -/// partitioning requirement. -pub(crate) fn aggregate_can_reuse_range_partitioning( - plan: &Arc, -) -> bool { - plan.downcast_ref::() - .is_some_and(|aggregate| { - matches!( - aggregate.mode(), - AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned - ) && !aggregate.group_expr().has_grouping_set() - }) -} - /// Checks whether the given operator is a limit; /// i.e. either a [`LocalLimitExec`] or a [`GlobalLimitExec`]. pub fn is_limit(plan: &Arc) -> bool { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 4f2c2e58bf791..725dc2cbd64fd 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -32,8 +32,8 @@ use crate::filter_pushdown::{ }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, InputOrderMode, - SendableRecordBatchStream, Statistics, check_if_same_properties, + DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + InputOrderMode, SendableRecordBatchStream, Statistics, check_if_same_properties, }; use datafusion_common::config::ConfigOptions; use datafusion_physical_expr::utils::collect_columns; @@ -1544,7 +1544,11 @@ impl ExecutionPlan for AggregateExec { } fn required_input_distribution(&self) -> Vec { - match &self.mode { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + let requirements = InputDistributionRequirements::new(match &self.mode { AggregateMode::Partial | AggregateMode::PartialReduce => { vec![Distribution::UnspecifiedDistribution] } @@ -1554,6 +1558,14 @@ impl ExecutionPlan for AggregateExec { AggregateMode::Final | AggregateMode::Single => { vec![Distribution::SinglePartition] } + }); + match &self.mode { + AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned + if !self.group_by.has_grouping_set() => + { + requirements.allow_range_satisfaction_for_key_partitioning() + } + _ => requirements, } } diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 582af8f1e3dae..708f6b4c95ac3 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -144,7 +144,13 @@ impl ExecutionPlan for AnalyzeExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn with_new_children( diff --git a/datafusion/physical-plan/src/distribution_requirements.rs b/datafusion/physical-plan/src/distribution_requirements.rs new file mode 100644 index 0000000000000..9c7a1336c06a3 --- /dev/null +++ b/datafusion/physical-plan/src/distribution_requirements.rs @@ -0,0 +1,510 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Input distribution requirements for physical execution plans. + +use std::sync::Arc; + +use datafusion_common::{Result, internal_err}; +use datafusion_physical_expr::{ + Distribution, EquivalenceProperties, Partitioning, PartitioningSatisfaction, + PhysicalExpr, physical_exprs_equal, +}; + +use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel}; + +/// Distribution requirements for an [`ExecutionPlan`]'s inputs. +/// +/// [`InputDistributionRequirements`] describes what distribution an operator +/// requires from each child. +/// +/// - [`Self::new`] describes independent per-child requirements. +/// - [`Self::co_partitioned`] additionally requires child partitions with the +/// same index to cover compatible key ranges. +/// +/// For a single-input aggregate: +/// +/// ```text +/// AggregateExec +/// child 0 requirement: KeyPartitioned(group_exprs) +/// ``` +/// +/// each input partition can aggregate its own key domain independently. +/// +/// For a partitioned join: +/// +/// ```text +/// HashJoinExec +/// child 0 requirement: KeyPartitioned(left_keys) +/// child 1 requirement: KeyPartitioned(right_keys) +/// +/// partition 0: join(left partition 0, right partition 0) +/// partition 1: join(left partition 1, right partition 1) +/// partition 2: join(left partition 2, right partition 2) +/// ``` +/// +/// each child must satisfy its own key requirement. In addition, matching +/// partition indexes must be safe to process together. +#[non_exhaustive] +#[derive(Debug, Clone)] +pub struct InputDistributionRequirements { + /// Per-child distribution requirements, indexed by child position. + children: Vec, + /// Child indexes that must also have compatible partition layouts. + co_partitioned: Option>, +} + +/// Options for checking child distribution satisfaction. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ChildSatisfactionOptions { + allow_subset: bool, +} + +impl ChildSatisfactionOptions { + /// Create default satisfaction options. + pub fn new() -> Self { + Self::default() + } + + /// Allow a child partitioning whose key expressions are a subset of the + /// required key expressions to satisfy the requirement. + pub fn with_allow_subset(mut self, allow_subset: bool) -> Self { + self.allow_subset = allow_subset; + self + } + + /// Whether subset satisfaction is enabled. + pub fn allow_subset(&self) -> bool { + self.allow_subset + } +} + +impl InputDistributionRequirements { + /// Create independent per-child requirements. + pub fn new(per_child: Vec) -> Self { + let children = per_child + .into_iter() + .map(|distribution| ChildDistributionRequirement { + distribution, + satisfaction: InputDistributionSatisfaction::Default, + }) + .collect(); + + Self { + children, + co_partitioned: None, + } + } + + /// Create a requirement that all children are co-partitioned. + /// + /// Each child must satisfy its own [`Distribution`]. Matching partition + /// indexes are processed together: + /// + /// ```text + /// left: Range(left.a ASC, split_points=[10, 20]) + /// right: Range(right.x ASC, split_points=[10, 20]) + /// + /// partition 0 from both sides contains keys before 10 + /// partition 1 from both sides contains keys in [10, 20) + /// partition 2 from both sides contains keys at/after 20 + /// ``` + /// + /// If the split points differ, partition `i` from one side no longer covers + /// the same key range as partition `i` from the other side. + pub fn co_partitioned(per_child: Vec) -> Self { + debug_assert!( + per_child.len() >= 2, + "co-partitioned distribution requirements need at least two children" + ); + let co_partitioned = (0..per_child.len()).collect(); + let mut result = Self::new(per_child); + result.co_partitioned = Some(co_partitioned); + result + } + + /// Return the per-child distribution requirements. + pub fn per_child_distributions( + &self, + ) -> impl ExactSizeIterator + '_ { + self.children.iter().map(|child| &child.distribution) + } + + /// Return the distribution requirement for a child. + pub fn child_distribution(&self, child_idx: usize) -> Option<&Distribution> { + self.children + .get(child_idx) + .map(|child| &child.distribution) + } + + /// Return the per-child distribution requirements. + /// + /// WARNING: This intentionally drops any grouped relationship. + pub fn into_per_child(self) -> Vec { + self.children + .into_iter() + .map(|child| child.distribution) + .collect() + } + + /// Returns how a child satisfies its distribution requirement. + /// + /// This preserves the requirement set's satisfaction policy. + pub fn child_satisfaction( + &self, + child_idx: usize, + child: &dyn ExecutionPlan, + options: ChildSatisfactionOptions, + ) -> Result { + let Some(requirement) = self.children.get(child_idx) else { + return internal_err!( + "missing distribution requirement for child {child_idx}" + ); + }; + + Ok(requirement.satisfaction.satisfaction( + child.output_partitioning(), + &requirement.distribution, + child.equivalence_properties(), + options.allow_subset(), + )) + } + + /// Return child indexes whose co-partitioning requirements are + /// unsatisfied by the provided candidate children. + /// + /// Independent per-child requirements are intentionally ignored here, use + /// [`Self::child_satisfaction`] for those checks. An empty result means all + /// co-partitioning requirements are satisfied. + #[doc(hidden)] + pub fn unsatisfied_co_partitioned_children( + &self, + plan_name: &str, + children: &[&dyn ExecutionPlan], + ) -> Result> { + self.validate_shape(plan_name, children.len())?; + + let Some(co_partitioned) = &self.co_partitioned else { + return Ok(vec![]); + }; + if self.co_partitioning_satisfied(co_partitioned, children) { + return Ok(vec![]); + } + + Ok(co_partitioned.clone()) + } + + /// TODO: remove this temporary bridge once [`Partitioning::Range`] + /// generally satisfies [`Distribution::KeyPartitioned`] through + /// [`Partitioning::satisfaction`]. + /// . + /// + /// Also allow compatible [`Partitioning::Range`] to satisfy + /// [`Distribution::KeyPartitioned`]. + #[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" + )] + pub(crate) fn allow_range_satisfaction_for_key_partitioning(mut self) -> Self { + for child in &mut self.children { + if matches!( + child.distribution, + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) + ) { + child.satisfaction = + InputDistributionSatisfaction::AllowRangeKeyPartitioning; + } + } + self + } + + /// Validate the requirements against a plan's children. + pub(crate) fn check_invariants( + &self, + plan: &P, + check: InvariantLevel, + ) -> Result<()> { + let children = plan.children(); + self.validate_shape(plan.name(), children.len())?; + + let children = children + .into_iter() + .map(|child| child.as_ref()) + .collect::>(); + if matches!(check, InvariantLevel::Executable) + && let Some(co_partitioned) = &self.co_partitioned + && !self.co_partitioning_satisfied(co_partitioned, &children) + { + return internal_err!( + "{} requires children {:?} to be co-partitioned", + plan.name(), + co_partitioned + ); + } + + Ok(()) + } + + fn validate_shape(&self, plan_name: &str, children_len: usize) -> Result<()> { + if self.children.len() != children_len { + return internal_err!( + "{plan_name}::input_distribution_requirements returned incorrect child count: {} != {}", + self.children.len(), + children_len + ); + } + + if let Some(co_partitioned) = &self.co_partitioned { + if co_partitioned.len() < 2 { + return internal_err!( + "{plan_name} has invalid co-partitioning requirement: at least two children are required" + ); + } + let mut seen = vec![false; self.children.len()]; + for &child in co_partitioned { + validate_child_index(plan_name, child, self.children.len(), &mut seen)?; + if matches!( + self.children[child].distribution, + Distribution::UnspecifiedDistribution + ) { + return internal_err!( + "{plan_name} has invalid co-partitioning requirement: child {child} has unspecified distribution" + ); + } + } + } + + Ok(()) + } + + fn co_partitioning_satisfied( + &self, + co_partitioned: &[usize], + children: &[&dyn ExecutionPlan], + ) -> bool { + let first_idx = co_partitioned[0]; + let first_requirement = &self.children[first_idx]; + let first = children[first_idx]; + let first_partitioning = first.output_partitioning(); + + if !first_requirement + .satisfaction + .satisfaction( + first_partitioning, + &first_requirement.distribution, + first.equivalence_properties(), + false, + ) + .is_satisfied() + { + return false; + } + + for &child_idx in co_partitioned.iter().skip(1) { + let requirement = &self.children[child_idx]; + let child = children[child_idx]; + if !requirement + .satisfaction + .satisfaction( + child.output_partitioning(), + &requirement.distribution, + child.equivalence_properties(), + false, + ) + .is_satisfied() + || !compatible_co_partitioning_layout( + first_requirement, + first_partitioning, + requirement, + child.output_partitioning(), + ) + { + return false; + } + } + + true + } +} + +/// A distribution requirement for a single child. +#[derive(Debug, Clone)] +struct ChildDistributionRequirement { + distribution: Distribution, + satisfaction: InputDistributionSatisfaction, +} + +/// TODO: remove this temporary bridge once [`Partitioning::Range`] +/// generally satisfies [`Distribution::KeyPartitioned`] through +/// [`Partitioning::satisfaction`]. +/// . +#[non_exhaustive] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +enum InputDistributionSatisfaction { + /// Use [`Partitioning::satisfaction`] as-is. + #[default] + Default, + /// Also allow [`Partitioning::Range`] to satisfy + /// [`Distribution::KeyPartitioned`]. + AllowRangeKeyPartitioning, +} + +impl InputDistributionSatisfaction { + /// Returns how `partitioning` satisfies `requirement`. + #[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" + )] + fn satisfaction( + self, + partitioning: &Partitioning, + requirement: &Distribution, + eq_properties: &EquivalenceProperties, + allow_subset: bool, + ) -> PartitioningSatisfaction { + let satisfaction = + partitioning.satisfaction(requirement, eq_properties, allow_subset); + if satisfaction.is_satisfied() { + return satisfaction; + } + + if !matches!(self, Self::AllowRangeKeyPartitioning) { + return PartitioningSatisfaction::NotSatisfied; + } + + let (Distribution::HashPartitioned(required_exprs) + | Distribution::KeyPartitioned(required_exprs)) = requirement + else { + return PartitioningSatisfaction::NotSatisfied; + }; + + range_satisfies_key_partitioning( + partitioning, + required_exprs, + eq_properties, + allow_subset, + ) + } +} + +fn validate_child_index( + plan_name: &str, + child_idx: usize, + child_count: usize, + seen: &mut [bool], +) -> Result<()> { + if child_idx >= child_count { + return internal_err!( + "{plan_name} has invalid distribution requirement: child index {child_idx} out of bounds" + ); + } + if seen[child_idx] { + return internal_err!( + "{plan_name} has invalid distribution requirement: child {child_idx} appears more than once" + ); + } + seen[child_idx] = true; + Ok(()) +} + +/// TODO: remove this temporary bridge once [`Partitioning::Range`] +/// generally satisfies [`Distribution::KeyPartitioned`] through +/// [`Partitioning::satisfaction`]. +/// . +fn range_satisfies_key_partitioning( + partitioning: &Partitioning, + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, +) -> PartitioningSatisfaction { + let Partitioning::Range(range) = partitioning else { + return PartitioningSatisfaction::NotSatisfied; + }; + + let partition_exprs = range + .ordering() + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + + if partition_exprs.is_empty() || required_exprs.is_empty() { + return PartitioningSatisfaction::NotSatisfied; + } + + let eq_group = eq_properties.eq_group(); + let normalized_partition_exprs = partition_exprs + .iter() + .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) + .collect::>(); + let normalized_required_exprs = required_exprs + .iter() + .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) + .collect::>(); + + if physical_exprs_equal(&normalized_required_exprs, &normalized_partition_exprs) { + return PartitioningSatisfaction::Exact; + } + + if allow_subset + && normalized_partition_exprs.len() < normalized_required_exprs.len() + && normalized_partition_exprs.iter().all(|partition_expr| { + normalized_required_exprs + .iter() + .any(|required_expr| partition_expr.eq(required_expr)) + }) + { + PartitioningSatisfaction::Subset + } else { + PartitioningSatisfaction::NotSatisfied + } +} + +fn compatible_co_partitioning_layout( + first: &ChildDistributionRequirement, + first_partitioning: &Partitioning, + other: &ChildDistributionRequirement, + other_partitioning: &Partitioning, +) -> bool { + if first_partitioning.partition_count() == 1 + && other_partitioning.partition_count() == 1 + { + return true; + } + + if first_partitioning.partition_count() != other_partitioning.partition_count() { + return false; + } + + match (first_partitioning, other_partitioning) { + (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, + (Partitioning::Range(left), Partitioning::Range(right)) + if first.satisfaction + == InputDistributionSatisfaction::AllowRangeKeyPartitioning + && other.satisfaction + == InputDistributionSatisfaction::AllowRangeKeyPartitioning => + { + left.split_points() == right.split_points() + && left.ordering().len() == right.ordering().len() + && left + .ordering() + .iter() + .zip(right.ordering()) + .all(|(left, right)| left.options == right.options) + } + _ => false, + } +} diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 9fc3725fc6457..5c20f85656e90 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -16,6 +16,7 @@ // under the License. pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; +use crate::distribution_requirements::InputDistributionRequirements; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -161,17 +162,30 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { check_default_invariants(self, check) } - /// Specifies the data distribution requirements for all the children for - /// this `ExecutionPlan`. + /// Specifies simple per-child input distribution requirements. + /// + /// Deprecated: override [`Self::input_distribution_requirements`] instead. /// /// By default, each child has [`Distribution::UnspecifiedDistribution`]. - /// Multi-input operators that use [`Distribution::KeyPartitioned`] must - /// use [`Partitioning::co_partitioned_with`] to verify that satisfied - /// children can be paired by partition index. + #[deprecated(since = "55.0.0", note = "Use input_distribution_requirements")] fn required_input_distribution(&self) -> Vec { vec![Distribution::UnspecifiedDistribution; self.children().len()] } + /// Specifies the input distribution requirements for this plan. + /// + /// The default implementation wraps [`Self::required_input_distribution`]. + /// Override this method for richer requirements, such as allowing alternate + /// satisfaction policies or requiring multiple children to be co-partitioned. + /// See [`InputDistributionRequirements`] for details. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + #[expect( + deprecated, + reason = "compatibility shim for external ExecutionPlan implementations" + )] + InputDistributionRequirements::new(self.required_input_distribution()) + } + /// Specifies the ordering required for all of the children of this /// `ExecutionPlan`. /// @@ -218,8 +232,8 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { fn benefits_from_input_partitioning(&self) -> Vec { // By default try to maximize parallelism with more CPUs if // possible - self.required_input_distribution() - .into_iter() + self.input_distribution_requirements() + .per_child_distributions() .map(|dist| !matches!(dist, Distribution::SinglePartition)) .collect() } @@ -1202,14 +1216,15 @@ macro_rules! check_len { /// Returns an error if the given node does not conform. pub fn check_default_invariants( plan: &P, - _check: InvariantLevel, + check: InvariantLevel, ) -> Result<(), DataFusionError> { let children_len = plan.children().len(); check_len!(plan, maintains_input_order, children_len); check_len!(plan, required_input_ordering, children_len); - check_len!(plan, required_input_distribution, children_len); check_len!(plan, benefits_from_input_partitioning, children_len); + plan.input_distribution_requirements() + .check_invariants(plan, check)?; Ok(()) } diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 6661d2782b212..fde1718001252 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -306,10 +306,14 @@ impl ExecutionPlan for CrossJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn execute( diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index f746ebcb6b9b1..1e859a045345b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -24,8 +24,8 @@ use std::vec; use crate::ExecutionPlanProperties; use crate::execution_plan::{ - EmissionType, InvariantLevel, boundedness_from_children, check_default_invariants, - has_same_children_properties, stub_properties, + EmissionType, boundedness_from_children, has_same_children_properties, + stub_properties, }; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, @@ -54,8 +54,9 @@ use crate::projection::{ use crate::repartition::REPARTITION_RANDOM_STATE; use crate::spill::get_record_batch_memory_size; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, - PlanProperties, SendableRecordBatchStream, Statistics, + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + InputDistributionRequirements, Partitioning, PlanProperties, + SendableRecordBatchStream, Statistics, common::can_project, joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, @@ -876,6 +877,15 @@ impl HashJoinExec { } } + if self.mode == PartitionMode::Partitioned + && !self.has_partitioned_dynamic_filter_routing() + { + // TODO: support partition-routed dynamic filters for compatible + // range co-partitioned joins. + // . + return false; + } + true } @@ -895,21 +905,6 @@ impl HashJoinExec { } } - fn partitioned_children_co_partitioned(&self) -> bool { - let requirements = self.required_input_distribution(); - let [left_requirement, right_requirement] = requirements.as_slice() else { - return false; - }; - - self.left.output_partitioning().co_partitioned_with( - left_requirement, - self.left.equivalence_properties(), - self.right.output_partitioning(), - right_requirement, - self.right.equivalence_properties(), - ) - } - /// left (build) side which gets hashed pub fn left(&self) -> &Arc { &self.left @@ -1248,46 +1243,41 @@ impl ExecutionPlan for HashJoinExec { "HashJoinExec" } - fn check_invariants(&self, check: InvariantLevel) -> Result<()> { - check_default_invariants(self, check)?; - - if matches!(check, InvariantLevel::Executable) - && self.mode == PartitionMode::Partitioned - && !self.partitioned_children_co_partitioned() - { - return plan_err!( - "Invalid HashJoinExec, partitioned children are not co-partitioned, consider using RepartitionExec" - ); - } - - Ok(()) - } - fn properties(&self) -> &Arc { &self.cache } fn required_input_distribution(&self) -> Vec { - match self.mode { - PartitionMode::CollectLeft => vec![ - Distribution::SinglePartition, - Distribution::UnspecifiedDistribution, - ], + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + let requirements = match self.mode { PartitionMode::Partitioned => { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - vec![ + InputDistributionRequirements::co_partitioned(vec![ Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), - ] + ]) } - PartitionMode::Auto => vec![ + PartitionMode::CollectLeft => InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, Distribution::UnspecifiedDistribution, + ]), + PartitionMode::Auto => InputDistributionRequirements::new(vec![ Distribution::UnspecifiedDistribution, - ], + Distribution::UnspecifiedDistribution, + ]), + }; + + if self.mode == PartitionMode::Partitioned && self.join_type == JoinType::Inner { + requirements.allow_range_satisfaction_for_key_partitioning() + } else { + requirements } } @@ -1351,12 +1341,6 @@ impl ExecutionPlan for HashJoinExec { consider using RepartitionExec" ); - assert_or_internal_err!( - self.mode != PartitionMode::Partitioned - || self.partitioned_children_co_partitioned(), - "Invalid HashJoinExec, partitioned children are not co-partitioned, consider using RepartitionExec" - ); - assert_or_internal_err!( self.mode != PartitionMode::CollectLeft || left_partitions == 1, "Invalid HashJoinExec, the output partition count of the left child must be 1 in CollectLeft mode,\ @@ -2273,6 +2257,7 @@ mod tests { } use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::execution_plan::Boundedness; use crate::joins::hash_join::stream::lookup_join_hashmap; use crate::test::{TestMemoryExec, assert_join_metrics}; use crate::{ @@ -2295,12 +2280,67 @@ mod tests { use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; - use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint}; + use datafusion_physical_expr::{ + EquivalenceProperties, PhysicalSortExpr, RangePartitioning, SplitPoint, + }; use hashbrown::HashTable; use insta::{allow_duplicates, assert_snapshot}; use rstest::*; use rstest_reuse::*; + #[derive(Debug)] + struct PartitionedTestExec { + cache: Arc, + } + + impl PartitionedTestExec { + fn try_new(schema: SchemaRef, partitioning: Partitioning) -> Result { + Ok(Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )), + }) + } + } + + impl DisplayAs for PartitionedTestExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PartitionedTestExec") + } + } + + impl ExecutionPlan for PartitionedTestExec { + fn name(&self) -> &'static str { + "PartitionedTestExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } + } + fn div_ceil(a: usize, b: usize) -> usize { a.div_ceil(b) } @@ -6594,8 +6634,9 @@ mod tests { let session_config = join_dynamic_filter_session_config(0); let join = partitioned_inner_hash_join(left, right, on)?; + let requirements = join.input_distribution_requirements().into_per_child(); assert!(matches!( - join.required_input_distribution().as_slice(), + requirements.as_slice(), [ Distribution::KeyPartitioned(_), Distribution::KeyPartitioned(_) @@ -6634,6 +6675,58 @@ mod tests { Ok(()) } + #[test] + fn test_partitioned_dynamic_filter_pushdown_rejects_range_partitioning() -> Result<()> + { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left_partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr { + expr: Arc::clone(&on[0].0), + options: Default::default(), + }] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?); + let right_partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr { + expr: Arc::clone(&on[0].1), + options: Default::default(), + }] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?); + let left = Arc::new(PartitionedTestExec::try_new( + left_schema, + left_partitioning, + )?); + let right = Arc::new(PartitionedTestExec::try_new( + right_schema, + right_partitioning, + )?); + + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + )?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + #[test] fn test_with_dynamic_filter_rejects_invalid_columns() -> Result<()> { let (_, _, on) = build_schema_and_on()?; diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 15af23b447836..20d05983cf4ed 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -565,10 +565,14 @@ impl ExecutionPlan for NestedLoopJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 50e9252a21131..a2a1d6c787bb7 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -508,10 +508,14 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 1fe6f3bda922c..ea6db4f77369a 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -423,15 +423,19 @@ impl ExecutionPlan for SortMergeJoinExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - vec![ + crate::InputDistributionRequirements::new(vec![ Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), - ] + ]) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 41caa8cba2d9a..24337cbe6a45d 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -426,7 +426,11 @@ impl ExecutionPlan for SymmetricHashJoinExec { } fn required_input_distribution(&self) -> Vec { - match self.mode { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(match self.mode { StreamJoinPartitionMode::Partitioned => { let (left_expr, right_expr) = self .on @@ -441,7 +445,7 @@ impl ExecutionPlan for SymmetricHashJoinExec { StreamJoinPartitionMode::SinglePartition => { vec![Distribution::SinglePartition, Distribution::SinglePartition] } - } + }) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 3fb090a406286..941a57d72bad4 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -144,10 +144,9 @@ pub fn adjust_right_output_partitioning( .collect::>()?; Partitioning::Hash(new_exprs, *size) } - // Range partitioning can satisfy join input requirements, but range - // output propagation needs broader join semantics coverage. - // https://github.com/apache/datafusion/issues/22395 Partitioning::Range(range) => { + // Range partitioning optimizer propagation is tracked in + // https://github.com/apache/datafusion/issues/22395 Partitioning::UnknownPartitioning(range.partition_count()) } result => result.clone(), diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index c7b1d4729e21d..8190663b2b23d 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -41,6 +41,9 @@ pub use datafusion_physical_expr::{ }; pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; +pub use crate::distribution_requirements::{ + ChildSatisfactionOptions, InputDistributionRequirements, +}; pub use crate::execution_plan::{ ExecutionPlan, ExecutionPlanProperties, PlanProperties, collect, collect_partitioned, displayable, execute_input_stream, execute_stream, execute_stream_partitioned, @@ -71,6 +74,7 @@ pub mod column_rewriter; pub mod common; pub mod coop; pub mod display; +pub mod distribution_requirements; pub mod empty; pub mod execution_plan; pub mod explain; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index 7f42c33a79ca0..94ef00d8567aa 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -162,7 +162,11 @@ impl ExecutionPlan for GlobalLimitExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index f34aac3744557..99767df73a734 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -163,10 +163,14 @@ impl ExecutionPlan for RecursiveQueryExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ crate::Distribution::SinglePartition, crate::Distribution::SinglePartition, - ] + ]) } fn with_new_children( diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 3bf16af36c62b..3eeefa3acd7c7 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -268,11 +268,15 @@ impl ExecutionPlan for PartialSortExec { } fn required_input_distribution(&self) -> Vec { - if self.preserve_partitioning { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.preserve_partitioning { vec![Distribution::UnspecifiedDistribution] } else { vec![Distribution::SinglePartition] - } + }) } fn benefits_from_input_partitioning(&self) -> Vec { diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 5596783caf2e4..17eb70ef12131 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -302,12 +302,18 @@ impl ExecutionPlan for PartitionedTopKExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { let partition_exprs: Vec> = self.expr [..self.partition_prefix_len] .iter() .map(|e| Arc::clone(&e.expr)) .collect(); - vec![Distribution::KeyPartitioned(partition_exprs)] + crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + partition_exprs, + )]) } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 929ff4f7dfc85..551f7e0ffa6dd 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1137,14 +1137,18 @@ impl ExecutionPlan for SortExec { } fn required_input_distribution(&self) -> Vec { - if self.preserve_partitioning { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.preserve_partitioning { vec![Distribution::UnspecifiedDistribution] } else { // global sort // TODO support range partitioning and OrderedDistribution. // See https://github.com/apache/datafusion/issues/22395 vec![Distribution::SinglePartition] - } + }) } fn children(&self) -> Vec<&Arc> { diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index eb9b5f09aa3ed..053c57a5a6a62 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -266,7 +266,13 @@ impl ExecutionPlan for SortPreservingMergeExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn benefits_from_input_partitioning(&self) -> Vec { diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index c31d0dd23fa68..4f345a6c2dc54 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -253,7 +253,13 @@ impl ExecutionPlan for UnnestExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn execute( diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 4b31ce9b9245d..2e6523d087c0a 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -331,12 +331,16 @@ impl ExecutionPlan for BoundedWindowAggExec { } fn required_input_distribution(&self) -> Vec { - if self.partition_keys().is_empty() { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() { debug!("No partition defined for BoundedWindowAggExec!!!"); vec![Distribution::SinglePartition] } else { vec![Distribution::KeyPartitioned(self.partition_keys().clone())] - } + }) } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 0c77adc40dd92..3b03c794051c1 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -240,11 +240,15 @@ impl ExecutionPlan for WindowAggExec { } fn required_input_distribution(&self) -> Vec { - if self.partition_keys().is_empty() { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() { vec![Distribution::SinglePartition] } else { vec![Distribution::KeyPartitioned(self.partition_keys())] - } + }) } fn with_new_children( diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 1cf55a9afee40..674bf5a2d60ed 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -4082,199 +4082,6 @@ fn test_sort_topk_with_dynamic_filter_roundtrip() -> Result<()> { Ok(()) } -/// A custom [`ExecutionPlan`] which stores [`PhysicalExpr`]s. -struct CustomExecWithExprs { - exprs: Vec>, - child: Arc, -} - -#[derive(Clone, PartialEq, Message)] -struct CustomExecWithExprsProto { - #[prost(message, repeated, tag = "1")] - exprs: Vec, -} - -impl std::fmt::Debug for CustomExecWithExprs { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_struct("CustomExecWithExprs") - .field("exprs", &self.exprs) - .field("child", &self.child) - .finish() - } -} - -impl CustomExecWithExprs { - fn new(exprs: Vec>, child: Arc) -> Self { - Self { exprs, child } - } -} - -impl DisplayAs for CustomExecWithExprs { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { - write!(f, "CustomExecWithExprs") - } -} - -impl ExecutionPlan for CustomExecWithExprs { - fn name(&self) -> &str { - "CustomExecWithExprs" - } - - fn schema(&self) -> SchemaRef { - self.child.schema() - } - - fn properties(&self) -> &Arc { - self.child.properties() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.child] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - unreachable!() - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - unreachable!() - } -} - -/// A [`PhysicalExtensionCodec`] for [`CustomExecWithExprs`]. -#[derive(Debug)] -struct CustomExecWithExprsCodec {} - -impl PhysicalExtensionCodec for CustomExecWithExprsCodec { - fn try_decode( - &self, - buf: &[u8], - inputs: &[Arc], - ctx: &TaskContext, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self); - let input_schema = inputs[0].schema(); - let proto = CustomExecWithExprsProto::decode(buf) - .map_err(|e| internal_datafusion_err!("Failed to decode custom exec: {e}"))?; - let exprs = proto - .exprs - .iter() - .map(|expr_proto| { - proto_converter.proto_to_physical_expr( - expr_proto, - input_schema.as_ref(), - &decode_ctx, - ) - }) - .collect::>>()?; - - Ok(Arc::new(CustomExecWithExprs::new(exprs, inputs[0].clone()))) - } - - fn try_encode( - &self, - node: Arc, - buf: &mut Vec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result<()> { - let custom = node - .downcast_ref::() - .ok_or_else(|| internal_datafusion_err!("Expected CustomExecWithExprs"))?; - let proto = CustomExecWithExprsProto { - exprs: custom - .exprs - .iter() - .map(|expr| proto_converter.physical_expr_to_proto(expr, self)) - .collect::>>()?, - }; - proto - .encode(buf) - .map_err(|e| internal_datafusion_err!("Failed to encode custom exec: {e}"))?; - - Ok(()) - } -} - -/// Tests that a custom [`ExecutionPlan`] with [`PhysicalExpr`] can -/// dedupe dynamic filters by using the proto converter in its -/// [`PhysicalExtensionCodec`] implementation. -#[test] -fn test_custom_node_with_dynamic_filter_dedup_roundtrip() -> Result<()> { - // Create the plan: - // - // FilterExec(dynamic_filter) - // -> CustomExecWithExprs(exprs: [dynamic_filter]) - // -> EmptyExec - // - // The same dynamic filter expression is saved in both the FilterExec and CustomExecWithExprs. - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( - vec![Arc::new(Column::new("a", 0)) as Arc], - lit(true), - )); - let dynamic_filter_expr: Arc = dynamic_filter; - - let empty = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let custom_exec = Arc::new(CustomExecWithExprs::new( - vec![Arc::clone(&dynamic_filter_expr)], - empty, - )); - let filter_exec = Arc::new(FilterExec::try_new( - Arc::clone(&dynamic_filter_expr), - custom_exec, - )?) as Arc; - - // Roundtrip with DeduplicatingProtoConverter - let codec = CustomExecWithExprsCodec {}; - let converter = DeduplicatingProtoConverter {}; - - let bytes = physical_plan_to_bytes_with_proto_converter( - Arc::clone(&filter_exec), - &codec, - &converter, - )?; - - let ctx = SessionContext::new(); - let deser_converter = DeduplicatingProtoConverter {}; - let deserialized = physical_plan_from_bytes_with_proto_converter( - bytes.as_ref(), - ctx.task_ctx().as_ref(), - &codec, - &deser_converter, - )?; - - // Extract the deserialized FilterExec's dynamic filter - let deser_filter = deserialized - .downcast_ref::() - .expect("Top-level should be FilterExec"); - let deser_filter_df = deser_filter.predicate(); - - // Extract the deserialized custom node's dynamic filter - let deser_custom = deser_filter - .input() - .downcast_ref::() - .expect("FilterExec child should be CustomExecWithExprs"); - assert_eq!(deser_custom.exprs.len(), 1, "Should have one expression"); - let [deser_custom_df] = deser_custom.exprs.as_slice() else { - return internal_err!("Custom node should have one expression"); - }; - - // Pass the un-remapped filter first so the helper's `with_new_children` - // rewrite can reconstruct the remapped form on the other side. - assert_dynamic_filters_equal(deser_custom_df, deser_filter_df); - assert_dynamic_filter_update_is_visible(deser_custom_df, deser_filter_df)?; - - Ok(()) -} - fn roundtrip_file_scan_config(scan_config: FileScanConfig) -> Result { let exec_plan: Arc = DataSourceExec::from_data_source(scan_config); let ctx = SessionContext::new(); diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index a3e16eefd881a..aa741dded77be 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -57,7 +57,7 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { "range_partitioned", Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_partitioned"), - schema, + Arc::clone(&schema), [ "1,1,10\n5,2,50\n", "10,1,100\n15,2,150\n", @@ -66,6 +66,33 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { ], Some(output_partitioning), ); + + let shifted_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(15))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_shifted", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_shifted"), + schema, + [ + "1,1,10\n5,2,50\n10,1,100\n", + "15,2,150\n", + "20,1,200\n25,2,250\n", + "30,1,300\n35,2,350\n", + ], + Some(shifted_output_partitioning), + ); } fn register_csv_listing_table( diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 2b7a2cfdf4083..0dfe399a4646d 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -37,10 +37,8 @@ query TT EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; ---- physical_plan -01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] -02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 -03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] -04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query II SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key ORDER BY range_key; @@ -69,7 +67,7 @@ physical_plan 01)AggregateExec: mode=FinalPartitioned, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] 02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 03)----AggregateExec: mode=Partial, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] -04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false +04)------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false query II SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key ORDER BY non_range_key; @@ -79,22 +77,566 @@ SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key O ########## -# TEST 3: Join on Range Partition Column -# Both inputs expose Range partitioning on range_key. Join planning currently -# reaches the unsupported Range output-partitioning path; later optimizer PRs -# can replace this baseline with a successful plan and result test. +# TEST 3: Aggregate Reuses Range Subset Partitioning +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies grouping by (range_key, non_range_key). +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key ORDER BY range_key, non_range_key; +---- +1 1 10 +5 2 50 +10 1 100 +15 2 150 +20 1 200 +25 2 250 +30 1 300 +35 2 350 + + +########## +# TEST 4: Exact Range Aggregate Below Subset Threshold +# Even when subset satisfaction is disabled, exact Range([range_key]) +# satisfies GROUP BY range_key when repartitioning would not increase +# partition count. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + + +########## +# TEST 5: Range Subset Aggregate Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of GROUP BY (range_key, non_range_key), +# so it should not satisfy the aggregate key when subset satisfaction is +# disabled. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + + +########## +# TEST 6: Aggregate Rehashes Below Subset Threshold +# With subset threshold 5 and only 4 input partitions, planning repartitions +# to increase parallelism instead of reusing Range partitioning. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + + +########## +# TEST 7: Aggregate Preserves Range When Preserve File Threshold Met +# With preserve-file threshold 1 and 4 input partitions, Range is preserved +# even though target_partitions is 5. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 8: Aggregate Rehashes When Preserve File Threshold Not Met +# With preserve-file threshold 5 and only 4 input partitions, planning can +# repartition to increase parallelism. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 5; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + + +########## +# TEST 9: Join on Range Partition Column +# A partitioned inner hash join requires co-partitioned KeyPartitioned inputs. +# Compatible Range layouts satisfy both the per-child key requirements and the +# cross-child layout requirement, so no Hash repartitioning is inserted. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 10: Incompatible Range Join Repartitions +# Both inputs are independently range partitioned on range_key, but their split +# points differ. The per-child key requirements can be satisfied by Range, but +# the co-partitioned layout requirement cannot, so Hash repartitioning repairs +# both sides. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 11: Non-Range Join Repartitions +# Range([range_key]) does not satisfy KeyPartitioned([non_range_key]), so +# planning inserts Hash repartitioning on the actual join key. +########## + +query TT +EXPLAIN SELECT l.non_range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(non_range_key@0, non_range_key@0)], projection=[non_range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT l.non_range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.non_range_key = r.non_range_key +ORDER BY l.non_range_key, l.value, r.value; +---- +1 10 10 +1 10 100 +1 10 200 +1 10 300 +1 100 10 +1 100 100 +1 100 200 +1 100 300 +1 200 10 +1 200 100 +1 200 200 +1 200 300 +1 300 10 +1 300 100 +1 300 200 +1 300 300 +2 50 50 +2 50 150 +2 50 250 +2 50 350 +2 150 50 +2 150 150 +2 150 250 +2 150 350 +2 250 50 +2 250 150 +2 250 250 +2 250 350 +2 350 50 +2 350 150 +2 350 250 +2 350 350 + +########## +# TEST 12: Non-Inner Range Join Repartitions +# Only inner partitioned hash joins opt in to Range satisfying KeyPartitioned +# requirements. Non-inner joins keep using Hash repartitioning. ########## -query error This feature is not implemented: Join output partitioning with range partitioning is not implemented +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III SELECT l.range_key, l.value, r.value FROM range_partitioned l +LEFT JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 13: Compatible Range Join Repartitions to Increase Parallelism +# Co-partitioning satisfaction does not prevent a repartition that increases +# parallelism. With target_partitions larger than the Range partition count, +# both sides are hash repartitioned. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 14: Preserve File Partitions Preserves Range Join Inputs +# preserve_file_partitions preserves compatible Range inputs for partitioned +# joins even when target_partitions is higher than the input partition count. +########## + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +########## +# TEST 15: Nested Range Joins +# Compatible Range partitioning satisfies the lower join inputs. The upper join +# still repairs the intermediate join output with Hash repartitioning because +# HashJoinExec does not currently expose Range output partitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value, s.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +JOIN range_partitioned s ON r.range_key = s.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@2, range_key@0)], projection=[range_key@0, value@1, value@3, value@5] +02)--RepartitionExec: partitioning=Hash([range_key@2], 4), input_partitions=4 +03)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +06)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +07)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT l.range_key, l.value, r.value, s.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +JOIN range_partitioned s ON r.range_key = s.range_key +ORDER BY l.range_key; +---- +1 10 10 10 +5 50 50 50 +10 100 100 100 +15 150 150 150 +20 200 200 200 +25 250 250 250 +30 300 300 300 +35 350 350 350 + +########## +# TEST 16: Range Aggregates Feed Range Join +# Aggregates on range_key preserve reusable partitioning for the downstream +# partitioned join. +########## + +query TT +EXPLAIN WITH + l AS ( + SELECT range_key, SUM(value) AS l_sum + FROM range_partitioned + GROUP BY range_key + ), + r AS ( + SELECT range_key, SUM(value) AS r_sum + FROM range_partitioned + GROUP BY range_key + ) +SELECT l.range_key, l.l_sum, r.r_sum +FROM l JOIN r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, l_sum@1, r_sum@3] +02)--ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value)@1 as l_sum] +03)----AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value)@1 as r_sum] +06)----AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +07)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +WITH + l AS ( + SELECT range_key, SUM(value) AS l_sum + FROM range_partitioned + GROUP BY range_key + ), + r AS ( + SELECT range_key, SUM(value) AS r_sum + FROM range_partitioned + GROUP BY range_key + ) +SELECT l.range_key, l.l_sum, r.r_sum +FROM l JOIN r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 17: Range Join Feeds Aggregate +# The join inputs avoid Hash repartitioning, but the aggregate above the join +# still repartitions because HashJoinExec does not currently expose Range +# output partitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, SUM(l.value + r.value) +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +GROUP BY l.range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)] +04)------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +06)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT l.range_key, SUM(l.value + r.value) +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +GROUP BY l.range_key +ORDER BY l.range_key; +---- +1 20 +5 100 +10 200 +15 300 +20 400 +25 500 +30 600 +35 700 + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +statement ok +reset datafusion.optimizer.repartition_joins; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 4: Union of Range Partitioned Inputs -# Each input exposes Range partitioning on range_key. This records current -# UNION ALL behavior before later PRs decide whether compatible range inputs can -# preserve Range partitioning across the union. +# TEST 18: Union of Range Partitioned Inputs +# Each input exposes Range partitioning on range_key. These changes do not add a +# cross-child Range relationship for UNION ALL. ########## query TT @@ -104,8 +646,8 @@ SELECT range_key, value FROM range_partitioned; ---- physical_plan 01)UnionExec -02)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false -03)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query II SELECT range_key, value FROM range_partitioned From 1248881780a5218df24f5892fe6e8182fd5e74d7 Mon Sep 17 00:00:00 2001 From: Mithun Chicklore Yogendra Date: Wed, 15 Jul 2026 00:11:38 +0530 Subject: [PATCH 10/22] feat: allow Partitioning::Range to satisfy window Distribution::KeyPartitioned requirements (#23416) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #23289. Rebased on main now that #23184 has merged. Overlaps with #23355, which also opts the window execs into range satisfaction as part of its PartitionedTopK work — happy to rebase whichever lands second. ## What changes are included in this PR? In one commit: * Opt `WindowAggExec` and `BoundedWindowAggExec` with partition keys into range satisfaction via `InputDistributionRequirements::allow_range_satisfaction_for_key_partitioning`, mirroring `AggregateExec`. Compatible range-partitioned inputs then satisfy the window key requirement without a hash repartition; subset satisfaction and the hash fallback for incompatible keys come from the existing satisfaction machinery. Windows without `PARTITION BY` keep requiring a single partition. ## Are these changes tested? Yes: new `slt` tests in `range_partitioning.slt` (exact and subset reuse, rehash on incompatible keys, `subset_repartition_threshold` / `preserve_file_partitions` / `target_partitions` behavior, `WindowAggExec` via an unbounded frame, and no-`PARTITION BY`), plus plan-shape tests in `enforce_distribution.rs` and acceptance/rejection tests in `sanity_checker.rs`. ## Are there any user-facing changes? No: only plan changes. --- .../enforce_distribution.rs | 67 +++++- .../physical_optimizer/sanity_checker.rs | 77 +++++- .../tests/physical_optimizer/test_utils.rs | 18 +- .../src/windows/bounded_window_agg_exec.rs | 18 +- .../src/windows/window_agg_exec.rs | 18 +- .../test_files/range_partitioning.slt | 225 ++++++++++++++++++ 6 files changed, 403 insertions(+), 20 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 3a40a2cd2fc86..a59b7e95e1550 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -20,8 +20,8 @@ use std::ops::Deref; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - check_integrity, coalesce_partitions_exec, parquet_exec_with_sort, - parquet_exec_with_stats, repartition_exec, schema, sort_exec, + bounded_window_exec_with_can_repartition, check_integrity, coalesce_partitions_exec, + parquet_exec_with_sort, parquet_exec_with_stats, repartition_exec, schema, sort_exec, sort_exec_with_preserve_partitioning, sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; @@ -1062,6 +1062,69 @@ fn range_inner_hash_join_rehashes_incompatible_range_partitioning() -> Result<() Ok(()) } +#[test] +fn range_window_reuses_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let window = bounded_window_exec_with_can_repartition( + "a", + vec![], + &[col("a", &schema())?], + input, + true, + ); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(window, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + "# + ); + + Ok(()) +} + +#[test] +fn range_window_rehashes_incompatible_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let window = bounded_window_exec_with_can_repartition( + "b", + vec![], + &[col("b", &schema())?], + input, + true, + ); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(window, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([b@1], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + "# + ); + + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index e759156282306..e5718f5b3d0f7 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -19,9 +19,10 @@ use insta::assert_snapshot; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - bounded_window_exec, global_limit_exec, hash_join_exec, local_limit_exec, - memory_exec, projection_exec, repartition_exec, sort_exec, sort_expr, - sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec, + bounded_window_exec, bounded_window_exec_with_can_repartition, global_limit_exec, + hash_join_exec, local_limit_exec, memory_exec, projection_exec, repartition_exec, + sort_exec, sort_exec_with_preserve_partitioning, sort_expr, sort_expr_options, + sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; use arrow::compute::SortOptions; @@ -501,6 +502,76 @@ async fn test_bounded_window_agg_no_sort_requirement() -> Result<()> { Ok(()) } +#[tokio::test] +/// Tests that a window over a compatible range-partitioned input satisfies +/// the window's key distribution requirement without a hash repartition. +async fn test_bounded_window_agg_range_partitioning() -> Result<()> { + let schema = create_test_schema2(); + let source = range_partitioned_exec(&schema, "a", [10, 20, 30])?; + let ordering: LexOrdering = [sort_expr_options( + "a", + &schema, + SortOptions { + descending: false, + nulls_first: false, + }, + )] + .into(); + let partition_by = vec![col("a", &schema)?]; + let sort = sort_exec_with_preserve_partitioning(ordering, source); + let bw = + bounded_window_exec_with_can_repartition("a", vec![], &partition_by, sort, true); + let plan_str = displayable(bw.as_ref()).indent(true).to_string(); + let actual = plan_str.trim(); + assert_snapshot!( + actual, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + "# + ); + assert_sanity_check(&bw, true); + Ok(()) +} + +#[tokio::test] +/// Tests that a window over an incompatible range-partitioned input fails +/// the window's key distribution requirement. +async fn test_bounded_window_agg_incompatible_range_partitioning() -> Result<()> { + let schema = create_test_schema2(); + let source = range_partitioned_exec(&schema, "a", [10, 20, 30])?; + let ordering: LexOrdering = [sort_expr_options( + "b", + &schema, + SortOptions { + descending: false, + nulls_first: false, + }, + )] + .into(); + let partition_by = vec![col("b", &schema)?]; + let sort = sort_exec_with_preserve_partitioning(ordering, source); + let bw = + bounded_window_exec_with_can_repartition("b", vec![], &partition_by, sort, true); + let plan_str = displayable(bw.as_ref()).indent(true).to_string(); + let actual = plan_str.trim(); + assert_snapshot!( + actual, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + "# + ); + // Range([a]) does not colocate `b` values, so the window's key + // distribution requirement is not satisfied. + assert_sanity_check(&bw, false); + Ok(()) +} + #[tokio::test] /// A valid when a single partition requirement /// is satisfied. diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 09225cb0385a7..1973813215ed6 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -263,6 +263,22 @@ pub fn bounded_window_exec_with_partition( sort_exprs: impl IntoIterator, partition_by: &[Arc], input: Arc, +) -> Arc { + bounded_window_exec_with_can_repartition( + col_name, + sort_exprs, + partition_by, + input, + false, + ) +} + +pub fn bounded_window_exec_with_can_repartition( + col_name: &str, + sort_exprs: impl IntoIterator, + partition_by: &[Arc], + input: Arc, + can_repartition: bool, ) -> Arc { let sort_exprs = sort_exprs.into_iter().collect::>(); let schema = input.schema(); @@ -285,7 +301,7 @@ pub fn bounded_window_exec_with_partition( vec![window_expr], Arc::clone(&input), InputOrderMode::Sorted, - false, + can_repartition, ) .unwrap(), ) diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 2e6523d087c0a..9e38902c520a3 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -35,8 +35,9 @@ use crate::windows::{ }; use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputOrderMode, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, + ExecutionPlanProperties, InputDistributionRequirements, InputOrderMode, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, + check_if_same_properties, }; use arrow::compute::take_record_batch; @@ -334,13 +335,16 @@ impl ExecutionPlan for BoundedWindowAggExec { self.input_distribution_requirements().into_per_child() } - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() { + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + if self.partition_keys().is_empty() { debug!("No partition defined for BoundedWindowAggExec!!!"); - vec![Distribution::SinglePartition] + InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } else { - vec![Distribution::KeyPartitioned(self.partition_keys().clone())] - }) + InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + self.partition_keys(), + )]) + .allow_range_satisfaction_for_key_partitioning() + } } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 3b03c794051c1..a66e0588ceb3f 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -31,8 +31,9 @@ use crate::windows::{ }; use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, PhysicalExpr, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, + ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, + check_if_same_properties, }; use arrow::array::ArrayRef; @@ -243,12 +244,15 @@ impl ExecutionPlan for WindowAggExec { self.input_distribution_requirements().into_per_child() } - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() { - vec![Distribution::SinglePartition] + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + if self.partition_keys().is_empty() { + InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } else { - vec![Distribution::KeyPartitioned(self.partition_keys())] - }) + InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + self.partition_keys(), + )]) + .allow_range_satisfaction_for_key_partitioning() + } } fn with_new_children( diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 0dfe399a4646d..d3976f024a0ef 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -672,5 +672,230 @@ ORDER BY range_key, value; 35 350 35 350 +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + + +########## +# TEST 19: Window on Range Partition Column +# Range([range_key]) colocates equal range_key values, so +# PARTITION BY range_key is satisfied without a hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 20: Unbounded-Frame Window on Range Partition Column +# The unbounded frame makes DataFusion use WindowAggExec instead of +# BoundedWindowAggExec, which likewise reuses Range partitioning without a +# hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] +02)--WindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 21: Window on Non-Range Column Rehashes +# Range([range_key]) does not colocate non_range_key values, so +# PARTITION BY non_range_key still requires a hash repartition. +########## + +query TT +EXPLAIN SELECT non_range_key, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[non_range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value) FROM range_partitioned ORDER BY non_range_key, value; +---- +1 10 10 +1 100 110 +1 200 310 +1 300 610 +2 50 50 +2 150 200 +2 250 450 +2 350 800 + + +########## +# TEST 22: Unbounded-Frame Window on Non-Range Column Rehashes +# The unbounded frame makes DataFusion use WindowAggExec; Range([range_key]) +# does not colocate non_range_key values, so PARTITION BY non_range_key +# still requires a hash repartition. +########## + +query TT +EXPLAIN SELECT non_range_key, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] +02)--WindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +03)----SortExec: expr=[non_range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned ORDER BY non_range_key, value; +---- +1 10 610 +1 100 610 +1 200 610 +1 300 610 +2 50 800 +2 150 800 +2 250 800 +2 350 800 + + +########## +# TEST 23: Window Subset Satisfaction on Range Partition Column +# With the subset threshold met, Range([range_key]) satisfies +# PARTITION BY (range_key, non_range_key): equal composite keys share the +# same range_key, so they are already colocated. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, non_range_key@1 ASC NULLS LAST, value@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 24: Window Subset Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of PARTITION BY +# (range_key, non_range_key), so it should not satisfy the window key when +# subset satisfaction is disabled. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, non_range_key@1 ASC NULLS LAST, value@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 25: Window Without Partition Keys Uses a Single Partition +# A window with no PARTITION BY requires a single partition; range +# partitioning is not applicable. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortPreservingMergeExec: [value@1 ASC NULLS LAST] +04)------SortExec: expr=[value@1 ASC NULLS LAST], preserve_partitioning=[true] +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 60 +10 160 +15 310 +20 510 +25 760 +30 1060 +35 1410 + statement ok reset datafusion.explain.physical_plan_only; From e98fb600c15e24e27ec8a2dfec9d48ae7dedca39 Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Wed, 15 Jul 2026 02:36:20 -0700 Subject: [PATCH 11/22] Allow Range partitioned inputs to PartitionedTopK (#23355) - Closes #23290. In two commits: * Improve the robustness of `WindowTopN` to intermediate `RepartitionExec` nodes, to ensure that `PartitionedTopK` can still be used when `target_partitions` mismatches physical partitions. * Use `required_input_distributions` to declare required partitioning for `PartitionedTopKExec`, `BoundedWindowAggExec`, and `WindowAggExec`. Yes, new `slt` tests are added for both the `WindowTopN` robustness fix, and for the use of `required_input_distributions` in `PartitionedTopK`. No: only plan changes. --- .../physical-optimizer/src/window_topn.rs | 74 ++++----- .../src/sorts/partitioned_topk.rs | 1 + .../test_files/range_partitioning.slt | 157 ++++++++++++++++++ .../sqllogictest/test_files/window_topn.slt | 45 +++++ 4 files changed, 238 insertions(+), 39 deletions(-) diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 40dbddfbdf9fb..9459376054fda 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -46,6 +46,7 @@ use datafusion_physical_expr::window::StandardWindowExpr; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; @@ -119,17 +120,19 @@ impl WindowTopN { // Step 2: Extract limit from predicate (rn <= K, rn < K, etc.) let (col_idx, limit_n) = extract_window_limit(filter.predicate())?; - // Step 3: Walk through optional ProjectionExec to find BoundedWindowAggExec + // Step 3: Walk through optional ProjectionExec and RepartitionExec to find BoundedWindowAggExec let child = filter.input(); - let (window_exec, proj_between) = find_window_below(child)?; + let (window_exec, intermediates) = find_window_below(child)?; // Step 4: Verify col_idx references a ROW_NUMBER window output column - let input_field_count = window_exec.input().schema().fields().len(); + let window_exec_typed = window_exec.downcast_ref::()?; + let sort_exec = window_exec_typed.input().downcast_ref::()?; + let input_field_count = window_exec_typed.input().schema().fields().len(); if col_idx < input_field_count { return None; // Filter is on an input column, not a window column } let window_expr_idx = col_idx - input_field_count; - let window_exprs = window_exec.window_expr(); + let window_exprs = window_exec_typed.window_expr(); if window_expr_idx >= window_exprs.len() { return None; } @@ -137,8 +140,7 @@ impl WindowTopN { return None; } - // Step 5: Verify child of window is SortExec - let sort_exec = window_exec.input().downcast_ref::()?; + // Step 5: child of window is SortExec (verified above) let sort_child = sort_exec.input(); // Step 6: Determine partition_prefix_len from the window expression @@ -161,28 +163,19 @@ impl WindowTopN { .ok()?; // Step 8: Rebuild window with new child - let new_window = Arc::clone(&child_as_arc(window_exec)) + let mut result = window_exec .with_new_children(vec![Arc::new(partitioned_topk)]) .ok()?; - // Step 9: If ProjectionExec was between Filter and Window, rebuild it - let result = match proj_between { - Some(proj) => Arc::clone(&child_as_arc(proj)) - .with_new_children(vec![new_window]) - .ok()?, - None => new_window, - }; + // Step 9: Rebuild intermediate nodes (ProjectionExec/RepartitionExec) + for node in intermediates.into_iter().rev() { + result = node.with_new_children(vec![result]).ok()?; + } Some(result) } } -/// Helper to get an `Arc` from a reference. -/// We need this because `with_new_children` takes `Arc`. -fn child_as_arc(plan: &T) -> Arc { - Arc::new(plan.clone()) -} - impl PhysicalOptimizerRule for WindowTopN { fn optimize( &self, @@ -303,29 +296,32 @@ fn is_row_number(expr: &Arc) - udf.fun().name() == "row_number" } +type PlanAndIntermediates = (Arc, Vec>); + /// Walk below a plan node looking for a [`BoundedWindowAggExec`]. /// -/// Handles two cases: -/// - Direct child: `FilterExec → BoundedWindowAggExec` -/// - With projection: `FilterExec → ProjectionExec → BoundedWindowAggExec` +/// Handles sequences of `ProjectionExec` and `RepartitionExec`. +/// This is safe because `PartitionedTopKExec` can be pushed below them: +/// projections only provide aliases, and pushing the limit below repartitions +/// is safe because the limit is computed per-partition. /// -/// Returns the window exec and an optional `ProjectionExec` in between, -/// or `None` if no `BoundedWindowAggExec` is found within one or two levels. -fn find_window_below( - plan: &Arc, -) -> Option<(&BoundedWindowAggExec, Option<&ProjectionExec>)> { - // Direct child is BoundedWindowAggExec - if let Some(window) = plan.downcast_ref::() { - return Some((window, None)); - } +/// Returns the window exec and a list of intermediate nodes to rebuild, +/// or `None` if no `BoundedWindowAggExec` is found. +fn find_window_below(plan: &Arc) -> Option { + let mut current = Arc::clone(plan); + let mut intermediates = Vec::new(); - // Child is ProjectionExec with BoundedWindowAggExec below - if let Some(proj) = plan.downcast_ref::() { - let proj_child = proj.input(); - if let Some(window) = proj_child.downcast_ref::() { - return Some((window, Some(proj))); + loop { + if current.downcast_ref::().is_some() { + return Some((current, intermediates)); + } else if current.downcast_ref::().is_some() + || current.downcast_ref::().is_some() + { + let next = Arc::clone(current.children().first()?); + intermediates.push(current); + current = next; + } else { + return None; } } - - None } diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 17eb70ef12131..6a33aa909e648 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -314,6 +314,7 @@ impl ExecutionPlan for PartitionedTopKExec { crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( partition_exprs, )]) + .allow_range_satisfaction_for_key_partitioning() } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index d3976f024a0ef..423246ee2eebc 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -897,5 +897,162 @@ SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER 30 1060 35 1410 + + +########## +# TEST 26: PartitionedTopK on Range Partition Column +# Exact Range([range_key]) satisfies the TopK partition key and avoids repartitioning. +########## + +statement ok +set datafusion.optimizer.enable_window_topn = true; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0], order=[value@1 DESC] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY range_key; +---- +1 10 1 +5 50 1 +10 100 1 +15 150 1 +20 200 1 +25 250 1 +30 300 1 +35 350 1 + + +########## +# TEST 27: PartitionedTopK on Non-Range Column +# Partitioning on a non-range key cannot reuse Range([range_key]) and +# requires hash repartitioning. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[non_range_key@0], order=[value@1 DESC] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT * FROM ( + SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY non_range_key; +---- +1 300 1 +2 350 1 + + +########## +# TEST 28: PartitionedTopK Reuses Range Subset Partitioning +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies partitioning by (range_key, non_range_key). +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY range_key, non_range_key; +---- +1 1 10 1 +5 2 50 1 +10 1 100 1 +15 2 150 1 +20 1 200 1 +25 2 250 1 +30 1 300 1 +35 2 350 1 + + +########## +# TEST 29: Range Subset PartitionedTopK Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key), +# so it should not satisfy the TopK partition key when subset satisfaction is +# disabled. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + statement ok reset datafusion.explain.physical_plan_only; + +statement ok +reset datafusion.optimizer.enable_window_topn; diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt index bf9ce26b35537..09b52daa2ea79 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -624,3 +624,48 @@ DROP TABLE window_topn_nulls; # Reset config to default (false) statement ok SET datafusion.optimizer.enable_window_topn = false; + +statement ok +create table t(c1 int, c2 int) as values (1, 2), (3, 4); + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.repartition_windows = false; + +statement ok +set datafusion.execution.batch_size = 1; + +statement ok +set datafusion.optimizer.enable_window_topn = true; + +query TT +EXPLAIN SELECT * FROM ( + SELECT c1, c2, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c2 DESC) as rn + FROM t +) WHERE rn <= 1; +---- +logical_plan +01)Projection: t.c1, t.c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: t projection=[c1, c2] +physical_plan +01)ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=1, maintains_sort_order=true +03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------PartitionedTopKExec: fn=row_number, fetch=1, partition=[c1@0], order=[c2@1 DESC] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_windows = true; + +statement ok +set datafusion.execution.batch_size = 8192; + +statement ok +set datafusion.optimizer.enable_window_topn = false; From cd3a21887773b31fde53d2b8f84c29cbde4f6e7e Mon Sep 17 00:00:00 2001 From: Edson Petry <124717297+EdsonPetry@users.noreply.github.com> Date: Wed, 15 Jul 2026 05:37:07 -0400 Subject: [PATCH 12/22] fix: preserve range partitioning through joins (#23584) - Closes #23450. partitioned inner joins without hash repartitioning. However, join output still downgraded `Partitioning::Range` to `UnknownPartitioning`, so downstream joins and aggregates could not reuse the preserved range layout and inserted avoidable hash repartitions. - Preserve `Partitioning::Range` in `adjust_right_output_partitioning`. - Shift right-side range-ordering column indexes into the joined schema while retaining split points and sort options. - Add unit coverage for compound range keys and non-default sort options. - Update SQL logic test plans for nested range joins and a range join feeding an aggregate. Yes. - `./dev/rust_lint.sh` - `cargo clippy --all-targets --all-features -- -D warnings` - `cargo test -p datafusion-physical-plan test_adjust_right_output_partitioning_preserves_range` - `cargo test --profile=ci --test sqllogictests -- range_partitioning.slt` - `ulimit -n 10240 && RUST_BACKTRACE=1 cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption` Yes. Physical plans can preserve proven range partitioning through joins, allowing compatible downstream joins and aggregates to avoid unnecessary hash repartitioning. There are no public API changes. --- datafusion/physical-plan/src/joins/utils.rs | 70 +++++++++++++++++-- .../test_files/range_partitioning.slt | 30 ++++---- 2 files changed, 76 insertions(+), 24 deletions(-) diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 941a57d72bad4..134d650a9cd6b 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -33,7 +33,8 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, ProjectionExpr}; use crate::{ - ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, Statistics, + ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, + RangePartitioning, Statistics, }; // compatibility pub use super::join_filter::JoinFilter; @@ -67,7 +68,7 @@ use datafusion_common::hash_utils::create_hashes; use datafusion_common::stats::Precision; use datafusion_common::{ DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, - not_impl_err, plan_err, + internal_datafusion_err, not_impl_err, plan_err, }; use datafusion_expr::Operator; use datafusion_expr::interval_arithmetic::Interval; @@ -145,9 +146,19 @@ pub fn adjust_right_output_partitioning( Partitioning::Hash(new_exprs, *size) } Partitioning::Range(range) => { - // Range partitioning optimizer propagation is tracked in - // https://github.com/apache/datafusion/issues/22395 - Partitioning::UnknownPartitioning(range.partition_count()) + let ordering = add_offset_to_physical_sort_exprs( + range.ordering().iter().cloned(), + left_columns_len as _, + )?; + let ordering = LexOrdering::new(ordering).ok_or_else(|| { + internal_datafusion_err!( + "Offsetting range partitioning produced an empty ordering" + ) + })?; + Partitioning::Range(RangePartitioning::new( + ordering, + range.split_points().to_vec(), + )) } result => result.clone(), }; @@ -2149,7 +2160,7 @@ mod tests { use arrow::datatypes::{DataType, Fields}; use arrow::error::{ArrowError, Result as ArrowResult}; use datafusion_common::stats::Precision::{Absent, Exact, Inexact}; - use datafusion_common::{ScalarValue, arrow_datafusion_err, arrow_err}; + use datafusion_common::{ScalarValue, SplitPoint, arrow_datafusion_err, arrow_err}; use datafusion_physical_expr::PhysicalSortExpr; use rstest::rstest; @@ -3244,6 +3255,53 @@ mod tests { Ok(()) } + #[test] + fn test_adjust_right_output_partitioning_preserves_range() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ + ScalarValue::Int32(Some(10)), + ScalarValue::Int32(Some(100)), + ]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(20)), + ScalarValue::Int32(Some(50)), + ]), + ]; + let range = RangePartitioning::try_new( + LexOrdering::new([ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 2)), + SortOptions::new(true, false), + ), + ]) + .unwrap(), + split_points.clone(), + )?; + + let adjusted = adjust_right_output_partitioning(&Partitioning::Range(range), 3)?; + let expected = Partitioning::Range(RangePartitioning::new( + LexOrdering::new([ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 3)), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 5)), + SortOptions::new(true, false), + ), + ]) + .unwrap(), + split_points, + )); + + assert_eq!(adjusted, expected); + Ok(()) + } + #[test] fn test_calculate_join_output_ordering() -> Result<()> { let left_ordering = LexOrdering::new(vec![ diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 423246ee2eebc..1fedb0840390d 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -496,9 +496,8 @@ set datafusion.optimizer.preserve_file_partitions = 0; ########## # TEST 15: Nested Range Joins -# Compatible Range partitioning satisfies the lower join inputs. The upper join -# still repairs the intermediate join output with Hash repartitioning because -# HashJoinExec does not currently expose Range output partitioning. +# Compatible Range partitioning is preserved through the lower join, allowing +# the upper join to consume it without Hash repartitioning either input. ########## query TT @@ -509,12 +508,10 @@ JOIN range_partitioned s ON r.range_key = s.range_key; ---- physical_plan 01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@2, range_key@0)], projection=[range_key@0, value@1, value@3, value@5] -02)--RepartitionExec: partitioning=Hash([range_key@2], 4), input_partitions=4 -03)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)] -04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false -05)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false -06)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 -07)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +02)--HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query IIII SELECT l.range_key, l.value, r.value, s.value @@ -589,9 +586,8 @@ ORDER BY l.range_key; ########## # TEST 17: Range Join Feeds Aggregate -# The join inputs avoid Hash repartitioning, but the aggregate above the join -# still repartitions because HashJoinExec does not currently expose Range -# output partitioning. +# The join preserves compatible Range partitioning on range_key, allowing the +# aggregate above it to avoid Hash repartitioning. ########## query TT @@ -601,12 +597,10 @@ JOIN range_partitioned r ON l.range_key = r.range_key GROUP BY l.range_key; ---- physical_plan -01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)] -02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 -03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)] -04)------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] -05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false -06)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)] +02)--HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query II SELECT l.range_key, SUM(l.value + r.value) From a9de07acf5b9450de05d8ac4e883d691d286b45d Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Wed, 15 Jul 2026 08:31:57 -0400 Subject: [PATCH 13/22] Enforce co-partitioning for sort merge and symmetric hash joins (#23480) ## Which issue does this PR close? - Closes #23451 - Closes #23478 - Closes #23479 ## Rationale for this change Partition-index-aware joins require compatible input layouts. Compatible Range layouts can satisfy that without repartitioning. ## What changes are included in this PR? - Require co-partitioned children for sort-merge and partitioned symmetric hash joins. - Allow compatible Range inputs to satisfy those requirements. - Let streaming tables declare output partitioning and preserve it through scan projection. - Add sanity checks and range-partitioning SLT coverage. ## Are these changes tested? Yes ## Are there any user-facing changes? Streaming table providers can declare output partitioning with `StreamingTable::with_output_partitioning()`. --- datafusion/catalog/src/streaming.rs | 41 ++++- .../physical_optimizer/sanity_checker.rs | 74 +++++++- .../src/joins/sort_merge_join/exec.rs | 8 +- .../src/joins/symmetric_hash_join.rs | 19 +- datafusion/physical-plan/src/streaming.rs | 48 ++++- .../src/test_context/range_partitioning.rs | 104 ++++++++++- .../test_files/range_partitioning.slt | 167 ++++++++++++++++-- 7 files changed, 428 insertions(+), 33 deletions(-) diff --git a/datafusion/catalog/src/streaming.rs b/datafusion/catalog/src/streaming.rs index e609877c2b778..5bfecef1fb2ed 100644 --- a/datafusion/catalog/src/streaming.rs +++ b/datafusion/catalog/src/streaming.rs @@ -24,7 +24,10 @@ use async_trait::async_trait; use datafusion_common::{DFSchema, Result, plan_err}; use datafusion_expr::{Expr, SortExpr, TableType}; use datafusion_physical_expr::equivalence::project_ordering; -use datafusion_physical_expr::{LexOrdering, create_physical_sort_exprs}; +use datafusion_physical_expr::projection::ProjectionMapping; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, Partitioning, create_physical_sort_exprs, +}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; use log::debug; @@ -38,6 +41,7 @@ pub struct StreamingTable { partitions: Vec>, infinite: bool, sort_order: Vec, + output_partitioning: Option, } impl StreamingTable { @@ -62,6 +66,7 @@ impl StreamingTable { partitions, infinite: false, sort_order: vec![], + output_partitioning: None, }) } @@ -76,6 +81,33 @@ impl StreamingTable { self.sort_order = sort_order; self } + + /// Declares the output partitioning of this streaming table. + /// + /// The partitioning expressions refer to the table schema before scan + /// projection. If a scan projection removes a partitioning expression, the + /// physical plan reports unknown partitioning. + pub fn with_output_partitioning(mut self, output_partitioning: Partitioning) -> Self { + self.output_partitioning = Some(output_partitioning); + self + } + + fn output_partitioning( + &self, + projection: Option<&Vec>, + ) -> Result { + let Some(output_partitioning) = &self.output_partitioning else { + return Ok(Partitioning::UnknownPartitioning(self.partitions.len())); + }; + let Some(projection) = projection else { + return Ok(output_partitioning.clone()); + }; + + let projection_mapping = + ProjectionMapping::from_indices(projection, &self.schema)?; + let eq_properties = EquivalenceProperties::new(Arc::clone(&self.schema)); + Ok(output_partitioning.project(&projection_mapping, &eq_properties)) + } } #[async_trait] @@ -119,13 +151,16 @@ impl TableProvider for StreamingTable { vec![] }; - Ok(Arc::new(StreamingTableExec::try_new( + let exec = StreamingTableExec::try_new( Arc::clone(&self.schema), self.partitions.clone(), projection, LexOrdering::new(physical_sort), self.infinite, limit, - )?)) + )? + .with_output_partitioning(self.output_partitioning(projection)?)?; + + Ok(Arc::new(exec)) } } diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index e5718f5b3d0f7..3c426e2b09059 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -30,12 +30,13 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; use datafusion::prelude::{CsvReadOptions, SessionContext}; use datafusion_common::config::ConfigOptions; -use datafusion_common::{JoinType, Result, ScalarValue}; +use datafusion_common::{JoinType, NullEquality, Result, ScalarValue}; use datafusion_physical_expr::expressions::{Literal, col}; use datafusion_physical_expr::{Partitioning, RangePartitioning, SplitPoint}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; +use datafusion_physical_plan::joins::{StreamJoinPartitionMode, SymmetricHashJoinExec}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::{ExecutionPlan, displayable}; @@ -444,6 +445,77 @@ fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> { Ok(()) } +#[test] +fn test_sort_merge_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + let ordering: LexOrdering = [sort_expr("a", &schema)].into(); + + let compatible_join = sort_merge_join_exec( + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&compatible_join, true); + + let incompatible_join = sort_merge_join_exec( + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + sort_exec_with_preserve_partitioning( + ordering, + range_partitioned_exec(&schema, "a", [20])?, + ), + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + +#[test] +fn test_symmetric_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = Arc::new(SymmetricHashJoinExec::try_new( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + join_on.clone(), + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?) as Arc; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = Arc::new(SymmetricHashJoinExec::try_new( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + join_on, + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?) as Arc; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + #[tokio::test] /// Tests that plan is valid when the sort requirements are satisfied. async fn test_bounded_window_agg_sort_requirement() -> Result<()> { diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index ea6db4f77369a..d7975d67ac2a6 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -40,7 +40,8 @@ use crate::projection::{ use crate::spill::spill_manager::SpillManager; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties, + InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use arrow::compute::SortOptions; @@ -426,16 +427,17 @@ impl ExecutionPlan for SortMergeJoinExec { self.input_distribution_requirements().into_per_child() } - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + fn input_distribution_requirements(&self) -> InputDistributionRequirements { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - crate::InputDistributionRequirements::new(vec![ + InputDistributionRequirements::co_partitioned(vec![ Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), ]) + .allow_range_satisfaction_for_key_partitioning() } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 24337cbe6a45d..e0ec6b4cdaaef 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -53,7 +53,8 @@ use crate::projection::{ use crate::stream::EmptyRecordBatchStream; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, + InputDistributionRequirements, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, joins::StreamJoinPartitionMode, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; @@ -429,23 +430,27 @@ impl ExecutionPlan for SymmetricHashJoinExec { self.input_distribution_requirements().into_per_child() } - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(match self.mode { + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + match self.mode { StreamJoinPartitionMode::Partitioned => { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) .unzip(); - vec![ + InputDistributionRequirements::co_partitioned(vec![ Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), - ] + ]) + .allow_range_satisfaction_for_key_partitioning() } StreamJoinPartitionMode::SinglePartition => { - vec![Distribution::SinglePartition, Distribution::SinglePartition] + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) } - }) + } } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index cdf4b08f718c6..61a9b9cc6d0de 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -35,6 +35,7 @@ use crate::{ExecutionPlan, Partitioning, SendableRecordBatchStream}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::{Result, internal_err, plan_err}; use datafusion_execution::TaskContext; +use datafusion_physical_expr::projection::ProjectionMapping; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; use async_trait::async_trait; @@ -100,7 +101,7 @@ impl StreamingTableExec { let cache = Self::compute_properties( Arc::clone(&projected_schema), projected_output_ordering.clone(), - &partitions, + Partitioning::UnknownPartitioning(partitions.len()), infinite, ); Ok(Self { @@ -115,6 +116,25 @@ impl StreamingTableExec { }) } + /// Declares the output partitioning of this stream. + /// + /// `output_partitioning` must describe this plan's current output and have + /// the same number of partitions as the stream. + pub fn with_output_partitioning( + mut self, + output_partitioning: Partitioning, + ) -> Result { + if output_partitioning.partition_count() != self.partitions.len() { + return plan_err!( + "Output partitioning has {} partitions but stream has {} partitions", + output_partitioning.partition_count(), + self.partitions.len() + ); + } + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; + Ok(self) + } + pub fn partitions(&self) -> &Vec> { &self.partitions } @@ -147,14 +167,12 @@ impl StreamingTableExec { fn compute_properties( schema: SchemaRef, orderings: Vec, - partitions: &[Arc], + output_partitioning: Partitioning, infinite: bool, ) -> PlanProperties { // Calculate equivalence properties: let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings); - // Get output partitioning: - let output_partitioning = Partitioning::UnknownPartitioning(partitions.len()); let boundedness = if infinite { Boundedness::Unbounded { requires_infinite_memory: false, @@ -204,6 +222,16 @@ impl DisplayAs for StreamingTableExec { if let Some(fetch) = self.limit { write!(f, ", fetch={fetch}")?; } + if !matches!( + self.cache.output_partitioning(), + Partitioning::UnknownPartitioning(_) + ) { + write!( + f, + ", output_partitioning={}", + self.cache.output_partitioning() + )?; + } display_orderings(f, &self.projected_output_ordering)?; @@ -306,6 +334,17 @@ impl ExecutionPlan for StreamingTableExec { }; lex_orderings.push(ordering); } + let projection_mapping = ProjectionMapping::try_new( + projection + .expr() + .iter() + .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())), + &self.schema(), + )?; + let output_partitioning = self + .cache + .output_partitioning() + .project(&projection_mapping, self.cache.equivalence_properties()); StreamingTableExec::try_new( Arc::clone(self.partition_schema()), @@ -315,6 +354,7 @@ impl ExecutionPlan for StreamingTableExec { self.is_infinite(), self.limit(), ) + .and_then(|exec| exec.with_output_partitioning(output_partitioning)) .map(|e| Some(Arc::new(e) as _)) } diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index aa741dded77be..4c8545fecaa16 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -19,13 +19,23 @@ use std::fs::{create_dir_all, remove_dir_all, write}; use std::path::Path; use std::sync::Arc; -use arrow::datatypes::{DataType, Field, Schema}; +use arrow::array::{ArrayRef, Int32Array}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion::catalog::streaming::StreamingTable; use datafusion::common::{ScalarValue, SplitPoint}; use datafusion::datasource::file_format::csv::CsvFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; use datafusion::logical_expr::{Partitioning, RangePartitioning, col}; +use datafusion::physical_expr::{ + Partitioning as PhysicalPartitioning, PhysicalSortExpr, + RangePartitioning as PhysicalRangePartitioning, expressions::col as physical_col, +}; +use datafusion::physical_plan::streaming::PartitionStream; +use datafusion::physical_plan::test::TestPartitionStream; use datafusion::prelude::SessionContext; // ============================================================================== @@ -52,11 +62,13 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { .expect("range partitioning should be valid"), ); + let range_table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned"); + register_csv_listing_table( ctx, "range_partitioned", - Path::new(env!("CARGO_MANIFEST_DIR")) - .join("test_files/scratch_range_partitioning/range_partitioned"), + &range_table_dir, Arc::clone(&schema), [ "1,1,10\n5,2,50\n", @@ -67,6 +79,31 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Some(output_partitioning), ); + register_unbounded_range_stream_table( + ctx, + "unbounded_range_like", + Arc::clone(&schema), + [10, 20, 30], + [ + vec![(1, 1, 10), (5, 2, 50)], + vec![(10, 1, 100), (15, 2, 150)], + vec![(20, 1, 200), (25, 2, 250)], + vec![(30, 1, 300), (35, 2, 350)], + ], + ); + register_unbounded_range_stream_table( + ctx, + "unbounded_range_like_shifted", + Arc::clone(&schema), + [15, 20, 30], + [ + vec![(1, 1, 10), (5, 2, 50), (10, 1, 100)], + vec![(15, 2, 150)], + vec![(20, 1, 200), (25, 2, 250)], + vec![(30, 1, 300), (35, 2, 350)], + ], + ); + let shifted_output_partitioning = Partitioning::Range( RangePartitioning::try_new( vec![col("range_key").sort(true, true)], @@ -133,3 +170,64 @@ fn register_csv_listing_table( ctx.register_table(name, Arc::new(table)) .expect("test listing table registration should succeed"); } + +fn register_unbounded_range_stream_table( + ctx: &SessionContext, + name: &str, + schema: Arc, + split_points: [i32; 3], + partition_rows: [Vec<(i32, i32, i32)>; 4], +) { + let output_partitioning = PhysicalPartitioning::Range( + PhysicalRangePartitioning::try_new( + [PhysicalSortExpr { + expr: physical_col("range_key", &schema) + .expect("range key should exist in stream schema"), + options: SortOptions::default(), + }] + .into(), + split_points + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(), + ) + .expect("range partitioning should be valid"), + ); + let partitions = partition_rows + .into_iter() + .map(|rows| range_stream_partition(Arc::clone(&schema), &rows)) + .collect(); + + ctx.register_table( + name, + Arc::new( + StreamingTable::try_new(schema, partitions) + .expect("range stream table should be valid") + .with_infinite_table(true) + .with_output_partitioning(output_partitioning), + ), + ) + .expect("test stream table registration should succeed"); +} + +fn range_stream_partition( + schema: SchemaRef, + rows: &[(i32, i32, i32)], +) -> Arc { + let range_key: Vec = rows.iter().map(|(range_key, _, _)| *range_key).collect(); + let non_range_key: Vec = rows + .iter() + .map(|(_, non_range_key, _)| *non_range_key) + .collect(); + let value: Vec = rows.iter().map(|(_, _, value)| *value).collect(); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(range_key)) as ArrayRef, + Arc::new(Int32Array::from(non_range_key)) as ArrayRef, + Arc::new(Int32Array::from(value)) as ArrayRef, + ], + ) + .expect("range stream batch should be valid"); + Arc::new(TestPartitionStream::new_with_batches(vec![batch])) +} diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 1fedb0840390d..d16a90b9a3384 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -618,6 +618,149 @@ ORDER BY l.range_key; 30 600 35 700 +########## +# TEST 18: Sort Merge Join Avoids Repartition for Compatible Range Inputs +# Compatible Range inputs satisfy SortMergeJoinExec's co-partitioned +# KeyPartitioned requirements. +########## + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] +03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +06)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 19: Sort Merge Join Repartitions Incompatible Range Inputs +# Different Range split points do not satisfy SortMergeJoinExec's +# co-partitioned KeyPartitioned requirements. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] +03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +06)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +07)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +08)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +########## +# TEST 20: Symmetric Hash Join Avoids Repartition for Compatible Range Inputs +# Compatible Range streams satisfy SymmetricHashJoinExec's co-partitioned +# KeyPartitioned requirements. +########## + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +03)----StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +04)----StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) + +query III rowsort +SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; +---- +1 10 10 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 +5 50 50 + +########## +# TEST 21: Symmetric Hash Join Repartitions Incompatible Range Inputs +# Different Range split points do not satisfy SymmetricHashJoinExec's +# co-partitioned KeyPartitioned requirements. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +04)------StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +05)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +06)------StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4) + +query III rowsort +SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; +---- +1 10 10 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 +5 50 50 + statement ok reset datafusion.optimizer.prefer_hash_join; @@ -628,7 +771,7 @@ statement ok reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 18: Union of Range Partitioned Inputs +# TEST 22: Union of Range Partitioned Inputs # Each input exposes Range partitioning on range_key. These changes do not add a # cross-child Range relationship for UNION ALL. ########## @@ -677,7 +820,7 @@ set datafusion.optimizer.preserve_file_partitions = 0; ########## -# TEST 19: Window on Range Partition Column +# TEST 23: Window on Range Partition Column # Range([range_key]) colocates equal range_key values, so # PARTITION BY range_key is satisfied without a hash repartition. ########## @@ -705,7 +848,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM r ########## -# TEST 20: Unbounded-Frame Window on Range Partition Column +# TEST 24: Unbounded-Frame Window on Range Partition Column # The unbounded frame makes DataFusion use WindowAggExec instead of # BoundedWindowAggExec, which likewise reuses Range partitioning without a # hash repartition. @@ -734,7 +877,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BE ########## -# TEST 21: Window on Non-Range Column Rehashes +# TEST 25: Window on Non-Range Column Rehashes # Range([range_key]) does not colocate non_range_key values, so # PARTITION BY non_range_key still requires a hash repartition. ########## @@ -763,7 +906,7 @@ SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER B ########## -# TEST 22: Unbounded-Frame Window on Non-Range Column Rehashes +# TEST 26: Unbounded-Frame Window on Non-Range Column Rehashes # The unbounded frame makes DataFusion use WindowAggExec; Range([range_key]) # does not colocate non_range_key values, so PARTITION BY non_range_key # still requires a hash repartition. @@ -793,7 +936,7 @@ SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER B ########## -# TEST 23: Window Subset Satisfaction on Range Partition Column +# TEST 27: Window Subset Satisfaction on Range Partition Column # With the subset threshold met, Range([range_key]) satisfies # PARTITION BY (range_key, non_range_key): equal composite keys share the # same range_key, so they are already colocated. @@ -825,7 +968,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER B ########## -# TEST 24: Window Subset Rehashes Below Subset Threshold +# TEST 28: Window Subset Rehashes Below Subset Threshold # Range([range_key]) is only a subset of PARTITION BY # (range_key, non_range_key), so it should not satisfy the window key when # subset satisfaction is disabled. @@ -864,7 +1007,7 @@ reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 25: Window Without Partition Keys Uses a Single Partition +# TEST 29: Window Without Partition Keys Uses a Single Partition # A window with no PARTITION BY requires a single partition; range # partitioning is not applicable. ########## @@ -894,7 +1037,7 @@ SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER ########## -# TEST 26: PartitionedTopK on Range Partition Column +# TEST 30: PartitionedTopK on Range Partition Column # Exact Range([range_key]) satisfies the TopK partition key and avoids repartitioning. ########## @@ -937,7 +1080,7 @@ ORDER BY range_key; ########## -# TEST 27: PartitionedTopK on Non-Range Column +# TEST 31: PartitionedTopK on Non-Range Column # Partitioning on a non-range key cannot reuse Range([range_key]) and # requires hash repartitioning. ########## @@ -973,7 +1116,7 @@ ORDER BY non_range_key; ########## -# TEST 28: PartitionedTopK Reuses Range Subset Partitioning +# TEST 32: PartitionedTopK Reuses Range Subset Partitioning # With subset threshold met and preserve-file disabled, Range([range_key]) # satisfies partitioning by (range_key, non_range_key). ########## @@ -1014,7 +1157,7 @@ ORDER BY range_key, non_range_key; ########## -# TEST 29: Range Subset PartitionedTopK Rehashes Below Subset Threshold +# TEST 33: Range Subset PartitionedTopK Rehashes Below Subset Threshold # Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key), # so it should not satisfy the TopK partition key when subset satisfaction is # disabled. From 6709f7257a09e77f3a5e031a3733974d8752a5f2 Mon Sep 17 00:00:00 2001 From: RIchard Baah <137434454+Rich-T-kid@users.noreply.github.com> Date: Thu, 16 Jul 2026 14:38:08 -0400 Subject: [PATCH 14/22] allow interleaveExec to support Range partioning (#23623) ## Which issue does this PR close? - works towards #22395 - Closes #23455. ## Rationale for this change see #23455 ## What changes are included in this PR? The goal of this PR is to allow range partitioning to propagate through `InterleaveExec`. Updated `can_interleave()` to accept `Partitioning::Range` when all children share an identical RangePartitioning (same ordering and split points), matching the existing behavior for `Partitioning::Hash`. Updated `range_partitioning.slt` to expect InterleaveExec where it previously expected UnionExec, since `can_interleave()` now accepts Partitioning::Range. ## Are these changes tested? yes, the `range_partioning.slt` file as well as the `union.slt` file sql logic test. This PR also includes three test. ## Are there any user-facing changes? physical plans may look different now. --- datafusion/physical-plan/src/union.rs | 138 ++++++++++++++++- .../test_files/range_partitioning.slt | 143 +++++++++++++++++- 2 files changed, 273 insertions(+), 8 deletions(-) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 3ea2eb5402fe5..08c47ab3aee93 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -462,7 +462,10 @@ impl ExecutionPlan for UnionExec { /// Combines multiple input streams by interleaving them. /// -/// This only works if all inputs have the same hash-partitioning. +/// All inputs must share an identical [`Partitioning::Hash`] or [`Partitioning::Range`] so that +/// partition `k` covers the same data across every input. Each output partition is the +/// interleaving of the same-indexed partition from all inputs: +/// `output[k] = input[0][k] + input[1][k] + ... + input[n-1][k]` /// /// # Data Flow /// ```text @@ -507,7 +510,7 @@ impl InterleaveExec { pub fn try_new(inputs: Vec>) -> Result { assert_or_internal_err!( can_interleave(inputs.iter()), - "Not all InterleaveExec children have a consistent hash partitioning" + "Not all InterleaveExec children have a consistent hash or range partitioning" ); let cache = Self::compute_properties(&inputs)?; Ok(InterleaveExec { @@ -655,8 +658,12 @@ impl ExecutionPlan for InterleaveExec { } } -/// If all the input partitions have the same Hash partition spec with the first_input_partition -/// The InterleaveExec is partition aware. +/// Returns true if all inputs have the same [`Partitioning::Hash`] or [`Partitioning::Range`] +/// spec, making them safe to interleave. Two inputs are interleave-compatible when partition +/// `k` covers the identical key range or hash bucket across every input. +/// +/// Note: compatibility is checked sequentially against the first input, so +/// `InputDistributionRequirements::co_partitioned` is not needed here. /// /// It might be too strict here in the case that the input partition specs are compatible but not exactly the same. /// For example one input partition has the partition spec Hash('a','b','c') and @@ -669,7 +676,7 @@ pub fn can_interleave>>( }; let reference = first.borrow().output_partitioning(); - matches!(reference, Partitioning::Hash(_, _)) + matches!(reference, Partitioning::Hash(_, _) | Partitioning::Range(_)) && inputs .map(|plan| plan.borrow().output_partitioning().clone()) .all(|partition| partition == *reference) @@ -833,10 +840,13 @@ mod tests { use arrow::compute::SortOptions; use arrow::datatypes::DataType; + use datafusion_common::SplitPoint; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_physical_expr::RangePartitioning; use datafusion_physical_expr::equivalence::convert_to_orderings; use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; // Generate a schema which consists of 7 columns (a, b, c, d, e, f, g) fn create_test_schema() -> Result { @@ -1273,6 +1283,124 @@ mod tests { ); } + fn make_hash_exec( + schema: &SchemaRef, + hash_cols: Vec<&str>, + buckets: usize, + ) -> Result> { + let exprs = hash_cols + .iter() + .map(|c| col(c, schema)) + .collect::>>()?; + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Hash(exprs, buckets), + )?)) + } + + fn make_range_exec( + schema: &SchemaRef, + split_values: Vec, + sort_options: SortOptions, + ) -> Result> { + let sort_expr = + PhysicalSortExpr::new(col(schema.field(0).name(), schema)?, sort_options); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let split_points = split_values + .into_iter() + .map(|v| SplitPoint::new(vec![ScalarValue::Int32(Some(v))])) + .collect(); + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?), + )?)) + } + + #[test] + fn test_can_interleave_matrix() -> Result<()> { + let name_column = "name"; + let age_column = "age"; + let schema = Arc::new(Schema::new(vec![ + Field::new(name_column, DataType::Int32, true), + Field::new(age_column, DataType::Int32, true), + ])); + + let ascending = SortOptions { + descending: false, + nulls_first: false, + }; + struct Case { + inputs: Vec>, + expected: bool, + label: &'static str, + } + + let cases = vec![ + // compatible + Case { + label: "matching hash on single column", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column], 3)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + Case { + label: "matching hash on multiple columns", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + ], + }, + Case { + label: "matching range same splits and order", + expected: true, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 20], ascending)?, + ], + }, + // incompatible + Case { + label: "subset range partition", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 15], ascending)?, + ], + }, + Case { + label: "range different split points", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 30], ascending)?, + ], + }, + Case { + label: "mixed range and hash", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + ]; + + for case in cases { + assert_eq!( + can_interleave(case.inputs.iter()), + case.expected, + "{}", + case.label + ); + } + Ok(()) + } + #[test] fn test_union_cardinality_effect() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index d16a90b9a3384..377e8266e4d2d 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -772,8 +772,8 @@ reset datafusion.optimizer.preserve_file_partitions; ########## # TEST 22: Union of Range Partitioned Inputs -# Each input exposes Range partitioning on range_key. These changes do not add a -# cross-child Range relationship for UNION ALL. +# Each input exposes the same Range partitioning on range_key, so the optimizer +# converts UnionExec to InterleaveExec to avoid redundant repartitioning. ########## query TT @@ -782,7 +782,7 @@ UNION ALL SELECT range_key, value FROM range_partitioned; ---- physical_plan -01)UnionExec +01)InterleaveExec 02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false 03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false @@ -1193,3 +1193,140 @@ reset datafusion.explain.physical_plan_only; statement ok reset datafusion.optimizer.enable_window_topn; + +########## +# TEST 34: Subset of Inputs Compatible Does Not Trigger InterleaveExec +# In a three-way union, two inputs share the same Range split points [10,20,30] +# while the third has a partially-overlapping but different set [15,20,30]. +# can_interleave requires ALL inputs to match, so UnionExec is kept. +########## + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted +ORDER BY range_key, value; +---- +1 10 +1 10 +1 10 +5 50 +5 50 +5 50 +10 100 +10 100 +10 100 +15 150 +15 150 +15 150 +20 200 +20 200 +20 200 +25 250 +25 250 +25 250 +30 300 +30 300 +30 300 +35 350 +35 350 +35 350 + +########## +# TEST 35: Incompatible Range Split Points Falls Back to UnionExec +# Two range-partitioned inputs with different split points cannot be interleaved, +# so the optimizer keeps UnionExec instead of converting to InterleaveExec. +########## + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted +ORDER BY range_key, value; +---- +1 10 +1 10 +5 50 +5 50 +10 100 +10 100 +15 150 +15 150 +20 200 +20 200 +25 250 +25 250 +30 300 +30 300 +35 350 +35 350 + +########## +# TEST 36: InterleaveExec Propagates Range Partitioning to Aggregate +# InterleaveExec outputs the same Range partitioning as its compatible inputs, +# allowing a downstream aggregate on range_key to run SinglePartitioned without +# a Hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM ( + SELECT range_key, value FROM range_partitioned + UNION ALL + SELECT range_key, value FROM range_partitioned +) GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(value)] +02)--InterleaveExec +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) FROM ( + SELECT range_key, value FROM range_partitioned + UNION ALL + SELECT range_key, value FROM range_partitioned +) GROUP BY range_key ORDER BY range_key; +---- +1 20 +5 100 +10 200 +15 300 +20 400 +25 500 +30 600 +35 700 + +statement ok +reset datafusion.explain.physical_plan_only; From 7d37d4e68188c4623eb5aa460ffd935d0ffb35d6 Mon Sep 17 00:00:00 2001 From: Matthew Patton Date: Mon, 20 Jul 2026 03:17:21 -0400 Subject: [PATCH 15/22] feat: allow Full joins to reuse range co-partitioning in HashJoinExec (#23583) ## Which issue does this PR close? - Closes #23454. ## Rationale for this change #23184 let compatible range-partitioned inputs satisfy inner partitioned hash joins without repartitioning. Full partitioned equi joins still always went through the conservative hash-repartition path, even when both inputs were already co-partitioned by range on the join key(s). The per-partition unmatched-row tracking in `HashJoinExec` is already partition-local under `PartitionMode::Partitioned` (not shared globally like in `CollectLeft`), so Full-join semantics generalize cleanly to range co-partitioning with no additional bookkeeping required. ## What changes are included in this PR? - Extend `HashJoinExec::input_distribution_requirements()` to opt `JoinType::Full` in to `allow_range_satisfaction_for_key_partitioning()`, alongside the existing `JoinType::Inner` case. The underlying `co_partitioned` / `compatible_co_partitioning_layout` / `co_partitioning_satisfied` logic in `distribution_requirements.rs` was already join-type-agnostic, so no changes were needed there. - Add planner unit tests in `enforce_distribution.rs` covering both the compatible-layout case (no repartition inserted) and the incompatible-split-points case (repartition still inserted) for `JoinType::Full`. - Add a `range_partitioned_sparse` sqllogictest fixture table with the same partition layout as `range_partitioned` but only partially overlapping keys, and add sqllogictest coverage in `range_partitioning.slt` for: a compatible Full join avoiding repartition, an incompatible Full join still repartitioning, and matched/left-only/right-only unmatched rows produced correctly by a co-partitioned Full join. ## Are these changes tested? Yes: - Two new Rust unit tests in `datafusion/core/tests/physical_optimizer/enforce_distribution.rs` assert on the physical plan shape (repartition inserted or not) for compatible and incompatible range layouts. - Three new sqllogictest cases in `datafusion/sqllogictest/test_files/range_partitioning.slt` exercise the feature end-to-end against real data, including matched rows, left-only unmatched rows, and right-only unmatched rows for a Full outer join. - Existing `enforce_distribution` tests, `range_partitioning.slt`, and proto roundtrip tests all continue to pass. ## Are there any user-facing changes? Yes: `EXPLAIN` output for `Full` joins over compatible range-partitioned inputs will no longer show a `RepartitionExec`, and such queries will avoid the associated hash-shuffle cost at execution time. No public API changes. --------- Co-authored-by: Matthew Patton --- .../enforce_distribution.rs | 84 +++++++++++ .../physical-plan/src/joins/hash_join/exec.rs | 4 +- .../src/test_context/range_partitioning.rs | 29 +++- .../test_files/range_partitioning.slt | 130 ++++++++++++++++-- 4 files changed, 230 insertions(+), 17 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index a59b7e95e1550..58563a30663a4 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1125,6 +1125,90 @@ fn range_window_rehashes_incompatible_range_partitioning() -> Result<()> { Ok(()) } +#[test] +fn range_full_hash_join_reuses_compatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Full); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_full_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 30, 40], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Full); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + RepartitionExec: partitioning=Hash([a1@0], 4), input_partitions=4 + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4), file_type=parquet + " + ); + + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 1e859a045345b..6369d0b4c4d0c 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1274,7 +1274,9 @@ impl ExecutionPlan for HashJoinExec { ]), }; - if self.mode == PartitionMode::Partitioned && self.join_type == JoinType::Inner { + if self.mode == PartitionMode::Partitioned + && matches!(self.join_type, JoinType::Inner | JoinType::Full) + { requirements.allow_range_satisfaction_for_key_partitioning() } else { requirements diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index 4c8545fecaa16..d60fdcecf4b55 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -121,7 +121,7 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { "range_partitioned_shifted", Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_partitioned_shifted"), - schema, + Arc::clone(&schema), [ "1,1,10\n5,2,50\n10,1,100\n", "15,2,150\n", @@ -130,6 +130,33 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { ], Some(shifted_output_partitioning), ); + + let sparse_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_sparse", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_sparse"), + schema, + [ + "5,2,50\n8,3,80\n", + "10,1,100\n", + "20,1,200\n", + "30,1,300\n40,4,400\n", + ], + Some(sparse_output_partitioning), + ); } fn register_csv_listing_table( diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 377e8266e4d2d..d2f65e31e6225 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -385,9 +385,9 @@ ORDER BY l.non_range_key, l.value, r.value; 2 350 350 ########## -# TEST 12: Non-Inner Range Join Repartitions -# Only inner partitioned hash joins opt in to Range satisfying KeyPartitioned -# requirements. Non-inner joins keep using Hash repartitioning. +# TEST 12: Unsupported +# Only Inner and Full partitioned hash joins opt in to Range satisfying +# KeyPartitioned requirements. Other join types keep using Hash repartitioning. ########## query TT @@ -761,6 +761,106 @@ FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; 35 350 350 5 50 50 +########## +# TEST 22: Full Outer Join on Range Partition Column +# Full partitioned hash joins also opt in to Range satisfying KeyPartitioned +# requirements, so compatible Range layouts avoid Hash repartitioning here too. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 23: Full Outer Join Incompatible Range Repartitions +# Same as TEST 10, but for Full: differing split points between the two +# Range-partitioned inputs still require Hash repartitioning to co-partition. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 24: Full Outer Join Produces Matched and Unmatched Rows +# `range_partitioned` and `range_partitioned_sparse` share the same Range +# split points/partition count but only partially overlapping range_key +# values, so this exercises matched rows, left-only unmatched rows (NULLs on +# the right), and right-only unmatched rows (NULLs on the left) while still +# avoiding Hash repartitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, r.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, range_key@2, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT l.range_key, r.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +ORDER BY l.range_key, r.range_key; +---- +1 NULL 10 NULL +5 5 50 50 +10 10 100 100 +15 NULL 150 NULL +20 20 200 200 +25 NULL 250 NULL +30 30 300 300 +35 NULL 350 NULL +NULL 8 NULL 80 +NULL 40 NULL 400 + statement ok reset datafusion.optimizer.prefer_hash_join; @@ -771,7 +871,7 @@ statement ok reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 22: Union of Range Partitioned Inputs +# TEST 25: Union of Range Partitioned Inputs # Each input exposes the same Range partitioning on range_key, so the optimizer # converts UnionExec to InterleaveExec to avoid redundant repartitioning. ########## @@ -820,7 +920,7 @@ set datafusion.optimizer.preserve_file_partitions = 0; ########## -# TEST 23: Window on Range Partition Column +# TEST 26: Window on Range Partition Column # Range([range_key]) colocates equal range_key values, so # PARTITION BY range_key is satisfied without a hash repartition. ########## @@ -848,7 +948,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM r ########## -# TEST 24: Unbounded-Frame Window on Range Partition Column +# TEST 27: Unbounded-Frame Window on Range Partition Column # The unbounded frame makes DataFusion use WindowAggExec instead of # BoundedWindowAggExec, which likewise reuses Range partitioning without a # hash repartition. @@ -877,7 +977,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BE ########## -# TEST 25: Window on Non-Range Column Rehashes +# TEST 28: Window on Non-Range Column Rehashes # Range([range_key]) does not colocate non_range_key values, so # PARTITION BY non_range_key still requires a hash repartition. ########## @@ -906,7 +1006,7 @@ SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER B ########## -# TEST 26: Unbounded-Frame Window on Non-Range Column Rehashes +# TEST 29: Unbounded-Frame Window on Non-Range Column Rehashes # The unbounded frame makes DataFusion use WindowAggExec; Range([range_key]) # does not colocate non_range_key values, so PARTITION BY non_range_key # still requires a hash repartition. @@ -936,7 +1036,7 @@ SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER B ########## -# TEST 27: Window Subset Satisfaction on Range Partition Column +# TEST 30: Window Subset Satisfaction on Range Partition Column # With the subset threshold met, Range([range_key]) satisfies # PARTITION BY (range_key, non_range_key): equal composite keys share the # same range_key, so they are already colocated. @@ -968,7 +1068,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER B ########## -# TEST 28: Window Subset Rehashes Below Subset Threshold +# TEST 31: Window Subset Rehashes Below Subset Threshold # Range([range_key]) is only a subset of PARTITION BY # (range_key, non_range_key), so it should not satisfy the window key when # subset satisfaction is disabled. @@ -1007,7 +1107,7 @@ reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 29: Window Without Partition Keys Uses a Single Partition +# TEST 32: Window Without Partition Keys Uses a Single Partition # A window with no PARTITION BY requires a single partition; range # partitioning is not applicable. ########## @@ -1037,7 +1137,7 @@ SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER ########## -# TEST 30: PartitionedTopK on Range Partition Column +# TEST 33: PartitionedTopK on Range Partition Column # Exact Range([range_key]) satisfies the TopK partition key and avoids repartitioning. ########## @@ -1080,7 +1180,7 @@ ORDER BY range_key; ########## -# TEST 31: PartitionedTopK on Non-Range Column +# TEST 34: PartitionedTopK on Non-Range Column # Partitioning on a non-range key cannot reuse Range([range_key]) and # requires hash repartitioning. ########## @@ -1116,7 +1216,7 @@ ORDER BY non_range_key; ########## -# TEST 32: PartitionedTopK Reuses Range Subset Partitioning +# TEST 35: PartitionedTopK Reuses Range Subset Partitioning # With subset threshold met and preserve-file disabled, Range([range_key]) # satisfies partitioning by (range_key, non_range_key). ########## @@ -1157,7 +1257,7 @@ ORDER BY range_key, non_range_key; ########## -# TEST 33: Range Subset PartitionedTopK Rehashes Below Subset Threshold +# TEST 36: Range Subset PartitionedTopK Rehashes Below Subset Threshold # Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key), # so it should not satisfy the TopK partition key when subset satisfaction is # disabled. From 34149d9a825ea0a848d01187fd76526819d4494c Mon Sep 17 00:00:00 2001 From: Gustavo Schneiter Date: Tue, 21 Jul 2026 04:06:08 -0300 Subject: [PATCH 16/22] feat: support co-partitioned range right-side equi hash joins (#23484) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #23453. - Part of #22395. ## Rationale for this change #23184 let compatible range-partitioned inputs satisfy **inner** partitioned hash joins without repartitioning. Right-side partitioned equi joins have the same locality guarantee: `Right`, `RightSemi`, `RightAnti` and `RightMark` all anchor every output row on the probe (right) partition (`on_lr_is_preserved` is probe-side for all four), so when both inputs are co-partitioned by the join keys, execution stays partition-local and the hash repartition is unnecessary. This removes unnecessary `RepartitionExec`s for already-co-located inputs, extending the inner-join behavior from #23184 to the right-side variants. No micro-benchmark included, consistent with #23184; correctness is demonstrated by matched/unmatched execution results below. ## What changes are included in this PR? The only production change is widening the existing inner-only gate in `HashJoinExec::input_distribution_requirements` from `join_type == JoinType::Inner` to `matches!(join_type, Inner | Right | RightSemi | RightAnti | RightMark)` (under `PartitionMode::Partitioned`). The `co_partitioned` / range-satisfaction machinery from #23184 is unchanged — the sanity checker and enforce_distribution consume range satisfaction join-type-agnostically. One behavior note for #23376: range co-partitioned right joins now stay `Range`/`Range`, so partitioned dynamic filters are disabled for them (`has_partitioned_dynamic_filter_routing` returns false), the same safe delta #23184 introduced for inner joins. These variants are probe-not-preserved for pruning, so dynamic filters were not eligible to prune their probe rows regardless. ## Are these changes tested? Yes. - Optimizer (`enforce_distribution.rs`): 4 reuse tests (one per join type) proving compatible range/range inputs keep `Range` partitioning with no `RepartitionExec`; 4 incompatibility tests proving that mismatched split points, sort options, partition counts, or join-key expressions still insert a hash repartition; sanity-check pairs mirroring #23184. - Execution (`range_partitioning.slt`): `EXPLAIN` plan pins plus matched/unmatched result checks for `Right` (incl. `NULL` left values), `RightSemi`, `RightAnti`, and an incompatible-layout `Right` join (repartitions, correct results). - The new reuse tests fail on `main` without the gate change (verified by reverting the production diff: exactly the 4 reuse tests fail, everything else passes). - `RightMark` testing note: `RightMark` is not reachable from SQL in sqllogictest (`IN`-subquery decorrelation emits `LeftMark`; physical `RightMark` only appears via a statistics-based swap), so its co-partitioning behavior is pinned at the optimizer/plan level, and mark null-marker semantics (matched/unmatched/`NULL` build keys) are pinned via the `LeftMark` path in the slt. ## Are there any user-facing changes? No API changes. Plans over compatible range-partitioned inputs avoid a hash repartition for right-side equi hash joins. --- .../enforce_distribution.rs | 73 ++++ .../physical_optimizer/sanity_checker.rs | 26 ++ .../physical-plan/src/joins/hash_join/exec.rs | 10 +- .../src/test_context/range_partitioning.rs | 28 ++ .../test_files/range_partitioning.slt | 338 ++++++++++++++++-- 5 files changed, 449 insertions(+), 26 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 58563a30663a4..09a85ab87cde2 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1062,6 +1062,79 @@ fn range_inner_hash_join_rehashes_incompatible_range_partitioning() -> Result<() Ok(()) } +#[test] +fn range_right_mark_hash_join_reuses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::RightMark); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=RightMark, on=[(a@0, a@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_right_semi_hash_join_rehashes_incompatible_sort_options() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [20], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [20], + SortOptions { + descending: true, + nulls_first: true, + }, + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::RightSemi); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(a@0, a@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 + DataSourceExec: file_groups={2 groups: [[p0], [p1]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(20)], 2), file_type=parquet + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 + DataSourceExec: file_groups={2 groups: [[p0], [p1]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 DESC], [(20)], 2), file_type=parquet + " + ); + + Ok(()) +} + #[test] fn range_window_reuses_range_partitioning() -> Result<()> { let input = parquet_exec_with_output_partitioning(range_partitioning( diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index 3c426e2b09059..184125dcbe180 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -445,6 +445,32 @@ fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> { Ok(()) } +#[test] +fn test_partitioned_right_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + join_on.clone(), + None, + &JoinType::Right, + )?; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + join_on, + None, + &JoinType::Right, + )?; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + #[test] fn test_sort_merge_join_requires_co_partitioned_children() -> Result<()> { let schema = create_test_schema2(); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 6369d0b4c4d0c..3218f50558a36 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1275,7 +1275,15 @@ impl ExecutionPlan for HashJoinExec { }; if self.mode == PartitionMode::Partitioned - && matches!(self.join_type, JoinType::Inner | JoinType::Full) + && matches!( + self.join_type, + JoinType::Inner + | JoinType::Full + | JoinType::Right + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightMark + ) { requirements.allow_range_satisfaction_for_key_partitioning() } else { diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index d60fdcecf4b55..4141e000145a8 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -131,6 +131,34 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Some(shifted_output_partitioning), ); + // Same rows as `range_partitioned` but split into only three range + // partitions on `range_key`. Used to exercise the co-partition check when + // two Range inputs disagree on partition count. + let narrow_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_narrow", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_narrow"), + Arc::clone(&schema), + [ + "1,1,10\n5,2,50\n", + "10,1,100\n15,2,150\n", + "20,1,200\n25,2,250\n30,1,300\n35,2,350\n", + ], + Some(narrow_output_partitioning), + ); + let sparse_output_partitioning = Partitioning::Range( RangePartitioning::try_new( vec![col("range_key").sort(true, true)], diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index d2f65e31e6225..5a35653a437b5 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -385,9 +385,10 @@ ORDER BY l.non_range_key, l.value, r.value; 2 350 350 ########## -# TEST 12: Unsupported -# Only Inner and Full partitioned hash joins opt in to Range satisfying -# KeyPartitioned requirements. Other join types keep using Hash repartitioning. +# TEST 12: Left Range Join Repartitions +# Only Inner, Full, and right-side (Right/RightSemi/RightAnti/RightMark) +# partitioned hash joins opt in to Range satisfying KeyPartitioned +# requirements. Other join types, such as Left, keep using Hash repartitioning. ########## query TT @@ -619,7 +620,294 @@ ORDER BY l.range_key; 35 700 ########## -# TEST 18: Sort Merge Join Avoids Repartition for Compatible Range Inputs +# TEST 18: Right Join on Range Partition Column +# Compatible Range inputs satisfy the join's partitioning requirements, so no +# Hash repartitioning is inserted. The left filter keeps its Range partitioning +# and the unmatched right rows above 150 are preserved. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--FilterExec: value@1 <= 150 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +NULL 20 200 +NULL 25 250 +NULL 30 300 +NULL 35 350 + +########## +# TEST 19: Right Semi Join on Range Partition Column +# Compatible Range inputs avoid Hash repartitioning for RightSemi joins. +# Only right rows with a match on the filtered left side are returned. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT SEMI JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(range_key@0, range_key@0)] +02)--FilterExec: value@1 <= 150, projection=[range_key@0] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT SEMI JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 + +########## +# TEST 20: Right Anti Join on Range Partition Column +# Compatible Range inputs avoid Hash repartitioning for RightAnti joins. +# Only right rows without a match on the filtered left side are returned. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT ANTI JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=RightAnti, on=[(range_key@0, range_key@0)] +02)--FilterExec: value@1 <= 150, projection=[range_key@0] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT ANTI JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +20 200 +25 250 +30 300 +35 350 + +########## +# TEST 21: Incompatible Range Right Join Repartitions +# The split points of the two inputs differ, so the co-partitioned layout +# requirement cannot be satisfied and Hash repartitioning repairs both sides +# of the right join. Results stay correct on the repartitioned path. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----FilterExec: value@1 <= 150 +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +06)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +NULL 20 200 +NULL 25 250 +NULL 30 300 +NULL 35 350 + +########## +# TEST 22: Composite-Key Right Join Repartitions +# Range([range_key]) does not satisfy a partitioned join on +# (range_key, non_range_key), so both sides repartition on the full key. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0), (non_range_key@1, non_range_key@1)], projection=[range_key@0, non_range_key@1, value@2, value@5] +02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key +ORDER BY l.range_key; +---- +1 1 10 10 +5 2 50 50 +10 1 100 100 +15 2 150 150 +20 1 200 200 +25 2 250 250 +30 1 300 300 +35 2 350 350 + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +########## +# TEST 23: Right Join with Mismatched Range Partition Counts Repartitions +# Both inputs are range partitioned on range_key, but declare a different number +# of partitions (four vs three). The per-child key requirements can be satisfied +# by Range, but the co-partitioned layout requirement cannot, so Hash +# repartitioning repairs both sides of the right join. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned_narrow r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=3 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20)], 3), file_type=csv, has_header=false + +query III +SELECT l.value, r.range_key, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned_narrow r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +200 20 200 +250 25 250 +300 30 300 +350 35 350 + +########## +# TEST 24: Right Join on Non-Range Key Repartitions +# Both inputs expose Range([range_key]), but the join key is non_range_key. +# Range([range_key]) does not satisfy KeyPartitioned([non_range_key]), so +# planning inserts Hash repartitioning on the actual join key for the right join. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT non_range_key, value FROM range_partitioned WHERE range_key < 10) l +RIGHT JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(non_range_key@0, non_range_key@1)], projection=[value@1, range_key@2, value@4] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----FilterExec: range_key@0 < 10, projection=[non_range_key@1, value@2] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--RepartitionExec: partitioning=Hash([non_range_key@1], 4), input_partitions=4 +06)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT non_range_key, value FROM range_partitioned WHERE range_key < 10) l +RIGHT JOIN range_partitioned r ON l.non_range_key = r.non_range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +10 10 100 +50 15 150 +10 20 200 +50 25 250 +10 30 300 +50 35 350 + +########## +# TEST 25: Mark Join Marker Semantics +# Mark joins preserve matched, unmatched, and NULL-key marker behavior over +# range-partitioned inputs. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150); +---- +physical_plan +01)FilterExec: non_range_key@1 = 2 OR mark@3, projection=[range_key@0, value@2] +02)--HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(range_key@0, range_key@0)] +03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +06)------FilterExec: value@1 <= 150, projection=[range_key@0] +07)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +# Matched rows have mark=true and are returned; unmatched rows have +# mark=false and are only returned when non_range_key = 2. +query II +SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150) +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +# NULL join keys on the build side never match: rows whose keys only "match" +# the NULL entries keep a non-true marker and are filtered out unless the +# non_range_key = 2 disjunct covers them. +query II +SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT CASE WHEN value <= 150 THEN range_key ELSE NULL END FROM range_partitioned) +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +########## +# TEST 26: Sort Merge Join Avoids Repartition for Compatible Range Inputs # Compatible Range inputs satisfy SortMergeJoinExec's co-partitioned # KeyPartitioned requirements. ########## @@ -656,7 +944,7 @@ ORDER BY l.range_key; 35 350 350 ########## -# TEST 19: Sort Merge Join Repartitions Incompatible Range Inputs +# TEST 27: Sort Merge Join Repartitions Incompatible Range Inputs # Different Range split points do not satisfy SortMergeJoinExec's # co-partitioned KeyPartitioned requirements. ########## @@ -695,7 +983,7 @@ statement ok reset datafusion.optimizer.prefer_hash_join; ########## -# TEST 20: Symmetric Hash Join Avoids Repartition for Compatible Range Inputs +# TEST 28: Symmetric Hash Join Avoids Repartition for Compatible Range Inputs # Compatible Range streams satisfy SymmetricHashJoinExec's co-partitioned # KeyPartitioned requirements. ########## @@ -729,7 +1017,7 @@ FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; 5 50 50 ########## -# TEST 21: Symmetric Hash Join Repartitions Incompatible Range Inputs +# TEST 29: Symmetric Hash Join Repartitions Incompatible Range Inputs # Different Range split points do not satisfy SymmetricHashJoinExec's # co-partitioned KeyPartitioned requirements. ########## @@ -762,7 +1050,7 @@ FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; 5 50 50 ########## -# TEST 22: Full Outer Join on Range Partition Column +# TEST 30: Full Outer Join on Range Partition Column # Full partitioned hash joins also opt in to Range satisfying KeyPartitioned # requirements, so compatible Range layouts avoid Hash repartitioning here too. ########## @@ -793,7 +1081,7 @@ ORDER BY l.range_key; 35 350 350 ########## -# TEST 23: Full Outer Join Incompatible Range Repartitions +# TEST 31: Full Outer Join Incompatible Range Repartitions # Same as TEST 10, but for Full: differing split points between the two # Range-partitioned inputs still require Hash repartitioning to co-partition. ########## @@ -826,7 +1114,7 @@ ORDER BY l.range_key; 35 350 350 ########## -# TEST 24: Full Outer Join Produces Matched and Unmatched Rows +# TEST 32: Full Outer Join Produces Matched and Unmatched Rows # `range_partitioned` and `range_partitioned_sparse` share the same Range # split points/partition count but only partially overlapping range_key # values, so this exercises matched rows, left-only unmatched rows (NULLs on @@ -871,7 +1159,7 @@ statement ok reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 25: Union of Range Partitioned Inputs +# TEST 33: Union of Range Partitioned Inputs # Each input exposes the same Range partitioning on range_key, so the optimizer # converts UnionExec to InterleaveExec to avoid redundant repartitioning. ########## @@ -920,7 +1208,7 @@ set datafusion.optimizer.preserve_file_partitions = 0; ########## -# TEST 26: Window on Range Partition Column +# TEST 34: Window on Range Partition Column # Range([range_key]) colocates equal range_key values, so # PARTITION BY range_key is satisfied without a hash repartition. ########## @@ -948,7 +1236,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM r ########## -# TEST 27: Unbounded-Frame Window on Range Partition Column +# TEST 35: Unbounded-Frame Window on Range Partition Column # The unbounded frame makes DataFusion use WindowAggExec instead of # BoundedWindowAggExec, which likewise reuses Range partitioning without a # hash repartition. @@ -977,7 +1265,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BE ########## -# TEST 28: Window on Non-Range Column Rehashes +# TEST 36: Window on Non-Range Column Rehashes # Range([range_key]) does not colocate non_range_key values, so # PARTITION BY non_range_key still requires a hash repartition. ########## @@ -1006,7 +1294,7 @@ SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER B ########## -# TEST 29: Unbounded-Frame Window on Non-Range Column Rehashes +# TEST 37: Unbounded-Frame Window on Non-Range Column Rehashes # The unbounded frame makes DataFusion use WindowAggExec; Range([range_key]) # does not colocate non_range_key values, so PARTITION BY non_range_key # still requires a hash repartition. @@ -1036,7 +1324,7 @@ SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER B ########## -# TEST 30: Window Subset Satisfaction on Range Partition Column +# TEST 38: Window Subset Satisfaction on Range Partition Column # With the subset threshold met, Range([range_key]) satisfies # PARTITION BY (range_key, non_range_key): equal composite keys share the # same range_key, so they are already colocated. @@ -1068,7 +1356,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER B ########## -# TEST 31: Window Subset Rehashes Below Subset Threshold +# TEST 39: Window Subset Rehashes Below Subset Threshold # Range([range_key]) is only a subset of PARTITION BY # (range_key, non_range_key), so it should not satisfy the window key when # subset satisfaction is disabled. @@ -1107,7 +1395,7 @@ reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 32: Window Without Partition Keys Uses a Single Partition +# TEST 40: Window Without Partition Keys Uses a Single Partition # A window with no PARTITION BY requires a single partition; range # partitioning is not applicable. ########## @@ -1137,7 +1425,7 @@ SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER ########## -# TEST 33: PartitionedTopK on Range Partition Column +# TEST 41: PartitionedTopK on Range Partition Column # Exact Range([range_key]) satisfies the TopK partition key and avoids repartitioning. ########## @@ -1180,7 +1468,7 @@ ORDER BY range_key; ########## -# TEST 34: PartitionedTopK on Non-Range Column +# TEST 42: PartitionedTopK on Non-Range Column # Partitioning on a non-range key cannot reuse Range([range_key]) and # requires hash repartitioning. ########## @@ -1216,7 +1504,7 @@ ORDER BY non_range_key; ########## -# TEST 35: PartitionedTopK Reuses Range Subset Partitioning +# TEST 43: PartitionedTopK Reuses Range Subset Partitioning # With subset threshold met and preserve-file disabled, Range([range_key]) # satisfies partitioning by (range_key, non_range_key). ########## @@ -1257,7 +1545,7 @@ ORDER BY range_key, non_range_key; ########## -# TEST 36: Range Subset PartitionedTopK Rehashes Below Subset Threshold +# TEST 44: Range Subset PartitionedTopK Rehashes Below Subset Threshold # Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key), # so it should not satisfy the TopK partition key when subset satisfaction is # disabled. @@ -1295,7 +1583,7 @@ statement ok reset datafusion.optimizer.enable_window_topn; ########## -# TEST 34: Subset of Inputs Compatible Does Not Trigger InterleaveExec +# TEST 45: Subset of Inputs Compatible Does Not Trigger InterleaveExec # In a three-way union, two inputs share the same Range split points [10,20,30] # while the third has a partially-overlapping but different set [15,20,30]. # can_interleave requires ALL inputs to match, so UnionExec is kept. @@ -1351,7 +1639,7 @@ ORDER BY range_key, value; 35 350 ########## -# TEST 35: Incompatible Range Split Points Falls Back to UnionExec +# TEST 46: Incompatible Range Split Points Falls Back to UnionExec # Two range-partitioned inputs with different split points cannot be interleaved, # so the optimizer keeps UnionExec instead of converting to InterleaveExec. ########## @@ -1393,7 +1681,7 @@ ORDER BY range_key, value; 35 350 ########## -# TEST 36: InterleaveExec Propagates Range Partitioning to Aggregate +# TEST 47: InterleaveExec Propagates Range Partitioning to Aggregate # InterleaveExec outputs the same Range partitioning as its compatible inputs, # allowing a downstream aggregate on range_key to run SinglePartitioned without # a Hash repartition. From e95b7a260dd01f767436916525e2936cab15537e Mon Sep 17 00:00:00 2001 From: Saad Tajwar <59696464+saadtajwar@users.noreply.github.com> Date: Tue, 21 Jul 2026 05:26:52 -0400 Subject: [PATCH 17/22] feat: complete range repartition physical planning (#23617) ## Which issue does this PR close? - Closes #23230 ## Rationale for this change After #23231 was merged in for supporting physical execution of the range repartitioning scheme, we still had a few methods on physical planning unimplemented, specifically `try_swapping_with_projection`, `try_pushdown_sort`, `repartitioned` - this PR finishes the implementation of those methods ## What changes are included in this PR? - `try_swapping_with_projection`: similar to the `Hash` scheme, for `Range` we call `update_expr` for each of the range key expressions to attempt rewriting based on the projection expressions - `try_pushdown_sort`: same as other variants, we delegate to the child and wrap with a new `RepartitionExec` - `repartitioned`: unable to support for Range, left comment in codebase with explanation ## Are these changes tested? Yes ## Are there any user-facing changes? No --- .../physical-plan/src/repartition/mod.rs | 346 +++++++++++++++++- 1 file changed, 329 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index f1c06ccaff89d..a20714d1e8647 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1548,10 +1548,29 @@ impl ExecutionPlan for RepartitionExec { } Partitioning::Hash(new_partitions, *size) } - Partitioning::Range(_) => { - // Range partitioning optimizer propagation is tracked in - // https://github.com/apache/datafusion/issues/23230 - return Ok(None); + Partitioning::Range(range_partitioning) => { + // Rewrite range key expressions through the projection. + let mut sort_exprs = + Vec::with_capacity(range_partitioning.ordering().len()); + for sort_expr in range_partitioning.ordering() { + let Some(new_expr) = + update_expr(&sort_expr.expr, projection.expr(), false)? + else { + return Ok(None); + }; + sort_exprs.push(PhysicalSortExpr::new(new_expr, sort_expr.options)); + } + + let Some(ordering) = LexOrdering::new(sort_exprs) else { + return internal_err!( + "failed to create LexOrdering for range partitioning" + ); + }; + + Partitioning::Range(RangePartitioning::try_new( + ordering, + range_partitioning.split_points().to_vec(), + )?) } others => others.clone(), }; @@ -1589,16 +1608,6 @@ 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/23230 - return Ok(SortOrderPushdownResult::Unsupported); - } - 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| { @@ -1621,12 +1630,11 @@ impl ExecutionPlan for RepartitionExec { new_properties.partitioning = match new_properties.partitioning { RoundRobinBatch(_) => RoundRobinBatch(target_partitions), Hash(hash, _) => Hash(hash, target_partitions), - UnknownPartitioning(_) => UnknownPartitioning(target_partitions), Range(_) => { - // Range repartition optimizations are tracked in - // https://github.com/apache/datafusion/issues/23230 + // Number of partitions is constrained by the split points and cannot be changed return Ok(None); } + UnknownPartitioning(_) => UnknownPartitioning(target_partitions), }; Ok(Some(Arc::new(Self { input: Arc::clone(&self.input), @@ -2078,6 +2086,8 @@ mod tests { use std::collections::HashSet; use super::*; + use crate::empty::EmptyExec; + use crate::projection::ProjectionExpr; use crate::test::TestMemoryExec; use crate::{ test::{ @@ -2538,6 +2548,281 @@ mod tests { Ok(()) } + #[test] + fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> { + // Three columns so the projection both narrows the schema (required for + // swap) and moves the range key from @0 to @1. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; + + let swapped = repartition + .try_swapping_with_projection(&projection)? + .expect("swap should succeed when projection keeps the range key"); + let swapped_repartition = swapped + .downcast_ref::() + .expect("top node should be RepartitionExec"); + + 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.split_points(), + &[SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])] + ); + + Ok(()) + } + + #[test] + fn range_repartition_does_not_swap_when_projection_drops_key() -> Result<()> { + // Drop a simple range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload"])?; + assert!( + repartition + .try_swapping_with_projection(&projection)? + .is_none() + ); + + // Drop part of a compound range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + Field::new("c", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["a", "b"], vec![vec![10, 1]])?, + )?); + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?; + assert!( + repartition + .try_swapping_with_projection(&projection)? + .is_none() + ); + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_when_maintains_order() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source with preserve_order: Range maintains input order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let repartition = Arc::new( + RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )? + .with_preserve_order(), + ); + assert!(repartition.maintains_input_order()[0]); + + match repartition.try_pushdown_sort(ordering.as_ref())? { + SortOrderPushdownResult::Exact { inner } => { + let pushed = inner + .downcast_ref::() + .expect("pushdown should keep RepartitionExec"); + + assert!(pushed.preserve_order()); + assert!(pushed.maintains_input_order()[0]); + + let range = expect_range_partitioning(pushed.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); + assert_eq!( + inner.properties().output_ordering().map(|o| o.to_string()), + Some(ordering.to_string()), + "pushed repartition output ordering should match the requested sort" + ); + } + other => panic!("expected Exact sort pushdown, got {other:?}"), + } + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_unsupported_without_order_maintenance() + -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source without preserve_order: Range does not maintain order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let repartition = Arc::new(RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + assert!(!repartition.maintains_input_order()[0]); + + assert!(matches!( + repartition.try_pushdown_sort(ordering.as_ref())?, + SortOrderPushdownResult::Unsupported + )); + + Ok(()) + } + + fn range_partitioning_on_columns( + schema: &SchemaRef, + key_columns: &[&str], + split_points: Vec>, + ) -> Result { + let Some(ordering) = LexOrdering::new( + key_columns + .iter() + .map(|name| { + Ok(PhysicalSortExpr::new( + col(name, schema)?, + SortOptions::default(), + )) + }) + .collect::>>()?, + ) else { + return exec_err!("range ordering must not be empty"); + }; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points + .into_iter() + .map(|values| { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::UInt32(Some(value))) + .collect(), + ) + }) + .collect(), + )?)) + } + + fn projection_on_columns( + input: &Arc, + names: &[&str], + ) -> Result { + let exprs = names + .iter() + .map(|name| { + Ok(ProjectionExpr { + expr: col(name, &input.schema())?, + alias: (*name).to_string(), + }) + }) + .collect::>>()?; + ProjectionExec::try_new(exprs, Arc::clone(input)) + } + + fn expect_range_partitioning(partitioning: &Partitioning) -> &RangePartitioning { + match partitioning { + Partitioning::Range(range) => range, + other => panic!("expected Range partitioning, got {other:?}"), + } + } + + /// Test source that claims Exact support for any sort pushdown request. + #[derive(Debug, Clone)] + struct ExactSortPushdownExec { + cache: Arc, + } + + impl ExactSortPushdownExec { + fn new(schema: SchemaRef, num_partitions: usize, ordering: LexOrdering) -> Self { + use crate::execution_plan::{Boundedness, EmissionType}; + Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new_with_orderings(schema, [ordering]), + Partitioning::UnknownPartitioning(num_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + )), + } + } + } + + impl DisplayAs for ExactSortPushdownExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ExactSortPushdownExec") + } + } + + impl ExecutionPlan for ExactSortPushdownExec { + fn name(&self) -> &str { + "ExactSortPushdownExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(EmptyRecordBatchStream::new(self.schema()))) + } + + fn try_pushdown_sort( + &self, + _order: &[PhysicalSortExpr], + ) -> Result>> { + Ok(SortOrderPushdownResult::Exact { + inner: Arc::new(self.clone()), + }) + } + } + #[tokio::test] async fn test_repartition_with_coalescing() -> Result<()> { let schema = test_schema(false); @@ -3666,6 +3951,33 @@ mod test { Ok(()) } + #[test] + fn test_range_repartitioned_returns_none() -> Result<()> { + let schema = test_schema(); + let source = memory_exec(&schema); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new( + col("c0", &schema)?, + SortOptions::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), + SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), + ], + )?); + let exec = RepartitionExec::try_new(source, partitioning)?; + + // Range partition count is fixed by split points, so repartitioned() + // cannot change it to an arbitrary target. + let result = exec.repartitioned(10, &Default::default())?; + assert!( + result.is_none(), + "range repartitioning should not support changing partition count" + ); + Ok(()) + } + fn test_schema() -> Arc { Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) } From 5d0aaf5f9cb3d4d78e61248a3d70b8557eacb5c7 Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Tue, 21 Jul 2026 17:59:33 +0200 Subject: [PATCH 18/22] feat(physical-plan): Allow co-partitioned Partitioning::Range inputs for left-side hash joins (#23487) ## Which issue does this PR close? - Closes #23452 ## Rationale for this change Allows compatible `Partitioning::Range` inputs to satisfy partitioned hash join distribution requirements for left-side joins, avoiding unnecessary hash repartitioning. ## What changes are included in this PR? - Enables range co-partitioning satisfaction for `Left`, `LeftSemi`, `LeftAnti`, and `LeftMark` hash joins. - Adds optimizer and sqllogictest coverage for compatible and incompatible range layouts. - Covers matched/unmatched rows and LeftMark null-related marker behavior. ## Are these changes tested? Yes. Added/updated physical optimizer tests and `range_partitioning.slt`. ## Are there any user-facing changes? --- .../enforce_distribution.rs | 73 +++++ .../physical-plan/src/joins/hash_join/exec.rs | 14 +- .../test_files/range_partitioning.slt | 271 +++++++++++++++--- 3 files changed, 300 insertions(+), 58 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 09a85ab87cde2..5ebc413d60d7b 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1282,6 +1282,79 @@ fn range_full_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> Ok(()) } +#[test] +fn range_left_mark_hash_join_reuses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftMark); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(a@0, a@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_left_anti_hash_join_rehashes_incompatible_null_options() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions { + descending: false, + nulls_first: false, + }, + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftAnti); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(a@0, a@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC NULLS LAST], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 3218f50558a36..c1d515c284ced 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1274,17 +1274,9 @@ impl ExecutionPlan for HashJoinExec { ]), }; - if self.mode == PartitionMode::Partitioned - && matches!( - self.join_type, - JoinType::Inner - | JoinType::Full - | JoinType::Right - | JoinType::RightSemi - | JoinType::RightAnti - | JoinType::RightMark - ) - { + if self.mode == PartitionMode::Partitioned { + // Compatible Range inputs co-locate equal join keys, which + // satisfies the co-partitioned requirement for hash joins. requirements.allow_range_satisfaction_for_key_partitioning() } else { requirements diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 5a35653a437b5..e53436c22e173 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -385,28 +385,145 @@ ORDER BY l.non_range_key, l.value, r.value; 2 350 350 ########## -# TEST 12: Left Range Join Repartitions -# Only Inner, Full, and right-side (Right/RightSemi/RightAnti/RightMark) -# partitioned hash joins opt in to Range satisfying KeyPartitioned -# requirements. Other join types, such as Left, keep using Hash repartitioning. +# TEST 12: Left-Side Range Hash Joins +# Compatible Range layouts satisfy left-side partitioned hash join +# requirements without Hash repartitioning. ########## query TT EXPLAIN SELECT l.range_key, l.value, r.value FROM range_partitioned l -LEFT JOIN range_partitioned r ON l.range_key = r.range_key; +LEFT JOIN (SELECT range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--FilterExec: value@1 <= 150 +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN (SELECT range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 NULL +25 250 NULL +30 300 NULL +35 350 NULL + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +LEFT SEMI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(range_key@0, range_key@0)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--FilterExec: value@1 <= 150, projection=[range_key@0] +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +LEFT SEMI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +LEFT ANTI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(range_key@0, range_key@0)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--FilterExec: value@1 <= 150, projection=[range_key@0] +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +LEFT ANTI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +20 200 +25 250 +30 300 +35 350 + +########## +# TEST 13: Left-Side Range Hash Joins With Incomplete Range Keys +# Range partitioning covers only range_key, so joins requiring additional +# or different keys are repaired with Hash repartitioning. +########## + +# Range([range_key]) is only a subset of the composite join key, so the +# co-partitioned hash join requirement is repaired with Hash repartitioning. +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN (SELECT range_key, non_range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0), (non_range_key@1, non_range_key@1)], projection=[range_key@0, non_range_key@1, value@2, value@5] +02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)----FilterExec: value@2 <= 150 +06)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +# Range([range_key]) does not satisfy a join keyed on non_range_key. +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(non_range_key@1, non_range_key@0)], projection=[range_key@0, non_range_key@1, value@2, value@4] +02)--RepartitionExec: partitioning=Hash([non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +########## +# TEST 14: Left-Side Range Hash Joins With Incompatible Range Layouts +# Different split points or partition counts do not satisfy the +# co-partitioned layout requirement. +########## + +# Different split points do not satisfy the co-partitioned layout requirement. +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned_shifted r ON l.range_key = r.range_key; ---- physical_plan 01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] 02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false 04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 -05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false query III SELECT l.range_key, l.value, r.value FROM range_partitioned l -LEFT JOIN range_partitioned r ON l.range_key = r.range_key +LEFT JOIN range_partitioned_shifted r ON l.range_key = r.range_key ORDER BY l.range_key; ---- 1 10 10 @@ -418,8 +535,70 @@ ORDER BY l.range_key; 30 300 300 35 350 350 +# Different partition counts do not satisfy the co-partitioned layout +# requirement. +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned_narrow r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=3 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20)], 3), file_type=csv, has_header=false + ########## -# TEST 13: Compatible Range Join Repartitions to Increase Parallelism +# TEST 15: LeftMark Subqueries Over Range Hash Joins +# SQL IN subqueries decorrelate to LeftMark joins. These queries pin matched, +# unmatched, and NULL marker behavior over compatible Range inputs. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150); +---- +physical_plan +01)FilterExec: non_range_key@1 = 2 OR mark@3, projection=[range_key@0, value@2] +02)--HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(range_key@0, range_key@0)] +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----FilterExec: value@1 <= 150, projection=[range_key@0] +05)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150) +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT CASE WHEN value <= 150 THEN range_key ELSE NULL END + FROM range_partitioned) +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +########## +# TEST 16: Compatible Range Join Repartitions to Increase Parallelism # Co-partitioning satisfaction does not prevent a repartition that increases # parallelism. With target_partitions larger than the Range partition count, # both sides are hash repartitioned. @@ -456,7 +635,7 @@ ORDER BY l.range_key; 35 350 350 ########## -# TEST 14: Preserve File Partitions Preserves Range Join Inputs +# TEST 17: Preserve File Partitions Preserves Range Join Inputs # preserve_file_partitions preserves compatible Range inputs for partitioned # joins even when target_partitions is higher than the input partition count. ########## @@ -496,7 +675,7 @@ statement ok set datafusion.optimizer.preserve_file_partitions = 0; ########## -# TEST 15: Nested Range Joins +# TEST 18: Nested Range Joins # Compatible Range partitioning is preserved through the lower join, allowing # the upper join to consume it without Hash repartitioning either input. ########## @@ -531,7 +710,7 @@ ORDER BY l.range_key; 35 350 350 350 ########## -# TEST 16: Range Aggregates Feed Range Join +# TEST 19: Range Aggregates Feed Range Join # Aggregates on range_key preserve reusable partitioning for the downstream # partitioned join. ########## @@ -586,7 +765,7 @@ ORDER BY l.range_key; 35 350 350 ########## -# TEST 17: Range Join Feeds Aggregate +# TEST 20: Range Join Feeds Aggregate # The join preserves compatible Range partitioning on range_key, allowing the # aggregate above it to avoid Hash repartitioning. ########## @@ -620,7 +799,7 @@ ORDER BY l.range_key; 35 700 ########## -# TEST 18: Right Join on Range Partition Column +# TEST 21: Right Join on Range Partition Column # Compatible Range inputs satisfy the join's partitioning requirements, so no # Hash repartitioning is inserted. The left filter keeps its Range partitioning # and the unmatched right rows above 150 are preserved. @@ -653,7 +832,7 @@ NULL 30 300 NULL 35 350 ########## -# TEST 19: Right Semi Join on Range Partition Column +# TEST 22: Right Semi Join on Range Partition Column # Compatible Range inputs avoid Hash repartitioning for RightSemi joins. # Only right rows with a match on the filtered left side are returned. ########## @@ -681,7 +860,7 @@ ORDER BY r.range_key; 15 150 ########## -# TEST 20: Right Anti Join on Range Partition Column +# TEST 23: Right Anti Join on Range Partition Column # Compatible Range inputs avoid Hash repartitioning for RightAnti joins. # Only right rows without a match on the filtered left side are returned. ########## @@ -709,7 +888,7 @@ ORDER BY r.range_key; 35 350 ########## -# TEST 21: Incompatible Range Right Join Repartitions +# TEST 24: Incompatible Range Right Join Repartitions # The split points of the two inputs differ, so the co-partitioned layout # requirement cannot be satisfied and Hash repartitioning repairs both sides # of the right join. Results stay correct on the repartitioned path. @@ -744,7 +923,7 @@ NULL 30 300 NULL 35 350 ########## -# TEST 22: Composite-Key Right Join Repartitions +# TEST 25: Composite-Key Right Join Repartitions # Range([range_key]) does not satisfy a partitioned join on # (range_key, non_range_key), so both sides repartition on the full key. ########## @@ -783,7 +962,7 @@ statement ok reset datafusion.optimizer.subset_repartition_threshold; ########## -# TEST 23: Right Join with Mismatched Range Partition Counts Repartitions +# TEST 26: Right Join with Mismatched Range Partition Counts Repartitions # Both inputs are range partitioned on range_key, but declare a different number # of partitions (four vs three). The per-child key requirements can be satisfied # by Range, but the co-partitioned layout requirement cannot, so Hash @@ -818,7 +997,7 @@ ORDER BY r.range_key; 350 35 350 ########## -# TEST 24: Right Join on Non-Range Key Repartitions +# TEST 27: Right Join on Non-Range Key Repartitions # Both inputs expose Range([range_key]), but the join key is non_range_key. # Range([range_key]) does not satisfy KeyPartitioned([non_range_key]), so # planning inserts Hash repartitioning on the actual join key for the right join. @@ -853,7 +1032,7 @@ ORDER BY r.range_key; 50 35 350 ########## -# TEST 25: Mark Join Marker Semantics +# TEST 28: Mark Join Marker Semantics # Mark joins preserve matched, unmatched, and NULL-key marker behavior over # range-partitioned inputs. ########## @@ -867,11 +1046,9 @@ WHERE r.non_range_key = 2 OR r.range_key IN ( physical_plan 01)FilterExec: non_range_key@1 = 2 OR mark@3, projection=[range_key@0, value@2] 02)--HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(range_key@0, range_key@0)] -03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 -04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false -05)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 -06)------FilterExec: value@1 <= 150, projection=[range_key@0] -07)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----FilterExec: value@1 <= 150, projection=[range_key@0] +05)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false # Matched rows have mark=true and are returned; unmatched rows have # mark=false and are only returned when non_range_key = 2. @@ -907,7 +1084,7 @@ ORDER BY r.range_key; 35 350 ########## -# TEST 26: Sort Merge Join Avoids Repartition for Compatible Range Inputs +# TEST 29: Sort Merge Join Avoids Repartition for Compatible Range Inputs # Compatible Range inputs satisfy SortMergeJoinExec's co-partitioned # KeyPartitioned requirements. ########## @@ -944,7 +1121,7 @@ ORDER BY l.range_key; 35 350 350 ########## -# TEST 27: Sort Merge Join Repartitions Incompatible Range Inputs +# TEST 30: Sort Merge Join Repartitions Incompatible Range Inputs # Different Range split points do not satisfy SortMergeJoinExec's # co-partitioned KeyPartitioned requirements. ########## @@ -983,7 +1160,7 @@ statement ok reset datafusion.optimizer.prefer_hash_join; ########## -# TEST 28: Symmetric Hash Join Avoids Repartition for Compatible Range Inputs +# TEST 31: Symmetric Hash Join Avoids Repartition for Compatible Range Inputs # Compatible Range streams satisfy SymmetricHashJoinExec's co-partitioned # KeyPartitioned requirements. ########## @@ -1017,7 +1194,7 @@ FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; 5 50 50 ########## -# TEST 29: Symmetric Hash Join Repartitions Incompatible Range Inputs +# TEST 32: Symmetric Hash Join Repartitions Incompatible Range Inputs # Different Range split points do not satisfy SymmetricHashJoinExec's # co-partitioned KeyPartitioned requirements. ########## @@ -1050,7 +1227,7 @@ FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; 5 50 50 ########## -# TEST 30: Full Outer Join on Range Partition Column +# TEST 33: Full Outer Join on Range Partition Column # Full partitioned hash joins also opt in to Range satisfying KeyPartitioned # requirements, so compatible Range layouts avoid Hash repartitioning here too. ########## @@ -1081,7 +1258,7 @@ ORDER BY l.range_key; 35 350 350 ########## -# TEST 31: Full Outer Join Incompatible Range Repartitions +# TEST 34: Full Outer Join Incompatible Range Repartitions # Same as TEST 10, but for Full: differing split points between the two # Range-partitioned inputs still require Hash repartitioning to co-partition. ########## @@ -1114,7 +1291,7 @@ ORDER BY l.range_key; 35 350 350 ########## -# TEST 32: Full Outer Join Produces Matched and Unmatched Rows +# TEST 35: Full Outer Join Produces Matched and Unmatched Rows # `range_partitioned` and `range_partitioned_sparse` share the same Range # split points/partition count but only partially overlapping range_key # values, so this exercises matched rows, left-only unmatched rows (NULLs on @@ -1159,7 +1336,7 @@ statement ok reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 33: Union of Range Partitioned Inputs +# TEST 36: Union of Range Partitioned Inputs # Each input exposes the same Range partitioning on range_key, so the optimizer # converts UnionExec to InterleaveExec to avoid redundant repartitioning. ########## @@ -1208,7 +1385,7 @@ set datafusion.optimizer.preserve_file_partitions = 0; ########## -# TEST 34: Window on Range Partition Column +# TEST 37: Window on Range Partition Column # Range([range_key]) colocates equal range_key values, so # PARTITION BY range_key is satisfied without a hash repartition. ########## @@ -1236,7 +1413,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM r ########## -# TEST 35: Unbounded-Frame Window on Range Partition Column +# TEST 38: Unbounded-Frame Window on Range Partition Column # The unbounded frame makes DataFusion use WindowAggExec instead of # BoundedWindowAggExec, which likewise reuses Range partitioning without a # hash repartition. @@ -1265,7 +1442,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BE ########## -# TEST 36: Window on Non-Range Column Rehashes +# TEST 39: Window on Non-Range Column Rehashes # Range([range_key]) does not colocate non_range_key values, so # PARTITION BY non_range_key still requires a hash repartition. ########## @@ -1294,7 +1471,7 @@ SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER B ########## -# TEST 37: Unbounded-Frame Window on Non-Range Column Rehashes +# TEST 40: Unbounded-Frame Window on Non-Range Column Rehashes # The unbounded frame makes DataFusion use WindowAggExec; Range([range_key]) # does not colocate non_range_key values, so PARTITION BY non_range_key # still requires a hash repartition. @@ -1324,7 +1501,7 @@ SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER B ########## -# TEST 38: Window Subset Satisfaction on Range Partition Column +# TEST 41: Window Subset Satisfaction on Range Partition Column # With the subset threshold met, Range([range_key]) satisfies # PARTITION BY (range_key, non_range_key): equal composite keys share the # same range_key, so they are already colocated. @@ -1356,7 +1533,7 @@ SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER B ########## -# TEST 39: Window Subset Rehashes Below Subset Threshold +# TEST 42: Window Subset Rehashes Below Subset Threshold # Range([range_key]) is only a subset of PARTITION BY # (range_key, non_range_key), so it should not satisfy the window key when # subset satisfaction is disabled. @@ -1395,7 +1572,7 @@ reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 40: Window Without Partition Keys Uses a Single Partition +# TEST 43: Window Without Partition Keys Uses a Single Partition # A window with no PARTITION BY requires a single partition; range # partitioning is not applicable. ########## @@ -1425,7 +1602,7 @@ SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER ########## -# TEST 41: PartitionedTopK on Range Partition Column +# TEST 44: PartitionedTopK on Range Partition Column # Exact Range([range_key]) satisfies the TopK partition key and avoids repartitioning. ########## @@ -1468,7 +1645,7 @@ ORDER BY range_key; ########## -# TEST 42: PartitionedTopK on Non-Range Column +# TEST 45: PartitionedTopK on Non-Range Column # Partitioning on a non-range key cannot reuse Range([range_key]) and # requires hash repartitioning. ########## @@ -1504,7 +1681,7 @@ ORDER BY non_range_key; ########## -# TEST 43: PartitionedTopK Reuses Range Subset Partitioning +# TEST 46: PartitionedTopK Reuses Range Subset Partitioning # With subset threshold met and preserve-file disabled, Range([range_key]) # satisfies partitioning by (range_key, non_range_key). ########## @@ -1545,7 +1722,7 @@ ORDER BY range_key, non_range_key; ########## -# TEST 44: Range Subset PartitionedTopK Rehashes Below Subset Threshold +# TEST 47: Range Subset PartitionedTopK Rehashes Below Subset Threshold # Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key), # so it should not satisfy the TopK partition key when subset satisfaction is # disabled. @@ -1583,7 +1760,7 @@ statement ok reset datafusion.optimizer.enable_window_topn; ########## -# TEST 45: Subset of Inputs Compatible Does Not Trigger InterleaveExec +# TEST 48: Subset of Inputs Compatible Does Not Trigger InterleaveExec # In a three-way union, two inputs share the same Range split points [10,20,30] # while the third has a partially-overlapping but different set [15,20,30]. # can_interleave requires ALL inputs to match, so UnionExec is kept. @@ -1639,7 +1816,7 @@ ORDER BY range_key, value; 35 350 ########## -# TEST 46: Incompatible Range Split Points Falls Back to UnionExec +# TEST 49: Incompatible Range Split Points Falls Back to UnionExec # Two range-partitioned inputs with different split points cannot be interleaved, # so the optimizer keeps UnionExec instead of converting to InterleaveExec. ########## @@ -1681,7 +1858,7 @@ ORDER BY range_key, value; 35 350 ########## -# TEST 47: InterleaveExec Propagates Range Partitioning to Aggregate +# TEST 50: InterleaveExec Propagates Range Partitioning to Aggregate # InterleaveExec outputs the same Range partitioning as its compatible inputs, # allowing a downstream aggregate on range_key to run SinglePartitioned without # a Hash repartition. From a785b2f16c5be3fe3bba3e0cd54e8d5a6fa9b273 Mon Sep 17 00:00:00 2001 From: Saad Tajwar <59696464+saadtajwar@users.noreply.github.com> Date: Wed, 22 Jul 2026 07:01:29 -0400 Subject: [PATCH 19/22] feat: Range Partitioning FFI (#23520) ## Which issue does this PR close? - Closes #22394 ## Rationale for this change Exposing range partition metadata via the FFI for external consumers. ## What changes are included in this PR? - Added FFI mirror struct for `RangePartitioning` and added new enum variant for range in `FFI_Partitioning` - For native -> FFI, added match arm for the new variant, same with FFI -> native but changed the approach of `From` -> `TryFrom` to utilize the validation for `RangePartitioning` and modified `plan_properties` to match - Added tests ## Are these changes tested? Yes ## Are there any user-facing changes? Yes, exposing Range partitioning over FFI. This exposes a new `Range` variant in the `FFI_Partitioning` enum, which may cause consumers of this enum to add another arm to match statements to handle the new enum. New `FFI_RangePartitioning` struct for the `Range` variant. --------- Co-authored-by: Tim Saucer --- .../ffi/src/physical_expr/partitioning.rs | 157 ++++++++++++++++-- datafusion/ffi/src/plan_properties.rs | 46 ++++- 2 files changed, 183 insertions(+), 20 deletions(-) diff --git a/datafusion/ffi/src/physical_expr/partitioning.rs b/datafusion/ffi/src/physical_expr/partitioning.rs index eec437639e156..2a9a8528c6c3e 100644 --- a/datafusion/ffi/src/physical_expr/partitioning.rs +++ b/datafusion/ffi/src/physical_expr/partitioning.rs @@ -17,20 +17,35 @@ use std::sync::Arc; -use datafusion_physical_expr::Partitioning; +use datafusion_common::{DataFusionError, ScalarValue, SplitPoint}; +use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, +}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use stabby::vec::Vec as SVec; +use crate::arrow_wrappers::WrappedArray; use crate::physical_expr::FFI_PhysicalExpr; +use crate::physical_expr::sort::FFI_PhysicalSortExpr; + +/// A stable struct for sharing [`RangePartitioning`] across FFI boundaries. +/// See [`RangePartitioning`] for the descriptions of each field. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_RangePartitioning { + split_points: SVec>, + ordering: SVec, +} /// A stable struct for sharing [`Partitioning`] across FFI boundaries. -/// See ['Partitioning'] for the meaning of each variant. +/// See [`Partitioning`] for the meaning of each variant. #[repr(C)] #[derive(Debug)] pub enum FFI_Partitioning { RoundRobinBatch(usize), Hash(SVec, usize), UnknownPartitioning(usize), + Range(FFI_RangePartitioning), } impl From<&Partitioning> for FFI_Partitioning { @@ -45,49 +60,130 @@ 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()) + // Producer-side conversion should be infallible at ABI boundary + let split_points = range + .split_points() + .iter() + .map(|split_point| { + split_point + .values() + .iter() + .map(|value| { + WrappedArray::try_from(value).expect( + "ScalarValue in RangePartitioning should convert to WrappedArray", + ) + }) + .collect() + }) + .collect(); + let ordering = range + .ordering() + .iter() + .map(FFI_PhysicalSortExpr::from) + .collect(); + Self::Range(FFI_RangePartitioning { + split_points, + ordering, + }) } Partitioning::UnknownPartitioning(size) => Self::UnknownPartitioning(*size), } } } -impl From<&FFI_Partitioning> for Partitioning { - fn from(value: &FFI_Partitioning) -> Self { - match value { +impl TryFrom for Partitioning { + type Error = DataFusionError; + + fn try_from(value: FFI_Partitioning) -> Result { + Ok(match value { FFI_Partitioning::RoundRobinBatch(size) => { - Partitioning::RoundRobinBatch(*size) + Partitioning::RoundRobinBatch(size) } FFI_Partitioning::Hash(exprs, size) => { let exprs = exprs.iter().map(>::from).collect(); - Self::Hash(exprs, *size) + Self::Hash(exprs, size) + } + FFI_Partitioning::Range(range) => { + let split_points = range + .split_points + .into_iter() + .map(|split_point| { + split_point + .into_iter() + .map(ScalarValue::try_from) + .collect::, _>>() + .map(SplitPoint::new) + }) + .collect::, _>>()?; + + let ordering = + LexOrdering::new(range.ordering.iter().map(PhysicalSortExpr::from)) + .ok_or_else(|| { + DataFusionError::Internal( + "FFI Range partitioning ordering must be non-empty" + .to_string(), + ) + })?; + + Self::Range(RangePartitioning::try_new(ordering, split_points)?) } FFI_Partitioning::UnknownPartitioning(size) => { - Self::UnknownPartitioning(*size) + Self::UnknownPartitioning(size) } - } + }) } } #[cfg(test)] mod tests { - use datafusion_physical_expr::Partitioning; - use datafusion_physical_expr::expressions::lit; + use std::sync::Arc; + + use arrow_schema::SortOptions; + use datafusion_common::{Result, ScalarValue, SplitPoint}; + use datafusion_physical_expr::expressions::{Column, lit}; + use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, + }; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use stabby::vec::Vec as SVec; - use crate::physical_expr::partitioning::FFI_Partitioning; + use crate::physical_expr::partitioning::{FFI_Partitioning, FFI_RangePartitioning}; + + fn range_partitioning() -> Result { + let a = Arc::new(Column::new("a", 0)) as Arc; + let b = Arc::new(Column::new("b", 1)) as Arc; + let ordering = LexOrdering::new([ + PhysicalSortExpr::new(a, SortOptions::default()), + PhysicalSortExpr::new(b, SortOptions::new(true, false)), + ]) + .expect("non-empty ordering"); + let split_points = vec![ + SplitPoint::new(vec![ + ScalarValue::Int64(Some(10)), + ScalarValue::Utf8(Some("a".to_string())), + ]), + SplitPoint::new(vec![ + ScalarValue::Int64(Some(20)), + ScalarValue::Utf8(Some("b".to_string())), + ]), + ]; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) + } #[test] - fn round_trip_ffi_partitioning() { + fn round_trip_ffi_partitioning() -> Result<()> { for partitioning in [ Partitioning::RoundRobinBatch(10), Partitioning::Hash(vec![lit(1)], 10), Partitioning::UnknownPartitioning(10), + range_partitioning()?, ] { let ffi_partitioning: FFI_Partitioning = (&partitioning).into(); - let returned: Partitioning = (&ffi_partitioning).into(); + let returned: Partitioning = ffi_partitioning.try_into()?; if let Partitioning::UnknownPartitioning(return_size) = returned { let Partitioning::UnknownPartitioning(original_size) = partitioning @@ -99,5 +195,32 @@ mod tests { assert_eq!(partitioning, returned); } } + + Ok(()) + } + + #[test] + fn round_trip_ffi_range_partitioning_compound_key() -> Result<()> { + let partitioning = range_partitioning()?; + + let ffi_partitioning: FFI_Partitioning = (&partitioning).into(); + let returned: Partitioning = ffi_partitioning.try_into()?; + assert_eq!(partitioning, returned); + + Ok(()) + } + + #[test] + fn ffi_range_partitioning_rejects_empty_ordering() { + let ffi_partitioning = FFI_Partitioning::Range(FFI_RangePartitioning { + split_points: SVec::new(), + ordering: SVec::new(), + }); + + let err = Partitioning::try_from(ffi_partitioning).unwrap_err(); + assert!( + err.to_string().contains("ordering must be non-empty"), + "{err}" + ); } } diff --git a/datafusion/ffi/src/plan_properties.rs b/datafusion/ffi/src/plan_properties.rs index b286ee2d7d30c..09ef26af32349 100644 --- a/datafusion/ffi/src/plan_properties.rs +++ b/datafusion/ffi/src/plan_properties.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use datafusion_common::error::{DataFusionError, Result}; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::PlanProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -172,6 +172,7 @@ impl TryFrom for PlanProperties { .unwrap_or_default(); let partitioning = unsafe { (ffi_props.output_partitioning)(&ffi_props) }; + let partitioning = Partitioning::try_from(partitioning)?; let eq_properties = if sort_exprs.is_empty() { EquivalenceProperties::new(Arc::new(schema)) @@ -187,7 +188,7 @@ impl TryFrom for PlanProperties { Ok(PlanProperties::new( eq_properties, - (&partitioning).into(), + partitioning, emission_type, boundedness, )) @@ -260,13 +261,15 @@ impl From for EmissionType { #[cfg(test)] mod tests { + use arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_expr::PhysicalSortExpr; use datafusion::physical_plan::Partitioning; + use datafusion_common::{ScalarValue, SplitPoint}; + use datafusion_physical_expr::{LexOrdering, RangePartitioning}; use super::*; fn create_test_props() -> Result { - use arrow::datatypes::{DataType, Field, Schema}; let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); @@ -282,6 +285,25 @@ mod tests { )) } + fn create_range_test_props() -> Result { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + 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![ + SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), + ]; + let range = RangePartitioning::try_new(ordering, split_points)?; + + Ok(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::Range(range), + EmissionType::Incremental, + Boundedness::Bounded, + )) + } + #[test] fn test_round_trip_ffi_plan_properties() -> Result<()> { let original_props = create_test_props()?; @@ -314,4 +336,22 @@ mod tests { Ok(()) } + + #[test] + fn test_round_trip_ffi_plan_properties_range_partitioning() -> Result<()> { + let original_props = create_range_test_props()?; + + let mut local_props_ptr = FFI_PlanProperties::from(&original_props); + local_props_ptr.library_marker_id = crate::mock_foreign_marker_id; + + let foreign_props: PlanProperties = local_props_ptr.try_into()?; + + assert_eq!( + format!("{:?}", foreign_props.output_partitioning()), + format!("{:?}", original_props.output_partitioning()) + ); + assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}")); + + Ok(()) + } } From 7b01c11f029da2e6de69e6e40db9558858d766b3 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Thu, 23 Jul 2026 03:20:07 -0400 Subject: [PATCH 20/22] allow range to satisfy key distribution generally (#23680) - Closes #23266. - Part of #22395. `Partitioning::Range` now satisfies `Distribution::KeyPartitioned` privately across all operators that require it: aggregates, windows, TopK, and co-partitioned joins. So now the temporary operator opt-ins can be consolidated. - Allow compatible `Partitioning::Range` to satisfy `Distribution::KeyPartitioned` via `Partitioning::satisfaction` - Remove the temporary range-satisfaction helpers and operator-specific opt-ins Yes No, this should not change any exisitng behavior just consolidation --- datafusion/physical-expr/src/partitioning.rs | 454 ++++++------------ .../src/enforce_distribution.rs | 37 +- datafusion/physical-optimizer/src/utils.rs | 59 +-- .../physical-plan/src/aggregates/mod.rs | 12 +- .../src/distribution_requirements.rs | 165 +------ .../physical-plan/src/joins/hash_join/exec.rs | 10 +- .../src/joins/sort_merge_join/exec.rs | 1 - .../src/joins/symmetric_hash_join.rs | 1 - .../src/sorts/partitioned_topk.rs | 1 - .../src/windows/bounded_window_agg_exec.rs | 1 - .../src/windows/window_agg_exec.rs | 1 - 11 files changed, 178 insertions(+), 564 deletions(-) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 75002a67e614e..886d47ff403c5 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -381,47 +381,6 @@ impl Partitioning { }) } - fn key_expr_satisfaction( - partition_exprs: &[Arc], - required_exprs: &[Arc], - eq_properties: &EquivalenceProperties, - allow_subset: bool, - ) -> PartitioningSatisfaction { - if partition_exprs.is_empty() || required_exprs.is_empty() { - return PartitioningSatisfaction::NotSatisfied; - } - - if equivalent_exprs(required_exprs, partition_exprs, eq_properties) { - return PartitioningSatisfaction::Exact; - } - - if !allow_subset { - return PartitioningSatisfaction::NotSatisfied; - } - - let eq_groups = eq_properties.eq_group(); - if eq_groups.is_empty() { - if Self::is_subset_partitioning(partition_exprs, required_exprs) { - PartitioningSatisfaction::Subset - } else { - PartitioningSatisfaction::NotSatisfied - } - } else { - let normalized_partition_exprs = - normalize_exprs(partition_exprs, eq_properties); - let normalized_required_exprs = - normalize_exprs(required_exprs, eq_properties); - if Self::is_subset_partitioning( - &normalized_partition_exprs, - &normalized_required_exprs, - ) { - PartitioningSatisfaction::Subset - } else { - PartitioningSatisfaction::NotSatisfied - } - } - } - #[deprecated(since = "52.0.0", note = "Use satisfaction instead")] pub fn satisfy( &self, @@ -455,8 +414,12 @@ impl Partitioning { { PartitioningSatisfaction::Exact } - Distribution::KeyPartitioned(required_exprs) => match self { - Partitioning::Hash(partition_exprs, _) => Self::key_expr_satisfaction( + Distribution::HashPartitioned(required_exprs) + | Distribution::KeyPartitioned(required_exprs) => match self { + // Here we do not check the partition count for hash partitioning and assumes the partition count + // and hash functions in the system are the same. In future if we plan to support storage partition-wise joins, + // then we need to have the partition count and hash functions validation. + Partitioning::Hash(partition_exprs, _) => Self::key_satisfaction( partition_exprs, required_exprs, eq_properties, @@ -464,11 +427,11 @@ impl Partitioning { ), Partitioning::Range(range) => { let partition_exprs = range - .ordering + .ordering() .iter() .map(|sort_expr| Arc::clone(&sort_expr.expr)) .collect::>(); - Self::key_expr_satisfaction( + Self::key_satisfaction( &partition_exprs, required_exprs, eq_properties, @@ -480,26 +443,47 @@ impl Partitioning { PartitioningSatisfaction::NotSatisfied } }, - Distribution::HashPartitioned(required_exprs) => match self { - // Here we do not check the partition count for hash partitioning and assumes the partition count - // and hash functions in the system are the same. In future if we plan to support storage partition-wise joins, - // then we need to have the partition count and hash functions validation. - Partitioning::Hash(partition_exprs, _) => Self::key_expr_satisfaction( - partition_exprs, - required_exprs, - eq_properties, - allow_subset, - ), - Partitioning::RoundRobinBatch(_) - | Partitioning::UnknownPartitioning(_) => { - PartitioningSatisfaction::NotSatisfied - } - Partitioning::Range(_) => PartitioningSatisfaction::NotSatisfied, - }, Distribution::SinglePartition => PartitioningSatisfaction::NotSatisfied, } } + fn key_satisfaction( + partition_exprs: &[Arc], + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, + ) -> PartitioningSatisfaction { + if partition_exprs.is_empty() || required_exprs.is_empty() { + return PartitioningSatisfaction::NotSatisfied; + } + + if equivalent_exprs(required_exprs, partition_exprs, eq_properties) { + return PartitioningSatisfaction::Exact; + } + + let eq_groups = eq_properties.eq_group(); + if !eq_groups.is_empty() { + if allow_subset { + let normalized_partition_exprs = + normalize_exprs(partition_exprs, eq_properties); + let normalized_required_exprs = + normalize_exprs(required_exprs, eq_properties); + if Self::is_subset_partitioning( + &normalized_partition_exprs, + &normalized_required_exprs, + ) { + return PartitioningSatisfaction::Subset; + } + } + } else if allow_subset + && Self::is_subset_partitioning(partition_exprs, required_exprs) + { + return PartitioningSatisfaction::Subset; + } + + PartitioningSatisfaction::NotSatisfied + } + /// Calculate the output partitioning after applying the given projection. pub fn project( &self, @@ -752,6 +736,26 @@ mod tests { } } + fn assert_satisfaction( + desc: &str, + partitioning: &Partitioning, + required: &Distribution, + eq_properties: &EquivalenceProperties, + expected_with_subset: PartitioningSatisfaction, + expected_without_subset: PartitioningSatisfaction, + ) { + assert_eq!( + partitioning.satisfaction(required, eq_properties, true), + expected_with_subset, + "Failed for {desc} with subset enabled" + ); + assert_eq!( + partitioning.satisfaction(required, eq_properties, false), + expected_without_subset, + "Failed for {desc} with subset disabled" + ); + } + #[test] #[expect( deprecated, @@ -835,294 +839,105 @@ mod tests { } #[test] - fn test_partitioning_satisfy_by_subset() -> Result<()> { + fn hash_partitioning_key_distribution_satisfaction() -> Result<()> { let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; + let unknown: Arc = Arc::new(UnKnownColumn::new("dropped")); let test_cases = vec![ ( - "KeyPartitioned([a, b]) satisfied by Hash([a])", - fixture.hash_partitioning([0], 4), + "exact: KeyPartitioned([a, b]) satisfied by Hash([a, b])", + fixture.hash_partitioning([0, 1], 4), fixture.key_distribution([0, 1]), - PartitioningSatisfaction::Subset, - PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, ), ( - "KeyPartitioned([a, b, c]) satisfied by Hash([a])", + "subset: KeyPartitioned([a, b]) satisfied by Hash([a])", fixture.hash_partitioning([0], 4), - fixture.key_distribution([0, 1, 2]), - PartitioningSatisfaction::Subset, - PartitioningSatisfaction::NotSatisfied, - ), - ( - "KeyPartitioned([a, b, c]) satisfied by Hash([a, b])", - fixture.hash_partitioning([0, 1], 4), - fixture.key_distribution([0, 1, 2]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([a, b, c]) satisfied by Hash([b])", + "subset: KeyPartitioned([a, b, c]) satisfied by Hash([b])", fixture.hash_partitioning([1], 4), fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([a, b, c]) satisfied by Hash([b, a])", + "subset reordered: KeyPartitioned([a, b, c]) satisfied by Hash([b, a])", fixture.hash_partitioning([1, 0], 4), fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - 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, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_current_superset() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; - - let test_cases = vec![ ( - "KeyPartitioned([a]) satisfied by Hash([a, b])", + "superset: KeyPartitioned([a]) not satisfied by Hash([a, b])", fixture.hash_partitioning([0, 1], 4), fixture.key_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([a]) satisfied by Hash([a, b, c])", + "superset: KeyPartitioned([a, b]) not satisfied by Hash([a, b, c])", fixture.hash_partitioning([0, 1, 2], 4), - fixture.key_distribution([0]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([a, b]) satisfied by Hash([a, b, c])", - fixture.hash_partitioning([0, 1, 2], 4), + "partial overlap: KeyPartitioned([a, b]) not satisfied by Hash([a, c])", + fixture.hash_partitioning([0, 2], 4), fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - 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, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_partial_overlap() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; - - let test_cases = vec![( - "Partial overlap: KeyPartitioned([a, b]) satisfied by Hash([a, c])", - fixture.hash_partitioning([0, 2], 4), - fixture.key_distribution([0, 1]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - )]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - 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, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_no_overlap() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; - - let test_cases = vec![ ( - "KeyPartitioned([b, c]) satisfied by Hash([a])", + "no overlap: KeyPartitioned([b, c]) not satisfied by Hash([a])", fixture.hash_partitioning([0], 4), fixture.key_distribution([1, 2]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([c]) satisfied by Hash([a, b])", - fixture.hash_partitioning([0, 1], 4), - fixture.key_distribution([2]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - 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, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_exact_match() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - - let test_cases = vec![ - ( - "KeyPartitioned([a, b]) satisfied by Hash([a, b])", - fixture.hash_partitioning([0, 1], 4), - fixture.key_distribution([0, 1]), - PartitioningSatisfaction::Exact, - PartitioningSatisfaction::Exact, - ), - ( - "KeyPartitioned([a]) satisfied by Hash([a])", - fixture.hash_partitioning([0], 4), - fixture.key_distribution([0]), - PartitioningSatisfaction::Exact, - PartitioningSatisfaction::Exact, - ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - 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, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_unknown() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let unknown: Arc = Arc::new(UnKnownColumn::new("dropped")); - - let test_cases = vec![ - ( - "KeyPartitioned([a, b]) satisfied by Hash([unknown])", + "unknown partition expr", Partitioning::Hash(vec![Arc::clone(&unknown)], 4), fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([unknown]) satisfied by Hash([a, b])", + "unknown required expr", fixture.hash_partitioning([0, 1], 4), Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([unknown]) satisfied by Hash([unknown])", + "same unknown expr", Partitioning::Hash(vec![Arc::clone(&unknown)], 4), Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - 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, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_empty_hash() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a"])?; - - let test_cases = vec![ ( - "KeyPartitioned([a]) satisfied by Hash([])", - Partitioning::Hash(vec![], 4), - fixture.key_distribution([0]), + "unknown partition expr is not a valid subset", + Partitioning::Hash(vec![Arc::clone(&unknown)], 4), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown), fixture.col(0)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([]) satisfied by Hash([a])", - fixture.hash_partitioning([0], 4), - Distribution::KeyPartitioned(vec![]), + "empty hash partitioning", + Partitioning::Hash(vec![], 4), + fixture.key_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([]) satisfied by Hash([])", - Partitioning::Hash(vec![], 4), + "empty key distribution", + fixture.hash_partitioning([0], 4), Distribution::KeyPartitioned(vec![]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, @@ -1132,16 +947,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - 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, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" + assert_satisfaction( + desc, + &partition, + &required, + &fixture.eq_properties, + expected_with_subset, + expected_without_subset, ); } @@ -1385,19 +1197,65 @@ mod tests { } #[test] - #[expect( - deprecated, - reason = "test intentionally covers the hash-specific requirement" - )] - fn test_multi_partition_range_does_not_satisfy_hash_distribution() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let range_partitioning = + fn range_partitioning_key_distribution_satisfaction() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; + let range_a = fixture.range_partitioning([0], vec![int_split_point([10])]); + let range_ab = fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); - let required = Distribution::HashPartitioned(fixture.cols([0, 1])); - assert_eq!( - range_partitioning.satisfaction(&required, &fixture.eq_properties, false), - PartitioningSatisfaction::NotSatisfied + assert_satisfaction( + "exact single key", + &range_a, + &fixture.key_distribution([0]), + &fixture.eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); + assert_satisfaction( + "exact compound key", + &range_ab, + &fixture.key_distribution([0, 1]), + &fixture.eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); + assert_satisfaction( + "subset key", + &range_a, + &fixture.key_distribution([0, 1]), + &fixture.eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + assert_satisfaction( + "incompatible key", + &range_a, + &fixture.key_distribution([1]), + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let mut eq_properties = fixture.eq_properties.clone(); + eq_properties.add_equal_conditions(fixture.col(0), fixture.col(2))?; + assert_satisfaction( + "equivalent subset key", + &range_a, + &fixture.key_distribution([1, 2]), + &eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + + let mut eq_properties = fixture.eq_properties.clone(); + eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?; + assert_satisfaction( + "equivalent exact key", + &range_a, + &fixture.key_distribution([1]), + &eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, ); Ok(()) diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index d8f9fc880861a..3b520fd3de60c 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -29,7 +29,7 @@ use crate::optimizer::PhysicalOptimizerRule; use crate::output_requirements::OutputRequirementExec; use crate::utils::{ add_sort_above_with_check, is_coalesce_partitions, is_repartition, - is_sort_preserving_merge, range_partitioning_satisfies_key_partitioning, + is_sort_preserving_merge, }; use arrow::compute::SortOptions; @@ -857,18 +857,13 @@ fn add_roundrobin_on_top( } } -// TODO: remove this temporary bridge once [`Partitioning::Range`] -// generally satisfies [`Distribution::KeyPartitioned`] through -// [`Partitioning::satisfaction`]. -// . -// -// Partial aggregates do not require key partitioning, but they preserve their -// input partitioning for the final aggregate. Until Range satisfies -// KeyPartitioned generally, this check keeps preserve_file_partitions from -// inserting RoundRobin between a reusable Range input and the partial aggregate. -fn partial_aggregate_preserves_reusable_partitioning( +// Partial aggregates require unspecified input distribution, but their output +// may already satisfy the final aggregate's key distribution because partial +// aggregation preserves/projects input partitioning. Keep that reusable output +// partitioning intact when preserve_file_partitions would otherwise insert +// RoundRobin below the partial aggregate. +fn partial_aggregate_output_satisfies_final_partitioning( plan: &Arc, - child: &Arc, allow_subset_satisfy_partitioning: bool, ) -> bool { let Some(aggregate) = plan.downcast_ref::() else { @@ -881,24 +876,15 @@ fn partial_aggregate_preserves_reusable_partitioning( return false; } - let group_exprs = aggregate.group_expr().input_exprs(); - let output_partitioning = child.output_partitioning(); - let eq_properties = child.equivalence_properties(); - let key_distribution = Distribution::KeyPartitioned(group_exprs.clone()); + let key_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr()); - output_partitioning + plan.output_partitioning() .satisfaction( &key_distribution, - eq_properties, + plan.equivalence_properties(), allow_subset_satisfy_partitioning, ) .is_satisfied() - || range_partitioning_satisfies_key_partitioning( - output_partitioning, - &group_exprs, - eq_properties, - allow_subset_satisfy_partitioning, - ) } /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator @@ -1424,9 +1410,8 @@ pub fn ensure_distribution( let preserve_partial_aggregate_partitioning = preserve_file_partition_threshold_met - && partial_aggregate_preserves_reusable_partitioning( + && partial_aggregate_output_satisfies_final_partitioning( &plan, - &child.plan, allow_subset_satisfy_partitioning, ); diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index 36e630c1f4ae3..a6b01637c970e 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -18,10 +18,7 @@ use std::sync::Arc; use datafusion_common::Result; -use datafusion_physical_expr::{ - EquivalenceProperties, LexOrdering, LexRequirement, Partitioning, PhysicalExpr, - physical_exprs_equal, -}; +use datafusion_physical_expr::{LexOrdering, LexRequirement}; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -111,60 +108,6 @@ pub fn is_repartition(plan: &Arc) -> bool { plan.is::() } -/// TODO: remove once Range generally satisfies KeyPartitioned requirements -/// through Partitioning::satisfaction. -/// See . -/// -/// Checks whether range partitioning satisfies a key partitioning requirement. -/// This is intentionally separate from general partitioning satisfaction while -/// range reuse is rolled out operator by operator. -pub(crate) fn range_partitioning_satisfies_key_partitioning( - partitioning: &Partitioning, - required_exprs: &[Arc], - eq_properties: &EquivalenceProperties, - allow_subset: bool, -) -> bool { - match partitioning { - Partitioning::Range(range) => { - let partition_exprs = range - .ordering() - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - - if partition_exprs.is_empty() || required_exprs.is_empty() { - return false; - } - - let eq_group = eq_properties.eq_group(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - let normalized_required_exprs = required_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - - if physical_exprs_equal( - &normalized_required_exprs, - &normalized_partition_exprs, - ) { - return true; - } - - allow_subset - && normalized_partition_exprs.len() < normalized_required_exprs.len() - && normalized_partition_exprs.iter().all(|partition_expr| { - normalized_required_exprs - .iter() - .any(|required_expr| partition_expr.eq(required_expr)) - }) - } - _ => false, - } -} - /// Checks whether the given operator is a limit; /// i.e. either a [`LocalLimitExec`] or a [`GlobalLimitExec`]. pub fn is_limit(plan: &Arc) -> bool { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 725dc2cbd64fd..61e88f0f75db3 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1548,7 +1548,7 @@ impl ExecutionPlan for AggregateExec { } fn input_distribution_requirements(&self) -> InputDistributionRequirements { - let requirements = InputDistributionRequirements::new(match &self.mode { + InputDistributionRequirements::new(match &self.mode { AggregateMode::Partial | AggregateMode::PartialReduce => { vec![Distribution::UnspecifiedDistribution] } @@ -1558,15 +1558,7 @@ impl ExecutionPlan for AggregateExec { AggregateMode::Final | AggregateMode::Single => { vec![Distribution::SinglePartition] } - }); - match &self.mode { - AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned - if !self.group_by.has_grouping_set() => - { - requirements.allow_range_satisfaction_for_key_partitioning() - } - _ => requirements, - } + }) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/distribution_requirements.rs b/datafusion/physical-plan/src/distribution_requirements.rs index 9c7a1336c06a3..6405b1f121ef7 100644 --- a/datafusion/physical-plan/src/distribution_requirements.rs +++ b/datafusion/physical-plan/src/distribution_requirements.rs @@ -17,13 +17,8 @@ //! Input distribution requirements for physical execution plans. -use std::sync::Arc; - use datafusion_common::{Result, internal_err}; -use datafusion_physical_expr::{ - Distribution, EquivalenceProperties, Partitioning, PartitioningSatisfaction, - PhysicalExpr, physical_exprs_equal, -}; +use datafusion_physical_expr::{Distribution, Partitioning, PartitioningSatisfaction}; use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel}; @@ -98,10 +93,7 @@ impl InputDistributionRequirements { pub fn new(per_child: Vec) -> Self { let children = per_child .into_iter() - .map(|distribution| ChildDistributionRequirement { - distribution, - satisfaction: InputDistributionSatisfaction::Default, - }) + .map(|distribution| ChildDistributionRequirement { distribution }) .collect(); Self { @@ -176,8 +168,7 @@ impl InputDistributionRequirements { ); }; - Ok(requirement.satisfaction.satisfaction( - child.output_partitioning(), + Ok(child.output_partitioning().satisfaction( &requirement.distribution, child.equivalence_properties(), options.allow_subset(), @@ -208,30 +199,6 @@ impl InputDistributionRequirements { Ok(co_partitioned.clone()) } - /// TODO: remove this temporary bridge once [`Partitioning::Range`] - /// generally satisfies [`Distribution::KeyPartitioned`] through - /// [`Partitioning::satisfaction`]. - /// . - /// - /// Also allow compatible [`Partitioning::Range`] to satisfy - /// [`Distribution::KeyPartitioned`]. - #[expect( - deprecated, - reason = "HashPartitioned is accepted during the KeyPartitioned migration" - )] - pub(crate) fn allow_range_satisfaction_for_key_partitioning(mut self) -> Self { - for child in &mut self.children { - if matches!( - child.distribution, - Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) - ) { - child.satisfaction = - InputDistributionSatisfaction::AllowRangeKeyPartitioning; - } - } - self - } - /// Validate the requirements against a plan's children. pub(crate) fn check_invariants( &self, @@ -301,10 +268,8 @@ impl InputDistributionRequirements { let first = children[first_idx]; let first_partitioning = first.output_partitioning(); - if !first_requirement - .satisfaction + if !first_partitioning .satisfaction( - first_partitioning, &first_requirement.distribution, first.equivalence_properties(), false, @@ -317,19 +282,16 @@ impl InputDistributionRequirements { for &child_idx in co_partitioned.iter().skip(1) { let requirement = &self.children[child_idx]; let child = children[child_idx]; - if !requirement - .satisfaction + if !child + .output_partitioning() .satisfaction( - child.output_partitioning(), &requirement.distribution, child.equivalence_properties(), false, ) .is_satisfied() || !compatible_co_partitioning_layout( - first_requirement, first_partitioning, - requirement, child.output_partitioning(), ) { @@ -345,60 +307,6 @@ impl InputDistributionRequirements { #[derive(Debug, Clone)] struct ChildDistributionRequirement { distribution: Distribution, - satisfaction: InputDistributionSatisfaction, -} - -/// TODO: remove this temporary bridge once [`Partitioning::Range`] -/// generally satisfies [`Distribution::KeyPartitioned`] through -/// [`Partitioning::satisfaction`]. -/// . -#[non_exhaustive] -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -enum InputDistributionSatisfaction { - /// Use [`Partitioning::satisfaction`] as-is. - #[default] - Default, - /// Also allow [`Partitioning::Range`] to satisfy - /// [`Distribution::KeyPartitioned`]. - AllowRangeKeyPartitioning, -} - -impl InputDistributionSatisfaction { - /// Returns how `partitioning` satisfies `requirement`. - #[expect( - deprecated, - reason = "HashPartitioned is accepted during the KeyPartitioned migration" - )] - fn satisfaction( - self, - partitioning: &Partitioning, - requirement: &Distribution, - eq_properties: &EquivalenceProperties, - allow_subset: bool, - ) -> PartitioningSatisfaction { - let satisfaction = - partitioning.satisfaction(requirement, eq_properties, allow_subset); - if satisfaction.is_satisfied() { - return satisfaction; - } - - if !matches!(self, Self::AllowRangeKeyPartitioning) { - return PartitioningSatisfaction::NotSatisfied; - } - - let (Distribution::HashPartitioned(required_exprs) - | Distribution::KeyPartitioned(required_exprs)) = requirement - else { - return PartitioningSatisfaction::NotSatisfied; - }; - - range_satisfies_key_partitioning( - partitioning, - required_exprs, - eq_properties, - allow_subset, - ) - } } fn validate_child_index( @@ -421,62 +329,8 @@ fn validate_child_index( Ok(()) } -/// TODO: remove this temporary bridge once [`Partitioning::Range`] -/// generally satisfies [`Distribution::KeyPartitioned`] through -/// [`Partitioning::satisfaction`]. -/// . -fn range_satisfies_key_partitioning( - partitioning: &Partitioning, - required_exprs: &[Arc], - eq_properties: &EquivalenceProperties, - allow_subset: bool, -) -> PartitioningSatisfaction { - let Partitioning::Range(range) = partitioning else { - return PartitioningSatisfaction::NotSatisfied; - }; - - let partition_exprs = range - .ordering() - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - - if partition_exprs.is_empty() || required_exprs.is_empty() { - return PartitioningSatisfaction::NotSatisfied; - } - - let eq_group = eq_properties.eq_group(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - let normalized_required_exprs = required_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - - if physical_exprs_equal(&normalized_required_exprs, &normalized_partition_exprs) { - return PartitioningSatisfaction::Exact; - } - - if allow_subset - && normalized_partition_exprs.len() < normalized_required_exprs.len() - && normalized_partition_exprs.iter().all(|partition_expr| { - normalized_required_exprs - .iter() - .any(|required_expr| partition_expr.eq(required_expr)) - }) - { - PartitioningSatisfaction::Subset - } else { - PartitioningSatisfaction::NotSatisfied - } -} - fn compatible_co_partitioning_layout( - first: &ChildDistributionRequirement, first_partitioning: &Partitioning, - other: &ChildDistributionRequirement, other_partitioning: &Partitioning, ) -> bool { if first_partitioning.partition_count() == 1 @@ -491,12 +345,7 @@ fn compatible_co_partitioning_layout( match (first_partitioning, other_partitioning) { (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, - (Partitioning::Range(left), Partitioning::Range(right)) - if first.satisfaction - == InputDistributionSatisfaction::AllowRangeKeyPartitioning - && other.satisfaction - == InputDistributionSatisfaction::AllowRangeKeyPartitioning => - { + (Partitioning::Range(left), Partitioning::Range(right)) => { left.split_points() == right.split_points() && left.ordering().len() == right.ordering().len() && left diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index c1d515c284ced..0ee918bb66e70 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1252,7 +1252,7 @@ impl ExecutionPlan for HashJoinExec { } fn input_distribution_requirements(&self) -> InputDistributionRequirements { - let requirements = match self.mode { + match self.mode { PartitionMode::Partitioned => { let (left_expr, right_expr) = self .on @@ -1272,14 +1272,6 @@ impl ExecutionPlan for HashJoinExec { Distribution::UnspecifiedDistribution, Distribution::UnspecifiedDistribution, ]), - }; - - if self.mode == PartitionMode::Partitioned { - // Compatible Range inputs co-locate equal join keys, which - // satisfies the co-partitioned requirement for hash joins. - requirements.allow_range_satisfaction_for_key_partitioning() - } else { - requirements } } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index d7975d67ac2a6..6a7064bf0417b 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -437,7 +437,6 @@ impl ExecutionPlan for SortMergeJoinExec { Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), ]) - .allow_range_satisfaction_for_key_partitioning() } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index e0ec6b4cdaaef..b798a166e9cc2 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -442,7 +442,6 @@ impl ExecutionPlan for SymmetricHashJoinExec { Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), ]) - .allow_range_satisfaction_for_key_partitioning() } StreamJoinPartitionMode::SinglePartition => { InputDistributionRequirements::new(vec![ diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 6a33aa909e648..17eb70ef12131 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -314,7 +314,6 @@ impl ExecutionPlan for PartitionedTopKExec { crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( partition_exprs, )]) - .allow_range_satisfaction_for_key_partitioning() } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 9e38902c520a3..1acda21196008 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -343,7 +343,6 @@ impl ExecutionPlan for BoundedWindowAggExec { InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( self.partition_keys(), )]) - .allow_range_satisfaction_for_key_partitioning() } } diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index a66e0588ceb3f..e909dd2017827 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -251,7 +251,6 @@ impl ExecutionPlan for WindowAggExec { InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( self.partition_keys(), )]) - .allow_range_satisfaction_for_key_partitioning() } } From 4aed15d078468eaf10ad00d9dc623d71a2fe9a99 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Fri, 24 Jul 2026 13:06:14 -0400 Subject: [PATCH 21/22] Adapt range partitioning tests for branch 54 --- .../enforce_distribution.rs | 28 +++++++++---------- .../tests/cases/roundtrip_logical_plan.rs | 6 ++-- .../test_files/range_partitioning.slt | 16 +++++------ .../sqllogictest/test_files/window_topn.slt | 2 +- 4 files changed, 26 insertions(+), 26 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 5ebc413d60d7b..54b80644a31a3 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1088,8 +1088,8 @@ fn range_right_mark_hash_join_reuses_range_partitioning() -> Result<()> { plan, @r" HashJoinExec: mode=Partitioned, join_type=RightMark, on=[(a@0, a@0)] - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) " ); @@ -1126,9 +1126,9 @@ fn range_right_semi_hash_join_rehashes_incompatible_sort_options() -> Result<()> @r" HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(a@0, a@0)] RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 - DataSourceExec: file_groups={2 groups: [[p0], [p1]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(20)], 2), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(20)], 2) RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 - DataSourceExec: file_groups={2 groups: [[p0], [p1]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 DESC], [(20)], 2), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 DESC], [(20)], 2) " ); @@ -1159,7 +1159,7 @@ fn range_window_reuses_range_partitioning() -> Result<()> { @r#" BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) "# ); @@ -1191,7 +1191,7 @@ fn range_window_rehashes_incompatible_range_partitioning() -> Result<()> { BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true] RepartitionExec: partitioning=Hash([b@1], 4), input_partitions=4 - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) "# ); @@ -1230,9 +1230,9 @@ fn range_full_hash_join_reuses_compatible_range_partitioning() -> Result<()> { plan, @r" HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) ProjectionExec: expr=[a@0 as a1, b@1 as b1] - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) " ); @@ -1272,10 +1272,10 @@ fn range_full_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> @r" HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) RepartitionExec: partitioning=Hash([a1@0], 4), input_partitions=4 ProjectionExec: expr=[a@0 as a1, b@1 as b1] - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4) " ); @@ -1308,8 +1308,8 @@ fn range_left_mark_hash_join_reuses_range_partitioning() -> Result<()> { plan, @r" HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(a@0, a@0)] - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) " ); @@ -1346,9 +1346,9 @@ fn range_left_anti_hash_join_rehashes_incompatible_null_options() -> Result<()> @r" HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(a@0, a@0)] RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 - DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC NULLS LAST], [(10), (20), (30)], 4), file_type=parquet + PartitionedTestExec: output_partitioning=Range([a@0 ASC NULLS LAST], [(10), (20), (30)], 4) " ); diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 3bb5b898677c7..555690d4db419 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -85,9 +85,9 @@ use datafusion_expr::logical_plan::{Extension, UserDefinedLogicalNodeCore}; use datafusion_expr::{ Accumulator, AggregateUDF, ColumnarValue, ExprFunctionExt, ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, - RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, Volatility, - WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, WindowUDF, - WindowUDFImpl, + PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, + Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, + WindowFunctionDefinition, WindowUDF, WindowUDFImpl, }; use datafusion_functions_aggregate::average::avg_udaf; use datafusion_functions_aggregate::expr_fn::{ diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index e53436c22e173..ec51888c24b1e 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -1624,8 +1624,8 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0], order=[value@1 DESC] -04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0], order=[value@1 DESC] +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query III SELECT * FROM ( @@ -1665,9 +1665,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[non_range_key@0 as non_range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[non_range_key@0], order=[value@1 DESC] +03)----PartitionedTopKExec: fetch=1, partition=[non_range_key@0], order=[value@1 DESC] 04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 -05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false query III SELECT * FROM ( @@ -1701,8 +1701,8 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] -04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query IIII SELECT * FROM ( @@ -1746,9 +1746,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] 04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 -05)--------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false statement ok reset datafusion.optimizer.subset_repartition_threshold; diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt index 09b52daa2ea79..2f885991e3497 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -655,7 +655,7 @@ physical_plan 01)ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] 02)--RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=1, maintains_sort_order=true 03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -04)------PartitionedTopKExec: fn=row_number, fetch=1, partition=[c1@0], order=[c2@1 DESC] +04)------PartitionedTopKExec: fetch=1, partition=[c1@0], order=[c2@1 DESC] 05)--------DataSourceExec: partitions=1, partition_sizes=[1] statement ok From b73fd3412456c9c1dbf269a22e214ae4da438307 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Sat, 1 Aug 2026 19:10:18 +0200 Subject: [PATCH 22/22] Fix KeyPartitioned co-partitioning docs --- datafusion/physical-expr/src/partitioning.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 886d47ff403c5..14065ab8853ad 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -557,8 +557,8 @@ pub enum Distribution { /// /// For multi-input operators, satisfaction alone is not enough: each input /// may satisfy its own key requirement while using incompatible partition - /// boundaries. Use [`Partitioning::co_partitioned_with`] before pairing - /// partitions by index. + /// boundaries. Such operators must separately require compatible + /// co-partitioning before pairing partitions by index. KeyPartitioned(Vec>), }