andygrove opened a new pull request, #6729: URL: https://github.com/apache/datafusion-comet/pull/6729
## Which issue does this PR close? Closes #6720. ## Rationale for this change `RowArrowReader` and `SparkColumnarArrowReader` call `ArrowWriter.create` for every Arrow batch, and it called `allocateNew` on every vector, releasing the previous batch's buffers and allocating new ones. For a fixed-width vector, Arrow's `allocateNew` also zeroes the validity and data buffers, which it slices out of one allocation rounded up to a power of two. A column of 8192 longs zeroes 128 KiB to hold 64 KiB of values, and the writers then overwrite every value. After #6566 this was about half of what the fixed-width columnar path cost. The fresh buffers were there because a batch exported to native keeps its buffers until native releases it. The export retains each buffer it hands to native and releases it from the batch's release callback, so the buffers' reference counts already show whether native still holds a batch. ## What changes are included in this PR? All of it is in `ArrowWriter.create`, which every Spark-to-Arrow conversion goes through: - A vector keeps its buffers for the next batch when nothing but the vector and its children references them. Arrow slices a vector's validity and data (or offset) buffers out of one allocation, with a reference for each slice, so the check compares each allocation's reference count with the number of the vector's own buffers on it. Validity and offsets are zeroed, as a new allocation would be. Values are not. A top-level fixed-width vector must also hold the batch size, because the row path writes those values without checking capacity. - When something still holds the buffers, as a native sort or shuffle writer does with the batches it buffers, the vector gets new buffers and leaves the old ones to the holder. A new fixed-width data buffer is not zeroed either, booleans aside: it is allocated directly and handed to the vector with `loadFieldBuffers`. This also covers `CometArrowConverters`, which the cache serializer uses. It gives every batch a root of its own, so reuse never applies there. - The writers now write a zero for each null row they write one value at a time: the row path, the per-value path that struct, array and map fields use, and the value-by-value columnar copies. Every batch has the same bytes as on `main`, including under nulls, so shuffle and cache compression and the determinism of the output do not change. A bulk columnar copy carries whatever Spark holds under a null, as it already did. - Reference counts are kept per allocator, so a JVM consumer that loads a batch into vectors of another allocator has to close the batch afterwards, as `ColumnarBatchArrowReader` does. Every consumer today either goes through the reader's own allocator (the C stream export) or closes or clears the batch (`ColumnarBatchArrowReader`, broadcast serialization, the cache serializer). The ownership section of `ffi.md` and the `ArrowWriter.create` doc say this. `CometArrowWriterBenchmark` gets a case that holds each batch until the next one is written, which measures the new-buffer path. Nanoseconds per row over 8192 rows, from two runs of `main` and one of this PR: | Case | `main` | this PR | | --- | --- | --- | | int, long and double columns, on heap | 2.3 | 0.9 | | the same, off heap | 1.8 | 0.4 | | the same, on heap, each batch held until the next is written | 2.4 | 1.2 | | the same, half the rows null | 3.3 to 3.5 | 2.0 | | decimal(18,2) column | 1.8 | 0.8 | | decimal(38,10) column | 6.9 | 5.9 | | struct<int,long,double,date> column | 5.6 | 3.8 | | string column | 2.9 | 2.9 | | int, long and double rows | 19.7 to 20.1 | 18.2 | | the same, half the rows null | 22.0 to 23.4 | 21.7 | | decimal(18,2) rows | 8.8 | 7.9 | | array<int> rows | 19.2 | 19.8 to 20.0 | Strings and the row path move little: Arrow never zeroed string data, and a row costs far more than zeroing its slots did. Rows of `array<int>` came out about 0.7 ns slower in every run of this PR. One run without the zero under null elements gave 18.7, but a focused rerun of that case moved by 6 ns between runs of the same code, so I could not resolve it further. End to end, this is TPC-H Q1 over 6M Parquet rows that Spark's vectorized reader reads and `CometSparkColumnarToColumnar` converts (`spark.comet.convert.parquet.enabled`, native scan off). Each number is the median of 7 runs after 2 warmups, over two rounds each of `main` and this PR, interleaved: | | `main` | this PR | | --- | --- | --- | | `conversionTime` | 109 to 119 ms | 80 to 95 ms | | wall time | 217 to 270 ms | 207 to 256 ms | Conversion time drops by about a quarter, and wall time moves within noise. Most of the remaining conversion time is the two dictionary-encoded string columns. With temporary counters, every batch after a reader's first one reused its buffers in this query: 46,116 vectors reused and 252 allocated. A sort over the same data, whose native shuffle writer buffers its input, allocated new buffers for about half of its batches, and its 6M sorted rows matched Spark's. Setup: Apple M3 Max, JDK 17.0.10, Spark 4.1.3 (the default profile), release native build, with other work running on the machine (load average 4 to 10). ## How are these changes tested? New tests in `CometArrowWriterSuite` cover every type the writer supports, written both as columns and as unsafe rows: - Batches written into memory that starts out holding `0xA5`, each into a root of its own and all three into one root, have the same bytes as batches written into zeroed memory. That covers validity and offsets for every row and the values those rows cover, through every child. `FilledMemoryAllocators` provides the memory. It is a test helper in Arrow's package, because the configuration that takes an allocation manager is not public. - A batch loaded into another root keeps its buffers while the next batch is written, and they are reused once that root lets go. - A C stream consumer that holds every batch keeps their values, and one that releases each batch before pulling the next, as native does, gets the same buffers back every time. - A batch larger than the previous one gets room for its rows before the row path writes them. These mutations fail the new tests: no zero under a null on the row path or in a value-by-value copy; reuse without checking reference counts, which also fails the existing string array, string map and C Data ownership tests; no zeroing of a list's validity on reuse; and reuse of a top-level vector too small for the batch. `fixed-width column write grows an undersized vector` sized its input from how Arrow rounds a combined allocation. It now sizes it from the vector's actual capacity. `CometArrowWriterSuite`, `CometStringWriterSuite`, `CometArrowStreamSuite`, `CachedBatchRowIteratorSuite`, the SparkToColumnar and LocalTableScan tests in `CometExecSuite`, the three in-memory cache suites, `CometTypedDatasetSuite`, `CometNativeShuffleSuite` and `CometRangeExecSuite` pass on Spark 4.1. The Arrow suites, the SparkToColumnar tests and `CometInMemoryCacheSuite` pass on Spark 3.5 (Scala 2.12). This PR has the `run-benchmark-check` label for the benchmark change. Spark's SQL suites would not exercise it: the only conversion on by default is OneRowRelation's, and that row has no columns. -- 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]
