andygrove opened a new pull request, #6988: URL: https://github.com/apache/datafusion-comet/pull/6988
## Which issue does this PR close? Closes #6950. ## Rationale for this change The columnar cache write stores each input batch as one cached batch, and nothing bounds how small those batches are. A shuffle read returns at least one batch per map output block, so a relation cached straight from a shuffle over many map tasks holds thousands of small batches. Each carries its own message metadata, compression framing and statistics row, and every read of the cache pays for the batch count. #6950 has the details. ## What changes are included in this PR? - A new `CacheBatchCoalescer` sits in front of the columnar write in `ArrowCachedBatchSerializer.convertColumnarBatchToCachedBatch`. It passes a batch through when two things hold: the batch is Arrow-backed and laid out the way the reader reads it, which is the test the write already used to store a batch as it stands, and it holds at least half of `spark.sql.inMemoryColumnarStorage.batchSize` rows. Full native batches of 8192 rows therefore still reach the writer without a copy. - Every other batch is copied with `ArrowWriter.writeColumns` into batches of `batchSize` rows, the size the row path cuts at, and is split where it crosses that size. A small batch is worth the copy. A batch of Spark's own column vectors, or one laid out differently from the reader, was already being copied into Arrow one batch at a time, so combining it adds no copy. The copy reads each column through the `ColumnVector` interface, so it handles dictionary-encoded and nested columns, and offsets that don't start at zero, as a native producer's sliced arrays can have. Arrow's `VectorAppender` assumes offsets start at zero, which is why this doesn't use it. - Row order is kept, so a sorted cache stays sorted: the rows copied so far are emitted before a batch that passes through. Each batch is copied before the input advances, because a native producer such as `CometExecIterator` releases the batch it returned last when it is asked for the next. The coalescer owns the batches it builds and releases each one when the next is requested, and anything still open when the task completes. - The row path is unchanged, since it already cuts at `batchSize`. - The in-memory cache guide's "Storage format" section describes the batch sizes. ## How are these changes tested? Three new tests in `CometInMemoryCacheSuite`: - "combines small columnar batches in order" feeds the serializer a mix of Arrow and Spark-vector batches at a batch size of 1000. The input releases each batch when it is asked for the next, as a native producer does. The test checks the exact cached batch sizes (600, 9000, 700, 1000, 1000, 1000 and 100) and every row in order. - "combines the small batches of a shuffle read" caches a 20,000-row relation with nested columns straight from a shuffle over 20 map tasks into 2 partitions. Without this change it caches 40 batches of about 500 rows. With it there are at most 4, and the rows match Spark's. - "keeps a sorted cache sorted when it combines its batches" caches a sort whose native batches hold 100 rows and checks that each partition reads back in order. Without this change it caches 201 batches of 100 rows. Each of five mutants fails at least one of these tests: removing the coalescer, passing through at the full batch size instead of half of it, not splitting at the batch size, emitting a pass-through batch ahead of the rows combined before it, and passing Spark-vector batches through. `CometInMemoryCacheSuite`, `CometInMemoryCachePruningSuite`, the three Kryo suites and `CachedBatchRowIteratorSuite` pass locally on Spark 3.4, 3.5, 4.0, 4.1 and 4.2. Scalafix and the strict-warnings build pass on 3.5. I measured #6950's shape with a probe that isn't part of this PR. It reads 5M rows of four columns (bigint, int, double and string) from 200 Parquet files as 200 splits, and caches them with `repartition(50, $"a")`. Both caches were built and read in the same JVM, interleaved. The first was built at a batch size of 2, so every one of these batches passes through, which is the layout `main` produces. The figures are medians of 3 builds and 9 reads, on an Apple M3 Max with JDK 17, Spark 4.1 and a debug native library: | | `main`'s layout | This PR | | --- | ---: | ---: | | Cached batches | 10000, of 250 to 775 rows | 510, of 5000 to 10000 rows | | Payload | 36.1 MiB | 32.8 MiB | | Storage memory | 44.0 MiB | 33.2 MiB | | Build | 3961 ms | 3815 ms | | Aggregate through `CometInMemoryTableScan` | 369 ms | 289 ms | | `noop` write of every row | 141 ms | 113 ms | | Aggregate by Spark operators | 315 ms | 299 ms | On a copy of the same rows in random order, which compresses less well, storage memory went from 77.3 to 66.0 MiB, the build from 4.4 to 3.9 seconds, and the aggregate from 394 to 322 ms. Separate runs of `main` and this PR gave the same layouts. From a single map task, the same relation caches the same 633 full batches with or without this PR. -- 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]
