[benbellick]

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.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 93.69592% with 51 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.93%. Comparing base (484183d) to head (c7a7f94).
⚠️ Report is 14 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-expr/src/partitioning.rs	96.02%	8 Missing and 17 partials ⚠️
datafusion/functions/src/datetime/date_trunc.rs	72.05%	14 Missing and 5 partials ⚠️
datafusion/expr/src/udf.rs	77.77%	4 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.93%   +0.03%     
==========================================
  Files        1154     1155       +1     
  Lines      457630   458894    +1264     
  Branches   457630   458894    +1264     
==========================================
+ Hits       379370   380592    +1222     
+ Misses      55256    55240      -16     
- Partials    23004    23062      +58     

☔ 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.

[NGA-TRAN]

The description and the plan outputs are for those test cases are exactly what we need. The code looks great. Thanks so much @benbellick

@jayzhan211 : since you review the earlier PR (that covers boarder cases), we hope you give this a look. Thanks

[jayzhan211]

Thanks @benbellick @NGA-TRAN! LGTM

[benbellick]

Thanks for the quick reviews! I will continue with the followups once this is merged 🚀