peterxcli opened a new pull request, #5713:
URL: https://github.com/apache/datafusion-comet/pull/5713
## Which issue does this PR close?
Closes #5317.
## Rationale for this change
Spark-to-Arrow readers and converters already know the logical number of
rows in each output batch, but `ArrowWriter.finish()` keeps a second count
derived from row or column writes. Pass that count explicitly so finishing a
batch, including one with no columns, does not depend on which encoding method
was called.
Most of #5317 has already landed: #5442 introduced
`ArrowWriter.writeColumns` and `ArrowFieldWriter.writeColumnSlice`, and #5051
landed fresh-batch conversion using the shared writer API. This PR completes
the explicit row-count API and adds the remaining focused regression coverage.
Zero-column conversion already works on the base revision; this makes its
row-count contract explicit.
## What changes are included in this PR?
- Replace `ArrowWriter.finish()` with `finish(rowCount)` and remove its
duplicate batch counter.
- Pass the known row count from both readers and both converter paths;
update the existing benchmark and direct writer tests.
- Keep per-field counts, which are needed for nested values, and preserve
existing allocation, bulk-copy dispatch, fallback, and ownership behavior.
- Add coverage for zero-column batch sizes, nested and constant-vector
fallback across slices, independent output lifetimes, and cleanup after
row/column encoding failures.
### Resulting relationships
The diagram shows the complete resulting architecture; the shared encoding
APIs and bulk-copy path already exist on the base revision. Readers and
converters control batching and ownership, while writers control value encoding.
```text
+------------------------------+ +------------------------------+
+------------------------------+
| RowArrowReader | | SparkColumnarArrowReader | |
CometArrowConverters |
|------------------------------| |------------------------------|
|------------------------------|
| Consumes InternalRows | | Selects source batch slices | | Fresh
root per output batch |
| Controls batch size | | Controls slice size | | Caller
owns returned batch |
| Reuses one root | | Reuses one root | | Closes
root on failure |
| | | | |
|
| write(row) | | writeColumns(batch, | |
rowToArrowBatchIter |
| finish(rowCount) | | startRow, numRows) | | ->
write(row) |
| | | finish(rowCount) | |
columnarBatchToArrowBatch |
| | | | | ->
writeColumns(...) |
+------------------------------+ +------------------------------+
+------------------------------+
| |
|
+---------------------------------+---------------------------------+
|
v
+--------------------------------------------------------------------------------------------------+
| ArrowWriter
|
|--------------------------------------------------------------------------------------------------|
| Coordinates encoding for one VectorSchemaRoot.
|
|
|
| write(row)
|
| -> each field.writeUnsafe(row, columnIndex)
|
|
|
| writeColumns(batch, startRow, numRows)
|
| -> each field.writeColumnSlice(source ColumnVector, startRow, numRows)
|
|
|
| finish(rowCount)
|
| -> root.setRowCount(rowCount)
|
| -> each field.finish()
|
+--------------------------------------------------------------------------------------------------+
|
v
+--------------------------------------------------------------------------------------------------+
| ArrowFieldWriter (one per Arrow field)
|
|--------------------------------------------------------------------------------------------------|
| Encodes values into its Arrow vector and tracks its own value count.
|
|
|
| Row input: writeUnsafe -> null handling / typed value write
|
|
|
| Column input: writeColumnSlice(original ColumnVector, offset, length)
|
| |
|
| +-----------+-----------+
|
| | |
|
| v v
|
| Eligible fixed-width Scalar fallback
|
| no-null input ColumnarArray slice
|
| | |
|
| v v
|
| Bulk copy into Arrow writeCol / writeColNoNull
|
|
|
| Nested types -> child ArrowFieldWriters
|
| finish() -> set vector value count; finish children
|
+--------------------------------------------------------------------------------------------------+
```
Both readers reuse a stable root but allocate fresh buffers per batch to
preserve FFI safety. `CometArrowConverters` creates independently owned output
batches. Arrow-backed input continues through the separate
`ColumnarBatchArrowReader` retain/transfer path. Field writers keep their own
counts because nested arrays and maps can contain more values than the root's
logical row count.
## How are these changes tested?
- `./mvnw test-compile -Pspark-3.5 -DskipTests`: reactor build passed,
including Scalastyle and Spotless checks. Formatting was applied with `./mvnw
spotless:apply -Pspark-3.5`.
- `org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite`, run
directly with ScalaTest Runner against this worktree's compiled classes: **15
passed, 0 failed**.
- Coverage includes nonzero offsets and split batches; zero-column and empty
batches; nullable/no-null and nested inputs; dictionary and constant-vector
fallbacks; capacity growth; independent batch ownership; and allocator cleanup
when encoding throws.
- `git diff --check`: passed.
- JVM Arrow tests and benchmarks were run; no native execution benchmark was
run for this JVM-only change.
### Benchmark results
Existing
`org.apache.spark.sql.comet.execution.arrow.CometArrowWriterBenchmark`,
comparing base `93c323466feba127c46c1dfcd5fa4c8b3cc87195` with this change.
Apple M4, macOS 26.6.2, OpenJDK 21.0.6, Spark 3.5.9 / Scala 2.12. The workload
has three fixed-width columns (`Int`, `Long`, `Double`).
Ran all **27 cases twice per version**, in separate JVMs and in
baseline/changed/changed/baseline order, without concurrent compilation. The
baseline used saved unmodified main and test classes; the changed runs used the
worktree's rebuilt classes. Both used the same dependency classpath and JVM
options. The only benchmark source change is the explicit argument to `finish`.
The table reports Spark Benchmark's best-case throughput in **million
rows/sec**, retaining both samples. Delta compares the median of each pair;
positive means higher throughput. These are local desktop microbenchmarks
subject to timer rounding and run-to-run variation.
Median case throughput was unchanged. At 8192 rows, on-heap/off-heap bulk
throughput changed **-1.03%/-0.50%**. The largest median slowdown was the
32-row on-heap scalar reference case (**-5.65%**), which bypasses production
bulk dispatch. The 32-row nullable fallback was **-3.73%**. These measurements
do not establish a speedup or that every case is regression-free.
| Workload | Case | Baseline 1 | Baseline 2 | Changed 1 | Changed 2 | Median
throughput delta |
|---|---|---:|---:|---:|---:|---:|
| Spark columnar to Arrow (1 rows) | on-heap optimized path | 1.4 | 1.4 |
1.4 | 1.4 | +0.00% |
| Spark columnar to Arrow (1 rows) | on-heap scalar copy | 1.4 | 1.4 | 1.4 |
1.4 | +0.00% |
| Spark columnar to Arrow (1 rows) | off-heap optimized path | 1.5 | 1.4 |
1.4 | 1.4 | -3.45% |
| Spark columnar to Arrow (1 rows) | off-heap scalar copy | 1.4 | 1.4 | 1.4
| 1.4 | +0.00% |
| Spark columnar to Arrow (1 rows) | nullable fallback | 1.4 | 1.4 | 1.4 |
1.4 | +0.00% |
| Spark columnar to Arrow (8 rows) | on-heap optimized path | 10.1 | 10.1 |
10.1 | 10.1 | +0.00% |
| Spark columnar to Arrow (8 rows) | on-heap scalar copy | 10.1 | 10.1 |
10.1 | 10.1 | +0.00% |
| Spark columnar to Arrow (8 rows) | off-heap optimized path | 10.1 | 10.1 |
10.1 | 10.1 | +0.00% |
| Spark columnar to Arrow (8 rows) | off-heap scalar copy | 10.1 | 10.1 |
10.1 | 10.1 | +0.00% |
| Spark columnar to Arrow (8 rows) | nullable fallback | 10.1 | 10.1 | 10.1
| 10.1 | +0.00% |
| Spark columnar to Arrow (32 rows) | on-heap optimized path | 34.9 | 36.6 |
36.6 | 36.6 | +2.38% |
| Spark columnar to Arrow (32 rows) | on-heap scalar copy | 30.7 | 29.5 |
28.4 | 28.4 | -5.65% |
| Spark columnar to Arrow (32 rows) | off-heap optimized path | 36.6 | 36.6
| 36.6 | 34.9 | -2.32% |
| Spark columnar to Arrow (32 rows) | off-heap scalar copy | 29.5 | 28.4 |
29.5 | 28.4 | +0.00% |
| Spark columnar to Arrow (32 rows) | nullable fallback | 29.5 | 29.5 | 28.4
| 28.4 | -3.73% |
| Spark columnar to Arrow (512 rows) | on-heap optimized path | 341.3 |
341.3 | 341.3 | 341.3 | +0.00% |
| Spark columnar to Arrow (512 rows) | on-heap scalar copy | 73.1 | 75.9 |
75.4 | 75.4 | +1.21% |
| Spark columnar to Arrow (512 rows) | off-heap optimized path | 341.3 |
332.3 | 332.0 | 341.3 | -0.04% |
| Spark columnar to Arrow (512 rows) | off-heap scalar copy | 76.8 | 73.1 |
72.7 | 72.3 | -3.27% |
| Spark columnar to Arrow (512 rows) | nullable fallback | 74.5 | 74.5 |
74.0 | 72.7 | -1.54% |
| Spark columnar to Arrow (8192 rows) | on-heap optimized path | 664.2 |
636.3 | 642.5 | 644.6 | -1.03% |
| Spark columnar to Arrow (8192 rows) | on-heap scalar copy | 80.0 | 84.3 |
83.9 | 79.6 | -0.49% |
| Spark columnar to Arrow (8192 rows) | off-heap optimized path | 646.8 |
664.2 | 662.0 | 642.5 | -0.50% |
| Spark columnar to Arrow (8192 rows) | off-heap scalar copy | 83.8 | 80.1 |
79.9 | 78.4 | -3.42% |
| Spark columnar to Arrow (8192 rows) | nullable fallback | 83.9 | 83.7 |
83.2 | 80.2 | -2.51% |
| Spark rows to Arrow | fixed-width, no nulls | 62.3 | 57.1 | 66.2 | 60.4 |
+6.03% |
| Spark rows to Arrow | fixed-width, 50% nulls | 62.5 | 62.8 | 63.5 | 62.7 |
+0.72% |
Median case throughput delta: +0.00%. Case range: -5.65% to +6.03%.
The existing target for a single-version benchmark run is:
```sh
make
benchmark-org.apache.spark.sql.comet.execution.arrow.CometArrowWriterBenchmark
PROFILES=-Pspark-3.5
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
