This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 63e25c15bc Support co-partitioned range inner equi joins (#23184)
63e25c15bc is described below
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]
