feat: add native spark_sequence kernel for integral element types - #5614
Conversation
andygrove
left a comment
There was a problem hiding this comment.
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.
| // Second pass: write elements straight into the child buffer and push offsets. The | ||
| // batch-total check above guarantees `values.len() <= i32::MAX` at every iteration, so | ||
| // the offset push cannot overflow. | ||
| let mut values: Vec<T::Native> = Vec::with_capacity(total); |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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().
| s"Illegal sequence boundaries: ${params("start")} to ${params("stop")} " + | ||
| s"by ${params("step")}")) | ||
|
|
||
| case "Internal" => |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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!
| -- Error paths: step direction contradicts bounds, or zero step with start != stop | ||
| -- ============================================================================ | ||
|
|
||
| query expect_error(Illegal sequence boundaries: 1 to 5 by -1) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Thanks, Added all three to sequence.sql: full byte/short range, Int32 step overflow, and CASE WHEN guarded branch.
sunchao
left a comment
There was a problem hiding this comment.
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.
| // `start <= stop ? 1 : -1`, which cannot be expressed as a plan-time literal. | ||
| val argProtos = Seq(startExprProto, stopExprProto) ++ | ||
| expr.stepOpt.map(exprToProto(_, inputs, binding)) | ||
| scalarFunctionExprToProtoWithReturnType( |
There was a problem hiding this comment.
[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?
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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.
| // Integral sequence with column-reference/literal args lowers to the native spark_sequence | ||
| // kernel; no codegen-dispatch marker should appear. | ||
| withSequenceTable { | ||
| val df = sql("SELECT sequence(a, b), sequence(a, b, 2) FROM t") |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
Thanks @sunchao, both P2 items addressed in the latest push.
-
Native gate is now Literal | Attribute | BoundReference only (argsAreLiteralsOrRefs). Nested calls, CASE WHEN, and zero-arg UDFs fall back to the dispatcher.
-
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
left a comment
There was a problem hiding this comment.
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.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…quence # Conflicts: # native/spark-expr/Cargo.toml
sunchao
left a comment
There was a problem hiding this comment.
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.
| ("seq_short_5_elems", "SELECT sequence(c_start, c_start + 4) FROM parquetV1Table"), | ||
| ("seq_spine_365_elems", "SELECT sequence(c_start, c_start + 364) FROM parquetV1Table"), |
There was a problem hiding this comment.
[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?
There was a problem hiding this comment.
Thanks for the feedback, I pushed the fix.
Stop endpoints are now materialized columns (c_stop_5, c_stop_365, c_stop_10000, c_null_stop_365), so every integral query passes only literals and column references and satisfies CometSequence.argsAreLiteralsOrRefs (spark/src/main/scala/org/apache/comet/serde/arrays.scala:957). The date case keeps the arithmetic form intentionally, because temporal Sequence unconditionally routes through the JVM dispatcher (arrays.scala:948-952); it stays in the list as the dispatcher control.
Benchmark on Apple M5 / OpenJDK 17.0.18, 8192 rows:
| case | Spark ns/row | Comet ns/row | Comet vs Spark |
|---|---|---|---|
| seq_short_5_elems | 1795 | 645 | 2.8× |
| seq_spine_365_elems | 1796 | 798 | 2.3× |
| seq_long_10000_elems | 7922 | 7135 | 1.1× |
| seq_descending_default_step | 1746 | 736 | 2.4× |
| seq_explicit_step_7 | 1421 | 489 | 2.9× |
| seq_sparse_nulls_365_elems | 1651 | 663 | 2.5× |
| seq_date_spine_dispatcher (control) | 1528 | 1554 | 1.0× |
The integral vs date gap (2–3× vs ~1.0×) is the path evidence: leaf-arg queries now hit native; the date control stays on the dispatcher. The old c_start + 4 shape would have been ~1.0× across the board because argsAreLiteralsOrRefs rejected it.
Dropped a plan-text spark_sequence check, that name only appears in the native proto, so it always reported native=false. Operator-level Comet coverage is already asserted by findFirstNonCometOperator; the timings cover the expression path.
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed e774be49 against 55ae4f20. The materialized endpoint columns address the benchmark's dispatch issue. By source inspection, all six integral cases now satisfy the native Sequence argument gate. The earlier null-guard and fixture fixes remain intact, and I found no remaining actionable P1/P2.
This was a source-only follow-up. I did not rerun tests or benchmarks, so the reported timings remain author-provided. The exact-head workflows still require action and do not validate this head.
…quence # Conflicts: # native/spark-expr/Cargo.toml # native/spark-expr/src/comet_scalar_funcs.rs
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Re-reviewed 281b5223 against 81d637b9. This merges the previously approved revision with the new base. Twelve authored files are byte-for-byte unchanged, including the sequence kernel, literal/reference argument gate, error shims, regression fixtures and benchmarks. The two conflict resolutions preserve the sequence import, registration and benchmark alongside the upstream additions.
The earlier null-short-circuit and sign-correct fixture fixes remain intact. I rechecked integral bounds, default and zero-step behavior, overflow/error contracts and temporal fallback against the maintained Spark 3.5 and 4.0 sources. The documentation change is preserved as well. I found no remaining or new verified P1/P2 in this increment.
Validation and CI
This was a source-only follow-up. Exact merge-parent and file comparisons, parsed Cargo manifest checks, registration invariants, preserved base changes and git diff --check passed. I did not rerun compilation, runtime tests, generated code, fuzzing or benchmarks.
At the complete 17:03 UTC refresh, GitHub returned no head checks. The current-head CI workflow, Delta Contrib Build Gate and CodeQL were queued. This head had not been validated by a completed CI run. I did not approve or rerun workflows.
Performance
The six integral benchmark queries retain materialized endpoint columns and satisfy the native argument gate by source inspection. The merge changes no sequence element loop, allocation strategy or output-size accounting. Timing claims remain author-provided, and the existing batch-memory limitation remains documented. No additional performance finding emerged from this update.
Design
The merge preserves the integral native path and whole-expression dispatcher fallback for temporal or computed arguments. Both sequence and upstream registrations survive the conflict resolution. The existing split of responsibilities remains intact, with no new design issue identified.
Abstraction & complexity
The typed kernel, explicit return type and shared registration pattern retain their existing roles. The update integrates those pieces with the new base rather than adding a new dispatch layer or helper hierarchy. No additional abstraction change is needed for this merge.
andygrove
left a comment
There was a problem hiding this comment.
I checked this out locally and ran it against a release build this time: the Rust unit tests, the sequence.sql fixture on both Spark 4.1 and 3.5, a 4000-row randomized fuzz biased toward the Long and Int boundaries, and both benchmarks.
The correctness side holds up well. I diffed sequenceLength against the 3.5.8 and 4.1.3 sources and the port is exact, including the Long.MinValue / -1 special case and the ordering of the size check ahead of internalError in the BigInt fallback. More usefully, I forced all five error paths through the native kernel using column arguments and compared against Spark on both 3.5 and 4.1. Exception class and message match byte for byte on every one, so the shim mappings are right, including _LEGACY_ERROR_TEMP_2161 on 3.5 and COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER carrying the function name on 4.1. The fuzz found no divergences across 3628 comparable rows, and fourteen downstream consumers of the produced list all match Spark.
The performance story is where I want to push back, in both directions.
First, seq_long_10000_elems is not parity. On this machine (M3 Max, release, local[1], 8192 rows) I measure 0.6X there, and extending the sweep it keeps getting worse:
| elems/row | Spark | native | native vs Spark |
|---|---|---|---|
| 5 | 59 ms | 41 ms | 1.44X |
| 365 | 42 ms | 34 ms | 1.24X |
| 1000 | 40 ms | 36 ms | 1.11X |
| 10000 | 51 ms | 72 ms | 0.71X |
| 50000 | 120 ms | 231 ms | 0.52X |
Spark's average time beats Comet's best time at 10000 elements, so this is not measurement noise. My seq_date_spine_dispatcher control reads 0.7X against your 1.0X, which suggests this machine sits roughly 1.4x in Spark's favour relative to yours, but even allowing for all of that the long shapes do not reach parity.
I do not think this is a defect in your loop. It reads like an inherent consequence of the representation. Spark allocates one long[] per row, which at 10000 elements is 80 KB and stays resident in L2 while it is written and immediately consumed, whereas a per-batch buffer has to stream to DRAM. Nothing in the element loop can recover that. What I would like is for the crossover to be stated rather than implied away. Could the audit entry in array_funcs.md and the description both say something like "faster than Spark below roughly a thousand elements per row, slower above it"? As the table stands a reader concludes the native path is never worse than Spark.
Second, and this is the part I think you are underselling. The comparison you have not shown is the one that matters most for anyone actually running Comet, because sequence goes through the JVM codegen dispatcher today. I added a third arm in the same session over the same data, with the routing asserted from ExtendedExplainInfo on each arm so I knew which path I was timing:
| elems/row | native vs dispatcher |
|---|---|
| 365 | 6.5X |
| 1000 | 14.4X |
| 10000 | 64X |
| 50000 | 101X |
That is a far better argument for this change than 2X to 3X against Spark, and it is the number a Comet user actually experiences. Would you consider adding a dispatcher arm to CometSequenceBenchmark and leading with it?
One practical note on benchmarking. My first run showed Comet 23x slower and it took me a while to spot that Maven had left the debug libcomet in spark/target/classes. Worth building with -Prelease and checking the library size before trusting any number from that harness.
Everything else is inline. Nothing I found is a correctness problem.
| // own `long[]`, so the user may hit this on a query Spark itself would run. Report it via | ||
| // a dedicated error that names `spark.comet.batchSize` as the actionable knob rather than | ||
| // Spark's per-array size limit. | ||
| if total > i32::MAX as usize { |
There was a problem hiding this comment.
The try_reserve_exact change does what I asked for, thank you. I think there is still a gap underneath it though.
This guard is on element count rather than bytes, so for bigint it only fires at around 17 GB, and the Vec comes from the global allocator rather than the DataFusion MemoryPool. That means the allocation is not counted against spark.comet.memory*, cannot be spilled, and applies no back-pressure.
At the default batch size I measured 50000 elements per row allocating 3.3 GB in a single reservation and completing fine, with peak process RSS 2852 MB above baseline against Spark's 1076 MB for the same query. On a Linux executor with overcommit the OOM killer arrives well before try_reserve_exact gets a chance to return Err, so the graceful path is the one a user is least likely to reach.
Could the ceiling be a byte budget as well as an element count, sized off the batch memory budget rather than i32::MAX? That would make SequenceBatchTooLarge fire while the executor is still healthy, which is the point at which its actionable message is worth something. This is also the peak-memory question from the earlier round, which I do not think has been answered with a measurement yet.
| -- Error paths: length exceeds MAX_ROUNDED_ARRAY_LENGTH | ||
| -- ============================================================================ | ||
|
|
||
| query expect_error(the array size limit 2147483632) |
There was a problem hiding this comment.
SequenceBatchTooLarge is the one behaviour in this PR that fails a query Spark completes, and I cannot find a test for it anywhere. It is also close to free to test, because the total > i32::MAX check runs in the first pass before anything is allocated.
Could a case go in around here? Something like this trips it at the default batch size and passes at half of it, so it also pins the remedy you documented:
statement
CREATE TABLE t_seq_ceiling(a INT, b INT) USING parquet
query expect_error(Lower `spark.comet.batchSize`)
SELECT sum(CAST(size(sequence(a, b)) AS BIGINT)) FROM t_seq_ceilingwith a = 0, b = 262143 over 8192 rows in a single partition. I ran that against this branch and it produces exactly the message you intended, and lowering spark.comet.batchSize to 4096 makes it return Spark's answer instead. That is the sequence(0, 262143) case from my first pass, now confirmed end to end.
| -- without evaluating the inner argument, so the inner `sequence(1, 5, -1)` | ||
| -- must not fire on the NULL row. Non-leaf argument shapes stay on the JVM | ||
| -- codegen dispatcher for this reason | ||
| -- (https://github.com/apache/datafusion-comet/pull/5614#discussion_r3910237757). |
There was a problem hiding this comment.
Could the #discussion_r3910237757 link come out of this comment? The two sentences before it already say why non-leaf argument shapes stay on the dispatcher, which is the part a future reader needs. A pointer into a review thread records how the code came to be rather than what it does, and it will not survive the next change to this reasoning. The #5349 link at the top of the file is the durable kind and is worth keeping.
| s"Illegal sequence boundaries: ${params("start")} to ${params("stop")} " + | ||
| s"by ${params("step")}")) | ||
|
|
||
| case "SequenceBatchTooLarge" => |
There was a problem hiding this comment.
This message now exists in four places, error.rs plus all three shims, and unlike SequenceIllegalBoundaries it has no version-specific behaviour to justify living in the shims at all. Would a shared constant work, with each shim interpolating params("totalElements") into it? The same question applies to the case "Internal" arm just below, which is character-identical in all three files. Four copies of one English sentence is the kind of thing where a later wording fix lands in three places and misses the fourth.
| | `flatten` | ✅ | Native | Binary/struct/map elements fall back | | ||
| | `get` | ✅ | — | | | ||
| | `sequence` | ✅ | Codegen dispatch | | | ||
| | `sequence` | ✅ | Hybrid | Integral types run natively; date/timestamp sequences use codegen dispatch | |
There was a problem hiding this comment.
This note describes the native versus dispatcher split, which a user cannot observe, and leaves out the per-batch ceiling, which is the one thing they can, since it is a query Spark runs that Comet fails. Could it mention that as well, something like "very large per-row sequences may exceed Comet's per-batch limit, lower spark.comet.batchSize"? The audit entry covers it well, but that is not where somebody who has just hit the error will be looking.
| Unsupported(Some(s"sequence with element type $other is not supported natively")) | ||
| } | ||
|
|
||
| private def argsAreLiteralsOrRefs(expr: Sequence): Boolean = { |
There was a problem hiding this comment.
The gate itself is the right conservative call and I am not arguing with it. It is worth being explicit about how narrow the resulting fast path is though. The native kernel only engages when both endpoints already exist as columns or literals, so the idiomatic spine sequence(x, x + n) stays on the dispatcher, and so does anything behind a coercion CAST. Your own benchmark is the evidence, since it needed c_stop_5 and friends materialised as stored columns before any integral case went native.
I also tried the workaround a user would reach for first, and it does not work. FROM (SELECT c_start, c_start + 364 AS c_stop FROM p) gets folded straight back by CollapseProject, and the explain still reports JVM codegen dispatcher: sequence. So there is no way to opt in short of rewriting the table.
Two things would help. Could the audit entry spell out which shapes reach the native path, since "leaf arguments only" is not something a user can map onto their own SQL? And separately, is a safe widening worth considering later, accepting an argument subtree that provably cannot throw and preserves nulls, which would cover x + n at least under non-ANSI? Happy for that to be a follow-up.
|
This merged while I was still working through the review above, so none of it was blocking. I have moved the actionable items into #5712 so they do not get lost: the perf crossover above roughly a thousand elements per row and the missing dispatcher benchmark arm, the unbounded per-batch allocation outside the memory pool, the missing To be clear about what I did verify, since it is the reassuring part: all five of Spark's |
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 twolong[]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?
spark_sequencekernel (native/spark-expr/src/array_funcs/sequence.rs) forByte/Short/Int/Long. Matches SparkSequence.sequenceLength, including the overflow-report paths (2^63,2^63+1, and the internal-error edge).CometSequencewithCodegenDispatchFallback:CASE WHEN, nestedsequence(...)) returnUnsupportedand stay on the JVM codegen dispatcher, preserving Spark's per-row null short-circuit.Unsupportedand stay on the dispatcher.SequenceBatchTooLargeerror when the batch's total generated elements exceed the Arrow i32 offset ceiling (i32::MAX) ortry_reserve_exactfails. The message namesspark.comet.batchSizeas the actionable knob. Spark itself has no equivalent limit because it stores each row as its ownlong[].ShimSparkErrorConverter: Spark 3.x throwsIllegalArgumentException("Illegal sequence boundaries: ..."); Spark 4.x throwsSparkIllegalArgumentException("_LEGACY_ERROR_TEMP_3243").CollectionSizeLimitExceedednow carries a decimalStringcount (can exceedi64) and afunction_namefor 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 ascountand rendered(array,N).case "Internal"inShimSparkErrorConvertermaps ~20 existingSparkError::Internalproducers (intemporal.rs,numeric.rs,conversion_funcs/string.rs,rlike.rs, etc.) fromSparkException(message, <text>)to[INTERNAL_ERROR] <text>.sequencemarkedHybridinexpressions.md; audit notes underarray_funcs.mddocument the per-batch ceiling and the non-leaf argument fallback.Limitations
Native integral
sequencematerializes 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 withSequenceBatchTooLarge; loweringspark.comet.batchSizeis the fix.How are these changes tested?
sequence.rs.spark/src/test/resources/sql-tests/expressions/array/sequence.sql: integral types, default/explicit step, nulls, explode, seven error cases, plus:sequence(-128Y, 127Y),sequence(-32768S, 32767S)),sequence(-2147483648, 2147483647, 1073741824)),sequence(s, size(sequence(1, 5, k)))),CometCodegenSuite: leaf-arg integral sequences run natively; non-leaf integral and temporal sequences show "JVM codegen dispatcher".(message,<text>)shape forSparkError::Internal.Criterion (
cargo bench --bench sequence, N=8192). Absolute numbers; the kernel is not onmain. Benchmarks use leaf-arg shapes only.short_2_elemsshort_5_elemslong_365_elemslong_10000_elemsdescending_365_elemszero_step_start_eq_stopsparse_nulls_365_elemsdense_nulls_365_elemserror_illegal_boundariesSpark (
CometSequenceBenchmark, 8192 rows, Apple M5, Spark 4.1.3 / Scala 2.13)seq_short_5_elemsseq_spine_365_elemsseq_long_10000_elemsseq_descending_default_stepseq_explicit_step_7seq_sparse_nulls_365_elemsseq_date_spine_dispatcherseq_date_spine_dispatcheris a control (date path unchanged).seq_long_10000_elemsis memory-bandwidth-bound at 82M elements/batch. Typical speedup is 2X–3X on the shorter-list shapes the issue targets.