benbellick opened a new pull request, #26449: URL: https://github.com/apache/datafusion/pull/26449
## Which issue does this PR close? Closes apache/datafusion#26448. Related to apache/datafusion#23569, apache/datafusion#25344, and apache/datafusion#25345. ## Rationale for this change This PR implements a focused subset of the optimization in [Nga's PR #24501](https://github.com/apache/datafusion/pull/24501) for a known-safe case. The plan is to iteratively expand support to more safe cases, following [Gene's review feedback](https://github.com/apache/datafusion/pull/24501#issuecomment-5589790676). Suppose `events` is range-partitioned on `timestamp`, with partitions split at hourly boundaries. ```sql 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** ```text 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** ```text 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. ```mermaid 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"] ``` **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 apache/datafusion#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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
