test: cover Spark-to-Arrow batch conversion edge cases

[peterxcli]

Which issue does this PR close?

Closes #5317.

Rationale for this change

Spark-to-Arrow conversion must preserve logical batch sizes even when there are no columns, apply slice offsets to nested and nullable values, and release output allocations if encoding fails. Add direct regression coverage for these contracts in the existing Arrow suite.

What changes are included in this PR?

Add three tests to CometArrowStreamSuite:

• Verify that row and columnar readers split a five-row, zero-column input into batches of 2, 2, and 1 rows. Check independently owned row conversion and full/empty columnar conversion too.
• Verify nullable array elements, nullable strings, and constant columns across nonzero slice offsets. Check that reader roots remain stable, separately allocated output batches remain valid, and closing converted output leaves source data usable.
• Inject failures during row and column encoding and verify that the original exception propagates and Arrow allocator memory returns to zero.

The final diff changes one test file. Production code, the parameterless ArrowWriter.finish() API, and the benchmark implementation remain unchanged.

Existing conversion paths covered by the tests

BATCHING AND OWNERSHIP: three independent entry points

+------------------------------+  +------------------------------+  +------------------------------+
| RowArrowReader               |  | SparkColumnarArrowReader     |  | CometArrowConverters         |
+------------------------------+  +------------------------------+  +------------------------------+
| Consumes InternalRows        |  | Consumes ColumnarBatches     |  | Converts rows or batches     |
| Chooses row batch size       |  | Chooses source slices        |  | Fresh root per output batch  |
| Stable root, fresh buffers   |  | Stable root, fresh buffers   |  | Cleans up on failure         |
|                              |  |                              |  |                              |
| write(row)                   |  | writeColumns(batch, off, n)  |  | write(row), writeColumns(..) |
+------------------------------+  +------------------------------+  +------------------------------+
               |                                 |                                 |
               +---------------------------------+---------------------------------+
                                                 |
                                                 v
                  +--------------------------------------------------------------+
                  | ArrowWriter: shared batch encoding                           |
                  +--------------------------------------------------------------+
                  | All three entry points call this writer directly.            |
                  |                                                              |
                  | write(row)       -> field.writeUnsafe(row, columnIndex)      |
                  | writeColumns(..) -> field.writeColumnSlice(vector, off, n)   |
                  |                                                              |
                  | Maintains the logical batch count, including zero columns.   |
                  | finish() -> root.setRowCount(count)                          |
                  |          -> each field.finish()                              |
                  +--------------------------------------------------------------+
                                                 |
                                                 | delegates to one writer per field
                                                 v
                  +--------------------------------------------------------------+
                  | ArrowFieldWriter: per-field value encoding                   |
                  +--------------------------------------------------------------+
                  | Writes typed values and nulls into its Arrow vector.         |
                  | Maintains its own value count.                               |
                  |                                                              |
                  | Row input     -> writeUnsafe(...)                            |
                  | Column slice  -> bulk copy when eligible                     |
                  |               -> ColumnarArray scalar fallback otherwise     |
                  |                                                              |
                  | Nested types delegate to child ArrowFieldWriters.            |
                  | finish() -> set vector value count; finish children          |
                  +--------------------------------------------------------------+

Readers and converters control batching and ownership. ArrowWriter maintains the logical batch count and passes it to the root in finish(). Field writers retain their own counts for nested values. Readers replace buffers behind a stable root, while converters allocate independently owned batches.

How are these changes tested?

./mvnw test -Pspark-3.5 -Dtest=none -Dsuites=org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite

CometArrowStreamSuite: 15 tests passed, 0 failed, including the three new tests, against the existing production implementation. Spark 3.5.9 / Scala 2.12, OpenJDK 21.0.6.

Scalastyle, Spotless, and git diff --check passed. Benchmarks are not applicable to this test-only diff.

[sunchao]

Correctness

Reviewed de1f04a6 against authoritative base 7190df63. I found no actionable P1/P2.

The readers and converters already determine each output batch's logical size. This change passes that size to ArrowWriter.finish(rowCount) and removes the writer's duplicate batch counter. Zero-column conversion already works on the merge base. The behavioral contract becomes explicit without changing which rows are encoded: row readers pass the number of successful writes, full-batch conversion passes batch.numRows(), and the columnar reader passes the selected slice length.

I compared the result with the maintained Spark 3.5 and 4.0 implementations. Spark's ColumnarBatch stores row count independently of its columns, so zero columns must still allow a positive count. Its row writer counts null-valued rows, and ColumnarArray applies the slice offset to value and null access. Those semantics remain intact here. Empty source batches are still skipped by the streaming columnar reader, while fresh conversion can return a zero-row batch. Every batch-writer caller was updated. The remaining parameterless finish() calls belong to field writers.

Per-field counts remain necessary for arrays, maps and structs, whose child lengths can differ from the root row count. Their encoding, null handling and finish/reset paths are unchanged, as are decimal overflow handling, supported types, timezone metadata and fallback dispatch. The readers still replace buffers behind a stable root, and fresh converters still produce separately owned batches and release allocations on encoding failure. The row-count change adds no selection, cast or ANSI-mode behavior. Maintained Spark 3.4 and 4.1 branches were unavailable, so those versions were not independently source-qualified.

Validation and CI

The head branches from 93c32346, three commits behind the authoritative base. Those upstream changes do not overlap the six PR files. I verified CI merge e1fa9af8 has the exact base/head parents, retains every changed file from the head, and adds exactly the authored patch to the base.

The Linux Spark 3.5 job and macOS Spark 4.0 job both checked out that merge and passed the new zero-column, nested/constant-vector slice, independent-lifetime and encoding-failure cases. The macOS log also shows all 15 CometArrowStreamSuite cases and the existing zero-column df.count() regression passing. At 2026-09-05T21:08:17Z, the snapshot contained 68 successful checks, eight skipped checks and one cancelled title check, with later successful title checks. No failed or pending check was reported. Local validation was source inspection and git diff --check, without another build or runtime rerun.

Performance

The change removes a duplicate increment per row and redundant batch-count assignments. It adds no scan, buffer copy, allocation or lock. Existing bulk-copy eligibility and scalar fallback stay unchanged, and callers reuse counts they already compute for batching.

The author reports two runs per version for 27 benchmark cases, with unchanged median case throughput and a range of -5.65% to +6.03%. The 8192-row on/off-heap bulk cases are reported at -1.03%/-0.50%. These are author-provided local measurements with timer rounding and run variation, not an independently reproduced speedup or proof that every case is regression-free. The source does not establish a new material cost requiring another performance finding.

Design

The responsibility split is appropriate: readers and converters choose batch boundaries and ownership, while the writer coordinates encoding. Passing the already-known logical count avoids coupling that boundary to the last encoding method used, especially for zero-column batches. Preserving child counters keeps nested-vector lengths independent from the batch size without introducing another state holder or allocation policy.

Abstraction & complexity

This is a small API change that removes state and updates the existing callers. It keeps the shared field-writer implementation, allocation paths and ownership mechanisms in place. The added tests exercise observable batch sizes, slice contents, independent lifetimes and failure cleanup rather than introducing a broader conversion framework.

[andygrove]

I'm not sure of the value of this change, tbh. This changes the API and introduces a new invariant that all the callers must respect where previously the behavior was encapsulated and guaranteed to be correct. What is the motivation? I'm assuming there is no performance benefit from removing some integer math?

[peterxcli]

I'm not sure of the value of this change, tbh. This changes the API and introduces a new invariant that all the callers must respect where previously the behavior was encapsulated and guaranteed to be correct. What is the motivation? I'm assuming there is no performance benefit from removing some integer math?

@andygrove make sense, I also think auto tracking the row count in arrow writer is safer.
reverted changes around the arrow writer, only keep the new test, title and description are rewritten, too. ptal!

[andygrove]

Thanks @peterxcli