Skip to content

feat: reuse range partitioning for date_trunc aggregation - #26449

Merged
jayzhan211 merged 50 commits into
apache:mainfrom
benbellick:feat/issue-25344-range-transform-locality
Oct 23, 2026
Merged

jayzhan211 merged 50 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
@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 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 NGA-TRAN left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
jayzhan211 self-requested a review October 21, 2026 12:40
@jayzhan211

Copy link
Copy Markdown
Contributor

Thanks @benbellick @NGA-TRAN! LGTM

@benbellick

Copy link
Copy Markdown
Contributor Author

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

@jayzhan211
jayzhan211 added this pull request to the merge queue Oct 23, 2026
Merged via the queue into apache:main with commit 8aad40b Oct 23, 2026
45 checks passed
gabotechs pushed a commit to moraistejerina/datafusion that referenced this pull request Oct 23, 2026
…26607)

## Which issue does this PR close?

Closes apache#26605.

## Rationale for this change

`date_bin` returns NULL for unrepresentable timestamp bins. Return an
error instead, matching
[PostgreSQL](https://www.db-fiddle.com/f/h3qSHUSRprbzkX9GBkSuY2/0).

I'm not sure whether DataFusion has a specific policy on returning NULL
versus an error here, but the PostgreSQL precedent seems clear.
Returning NULL because we cannot compute a bin for a non-null timestamp
silently hides a computation failure rather than representing missing
input.

This prepares `date_bin` for avoiding aggregate repartitioning in
[apache#25345](apache#25345), similar to
the `date_trunc` optimization in
[apache#26449](apache#26449). It follows
the error behavior established for `date_trunc` in
[apache#25131](apache#25131).

## What changes are included in this PR?

Make scalar and array timestamp-binning failures return `DATE_BIN
timestamp out of range`. Preserve NULL inputs and successful
wider-arithmetic results.

## What is the testing strategy for this PR?

Unit tests and SQLLogicTests cover overflow errors, scalar/array
consistency, NULL preservation, and representable results across
timestamp precisions.

## Are there any user-facing changes?

Previously NULL-producing timestamp-binning failures now fail the query.

---
This PR was developed with AI assistance. All changes have been
reviewed, and I take full responsibility for this contribution.
jayzhan211 pushed a commit to CuteChuanChuan/datafusion that referenced this pull request Oct 31, 2026
…he#26835)

## Which issue does this PR close?

Closes apache#25344.
Closes apache#25345.

This completes the initial single-key work alongside apache#26449. Explicit
origins (apache#26913) and compound range keys (apache#24644) remain follow-ups.

## Rationale for this change

Extend the aggregate-locality optimization from apache#26449 to `date_bin`.
When bins stay within timestamp range partitions, aggregation can skip
the hash shuffle.

## What changes are included in this PR?

Pass literal argument values to the existing admission hook. `INTERVAL
'1 hour'` and `INTERVAL '1 month'` have the same `IntervalMonthDayNano`
type, so types alone cannot distinguish them. Inspecting the literal
value lets `date_bin` require a positive fixed-duration stride without
months, treating each day as 24 hours. `date_trunc` continues to rely
only on argument types.

Initially, admit only two-argument calls over timezone-less nanosecond
timestamps, using that literal stride. Explicit origins and other types
remain excluded.

For example:

```sql
SELECT
    date_bin(INTERVAL '1 hour', event_time) AS hour_bin,
    COUNT(*)
FROM events
GROUP BY hour_bin;
```

For ascending timestamp ranges:

- A split at `01:00` separates hourly bins, so the shuffle can be
skipped.
- A split at `01:30` cuts through the `01:00` bin, so the shuffle
remains.

Every split must pass the existing boundary proof. Evaluation errors
conservatively retain the shuffle. This follows the NULL-to-error fix in
apache#26607.

## What is the testing strategy for this PR?

Unit tests cover admission, hook forwarding, ascending/descending
boundaries, and evaluation errors. SQLLogicTests verify plans and
results for aligned bins, crossing bins, and excluded explicit origins.

## Are there any user-facing changes?

Eligible `date_bin` aggregations can avoid a hash shuffle. Function
results are unchanged.

---

Note: This PR was developed with AI assistance. All changes have been
reviewed, and I take full responsibility for this contribution.
vim89 pushed a commit to vim89/datafusion that referenced this pull request Nov 13, 2026
…pache#26716)

## Which issue does this PR close?

- Part of apache#26299

## Rationale for this change

Step 1 from the issue. The per-child closure in
`ensure_distribution_with_stats` works out several flags inline
(`allow_subset_satisfy_partitioning`, `add_roundrobin`,
`needs_repartition`, ...). They depend on each other and can only be
tested through full plan tests right now. This moves them into one place
so they can be unit tested. No behaviour change.

## What changes are included in this PR?

- `ChildPolicyConfig`: the settings that are the same for every child.
- `ChildPolicy::new`: works out the flags for one child. The partial
aggregate check (including the
`partial_aggregate_input_keeps_groups_local` case from apache#26449) is passed
as a closure, so it still only runs when it did before.
- `ChildPolicy::needs_hash_repartition`: the repartition decision from
the hash arm.
- The closure now uses these. The rest of the function is unchanged.

## What is the testing strategy for this PR?

New unit tests in `enforce_distribution.rs` for the flags. The existing
`physical_optimizer` tests and sqllogictests pass without changes.

<!-- screenshot: cargo test -p datafusion-physical-optimizer --lib --
enforce_distribution -->

## Are there any user-facing changes?

No
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) v56.0.0

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

4 participants