feat: skip hash shuffle for date_bin/date_trunc on Range([timestamp]) - #24501

Codecov Report

❌ Patch coverage is 94.75446% with 47 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.63%. Comparing base (8332cfa) to head (4b73736).
⚠️ Report is 934 commits behind head on main.

@@            Coverage Diff             @@
##             main   #24501      +/-   ##
==========================================
+ Coverage   81.32%   81.63%   +0.30%     
==========================================
  Files        1117     1123       +6     
  Lines      396269   412385   +16116     
  Branches   396269   412385   +16116     
==========================================
+ Hits       322260   336638   +14378     
- Misses      55186    55954     +768     
- Partials    18823    19793     +970     

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

Tag @gene-bordegaray and @jayshrivastava for review

Thanks @NGA-TRAN! I have two suggestions, the 2nd one is non-blocking

project applies the monotonic fallback to non-final range keys, where the disjointness check does not hold

Partition membership for a multi-key Range is a lexicographic comparison of the whole tuple, but the disjointness check only proves something about one key. That gap is reachable.

Setup

Range([ts, key]) with split (S, 'k5'), where S = 2024-01-01T01:00:00. Real membership:

Project ts away, keep h = date_trunc('hour', ts). The check passes, since f(S) = 01:00 and f(S-1ns) = 00:00, so project rewrites this to Range([h, key]) with split (01:00, 'k5'). I ran it on the branch to confirm:

Range([timestamp, key], [(2024-01-01T01:00, 'k5')])
  projected through: SELECT date_trunc('hour', timestamp) AS h, key
  => Range([h@0 ASC, key@1 ASC], [(1704070800000000000, k5)], 2)

Where the claim breaks

f is non-injective. It maps all of [01:00, 02:00) onto 01:00, so "h equals 01:00" is not the same as "ts equals S". Two rows, both with key = 'k1':

ts = 01:00:00   real: ts == S, 'k1' < 'k5'        -> partition 0
                claimed: h == 01:00, 'k1' < 'k5'  -> partition 0   agrees

ts = 01:30:00   real: ts > S                      -> partition 1
                claimed: h == 01:00, 'k1' < 'k5'  -> partition 0   WRON

The second row is above the split in reality. After the transform it ties with the split on the first key, the comparison falls through to key, and 'k1' < 'k5' drags
it back below the boundary.

Both rows are the same group under GROUP BY h, key. So key_satisfactshuffle is dropped, and SinglePartitionedemits(01:00, 'k1')` twice,once per partition.

Why the last key is safe and an earlier key is not

The boundary survives only if this equivalence holds:

row < split (source domain)  <=>  projected row < projected split

The check f(split) != f(split-1) gives only the below half:

ts < S  =>  f(ts) <= f(S-1) < f(S)

Transform on the last key. Earlier keys were projected identity, so ties on them are exact and the comparison always terminates at this key. It reduces to comparing
f(b) against f(b_split). The below half is covered above, and plainther side:

b >= b_split  =>  f(b) >= f(b_split)

That is the full equivalence. Sound.

Transform on an earlier key. The comparison can continue past it, s be strict:

ts > S  =>  f(ts) > f(S)

That is injectivity around the split, which flooring transforms like `dnever have. Nothing in the check tests this direction.

Fix

range_monotonic_fn_satisfaction already sidesteps all of this by requn() == 1(the#24644comment).project` needs the equivalentrestriction:

     ) -> Option<Self> {
         let mut split_points = self.split_points.clone();
         let mut sort_exprs = Vec::with_capacity(self.ordering.len());
+        let last_key_idx = self.ordering.len().checked_sub(1)?;
         for (key_idx, sort_expr) in self.ordering.iter().enumerate() {
             if let Some(projected) =
                 input_eq_properties.project_expr(&sort_expr.expr, mapp
             {
                 sort_exprs.push(PhysicalSortExpr::new(projected, sort_
                 continue;
             }

+            // A non-injective transform preserves the partition bound
+            // only for the final range key. For an earlier key, rows above
+            // the split can collapse onto `f(split)` and are then ord
+            // by the trailing keys, crossing the projected boundary.
+            if key_idx != last_key_idx {
+                return None;
+            }
+
             let (target, source) =
                 monotonic_range_key_projection(sort_expr, mapping, input_eq_properties)

If you would rather keep non-final keys working, the precise condition codebase as ExprProperties::strictly_order_preserving, whichcast_expr_properties sets only for widening, one-to-one casts. Gating non-final keys on that flag would also be correct, just more machinery than this PR needs, since
date_bin and date_trunc never qualify.

Either way, a regression test mirroring test_range_partitioning_projecn with the transform on key 0 instead of key 1 would pin this.

Satisfaction is granted in the transformed key domain, but the advertised split points stay in the source domain

Not a live bug as far as I can tell, but it is an invariant that other code already depends on, so it is worth pinning down before this lands.

The mismatch

project and range_monotonic_fn_satisfaction treat the same situation differently:

project:       Range([ts], [S])  ->  Range([f(ts)], [f(S)])   split translated
satisfaction:  Range([ts], [S])  vs  KeyPartitioned([f(ts)])  ->  Exact
               ... and the partitioning still advertises [S]

After Exact, the plan says "this input is partitioned by f(ts)", while Partitioning::Range still reports ordering ts and split point S. The effective boundary in required-key space is f(S), not S.

For the disjointness question this does not matter, and the aggregation case in this PR is sound. It matters for consumers that read split-point values and compare them across two partitionings:

// datafusion/physical-plan/src/distribution_requirements.rs:348
(Partitioning::Range(left), Partitioning::Range(right)) => {
    left.split_points() == right.split_points() && ...

co_partitioning_satisfied calls satisfaction(..., allow_subset = false), which is exactly the path that returns Exact here. So a join can conclude two inputs are co-partitioned by comparing boundaries that were never in the same domain as the keys being joined. union.rs:1726 builds a Range from children's partitionings on the same kind of equality.

Why it does not bite today

I probed the transforms that actually clear check_monotonic_transform with an unbounded source range:

Two things save it:

1. For date_bin and date_trunc, f(split) != f(split-1) means the split starts its bin, so f(split) == split. The translated and untranslated boundaries coincide, and the layout comparison is accidentally correct.
2. Shift-style transforms never get here. BinaryExpr::get_properties sends Plus through arithmetic_sort_properties, which returns Unordered for unbounded operand ranges, and get_expr_properties always seeds the dependency unbounded.

The cast row is the one that already breaks the invariant: CAST(ts AS Timestamp(Second)) yields Exact with f(split) a TimestampSecond while the advertised split stays TimestampNanosecond. It happens to be harmless because those variants never compare equal, so the layout check just returns false. That is luck, not design.

What I would do

Cheapest option, document the assumption where it is relied upon:

// Callers may compare `RangePartitioning::split_points()` across two
// partitionings (see `compatible_co_partitioning_layout`). Those values
// stay in the range-key domain, while satisfaction here is granted in the
// transformed domain. That is only consistent because every transform
// reaching this point is a flooring function, i.e. disjointness at `split`
// implies `f(split) == split`. A monotone transform without that fixed
// point would make the advertised boundary wrong for those callers.

Conservative option, enforce it, at the cost of dropping the narrowing-cast case:

         let Some(at_split) = evaluate_expr_on_key(fn_expr, range_key, split_value) else {
             return false;
         };
+        // The split must be a fixed point of `f`. Otherwise the boundary in
+        // the transformed domain is `f(split)`, not `split`, and callers that
+        // compare raw split points are reasoning about the wrong value.
+        if at_split != *split_value {
+            return false;
+        }
         let Some(below_split) = evaluate_expr_on_key(fn_expr, range_key, &predecessor)
         else {
             return false;
         };
         at_split != below_split

Every test in the PR passes with that guard, since the 60s date_bin and date_trunc('hour') splits are both fixed points.

Either way, a test asserting the current CAST(ts AS Timestamp(Second)) behaviour would make a future regression visible. Happy to see this as a follow-up issue rather than in this PR.

Thanks so much, @jayzhan211, for taking the time to verify the feature and suggest improvements to avoid regressions — really appreciate the effort ❤️.

I’d also love to hear from @gene-bordegaray, given his pioneering work on range partitioning, superset/subset group‑by keys, and multi‑key partitioning. With input from @jayshrivastava, @jayzhan211 and @gene-bordegaray, I’m hopeful we can avoid any regressions.

Thanks in advance, everyone.

project applies the monotonic fallback to non-final range keys, where the disjointness check does not hold

Partition membership for a multi-key Range is a lexicographic comparison of the whole tuple, but the disjointness check only proves something about one key. That gap is reachable.

read through this and yes, can confirm that this is possible and correctness bug. We could get duplicate group rows or partial sums as final results. I also agree with @jayzhan211 suggested approach for a solution

thank you @jayzhan211 🙇

Thanks @jayzhan211 and @gene-bordegaray — review comments from both of you are addressed in 4b73736.

Jay (@jayzhan211)

1. Monotonic project fallback is restricted to the last range key. Range([timestamp, key]) + date_trunc/date_bin now drops to UnknownPartitioning. Covered by test_range_partitioning_project_rejects_monotonic_transform_on_non_final_key.
2. Split-point domain vs satisfaction (non-blocking) is tracked as follow-up Pin Range satisfaction vs advertised split-point domain #24947.

Gene (@gene-bordegaray)

1. DESC disjointness uses the successor; hour-aligned DESC splits no longer look disjoint. Regression: test_range_partitioning_descending_uses_successor_for_disjointness.
2. Non-finite float splits return None from checked_predecessor / checked_successor instead of panicking.
3. Failed split-point projection now fail-closes (?) instead of keeping source-domain bounds.

I may create a stack of PRs that implments this gradually 👍 @NGA-TRAN

@gene-bordegaray and I talked, and we may need to break this into sub‑tasks since it touches many different areas.

Ok here is a more composable stack of work to introduce this feature in a way that is more digesteable and allows more of the community to have eyes and an opinion on this as there are many hidden implications in this change:

1. date_trunc correctness fix (pre-requisite)

There is unchecked arithmetic in general_date_trunc_array_fine_granularity. This is a prerequisite because the optimizer evaluates scalar boundary values when deciding whether date_trunc preserves a range partitioning. The scalar path uses checked arithmetic, but the array fast path can overflow or panic for i64::MIN.

This is an independent correctness and we need to fix this unless there is a motivation for this in the fast-path I am unaware of.

2. Audit Range transformation sand projections

I believe that preservation for functions should be default-deny. An ordinary Ordered property is not enough to confidently declare a function preserves partition boundaries. There can be nasty edge cases and opting into everything at once is scary.

I think we should:

• thoroughly audit date_bin and date_trunc timestamp units, time zones, NULLs, and all these weird cases
• only project split points through thoroughly tested transformations, otherwise reject

3. Distinguish transformed satisfaction from Exact, then add the first aggregation consumer

Satisfaction::Exact means that the partition key and required key are exact matches. This is strict enough for co-partitioned operators, like partitioned joins, because their split points can be compared directly to see, well if they are exactly matching.

Now, a transformed key is not as strong. Its weaker guarantee is that equal f(key) values may stay local to one input partition. This makes things like not shuffling on aggregation safe, but now the boundary is f(split), not the original split points.

For example:

  left:  Range(x), split 10, join key x + 1
  right: Range(y), split 10, join key y

The effective boundaries are 11 and 10. Without repartitioning, x = 9 and y = 10 are placed in different partition indexes even though x + 1 = y. Co-partitioning is quite tricky and we need to be careful when eliminating shuffles on opertors that require it, it is much stricter than just satisfying Distributions

So there may be need for discussion around if this needs a new variant of satisfaction. Single-input aggregates may accept it, while co-partitioned joins should continue to require Exact layout compatibility.

I also think this PR should also add the first consumer, which can be the no-repartition aggregate optimization as I have found it is usually better to have a direct consumer of an API change rather than writing speculative code.

cc: @NGA-TRAN @jayzhan211

👋 hey! I spoke offline with @NGA-TRAN and I am interested in taking over this one. Does that work? @gene-bordegaray I will proceed with the path suggested above.

that would be great thank you @benbellick 🙇

making an issuie to track this progress would be great too 👍

also much of this can be reused and is great work. Just smalle scoped changes and ensure we cross our t's and dot our i's

Issue for the first item above here: #25129

And here is a PR for that first issue: #25131

very small fix!

What is the status of this PR? Shall we resolve conflicts and get it merged?

What is the status of this PR? Shall we resolve conflicts and get it merged?

@alamb There is still some work to be done. This version of the PR currently opts queries in liberally. I am working on another PR which instead focuses on just adding date_trunc as a conservative first optimization and then following up with expanding out to date_bin as well.

I am working on it in a separate PR, however much of the implementation will be similar. The work for it is tracked here #25344 I have not yet opened the PR but will soon.

Let me know if anything is unclear 🙂

Let me know if anything is unclear 🙂

Great -- thank you @benbellick -- I was going through open and approved PRs trying to work through the list that needs to be merged.

Marking this one as Draft as it sounds like it is being reworked into smaller pieces

Thank you 🙏

@NGA-TRAN Following the safety-first stack discussed above, these PRs have now merged:

• #25131 — fix date_trunc array underflow.
• #25430 — floor negative scalar timestamps correctly in date_trunc.
• #26607 — error on out-of-range date_bin results instead of returning NULL.
• #26449 — audited date_trunc aggregate locality.
• #26835 — audited date_bin aggregate locality.
• #27135 — support literal origins for date_bin locality.

These cover the single-key aggregate shuffle optimization using a separate locality proof, without changing general range projection or join satisfaction. Compound-range support remains tracked in #24644.

I think we can close this PR in favor of the merged approach and track any broader transformed-range projection work separately. Does that sound reasonable?

@benbellick : Thanks for all the work. Yes, please close this PR