Skip to content

feat: support timestampadd, timestampdiff and make_interval via codegen dispatch - #5030

Open
andygrove wants to merge 3 commits into
apache:mainfrom
andygrove:feat/datetime-codegen-dispatch
Open

feat: support timestampadd, timestampdiff and make_interval via codegen dispatch#5030
andygrove wants to merge 3 commits into
apache:mainfrom
andygrove:feat/datetime-codegen-dispatch

Conversation

@andygrove

@andygrove andygrove commented Jul 24, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

There is no dedicated issue. This is part of the ongoing effort to provide native support for expressions that currently force a full fallback to Spark. Related to the expression coverage epic #240, and to #4540 for the interval piece.

Rationale for this change

timestampadd, timestampdiff and make_interval have no Comet handler today, so any query using them falls the entire operator back to Spark. All three are regular expressions, not RuntimeReplaceable, so they reach serde directly.

Rather than native implementations, this PR routes them through the JVM codegen dispatcher. The dispatcher runs Spark's own generated code inside the native Comet pipeline, which keeps the operator native while guaranteeing bit-for-bit Spark compatibility across all supported versions. This avoids the datetime edge-case divergences (timezone and calendar handling) that a chrono-based native implementation would be prone to.

make_interval produces CalendarIntervalType, which Comet's columnar layer gained support for in #4898. With that in place the dispatcher can carry its output, so it is wired up here alongside the other two.

What changes are included in this PR?

  • CometTimestampAdd, CometTimestampDiff and CometMakeInterval codegen-dispatch serdes in datetime.scala, registered in QueryPlanSerde's temporalExpressions map.
  • Comet SQL file tests covering all time units, negative quantities, month-end and leap-day rollover, whole-unit truncation, DST transitions, TIMESTAMP_NTZ inputs, ANSI and non-ANSI overflow, and NULL inputs.
  • Support status for the new functions in the Supported Spark Expressions guide, plus corrections to four existing rows.

Grammar aliases

timestampadd and timestampdiff are not the only spellings that build these nodes. The grammar routes TIMESTAMPADD | DATEADD | DATE_ADD to TimestampAdd, and TIMESTAMPDIFF | DATEDIFF | DATE_DIFF (plus TIMEDIFF on 4.0+) to TimestampDiff, whenever the first argument is a datetimeUnit keyword. So dateadd(DAY, 1, ts) is a TimestampAdd, not a DateAdd, and before this PR it fell back.

That made four rows in the expression guide misleading: date_add, date_diff, dateadd and datediff all claimed Native, which is true only of the two-argument form. Those rows are corrected here, and timediff, timestampadd and timestampdiff rows are added.

The alias spellings are covered by tests, split by the Spark version that introduced each:

Spelling Available Fixture
dateadd / datediff 3.4+ inline in timestampadd.sql / timestampdiff.sql
date_add / date_diff (unit form) 3.5+ date_add_unit_alias.sql
timediff 4.0+ timediff.sql

Spark 3.4's rule is only TIMESTAMPADD | DATEADD and TIMESTAMPDIFF | DATEDIFF; DATE_ADD and DATE_DIFF joined in 3.5, which is why the fixtures are gated with MinSparkVersion rather than kept in one file.

The string-literal unit form needs no fixture: it is rejected at parse time, not runtime. visitTimestampadd throws invalidDatetimeUnitError when ctx.invalidUnit is set, so timestampadd('MONTH', 1, ts) never reaches Comet.

How are these changes tested?

New Comet SQL file tests run each query through both Spark and Comet and verify the results match and that Comet executes the expression natively (through the dispatcher) rather than falling back.

Coverage includes every unit accepted by DateTimeUtils.timestampAdd and timestampDiffMap (including DAYOFYEAR, which shares a case arm with DAY, and QUARTER, the one entry with arithmetic of its own), positive and negative quantities, month-end and leap-day boundaries, whole-unit truncation toward zero, and NULL inputs.

Because both functions do calendar arithmetic on local time rather than elapsed time, the fixtures pin America/Los_Angeles and assert across real DST transitions: spring-forward, fall-back with its ambiguous local hour, and a result landing on the nonexistent 2024-03-10 02:30. These are exactly the cases a chrono-based native implementation would be most likely to get wrong.

TIMESTAMP_NTZ is covered separately for timestampadd, where dataType follows the input and an NTZ argument therefore compiles a distinct kernel with a different Arrow vector class and resolves zoneIdForType to UTC, and once for timestampdiff, where inputTypes casts NTZ up to TimestampType.

Error paths are covered on both sides: DATETIME_OVERFLOW from timestampAddOverflowError, and make_interval overflow in ANSI mode (make_interval_ansi.sql, where failOnError is true) against the non-ANSI path where it returns NULL. Each error fixture sits alongside valid-input queries so the case cannot pass vacuously through a fallback.

Verified green on all four supported Spark profiles: 3.4, 3.5, 4.0 and 4.1.

@andygrove andygrove added this to the 1.0.0 milestone Jul 24, 2026
@peterxcli

peterxcli commented Jul 25, 2026

Copy link
Copy Markdown
Member

make_interval was considered but left for a follow-up: its output type is CalendarIntervalType, which Comet's columnar layer does not yet support (the dispatcher reports unsupported output type CalendarIntervalType and falls back). It can be enabled once CalendarIntervalType support lands (#4898).

I think you're refer to #3099? I saw the native SparkMakeInterval implementation is ready and registered as comet udf, what we need to do is adding a new catalyst mapping classOf[MakeInterval] -> CometMakeInterval and CometMakeInterval scala implementation to convert expr to make_interval proto?

@peterxcli peterxcli left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

LGTM, only need to merge upstream and resolve conflict in expressions.md.

ai review: non-blocker review:

  1. add MILLISECOND to both tests and MICROSECOND to timestampdiff.
  2. Add one TIMESTAMP_NTZ or DST-boundary case; current tests only exercise timestamp-LTZ away from DST transitions.

@mbutrovich

Copy link
Copy Markdown
Contributor

The description defers make_interval because "the dispatcher reports unsupported output type CalendarIntervalType and falls back". That is no longer true on main: commit 69df0f1 (#4898) added CalendarIntervalType => true to CometBatchKernelCodegen.isSupportedDataType (spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala:92), and it is wired end to end, not just in the predicate: serializeDataType handles it (QueryPlanSerde.scala:530 and :571), the output side maps it to IntervalMonthDayNanoVector and emits a writer (CometBatchKernelCodegenOutput.scala:178 and :222), and the input side emits a getInterval getter (CometBatchKernelCodegenInput.scala:610 and :717). This PR is based on a pre-#4898 main, which is also why it conflicts. @peterxcli asked about exactly this. MakeInterval is a plain SeptenaryExpression with nullIntolerant and is not RuntimeReplaceable (Spark intervalExpressions.scala:338-348), so after the rebase object CometMakeInterval extends CometCodegenDispatch[MakeInterval] plus a classOf[MakeInterval] entry should be all it takes. Either add it here or correct the description. Note that MakeInterval.failOnError defaults to SQLConf.get.ansiEnabled, so if you do add it, the ANSI overflow path needs an expect_error case of its own.

@mbutrovich mbutrovich 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.

First round, thanks @andygrove!

| `timestamp_micros` | ✅ | |
| `timestamp_millis` | ✅ | |
| `timestamp_seconds` | ✅ | |
| `timestampadd` | ✅ | Runs through the JVM codegen dispatcher |

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.

timestampadd and timestampdiff are not the only way to build these nodes. The grammar routes TIMESTAMPADD | DATEADD | DATE_ADD to TimestampAdd and TIMESTAMPDIFF | DATEDIFF | DATE_DIFF (plus TIMEDIFF on 4.0) to TimestampDiff whenever the first argument is a datetimeUnit keyword (SqlBaseParser.g4:939-940 on branch-3.5, :1158-1159 on branch-4.0). So dateadd(DAY, 1, ts) is a TimestampAdd, not a DateAdd, and before this PR it fell back.

That makes four existing rows in docs/source/user-guide/latest/expressions.md misleading: date_add (:257), date_diff (:258), dateadd (:264) and datediff (:265) all claim Native, which is true only of the two-argument form. timediff has no row at all. Since this PR is what makes the unit form work, it is the right place to correct those and to add at least one query per fixture using an alias spelling, so the parse path is proven to land on the same serde.

Worth noting so nobody adds a test that cannot fail: the string-literal unit form is rejected at parse time, not runtime. visitTimestampadd throws invalidDatetimeUnitError when ctx.invalidUnit is set (Spark AstBuilder.scala:6012-6021), so timestampadd('MONTH', 1, ts) never reaches Comet and needs no fixture.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Corrected. All four rows now note that the Native implementation applies only to the two-argument form, and timediff, timestampadd and timestampdiff rows are added.

Each fixture now proves the parse path. One caveat on the grammar: 3.4 has only TIMESTAMPADD | DATEADD and TIMESTAMPDIFF | DATEDIFF; DATE_ADD and DATE_DIFF joined the rule in 3.5. Keeping them inline failed 3.4 with UNRESOLVED_COLUMN.WITH_SUGGESTION on DAY, since date_add fell through to the registered two-argument function. So the coverage is split by introducing version: dateadd/datediff inline (3.4+), date_add/date_diff in date_add_unit_alias.sql (MinSparkVersion: 3.5), timediff in timediff.sql (4.0+).

Noted on the string-literal unit form being a parse-time rejection; no fixture added for it.

| `timestamp_millis` | ✅ | |
| `timestamp_seconds` | ✅ | |
| `timestampadd` | ✅ | Runs through the JVM codegen dispatcher |
| `timestampdiff` | ✅ | Runs through the JVM codegen dispatcher |

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 expressions.md hunk here adds three-column rows, but main gained an Implementation column in 376676c, which is the conflict. That column is generated by GenerateDocs from the serde's trait mixins, so please regenerate with dev/generate-release-docs.sh rather than hand-merging the rows. The value will not be uniform across profiles: TimestampAdd and TimestampDiff are absent from the FunctionRegistry on 3.4 and 3.5 (reachable only through the grammar above) and are registerInternalExpression on 4.0 (FunctionRegistry.scala:927-928), and unregistered names get the placeholder.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Rebased onto main, so the rows now carry the Implementation column.

On regenerating: I ran GenerateDocs to check my values rather than hand-guessing, and it agrees on every row this PR touches. make_interval resolves to Codegen dispatch; timestampadd, timestampdiff and timediff all get the em-dash placeholder, as you predicted, since registerInternalExpression writes to the internal registry rather than FunctionRegistry.expressions, which is what buildFunctionNameToKind iterates. So the value is uniform across profiles here, just uniformly the placeholder.

I did not commit the generator output. On this checkout it also flipped about eight unrelated rows (hour, minute, second, make_timestamp to the placeholder; <<, >>, >>> to Native), which looks like main was last generated from a different profile tail. That churn does not belong in this PR, so the committed expressions.md diff is just the rows I edited. Happy to fold the regeneration in as a separate commit if you would rather have main resynced.

Comment on lines +43 to +51
SELECT
timestampadd(YEAR, 1, ts),
timestampadd(QUARTER, 1, ts),
timestampadd(WEEK, 2, ts),
timestampadd(DAY, -10, ts),
timestampadd(MINUTE, 90, ts),
timestampadd(SECOND, 30, ts),
timestampadd(MICROSECOND, 500, ts)
FROM test_timestampadd

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.

DateTimeUtils.timestampAdd handles MICROSECOND, MILLISECOND, SECOND, MINUTE, HOUR, DAY, DAYOFYEAR, WEEK, MONTH, QUARTER, YEAR (Spark DateTimeUtils.scala:670-708). This block covers all of them except MILLISECOND and DAYOFYEAR, and both are valid grammar keywords on every supported version (SqlBaseParser.g4:931-935 on branch-3.5, :1150-1154 on branch-4.0), so both are testable as written. DAYOFYEAR is worth adding because it is an alias arm sharing a case with DAY (case "DAY" | "DAYOFYEAR", :687), so nothing currently proves the alias both parses and dispatches.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Added. MILLISECOND and DAYOFYEAR are in the all-units block now, with a comment noting DAYOFYEAR shares the case "DAY" | "DAYOFYEAR" arm so the alias is proven to parse and dispatch.

Comment on lines +36 to +44
SELECT
timestampdiff(YEAR, a, b),
timestampdiff(MONTH, a, b),
timestampdiff(WEEK, a, b),
timestampdiff(DAY, a, b),
timestampdiff(HOUR, a, b),
timestampdiff(MINUTE, a, b),
timestampdiff(SECOND, a, b)
FROM test_timestampdiff

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.

timestampDiffMap covers MICROSECOND, MILLISECOND, SECOND, MINUTE, HOUR, DAY, WEEK, MONTH, QUARTER, YEAR (Spark DateTimeUtils.scala:719-730). This block omits MICROSECOND, MILLISECOND and QUARTER. QUARTER is the one entry in that map with arithmetic of its own (ChronoUnit.MONTHS.between(...) / 3, integer division toward zero), so it is both the most likely to diverge and the only one untested.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Added MICROSECOND, MILLISECOND and QUARTER in a second block, with QUARTER also exercised on literals. Agreed it was the risky one, being the only entry with its own arithmetic.

-- under the License.

-- timestampadd runs through the codegen dispatcher so results match Spark exactly.
-- Config: spark.sql.session.timeZone=America/Los_Angeles

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.

Same applies to timestampdiff.sql:

Pinning America/Los_Angeles only pays off if a fixture straddles a transition. Both functions do calendar arithmetic on local time, not elapsed time: timestampadd goes through timestampAddInterval, which is .atZone(zoneId).plusDays(...).plus(micros) (Spark DateTimeUtils.scala:269-274), and timestampdiff converts both sides with getLocalDateTime and then calls ChronoUnit.X.between on the local values (:745-747). So timestampdiff(HOUR, timestamp'2024-03-09 12:00:00', timestamp'2024-03-10 12:00:00') is 24 despite 23 real hours elapsing, and timestampadd(DAY, 1, timestamp'2024-03-09 12:00:00') is a local plus-one-day rather than plus-24-hours. Those are precisely the results a chrono-based native implementation gets wrong, which is the stated reason for choosing the dispatcher, and neither fixture asserts them. Please add a spring-forward row, a fall-back row (2024-11-03 01:30 is ambiguous), and one landing on the nonexistent local hour, timestampadd(HOUR, 1, timestamp'2024-03-10 01:30:00').

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Added to both fixtures. timestampadd.sql has a spring-forward block, a fall-back block using the ambiguous 2024-11-03 01:30, and timestampadd(HOUR, 1, timestamp'2024-03-10 01:30:00') landing on the nonexistent local hour, plus the plus-one-day versus plus-24-hours pairing on both transitions. timestampdiff.sql has the matching cases, including timestampdiff(HOUR, ...) across spring-forward reporting 24 despite 23 elapsed hours.

The comments in both files spell out why: local-time calendar arithmetic via timestampAddInterval on one side and getLocalDateTime plus ChronoUnit.X.between on the other.

SELECT
timestampadd(HOUR, 3, timestamp'2024-01-01 10:00:00'),
timestampadd(MONTH, 1, timestamp'2024-01-31 00:00:00'),
timestampadd(YEAR, 1, timestamp'2024-02-29 00:00:00')

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.

DateTimeUtils.timestampAdd wraps its whole body in a try that maps ArithmeticException and DateTimeException to timestampAddOverflowError (Spark DateTimeUtils.scala:709-713), which raises error class DATETIME_OVERFLOW (QueryExecutionErrors.scala:2484-2492). Nothing here reaches it, and this is the only path in the PR where an exception has to cross out of the generated kernel and surface as the same Spark error. Add query expect_error(DATETIME_OVERFLOW) over something like timestampadd(YEAR, 1000000000, timestamp'2024-01-15 10:30:45'); the existing valid-input blocks already satisfy the sentinel requirement.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Added: query expect_error(DATETIME_OVERFLOW) over timestampadd(YEAR, 1000000000, timestamp'2024-01-15 10:30:45'), which overflows in Math.multiplyExact(quantity, MONTHS_PER_YEAR). Confirmed it passes on 3.4, 3.5, 4.0 and 4.1, so the exception does cross out of the generated kernel with the condition name intact.

Worth recording for anyone adding a similar case: this is version-sensitive. The make_interval ANSI overflow I added in the same pass could not match on ARITHMETIC_OVERFLOW, because 3.5 renders that one as a bare "integer overflow. If necessary set ..." with no condition prefix. That fixture matches on overflow instead. DATETIME_OVERFLOW happens to carry its prefix on every supported version.

-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true

statement
CREATE TABLE test_timestampadd(ts timestamp, q int) USING parquet

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.

TimestampAdd.dataType is timestamp.dataType (Spark datetimeExpressions.scala:3443), so an NTZ input yields an NTZ output. That is a different Arrow vector class in the kernel (TimeStampMicroVector instead of TimeStampMicroTZVector, CometBatchKernelCodegen.scala:71-72), hence a separately compiled kernel, and zoneIdForType resolves to UTC instead of the session zone. This is a distinct code path rather than a variation, so please add an NTZ column here. For timestampdiff, inputTypes is Seq(TimestampType, TimestampType) (datetimeExpressions.scala:3531) so NTZ gets cast up; one case is enough to lock that in.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Added an NTZ column to timestampadd.sql with its own query block covering HOUR, MONTH, DAY and MICROSECOND, and an NTZ block in timestampdiff.sql to lock in the cast-up behavior.

@andygrove andygrove removed this from the 1.0.0 milestone Jul 27, 2026
Route TimestampAdd and TimestampDiff through the JVM codegen dispatcher so
they run inside the native Comet pipeline instead of forcing a full operator
fallback to Spark. The dispatcher executes Spark's own generated code, so
results match Spark exactly across all supported versions and avoid the
datetime incompatibilities that a chrono-based native implementation would
risk.

make_interval was considered but is left for a follow-up: its output type
CalendarIntervalType is not yet supported by Comet's columnar layer, so the
dispatcher cannot produce it.
CalendarIntervalType support landed in apache#4898, so MakeInterval can now go
through the JVM codegen dispatcher alongside TimestampAdd and TimestampDiff
instead of falling the projection back to Spark.

Also broadens the datetime test coverage:

- timestampadd gains MILLISECOND and DAYOFYEAR, a TIMESTAMP_NTZ column
  (a separately compiled kernel, since TimestampAdd.dataType follows the
  input), DST spring-forward, fall-back and nonexistent-local-hour cases,
  and a DATETIME_OVERFLOW expect_error case.
- timestampdiff gains MICROSECOND, MILLISECOND and QUARTER, NTZ inputs,
  and the matching DST cases.
- The grammar routes DATEADD/DATE_ADD and DATEDIFF/DATE_DIFF/TIMEDIFF to
  TimestampAdd/TimestampDiff when the first argument is a datetimeUnit
  keyword. Those alias spellings are now covered, split by the version
  that introduced them: dateadd/datediff inline, date_add/date_diff in
  date_add_unit_alias.sql (Spark 3.5+), timediff in timediff.sql (4.0+).
- make_interval and make_interval_ansi fixtures, the latter covering the
  ANSI overflow path where MakeInterval.failOnError is true.

Corrects the date_add, date_diff, dateadd and datediff rows in the
expression guide, which claimed Native without noting that this is true
only of the two-argument form, and adds rows for the unit spellings.
@andygrove andygrove changed the title feat: support timestampadd and timestampdiff via codegen dispatch feat: support timestampadd, timestampdiff and make_interval via codegen dispatch Jul 28, 2026
@andygrove

Copy link
Copy Markdown
Member Author

Rebased onto main and addressed both rounds. Replies are inline on each thread; summarizing the cross-cutting items here.

make_interval is now in this PR rather than deferred, since @mbutrovich is right that #4898 makes the stated reason stale. It is object CometMakeInterval extends CometCodegenDispatch[MakeInterval] plus the classOf[MakeInterval] entry, exactly as you both described, and the PR title and description are updated.

@peterxcli, on the native route: I looked for SparkMakeInterval and found no make_interval implementation anywhere under native/, so I went with codegen dispatch to match the rest of this PR. If you were thinking of an upstream datafusion-spark function, a native path would be a reasonable follow-up, but it would want its own PR given the CalendarInterval ANSI semantics.

Per @mbutrovich's note, the ANSI overflow path has its own coverage in make_interval_ansi.sql, where failOnError is true, alongside make_interval.sql for the non-ANSI path where the same input returns NULL instead.

A pre-existing bug surfaced while writing those tests, filed as #5058. A null CalendarIntervalType literal in a projection throws CometNativeException: Interval(MonthDayNano) is not supported in Comet instead of falling back or evaluating. #4898 added the type to the support checks, codegen kernels and serialization, but the null-literal match in planner.rs has no Interval arm, so the plan is accepted and then fails natively. YearMonthIntervalType is not in the supported set and falls back cleanly, which is the behavior a null calendar interval should have at minimum.

It is easy to hit indirectly: MakeInterval is NullIntolerant, so NullPropagation rewrites make_interval(NULL, 2, 3, ...) into a null interval literal and the expression vanishes before Comet sees it. Not caused by this PR, and since literal NULL arguments never exercise the expression anyway, the fixture drives its NULL coverage from nullable columns and points at #5058.

@peterxcli, on the non-blocking suggestions: MILLISECOND is in both fixtures, MICROSECOND is in timestampdiff, and both DST boundaries and TIMESTAMP_NTZ are covered rather than just one of them.

Verified green on all four supported profiles (3.4, 3.5, 4.0, 4.1). One thing worth flagging from that sweep: the alias coverage had to be split by Spark version, because 3.4's grammar rule is only TIMESTAMPADD | DATEADD and TIMESTAMPDIFF | DATEDIFF, with DATE_ADD and DATE_DIFF joining in 3.5. Details are in the thread on expressions.md.

@peterxcli peterxcli left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

@andygrove thanks for the update, LGTM

on the native route: I looked for SparkMakeInterval and found no make_interval implementation anywhere under native/,

yes, the native implementation would be handled in #5039 + #5131

so I went with codegen dispatch to match the rest of this PR.

yes, make sense, codegen dispatch support minimize the change in this PR as first step.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants