[tlfua]

Which issue does this PR close?

Part of #7246 (the top-K case). The OR/AND/IN cases were fixed by #7326; the arithmetic and concat cases are handled in a separate PR.

Rationale for this change

With spark.sql.ansi.enabled=true, SELECT a, b FROM topk ORDER BY a, b + 1 LIMIT 2 over (1, MaxInt), (2, 5), (3, MaxInt), (4, 7), (5, 9) returns (1, 2147483647), (2, 5) in Spark but fails with ARITHMETIC_OVERFLOW in Comet.

Spark's TakeOrderedAndProjectExec compares rows with LazilyGeneratedOrdering, whose generated comparator (GenerateOrdering.genComparisons) evaluates sort key i for a pair of rows only when keys 0..i-1 tie. Comet's native TopK (DataFusion SortExec with fetch) evaluates every sort key over the whole batch, so a later key that can fail raises on rows Spark never evaluates it for.

Making the native TopK evaluate later keys lazily was considered and rejected for this PR: which row pairs Spark compares depends on Guava's Ordering.leastOf, and DataFusion's TopK keeps a heap across batches, so per-batch masking cannot reproduce Spark's evaluation set without a custom operator. A scoped fallback is correct and small, in line with #6006.

What changes are included in this PR?

• CometTakeOrderedAndProjectExec.getSupportLevel returns Incompatible when any sort key after the first can fail. The first key is evaluated for every row in Spark too when the input has at least two rows, so it is not restricted; the single-row case is tracked in Native TopK evaluates sort keys for a single-row input, raising false ANSI errors #7527 . spark.comet.operator.TakeOrderedAndProjectExec.allowIncompatible=true still runs it natively.

• QueryPlanSerde.laterSortKeyThatCanFail(sortOrder) returns the first key after the leading one that can fail; the plain sort follow-up ( Native sort evaluates later sort keys on rows Spark's ordering skips, raising false ANSI errors #7470 ) can reuse it.

• QueryPlanSerde.cannotFail(expr): a conservative whitelist judged from Spark semantics (read for 3.4.3, 3.5.8, 4.0.1 and 4.1.1), independent of the native cannot_fail, with each native implementation also checked for data-dependent errors: columns, literals, comparisons, IN, boolean logic, null checks, If/CaseWhen/Coalesce, float normalization, struct field access (s.x), float/double +/-/*, negation and abs in any mode, other numeric arithmetic / negation / abs outside ANSI mode, casts that cannot fail, and:

  • string: substring, upper, lower, length, octet_length, bit_length, trim/ltrim/rtrim, replace, initcap, concat_ws, concat (string, at least one argument), reverse (string)
  • datetime: year, quarter, month, day, dayofweek, weekday, dayofyear, hour, minute, second
  • math: floor, ceil, sqrt, exp, log, signum, greatest, least, isnan, nanvl; round/bround for float/double, and for integral types only with ANSI off (never decimal)
  • hash: hash, xxhash64, md5, sha1, sha2, crc32

  Anything else, including lpad/rpad/repeat, element_at, date_add and to_date, is treated as able to fail. arrays.scala still has its own narrower cannotThrow for array_append; moving it onto cannotFail is tracked in Consolidate the Scala "cannot fail" checks: arrays.cannotThrow and QueryPlanSerde.cannotFail #7585 .

• CometCast.cannotFail(cast): TRY casts, LEGACY casts between atomic types, string to binary, and ANSI casts that keyCastCannotFail already proves exact. A cast to or from a timestamp never qualifies: the native kernels go through chrono, which rejects years Spark accepts, and raise even for TRY ( Native timestamp casts fail or panic for years beyond chrono's range #7471 ). Nor does a cast from a date to a string: arrow-cast's formatter ignores safe and raises for such years in every mode. Both are checked at any depth inside arrays, structs and maps with SupportLevel.containsType, since a nested cast applies the same kernels to each element, field, key and value.

• docs/source/user-guide/latest/compatibility/operators.md documents the fallback, notes that the check does not depend on ANSI, and that a sort without a LIMIT and with no exchange beneath it (such as SORT BY) still runs natively.

• expressions/array/array_null_short_circuit.sql: its two-key TopK query has a later key array_position(ai, CAST(s AS INT)), whose ANSI cast can fail, so it now falls back; the query is wrapped with spark.comet.operator.TakeOrderedAndProjectExec.allowIncompatible=true to keep testing the native array path.

Known trade-off: the whitelist is conservative, so a later key built from a function outside it also falls back even though it may not fail, with ANSI on or off. Global ORDER BY and window ordering are not affected in Spark either (the range exchange or window pre-projection evaluates every key); the plain-sort divergence for SORT BY and single-partition sorts is tracked in #7470 .

How are these changes tested?

New SQL file test spark/src/test/resources/sql-tests/operators/topk_later_key_ansi.sql (ANSI on, AQE off):

• the reproduction from the issue and two more fallible later keys (expect_fallback),
• later keys that cannot fail, and a fallible first key, stay native (query),
• rows that tie on the first key still raise ARITHMETIC_OVERFLOW (expect_error),
• with ANSI off, the same queries plus -b/abs(b) stay native,
• string later keys (substring/upper/lower/length) stay native. Without this, TPC-DS q79 plan stability changed on Spark 4.x,
• a fallible key nested under an allowed one (coalesce(b + 1, 0), IF(a > 3, b + 1, 0)) and round(b, -1) under ANSI fall back; try_cast stays native, timestamp casts (including TRY) fall back,
• a topk_types table with NULL, NaN, -0.0, int/long bounds, empty strings and boundary dates runs the string, datetime, math and hash keys natively,
• under ANSI, float/double *, +, -, negation and abs stay native, and CAST(dt AS STRING) falls back,
• a struct field key s.x stays native with ANSI on and off,
• with ANSI off, a LEGACY string-to-int cast, round(b, -1), b / 0, b % 0, b DIV 0, l + 1 and l * 2 stay native.
• a topk_nested table with out-of-range values (timestamp_micros(9000000000000000000), date_from_unix_date(100000000)): try_cast over nested timestamps and dates (array, struct, map, and array<struct<timestamp>>) falls back. The file writes timestamps as TIMESTAMP_MICROS.

Commands run locally (Spark 4.1 default profile, macOS):

• ./mvnw test -Dsuites="org.apache.comet.CometSqlFileTestSuite topk_later_key" -Dtest=none -Dmaven.gitcommitid.skip=true without the fix: failed with ARITHMETIC_OVERFLOW on ORDER BY a, b + 1 LIMIT 2, as expected.
• Same command with the fix: passed (1 succeeded, 0 failed).
• Same command for topk_later_key array_null_short_circuit and four other related SQL files after the second commit: passed (6 succeeded, 0 failed).
• After the review changes, topk_later_key (1 succeeded) and array_null_short_circuit (5 succeeded): passed.
• After coderfender's review, merged with the latest main: topk_later_key (1 succeeded) and array_null_short_circuit (5 succeeded): passed.
• After rich7420's review (nested types): topk_later_key (1 succeeded, 0 failed): passed.
• ./mvnw spotless:apply -Dmaven.gitcommitid.skip=true: no changes.
• After the second round from rich7420 and coderfender (reuse SupportLevel.containsType, no behavior change): make format PROFILES="-Dmaven.gitcommitid.skip=true" (formats and compiles, Spark 4.1): passed, no changes. The CI preflight checks (./dev/release/check-rat.sh HEAD, prettier --check "**/*.md" and the dev/ci scripts): passed.

Fork CI ( https://github.com/tlfua/datafusion-comet/actions/runs/42730158514 ): all jobs passed on the latest commit, including lint for all Spark profiles, TPC-H/TPC-DS result verification, and the Comet suites (expressions, exec, shuffle, scans) on Spark 4.1 / JDK 17. Since this touches the serde, the run-spark-3.5-tests, run-spark-4.1-tests and run-all-spark-profiles labels run the Spark SQL suites and the other Spark profiles.

[tlfua]

Fork CI is green ( https://github.com/tlfua/datafusion-comet/actions/runs/41611703653 ). Since this changes the serde/planner support level for TopK, could a maintainer please add the run-spark-4.1-tests label so the Spark SQL suite runs? Thanks!

[comphead]

Thanks for the clear write-up of how LazilyGeneratedOrdering evaluates later keys. The first-key and later-key description matches TakeOrderedAndProjectExec and GenerateOrdering.genComparisons at v3.4.3, v3.5.8, v4.0.1 and v4.1.1. I have not run the fixture. The fork CI covers Spark 4.1 only, so the Spark 3.5 and 4.0 plans and the Spark SQL suite still need the run-* labels.

[comphead]

The description leaves the plain sort for a follow-up. Reading Spark, it has the same difference: SortExec.createSorter builds its comparator with RowOrdering.create, which goes through GenerateOrdering, so only the first key is computed for every row, as the prefix. CometSortExec has no guard, so I expect ORDER BY a, b + 1 without a LIMIT to raise in Comet on the fixture data. I have not run it. Window ordering goes the same way. Since a plain sort is far more common than a top-K, could you file an issue for it and link it here and in the docs? If the guard will be reused there, a helper over Seq[SortOrder] would save copying drop(1).find into each operator.

[tlfua]

Thanks. I checked each plain-sort shape against Spark with ANSI on and filed #7470 .

Global ORDER BY a, b + 1 and window ordering (rank, row_number with PARTITION BY) throw in Spark too, because the range exchange or the window pre-projection evaluates every key. The divergence shows only when nothing beneath the sort evaluates all keys first, for example SORT BY or a sort over a single-partition child such as LIMIT. I extracted the check into QueryPlanSerde.laterSortKeyThatCanFail(Seq[SortOrder]) so the plain sort can reuse it in that follow-up.

[comphead]

This set decides the fallback for every top-K with more than one key, with ANSI on or off. Later keys such as year(d), trim(s), round(x, 2) or concat fall back even though they cannot fail, and substring, upper, lower and length came in to fix the q79 plan change. Is that narrow set intended? If not, the common infallible functions could be added with tests.

[tlfua]

Widened the list after reading the Spark source for 3.4.3, 3.5.8, 4.0.1 and 4.1.1. Added:

• String: octet_length, bit_length, trim/ltrim/rtrim, replace, initcap, concat_ws, concat (string, at least one argument), reverse (string)
• Datetime: year, quarter, month, day, dayofweek, weekday, dayofyear, hour, minute, second
• Math: floor, ceil, sqrt, exp, log, signum, greatest, least, isnan, nanvl, IN; round/bround for float/double always, and for integral types only when ANSI is off
• Hash: hash, xxhash64, md5, sha1, sha2, crc32
• Cast: string to binary (inserted implicitly by md5(s) and friends)

I also checked each one's native implementation (the Comet kernels plus DataFusion / datafusion-spark 55.2) for errors that depend on the data. The only real ones were timestamp casts: they go through chrono, which rejects years Spark accepts, and raise even for TRY (filed as #7471 ). So any cast to or from TIMESTAMP/TIMESTAMP_NTZ now counts as able to fail. The rest error only at planning time or beyond 2GB of string offsets per batch.

Left out on purpose: timestamp casts (above), decimal round (the negative-scale path builds a Decimal that can fail the precision check), integral round under ANSI (overflow), lpad/rpad/repeat (size limits), element_at, date_add, to_date, and anything that parses input.

On upkeep: the list fails safe. A function missing from it only falls back to Spark, so leaving one out costs performance, never correctness. Each entry was checked against Spark 3.4.3 to 4.1.1 and the native kernel, and topk_later_key_ansi.sql runs each one natively over NULL and boundary values.

[comphead]

The comment says this works like the native cannot_fail, but the two sets differ. The native one in conditional_funcs/case_when.rs (line 449) admits no functions and no integer or decimal Divide, while this one admits Substring, Upper, Lower, Length, Coalesce, all non-ANSI arithmetic and every LEGACY atomic cast. The difference may be right, but the wording should say they are independent. Also, arrays.scala already has a private cannotThrow (line 86) answering the same question for array_append. Could it call this helper, so there is one Scala notion of an expression that cannot fail?

[tlfua]

Good catch. I reworded the comment: the list now states that it answers from Spark semantics alone and is independent of the native cannot_fail in case_when.rs. I left cannotThrow in arrays.scala as is, because switching it to this list would change when array_append runs natively; I'd rather do that in a separate PR with its own tests.

[tlfua]

Following up on the cannotThrow part: opened #7585 to have arrays.scala (and NativeLambdaSerde.mayRaise from #7322 ) reuse QueryPlanSerde.cannotFail. It changes when array_append runs natively, so it needs its own tests and stays out of this PR.

[comphead]

The native cases here cover only int columns without NULLs. Could you add a fallible later key nested under an allowed parent (coalesce(b + 1, 0), IF(a > 3, b + 1, 0)), which is the only coverage of children.forall(cannotFail)? Also try_cast and a LEGACY cast with ANSI off for the two CometCast.cannotFail branches, later keys over bigint, decimal, double, string, date and timestamp with NULLs, NaN, -0.0 and bound values, and a non-ANSI /, % and DIV by zero.

[tlfua]

Added:

• Fallback: a fallible key nested under an allowed one (coalesce(b + 1, 0), IF(a > 3, b + 1, 0)), and round(b, -1) under ANSI.
• Native: try_cast, and a topk_types table with NULL, NaN, -0.0, int/long bounds, empty strings, and dates such as 1970-01-01, 2000-02-29 and 9999-12-31, run through the string, datetime, math and hash keys.
• ANSI off, running natively: a string-to-int cast in LEGACY mode, round(b, -1), b / 0, b % 0, b DIV 0, l + 1 and l * 2 as later keys.

[comphead]

The check does not read spark.sql.ansi.enabled. With ANSI off, it also falls back for any later key outside the allowed set. Could the text say that a later key Comet cannot show to be infallible runs on Spark in either mode, and that Sort and window ordering still run natively with the same difference, linking the issue?

[tlfua]

Reworded. The section now says the check does not depend on ANSI: Comet falls back whenever it cannot show that a later key never fails. It also notes that a sort without a LIMIT and with no exchange beneath it, such as SORT BY, still runs natively and can raise, and links #7470 .

[tlfua]

@comphead thanks for the review, I've addressed the comments above. Could a maintainer add run-all-spark-profiles and run-spark-4.1-tests? Spark 3.x defaults to ANSI off, and a later key outside the whitelist now falls back with ANSI off too, so plan-stability goldens on 3.x may change; the fork CI only covers Spark 4.1.

[coderfender]

Thanks for this. I checked Spark's TakeOrderedAndProjectExec and GenerateOrdering in 3.4, 3.5, 4.0 and 4.1, and the description of later keys is right. A later key is only compared when the earlier ones tie. The check sits on the one entry point, so it covers the local TopK, the multi-partition case and AQE. Reading evalMode per expression rather than spark.sql.ansi.enabled keeps TRY and LEGACY native, which is the right call. The CometCast addition is self-contained too, and it doesn't touch any existing support level.

There are two places where the list is stricter than Spark, and TopK queries that ran natively and correctly in 1.2 now fall back by default. I'd like these fixed before merge.

1. spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala:1432-1436. Under ANSI, Spark's float and double +, -, *, unary minus and abs can't raise. DoubleExactNumeric and FloatExactNumeric only override toInt, toLong and compare, and the native checked kernels don't error for floats either. Right now ORDER BY k, price * 1.1 over a double falls back on Spark 4.x defaults, even though it ran natively before. Could these five allow FloatType and DoubleType regardless of eval mode, with a test under ANSI? Divide, Remainder and Pmod would stay excluded, since division by zero raises for doubles too.

2. QueryPlanSerde.scala:1405. GetStructField isn't in the list, so ORDER BY a, s.x LIMIT n now falls back in both modes. It ran natively in 1.2. Spark's nullSafeEval there is just a field read. Could it join the list, with a struct column case in topk_later_key_ansi.sql?

Two more that the PR didn't cause, but that need a response.

3. spark/src/main/scala/org/apache/comet/expressions/CometCast.scala:582. The timestamp exclusion makes sense, but DATE to STRING has the same chrono limit. It goes through the DataFusion fall-through in cast.rs to arrow's value_to_string, and ValueFormatter::write ignores safe and returns Failed to convert ... to temporal for Date32 in every mode. One example is #7471's date_from_unix_date(100000000). As written, ORDER BY a, CAST(dt AS STRING) and implicit casts like concat(s, dt) stay native. Could a date source be excluded for the cast to string as well, and the case added to #7471?

4. spark/src/test/resources/sql-tests/operators/topk_later_key_ansi.sql:110. A small correction on "The first key is evaluated for every row in Spark too". Spark's takeOrdered goes through Guava's TopKSelector, which stores the first row without comparing it. So when the TopK's whole input is a single row, Spark evaluates no key at all, and WHERE id = 5 ORDER BY b + 1 LIMIT 10 on one overflowing row returns the row in Spark but raises in Comet. That isn't this PR's doing. Could the comment say "for every row when the input has at least two rows", and could you file a tracking issue for the single row case?

The check applies with ANSI off too, so Spark 3.4 and 3.5 plans are the ones that can move. I'll make sure all profiles and the Spark SQL suite run before it goes in the queue.

[tlfua]

Thanks @coderfender, all four points were valid. Addressed in c181b6e (also merged the latest main):

1. Float/double +, -, *, negation and abs are now allowed in any eval mode. Divide, IntegralDivide, Remainder and Pmod stay excluded under ANSI, and integral types are unchanged. Added an ANSI test ORDER BY a, d * 1.1D, d + 1, d - 1, -d, abs(d) that stays native.
2. GetStructField is allowed (its child must also pass). Added a struct table with ORDER BY a, s.x, s.y staying native with ANSI on and off.
3. A cast from DATE to STRING is now treated as able to fail in every mode, with an expect_fallback test for CAST(dt AS STRING). Reproduced locally (Cast error: Failed to convert 100000000 to temporal for Date32 for date_from_unix_date(100000000)) and added it to Native timestamp casts fail or panic for years beyond chrono's range #7471 (comment) .
4. The test comment now says the first key is evaluated for every row only when the input has at least two rows. Reproduced the single-row case locally (Spark returns the row, Comet raises ARITHMETIC_OVERFLOW) and filed Native TopK evaluates sort keys for a single-row input, raising false ANSI errors #7527 .

Fork CI on the new head: all jobs passed ( https://github.com/tlfua/datafusion-comet/actions/runs/42145536518 ).

[rich7420]

try_cast(ts_array AS STRING) bypasses this timestamp check. With Long.MaxValue microseconds, the native cast errors while Spark formats it. Check nested timestamp types too, with a TopK regression.

[tlfua]

Thanks, good catch. Fixed in 5a5b423. CometCast.cannotFail now looks inside arrays, structs and maps: a timestamp anywhere in the source or target type makes the cast able to fail, and so does a date in the source combined with a string in the target. A nested cast applies the same chrono-based kernels to each element, field, key and value, so they raise under TRY just like the top-level cast ( #7471 ).

topk_later_key_ansi.sql adds a topk_nested table with timestamp_micros(9000000000000000000) and date_from_unix_date(100000000), and checks that these later keys fall back: try_cast of array<timestamp> and array<date> to STRING and to array<string>, struct<timestamp> and struct<date> to STRING, struct<date> to struct<string> and struct<timestamp>, struct<timestamp> to struct<date>, map<int, date> to map<int, string>, map<timestamp, int> to map<string, int>, and array<struct<timestamp>> to STRING.

[rich7420]

@tlfua thanks for the patch!

[rich7420]

Could we reuse SupportLevel.containsType here? It walks arrays, maps, and structs the same way, e.g. containsType(from, classOf[TimestampType], classOf[TimestampNTZType]). topk_later_key_ansi.sql still passed locally with existsNested replaced.

[tlfua]

Thanks, done in 37be546 . cannotFail now uses SupportLevel.containsType for both the timestamp check and the date-to-string check, and the private existsNested is gone. Behavior is unchanged, so the existing topk_nested queries in topk_later_key_ansi.sql still cover it.

[rich7420]

nit: hasTimestamp already rules out any timestamp, so the TimestampType/TimestampNTZType cases here never affect the result. Could we remove them?

[tlfua]

Good catch, removed in 37be546 . isAtomicCastType now lists only DateType besides the numeric, string, boolean and binary types, with a comment that timestamps never reach it because cannotFail rejects them first.

[coderfender]

Thanks for working through all of these. The float arithmetic and GetStructField cases now stay native, and the date to string exclusion and the nested check from rich7420's comment look right to me. #7471 and #7527 track the rest, so I'm dismissing my earlier request for changes. Two things I'd still like to see before merge.

1. spark/src/main/scala/org/apache/comet/expressions/CometCast.scala:588-614. SupportLevel.containsType already walks array elements, struct fields and map keys and values the same way existsNested does, and SupportLevel is imported here. Could cannotFail use it and drop existsNested and the local isTimestamp? Something like

   val hasTimestamp =
     Seq(from, to).exists(containsType(_, classOf[TimestampType], classOf[TimestampNTZType]))
   val dateToString =
     containsType(from, classOf[DateType]) && containsType(to, classOf[StringType])

   The result is the same, since the date and timestamp types are case objects with private constructors. It leaves one recursive type walk instead of a third copy. Spark's DataType.existsRecursively would have been the obvious choice, but it is private[spark] in every version we support.

   While you're here, isAtomicCastType still lists TimestampType and TimestampNTZType. hasTimestamp rules them out before the LEGACY arm runs, so those arms can never match. Could they go? As written the list reads as if a LEGACY timestamp cast can't fail, which is the opposite of what #7471 found, and anyone who later relaxes the timestamp check would admit them without noticing.

2. On comphead's cannotThrow thread. Keeping cannotThrow as it is makes sense for this PR. Could you file an issue for moving it onto cannotFail and link it there? #7322 adds a third list in NativeLambdaSerde.mayRaise, and #7535 is already about that one drifting, so it would be good to have the consolidation tracked.

All four points are addressed in c181b6e and 5a5b423. Remaining comments are in my latest review.

[coderfender]

@tlfua I added run-all-spark-profiles, run-spark-4.1-tests and run-spark-3.5-tests. The check also applies with ANSI off, so the Spark 3.5 plans are worth seeing too.

[tlfua]

@coderfender thanks for the review and for adding the labels.

1. containsType / dead timestamp arms: done in 37be546 . cannotFail uses SupportLevel.containsType as you suggested, the local isTimestamp and existsNested are gone, and isAtomicCastType no longer lists the timestamp types.
2. cannotThrow: opened Consolidate the Scala "cannot fail" checks: arrays.cannotThrow and QueryPlanSerde.cannotFail #7585 and linked it in fix: fall back TopK whose later sort key can fail under ANSI #7457 (comment) .

Fork CI on the new commit is green: https://github.com/tlfua/datafusion-comet/actions/runs/42730158514 . Could you approve the workflows on this commit so the labelled Spark 3.5 / 4.1 / all-profiles runs start? They are waiting for approval.