From 29ab3a222bf1cf3b3a2ce1f83a841a2163e3ac70 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Tue, 7 Jul 2026 15:24:45 -0400 Subject: [PATCH 1/2] Support grouped distribution requirements --- .../enforce_distribution.rs | 51 +- .../physical_optimizer/ensure_requirements.rs | 8 +- .../physical_optimizer/projection_pushdown.rs | 4 +- .../physical_optimizer/sanity_checker.rs | 51 +- .../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 | 231 -------- .../enforce_distribution.rs | 305 +++++++---- .../enforce_sorting/mod.rs | 22 +- .../enforce_sorting/sort_pushdown.rs | 20 +- .../src/output_requirements.rs | 33 +- .../physical-optimizer/src/sanity_checker.rs | 49 +- 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 | 32 +- .../physical-plan/src/joins/cross_join.rs | 8 +- .../physical-plan/src/joins/hash_join/exec.rs | 167 +++++- .../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 | 6 +- 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 +- .../src/test_context/range_partitioning.rs | 29 +- .../test_files/range_partitioning.slt | 385 ++++++++++++- 37 files changed, 1625 insertions(+), 463 deletions(-) create mode 100644 datafusion/physical-plan/src/distribution_requirements.rs diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index cab8dc67f90d2..5a9a41f1153ac 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -265,8 +265,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 { @@ -823,6 +827,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 + 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/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 3fdbc9d312151..ef999cc4c1161 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -974,8 +974,12 @@ impl ExecutionPlan for MockReqExec { fn children(&self) -> Vec<&Arc> { vec![&self.input] } - fn required_input_distribution(&self) -> Vec { - vec![self.dist.clone()] + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + self.dist.clone(), + ]) } fn required_input_ordering(&self) -> Vec> { vec![ diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 24ec633d48d23..3a8d82f111145 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 217570846d56e..e759156282306 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,49 @@ 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, + )?); + 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("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::Inner, + )?; + 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::Inner, + )?; + 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/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 67419944cfde6..80e9f88b510ed 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -63,7 +63,9 @@ pub use equivalence::{ AcrossPartitions, ConstExpr, EquivalenceProperties, calculate_union, }; pub use expressions::{DynamicFilterTracker, DynamicFilterTracking}; -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 b662207e383a0..5afceacb8091b 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -244,50 +244,6 @@ impl RangePartitioning { self.split_points.len() + 1 } - /// Returns true when `self` and `other` describe the same range partition - /// map. - /// - /// Single-partition range partitionings are always compatible. Otherwise, - /// the two partitionings must have identical split points and equivalent - /// ordering expressions with the same sort options. - pub fn compatible_with( - &self, - other: &Self, - eq_properties: &EquivalenceProperties, - ) -> 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; - } - - let left_exprs = self - .ordering - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - let right_exprs = other - .ordering - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - - equivalent_exprs(&left_exprs, &right_exprs, eq_properties) - } - /// Calculates the range partitioning after applying the given projection. /// /// Returns `None` if any range key cannot be projected or if projection @@ -406,42 +362,6 @@ impl Partitioning { } } - /// Returns true when `self` and `other` describe compatible partition maps. - /// - /// Compatible partition maps can be used for partition-local behavior: if - /// this returns true, partition `i` from both partitionings can be treated - /// as covering the same partition domain. This is stricter than - /// [`Self::satisfaction`], which only answers whether this partitioning can - /// satisfy a required distribution. - pub fn compatible_with( - &self, - other: &Self, - eq_properties: &EquivalenceProperties, - ) -> bool { - if self.partition_count() == 1 && other.partition_count() == 1 { - return true; - } - - match (self, other) { - ( - Partitioning::Hash(left_exprs, left_count), - Partitioning::Hash(right_exprs, right_count), - ) => { - if left_count != right_count { - return false; - } - if left_exprs.is_empty() || right_exprs.is_empty() { - return false; - } - equivalent_exprs(left_exprs, right_exprs, eq_properties) - } - (Partitioning::Range(left), Partitioning::Range(right)) => { - left.compatible_with(right, eq_properties) - } - _ => 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). @@ -1304,157 +1224,6 @@ mod tests { Ok(()) } - #[test] - fn test_range_partitioning_compatible_with() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let mut eq_properties = fixture.eq_properties.clone(); - eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?; - - let split_points = vec![int_split_point([10]), int_split_point([20])]; - let range_a = fixture.range([0], split_points.clone()); - let range_a_same = fixture.range([0], split_points.clone()); - let range_b_equivalent = fixture.range([1], split_points.clone()); - let range_b_different_split = fixture.range([1], vec![int_split_point([30])]); - let range_a_desc = RangePartitioning::try_new( - [fixture.range_sort_expr(0, SortOptions::new(true, false))].into(), - vec![int_split_point([10])], - )?; - let single_partition_range_a = fixture.range([0], vec![]); - let single_partition_range_b = fixture.range([1], vec![]); - - assert!(range_a.compatible_with(&range_a_same, &fixture.eq_properties)); - assert!(range_a.compatible_with(&range_b_equivalent, &eq_properties)); - assert!(!range_a.compatible_with(&range_b_equivalent, &fixture.eq_properties)); - assert!(!range_a.compatible_with(&range_b_different_split, &eq_properties)); - assert!(!range_a.compatible_with(&range_a_desc, &eq_properties)); - assert!( - single_partition_range_a - .compatible_with(&single_partition_range_b, &fixture.eq_properties) - ); - - assert!( - fixture - .range_partitioning([0], vec![int_split_point([10])]) - .compatible_with( - &fixture.range_partitioning([1], vec![int_split_point([10])]), - &eq_properties - ) - ); - assert!( - !fixture - .range_partitioning([0], vec![int_split_point([10])]) - .compatible_with( - &fixture.range_partitioning([0], vec![int_split_point([20])]), - &fixture.eq_properties - ) - ); - assert!( - !fixture - .range_partitioning([0], vec![int_split_point([10])]) - .compatible_with( - &fixture.hash_partitioning([0], 2), - &fixture.eq_properties - ) - ); - - Ok(()) - } - - #[test] - fn test_hash_partitioning_compatible_with() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let mut eq_properties = fixture.eq_properties.clone(); - eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?; - - assert!( - fixture.hash_partitioning([0], 2).compatible_with( - &fixture.hash_partitioning([0], 2), - &fixture.eq_properties - ) - ); - assert!( - fixture - .hash_partitioning([0], 2) - .compatible_with(&fixture.hash_partitioning([1], 2), &eq_properties) - ); - assert!( - !fixture.hash_partitioning([0], 2).compatible_with( - &fixture.hash_partitioning([1], 2), - &fixture.eq_properties - ) - ); - assert!( - !fixture.hash_partitioning([0], 2).compatible_with( - &fixture.hash_partitioning([0], 3), - &fixture.eq_properties - ) - ); - assert!(!fixture.hash_partitioning([0], 2).compatible_with( - &fixture.hash_partitioning([0, 1], 2), - &fixture.eq_properties - )); - assert!( - !Partitioning::Hash(vec![], 2) - .compatible_with(&Partitioning::Hash(vec![], 2), &fixture.eq_properties) - ); - assert!(!fixture.hash_partitioning([0], 2).compatible_with( - &fixture.range_partitioning([0], vec![int_split_point([10])]), - &fixture.eq_properties - )); - assert!( - fixture.hash_partitioning([0], 1).compatible_with( - &Partitioning::RoundRobinBatch(1), - &fixture.eq_properties - ) - ); - - Ok(()) - } - - #[test] - fn test_round_robin_partitioning_compatible_with() { - let eq_properties = EquivalenceProperties::new(Arc::new(Schema::empty())); - - assert!( - Partitioning::RoundRobinBatch(1) - .compatible_with(&Partitioning::RoundRobinBatch(1), &eq_properties) - ); - assert!( - !Partitioning::RoundRobinBatch(2) - .compatible_with(&Partitioning::RoundRobinBatch(2), &eq_properties) - ); - assert!( - Partitioning::RoundRobinBatch(1) - .compatible_with(&Partitioning::UnknownPartitioning(1), &eq_properties) - ); - assert!( - !Partitioning::RoundRobinBatch(2) - .compatible_with(&Partitioning::UnknownPartitioning(2), &eq_properties) - ); - } - - #[test] - fn test_unknown_partitioning_compatible_with() { - let eq_properties = EquivalenceProperties::new(Arc::new(Schema::empty())); - - assert!( - Partitioning::UnknownPartitioning(1) - .compatible_with(&Partitioning::UnknownPartitioning(1), &eq_properties) - ); - assert!( - !Partitioning::UnknownPartitioning(2) - .compatible_with(&Partitioning::UnknownPartitioning(2), &eq_properties) - ); - assert!( - Partitioning::UnknownPartitioning(1) - .compatible_with(&Partitioning::RoundRobinBatch(1), &eq_properties) - ); - assert!( - !Partitioning::UnknownPartitioning(2) - .compatible_with(&Partitioning::RoundRobinBatch(2), &eq_properties) - ); - } - #[test] fn test_multi_partition_range_does_not_satisfy_hash_distribution() -> Result<()> { let fixture = PartitioningTestFixture::int64(&["a", "b"])?; diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 7f3a63ed91c48..804670f2bde59 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -34,9 +34,8 @@ use std::sync::Arc; 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; @@ -48,7 +47,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::{ @@ -68,7 +68,8 @@ 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, with_new_children_if_necessary, + ChildSatisfactionOptions, Distribution, ExecutionPlan, InputDistributionRequirements, + Partitioning, with_new_children_if_necessary, }; use itertools::izip; @@ -700,9 +701,10 @@ fn add_roundrobin_on_top( } } -// 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 @@ -946,6 +948,14 @@ struct RepartitionRequirementStatus { hash_necessary: 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 @@ -989,7 +999,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 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) @@ -1041,6 +1051,90 @@ fn get_repartition_requirement_status( .collect()) } +/// Enforce cross-child distribution relationships after each child has already +/// satisfied its own distribution requirement. +/// +/// 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 unsatisfied_children.is_empty() { + return Ok(()); + } + + 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; + } + + if !changed { + return datafusion_common::internal_err!( + "{plan_name} has distribution relationships that could not be enforced" + ); + } + } +} + /// 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 @@ -1148,6 +1242,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: @@ -1155,7 +1250,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(), @@ -1163,6 +1259,7 @@ pub fn ensure_distribution( ) .map( |( + child_idx, mut child, required_input_ordering, maintains, @@ -1177,6 +1274,10 @@ pub fn ensure_distribution( // 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 @@ -1247,26 +1348,19 @@ pub fn ensure_distribution( | Distribution::KeyPartitioned(exprs) => { let child_partitions = child.plan.output_partitioning().partition_count(); - let distribution_satisfied = child - .plan - .output_partitioning() - .satisfaction( - &requirement, - child.plan.equivalence_properties(), - allow_subset_satisfy_partitioning, - ) + 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 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 partitioning_satisfied = - distribution_satisfied || range_satisfied_for_aggregate; + let preserve_satisfying_file_partitioning = + preserve_file_partition_threshold_met + && !requires_grouping_id + && partitioning_satisfied + && target_partitions > child_partitions; // When subset satisfaction is enabled, preserve an // already-satisfying partitioning. Otherwise, hash @@ -1274,7 +1368,9 @@ pub fn ensure_distribution( let needs_hash_repartition = if allow_subset_satisfy_partitioning { !partitioning_satisfied } else { - !partitioning_satisfied || target_partitions > child_partitions + !partitioning_satisfied + || (target_partitions > child_partitions + && !preserve_satisfying_file_partitioning) }; let should_add_hash_repartition = hash_necessary && needs_hash_repartition; @@ -1309,75 +1405,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() @@ -1452,7 +1574,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)?; } @@ -676,7 +684,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/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 159c5a9e502f8..53c2dffd76624 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -80,7 +80,10 @@ pub type SortPushDown = PlanContext; /// Assigns the ordering requirement of the root node to the its children. pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) { let reqs = sort_push_down.plan.required_input_ordering(); - let dists = sort_push_down.plan.required_input_distribution(); + let dists = sort_push_down + .plan + .input_distribution_requirements() + .into_per_child(); for (idx, (child, requirement)) in sort_push_down.children.iter_mut().zip(reqs).enumerate() { @@ -165,7 +168,10 @@ fn pushdown_sorts_helper( } sort_push_down.plan = plan; // No ordering is being pushed; use each child's own distribution requirement - let dists = sort_push_down.plan.required_input_distribution(); + let dists = sort_push_down + .plan + .input_distribution_requirements() + .into_per_child(); for (idx, child) in sort_push_down.children.iter_mut().enumerate() { child.data.distribution_requirement = dists .get(idx) @@ -242,7 +248,10 @@ fn pushdown_sorts_helper( if satisfy_parent { // For non-sort operators which satisfy ordering: let reqs = sort_push_down.plan.required_input_ordering(); - let dists = sort_push_down.plan.required_input_distribution(); + let dists = sort_push_down + .plan + .input_distribution_requirements() + .into_per_child(); // If this node already outputs single partition, don't push SinglePartition // requirement to children (they're below the merge point). @@ -274,7 +283,10 @@ fn pushdown_sorts_helper( // requirements. If this node already outputs single partition (e.g. SPM), // don't push SinglePartition to children. let current_fetch = sort_push_down.plan.fetch(); - let dists = sort_push_down.plan.required_input_distribution(); + let dists = sort_push_down + .plan + .input_distribution_requirements() + .into_per_child(); let effective_dist = if sort_push_down.plan.output_partitioning().partition_count() == 1 { Distribution::UnspecifiedDistribution diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 4391c5ff6c981..c6f5f87622bea 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; @@ -208,7 +208,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 { @@ -275,9 +283,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)? @@ -288,7 +299,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| { @@ -371,7 +387,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 936bc8271a459..713213b70612d 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -24,21 +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::execution_plan::{ + Boundedness, EmissionType, InvariantLevel, +}; use datafusion_physical_plan::joins::SymmetricHashJoinExec; -use datafusion_physical_plan::{ExecutionPlanProperties, get_plan_string}; +use datafusion_physical_plan::{ + ChildSatisfactionOptions, 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; @@ -140,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 { @@ -170,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: {}", @@ -201,6 +186,8 @@ pub fn check_plan_sanity( } } + plan.check_invariants(InvariantLevel::Executable)?; + Ok(()) } diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index e7a19ca0b0012..1fbf8c6fb78cd 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -22,7 +22,6 @@ use datafusion_physical_expr::{ Distribution, 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; @@ -216,24 +215,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 11446137f3ca1..a0137f16b109d 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -38,8 +38,8 @@ use crate::filter_pushdown::{ use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::statistics::StatisticsArgs; 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; @@ -1783,7 +1783,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] } @@ -1793,6 +1797,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 27e0f5e923d85..72cd24ef95673 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -210,7 +210,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 76abf73e0ebbe..638746c267e9f 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, @@ -164,12 +165,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`, By default it's [[Distribution::UnspecifiedDistribution]] for each child, + /// Specifies simple per-child input distribution requirements. + /// + /// Deprecated: override [`Self::input_distribution_requirements`] instead. + /// + /// By default, each child has [`Distribution::UnspecifiedDistribution`]. + #[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`. /// @@ -216,8 +235,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() } @@ -1228,14 +1247,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 79295ba2fb556..d1f43ebc77f40 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -307,10 +307,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 50a90b0f54633..56e7132dc4df7 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -54,8 +54,9 @@ use crate::projection::{ use crate::repartition::REPARTITION_RANDOM_STATE; use crate::statistics::StatisticsArgs; 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, @@ -882,9 +883,34 @@ impl HashJoinExec { return false; } + 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 } + 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 + } + } + } + /// left (build) side which gets hashed pub fn left(&self) -> &Arc { &self.left @@ -1241,26 +1267,36 @@ impl ExecutionPlan for HashJoinExec { } 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 } } @@ -2190,6 +2226,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::{ @@ -2212,11 +2249,67 @@ mod tests { use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; + 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) } @@ -6701,6 +6794,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 3cae05a3a815a..fe17ee1838e59 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -566,10 +566,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 2dc7065eee04e..e43606ec4222e 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -424,15 +424,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 52a1aa056d244..0f79497d27929 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 39a4c178ca4b6..2a7759a8abeec 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -145,12 +145,10 @@ pub fn adjust_right_output_partitioning( .collect::>()?; Partitioning::Hash(new_exprs, *size) } - Partitioning::Range(_) => { + Partitioning::Range(range) => { // Range partitioning optimizer propagation is tracked in // https://github.com/apache/datafusion/issues/22395 - return not_impl_err!( - "Join output partitioning with range partitioning is not implemented" - ); + 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 6cc6e44c32cc3..8f40dde22ad2a 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, @@ -72,6 +75,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 1e4b5e5bb6426..0fe19f962fbe3 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -163,7 +163,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 7289ac43e510c..00df227cb87db 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -164,10 +164,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 d215e5296f91d..375dcae356f39 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -299,11 +299,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 e09147ed90274..aee9e52568b0d 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -299,12 +299,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 792c432155a8b..df6ff378887d8 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1258,14 +1258,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 77a7d8f8f2e11..f80b56b04ec77 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -267,7 +267,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 a9d580f4c687d..baafaba8058ab 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -332,12 +332,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 72474c6a55483..a10ba2cd87d18 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -241,11 +241,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/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 7904f92310957..dac9a95a16ca6 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -260,24 +260,393 @@ 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 -# 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. +# 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 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 +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={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 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={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 +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned_shifted/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned_shifted/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned_shifted/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned_shifted/part-3.csv]]}, 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={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)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +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 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 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={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 +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +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, 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={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 +04)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +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, 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={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 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={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 +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, 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={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 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={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 +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={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 +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={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 +06)--------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 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 10: 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 From 3a4b586b995157a221d942581996cf2ce9a990f9 Mon Sep 17 00:00:00 2001 From: Mithun Chicklore Yogendra Date: Wed, 8 Jul 2026 02:43:57 +0530 Subject: [PATCH 2/2] feat: allow Partitioning::Range to satisfy window Distribution::KeyPartitioned requirements Windows got a hash repartition even when the input was already range partitioned on the PARTITION BY keys. Opt WindowAggExec and BoundedWindowAggExec into range satisfaction for KeyPartitioned, the same way AggregateExec does. Incompatible keys still fall back to hash repartitioning, and windows without PARTITION BY keep requiring a single partition. Covered by enforce_distribution and sanity_checker tests plus a window section in range_partitioning.slt. Closes #23289 --- .../enforce_distribution.rs | 67 ++++- .../physical_optimizer/sanity_checker.rs | 77 ++++- .../tests/physical_optimizer/test_utils.rs | 18 +- .../src/windows/bounded_window_agg_exec.rs | 11 +- .../src/windows/window_agg_exec.rs | 11 +- .../test_files/range_partitioning.slt | 270 ++++++++++++++++++ 6 files changed, 440 insertions(+), 14 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 5a9a41f1153ac..6f730e4f92dc3 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, }; @@ -870,6 +870,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 d43a4a4cb9c26..c43ab016cebaa 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -274,6 +274,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(); @@ -296,7 +312,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 baafaba8058ab..2f83d8c290f5a 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -336,12 +336,15 @@ impl ExecutionPlan for BoundedWindowAggExec { } fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() { + if self.partition_keys().is_empty() { debug!("No partition defined for BoundedWindowAggExec!!!"); - vec![Distribution::SinglePartition] + crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } else { - vec![Distribution::KeyPartitioned(self.partition_keys().clone())] - }) + crate::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 a10ba2cd87d18..796356c433730 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -245,11 +245,14 @@ impl ExecutionPlan for WindowAggExec { } fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(if self.partition_keys().is_empty() { - vec![Distribution::SinglePartition] + if self.partition_keys().is_empty() { + crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } else { - vec![Distribution::KeyPartitioned(self.partition_keys())] - }) + crate::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 dac9a95a16ca6..fb3d9a58bf4f1 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -682,5 +682,275 @@ 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={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) 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={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) 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={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 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: 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={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 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 23: 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={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 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; + + +########## +# TEST 24: Window 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) 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={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) 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 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 25: Window 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) 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)------RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +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, 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; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 26: 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={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) 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;