Skip to content

feat: reuse range partitioning for date_trunc aggregation - #26449

Open
benbellick wants to merge 49 commits into
apache:mainfrom
benbellick:feat/issue-25344-range-transform-locality
Open

benbellick wants to merge 49 commits into
apache:mainfrom
benbellick:feat/issue-25344-range-transform-locality

Conversation

@benbellick

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #26448.

Related to #23569, #25344, and #25345.

Rationale for this change

This PR implements a focused subset of the optimization in Nga's PR #24501 for a known-safe case. The plan is to iteratively expand support to more safe cases, following Gene's review feedback.

Suppose events is range-partitioned on timestamp, with partitions split at hourly boundaries.

SELECT
    key,
    date_trunc('hour', timestamp) AS hour,
    SUM(value)
FROM events
GROUP BY key, date_trunc('hour', timestamp);

Each (key, hour) group belongs entirely to one input partition. The simplified physical plans below omit projections and use hour for the grouping expression.

Without the optimization

AggregateExec: mode=FinalPartitioned, gby=[key, hour], aggr=[sum(value)]
  RepartitionExec: Hash([key, hour])
    AggregateExec: mode=Partial, gby=[key, hour], aggr=[sum(value)]
      Scan: Range(timestamp), hourly splits

With the optimization

AggregateExec: mode=SinglePartitioned, gby=[key, hour], aggr=[sum(value)]
  Scan: Range(timestamp), hourly splits

The new locality check is specific to aggregates. General distribution satisfaction and join-layout checks are unchanged.

What changes are included in this PR?

Add a default-deny capability that identifies functions and argument types safe for range-key analysis. Check expression ordering and arguments, then evaluate both sides of each split to prove groups stay within partitions.

The optimizer calls it through this path.

flowchart TD
    A["EnsureRequirements checks aggregate input"] --> B["AggregateExec::input_keeps_groups_local"]
    B --> C["Partitioning::keeps_keys_local"]
    C --> D{"Function admitted, expression supported, and split boundaries safe?"}
    D -->|Yes| E["Can skip the aggregate hash shuffle"]
    D -->|No or unknown| F["Retain normal repartitioning"]
Loading

Support is intentionally limited to timezone-less date_trunc over single-key timestamp ranges, with literal non-range arguments. date_bin remains excluded pending its safety audit, including null behavior. Compound range keys remain follow-up work in #24644.

Keep the shuffle when buckets cross partitions or the checks fail. Grouping sets still use the existing distribution checks.

What is the testing strategy for this PR?

Added multiple unit and integration tests.

I'm happy to add a benchmark if maintainers would find it useful.

Are there any user-facing changes?

Eligible date_trunc aggregations can skip a hash shuffle. Results are unchanged.


Note: This PR was developed with AI assistance. All changes have been reviewed, and I take full responsibility for this contribution.

@github-actions github-actions Bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates optimizer Optimizer rules core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation physical-plan Changes to the physical-plan crate labels Oct 17, 2026
Mark intentional expression hash keys with narrowly scoped lint expectations. Interior-mutable memory accounting does not affect expression hashing or equality.
Reject malformed effective split points before proving transformed-key locality. Cover uncombined aggregation, repeated optimization, incompatible join layouts, and grouping-set subtotals with executable regressions.
@benbellick
benbellick force-pushed the feat/issue-25344-range-transform-locality branch from 2a0b4f5 to a6482e5 Compare October 17, 2026 16:32
@codecov-commenter

codecov-commenter commented Oct 17, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.79618% with 33 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.92%. Comparing base (484183d) to head (a6482e5).
⚠️ Report is 12 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/physical-expr/src/partitioning.rs 95.84% 9 Missing and 16 partials ⚠️
datafusion/expr/src/udf.rs 77.77% 4 Missing ⚠️
datafusion/functions/src/datetime/date_trunc.rs 98.59% 1 Missing ⚠️
...on/physical-expr/src/equivalence/properties/mod.rs 94.11% 1 Missing ⚠️
datafusion/physical-expr/src/scalar_function.rs 92.30% 0 Missing and 1 partial ⚠️
...atafusion/physical-optimizer/src/sanity_checker.rs 90.90% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #26449      +/-   ##
==========================================
+ Coverage   82.89%   82.92%   +0.02%     
==========================================
  Files        1154     1154              
  Lines      457630   458406     +776     
  Branches   457630   458406     +776     
==========================================
+ Hits       379370   380114     +744     
- Misses      55256    55267      +11     
- Partials    23004    23025      +21     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate functions Changes to functions implementation logical-expr Logical plan and expressions optimizer Optimizer rules physical-expr Changes to the physical-expr crates physical-plan Changes to the physical-plan crate sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Avoid aggregate repartitioning for audited date_trunc on single-key timestamp ranges

2 participants