Aggregations Support `Partitioning::Range` by gene-bordegaray · Pull Request #23239 · apache/datafusion

gene-bordegaray · 2026-06-29T13:35:46Z

Which issue does this PR close?

Closes Allow Partitioning::Range to satisfy aggregation Distribution::KeyPartitioned requirements #23191.
Related discussion: Support co-partitioned range inner equi joins #23184, Replace / rename HashPartitioned distribution as KeyPartitioned #23236.

Rationale for this change

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?

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?

Yes.

Are there any user-facing changes?

Yes. Range-partitioned aggregate plans can now avoid hash repartitioning.

gene-bordegaray · 2026-06-29T14:00:32Z

cc: @2010YOUY01 @gabotechs @stuhood

gabotechs · 2026-06-29T15:20:46Z

+    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


🤔 I think we don't need the Partial there right? we should be fine with just a FinalPartitioned aggregation

this happens in a later optimizer rule that collapses these partial -> finals where approacpriate: datafusion/physical-optimizer/src/combine_partial_final_agg.rs

that is why it shows up correctly in the slt tests 👍

Is it possible to move the unit-test coverage added in this file to end-to-end SLT tests instead?

It seems the same test goal can still be achieved at the SLT level, and those sql tests should be more stable across optimizer refactors.

For example, whether the initial physical plan uses a two-stage aggregation or a single-stage aggregation feels implementation-specific. A future refactor might legitimately change that plan shape, which would require updating these unit tests. At that point, it may be harder to recover the original intent of each assertion, and some coverage could accidentally be lost during the refactor. (while SLT behavior won't change a lot even after aggressive refactors)

Indeed, tweaking the plan like this:

let plan = CombinePartialFinalAggregate::new().optimize(plan, &ConfigOptions::new())?;

Produces the following plan:

AggregateExec: mode=SinglePartitioned, 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

I don't have a strong opinion whether porting this to SLT tests or also leaving them here. If there are things that are covered here but not in SLT, it's probably worth also covering them in SLT.

I prefer to have coverage directly in the enforce_distribution.rs tests but can add end to end coerage as well to SLT 👍. Let me know if that is alrigth with you guys

Thanks, adding slts would be enough. But I would lean toward removing most of the unit tests added here after moving the equivalent coverage to SLT.

I’ll try to explain the reasoning a bit better, since I realize it is uncommon to suggest removing tests.

The main concern is that internal module-level tests create maintenance overhead while providing weaker guarantees:

If the coverage is moved to SLT, we can also test whether this optimizer rule works correctly with other optimizer rules in most cases.

If future features change these unit tests, it is often hard to figure out the test goal of those internal tests, and update them correctly. Especially, I think this optimizer rule is likely to change a lot in the future.

So in practice, I would suggest moving most of the test coverage, especially edge cases, to SLT.

I would only keep module-level unit tests for:

1–2 demo-like tests, since they can catch errors earlier, to make development easier

Cases that are difficult to test with end-to-end tests.

Reference: a DuckDB author trying to convince us to avoid most UTs in favor of e2e tests: https://www.youtube.com/watch?v=BgC79Zt2fPs&t=940s

ah ok ya that makes sense. I think a nice split corresponding with this logic would be something like:

Unit

2 demo tests as smoke tests to catch errors for range specifically in this file

The grouping set test as this would be hard to repro. For context this is needed and was introduced in Fix grouping set subset satisfaction #19853

SLT

set of more end to end nehaavior covering: basic satisfaction, negative on satisfaction, subset, file preservation

2010YOUY01

Thank you! The high-level shape of the PR LGTM. I just need a bit more time to understand what EnforceDistribution is doing before finishing the review — that code looks a little intimidating 😅

2010YOUY01 · 2026-06-30T02:05:46Z

+    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


Is it possible to move the unit-test coverage added in this file to end-to-end SLT tests instead?

It seems the same test goal can still be achieved at the SLT level, and those sql tests should be more stable across optimizer refactors.

For example, whether the initial physical plan uses a two-stage aggregation or a single-stage aggregation feels implementation-specific. A future refactor might legitimately change that plan shape, which would require updating these unit tests. At that point, it may be harder to recover the original intent of each assertion, and some coverage could accidentally be lost during the refactor. (while SLT behavior won't change a lot even after aggressive refactors)

2010YOUY01 · 2026-06-30T02:12:10Z

-            .is_satisfied()
-        {
+            .is_satisfied();
+        let range_satisfies_aggregate_distribution =


I feel this PR is already treating HashPartitioned as KeyPartitioned logically, but the formal renaming PR is left to be done in a future PR 🤔

If that's the case, should we directly implement this logic into range_partitioning.satisfaction()? Otherwise we still have to do it after the formal renaming.

Following on our discussion in #23236, I think we'll maintain the two HashPartitioned(deprecated) and KeyPartitioned and treat them as equal.

Probably that can be done in a preliminary PR.

update: the new/deprecation is done in

Add Distribution::HashPartitioned to Distribution::KeyPartitioned API bridge #23259

Yes, want to hold off on doing general satisfaction until a few operators are implemented privately to ensure we have the correct semantics API contracts fleshed out a but more 👍

I also made an official issue for when this will be done and linked it throughout the PR where we have private helpers for this work: #23266

gabotechs

Approach looks good to me, besides a bit more test coverage and other suggestions, I think this is good.

gabotechs · 2026-06-30T08:13:49Z

+                    let should_add_hash_repartition = hash_necessary
+                        && needs_hash_repartition
+                        && !range_satisfied_for_aggregate;
+


The bool threading in this file is pretty mind-bending, although it seems like it's just how it is, I cannot think of a way of simplifying it...

This PR actually manages that complexity relatively well. I'll think a bit more about it and see if there's a way we can make it easier, but I don't think this is a blocker as long as we have good test coverage.

I know... this rule was the first optimizer rule I spent time in and man, its a hard one to digest. It is on my bucket list to clean this guy up more too

I think it’s good to proceed if someone is confident they fully understand the related implementation and can confirm this complexity is hard to avoid.

Otherwise, it feels a bit risky to rely only on tests here. The code may become harder to evolve if we keep adding more logic without a clearer model.

I understand this rule well and have spent a good amount of time in this. Without huge churn, the complexity is mostly unavoidable but I do think there is room for a separete PR / effort to reshape the logic around enforce distribution but should not be here.

gabotechs · 2026-06-30T08:19:55Z

 # 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.
+# Planning does not need Hash repartitioning because Range partitioning
+# colocates rows with equal range_key values.
 ##########



It seems like we can add a bit more coverage to this file?

For example, adding some positive and negative tests that play with subset satisfaction + range partitioning, and trying to get all the boolean code paths in enforce_distribution.rs to execute.

I have expanded the slt tests much more now 👍

THese are acting ad the main end to end coverage for the range behavior as they pass throughn all the optmizer rules. I left small smike and "speciific" cases to units

gabotechs · 2026-06-30T08:21:47Z

-            .is_satisfied()
-        {
+            .is_satisfied();
+        let range_satisfies_aggregate_distribution =


Following on our discussion in #23236, I think we'll maintain the two HashPartitioned(deprecated) and KeyPartitioned and treat them as equal.

Probably that can be done in a preliminary PR.

gene-bordegaray · 2026-06-30T12:41:58Z

Approach looks good to me, besides a bit more test coverage and other suggestions, I think this is good.

@gabotechs awesome, going to get that preliminary one in first and rebase it, thank you 👍

alamb · 2026-06-30T17:10:44Z

I am happy to help review this too if you think it would be helpful, but I took a quick skim and it looks good to me. Please don't wait for my review to merge it unless you have something specific I can help with

Very excited to see this moving

2010YOUY01 · 2026-07-01T09:20:32Z

+    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


Thanks, adding slts would be enough. But I would lean toward removing most of the unit tests added here after moving the equivalent coverage to SLT.

I’ll try to explain the reasoning a bit better, since I realize it is uncommon to suggest removing tests.

The main concern is that internal module-level tests create maintenance overhead while providing weaker guarantees:

If the coverage is moved to SLT, we can also test whether this optimizer rule works correctly with other optimizer rules in most cases.

If future features change these unit tests, it is often hard to figure out the test goal of those internal tests, and update them correctly. Especially, I think this optimizer rule is likely to change a lot in the future.

So in practice, I would suggest moving most of the test coverage, especially edge cases, to SLT.

I would only keep module-level unit tests for:

1–2 demo-like tests, since they can catch errors earlier, to make development easier

Cases that are difficult to test with end-to-end tests.

Reference: a DuckDB author trying to convince us to avoid most UTs in favor of e2e tests: https://www.youtube.com/watch?v=BgC79Zt2fPs&t=940s

2010YOUY01 · 2026-07-01T09:22:42Z

+
+    let plan = TestConfig::default()
+        .with_query_execution_partitions(4)
+        .with_preserve_file_partitions(1)


It would be great to add some comment to explain this test case, especially the implication of this config

2010YOUY01 · 2026-07-01T09:26:36Z

+}
+
+#[test]
+fn range_grouping_set_aggregate_rehashes_with_grouping_id() -> Result<()> {


Is it possible to disable this optimization when grouping sets are present?

I’m finding it a bit hard to convince myself this is implemented safely. I think it would be better to ignore this case for now, and enable the optimization in a separate PR with more edge-case tests targeting grouping sets.

(We have seen quite a few tricky bugs around grouping sets before.)

We cannot disable this, I just extracted is logic and thought it was interesting there was no uinit tests for this since it is quite a weird edge case. This was introduced in #19853 and without it the subset satisfaction can create correctness errors 👍

alamb · 2026-07-01T17:51:09Z

I agree with @2010YOUY01 -- the more SLT coverage we have, the better (otherwise we end up not testing that all the paths are correctly hooked up and we convince ourselves (or the LLMs convince us) that the code does what we want, when maybe it doesn't

We already invested a bunch of time getting the SLT support for range partitioing in https://github.com/apache/datafusion/blob/main/datafusion/sqllogictest/test_files/range_partitioning.slt so this should be a fairly simple addition

gene-bordegaray · 2026-07-01T18:34:28Z

👍 the most recent code added the coverage I think @2010YOUY01 and @alamb are talking about, let me know if this is off the mark

alamb

SLT looks good to me, other than I think we need to run a few more of the queries (not just EXPLAIN)

alamb · 2026-07-01T19:06:15Z

+set datafusion.optimizer.preserve_file_partitions = 0;
+
+query TT
+EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key;


we should also run this query I think , not just do the explain

I didnt run these because the physical plan is the same as test 1, so I didnt want the churn.

Out of curiosity for why, is the reason to duplicate to see if changes in physical plan result in a query result mismatch?

I guess in my mind if there is value in explaining the same query multipel times with different settings, I would expect that we should verify the same results come out too.

Maybe we don't need so many explains 🤔

Hm, I think tht there is value in explaining the query has the same physicacl plan shape across different settings that are tied to partitioning. My reasoning for asserting all the plans across these is to make it clear what the expected behavior is in inserting repartitions with these settings and that rules dont drift from this.

I think that if the physical plans are the same I am ok with asserting the results as well but thought I would eliminate the test overhead as the plans will catch any drift

ok, let's go with this and we can improve the tests as a follow on

alamb · 2026-07-01T19:06:26Z

+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;


likewise here (and below) also we should run the queries in addition to the explain

alamb · 2026-07-02T18:09:38Z

Putting this one in the merge queue

github-actions Bot added physical-expr Changes to the physical-expr crates optimizer Optimizer rules core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) labels Jun 29, 2026

gene-bordegaray changed the title ~~Support range partitioned aggregations~~ Aggregations Support Partitioning::Range Jun 29, 2026

gene-bordegaray commented Jun 29, 2026

View reviewed changes

Comment thread datafusion/physical-optimizer/src/utils.rs Outdated

gene-bordegaray marked this pull request as ready for review June 29, 2026 14:00

This was referenced Jun 29, 2026

Support co-partitioned range inner equi joins #23184

Draft

Replace / rename HashPartitioned distribution as KeyPartitioned #23236

Closed

gabotechs reviewed Jun 29, 2026

View reviewed changes

2010YOUY01 reviewed Jun 30, 2026

View reviewed changes

gabotechs approved these changes Jun 30, 2026

View reviewed changes

gene-bordegaray mentioned this pull request Jun 30, 2026

Allow Partitioning::Range to satisfy Distribution::KeyPartitioned generally #23266

Open

2010YOUY01 reviewed Jul 1, 2026

View reviewed changes

satisfy aggregation hash distribution with range

251b76f

gene-bordegaray force-pushed the gene.bordegaray/2026/06/range-partitioned-aggregations-main branch from ee5e61f to 251b76f Compare July 1, 2026 14:01

github-actions Bot removed the physical-expr Changes to the physical-expr crates label Jul 1, 2026

alamb reviewed Jul 1, 2026

View reviewed changes

This was referenced Jul 2, 2026

Allow Partitioning::Range to satisfy window Distribution::KeyPartitioned requirements #23289

Open

Allow Partitioning::Range to satisfy PartitionedTopK Distribution::KeyPartitioned requirements #23290

Open

alamb added this pull request to the merge queue Jul 2, 2026

Merged via the queue into apache:main with commit 044a85c Jul 2, 2026
38 checks passed

Uh oh!

Conversation

gene-bordegaray commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Uh oh!

gene-bordegaray commented Jun 29, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

2010YOUY01 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gene-bordegaray commented Jun 30, 2026

Uh oh!

alamb commented Jun 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Jul 1, 2026

Uh oh!

gene-bordegaray commented Jul 1, 2026

Uh oh!

alamb left a comment

gene-bordegaray commented Jun 29, 2026 •

edited

Loading

gene-bordegaray Jul 1, 2026 •

edited

Loading