[nat-openai]

Which issue does this PR close?

Partially addresses #5712: large integral-sequence performance and unbounded output allocations. The issue’s remaining follow-ups stay open.

Rationale for this change

Integral sequence currently limits the number of generated elements, but has no practical byte-based allocation limit. A single batch can allocate several gigabytes, and concurrent tasks or retained outputs can multiply that memory use. The native implementation also loses to Spark on the large shapes reported in #5712.

This change bounds live sequence-buffer allocations and reduces sizing and generation overhead.

Sequence buffers can outlive the producing native plan through Arrow clones, slices and JVM exports. Integrating these allocations into the general memory pool would require broader changes to propagate shared reservations across those ownership boundaries, charge each backing allocation precisely once, and release its reservation only after the final owner disappears. A dedicated executor-wide allowance provides a bounded solution without requiring that more invasive reservation-propagation work.

The 256 MiB default is configurable independently of Spark’s execution-memory budget. Requests that exceed it fail before allocation; applications can reduce the batch size or raise the allowance at startup. It bounds live sequence backing buffers, rather than total native memory or process RSS.

What changes are included in this PR?

• Add spark.comet.exec.sequence.maxBytesPerExecutor, read from executor startup SparkConf, with a default of 256 MiB. Concurrent tasks share the outstanding-byte counter; SQL session changes cannot enlarge the allowance.
• Validate and size the complete invocation before allocating. Account for values, offsets, optional validity and alignment padding, and retain each allocation’s charge through Arrow sharing and JVM exports.
• Reduce kernel overhead by borrowing scalar inputs, using constant-size sizing scratch, allocating the final output capacity once, inlining small input-access helpers, and copying previously generated rows for larger repeated scalar outputs.
• Generate and admit the full requested batch for scalar-only calls. Prevent constant folding into a list scalar whose later expansion could bypass admission.
• Check the captured Spark task context during sizing and generation so cancellation releases partially constructed output.
• Extend behavioural tests, benchmark coverage and documentation. Native eligibility remains limited to integral literal/column arguments; temporal sequences continue to use the dispatcher.

The allowance fails immediately when a request does not fit; it does not spill or wait for another task to release memory. Existing Spark row-error behaviour and the batch list-offset ceiling are preserved.

And LLM was used for for exploring potential solutions and implementing these commits.

How are these changes tested?

Native tests: 18 kernel/memory tests and the scalar-folding registration regression pass in release mode. Coverage includes all four integral widths, null/scalar inputs, arithmetic boundaries, error precedence, admission/refusal/recovery, concurrent reservations, clone/slice/FFI lifetimes, and cancellation/unwind cleanup.

JVM tests: nine targeted tests pass on Spark 4.1.3, covering executor configuration, allocation refusal and recovery, native/fallback routing, captured-task cancellation, and the existing sequence.sql fixture.

Build and checks: native release build, workspace all-target Clippy with warnings denied, Rust formatting, Spotless and whitespace checks pass.

Before merge: the broader Spark SQL suite still needs the run-spark-4.1-tests CI gate.

Performance

CometSequenceBenchmark completed on Spark 4.1.3, Scala 2.13.17 and Java 17 with local[1] and 8,192 rows. Integral endpoints were materialised Parquet columns. The benchmark explicitly raises the sequence allowance to 4 GiB to measure large outputs.

These are Spark Benchmark best/average timings, not medians:

Shape	Spark best / average (ms)	Comet best / average (ms)
2 elements	13 / 17	7 / 9
5 elements	10 / 13	5 / 7
365 elements	13 / 14	7 / 8
1,000 elements	17 / 18	9 / 11
10,000 elements	67 / 69	39 / 40
50,000 elements	270 / 272	167 / 169
Descending, 365 elements	12 / 13	6 / 8
Explicit step 7	10 / 11	5 / 6
Sparse nulls, 365 elements	12 / 13	6 / 8
Date dispatcher control	11 / 13	16 / 18

Comet is faster on all integral shapes in this run. These are fresh measurements from the current harness and Spark version; absolute timings are not directly comparable with the original issue’s report. Concurrent-task peak RSS was not measured.

The native performance gains compared to the main branch come from simplifying the sizing and generation loops: values are written directly into fixed-capacity buffers, scalar inputs are borrowed instead of expanded into temporary arrays, and small value, null and length helpers are inlined to avoid repeated calls and intermediate structures. Larger repeated scalar sequences also reuse the first generated row through bulk copying. The largest improvements occur with column-valued endpoints producing hundreds or thousands of elements per row, where generation costs dominate: 365- and 10,000-element sequences are approximately 5× faster. Ordinary two- and five-element column sequences inly improve by about 10%.

[ajsquared]

Reviewed the current changes and feedback; no independently confirmed unresolved P1 remains.

[andygrove]

The allowance makes queries fail that Spark and the 1.1 release both run. It counts live output across every task on the executor, and a buffer keeps its charge for as long as anything downstream holds it. The operators that hold batches never give that credit back. The native shuffle writer keeps its input in buffered_batches, SortExec keeps its input batches, and explode's contiguous-run path in unnest_list_array returns slices that alias the sequence values. All of them spill on pressure from their own pool, never because of this counter, so the next sequence batch fails instead.

I ran the same probes against main and this branch, in on-heap mode and in off-heap mode with a 2g pool. A single task producing 32 MiB batches, about 470 MiB in total, is fine when the arrays are consumed as they stream. The same output fails on its eighth batch once it goes through repartition(4, k), ORDER BY k, or LATERAL VIEW explode(sequence(a, b)) followed by a repartition. Five concurrent tasks with 64 MiB batches failed five runs out of five with nothing retained downstream. One 8192-row batch fails from 4,096 BIGINT elements per row, so the 10,000 and 50,000 element shapes from #5712 now fail by default. That's why both benchmarks raise the allowance, and why the existing ceiling test had to drop its recovery batch size from 256 to 128. Every one of these passes on main, and on this branch with the allowance raised to 16g.

What #5712 asked for was a ceiling sized from the batch memory budget that fires while the executor is still healthy. This one fires on healthy executors, and whether a query fails depends on how many tasks share the executor and on the plan downstream. Could the ceiling apply per invocation instead, so it depends only on the batch being built, with a default derived from the task's share of the pool in getMemoryConfig? Another option is a try_grow against the task's memory pool for the batch being built, released on return. A refusal would then mean the pool really has no room, and buffering would stay with the operators that already reserve for what they hold.

The kernel rewrite is worth landing on its own. In my criterion runs against main the 365 and 10,000 element shapes are about 5x faster. Would you be open to splitting it into its own PR and taking the ceiling design back to #5712?

[andygrove]

Rows where start == stop are 20 to 29 percent slower than main in interleaved criterion runs, whatever the step is. zero_step_start_eq_stop goes from 45.5 to 57.2 µs, and a one-element shape with the default step goes from 37.9 to 48.9 µs. Two- and five-element rows are faster. dense_nulls_2_elems, dense_nulls_5_elems and sparse_nulls_5_elems are also 6 to 8 percent slower, and all_scalar_8192_rows_365_elems about 5 percent. optimizing_expressions.md asks for no meaningful regression on any shape. Could you add a one-element shape to benches/sequence.rs, look at that case, and share the criterion comparison against main?

[andygrove]

This would be the only native code that polls TaskContext.isInterrupted. If cancellation latency matters for long native batches, it matters for sorts, joins and scans too, so it seems better solved once in the execution loop than inside one kernel. At the default allowance a single call takes tens of milliseconds at most, so the polling buys little here. It adds at least six JNI calls per batch plus one per 65,536 generated values, and the criterion numbers don't include them because the benchmark leaves check unset. Could we drop it from this PR and open an issue for a general mechanism if it turns out to be needed?

[andygrove]

This section adds a budget, but the Who allocates what table above still lists the operator pool as the only thing bounding native allocations. Could it get a row for the sequence allowance, in whatever form it ends up?

[nat-openai]

Thanks for the review @andygrove! The guidance around the batch memory ceiling is particularly helpful... I don't think I appreciated exactly what was being proposed in #5712. I'll iterate on this a little and split the kernel bits into a separate PR.

[sunchao]

Reviewed all 17 changed files across the eight commits from 00a4b422f0ed6560ef76dce004b94c3697613108 to cb2d92100009291950164186b524f5f7eb31efc2. The PR is not a draft. Existing reviews, conversation comments, inline comments, and unresolved threads were read before forming conclusions.

No additional introduced P1/P2 issues found within this review. The already-reported allocation-failure regression remains substantiated and unresolved, so this is not a clean verdict. No duplicate findings are submitted.

Summary

• Prior state and problem: Integral sequence previously bounded element counts but could allocate gigabytes outside the operator pool. Issue #5712 also documented poor performance for long sequences.
• Design approach: The PR adds a startup-configured, executor-wide allowance of 256 MiB for live sequence buffers. It combines preflight sizing with fixed-capacity generation and reservations that follow buffer ownership.
• Correctness: Compared Spark sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Null propagation, default direction, equal endpoints, illegal boundaries, length limits, and overflow behavior remain consistent. Focused tests and differential comparisons found no additional result or error divergence. Allocation refusals remain the existing blocker.
• Compatibility analysis: The latest branch-1.1, at e9efd9f764ee0a59b7898ff028d6985d4a7a28e1, has the same sequence kernel as this PR’s base and the same routing as its head. Integral literal/column eligibility and temporal dispatch are unchanged. The new allowance deliberately changes resource behavior, but its rejection of otherwise successful workloads is the unresolved regression discussed in the existing review.
• Key design decisions: Startup SparkConf overrides session settings. Charges survive clones, child slices, and FFI exports. Volatility::Volatile prevents DataFusion from folding native calls into list scalars whose expansion would bypass admission. Cancellation uses the captured task context.
• Implementation sketch: CometExecIterator serializes the allowance, JNI creates the shared-counter pool, and the planner injects it into SparkSequence. Typed kernels validate rows, calculate aligned capacities, reserve atomically, and fill values, offsets, and optional validity buffers. Buffer owners release their individual charges on final destruction.
• Performance: Borrowed scalars, constant-size sizing scratch, inlined access, and repeated-row copying remove identifiable work. The PR reports roughly fivefold gains for long column sequences. The existing review reports 20–29% regressions for equal endpoints and smaller regressions for other shapes. Those remain unresolved. Native Criterion measurements exclude the production JNI cancellation callback, and I did not independently rerun performance benchmarks.
• Design: The two-pass kernel is straightforward. The independent allowance cannot ask downstream consumers to spill, so admission can fail while the operator pool still has capacity. This is the central existing design concern, supported by both the ownership paths and the bounded reproduction.
• Abstraction & complexity: Reservation and SequenceBuffer coherently implement allocation-lifetime accounting, including partial construction and FFI ownership. They also introduce custom allocation and unsafe buffer construction to support the disputed lifetime-wide policy. The existing proposal to separate kernel optimization from admission policy would make those choices easier to evaluate.
• Behavioral changes worth calling out: Scalar-only native calls now generate the requested batch directly. Requests fail immediately without waiting or spilling, and session configuration cannot enlarge the allowance. An 8,192-row invocation producing 4,096 BIGINT elements per row requires 268,468,352 bytes and fails against the default 268,435,456-byte limit before allocation.
• Suggested improvements: Address the existing request to reconsider executor-wide lifetime admission and separate the kernel improvements, resolve the measured short-shape regressions, and obtain the outstanding Spark SQL verdict. No additional suggestion met the P1/P2 reporting bar.

Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-memory-pr, and review-comet-ffi-pr, with repository AGENTS.md and relevant contributor guidance.

Exact-head CI: The labeling check passed. Comet CI and CodeQL both concluded action_required. There is no successful build/test verdict for this head, and run-spark-4.1-tests has not run.

Validation: Rebuilt and passed 18 native sequence/kernel-memory tests plus the scalar-folding registration regression. A freshly linked bounded harness passed 80 base/head comparisons covering sliced nullable arrays, scalar/column combinations, and explicit/default steps. With a 64 KiB allowance, the base retained eight small outputs while the head refused the eighth after retaining seven. Releasing those outputs restored admission, independently confirming the existing retention concern.

Validation limits: I did not run the JVM suites, native core planner test, full Spark SQL matrix, release benchmarks, or concurrent-task RSS measurements. Native FFI ownership tests do not substitute for a live JVM export test. Project source contents remain unchanged, and no GitHub state was modified.