perf: project cached batches by buffer selection, prune on collated strings

[andygrove]

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	363 ms	2 MiB	56 ms	62 ms
none	1776 ms	13 MiB	78 ms	81 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. tracksBounds widens to any StringType, and bounds are compared with TypeUtils.getInterpretedOrdering(dataType) — Spark's own interpreted ordering for the type, which on Spark 4 resolves a StringType through CollationFactory.fetchCollation(collationId).comparator. So bounds are recorded with the same ordering the partition filter Spark generates over that column uses, the collation awareness comes from Spark at every supported version with no shim, and the ordering is resolved once per column instead of re-dispatching on the DataType per row. This also replaces the hand-rolled per-type compare.

Reading builds one CachedBatchIpc.Projection per partition, holding the projected schema and the node/buffer/variadic index layout. That arithmetic walks every field of the cached relation, so computing it per batch would make the bookkeeping O(total columns) against O(selected columns) of useful work — worst in exactly the wide-relation, narrow-projection case this format exists for.

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.

Nested columns get their own projection coverage, over a six-column relation whose middle four are a struct, an array, a map and a struct wrapping an array. They were previously round-tripped only under SELECT *, which cannot exercise the span arithmetic: a flat column always owns one field node and two or three buffers, a nested one owns a run as long as its subtree, and a full projection covers the whole buffer sequence however it is partitioned. So one test gives each column a turn as the sole projection with the other five corrupted, and another compares values against the uncached query across single-column, paired and out-of-order projections — a row count comes from the record batch header, so it cannot catch a window that is misaligned but still decompresses. The per-column statistics test runs over the nested relation too. Both new tests fail if fieldNodeCount stops recursing into children.

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 (Apple M3 Max, JDK 17, Spark 4.1, release build), over a 5M-row relation of six flat columns:

Query shape	Spark cache scan + convert	CometInMemoryTableScan	Relative
Repeated scan (3 of 6 columns)	201 ms	167 ms	1.2x
Selective filter	69 ms	61 ms	1.1x
Row count only (0 of 6)	45 ms	47 ms	1.0x
Narrow projection (1 of 6)	70 ms	57 ms	1.2x
Full projection (6 of 6)	556 ms	290 ms	1.9x

And over a 1M-row relation of six columns whose middle three are structs, one nested two levels deep:

Query shape	Spark cache scan + convert	CometInMemoryTableScan	Relative
Row count only (0 of 6)	39 ms	35 ms	1.1x
Narrow projection (1 of 6)	109 ms	61 ms	1.8x
Full projection (6 of 6)	282 ms	126 ms	2.2x

verifyPlan now asserts the projection width each case claims, which is what corrected the flat table above: the previous full projection (6 of 6 columns) row was reading three columns, because count() over a non-nullable column is rewritten to count(1) by NullPropagation and the column is then pruned out of the scan. Only k, s1 and s2 were nullable, and only incidentally, because Remainder can divide by zero. Every column of both relations is nullable now so count(c) genuinely reads c. A real six-column read is 1.9x, not the 2.2x previously published for a three-column one.

Arrays and maps are deliberately not in the benchmark. The baseline arm needs Spark's cache scan to bridge into Comet operators, and CometSparkToColumnarExec declines ArrayType and MapType, so for a query projecting one of those that arm does not exist and the two cases stop measuring the same boundary. They are covered in the suite instead.

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 feature is still disabled by default. spark.comet.exec.inMemoryCache.enabled stays false, as it has been since feat: add experimental native support for in-memory cache, disabled by default #5051, so none of this reaches a default configuration. Everything below the config, including the change of cached format, only affects users who have opted in to the experimental native cache. The config is static: its value at startup is what decides whether Comet's cache serializer is installed at all.
• 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 492dd6f, so it predates this PR and is untouched here; happy to file it separately.

[comphead]

Thanks @andygrove should we also test/benchmark nested data?

[andygrove]

Good call — this turned out to be worth doing, because the nested coverage that existed was in the one shape that can't fail.

Nested columns were already round-tripped, but only under SELECT *. That can't exercise the part of this format that is actually nontrivial for them: a flat column always owns one field node and two or three buffers, whereas a nested one owns a run as long as its whole subtree, and a full projection covers the entire buffer sequence however it happens to be partitioned. So the span arithmetic was only ever tested where getting it wrong doesn't show.

There are now two tests over a six-column relation whose middle four are a struct, an array, a map and a struct wrapping an array. The first gives each column a turn as the sole projection with the other five corrupted, so a run computed short or long by a buffer gets caught reaching into a corrupted neighbour. The second compares values against the uncached query across single-column, paired and out-of-order projections — row counts come from the record batch header, so they can't catch a window that is misaligned but still decompresses, and an out-of-order projection is the case a full one genuinely cannot stand in for. The per-column statistics test now runs over the nested relation too, since a nested column's recorded size is the sum of its subtree. Both new tests fail if fieldNodeCount stops recursing into children, which is how I checked they aren't passing for free.

The format itself needed no changes, so this is coverage rather than a fix.

On the benchmark, I added the three projection widths over a relation of struct columns:

Query shape	Spark cache scan + convert	CometInMemoryTableScan	Relative
Row count only (0 of 6)	39 ms	35 ms	1.1x
Narrow projection (1 of 6)	109 ms	61 ms	1.8x
Full projection (6 of 6)	282 ms	126 ms	2.2x

The gap is wider than the flat relation's at every width, which is what you'd expect — the conversion the left column pays scales with values per row, not with columns.

I left arrays and maps out of the benchmark on purpose, and said so in the docs. The left column needs Spark's cache scan to bridge into Comet operators, and CometSparkToColumnarExec declines ArrayType and MapType outright, so for a query projecting one of those that arm doesn't exist — the partial aggregate stays on Spark and the two cases stop measuring the same boundary. Tests are the right place for those, and they're covered there.

One thing your question shook out that I should flag: while adding the nested cases I made verifyPlan assert the projection width each case claims, and the existing full projection (6 of 6 columns) case was reading three. count() over a non-nullable column gets rewritten to count(1) by NullPropagation, which then prunes the column out of the scan, and only k, s1 and s2 were nullable — incidentally at that, because Remainder can divide by zero. Every column of both relations is nullable now so count(c) really reads c, and I've regenerated the numbers in the docs. The genuine 6-of-6 read is 1.9x rather than the 2.2x that was published for what was really a 3-column read.