This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23184-24d8957cf4b19bfa45568a9ef3dec82109e5d05e in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 63e25c15bc72e731c174ba03d1143b243cc1aaae Author: Gene Bordegaray <[email protected]> AuthorDate: Fri Jul 10 12:53:22 2026 -0400 Support co-partitioned range inner equi joins (#23184) ## Which issue does this PR close? - Closes #23183. - Part of #22395. ## Rationale for this change DataFusion can represent source-declared range partitioning, but partitioned hash joins still required hash partitioned inputs. So an inner join on compatible range-partitioned keys would insert unnecessary hash repartitions, even when each left/right partition already covered the same key domain. This PR adds a partitioning requirement that means "equal key values are co-located" . I was calling this "compatibility" but found we can satisfy the requirement with looser conditions. Other systems call this "co-location" or "co-partitioning" ([trino](https://trino.io/docs/current/admin/properties-optimizer.html#optimizer-colocated-joins-enabled), [spark](https://spark.apache.org/docs/latest/sql-performance-tuning.html#storage-partition-join)). Which they (and now I am proposing) define as when both sides of a join are already partitioned so matching key values appear in corresponding partitions, so we can join partition pairs directly without repartitioning the sides. This lets "co-partitioned" range inputs satisfy inner partitioned hash joins. This will also be applicable to other join types and operators but kept the first PR thin to keep scope more reviewable. ## What changes are included in this PR? - Adds `Distribution::KeyPartitioned(Vec<Arc<dyn PhysicalExpr>>)` as a public distribution requirement. - `HashPartitioned([a])` means rows must be partitioned by hash on `a`. - `KeyPartitioned([a])` means rows with equal `a` values must be co-located, but the partitioning algorithm may be hash, range, or another compatible scheme. - Example: ```text Hash([left.a], 3) satisfies KeyPartitioned([left.a]) Range([right.b ASC], [(10), (20)], 3) satisfies KeyPartitioned([right.b]) ``` - Adds `Partitioning::co_partitioned_with(...)` to validate that two independently satisfying partitionings also can be paired by partition index. - Examples: - Accepted: both sides satisfy their own key requirement and have matching range boundaries. ```text left: Range([a ASC], [(10), (20)], 3), required KeyPartitioned([a]) right: Range([b ASC], [(10), (20)], 3), required KeyPartitioned([b]) ``` - Accepted: both sides satisfy their own key requirement and have matching hash partition counts. ```text left: Hash([a], 3), required KeyPartitioned([a]) right: Hash([b], 3), required KeyPartitioned([b]) ``` - Rejected: both sides satisfy their own key requirement, but range boundaries differ. ```text left: Range([a ASC], [(10), (20)], 3), required KeyPartitioned([a]) right: Range([b ASC], [(15), (20)], 3), required KeyPartitioned([b]) ``` - Rejected: both sides satisfy their own key requirement, but partition counts differ. ```text left: Hash([a], 3), required KeyPartitioned([a]) right: Hash([b], 4), required KeyPartitioned([b]) ``` - Changes inner partitioned `HashJoinExec` requirements from `HashPartitioned` to `KeyPartitioned`. - All other hash joins still require `HashPartitioned` for now. - Updates `EnforceDistribution` so co-partitioned range inner joins avoid repartitioning. - Examples: - Compatible range partitioning: no repartition is inserted because partitions can be joined by index. ```text HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a, b)] DataSourceExec: output_partitioning=Range([a ASC], [(10), (20)], 3) DataSourceExec: output_partitioning=Range([b ASC], [(10), (20)], 3) ``` - Incompatible range boundaries: both sides are repartitioned by hash because partition `i` does not represent the same key domain on both sides. ```text HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a, b)] RepartitionExec: partitioning=Hash([a], target_partitions) DataSourceExec: output_partitioning=Range([a ASC], [(10), (20)], 3) RepartitionExec: partitioning=Hash([b], target_partitions) DataSourceExec: output_partitioning=Range([b ASC], [(15), (20)], 3) ``` - Mismatched hash partition counts: both sides are forced to the target hash partition count so partition indexes line up. ```text HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a, b)] RepartitionExec: partitioning=Hash([a], target_partitions) DataSourceExec: output_partitioning=Hash([a], 11) RepartitionExec: partitioning=Hash([b], target_partitions) DataSourceExec: output_partitioning=Hash([b], 12) ``` - Non-inner joins: range inputs still get hash repartitioning because only inner partitioned hash joins use `KeyPartitioned` in this PR. ```text HashJoinExec: mode=Partitioned, join_type=Left, on=[(a, b)] RepartitionExec: partitioning=Hash([a], target_partitions) DataSourceExec: output_partitioning=Range([a ASC], [(10), (20)], 3) RepartitionExec: partitioning=Hash([b], target_partitions) DataSourceExec: output_partitioning=Range([b ASC], [(10), (20)], 3) ``` - Keeps partitioned dynamic filter pushdown restricted to hash-compatible routing. - Compatible range partitioning can satisfy the join, but dynamic filters still route by hash, so range/range partitioned joins disable dynamic filters. - Degrades range join output partitioning to `UnknownPartitioning(n)` rather than erroring. Adding this behavior would need more tests and careful thought about, I think its safert o just degrade for first PR. ## Are these changes tested? Yes. - `KeyPartitioned` satisfaction for hash and range partitioning. - `co_partitioned_with` for compatible and incompatible range/hash partitioning. - `EnforceDistribution` behavior for: - compatible range joins avoiding hash repartitioning - incompatible range bounds rehashing - mismatched hash partition counts rehashing - non-inner range joins rehashing - sanity checking for invalid partitioned hash joins. - dynamic filter rejection for range partitioning, preserved file partitions, and mismatched hash counts. - sqllogictest coverage for range-partitioned joins avoiding hash repartitioning and non-range joins still repartitioning. ## Are there any user-facing changes? Yes. This PR changes public physical planning APIs: - Adds `Distribution::KeyPartitioned`. - Adds `Partitioning::co_partitioned_with`. - **NOTE**: This replaces the previous partition compatibility API with the new co-partitioning API. Since the compatibility API was never in a release I believe this is ok to do (lesson learned to not make API change until ew have definitive consumer). - Affects users matching exhaustively on `Distribution`. --- .../physical_optimizer/enforce_distribution.rs | 51 ++- .../physical_optimizer/ensure_requirements.rs | 8 +- .../physical_optimizer/projection_pushdown.rs | 4 +- .../tests/physical_optimizer/sanity_checker.rs | 51 ++- .../core/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 ---------- .../ensure_requirements/enforce_distribution.rs | 305 ++++++++---- .../src/ensure_requirements/enforce_sorting/mod.rs | 22 +- .../enforce_sorting/sort_pushdown.rs | 20 +- .../physical-optimizer/src/output_requirements.rs | 33 +- .../physical-optimizer/src/sanity_checker.rs | 49 +- datafusion/physical-optimizer/src/utils.rs | 19 - datafusion/physical-plan/src/aggregates/mod.rs | 18 +- datafusion/physical-plan/src/analyze.rs | 8 +- .../physical-plan/src/distribution_requirements.rs | 510 +++++++++++++++++++++ datafusion/physical-plan/src/execution_plan.rs | 32 +- datafusion/physical-plan/src/joins/cross_join.rs | 8 +- .../physical-plan/src/joins/hash_join/exec.rs | 167 ++++++- .../physical-plan/src/joins/nested_loop_join.rs | 8 +- .../src/joins/piecewise_merge_join/exec.rs | 8 +- .../src/joins/sort_merge_join/exec.rs | 8 +- .../physical-plan/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 +- datafusion/physical-plan/src/recursive_query.rs | 8 +- datafusion/physical-plan/src/sorts/partial_sort.rs | 8 +- .../physical-plan/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 +- .../physical-plan/src/windows/window_agg_exec.rs | 8 +- .../src/test_context/range_partitioning.rs | 29 +- .../sqllogictest/test_files/range_partitioning.slt | 405 +++++++++++++++- 37 files changed, 1635 insertions(+), 473 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 462807e436..e01311e25b 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<Distribution> { - 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<bool> { @@ -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 4e2f2ce601..2c6c46c829 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<dyn ExecutionPlan>> { vec![&self.input] } - fn required_input_distribution(&self) -> Vec<Distribution> { - 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<Option<OrderingRequirements>> { vec![ diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 24ec633d48..3a8d82f111 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::<OutputRequirementExec>() .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 217570846d..e759156282 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<dyn ExecutionPlan>, is_sane: bool) { ); } +fn range_partitioned_exec( + schema: &SchemaRef, + key: &str, + split_points: impl IntoIterator<Item = i32>, +) -> Result<Arc<dyn ExecutionPlan>> { + 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<dyn ExecutionPlan>) +} + +#[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 e8ff6758cc..b837373632 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<Distribution> { - 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<dyn ExecutionPlan>> { diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index e3df1ad638..18ebe80773 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<Distribution> { + 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<Option<OrderingRequirements>> { diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index 67419944cf..80e9f88b51 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 b9ec312e94..61492934eb 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::<Vec<_>>(); - let right_exprs = other - .ordering - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::<Vec<_>>(); - - 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). @@ -1309,157 +1229,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 0d52dc5614..3cf79619cd 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; @@ -697,9 +698,10 @@ fn add_roundrobin_on_top( } } -// TODO: remove this private helper once Range generally satisfies -// KeyPartitioned requirements through Partitioning::satisfaction. -// See <https://github.com/apache/datafusion/issues/23266>. +// TODO: remove this temporary bridge once [`Partitioning::Range`] +// generally satisfies [`Distribution::KeyPartitioned`] through +// [`Partitioning::satisfaction`]. +// <https://github.com/apache/datafusion/issues/23266>. // // Partial aggregates do not require key partitioning, but they preserve their // input partitioning for the final aggregate. Until Range satisfies @@ -943,6 +945,14 @@ struct RepartitionRequirementStatus { hash_necessary: bool, } +/// Per-child state while enforcing a parent's distribution requirements. +struct DistributionChildState { + context: DistributionContext, + required_input_ordering: Option<OrderingRequirements>, + 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 @@ -986,7 +996,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) @@ -1038,6 +1048,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::<Vec<_>>(); + 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 @@ -1145,6 +1239,7 @@ pub fn ensure_distribution( .is_some_and(|join| join.mode == PartitionMode::Partitioned) || plan.is::<SortMergeJoinExec>(); + 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: @@ -1152,7 +1247,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(), @@ -1160,6 +1256,7 @@ pub fn ensure_distribution( ) .map( |( + child_idx, mut child, required_input_ordering, maintains, @@ -1174,6 +1271,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 @@ -1244,26 +1345,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 @@ -1271,7 +1365,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; @@ -1306,75 +1402,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::<Result<Vec<_>>>()?; - // 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::<OutputRequirementExec>() - .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::<OutputRequirementExec>() + .map(|output| output.fetch()) + .unwrap_or(None), + )?; } } - Distribution::UnspecifiedDistribution => { - // Since ordering is lost, trying to preserve ordering is pointless - if !maintains || plan.is::<OutputRequirementExec>() { - 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::<OutputRequirementExec>() + { + context = replace_order_preserving_variants(context)?; + } } } } - } - Ok(child) - }, - ) - .collect::<Result<Vec<_>>>()?; + Ok(context) + }, + ) + .collect::<Result<Vec<_>>>()?; let children_plans = children .iter() @@ -1449,7 +1571,8 @@ fn update_children(mut dist_context: DistributionContext) -> Result<Distribution || child_context.children[0].data || child_context .plan - .required_input_distribution() + .input_distribution_requirements() + .into_per_child() .iter() .zip(child_context.children.iter()) .any(|(required_dist, child_context)| { diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs index 53917a5108..4dce4691f0 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -178,8 +178,11 @@ fn update_coalesce_ctx_children( // and connected to some `CoalescePartitionsExec`: node.data && !matches!( - coalesce_context.plan.required_input_distribution()[idx], - Distribution::SinglePartition + coalesce_context + .plan + .input_distribution_requirements() + .child_distribution(idx), + Some(Distribution::SinglePartition) ) }) }; @@ -539,8 +542,11 @@ fn adjust_window_sort_removal( let child_node = remove_corresponding_sort_from_sub_plan( window_tree.children.swap_remove(0), matches!( - window_tree.plan.required_input_distribution()[0], - Distribution::SinglePartition + window_tree + .plan + .input_distribution_requirements() + .child_distribution(0), + Some(Distribution::SinglePartition) ), )?; window_tree.children.push(child_node); @@ -654,8 +660,10 @@ fn update_child_to_remove_unnecessary_sort( ) -> Result<PlanWithCorrespondingSort> { 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 d69ae34610..c1e42a7c9a 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<ParentRequirements>; /// 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 4391c5ff6c..c6f5f87622 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<Distribution> { - 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<bool> { @@ -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 936bc8271a..713213b706 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<dyn ExecutionPlan>, 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 e7a19ca0b0..1fbf8c6fb7 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 <https://github.com/apache/datafusion/issues/23266>. -/// -/// Checks whether an aggregate can reuse range partitioning to satisfy its key -/// partitioning requirement. -pub(crate) fn aggregate_can_reuse_range_partitioning( - plan: &Arc<dyn ExecutionPlan>, -) -> bool { - plan.downcast_ref::<AggregateExec>() - .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<dyn ExecutionPlan>) -> bool { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index aa42f7a01b..e7832629b7 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; @@ -1772,7 +1772,11 @@ impl ExecutionPlan for AggregateExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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] } @@ -1782,6 +1786,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 27e0f5e923..72cd24ef95 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<Distribution> { - 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 0000000000..9c7a1336c0 --- /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<ChildDistributionRequirement>, + /// Child indexes that must also have compatible partition layouts. + co_partitioned: Option<Vec<usize>>, +} + +/// 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<Distribution>) -> 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<Distribution>) -> 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<Item = &Distribution> + '_ { + 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<Distribution> { + 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<PartitioningSatisfaction> { + 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<Vec<usize>> { + 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`]. + /// <https://github.com/apache/datafusion/issues/23266>. + /// + /// 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<P: ExecutionPlan + ?Sized>( + &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::<Vec<_>>(); + 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`]. +/// <https://github.com/apache/datafusion/issues/23266>. +#[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`]. +/// <https://github.com/apache/datafusion/issues/23266>. +fn range_satisfies_key_partitioning( + partitioning: &Partitioning, + required_exprs: &[Arc<dyn PhysicalExpr>], + 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::<Vec<_>>(); + + 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::<Vec<_>>(); + let normalized_required_exprs = required_exprs + .iter() + .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) + .collect::<Vec<_>>(); + + 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 5837e4d07b..5f92ff7659 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<Distribution> { 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<bool> { // 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() } @@ -1249,14 +1268,15 @@ macro_rules! check_len { /// Returns an error if the given node does not conform. pub fn check_default_invariants<P: ExecutionPlan + ?Sized>( 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 f8a9b80179..d7ff07d00f 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<Distribution> { - 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 50a90b0f54..56e7132dc4 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. + // <https://github.com/apache/datafusion/issues/23376>. + 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<dyn ExecutionPlan> { &self.left @@ -1241,26 +1267,36 @@ impl ExecutionPlan for HashJoinExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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<PlanProperties>, + } + + impl PartitionedTestExec { + fn try_new(schema: SchemaRef, partitioning: Partitioning) -> Result<Self> { + 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<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn with_new_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + 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 c120654319..2d1a3ae62d 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -544,10 +544,14 @@ impl ExecutionPlan for NestedLoopJoinExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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<bool> { 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 bec91fdb62..5ec564295e 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -483,10 +483,14 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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<Option<OrderingRequirements>> { 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 3c0743ad89..82d9c900e8 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -410,15 +410,19 @@ impl ExecutionPlan for SortMergeJoinExec { } fn required_input_distribution(&self) -> Vec<Distribution> { + 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<Option<OrderingRequirements>> { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 0cef469071..f33d9b1d07 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -412,7 +412,11 @@ impl ExecutionPlan for SymmetricHashJoinExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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 @@ -427,7 +431,7 @@ impl ExecutionPlan for SymmetricHashJoinExec { StreamJoinPartitionMode::SinglePartition => { vec![Distribution::SinglePartition, Distribution::SinglePartition] } - } + }) } fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 39a4c178ca..2a7759a8ab 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::<Result<_>>()?; 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 6cc6e44c32..8f40dde22a 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 4f1abf288d..3327098040 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -152,7 +152,11 @@ impl ExecutionPlan for GlobalLimitExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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<bool> { diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 7289ac43e5..00df227cb8 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<crate::Distribution> { - 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 95c47c5d01..916cf1bcbb 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -288,11 +288,15 @@ impl ExecutionPlan for PartialSortExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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<bool> { diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index e09147ed90..aee9e52568 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<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { let partition_exprs: Vec<Arc<dyn PhysicalExpr>> = 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<bool> { diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 792c432155..df6ff37888 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<Distribution> { - 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<dyn ExecutionPlan>> { diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 6d377e64a6..b6625885eb 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -256,7 +256,13 @@ impl ExecutionPlan for SortPreservingMergeExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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<bool> { diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 632306b80c..01c2f3ae27 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<Distribution> { - 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 f0e18ee818..cc3d70a1ae 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -321,12 +321,16 @@ impl ExecutionPlan for BoundedWindowAggExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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<bool> { diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 0a99164144..f1b78ef5c1 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -230,11 +230,15 @@ impl ExecutionPlan for WindowAggExec { } fn required_input_distribution(&self) -> Vec<Distribution> { - 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 a3e16eefd8..aa741dded7 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 7904f92310..5d004ca7fd 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -43,7 +43,7 @@ EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; ---- physical_plan 01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] -02)--DataSourceExec: file_groups={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]]}, projectio [...] +02)--DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query II SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key ORDER BY range_key; @@ -77,7 +77,7 @@ physical_plan 01)AggregateExec: mode=FinalPartitioned, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] 02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 03)----AggregateExec: mode=Partial, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] -04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, proje [...] +04)------DataSourceExec: file_groups=<slt:ignore>, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false query II SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key ORDER BY non_range_key; @@ -103,7 +103,7 @@ EXPLAIN SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP ---- physical_plan 01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] -02)--DataSourceExec: file_groups={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]]}, projectio [...] +02)--DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query III SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key ORDER BY range_key, non_range_key; @@ -139,7 +139,7 @@ EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; ---- physical_plan 01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] -02)--DataSourceExec: file_groups={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]]}, projectio [...] +02)--DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false ########## @@ -165,7 +165,7 @@ physical_plan 01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] 02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] -04)------DataSourceExec: file_groups={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]]}, proje [...] +04)------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false ########## @@ -191,7 +191,7 @@ physical_plan 02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5 03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] 04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 -05)--------DataSourceExec: file_groups={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]]}, pro [...] +05)--------DataSourceExec: file_groups=<slt:ignore>, 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; @@ -220,7 +220,7 @@ EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; ---- physical_plan 01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] -02)--DataSourceExec: file_groups={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]]}, projectio [...] +02)--DataSourceExec: file_groups=<slt:ignore>, 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; @@ -252,7 +252,7 @@ physical_plan 02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5 03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] 04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 -05)--------DataSourceExec: file_groups={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]]}, pro [...] +05)--------DataSourceExec: file_groups=<slt:ignore>, 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; @@ -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 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false -query error This feature is not implemented: Join output partitioning with range partitioning is not implemented +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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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 @@ -287,8 +656,8 @@ SELECT range_key, value FROM range_partitioned; ---- physical_plan 01)UnionExec -02)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projectio [...] -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]]}, projectio [...] +02)--DataSourceExec: file_groups=<slt:ignore>, 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=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false query II SELECT range_key, value FROM range_partitioned --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
