feat: add native spark_sequence kernel for integral element types

[0lai0]

Which issue does this PR close?

Closes #5349.

Rationale for this change

Integral sequence(start, stop[, step]) currently runs through the JVM codegen dispatcher, which allocates two long[] per row. This PR adds a native kernel that reserves the Arrow child buffer once per batch. Date and timestamp sequences stay on the dispatcher (timezone / DST / legacy calendar).

What changes are included in this PR?

• Native spark_sequence kernel (native/spark-expr/src/array_funcs/sequence.rs) for Byte / Short / Int / Long. Matches Spark Sequence.sequenceLength, including the overflow-report paths (2^63, 2^63+1, and the internal-error edge).
• CometSequence with CodegenDispatchFallback:
  • Integral types with leaf arguments only (literals or column references) route to the native kernel.
  • Integral types with non-leaf sub-expressions (e.g. CASE WHEN, nested sequence(...)) return Unsupported and stay on the JVM codegen dispatcher, preserving Spark's per-row null short-circuit.
  • Date/timestamp/timestamp_ntz sequences return Unsupported and stay on the dispatcher.
• Comet-specific SequenceBatchTooLarge error when the batch's total generated elements exceed the Arrow i32 offset ceiling (i32::MAX) or try_reserve_exact fails. The message names spark.comet.batchSize as the actionable knob. Spark itself has no equivalent limit because it stores each row as its own long[].
• Error mapping via ShimSparkErrorConverter: Spark 3.x throws IllegalArgumentException("Illegal sequence boundaries: ..."); Spark 4.x throws SparkIllegalArgumentException("_LEGACY_ERROR_TEMP_3243").
• CollectionSizeLimitExceeded now carries a decimal String count (can exceed i64) and a function_name for Spark 4.x. This is the first producer of that error, and it also fixes a latent Spark 3.5 shim bug that passed a Scala tuple as count and rendered (array,N).
• User-visible error format change (unrelated to sequence): new case "Internal" in ShimSparkErrorConverter maps ~20 existing SparkError::Internal producers (in temporal.rs, numeric.rs, conversion_funcs/string.rs, rlike.rs, etc.) from SparkException(message, <text>) to [INTERNAL_ERROR] <text>.
• Docs: sequence marked Hybrid in expressions.md; audit notes under array_funcs.md document the per-batch ceiling and the non-leaf argument fallback.

Limitations

Native integral sequence materializes every row's output into one Arrow child buffer per batch. The sum of all row lengths in a batch must fit in an i32 offset buffer. A query that Spark runs fine (e.g. sequence(0, 262143) over a full 8192-row batch) may fail in Comet with SequenceBatchTooLarge; lowering spark.comet.batchSize is the fix.

How are these changes tested?

• Five unit tests in sequence.rs.
• spark/src/test/resources/sql-tests/expressions/array/sequence.sql: integral types, default/explicit step, nulls, explode, seven error cases, plus:
  • full byte/short range (sequence(-128Y, 127Y), sequence(-32768S, 32767S)),
  • Int32 step overflow boundary (sequence(-2147483648, 2147483647, 1073741824)),
  • nested null short-circuit (sequence(s, size(sequence(1, 5, k)))),
  • CASE WHEN guarded branch,
  • date/timestamp dispatcher coverage.
• CometCodegenSuite: leaf-arg integral sequences run natively; non-leaf integral and temporal sequences show "JVM codegen dispatcher".
• Confirmed no Spark SQL suite tests match the old (message,<text>) shape for SparkError::Internal.

Criterion (cargo bench --bench sequence, N=8192). Absolute numbers; the kernel is not on main. Benchmarks use leaf-arg shapes only.

Shape	Time	Elems/batch	ns / elem
short_2_elems	49.83 µs	16,384	3.04
short_5_elems	55.35 µs	40,960	1.35
long_365_elems	1.318 ms	2,990,080	0.44
long_10000_elems	34.87 ms	81,920,000	0.43
descending_365_elems	1.324 ms	2,990,080	0.44
zero_step_start_eq_stop	35.70 µs	8,192	4.36
sparse_nulls_365_elems	1.200 ms	2,691,072	0.45
dense_nulls_365_elems	673.2 µs	1,495,040	0.45
error_illegal_boundaries	673 ns	(errors)	-

Spark (CometSequenceBenchmark, 8192 rows, Apple M5, Spark 4.1.3 / Scala 2.13)

Shape	Spark best (ms)	Comet best (ms)	speedup
seq_short_5_elems	15	5	2.9X
seq_spine_365_elems	14	6	2.2X
seq_long_10000_elems	59	57	1.0X
seq_descending_default_step	13	6	2.3X
seq_explicit_step_7	11	3	3.2X
seq_sparse_nulls_365_elems	12	5	2.4X
seq_date_spine_dispatcher	12	12	1.0X

seq_date_spine_dispatcher is a control (date path unchanged). seq_long_10000_elems is memory-bandwidth-bound at 82M elements/batch. Typical speedup is 2X–3X on the shorter-list shapes the issue targets.

[andygrove]

I read this against the Spark sources for 3.4.3, 3.5.8, 4.0.1 and 4.1.1 and ran it locally: the Rust unit tests, clippy, CometSqlFileTestSuite sequence and CometCodegenSuite on both Spark 4.1 and 3.5, plus a throwaway fixture covering the edge cases in my third comment. All green. The sequenceLength port is faithful, including which of the three failure paths fires and the exact count each reports.

Three things below that I would like addressed before this goes in.

[andygrove]

Nice work on the error-parity side of this, the sequenceLength port matches Spark on all three failure paths and I checked it against 3.4.3 through 4.1.1.

The thing I keep coming back to is Vec::with_capacity(total). sequence is the first expression we have made native where the output size is unbounded relative to the input size. Every other with_capacity in array_funcs/ is sized by row_count or args.len(), but here total is the sum of every row's generated length, so a single batch can ask for up to i32::MAX elements, which is 16 GiB for bigint. Your own benchmark shows the shape: seq_long_10000_elems materializes 8192 x 10000 x 8 bytes, so 655 MB in one allocation, where Spark holds one row's long[] at a time. That is also the one row in your table with no speedup, which makes me wonder whether the large-per-row case is paying for itself at all.

Two things I would like to see. Could the allocation go through try_reserve so an oversized batch surfaces as a query error rather than an allocator abort that takes the executor down? And could the batch ceiling be documented somewhere the user can find it?

On the ceiling specifically, the message a user gets today is misleading. sequence(0, 262143) over a full 8192-row batch lands on exactly 2147483648 total elements and trips the check, even though every individual array is well inside Spark's limit and Spark itself would run the query. The shim ignores the max_elements you pass, so the message reads "Can't create array with 2147483648 elements which exceeding the array size limit 2147483632", which points the user at a per-array limit that they have not actually exceeded and gives them nothing actionable. The real fix on their side is to lower spark.comet.batchSize. Given that, is Compatible() the right support level, or should this at least get a compatible note and a line in the audit entry?

[0lai0]

Thanks @andygrove for review.
I switched to try_reserve_exact and added SequenceBatchTooLarge pointing at spark.comet.batchSize. Documented the per-batch ceiling in array_funcs.md and a new Limitations section. Leaf-arg integral sequences stay Compatible().

[andygrove]

The case "Internal" arm changes behavior well beyond sequence. There are around twenty SparkError::Internal producers today in temporal.rs, numeric.rs, conversion_funcs/string.rs and rlike.rs, and all of them previously fell through to the None branch in SparkErrorConverter, which renders as new SparkException(msgParams.mkString(", ")), so users saw (message,<text>). After this they all become [INTERNAL_ERROR] <text>.

That is a clear improvement and I am not asking you to revert it. Could you call it out in the PR description though? It is a user-visible message change for a set of expressions that have nothing to do with sequence, and right now the "How are these changes tested?" section does not mention it. It would also be worth a pass over the Spark SQL suite diffs to confirm nothing was matching on the old shape.

The same arm is added to the 3.5 and 4.x shims, so this applies to all three.

[0lai0]

Oh thats a good point, this is indeed a user-visible change beyond sequence.
I have added a bullet to the PR description under What changes are included calling out that the new Internal arm affects ~20 existing native expressions and changes the message from SparkException(message, ) to [INTERNAL_ERROR] . Thanks!

[andygrove]

The fixture is thorough on the error paths. There are three shapes I think are worth pinning down that it does not reach today.

The first is full narrow-type range, sequence(-128Y, 127Y) and sequence(-32768S, 32767S). That is the one place where Spark's arr(i) = start + step * num.fromInt(i) genuinely wraps at 8 and 16 bits while your kernel accumulates in i64 and truncates on the way out. The two agree because the true value is always in range, but it is the case I would most want a regression test on, and sequence(-128Y, -120Y) does not get there.

The second is a step whose product with the index overflows int, something like sequence(-2147483648, 2147483647, 1073741824), which exercises the Int32 monomorphization at the boundary.

The third is sequence under a CASE WHEN where the throwing branch is not taken, for example SELECT CASE WHEN step > 0 THEN sequence(1, 5, step) ELSE array(-1) END FROM t with rows carrying negative and zero steps. DataFusion filters the batch before evaluating each then branch so this works today, but it is the one construct where an eagerly evaluated throwing expression would diverge from Spark, and it would be cheap insurance.

I ran all three locally against this branch and they pass on both 4.1 and 3.5, so this is about locking the behavior in rather than chasing a suspected bug.

[0lai0]

Thanks, Added all three to sequence.sql: full byte/short range, Int32 step overflow, and CASE WHEN guarded branch.

[sunchao]

Reviewed 9d2f15751086e5b40f0464c84375bcb21ec170eb against 2949fd0d244ef0b201708820efbf2e07b7092156. One P2 in the native sequence path's argument null short-circuiting, detailed below. The witness is source-derived, not an executed query. I did not run the suites.

Current checks show 63 successful, 8 skipped and 1 failed. The failed Spark 4.0/JDK 21 execution job reports 780 tests passed before Maven dependency resolution failed with HTTP 403, rather than a demonstrated expression-test failure.

For the existing allocation/performance discussion, could we also compare this kernel with the pre-PR dispatcher using matched data and batch sizes, including concurrent tasks and the 10,000-element shape, reporting peak memory and checking output equality? The benchmark currently forces local[1], so the published timings do not cover concurrent allocation pressure.

[sunchao]

[P2] Preserve null short-circuiting before evaluating later arguments

Could this lowering preserve Sequence's left-to-right null guards? For a Parquet table t(s INT, k INT) containing (NULL, -1) and (1, 1), consider SELECT sequence(s, size(sequence(1, 5, k))) FROM t. Spark returns NULL for the first row without evaluating the inner sequence, and [1,2,3,4,5] for the second. Here both sequences become scalar UDFs, whose arguments DataFusion evaluates over the batch before calling the outer kernel. The inner sequence(1, 5, -1) therefore throws before the outer row_is_null check can discard that row. The previous dispatcher kept the whole expression tree inside Spark's guarded evaluation. Could we retain those guards or dispatch such shapes, and add a composed null/error regression case?

[0lai0]

Agreed, the native UDF path can't short-circuit per row. Shapes with non-leaf arguments now fall back to the codegen dispatcher via hasLeafArgsOnly. I added that query to sequence.sql and a routing check in CometCodegenSuite. Thanks @sunchao for reiview.

[sunchao]

Re-reviewed 51007cd45d57eb2fcb5ca78135a75bc5c3738c1b. The nested-sequence witness now routes through the dispatcher, but the existing P2 null-guard issue remains for a zero-argument Scala UDF: children.isEmpty admits more than safe references and literals.

With a scanned nullable s INT and a throwing, deterministic boom(): Int UDF, sequence(s, boom()) still evaluates boom() before the outer null check when spark.comet.exec.scalaUDF.codegen.enabled=true; Spark skips it for a null s. This is a source-derived residual case, not an executed reproduction. Could the native gate accept only safe reference/literal forms, or preserve the whole expression's guards?

One new test-fixture issue is noted inline. I did not run the suites, and current CI has not validated this head.

[sunchao]

[P2] Use legal bounds in the native-path fixture

withSequenceTable also inserts (a, b) = (9, 2), so the second expression becomes sequence(9, 2, 2). Spark rejects a positive step with descending bounds, and checkSparkAnswerAndOperator first collects the Spark reference result with Comet disabled. This test therefore raises before the output comparison or native-path assertion. Could the explicit-step case use a sign-correct step column or separate ascending/descending inputs, keeping its arguments as leaves so it still tests the native path? This conclusion is source-derived; I have not run the suite.

[0lai0]

Thanks @sunchao, both P2 items addressed in the latest push.

1. Native gate is now Literal | Attribute | BoundReference only (argsAreLiteralsOrRefs). Nested calls, CASE WHEN, and zero-arg UDFs fall back to the dispatcher.

2. Fixture uses sign-correct stp so sequence(a, b, stp) is legal on both rows.

Added sequence with zero-arg UDF stop routes through the dispatcher (comet_seq_stopper()).
Null-short-circuit / CASE cases stay in sequence.sql.

Local validation:
CometCodegenSuite: 178/178 pass (4 sequence tests, incl. zero-arg UDF → dispatcher)
CometSqlFileTestSuite: expressions/array/sequence.sql + sequence_ansi.sql pass
cargo test -p datafusion-comet-spark-expr array_funcs::sequence: 5/5 pass
test-compile green on -Pspark-3.4, -Pspark-3.5, -Pspark-4.0 (shim SequenceBatchTooLarge)

[sunchao]

Re-reviewed 645fedd5572f32bcb71e3ac44196f2ba406af17b. Both previous P2s are addressed: the explicit literal/reference whitelist keeps zero-argument UDFs and other computed arguments inside whole-expression dispatch, and the sign-correct step column fixes the native-path fixture. The added test checks the whole Sequence's dispatcher routing. No new actionable P1/P2 findings in this increment.

This was a focused source re-review; I did not rerun the Scala/Rust suites, generated code or benchmarks. The current workflows require action, so the author's reported local runs are not independent current-head CI validation.

[sunchao]

Re-reviewed a6191e420369ba1e60b8617183a8a88240ec56a3. Both earlier P2 fixes remain intact after the merge. One additional P2 is detailed inline: the committed integral benchmark queries still select whole-expression dispatch, so they do not validate native sequence performance.

This is a source-derived finding. I did not rerun the Scala/Rust suites, generated code or benchmarks. The current workflows require action and do not validate this head.

[sunchao]

[P2] Make the integral benchmark exercise native sequence

Could you materialize the integral endpoints as columns in the prepared Parquet table, then verify that these cases use spark_sequence before timing them? Every integral query here passes an arithmetic expression such as c_start + 4 or c_null_start + 364. argsAreLiteralsOrRefs rejects those arguments, so the complete Sequence goes through the JVM dispatcher, or falls back to Spark if dispatch is unavailable. runExpressionBenchmark only checks Comet operators and does not catch expression dispatch. These queries therefore cannot measure this native kernel's Spark-versus-Comet benefit. Could you refresh the comparison with leaf arguments and retain the date case as a dispatcher control?