andygrove opened a new pull request, #5543: URL: https://github.com/apache/datafusion-comet/pull/5543
## Which issue does this PR close? Part of #5487. Ticks these boxes: - **Drop the schema message from each cached stream** — subsumed by the buffer-selection change below, which removes per-column framing entirely. - **Evaluate projection by buffer selection instead of per-column streams** - **Prune on collated string columns** - **Documentation** — the committed-benchmark-results half is not done; see the note at the end. Not addressed here: typed readers for the row path (#5485 has no established cause yet, so this would be optimising ahead of a diagnosis) and background prefetch. ## Rationale for this change #5051 stores each cached column as its own compressed Arrow IPC stream, so a scan decodes only the columns it projects. That works, but it pays an Arrow schema block and compression framing per column per batch, and gives up cross-column compression: footprint grows 2.5% at 6 columns and 32% at 60. Spark's `ArrowCachedBatchSerializer` (SPARK-57268) reaches the same projection-proportional decode with no per-column framing at all. It keeps one RecordBatch per cached batch and parses the IPC message flatbuffer, which lists every buffer's offset and length within the body, to copy out only the byte ranges belonging to the selected columns. It depends on Arrow's native per-buffer compression rather than wrapping the whole payload in a Spark `CompressionCodec`. That approach dominates the current design on footprint while keeping the projection win, so this PR adopts it. Separately, `tracksBounds` matches `case StringType`, which a collated `StringType` does not equal. Collated columns therefore get null bounds and `buildFilter` declines to push predicates on them. That is correct but loses pruning that Spark manages. ## What changes are included in this PR? **Cached batch format.** `CometCachedBatch.columns: Array[ChunkedByteBuffer]` becomes `bytes: Array[Byte]`: one encapsulated Arrow IPC record batch message and its body, with no Schema message and no end-of-stream marker. The reader rebuilds the schema from the cached relation's attributes via `Utils.toArrowSchema`, so a wide relation no longer repeats the same schema bytes once per cached batch. New `CachedBatchIpc` owns the format — the field-node, buffer-span and variadic-count arithmetic, and the projected read that copies only the selected columns' buffers into a single off-heap allocation. **Compression** moves from a whole-payload Spark codec to Arrow's per-buffer IPC compression, which is what lets a projected read decompress only what it selected. New `spark.comet.exec.inMemoryCache.compression.codec` (`zstd` default, `none` available) and `...compression.zstd.level`. Arrow's lz4 is deliberately not offered. It is commons-compress's pure-Java implementation, unrelated to the JNI-accelerated lz4-java behind `spark.io.compression.codec`. Over a 200k-row six-column relation: | Codec | Materialize | Footprint | Read 1 of 6 | Read 6 of 6 | |-------|------------:|----------:|------------:|------------:| | `zstd` | 347 ms | 2 MiB | 52 ms | 63 ms | | `none` | 1743 ms | 13 MiB | 74 ms | 79 ms | | ~~`lz4`~~ | 205373 ms | 6 MiB | 51 ms | 107 ms | lz4 is dominated on both speed and size, so nothing prefers it. zstd also beats storing batches uncompressed on both axes, because the bytes it saves cost more to copy and store than compressing them costs. Reads still accept any codec a batch records. **Dictionary-encoded columns are decoded before being stored.** A payload with no schema message has nowhere to record either the index type or the dictionary. Comet's native scans do produce such columns, so this is a real path rather than a defensive one. **Decompression is done in `CachedBatchIpc` rather than left to `VectorLoader`.** arrow-java 18.3.0 leaks on the failure path: `VectorLoader.loadBuffers` collects a field's decompressed buffers into a local list and releases them only after the whole field has loaded, so a buffer that fails to decompress strands every buffer of that field decompressed before it. A string column reaches this — its offsets buffer decompresses, then its data buffer throws — so a single corrupt cached batch leaks off-heap for the life of the executor. **Collated string pruning.** New `CometTypeShim.compareStrings` compares under the column's collation (`UTF8String.semanticCompare` on Spark 4.x, byte order on 3.x, where collations do not exist). `tracksBounds` widens to any `StringType`, so bounds are recorded with the same ordering the partition filter Spark generates over that column uses. **Dependency.** Adds `org.apache.arrow:arrow-compression`. Already covered by the existing `org.apache.arrow:*` shade include, so it relocates with the rest of Arrow; its `commons-compress` and `zstd-jni` are excluded and come from Spark, which ships both on every supported version. ## How are these changes tested? `CometInMemoryCacheSuite` keeps its existing coverage, with the format-dependent tests rewritten against the new layout: - The projection tests no longer corrupt per-column streams. They scramble the compressed bytes of the columns a read must not touch, leaving every other byte of the payload identical, and assert the read still succeeds — which it only can if those buffers were never copied out of the payload. Each asserts as a precondition that the column it corrupts is genuinely stored compressed, since Arrow stores a buffer verbatim when compressing would not shrink it. - A new test asserts the payload begins with its record batch rather than a schema message, so a regression to a self-describing stream shows up directly rather than only as a footprint number. - Per-column sizes in the statistics row are checked against the message's own buffer layout. - A new test round-trips every codec the config accepts through a full read, a projected read, a row-count-only read and stats pruning. `none` takes a different path on read and was broken until this test was written. - The collation test is inverted: it now asserts pruning happens, and adds a UTF8_LCASE case that returns no rows if bounds are recorded with byte-order comparison. The no-bounds case moves to `BinaryType`. - Two leak tests cover a failure at a column's first buffer and a failure part way through a column, the latter reaching the `VectorLoader` behaviour above. Both assert the decode error surfaces as itself rather than as an Arrow reference-count error, which is what catches a cleanup path releasing the shared body twice. - The dictionary tests assert values read back correctly, which is what proves the writer decoded them — a row count would not. `CometCachedBatchHelper` re-derives the IPC buffer arithmetic independently rather than calling into `CachedBatchIpc`, so the assertions built on it cannot pass by inheriting a bug from the code under test. Also run: `CometInMemoryCacheKryoSuite`, `CometExecSuite`, `UtilsSuite`. Compiles clean against Spark 3.5, 4.0 and 4.1. The shaded jar was checked to confirm arrow-compression relocates and that commons-compress and zstd-jni are not bundled. `CometInMemoryCacheBenchmark` over a 5M-row six-column relation (Apple M3 Ultra, JDK 17, Spark 4.1, release build): | Query shape | Spark cache scan + convert | `CometInMemoryTableScan` | Relative | |-------------|---------------------------:|-------------------------:|---------:| | Repeated scan (3 of 6 columns) | 156 ms | 118 ms | 1.3x | | Selective filter | 44 ms | 39 ms | 1.1x | | Row count only (0 of 6) | 32 ms | 28 ms | 1.1x | | Narrow projection (1 of 6) | 50 ms | 39 ms | 1.3x | | Full projection (6 of 6) | 316 ms | 135 ms | 2.3x | As with #5051, both columns read the same Comet-written `CometCachedBatch` and Comet execution is on in both, so this measures keeping the cached scan native against falling back to a Spark cache scan and converting — not Comet against Spark execution, and not a comparison with Spark's own cache format. ## Notes for reviewers - **The committed benchmark results half of the documentation item is not done.** `spark/benchmarks` is in `.gitignore`, so Comet does not currently commit results files the way Spark does. Doing it properly needs that directory un-ignored plus a workflow to regenerate them, or the numbers rot — worth its own decision rather than being slipped in here. The measured tables live in the new docs page instead. - Cache materialization retains about 800 bytes per partition, independent of row count. This reproduces identically on 492dd6f2a, so it predates this PR and is untouched here; happy to file it separately. -- 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]
