andygrove opened a new pull request, #6730: URL: https://github.com/apache/datafusion-comet/pull/6730
## Which issue does this PR close? Closes #6721. ## Rationale for this change `ArrowWriter`'s row path serves `CometSparkRowToColumnar`, `CometLocalTableScanExec` and the Comet cache serializer's row input, and #6607 and #5634 add more callers. Nearly all of those rows are `UnsafeRow`s, but the writer read them through Spark's generic getters. A string went through a `UTF8String` and Arrow's per-value setter. Each array, map and struct allocated a view, and every element took the generic `write`, which makes two or three virtual or interface calls per value. #6721 has the per-type comparison with the native row converter. Per column type, in ns per row with one column of 8192 rows and every tenth row and element null, in the shapes `CometArrowWriterBenchmark` uses. A harness outside the repo ran `main`'s writer and this one in alternating JVMs, three times each. It warms every type up before measuring, so the JIT sees a mix of types as a real workload does. Each number is the median of the three runs' best. | Row input | `main` | This PR | |---|---|---| | int | 4.6 | 4.3 | | long, double, date, boolean | 4.5-4.8 | 4.3-4.7 | | decimal(18,2) | 11.1 | 5.4 | | decimal(38,10) | 15.2 | 9.4 | | string, 11 bytes | 14.1 | 8.4 | | string, 48 bytes | 14.6 | 9.8 | | binary | 15.9 | 10.8 | | struct<int,long,double,date> | 31.1 | 23.0 | | struct<int,string> | 26.7 | 17.6 | | array<int> | 22.8 | 13.2 | | array<boolean> | 19.6 | 13.7 | | array<decimal(18,2)> | 28.4 | 21.7 | | array<string> | 36.8 | 21.0 | | array<struct<int,string>> | 70.5 | 44.2 | | array<array<int>> | 74.4 | 32.7 | | map<string,string> | 82.1 | 40.5 | | map<int,long> | 54.4 | 23.6 | | 8 columns: int, long, double, string, date, decimal(18,2), boolean, timestamp | 52.7 | 37.7 | `CometArrowWriterBenchmark` itself warms one type at a time, so `main`'s virtual calls stay inlined there and structs and `array<int>` gain less. Best of two runs each: | Spark rows to Arrow by type | `main` | This PR | |---|---|---| | int | 4.9 | 3.6 | | decimal(18,2) | 7.4 | 5.2 | | decimal(38,10) | 13.4 | 8.2 | | string | 9.8 | 6.4 | | struct<int,long,double,date> | 21.7 | 19.7 | | array<int> | 11.0 | 9.9 | | array<string> | 35.1 | 18.0 | | map<string,string> | 66.3 | 34.8 | | int, long, double (the "Spark rows to Arrow" case) | 14.7 | 9.3 | End to end, aggregates over a table of 4M rows cached in Spark's default format, which `CometSparkToColumnarExec` converts from rows, with a release `libcomet` at `local[1]` on the Spark 4.1 profile. Every column has 10% nulls. Each number is the median of three alternating runs, each the best of 7. `conversionTime` also counts decoding the cached rows. | Columns read | Wall time (ms) | `conversionTime` (ms) | |---|---|---| | int, long, double, decimal(18,2), date | 233 → 208 | 186 → 162 | | string | 123 → 107 | 94 → 77 | | struct<int,string> | 223 → 144 | 190 → 113 | | array<string>, map<string,string> | 678 → 452 | 539 → 316 | | all nine | 1157 → 879 | 975 → 706 | All of this is on an M3 Max with JDK 17, on a machine shared with other jobs, so take the numbers as indicative. ## What changes are included in this PR? `ArrowWriter.write` hands an `UnsafeRow` to a new `writeUnsafeRowField`, and unsafe arrays go to a new `writeArrayElements`. Other rows take the generic path as before. Both read Spark's memory directly: - **Fixed-width values** are read from their slot with `UnsafeRow`'s own getters and written at the vector's width, without the virtual call to `setValueUnsafe`. Arrays of these types are copied in one block, and their validity is the array's null words inverted, written at any bit offset. Booleans are converted from a byte to a bit per element. - **Decimals** up to 18 digits are written from the unscaled long, which `UnsafeRow.getDecimal` reads without checking it against the precision either. Wider ones are range-checked in place instead of being copied to a `byte[]` first. Arrays of decimals up to 18 digits still go through `getDecimal`, because `UnsafeArrayData.getDecimal` does check the precision. - **Strings and binaries** are copied using the offset and size in their slot, without a `UTF8String` or `setSafe`. An array's elements are all checked first, then the buffers are grown once and the bytes copied. Values up to 64 bytes are copied a word at a time, which beats `Unsafe.copyMemory`'s checks and call on short strings. - **Arrays, maps and structs** each point one reused `UnsafeArrayData`, `UnsafeMapData` or `UnsafeRow` at every value instead of allocating one per row, then write their elements, keys and values, or fields through these same paths. A map's entries are marked valid in one call. A generic row holding an unsafe value, as a copied `ColumnarBatchRow` holds unsafe primitive arrays, takes these paths too. `ArrowFieldWriter.writeUnsafe` is gone. Its fixed-width override skipped the capacity check and relied on `ArrowWriter.create` allocating top-level vectors at the batch size. Arrow 18 keeps a fixed-width vector's capacity in a field, so the new path checks it per value for nothing measurable. Struct fields can then use the same path as top-level fields, which addresses the issue's point about structs without allocating nested vectors at the batch size. `FixedWidthArrowFieldWriter` now takes its vector as a constructor parameter, so the per-value code reads it from a field instead of through each subclass's accessor. The two failure modes that `CometStringWriterSuite` pins still hold on the new path, and more strictly than on Arrow's setters. An offset overflow throws Arrow's `OversizedAllocationException` and a negative length throws an `IllegalArgumentException`, before anything is written or allocated, for a single value and for every element of an array. I also tried binding fixed-width fields statically from the row and struct loops. That gained another 12-16% on structs and wide rows in the harness, but it was a wash elsewhere, so I left it out. This overlaps #6729, which also changes `ArrowWriters.scala`. Git merges the two apart from the tests both append to `CometArrowWriterSuite`, and with that resolved the merged tree passes `CometArrowWriterSuite`, `CometStringWriterSuite` and `CometArrowStreamSuite`. A null fixed-width value on the new path goes through `setNullUnsafe`, so the zero that #6729 writes under each null applies to it. ## How are these changes tested? - New `CometArrowWriterSuite` tests write unsafe rows, both on heap and copied to off-heap memory, and check that every vector matches what the generic row path writes for the same rows. They cover the suite's primitive types (18 on Spark 4.1), the 11 nested types already there and 8 more (collections inside collections, element types converted one at a time, and a struct of 70 fields), plus a row with a column of every type three times over. Null fractions are 0%, 20% and 100%. A third writer alternates between unsafe and generic rows, so each path appends where the other left off. A fourth writes generic rows that hold unsafe values. - Collections of up to 150 elements cover null words past the first and every bit offset in the Arrow validity buffer. Strings and binaries of every length from 0 to 140 bytes cover both sides of the word-copy threshold. - The vector comparison now also checks that variable-width offsets never decrease, null rows included, since an unfilled hole would leave one smaller than the last. - New `CometStringWriterSuite` tests check the overflow and negative-length cases for an unsafe row and for an unsafe array, including that nothing is allocated. - Each of twelve deliberate bugs in the new code makes these tests fail: inverted validity, offset holes not filled, a missing sign extension, a wrong value width, an unshifted bit mask, map entries not marked valid, a short-copy tail, either check dropped, and a stale reused view in each nested writer. - `CometArrowWriterSuite`, `CometStringWriterSuite` and `CometArrowStreamSuite` pass on Spark 4.1 and on Spark 3.4 / Scala 2.12. On 4.1, the `SparkToColumnar` and `LocalTableScan` tests in `CometExecSuite`, `CometInMemoryCacheSuite`, `CometInMemoryCacheKryoSuite` and `CometInMemoryCachePruningSuite` pass too. - The strict Scala warnings build (`./mvnw test-compile -Pspark-3.5 -Pstrict-warnings`) and the semantic scalafix check on the 3.5 profile both pass. -- 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]
