fix: truncate per-row date_trunc formats with the literal-format kernel

[andygrove]

Which issue does this PR close?

Closes #5633.

Rationale for this change

date_trunc in a non-UTC session could panic on timestamps around DST transitions (#5633). #5956 fixed the literal-format path, which now truncates in local time the way Spark does. HOUR, DAY and MINUTE keep the input's offset in an overlap and move a nonexistent local time forward by the gap. The date levels follow LocalDate.atStartOfDay, taking the earlier midnight in an overlap and the end of the gap when midnight falls inside one.

When the format comes from a column, date_trunc still went through the old chrono helpers. It still panicked in as_micros_from_unix_epoch_utc, and a NULL format raised Unsupported format, where Spark returns NULL.

What changes are included in this PR?

timestamp_trunc_array_fmt_dyn now groups the rows by format and truncates each group with timestamp_trunc_upstream, the literal-format kernel, with the other rows set to NULL. Every row therefore follows the literal-format rules, including the DST handling and the integer arithmetic for UTC and TIMESTAMP_NTZ, and a value in one group can't make another group fail. Dictionary inputs are unpacked first, so one path replaces the four dictionary/plain variants and the macro. A NULL format gives NULL. An invalid format still raises an error, which is why a non-literal format stays Incompatible. Each distinct spelling in a batch is parsed once.

timestamp_trunc_legacy, the fallback for non-UTC HOUR and DAY outside DataFusion's nanosecond range, still uses the old helpers. It only sees years before 1678 and after 2262, where tzdata has no DST transitions.

#6740 adds a flag to timestamp_trunc_upstream for Spark 4.2's checked SECOND and MILLISECOND truncation. This path calls that function, so whichever of the two merges second only needs to pass the flag through.

How are these changes tested?

• Rust: row_formats_truncate_like_literal_formats rotates all 15 format spellings across rows around DST and offset transitions in Toronto, Havana, Sao Paulo, Los Angeles, Monrovia, Asuncion and Apia, plus UTC and TIMESTAMP_NTZ edge values, and checks every row against the literal-format result. Other tests cover the Toronto 1919 WEEK gap from the earlier review, NULL and invalid formats, dictionaries, and an overflowing value in one group. All five fail on main, either with the panic in as_micros_from_unix_epoch_utc or with an error.
• SQL: trunc_timestamp_dst_format_column.sql crosses the DST rows from trunc_timestamp_dst_ambiguous.sql with every spelling and NULL, in UTC and the same eight zones, under expect_native(date_trunc). It fails on main in all nine sessions and passes on Spark 4.1 and 3.5.
• I also compared the format-column path with Spark around every transition from 1800 to 2037 in all 590 JVM timezones, about 1.08 million rows. Every difference traced to the JVM's older tzdata (2023c, against chrono-tz's 2025b).
• A Criterion comparison against main over 8192 rows (release build without LTO): in UTC it is 2.7x faster with mixed formats, 11x for YEAR and 27x for HOUR. In America/Los_Angeles it is 1.5x, 4.2x and 1.1x.

[sunchao]

Summary

• Prior state and problem: Timezone-aware truncation could panic at DST transitions or select the wrong offset.
• Design approach: Separate instant, local-time, and local-date truncation rules while reusing the existing naive datetime helpers.
• Correctness / compatibility analysis: The common DST cases match Spark. One P2 remains: a gap spanning midnight now produces a silent incorrect weekly boundary where the base panicked.
• Key design decisions: Non-UTC execution remains opt-in through allowIncompatible. Timezone parsing and builder allocation stay outside the row loop.
• Implementation sketch: Replace the timezone-aware helper chain with TzTrunc, add transition assertions, and enable or add SQL fixtures.
• Behavioral changes worth calling out: Ambiguous local-time truncations preserve the input offset. Subsecond truncation uses instant arithmetic. Date-level gaps need a distinct resolution rule.
• Suggested improvements: Resolve date-level gaps at the first valid instant and add the Toronto regression described below.

Reviewed all three changed files at 933a68910dc53183f51103941e04f83d305bcb7c against e1d2c11729c2fc60a5def4e87bb17e5b28df2a29. The PR remains non-draft. Routed skills: review-comet-pr and review-comet-expression-pr. No existing reviews, comments, or threads were present.

Exact-head CI: 31 successful checks and 44 skipped, with no failures. Logs confirm both DST SQL fixtures and the Rust truncation tests passed. Spark SQL, macOS, and benchmark jobs were skipped.

Validation: Four focused Rust tests passed locally. A standalone comparison of the unchanged base/head kernels against java.time covered 593,520 cases and isolated the reported behavior from unchanged mismatches. Relevant Spark sources were checked across 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. No local JVM suite or performance benchmark was run.

[sunchao]

[P2] Resolve LocalDate gaps at the first valid instant. With session timezone America/Toronto and spark.comet.expression.TruncTimestamp.allowIncompatible=true, apply date_trunc('WEEK', ts) to a timestamp column containing 1919-03-31 00:45:00-04:00. Spark returns 1919-03-31 00:30:00-04:00, but this branch returns 01:00:00-04:00. Toronto's gap ran from the previous day's 23:30 to 00:30. resolve_local_datetime shifts midnight by the full gap length, whereas Spark's LocalDate.atStartOfDay selects the gap's end. This replaces the base's panic with silently incorrect weekly boundaries, even producing a result later than this input. Keep gap-length shifting for LocalTime, but resolve LocalDate to the transition boundary and add this regression case.

Evidence: Compiled the unmodified head/base kernels against locked Arrow 59.3.0 and chrono 0.4.45. For a timezone-annotated TimestampMicrosecondArray containing -1601752500000000, timestamp_trunc(..., "WEEK") returns -1601751600000000 at head. The base panics. Java 21's LocalDate.of(1919, 3, 31).atStartOfDay(ZoneId.of("America/Toronto")) returns -1601753400000000, corresponding to 00:30-04:00. Spark's date-level truncation uses this atStartOfDay rule through daysToMicros across the checked supported versions.

[comphead]

Thanks for the careful write-up and for checking the expected values against java.time. From reading ZonedDateTime.ofLocal and LocalDate.atStartOfDay, the offset choice for ambiguous local times looks right to me for all three level groups. I left a few small suggestions inline. I haven't run anything locally, so these come from reading the code.

[comphead]

Would it make sense to reuse micros_to_naive and naive_to_micros here? This line and the last line of the function look like the same conversions, and the NTZ path in this file already uses them.

[comphead]

Small thing: the new test seems to have landed between the existing test_timestamp_trunc_dst_boundary doc comment and its #[test] attribute. The Denver QUARTER lines above now read as part of this test's docs, and test_timestamp_trunc_dst_boundary has none. Would it make sense to move the new test above that comment block?

[comphead]

Would it make sense to add a small America/Havana case where local midnight is ambiguous? Clocks fell back from 01:00 to 00:00 on 2020-11-01, and I don't think any current row makes WEEK, MONTH, QUARTER or YEAR land on a repeated local time, so the earlier-offset rule for those levels isn't pinned. I haven't run it, but for timestamp('2020-11-01 00:30:00-05:00') I'd expect DAY to give 00:00 at -05:00 and MONTH to give 00:00 at -04:00. It would need its own session timezone, so probably a small new file.

[parthchandra]

A few review notes, comparing against Spark's DateTimeUtils.truncTimestamp:

• [correctness] native/spark-expr/src/kernels/temporal.rs:143 — the LocalResult::None => resolve_local_datetime(...) branch is still wrong for date-level truncation (WEEK/MONTH/QUARTER/YEAR) when a spring-forward gap straddles midnight. It shifts local time forward by the whole gap length, which is right for DAY/HOUR/MINUTE but not the date levels. Spark resolves those with atStartOfDay, which returns the gap's end instant, not midnight minus the pre-gap offset. They agree only when the gap begins exactly at midnight. Example: America/Toronto, date_trunc('WEEK', ts) on 1919-03-31 00:45:00-04:00 (gap ran 23:30 Mar 30 to 00:30 Mar 31). Spark returns 00:30:00-04:00; this returns 01:00:00-04:00, which is 30 minutes past the gap and actually later than the input instant. For the date levels, please resolve a gap to the transition boundary (the first valid instant at or after the truncated midnight) instead of the gap-length shift. (The bot flagged the same case.)

• [tests] native/spark-expr/src/kernels/temporal.rs:1235 — every DST case here (LA, Sao Paulo) has a gap that starts at midnight, so the date-level gap rule is never tested where it differs. Please add a straddles-midnight case (the Toronto 1919 one is clean) for WEEK/MONTH/QUARTER/YEAR, and an ambiguous-midnight zone like America/Havana (fall-back on 2020-11-01) to pin the earlier-offset rule for date-level overlaps.

• [tests] native/spark-expr/src/kernels/temporal.rs:1236 — test_timestamp_trunc_matches_spark_at_dst_transitions landed between the Denver doc comment and the #[test] it documents, so test_timestamp_trunc_dst_boundary lost its doc and the Denver comment now describes the wrong test. Move the new test above that comment block.

[sunchao]

Summary

• Prior state and problem: Timezone-aware date_trunc could panic around DST transitions or choose the wrong offset.
• Design approach: Separate instant, local-time and local-date truncation using TzTrunc, while sharing the existing naive datetime helpers.
• Correctness / compatibility analysis: Common DST cases follow Spark. The previously reported P2 remains unresolved: truncating 1919-03-31 00:45:00-04:00 to WEEK in America/Toronto returns 01:00-04:00, whereas Spark returns 00:30-04:00. I reproduced this against the unchanged head and confirmed that the base panics. See the existing finding.
• Key design decisions: Non-UTC native execution remains opt-in. The level grouping is straightforward, and timezone parsing and builder allocation remain outside the row loop.
• Implementation sketch: Replace the timezone-aware setter chain with explicit resolution rules, add Rust transition assertions, and enable or add SQL regression fixtures.
• Behavioral changes worth calling out: Local-time overlaps preserve the input offset. Date-level overlaps select the earlier offset. Subsecond truncation operates directly on the instant.
• Suggested improvements: Address the existing blocker by resolving LocalDate gaps to the first valid instant, retaining gap-length shifting for LocalTime, and adding the Toronto regression.

Reviewed the complete three-file diff at 933a68910dc53183f51103941e04f83d305bcb7c against e1d2c11729c2fc60a5def4e87bb17e5b28df2a29. The PR is not a draft. Routed skills: review-comet-pr and review-comet-expression-pr. Existing discussions and threads were read. No additional introduced P1/P2 issues were found within this review.

Exact-head CI: 31 successful checks, 44 skipped, no failures. Logs confirm both DST SQL fixtures and the timestamp-kernel tests passed. Spark SQL matrix, macOS and benchmark jobs were skipped.

Validation: Four focused timestamp-kernel tests passed in a standalone harness compiling the unchanged source against locked Arrow 59.3.0 and chrono 0.4.45. The existing blocker was independently reproduced against Java 21 java.time. Relevant Spark sources were checked across 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. No full local crate build, JVM suite or performance benchmark was run.

[andygrove]

Thanks @sunchao, @parthchandra and @comphead. You're right about the date-level gap. resolve_local_datetime shifts midnight by the length of the gap, which only matches LocalDate.atStartOfDay when the gap starts exactly at midnight, so Toronto's 1919-03-31 comes out as 01:00 instead of 00:30.

sam-1112's #5956 rewrites the same kernel for literal formats. It already resolves a date-level gap to the transition instant and takes the earlier midnight in an overlap, with Toronto, Asuncion and Havana tests. Rather than fix the same thing twice and keep the two PRs conflicting, I'd like #5956 to land first. After it, the only path that still panics on DST transitions is the per-row format path (timestamp_trunc_array_fmt_helper, through as_micros_from_unix_epoch_utc). I'll then cut this PR down to that path, reuse the helpers from #5956, and pick up comphead's suggestions there. I'm moving this to draft until then.

[andygrove]

I've reworked this on top of #5956. Instead of carrying its own truncation rules, the path for a format column now groups the rows by format and runs each group through the literal-format kernel. The Toronto date-level gap @sunchao and @parthchandra found is handled by the atStartOfDay gap search from #5956, and @comphead's Havana case is in both DST fixtures. The code those inline comments were on is gone, so they no longer apply. The description has the details and the test results.

[sunchao]

Summary

• Prior state and problem: Column-format date_trunc still used helpers that could panic around DST transitions and mishandled NULL formats after the literal-format path had been fixed.
• Design approach: Decode dictionaries, group rows by normalized granularity, mask unrelated rows, and reuse timestamp_trunc_upstream.
• Correctness: The DST, null, dictionary and group-isolation checks passed. Spark source confirms the distinction between instant, local-time and date-boundary truncation. One introduced P2 remains: the column path now silently wraps overflowing SECOND/MILLISECOND truncations on Spark 4.2, where Spark raises an error.
• Compatibility analysis: Checked Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0 sources. Nonliteral formats remain Incompatible, and output relabelling remains unchanged. Compared with branch-1.1 at e9efd9f764ee0a59b7898ff028d6985d4a7a28e1, the old column-format implementation matches the supplied base. DST/null fixes are intended changes. The Spark 4.2 error-to-result change needs correction.
• Key design decisions: Canonical groups share truncation semantics, masking prevents unrelated values from raising errors, and a single-group fast path avoids masking when safe.
• Implementation sketch: Replace four dictionary/plain variants and their macro with dictionary decoding and one grouped kernel. Add five Rust tests and a SQL fixture covering nine session timezones.
• Performance: Grouping avoids repeated format parsing and reuses faster arithmetic paths. Mixed formats require batch-sized masks and results for up to ten granularities. Bounded UTC/Los Angeles timings supported the reported improvement direction. Cached dependency builds and kernel-only timings do not establish production performance. No additional performance concern met the P1/P2 bar.
• Design: Sharing the literal-format implementation is simpler and reduces the risk of separate DST rules drifting. The existing Toronto gap finding is fixed: WEEK now returns 1919-03-31 00:30-04:00.
• Abstraction & complexity: The dictionary adapter and grouped helper replace duplicated dispatch without introducing a broader framework. The remaining Spark-version distinction should be carried through the shared kernel rather than implemented separately here.
• Behavioral changes worth calling out: NULL formats now produce NULL, valid column formats inherit the literal path's DST and wide-range handling, and invalid non-null formats still raise an error. Spark 4.2 overflow behavior is the reported regression.
• Suggested improvements: Thread the version-aware overflow policy described for #6740 through both shared-kernel call sites and add a Spark 4.2 column-format regression test.

Reviewed the complete two-file diff at 81a37853015a317a911d12eb98200d964737bfd9 against b56349697b786ff2ad1c1bcf6ecf45b809af5f30. The PR is not a draft. Routed skills: review-comet-pr and review-comet-expression-pr, with the timezone and contributor guidance. Existing reviews, comments and threads were read and checked for updates. The earlier Toronto P2 no longer remains substantiated.

Exact-head CI: 26 successful checks and 15 skipped, with no failures or pending checks. Rust logs confirm all five added tests passed. Spark 4.1 expression logs confirm all nine sessions of the new SQL fixture passed. Spark SQL matrix, macOS and benchmark checks were skipped.

Validation: All 30 temporal-kernel tests passed in a standalone harness compiling the unchanged head source against locked Arrow 59.3.0, chrono 0.4.45 and DataFusion 55.1.0. Two additional disposable tests covered sliced dictionaries, null keys/values, empty input and all-null formats. Unchanged base/head kernels and Spark 4.2's actual TruncTimestamp expression reproduced the overflow finding. No full local Comet build, end-to-end JVM suite or production benchmark was run.

[sunchao]

[P2] Preserve Spark 4.2's checked overflow behavior when routing column formats here. With a UTC session, spark.comet.expression.TruncTimestamp.allowIncompatible=true, and columns containing us=-9223372036854775808 and fmt='SECOND', unix_micros(date_trunc(fmt, timestamp_micros(us))) now yields 9223372036854551616. Spark 4.2 raises ArithmeticException: long overflow, and the base column path also errors. MILLISECOND similarly returns 9223372036854775616. Although the scalar mismatch predates this PR, this new delegation materially worsens the column path by turning an error into a silently wrapped timestamp. Could you thread the version-aware overflow policy described for #6740 through both the single-group and masked-group calls, retaining wrapping for Spark 3.4–4.1, and add a Spark 4.2 column-format regression?

Evidence: Compiled the unchanged base/head temporal kernels against locked Arrow 59.3.0, chrono 0.4.45 and DataFusion 55.1.0. For TimestampMicrosecondArray([i64::MIN]).with_timezone("UTC") and StringArray(["SECOND"]), the base returns Err("Unable to read value as datetime") and head returns 9223372036854551616. MILLISECOND returns 9223372036854775616 at head. Spark 4.2.0 DateTimeUtils.truncTimestamp and TruncTimestamp.eval with BoundReference column inputs both throw ArithmeticException: long overflow for both formats. Spark 4.2 source uses Math.subtractExact, while 3.4–4.1 use wrapping subtraction.