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/02 18:09:58 UTC
(datafusion) 01/01: Aggregations Support `Partitioning::Range` (#23239)
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-23239-457fc532f162512deb7e3fe2368089a9fb5b896f
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 044a85c57471af68cb819379d280b3fa3017707c
Author: Gene Bordegaray <ge...@datadoghq.com>
AuthorDate: Thu Jul 2 14:09:34 2026 -0400
Aggregations Support `Partitioning::Range` (#23239)
## 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 #23191.
- Related discussion: #23184, #23236.
## 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.
-->
Range partitioning can satisfy aggregate hash partitioning: equal group
keys are already partitioned, even though the partitioning is not
hash-based.
This is the first unary-operator implementation from the range
partitioning discussion before making broader public API changes around
`HashPartitioned` / `KeyPartitioned`.
## 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.
-->
- Let compatible range partitioning satisfy aggregate hash distribution
requirements in `EnforceDistribution`
- Keep this private to aggregate planning for now to not make public API
changes to `Distribution` enum variants yet until more operators are
supported
## 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.
-->
Yes. Range-partitioned aggregate plans can now avoid hash
repartitioning.
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---
.../physical_optimizer/enforce_distribution.rs | 125 ++++++++++++-
.../ensure_requirements/enforce_distribution.rs | 195 +++++++++++--------
.../physical-optimizer/src/sanity_checker.rs | 29 ++-
datafusion/physical-optimizer/src/utils.rs | 78 +++++++-
.../sqllogictest/test_files/range_partitioning.slt | 207 +++++++++++++++++++--
5 files changed, 535 insertions(+), 99 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
index 9424322396..cab8dc67f9 100644
--- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
+++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
@@ -59,7 +59,9 @@ use datafusion_physical_plan::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy,
};
-use datafusion_physical_expr::Distribution;
+use datafusion_physical_expr::{
+ Distribution, Partitioning, RangePartitioning, SplitPoint,
+};
use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion_physical_plan::execution_plan::ExecutionPlan;
use datafusion_physical_plan::expressions::col;
@@ -325,6 +327,47 @@ fn parquet_exec_multiple_sorted(
DataSourceExec::from_data_source(config)
}
+fn parquet_exec_with_output_partitioning(
+ output_partitioning: Partitioning,
+) -> Arc<DataSourceExec> {
+ let file_groups = (0..output_partitioning.partition_count())
+ .map(|partition| {
+ FileGroup::new(vec![PartitionedFile::new(format!("p{partition}"), 100)])
+ })
+ .collect::<Vec<_>>();
+
+ let config = FileScanConfigBuilder::new(
+ ObjectStoreUrl::parse("test:///").unwrap(),
+ Arc::new(ParquetSource::new(schema())),
+ )
+ .with_file_groups(file_groups)
+ .with_output_partitioning(Some(output_partitioning))
+ .build();
+
+ DataSourceExec::from_data_source(config)
+}
+
+fn range_partitioning(
+ column: &str,
+ split_values: impl IntoIterator<Item = i64>,
+ options: SortOptions,
+) -> Result<Partitioning> {
+ let ordering: LexOrdering = [PhysicalSortExpr {
+ expr: col(column, &schema())?,
+ options,
+ }]
+ .into();
+ let split_points = split_values
+ .into_iter()
+ .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))]))
+ .collect::<Vec<_>>();
+
+ Ok(Partitioning::Range(RangePartitioning::try_new(
+ ordering,
+ split_points,
+ )?))
+}
+
fn csv_exec() -> Arc<DataSourceExec> {
csv_exec_with_sort(vec![])
}
@@ -700,6 +743,86 @@ impl TestConfig {
}
}
+#[test]
+fn range_aggregate_reuses_range_partitioning() -> Result<()> {
+ let input = parquet_exec_with_output_partitioning(range_partitioning(
+ "a",
+ [10, 20, 30],
+ SortOptions::default(),
+ )?);
+ let aggregate =
+ aggregate_exec_with_alias(input, vec![("a".to_string(), "a".to_string())]);
+
+ let plan = TestConfig::default()
+ .with_query_execution_partitions(4)
+ .to_plan(aggregate, &DISTRIB_DISTRIB_SORT);
+
+ assert_plan!(
+ plan,
+ @r"
+ AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]
+ AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]
+ DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+ "
+ );
+
+ Ok(())
+}
+
+#[test]
+fn range_grouping_set_aggregate_rehashes_with_grouping_id() -> Result<()> {
+ let input = parquet_exec_with_output_partitioning(range_partitioning(
+ "a",
+ [10, 20],
+ SortOptions::default(),
+ )?);
+ let input_schema = input.schema();
+ let group_by = PhysicalGroupBy::new(
+ vec![
+ (col("a", &input_schema)?, "a".to_string()),
+ (col("b", &input_schema)?, "b".to_string()),
+ ],
+ vec![
+ (lit(ScalarValue::Int64(None)), "a".to_string()),
+ (lit(ScalarValue::Int64(None)), "b".to_string()),
+ ],
+ vec![vec![false, true], vec![false, false]],
+ true,
+ );
+ let partial = Arc::new(AggregateExec::try_new(
+ AggregateMode::Partial,
+ group_by.clone(),
+ vec![],
+ vec![],
+ input,
+ Arc::clone(&input_schema),
+ )?);
+ let aggregate = Arc::new(AggregateExec::try_new(
+ AggregateMode::FinalPartitioned,
+ group_by.as_final(),
+ vec![],
+ vec![],
+ Arc::clone(&partial) as _,
+ partial.schema(),
+ )?);
+
+ let plan = TestConfig::default()
+ .with_query_execution_partitions(3)
+ .to_plan(aggregate, &DISTRIB_DISTRIB_SORT);
+
+ assert_plan!(
+ plan,
+ @r"
+ AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, __grouping_id@2 as __grouping_id], aggr=[]
+ RepartitionExec: partitioning=Hash([a@0, b@1, __grouping_id@2], 3), input_partitions=3
+ AggregateExec: mode=Partial, gby=[(a@0 as a, NULL as b), (a@0 as a, b@1 as b)], aggr=[]
+ DataSourceExec: file_groups={3 groups: [[p0], [p1], [p2]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20)], 3), file_type=parquet
+ "
+ );
+
+ Ok(())
+}
+
#[test]
fn multi_hash_joins() -> Result<()> {
let left = parquet_exec();
diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
index 305d811fa2..7f3a63ed91 100644
--- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
@@ -34,8 +34,9 @@ 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,
+ add_sort_above_with_check, aggregate_can_reuse_range_partitioning,
+ is_coalesce_partitions, is_repartition, is_sort_preserving_merge,
+ range_partitioning_satisfies_key_partitioning,
};
use arrow::compute::SortOptions;
@@ -699,72 +700,47 @@ fn add_roundrobin_on_top(
}
}
-/// Adds a hash repartition operator:
-/// - to increase parallelism, and/or
-/// - to satisfy requirements of the subsequent operators.
-///
-/// Repartition(Hash) is added on top of operator `input`.
-///
-/// # Arguments
-///
-/// * `input`: Current node.
-/// * `hash_exprs`: Stores Physical Exprs that are used during hashing.
-/// * `n_target`: desired target partition number, if partition number of the
-/// current executor is less than this value. Partition number will be increased.
-/// * `allow_subset_satisfy_partitioning`: Whether to allow subset partitioning logic in satisfaction checks.
-/// Set to `false` for partitioned hash joins to ensure exact hash matching.
-///
-/// # Returns
-///
-/// A [`Result`] object that contains new execution plan where the desired
-/// distribution is satisfied by adding a Hash repartition.
-fn add_hash_on_top(
- input: DistributionContext,
- hash_exprs: Vec<Arc<dyn PhysicalExpr>>,
- n_target: usize,
+// TODO: remove this private helper once Range generally satisfies
+// KeyPartitioned requirements through Partitioning::satisfaction.
+// See <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(
+ plan: &Arc<dyn ExecutionPlan>,
+ child: &Arc<dyn ExecutionPlan>,
allow_subset_satisfy_partitioning: bool,
-) -> Result<DistributionContext> {
- // Early return if hash repartition is unnecessary
- // `RepartitionExec: partitioning=Hash([...], 1), input_partitions=1` is unnecessary.
- if n_target == 1 && input.plan.output_partitioning().partition_count() == 1 {
- return Ok(input);
- }
-
- let dist = Distribution::KeyPartitioned(hash_exprs);
- let satisfaction = input.plan.output_partitioning().satisfaction(
- &dist,
- input.plan.equivalence_properties(),
- allow_subset_satisfy_partitioning,
- );
-
- // Add hash repartitioning when:
- // - When subset satisfaction is enabled (current >= threshold): only repartition if not satisfied
- // - When below threshold (current < threshold): repartition if expressions don't match OR to increase parallelism
- let needs_repartition = if allow_subset_satisfy_partitioning {
- !satisfaction.is_satisfied()
- } else {
- !satisfaction.is_satisfied()
- || n_target > input.plan.output_partitioning().partition_count()
+) -> bool {
+ let Some(aggregate) = plan.downcast_ref::<AggregateExec>() else {
+ return false;
};
-
- if needs_repartition {
- // When there is an existing ordering, we preserve ordering during
- // repartition. This will be rolled back in the future if any of the
- // following conditions is true:
- // - Preserving ordering is not helpful in terms of satisfying ordering
- // requirements.
- // - Usage of order preserving variants is not desirable (per the flag
- // `config.optimizer.prefer_existing_sort`).
- let partitioning = dist.create_partitioning(n_target);
- let repartition =
- RepartitionExec::try_new(Arc::clone(&input.plan), partitioning)?
- .with_preserve_order();
- let plan = Arc::new(repartition) as _;
-
- return Ok(DistributionContext::new(plan, true, vec![input]));
+ if aggregate.mode() != &AggregateMode::Partial
+ || aggregate.group_expr().is_empty()
+ || aggregate.group_expr().has_grouping_set()
+ {
+ return false;
}
- Ok(input)
+ let group_exprs = aggregate.group_expr().input_exprs();
+ let output_partitioning = child.output_partitioning();
+ let eq_properties = child.equivalence_properties();
+ let key_distribution = Distribution::KeyPartitioned(group_exprs.clone());
+
+ output_partitioning
+ .satisfaction(
+ &key_distribution,
+ eq_properties,
+ allow_subset_satisfy_partitioning,
+ )
+ .is_satisfied()
+ || range_partitioning_satisfies_key_partitioning(
+ output_partitioning,
+ &group_exprs,
+ eq_properties,
+ allow_subset_satisfy_partitioning,
+ )
}
/// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator
@@ -1197,21 +1173,14 @@ pub fn ensure_distribution(
hash_necessary,
},
)| {
- let increases_partition_count =
- child.plan.output_partitioning().partition_count() < target_partitions;
-
- let add_roundrobin = enable_round_robin
- // Operator benefits from partitioning (e.g. filter):
- && roundrobin_beneficial
- && roundrobin_beneficial_stats
- // Unless partitioning increases the partition count, it is not beneficial:
- && increases_partition_count;
-
// Allow subset satisfaction when:
// 1. Current partition count >= threshold
// 2. Not a partitioned join since must use exact hash matching for joins
// 3. Not a grouping set aggregate (requires exact hash including __grouping_id)
let current_partitions = child.plan.output_partitioning().partition_count();
+ let preserve_file_partition_threshold_met =
+ config.optimizer.preserve_file_partitions > 0
+ && current_partitions >= config.optimizer.preserve_file_partitions;
// Check if the hash partitioning requirement includes __grouping_id column.
// Grouping set aggregates (ROLLUP, CUBE, GROUPING SETS) require exact hash
@@ -1232,11 +1201,29 @@ pub fn ensure_distribution(
// partitioning to the optimizer. Respect it when the only
// reason to repartition would be to increase partition count
// beyond the preserved file-group count.
- || (config.optimizer.preserve_file_partitions > 0
+ || (preserve_file_partition_threshold_met
&& current_partitions < target_partitions))
&& !is_partitioned_join
&& !requires_grouping_id;
+ let increases_partition_count = current_partitions < target_partitions;
+
+ let preserve_partial_aggregate_partitioning =
+ preserve_file_partition_threshold_met
+ && partial_aggregate_preserves_reusable_partitioning(
+ &plan,
+ &child.plan,
+ allow_subset_satisfy_partitioning,
+ );
+
+ let add_roundrobin = enable_round_robin
+ // Operator benefits from partitioning (e.g. filter):
+ && roundrobin_beneficial
+ && roundrobin_beneficial_stats
+ // Unless partitioning increases the partition count, it is not beneficial:
+ && increases_partition_count
+ && !preserve_partial_aggregate_partitioning;
+
// When `repartition_file_scans` is set, attempt to increase
// parallelism at the source.
//
@@ -1258,15 +1245,59 @@ pub fn ensure_distribution(
}
Distribution::HashPartitioned(exprs)
| 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,
+ )
+ .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;
+
+ // When subset satisfaction is enabled, preserve an
+ // already-satisfying partitioning. Otherwise, hash
+ // repartition may also increase parallelism.
+ let needs_hash_repartition = if allow_subset_satisfy_partitioning {
+ !partitioning_satisfied
+ } else {
+ !partitioning_satisfied || target_partitions > child_partitions
+ };
+ let should_add_hash_repartition =
+ hash_necessary && needs_hash_repartition;
+
// See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background
// When inserting hash is necessary to satisfy hash requirement, insert hash repartition.
- if hash_necessary {
- child = add_hash_on_top(
- child,
- exprs.to_vec(),
- target_partitions,
- allow_subset_satisfy_partitioning,
- )?;
+ if should_add_hash_repartition {
+ // When there is an existing ordering, we preserve ordering during
+ // repartition. This will be rolled back in the future if any of the
+ // following conditions is true:
+ // - Preserving ordering is not helpful in terms of satisfying ordering
+ // requirements.
+ // - Usage of order preserving variants is not desirable (per the flag
+ // `config.optimizer.prefer_existing_sort`).
+ let partitioning = Distribution::KeyPartitioned(exprs.to_vec())
+ .create_partitioning(target_partitions);
+ let repartition = RepartitionExec::try_new(
+ Arc::clone(&child.plan),
+ partitioning,
+ )?
+ .with_preserve_order();
+ let plan = Arc::new(repartition) as _;
+ child = DistributionContext::new(plan, true, vec![child]);
}
}
Distribution::UnspecifiedDistribution => {
diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs
index 40c6245d89..936bc8271a 100644
--- a/datafusion/physical-optimizer/src/sanity_checker.rs
+++ b/datafusion/physical-optimizer/src/sanity_checker.rs
@@ -24,6 +24,7 @@
use std::sync::Arc;
use datafusion_common::Result;
+use datafusion_physical_expr::Distribution;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_common::config::{ConfigOptions, OptimizerOptions};
@@ -35,6 +36,9 @@ use datafusion_physical_plan::joins::SymmetricHashJoinExec;
use datafusion_physical_plan::{ExecutionPlanProperties, get_plan_string};
use crate::PhysicalOptimizerRule;
+use crate::utils::{
+ aggregate_can_reuse_range_partitioning, range_partitioning_satisfies_key_partitioning,
+};
use datafusion_physical_expr_common::sort_expr::format_physical_sort_requirement_list;
use itertools::izip;
@@ -136,6 +140,10 @@ fn is_prunable(join: &SymmetricHashJoinExec) -> bool {
/// Ensures that the plan is pipeline friendly and the order and
/// distribution requirements from its children are satisfied.
+#[expect(
+ deprecated,
+ reason = "HashPartitioned is accepted during the KeyPartitioned migration"
+)]
pub fn check_plan_sanity(
plan: &Arc<dyn ExecutionPlan>,
optimizer_options: &OptimizerOptions,
@@ -162,11 +170,26 @@ pub fn check_plan_sanity(
}
}
- if !child
+ let child_satisfies_distribution = child
.output_partitioning()
.satisfaction(&dist_req, child_eq_props, true)
- .is_satisfied()
- {
+ .is_satisfied();
+ let range_satisfies_aggregate_distribution =
+ aggregate_can_reuse_range_partitioning(plan)
+ && match &dist_req {
+ Distribution::HashPartitioned(exprs)
+ | Distribution::KeyPartitioned(exprs) => {
+ range_partitioning_satisfies_key_partitioning(
+ child.output_partitioning(),
+ exprs,
+ child_eq_props,
+ true,
+ )
+ }
+ _ => false,
+ };
+
+ if !(child_satisfies_distribution || range_satisfies_aggregate_distribution) {
let plan_str = get_plan_string(plan);
return plan_err!(
"Plan: {:?} does not satisfy distribution requirements: {}. Child-{} output partitioning: {}",
diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs
index 04229e1cc2..e7a19ca0b0 100644
--- a/datafusion/physical-optimizer/src/utils.rs
+++ b/datafusion/physical-optimizer/src/utils.rs
@@ -18,7 +18,11 @@
use std::sync::Arc;
use datafusion_common::Result;
-use datafusion_physical_expr::{Distribution, LexOrdering, LexRequirement};
+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;
@@ -158,6 +162,78 @@ 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,
+ }
+}
+
+/// 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/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt
index 2b7a2cfdf4..7904f92310 100644
--- a/datafusion/sqllogictest/test_files/range_partitioning.slt
+++ b/datafusion/sqllogictest/test_files/range_partitioning.slt
@@ -28,19 +28,22 @@ set datafusion.explain.physical_plan_only = true;
##########
# TEST 1: Aggregate on Range Partition Column
-# Scanning range_key preserves source Range partitioning metadata.
-# Planning still inserts Hash repartitioning today; later optimizer PRs can
-# use this baseline to show when the repartition is removed.
+# With subset threshold met and preserve-file disabled, Range([range_key])
+# satisfies the aggregate key and avoids repartitioning.
##########
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 4;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
query TT
EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key;
----
physical_plan
-01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)]
-02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4
-03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)]
-04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, proje [...]
+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 [...]
query II
SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key ORDER BY range_key;
@@ -57,11 +60,16 @@ SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key ORDER BY
##########
# TEST 2: Aggregate on Non-Range Column
-# Projecting away range_key means the scan output no longer contains the
-# expression needed to describe range partitioning, so it reports
-# UnknownPartitioning with the same partition count.
+# With subset threshold met and preserve-file disabled, grouping on a non-range
+# key cannot reuse Range([range_key]) and requires hash repartitioning.
##########
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 4;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
query TT
EXPLAIN SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key;
----
@@ -79,7 +87,182 @@ SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key O
##########
-# TEST 3: Join on Range Partition Column
+# TEST 3: Aggregate Reuses Range Subset Partitioning
+# With subset threshold met and preserve-file disabled, Range([range_key])
+# satisfies grouping by (range_key, non_range_key).
+##########
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 4;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
+query TT
+EXPLAIN SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key;
+----
+physical_plan
+01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)]
+02)--DataSourceExec: file_groups={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 [...]
+
+query III
+SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key ORDER BY range_key, non_range_key;
+----
+1 1 10
+5 2 50
+10 1 100
+15 2 150
+20 1 200
+25 2 250
+30 1 300
+35 2 350
+
+
+##########
+# TEST 4: Exact Range Aggregate Below Subset Threshold
+# Even when subset satisfaction is disabled, exact Range([range_key])
+# satisfies GROUP BY range_key when repartitioning would not increase
+# partition count.
+##########
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 5;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
+query TT
+EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key;
+----
+physical_plan
+01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)]
+02)--DataSourceExec: file_groups={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 [...]
+
+
+##########
+# TEST 5: Range Subset Aggregate Rehashes Below Subset Threshold
+# Range([range_key]) is only a subset of GROUP BY (range_key, non_range_key),
+# so it should not satisfy the aggregate key when subset satisfaction is
+# disabled.
+##########
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 5;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
+query TT
+EXPLAIN SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key;
+----
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)]
+02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4
+03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)]
+04)------DataSourceExec: file_groups={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 [...]
+
+
+##########
+# TEST 6: Aggregate Rehashes Below Subset Threshold
+# With subset threshold 5 and only 4 input partitions, planning repartitions
+# to increase parallelism instead of reusing Range partitioning.
+##########
+
+statement ok
+set datafusion.execution.target_partitions = 5;
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 5;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
+query TT
+EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key;
+----
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)]
+02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5
+03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)]
+04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4
+05)--------DataSourceExec: file_groups={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 [...]
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+reset datafusion.optimizer.subset_repartition_threshold;
+
+
+##########
+# TEST 7: Aggregate Preserves Range When Preserve File Threshold Met
+# With preserve-file threshold 1 and 4 input partitions, Range is preserved
+# even though target_partitions is 5.
+##########
+
+statement ok
+set datafusion.execution.target_partitions = 5;
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 4;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 1;
+
+query TT
+EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key;
+----
+physical_plan
+01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)]
+02)--DataSourceExec: file_groups={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 [...]
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+reset datafusion.optimizer.preserve_file_partitions;
+
+
+##########
+# TEST 8: Aggregate Rehashes When Preserve File Threshold Not Met
+# With preserve-file threshold 5 and only 4 input partitions, planning can
+# repartition to increase parallelism.
+##########
+
+statement ok
+set datafusion.execution.target_partitions = 5;
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 4;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 5;
+
+query TT
+EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key;
+----
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)]
+02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5
+03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)]
+04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4
+05)--------DataSourceExec: file_groups={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 [...]
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+reset datafusion.optimizer.preserve_file_partitions;
+
+
+##########
+# 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.
@@ -91,7 +274,7 @@ FROM range_partitioned l
JOIN range_partitioned r ON l.range_key = r.range_key;
##########
-# TEST 4: Union of Range Partitioned Inputs
+# 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.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org