You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/23 07:20:30 UTC
(datafusion) 01/01: allow range to satisfy key distribution generally (#23680)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23680-47bd094c7fd2fc22cb2b94ee58458b87e365a56f
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit e4f3df7af8ada8c7f988c1398e2f2a1bb2dfa0fe
Author: Gene Bordegaray <ge...@datadoghq.com>
AuthorDate: Thu Jul 23 03:20:07 2026 -0400
allow range to satisfy key distribution generally (#23680)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Closes #23266.
- Part of #22395.
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
`Partitioning::Range` now satisfies `Distribution::KeyPartitioned`
privately across all operators that require it: aggregates, windows,
TopK, and co-partitioned joins. So now the temporary operator opt-ins
can be consolidated.
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
- Allow compatible `Partitioning::Range` to satisfy
`Distribution::KeyPartitioned` via `Partitioning::satisfaction`
- Remove the temporary range-satisfaction helpers and operator-specific
opt-ins
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
Yes
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
No, this should not change any exisitng behavior just consolidation
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---
datafusion/physical-expr/src/partitioning.rs | 428 ++++++++-------------
.../ensure_requirements/enforce_distribution.rs | 37 +-
datafusion/physical-optimizer/src/utils.rs | 59 +--
datafusion/physical-plan/src/aggregates/mod.rs | 12 +-
.../physical-plan/src/distribution_requirements.rs | 165 +-------
.../physical-plan/src/joins/hash_join/exec.rs | 10 +-
.../src/joins/sort_merge_join/exec.rs | 1 -
.../physical-plan/src/joins/symmetric_hash_join.rs | 1 -
.../physical-plan/src/sorts/partitioned_topk.rs | 1 -
.../src/windows/bounded_window_agg_exec.rs | 1 -
.../physical-plan/src/windows/window_agg_exec.rs | 1 -
11 files changed, 183 insertions(+), 533 deletions(-)
diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs
index 61492934eb..59d36c4efc 100644
--- a/datafusion/physical-expr/src/partitioning.rs
+++ b/datafusion/physical-expr/src/partitioning.rs
@@ -417,40 +417,26 @@ impl Partitioning {
// Here we do not check the partition count for hash partitioning and assumes the partition count
// and hash functions in the system are the same. In future if we plan to support storage partition-wise joins,
// then we need to have the partition count and hash functions validation.
- Partitioning::Hash(partition_exprs, _) => {
- // Empty hash partitioning is invalid
- if partition_exprs.is_empty() || required_exprs.is_empty() {
- return PartitioningSatisfaction::NotSatisfied;
- }
-
- if equivalent_exprs(required_exprs, partition_exprs, eq_properties) {
- return PartitioningSatisfaction::Exact;
- }
-
- let eq_groups = eq_properties.eq_group();
- if !eq_groups.is_empty() {
- if allow_subset {
- let normalized_partition_exprs =
- normalize_exprs(partition_exprs, eq_properties);
- let normalized_required_exprs =
- normalize_exprs(required_exprs, eq_properties);
- if Self::is_subset_partitioning(
- &normalized_partition_exprs,
- &normalized_required_exprs,
- ) {
- return PartitioningSatisfaction::Subset;
- }
- }
- } else if allow_subset
- && Self::is_subset_partitioning(partition_exprs, required_exprs)
- {
- return PartitioningSatisfaction::Subset;
- }
-
- PartitioningSatisfaction::NotSatisfied
+ Partitioning::Hash(partition_exprs, _) => Self::key_satisfaction(
+ partition_exprs,
+ required_exprs,
+ eq_properties,
+ allow_subset,
+ ),
+ Partitioning::Range(range) => {
+ let partition_exprs = range
+ .ordering()
+ .iter()
+ .map(|sort_expr| Arc::clone(&sort_expr.expr))
+ .collect::<Vec<_>>();
+ Self::key_satisfaction(
+ &partition_exprs,
+ required_exprs,
+ eq_properties,
+ allow_subset,
+ )
}
Partitioning::RoundRobinBatch(_)
- | Partitioning::Range(_)
| Partitioning::UnknownPartitioning(_) => {
PartitioningSatisfaction::NotSatisfied
}
@@ -459,6 +445,43 @@ impl Partitioning {
}
}
+ fn key_satisfaction(
+ partition_exprs: &[Arc<dyn PhysicalExpr>],
+ required_exprs: &[Arc<dyn PhysicalExpr>],
+ eq_properties: &EquivalenceProperties,
+ allow_subset: bool,
+ ) -> PartitioningSatisfaction {
+ if partition_exprs.is_empty() || required_exprs.is_empty() {
+ return PartitioningSatisfaction::NotSatisfied;
+ }
+
+ if equivalent_exprs(required_exprs, partition_exprs, eq_properties) {
+ return PartitioningSatisfaction::Exact;
+ }
+
+ let eq_groups = eq_properties.eq_group();
+ if !eq_groups.is_empty() {
+ if allow_subset {
+ let normalized_partition_exprs =
+ normalize_exprs(partition_exprs, eq_properties);
+ let normalized_required_exprs =
+ normalize_exprs(required_exprs, eq_properties);
+ if Self::is_subset_partitioning(
+ &normalized_partition_exprs,
+ &normalized_required_exprs,
+ ) {
+ return PartitioningSatisfaction::Subset;
+ }
+ }
+ } else if allow_subset
+ && Self::is_subset_partitioning(partition_exprs, required_exprs)
+ {
+ return PartitioningSatisfaction::Subset;
+ }
+
+ PartitioningSatisfaction::NotSatisfied
+ }
+
/// Calculate the output partitioning after applying the given projection.
pub fn project(
&self,
@@ -685,6 +708,26 @@ mod tests {
}
}
+ fn assert_satisfaction(
+ desc: &str,
+ partitioning: &Partitioning,
+ required: &Distribution,
+ eq_properties: &EquivalenceProperties,
+ expected_with_subset: PartitioningSatisfaction,
+ expected_without_subset: PartitioningSatisfaction,
+ ) {
+ assert_eq!(
+ partitioning.satisfaction(required, eq_properties, true),
+ expected_with_subset,
+ "Failed for {desc} with subset enabled"
+ );
+ assert_eq!(
+ partitioning.satisfaction(required, eq_properties, false),
+ expected_without_subset,
+ "Failed for {desc} with subset disabled"
+ );
+ }
+
#[test]
#[expect(
deprecated,
@@ -768,320 +811,121 @@ mod tests {
}
#[test]
- fn test_partitioning_satisfy_by_subset() -> Result<()> {
+ fn hash_partitioning_key_distribution_satisfaction() -> Result<()> {
let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?;
+ let unknown: Arc<dyn PhysicalExpr> = Arc::new(UnKnownColumn::new("dropped"));
let test_cases = vec![
(
- "KeyPartitioned([a, b]) satisfied by Hash([a])",
- fixture.hash_partitioning([0], 4),
+ "exact: KeyPartitioned([a, b]) satisfied by Hash([a, b])",
+ fixture.hash_partitioning([0, 1], 4),
fixture.key_distribution([0, 1]),
- PartitioningSatisfaction::Subset,
- PartitioningSatisfaction::NotSatisfied,
+ PartitioningSatisfaction::Exact,
+ PartitioningSatisfaction::Exact,
),
(
- "KeyPartitioned([a, b, c]) satisfied by Hash([a])",
+ "subset: KeyPartitioned([a, b]) satisfied by Hash([a])",
fixture.hash_partitioning([0], 4),
- fixture.key_distribution([0, 1, 2]),
- PartitioningSatisfaction::Subset,
- PartitioningSatisfaction::NotSatisfied,
- ),
- (
- "KeyPartitioned([a, b, c]) satisfied by Hash([a, b])",
- fixture.hash_partitioning([0, 1], 4),
- fixture.key_distribution([0, 1, 2]),
+ fixture.key_distribution([0, 1]),
PartitioningSatisfaction::Subset,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([a, b, c]) satisfied by Hash([b])",
+ "subset: KeyPartitioned([a, b, c]) satisfied by Hash([b])",
fixture.hash_partitioning([1], 4),
fixture.key_distribution([0, 1, 2]),
PartitioningSatisfaction::Subset,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([a, b, c]) satisfied by Hash([b, a])",
+ "subset reordered: KeyPartitioned([a, b, c]) satisfied by Hash([b, a])",
fixture.hash_partitioning([1, 0], 4),
fixture.key_distribution([0, 1, 2]),
PartitioningSatisfaction::Subset,
PartitioningSatisfaction::NotSatisfied,
),
- ];
-
- for (desc, partition, required, expected_with_subset, expected_without_subset) in
- test_cases
- {
- let result = partition.satisfaction(&required, &fixture.eq_properties, true);
- assert_eq!(
- result, expected_with_subset,
- "Failed for {desc} with subset enabled"
- );
-
- let result = partition.satisfaction(&required, &fixture.eq_properties, false);
- assert_eq!(
- result, expected_without_subset,
- "Failed for {desc} with subset disabled"
- );
- }
-
- Ok(())
- }
-
- #[test]
- fn test_partitioning_current_superset() -> Result<()> {
- let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?;
-
- let test_cases = vec![
(
- "KeyPartitioned([a]) satisfied by Hash([a, b])",
+ "superset: KeyPartitioned([a]) not satisfied by Hash([a, b])",
fixture.hash_partitioning([0, 1], 4),
fixture.key_distribution([0]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([a]) satisfied by Hash([a, b, c])",
+ "superset: KeyPartitioned([a, b]) not satisfied by Hash([a, b, c])",
fixture.hash_partitioning([0, 1, 2], 4),
- fixture.key_distribution([0]),
+ fixture.key_distribution([0, 1]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([a, b]) satisfied by Hash([a, b, c])",
- fixture.hash_partitioning([0, 1, 2], 4),
+ "partial overlap: KeyPartitioned([a, b]) not satisfied by Hash([a, c])",
+ fixture.hash_partitioning([0, 2], 4),
fixture.key_distribution([0, 1]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
- ];
-
- for (desc, partition, required, expected_with_subset, expected_without_subset) in
- test_cases
- {
- let result = partition.satisfaction(&required, &fixture.eq_properties, true);
- assert_eq!(
- result, expected_with_subset,
- "Failed for {desc} with subset enabled"
- );
-
- let result = partition.satisfaction(&required, &fixture.eq_properties, false);
- assert_eq!(
- result, expected_without_subset,
- "Failed for {desc} with subset disabled"
- );
- }
-
- Ok(())
- }
-
- #[test]
- fn test_partitioning_partial_overlap() -> Result<()> {
- let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?;
-
- let test_cases = vec![(
- "Partial overlap: KeyPartitioned([a, b]) satisfied by Hash([a, c])",
- fixture.hash_partitioning([0, 2], 4),
- fixture.key_distribution([0, 1]),
- PartitioningSatisfaction::NotSatisfied,
- PartitioningSatisfaction::NotSatisfied,
- )];
-
- for (desc, partition, required, expected_with_subset, expected_without_subset) in
- test_cases
- {
- let result = partition.satisfaction(&required, &fixture.eq_properties, true);
- assert_eq!(
- result, expected_with_subset,
- "Failed for {desc} with subset enabled"
- );
-
- let result = partition.satisfaction(&required, &fixture.eq_properties, false);
- assert_eq!(
- result, expected_without_subset,
- "Failed for {desc} with subset disabled"
- );
- }
-
- Ok(())
- }
-
- #[test]
- fn test_partitioning_no_overlap() -> Result<()> {
- let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?;
-
- let test_cases = vec![
(
- "KeyPartitioned([b, c]) satisfied by Hash([a])",
+ "no overlap: KeyPartitioned([b, c]) not satisfied by Hash([a])",
fixture.hash_partitioning([0], 4),
fixture.key_distribution([1, 2]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([c]) satisfied by Hash([a, b])",
- fixture.hash_partitioning([0, 1], 4),
- fixture.key_distribution([2]),
- PartitioningSatisfaction::NotSatisfied,
- PartitioningSatisfaction::NotSatisfied,
- ),
- ];
-
- for (desc, partition, required, expected_with_subset, expected_without_subset) in
- test_cases
- {
- let result = partition.satisfaction(&required, &fixture.eq_properties, true);
- assert_eq!(
- result, expected_with_subset,
- "Failed for {desc} with subset enabled"
- );
-
- let result = partition.satisfaction(&required, &fixture.eq_properties, false);
- assert_eq!(
- result, expected_without_subset,
- "Failed for {desc} with subset disabled"
- );
- }
-
- Ok(())
- }
-
- #[test]
- fn test_partitioning_exact_match() -> Result<()> {
- let fixture = PartitioningTestFixture::int64(&["a", "b"])?;
-
- let test_cases = vec![
- (
- "KeyPartitioned([a, b]) satisfied by Hash([a, b])",
- fixture.hash_partitioning([0, 1], 4),
- fixture.key_distribution([0, 1]),
- PartitioningSatisfaction::Exact,
- PartitioningSatisfaction::Exact,
- ),
- (
- "KeyPartitioned([a]) satisfied by Hash([a])",
- fixture.hash_partitioning([0], 4),
- fixture.key_distribution([0]),
- PartitioningSatisfaction::Exact,
- PartitioningSatisfaction::Exact,
- ),
- ];
-
- for (desc, partition, required, expected_with_subset, expected_without_subset) in
- test_cases
- {
- let result = partition.satisfaction(&required, &fixture.eq_properties, true);
- assert_eq!(
- result, expected_with_subset,
- "Failed for {desc} with subset enabled"
- );
-
- let result = partition.satisfaction(&required, &fixture.eq_properties, false);
- assert_eq!(
- result, expected_without_subset,
- "Failed for {desc} with subset disabled"
- );
- }
-
- Ok(())
- }
-
- #[test]
- fn test_partitioning_unknown() -> Result<()> {
- let fixture = PartitioningTestFixture::int64(&["a", "b"])?;
- let unknown: Arc<dyn PhysicalExpr> = Arc::new(UnKnownColumn::new("dropped"));
-
- let test_cases = vec![
- (
- "KeyPartitioned([a, b]) satisfied by Hash([unknown])",
+ "unknown partition expr",
Partitioning::Hash(vec![Arc::clone(&unknown)], 4),
fixture.key_distribution([0, 1]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([unknown]) satisfied by Hash([a, b])",
+ "unknown required expr",
fixture.hash_partitioning([0, 1], 4),
Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([unknown]) satisfied by Hash([unknown])",
+ "same unknown expr",
Partitioning::Hash(vec![Arc::clone(&unknown)], 4),
Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([unknown, a]) satisfied by Hash([unknown])",
+ "unknown partition expr is not a valid subset",
Partitioning::Hash(vec![Arc::clone(&unknown)], 4),
Distribution::KeyPartitioned(vec![Arc::clone(&unknown), fixture.col(0)]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
- ];
-
- for (desc, partition, required, expected_with_subset, expected_without_subset) in
- test_cases
- {
- let result = partition.satisfaction(&required, &fixture.eq_properties, true);
- assert_eq!(
- result, expected_with_subset,
- "Failed for {desc} with subset enabled"
- );
-
- let result = partition.satisfaction(&required, &fixture.eq_properties, false);
- assert_eq!(
- result, expected_without_subset,
- "Failed for {desc} with subset disabled"
- );
- }
-
- Ok(())
- }
-
- #[test]
- fn test_partitioning_empty_hash() -> Result<()> {
- let fixture = PartitioningTestFixture::int64(&["a"])?;
-
- let test_cases = vec![
(
- "KeyPartitioned([a]) satisfied by Hash([])",
+ "empty hash partitioning",
Partitioning::Hash(vec![], 4),
fixture.key_distribution([0]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
(
- "KeyPartitioned([]) satisfied by Hash([a])",
+ "empty key distribution",
fixture.hash_partitioning([0], 4),
Distribution::KeyPartitioned(vec![]),
PartitioningSatisfaction::NotSatisfied,
PartitioningSatisfaction::NotSatisfied,
),
- (
- "KeyPartitioned([]) satisfied by Hash([])",
- Partitioning::Hash(vec![], 4),
- Distribution::KeyPartitioned(vec![]),
- PartitioningSatisfaction::NotSatisfied,
- PartitioningSatisfaction::NotSatisfied,
- ),
];
for (desc, partition, required, expected_with_subset, expected_without_subset) in
test_cases
{
- let result = partition.satisfaction(&required, &fixture.eq_properties, true);
- assert_eq!(
- result, expected_with_subset,
- "Failed for {desc} with subset enabled"
- );
-
- let result = partition.satisfaction(&required, &fixture.eq_properties, false);
- assert_eq!(
- result, expected_without_subset,
- "Failed for {desc} with subset disabled"
+ assert_satisfaction(
+ desc,
+ &partition,
+ &required,
+ &fixture.eq_properties,
+ expected_with_subset,
+ expected_without_subset,
);
}
@@ -1230,15 +1074,65 @@ mod tests {
}
#[test]
- fn test_multi_partition_range_does_not_satisfy_hash_distribution() -> Result<()> {
- let fixture = PartitioningTestFixture::int64(&["a", "b"])?;
- let range_partitioning =
+ fn range_partitioning_key_distribution_satisfaction() -> Result<()> {
+ let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?;
+ let range_a = fixture.range_partitioning([0], vec![int_split_point([10])]);
+ let range_ab =
fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]);
- let required = fixture.key_distribution([0, 1]);
- assert_eq!(
- range_partitioning.satisfaction(&required, &fixture.eq_properties, false),
- PartitioningSatisfaction::NotSatisfied
+ assert_satisfaction(
+ "exact single key",
+ &range_a,
+ &fixture.key_distribution([0]),
+ &fixture.eq_properties,
+ PartitioningSatisfaction::Exact,
+ PartitioningSatisfaction::Exact,
+ );
+ assert_satisfaction(
+ "exact compound key",
+ &range_ab,
+ &fixture.key_distribution([0, 1]),
+ &fixture.eq_properties,
+ PartitioningSatisfaction::Exact,
+ PartitioningSatisfaction::Exact,
+ );
+ assert_satisfaction(
+ "subset key",
+ &range_a,
+ &fixture.key_distribution([0, 1]),
+ &fixture.eq_properties,
+ PartitioningSatisfaction::Subset,
+ PartitioningSatisfaction::NotSatisfied,
+ );
+ assert_satisfaction(
+ "incompatible key",
+ &range_a,
+ &fixture.key_distribution([1]),
+ &fixture.eq_properties,
+ PartitioningSatisfaction::NotSatisfied,
+ PartitioningSatisfaction::NotSatisfied,
+ );
+
+ let mut eq_properties = fixture.eq_properties.clone();
+ eq_properties.add_equal_conditions(fixture.col(0), fixture.col(2))?;
+ assert_satisfaction(
+ "equivalent subset key",
+ &range_a,
+ &fixture.key_distribution([1, 2]),
+ &eq_properties,
+ PartitioningSatisfaction::Subset,
+ PartitioningSatisfaction::NotSatisfied,
+ );
+
+ let mut eq_properties = fixture.eq_properties.clone();
+ eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?;
+ assert_satisfaction(
+ "equivalent exact key",
+ &range_a,
+ &fixture.key_distribution([1]),
+ &eq_properties,
+ PartitioningSatisfaction::Exact,
+ PartitioningSatisfaction::Exact,
);
Ok(())
diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
index 55769fce01..952aae9846 100644
--- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
@@ -35,7 +35,7 @@ use std::sync::Arc;
use crate::output_requirements::OutputRequirementExec;
use crate::utils::{
add_sort_above_with_check, is_coalesce_partitions, is_repartition,
- is_sort_preserving_merge, range_partitioning_satisfies_key_partitioning,
+ is_sort_preserving_merge,
};
use arrow::compute::SortOptions;
@@ -698,18 +698,13 @@ fn add_roundrobin_on_top(
}
}
-// 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
-// KeyPartitioned generally, this check keeps preserve_file_partitions from
-// inserting RoundRobin between a reusable Range input and the partial aggregate.
-fn partial_aggregate_preserves_reusable_partitioning(
+// Partial aggregates require unspecified input distribution, but their output
+// may already satisfy the final aggregate's key distribution because partial
+// aggregation preserves/projects input partitioning. Keep that reusable output
+// partitioning intact when preserve_file_partitions would otherwise insert
+// RoundRobin below the partial aggregate.
+fn partial_aggregate_output_satisfies_final_partitioning(
plan: &Arc<dyn ExecutionPlan>,
- child: &Arc<dyn ExecutionPlan>,
allow_subset_satisfy_partitioning: bool,
) -> bool {
let Some(aggregate) = plan.downcast_ref::<AggregateExec>() else {
@@ -722,24 +717,15 @@ fn partial_aggregate_preserves_reusable_partitioning(
return false;
}
- let group_exprs = aggregate.group_expr().input_exprs();
- let output_partitioning = child.output_partitioning();
- let eq_properties = child.equivalence_properties();
- let key_distribution = Distribution::KeyPartitioned(group_exprs.clone());
+ let key_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr());
- output_partitioning
+ plan.output_partitioning()
.satisfaction(
&key_distribution,
- eq_properties,
+ plan.equivalence_properties(),
allow_subset_satisfy_partitioning,
)
.is_satisfied()
- || range_partitioning_satisfies_key_partitioning(
- output_partitioning,
- &group_exprs,
- eq_properties,
- allow_subset_satisfy_partitioning,
- )
}
/// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator
@@ -1308,9 +1294,8 @@ pub fn ensure_distribution(
let preserve_partial_aggregate_partitioning =
preserve_file_partition_threshold_met
- && partial_aggregate_preserves_reusable_partitioning(
+ && partial_aggregate_output_satisfies_final_partitioning(
&plan,
- &child.plan,
allow_subset_satisfy_partitioning,
);
diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs
index 1fbf8c6fb7..04229e1cc2 100644
--- a/datafusion/physical-optimizer/src/utils.rs
+++ b/datafusion/physical-optimizer/src/utils.rs
@@ -18,10 +18,7 @@
use std::sync::Arc;
use datafusion_common::Result;
-use datafusion_physical_expr::{
- Distribution, EquivalenceProperties, LexOrdering, LexRequirement, Partitioning,
- PhysicalExpr, physical_exprs_equal,
-};
+use datafusion_physical_expr::{Distribution, LexOrdering, LexRequirement};
use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion_physical_plan::repartition::RepartitionExec;
@@ -161,60 +158,6 @@ pub fn is_repartition(plan: &Arc<dyn ExecutionPlan>) -> bool {
plan.is::<RepartitionExec>()
}
-/// TODO: remove once Range generally satisfies KeyPartitioned requirements
-/// through Partitioning::satisfaction.
-/// See <https://github.com/apache/datafusion/issues/23266>.
-///
-/// Checks whether range partitioning satisfies a key partitioning requirement.
-/// This is intentionally separate from general partitioning satisfaction while
-/// range reuse is rolled out operator by operator.
-pub(crate) fn range_partitioning_satisfies_key_partitioning(
- partitioning: &Partitioning,
- required_exprs: &[Arc<dyn PhysicalExpr>],
- eq_properties: &EquivalenceProperties,
- allow_subset: bool,
-) -> bool {
- match partitioning {
- Partitioning::Range(range) => {
- let partition_exprs = range
- .ordering()
- .iter()
- .map(|sort_expr| Arc::clone(&sort_expr.expr))
- .collect::<Vec<_>>();
-
- if partition_exprs.is_empty() || required_exprs.is_empty() {
- return false;
- }
-
- let eq_group = eq_properties.eq_group();
- let normalized_partition_exprs = partition_exprs
- .iter()
- .map(|expr| eq_group.normalize_expr(Arc::clone(expr)))
- .collect::<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 true;
- }
-
- allow_subset
- && normalized_partition_exprs.len() < normalized_required_exprs.len()
- && normalized_partition_exprs.iter().all(|partition_expr| {
- normalized_required_exprs
- .iter()
- .any(|required_expr| partition_expr.eq(required_expr))
- })
- }
- _ => false,
- }
-}
-
/// Checks whether the given operator is a limit;
/// i.e. either a [`LocalLimitExec`] or a [`GlobalLimitExec`].
pub fn is_limit(plan: &Arc<dyn ExecutionPlan>) -> bool {
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 732da32ab0..30e0ad2469 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -1934,7 +1934,7 @@ impl ExecutionPlan for AggregateExec {
}
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
- let requirements = InputDistributionRequirements::new(match &self.mode {
+ InputDistributionRequirements::new(match &self.mode {
AggregateMode::Partial | AggregateMode::PartialReduce => {
vec![Distribution::UnspecifiedDistribution]
}
@@ -1944,15 +1944,7 @@ impl ExecutionPlan for AggregateExec {
AggregateMode::Final | AggregateMode::Single => {
vec![Distribution::SinglePartition]
}
- });
- match &self.mode {
- AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned
- if !self.group_by.has_grouping_set() =>
- {
- requirements.allow_range_satisfaction_for_key_partitioning()
- }
- _ => requirements,
- }
+ })
}
fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
diff --git a/datafusion/physical-plan/src/distribution_requirements.rs b/datafusion/physical-plan/src/distribution_requirements.rs
index 9c7a1336c0..6405b1f121 100644
--- a/datafusion/physical-plan/src/distribution_requirements.rs
+++ b/datafusion/physical-plan/src/distribution_requirements.rs
@@ -17,13 +17,8 @@
//! Input distribution requirements for physical execution plans.
-use std::sync::Arc;
-
use datafusion_common::{Result, internal_err};
-use datafusion_physical_expr::{
- Distribution, EquivalenceProperties, Partitioning, PartitioningSatisfaction,
- PhysicalExpr, physical_exprs_equal,
-};
+use datafusion_physical_expr::{Distribution, Partitioning, PartitioningSatisfaction};
use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel};
@@ -98,10 +93,7 @@ impl InputDistributionRequirements {
pub fn new(per_child: Vec<Distribution>) -> Self {
let children = per_child
.into_iter()
- .map(|distribution| ChildDistributionRequirement {
- distribution,
- satisfaction: InputDistributionSatisfaction::Default,
- })
+ .map(|distribution| ChildDistributionRequirement { distribution })
.collect();
Self {
@@ -176,8 +168,7 @@ impl InputDistributionRequirements {
);
};
- Ok(requirement.satisfaction.satisfaction(
- child.output_partitioning(),
+ Ok(child.output_partitioning().satisfaction(
&requirement.distribution,
child.equivalence_properties(),
options.allow_subset(),
@@ -208,30 +199,6 @@ impl InputDistributionRequirements {
Ok(co_partitioned.clone())
}
- /// TODO: remove this temporary bridge once [`Partitioning::Range`]
- /// generally satisfies [`Distribution::KeyPartitioned`] through
- /// [`Partitioning::satisfaction`].
- /// <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,
@@ -301,10 +268,8 @@ impl InputDistributionRequirements {
let first = children[first_idx];
let first_partitioning = first.output_partitioning();
- if !first_requirement
- .satisfaction
+ if !first_partitioning
.satisfaction(
- first_partitioning,
&first_requirement.distribution,
first.equivalence_properties(),
false,
@@ -317,19 +282,16 @@ impl InputDistributionRequirements {
for &child_idx in co_partitioned.iter().skip(1) {
let requirement = &self.children[child_idx];
let child = children[child_idx];
- if !requirement
- .satisfaction
+ if !child
+ .output_partitioning()
.satisfaction(
- child.output_partitioning(),
&requirement.distribution,
child.equivalence_properties(),
false,
)
.is_satisfied()
|| !compatible_co_partitioning_layout(
- first_requirement,
first_partitioning,
- requirement,
child.output_partitioning(),
)
{
@@ -345,60 +307,6 @@ impl InputDistributionRequirements {
#[derive(Debug, Clone)]
struct ChildDistributionRequirement {
distribution: Distribution,
- satisfaction: InputDistributionSatisfaction,
-}
-
-/// TODO: remove this temporary bridge once [`Partitioning::Range`]
-/// generally satisfies [`Distribution::KeyPartitioned`] through
-/// [`Partitioning::satisfaction`].
-/// <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(
@@ -421,62 +329,8 @@ fn validate_child_index(
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
@@ -491,12 +345,7 @@ fn compatible_co_partitioning_layout(
match (first_partitioning, other_partitioning) {
(Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true,
- (Partitioning::Range(left), Partitioning::Range(right))
- if first.satisfaction
- == InputDistributionSatisfaction::AllowRangeKeyPartitioning
- && other.satisfaction
- == InputDistributionSatisfaction::AllowRangeKeyPartitioning =>
- {
+ (Partitioning::Range(left), Partitioning::Range(right)) => {
left.split_points() == right.split_points()
&& left.ordering().len() == right.ordering().len()
&& left
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index e941bb0898..c746aba028 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -1271,7 +1271,7 @@ impl ExecutionPlan for HashJoinExec {
}
fn input_distribution_requirements(&self) -> InputDistributionRequirements {
- let requirements = match self.mode {
+ match self.mode {
PartitionMode::Partitioned => {
let (left_expr, right_expr) = self
.on
@@ -1291,14 +1291,6 @@ impl ExecutionPlan for HashJoinExec {
Distribution::UnspecifiedDistribution,
Distribution::UnspecifiedDistribution,
]),
- };
-
- if self.mode == PartitionMode::Partitioned {
- // Compatible Range inputs co-locate equal join keys, which
- // satisfies the co-partitioned requirement for hash joins.
- requirements.allow_range_satisfaction_for_key_partitioning()
- } else {
- requirements
}
}
diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
index 304f547bdd..f15820e6da 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
@@ -424,7 +424,6 @@ impl ExecutionPlan for SortMergeJoinExec {
Distribution::KeyPartitioned(left_expr),
Distribution::KeyPartitioned(right_expr),
])
- .allow_range_satisfaction_for_key_partitioning()
}
fn required_input_ordering(&self) -> Vec<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 eb781e633c..b95e5e1e3d 100644
--- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
+++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
@@ -427,7 +427,6 @@ impl ExecutionPlan for SymmetricHashJoinExec {
Distribution::KeyPartitioned(left_expr),
Distribution::KeyPartitioned(right_expr),
])
- .allow_range_satisfaction_for_key_partitioning()
}
StreamJoinPartitionMode::SinglePartition => {
InputDistributionRequirements::new(vec![
diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs
index c250130341..730440a429 100644
--- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs
+++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs
@@ -355,7 +355,6 @@ impl ExecutionPlan for PartitionedTopKExec {
crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
partition_exprs,
)])
- .allow_range_satisfaction_for_key_partitioning()
}
fn maintains_input_order(&self) -> Vec<bool> {
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 97cafd24c0..3ca612bbdb 100644
--- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
+++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
@@ -333,7 +333,6 @@ impl ExecutionPlan for BoundedWindowAggExec {
InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
self.partition_keys(),
)])
- .allow_range_satisfaction_for_key_partitioning()
}
}
diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs
index 4e8dbc06f0..3eb8edd298 100644
--- a/datafusion/physical-plan/src/windows/window_agg_exec.rs
+++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs
@@ -241,7 +241,6 @@ impl ExecutionPlan for WindowAggExec {
InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
self.partition_keys(),
)])
- .allow_range_satisfaction_for_key_partitioning()
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org