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.

[comphead]

Reviewed for Spark semantics, the Arrow/JVM boundary, memory ownership, and scope.

The format change reads as sound. I checked the reference counting in load and decompressed against arrow-java 18.3.0: the 6-arg ArrowRecordBatch constructor is (.., variadicBufferCounts, alignBuffers) with retainBuffers defaulting to true, extractUncompressedBuffer returns a non-retained slice without closing its input, and ZstdCompressionCodec.doDecompress closes its allocation on both error paths. The retain-before-decompress ordering is correct for all four cases (compressed, stored-verbatim, zero-length, no codec). The collated-string change also matches Spark: TypeUtils.getInterpretedOrdering resolves through CollationFactory on 4.x, which is the same ordering the generated partition filter uses, so bounds and predicate agree with no shim.

One finding worth acting on before merge. The rest are cleanups.

Major. The read path derives every buffer window from Utils.toArrowSchema(cacheAttributes), but nothing checks the writer produced that layout, and isArrowBacked admits at least one vector whose buffer count differs. Details inline on CachedBatchIpc.scala. The fix is a length require in load using prefix sums selectedRange already computes.

Scope. Appropriately scoped. The getInterpretedOrdering switch replaces a hand-rolled per-type comparison rather than adding one, so it is not a drive-by.

Tests. The corruption-based projection tests are the right assertion and do distinguish the new behaviour from the old. Two gaps: no coverage of a BinaryType column backed by a fixed-size vector (see the Major), and the new tests cannot move to sql-tests/ because the serializer comes from a suite-level SparkConf and spark.sql.cache.serializer is static.

Performance. Benchmark evidence is present, and the verifyPlan projection-width assertion is a real improvement over the previously published numbers.

Two nits not worth their own thread. Projection(arrowFields: Seq[Field], ...) is indexed positionally and both call sites already pass .toIndexedSeq, so IndexedSeq[Field] would make that a compile-time guarantee. hydrateDictionaries reads oddly beside a scaladoc and a PR description that both say "decode".

[comphead]

Major: the read path trusts an unchecked layout invariant.

Projection derives every node and buffer window from Utils.toArrowSchema(cacheAttributes), but nothing verifies the writer produced that layout. encodeBatches takes the Utils.isArrowBacked(batch) fast path without converting, and Utils.isSupportedFieldVector admits FixedSizeBinaryVector, which Utils.fromArrowType maps to Spark BinaryType while Utils.toArrowType(BinaryType) is ArrowType.Binary. TypeLayout.getTypeBufferCount is 2 for the former and 3 for the latter.

So a cached BinaryType column backed by a fixed-size vector (mapInArrow output, an Iceberg fixed[N] read) is stored one buffer short. Every buffer index from that column onward shifts, and batch.buffers(j) is an unchecked flatbuffer accessor, so the result is wrong values or an ArrayIndexOutOfBoundsException inside setBytes rather than an error naming the cause. The previous per-column streams were self-describing, so they could not drift this way.

selectedRange already computes the full prefix sum. Keeping starts.last for each of the three sequences and checking it in load is O(1) per batch:

require(
  batch.nodesLength() == totalNodes && batch.buffersLength() == totalBuffers,
  s"cached batch layout does not match the cached schema: ...")

Dropping FixedSizeBinaryVector from the isArrowBacked fast path, the way LargeVarCharVector already is, would close the known case too.

[andygrove]

You're right, and the FixedSizeBinaryVector case is real: toArrowType(BinaryType) is Binary at three buffers against that vector's two, so every buffer index from such a column on shifts.

The length check is in, built the way you suggested. selectedRange now returns starts.last alongside the indices, and load compares both totals against nodesLength() and buffersLength() before it touches batch.buffers(j). I mutated the check away to confirm it earns its place: the new test for it then passes with no exception at all, which is exactly the silent-wrong-answer failure you described.

I went wider than dropping FixedSizeBinaryVector from isArrowBacked, because that only closes the top-level case. isArrowBacked answers for the top-level vector and never looks at children, so a struct whose child is a LargeVarCharVector passes it today and is stored with 64-bit offsets and read back with 32-bit. The length check can't catch that one either, since LargeUtf8 and Utf8 are both three buffers. So the write path now asks the direct question: do this batch's vectors already carry the Arrow types the reader will rebuild, recursively? That's CachedBatchIpc.matchesReaderLayout, and a batch that disagrees takes the conversion path it was already taking for non-Arrow input rather than being written unreadable. Names, nullability and a timestamp's timezone are excluded from the comparison — the last because a Comet scan labels with the session zone where the reader rebuilds UTC, and that's a label rather than a layout, so comparing it would send every timestamp column down the conversion path for nothing.

One wrinkle: Arrow-Java puts the index type on a dictionary-encoded vector's own field, so the field to compare is the dictionary's, resolved through the same Utils.lookupDictionary the writer uses.

On coverage, a unit test builds a FixedSizeBinaryVector-backed batch and asserts isArrowBacked accepts it while the write path declines it, with a VarBinaryVector beside it as the control so the predicate can't pass by refusing everything. The reader's check gets its own test that hands the reader one more attribute than the writer stored.

[comphead]

fieldVariadicCount has no effect today, and the stated reason for keeping it ("the span arithmetic above has to stay correct if that changes") does not hold. fieldBufferCount goes through TypeLayout.getTypeBufferCount, which returns 2 for Utf8View and BinaryView and ignores the variadic data buffers entirely, so if view vectors ever reach here the buffer spans misalign whether or not the variadic counts are tracked.

Suggest dropping this and the variadicCounts plumbing in load, and letting the length check above fail loudly instead.

[andygrove]

Agreed, dropped — both fieldVariadicCount and the variadicCounts plumbing in load. Your reasoning holds: getTypeBufferCount is 2 for Utf8View and ignores the data buffers entirely, so the spans misalign whether or not the counts are tracked. With the length check in place a view vector now fails there loudly instead, which is the right outcome given nothing on the write path can produce one.

[comphead]

fieldBufferCount(fields(i)) re-walks the column's subtree even though scanLeft already accumulated it. starts(i) until starts(i + 1) is the same range. This runs per batch on the write path, and the .map(...).sum allocates an intermediate per column, so a while over buffers between the two bounds does the whole thing in one pass.

[andygrove]

Done — a single while over buffers between starts(i) and starts(i + 1), so no re-walk and no intermediate collection per column.

[comphead]

readProjected does not exist. The method is Projection.load, so the [[readProjected]] link will not resolve. Same name appears in the alignBuffers comment at line 131 and in the CometCachedBatch scaladoc in ArrowCachedBatchSerializer.scala.

[andygrove]

Fixed, all three: the class scaladoc and the alignBuffers comment now say Projection.load, and so does the CometCachedBatch scaladoc in ArrowCachedBatchSerializer.scala.

[comphead]

Five lines of comment to say the line does nothing ("dropping this line leaks nothing"). Either drop both, or keep one line noting it matches what serializeBatches leaves behind.

Same pattern elsewhere in this file. decompressed, Projection and load each carry multi-paragraph rationale where the invariant is a sentence. The reference-counting notes earn their space. The reconstructions of what the code does do not.

[andygrove]

Fair. That one is three lines now and says only what it's for. I also cut the first paragraph of load's doc, which restated the class doc, and tightened Projection's. I left decompressed alone — both of its paragraphs are the reference-counting rationale you said earns its space.

[comphead]

Question: CompressionUtil.CodecType.fromCompressionType falls back to NO_COMPRESSION for an unrecognized byte rather than throwing, so this returns None and the payload is read as plain bytes. Is that intended? A batch written by a codec this build does not know would decode to garbage rather than fail.

[andygrove]

Not intended, no. fromCompressionType falls through to NO_COMPRESSION for anything outside its enum, so an unrecognized byte read the body as plain bytes and returned garbage rather than failing. readCodec now checks NoCompressionCodec.COMPRESSION_TYPE directly for the genuine uncompressed case and throws for anything else.

The realistic way to reach it is a corrupt payload rather than a writer from a newer build, since a cache never outlives the application that wrote it. Either way it's the same class of thing as the Major above: better to fail than to answer wrongly.

[comphead]

For a SELECT count(*) read, indices is empty and projection is never used, but a wide relation still pays toArrowSchema plus three full field-tree walks per partition. Hoisting the indices.isEmpty branch above this, or making projection a lazy val, keeps the cheapest read cheap. Worth it given this is exactly the wide-relation case the format targets.

[andygrove]

Made it a lazy val. The indices.isEmpty branch below already short-circuits the decode, so that alone is enough to keep a count(*) read from paying toArrowSchema plus the three field-tree walks.

[comphead]

A bare (String, Int) read back as codecSetting._1 and ._2 in encodeBatches. Two parameters, or a small case class, costs nothing and is equally serializable into the closure.

[andygrove]

Replaced with a CacheCodecSettings(name, zstdLevel) case class, with the "resolved on the driver" note moved onto it. Doing that shook out something else: the comment block above codecSettings actually describes encodeBatches and had been orphaned when that scaladoc went in, so I've moved it back and corrected its last sentence, which still named Utils.serializeBatches as the writer.

[comphead]

withNestedProjectionCache is inside the loop, so this materializes the whole relation six times. Caching once and restoring each batch's bytes from a clone() between iterations does the same work in a sixth of the time.

Also, the flat "decodes only the projected columns" above is the selectedIdx = 1 special case of this loop. One per-column loop parameterized over both relations would drop a test without dropping coverage.

[andygrove]

Both done. The fixture is out of the loop, with each batch's payload snapshotted once up front and restored between turns — the nested case runs in 213 ms where it was materializing the relation six times. And the flat test is folded into the same loop, so there's now one test parameterized over both relations rather than a general case and its special case.

[comphead]

This and "releases its vectors when a column fails to decode" assert the same thing at two corruption points, and the comment here says this is "the one that actually catches a leak". Worth merging into one test with two corruption modes rather than two fixtures and two allocator snapshots.

[andygrove]

Merged — one test, two corruption points, one fixture and one allocator snapshot per point. The commentary about what each point catches moved into the test's own comment.

[andygrove]

Thanks @comphead — the Major is real and the fix is in, along with everything else. Replies are on the individual threads; the two nits from your summary are done too (Projection takes an IndexedSeq[Field], and hydrateDictionaries is now decodeDictionaries).

Two things worth pulling up here rather than leaving buried in a thread.

The first is that I went wider than dropping FixedSizeBinaryVector from isArrowBacked. That closes the case you found, but not the class it belongs to: isArrowBacked answers for the top-level vector and never looks at children, so a struct whose child is a LargeVarCharVector passes it today and is stored with 64-bit offsets and read back with 32-bit. The length check you proposed can't catch that one either, because LargeUtf8 and Utf8 are both three buffers, so the totals agree and the values are still wrong. The write path now asks the question directly — do this batch's vectors already carry the Arrow types the reader will rebuild, recursively — and converts the batch when they don't, which is the path it already takes for non-Arrow input. That makes the fast path safe by construction rather than safe for the vectors we happened to enumerate.

The second is that the length check does carry its weight. I mutated it away and re-ran the test written for it: it fails with "no exception was thrown", which is exactly the silent-wrong-answer mode rather than a crash. So the read path really was returning an answer from a layout it had never checked.

Both checks also make the failure modes distinct in a way that should help later: a writer that produces an unexpected layout is now converted at cache time, and a payload that somehow still disagrees at read time fails with an error naming the mismatch rather than an ArrayIndexOutOfBoundsException from inside a buffer copy.

Also changed while I was in here: an unrecognized body-compression byte is now rejected rather than read as plain bytes — fromCompressionType falls back to NO_COMPRESSION, so a corrupt payload was decoding to garbage rather than failing.

Verified on the default profile: CometInMemoryCacheSuite 44/44, CometInMemoryCacheKryoSuite and UtilsSuite green, spotless:check and scalastyle:check clean, and the whole thing test-compiles against Spark 3.5 and 4.0 as well as 4.1.

[sunchao]

Reviewed 3e74e9dcde021a043a9ad6ccd0322ade2279151c across the IPC/vector boundary, Spark semantics, memory ownership, compatibility, and tests. The recursive writer-layout check and reader length guard address the earlier layout concern. I found two remaining issues, detailed inline: a write-side Arrow memory leak on recoverable compression failure, and an activation example that configures the cache after the serializer has already been selected.

Validation: JDK 17 root-reactor JVM test compilation passed. Exact-head IPC component checks passed 73,800 row comparisons across 19 types and 24 projections, including nested arrays/maps/structs, reordered/repeated/empty projections, nulls, timestamps, and uncompressed/zstd payloads. Normal round trips returned allocator usage to zero. I separately reproduced the compression-failure leak and checked the documentation example with actual Spark 4.1.3 contexts.

Limits: These were component and configuration probes, not full native-dependent Spark suites. Local native compilation was blocked because the configured dependency mirror lacks locked DataFusion 55.1.0. Benchmark performance was not remeasured. The previous-write-path comparison mentioned inline uses its single-column method sequence with an explicit SparkConf, not a full base-revision integration run.

Current CI: Required Checks is failing. Five lint checks failed; inspected logs request removing the redundant s prefix from the error string at CachedBatchIpc.scala:295. The Rust test check and Spark 4.1/JDK 17 build passed, while native-build and downstream jobs were cancelled and several optional jobs were skipped. CI's merge commit ca15f7da281b5a301e7498e9295d942537ec3d17 contains unrelated changes relative to this head, so those results are not exact-head validation.

[sunchao]

[P2] Release partial Arrow allocations when compression fails

VectorUnloader.getRecordBatch retains each input buffer and accumulates compressed buffers without cleaning them up if a later compression fails. That failure happens before this method's recordBatch.close() finally becomes active. Closing the input batch afterwards does not release those extra references or the earlier compressed buffers, so a failed cache materialization can leave batch-sized off-heap allocations behind.

I reproduced this against the current CachedBatchIpc.serialize and Arrow 18.3.0 with 4,194,304 integer values, an unlimited RootAllocator, and valid zstd level 22. After warming the child JVM, limiting its virtual address space to its current size plus 64 MiB caused zstd's native workspace allocation to fail with an ordinary RuntimeException: Error compressing: Allocation error : not enough memory. After closing the input batch, 18,350,080 bytes (17.5 MiB) remained allocated. The JVM remained usable. Under equivalent pressure, a component using the previous single-column stream-writing sequence threw ZstdIOException and returned Arrow allocation to zero after input close.

Please make write-side compression explicitly own and release partial buffers on failure, similar to the guarded decompression path, and add a failure-path regression test. Merely closing the input/root in a finally will not undo the unloader's retained references.

[andygrove]

You're right, and the repro made this quick to confirm. appendNodes retains each input buffer and then does buffers.add(codec.compress(...)) into a list that lives inside getRecordBatch, so when a compress throws, that retain is stranded and every buffer compressed before it is unreachable from anywhere a caller can get to. Closing the input batch afterwards releases the input's own reference, not the extra retain, and it never sees the compressed buffers at all.

Fixed by not handing the codec to the unloader. serialize now unloads with NoCompressionCodec, which retains and hands each buffer straight back without allocating, and a new CachedBatchIpc.compressed does the compression in the same shape decompressed already uses on the read side: retain, compress inside a try that releases on a throw, and an outer catch that closes whatever is already in the list. Peak memory is unchanged, since the unloader was doing the same thing one buffer at a time.

The retain is worth a note because it earns its keep differently in the two cases. A codec that allocates consumes the retained reference and returns a buffer of its own; NoCompressionCodec returns the input, and the retain is then the reference the result batch ends up owning. Either way the input buffer is back at its original count once the result is closed.

There is a regression test, and I mutated the fix away to check it isn't passing for free: it then fails with 2,176 bytes still allocated after the input is closed, at 256 rows, which is the same thing you measured at 17.5 MiB with a real workload. The codec it uses is a real zstd codec that throws on the third buffer, so the int column's two are genuine allocations by the time it does. Failing at the first buffer would have passed with no cleanup at all.

[sunchao]

[P3] Configure the cache before SparkContext startup

Starting from the documented default of false, this runtime setting cannot install Comet's cache serializer. CometDriverPlugin chooses spark.sql.cache.serializer once during SparkContext initialization, so subsequent cached relations still use Spark's default format and cannot take the native cache path.

I checked this with the current plugin on Spark 4.1.3: startup false followed by this runtime setting kept DefaultCachedBatchSerializer; startup true installed ArrowCachedBatchSerializer.

Please replace this activation example with --conf spark.comet.exec.inMemoryCache.enabled=true on startup, or builder configuration before creating a fresh SparkContext.

[andygrove]

Right, and the page contradicted itself: the paragraph immediately below the example said the config is read at startup. It is now a startup --conf on spark-shell, with a sentence saying why -- the driver plugin picks spark.sql.cache.serializer while the SparkContext is initializing, so a session that started with the default goes on using Spark's format however the config is set afterwards. The "is read at startup" sentence in the next section came out, since it was then saying the same thing twice.

[andygrove]

Thanks @sunchao — both are fixed in b978e54, with the details on the individual threads. Two things worth pulling up here.

The leak is real, and the fix is not local to the failure point: VectorUnloader accumulates the compressed buffers into a list inside getRecordBatch, so nothing a caller holds can reach them once it throws. serialize now unloads plain and compresses in a new CachedBatchIpc.compressed, the same shape decompressed already uses on the read side. I mutated the fix away to check the new test earns its place: it then fails with 2,176 bytes still allocated after the input batch is closed, which is your 17.5 MiB at a test-sized batch rather than a 4M-value one.

The lint failure was the redundant s prefix you found. All five failing checks were the same scalafix run across profiles, so that one line covers them; I scanned the rest of the PR's files for a second instance and there isn't one.

Verified on the default profile: CometInMemoryCacheSuite 45/45, CometInMemoryCacheKryoSuite and UtilsSuite green, spotless:check and scalastyle:check clean, and test-compiles against Spark 3.5 and 4.0 as well as 4.1.