Adopt techniques from Spark's ArrowCachedBatchSerializer (SPARK-57268) in Comet's cache format

[andygrove]

What is the problem the feature request solves?

Spark added its own Arrow in-memory cache format in SPARK-57268 (org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer), solving the same problem as the serializer added in #5051. It is present in branch-4.3 and master only, so it is not available in any Spark version Comet supports today (3.4 through 4.2) and does not remove the need for Comet's own serializer. It is, however, a more mature implementation of the same idea, and several of its decisions are worth adopting.

Recorded here so the comparison is not lost. Raised on #5051 by @sunchao, who pointed at the Spark work; @viirya wrote it.

Where the two implementations already agree, for the record: empty projections emit row counts without touching the payload, statistics use the same five-field SimpleMetricsCachedBatch layout, the payload is schema-agnostic with respect to session timezone, cleanup on early termination goes through a TaskCompletionListener, and the format is opt-in via spark.sql.cache.serializer with the default unchanged. Spark reached the empty-projection optimization in SPARK-58390, a follow-up eight days after the main PR, which is the same order we found it in.

Describe the potential solution

Roughly in order of value for effort.

• Drop the schema message from each cached stream. Spark's ArrowCachedBatch deliberately stores an encapsulated Arrow RecordBatch message with no Schema message and no end-of-stream marker, reconstructing the schema from the relation's attributes on read, explicitly to avoid repeating the schema bytes in every cached batch. Comet writes a full ArrowStreamWriter stream per column per batch, so a 60-column relation cached in 500 batches writes about 30,000 schema messages. That is a direct contributor to the framing overhead measured in Comet cache decode ignores column projection, making narrow reads of wide cached relations slower than Spark #5484 (2.5% more footprint at 6 columns, 32% at 60). Rough arithmetic puts the schema bytes at about a third of that. Contained change: MessageSerializer.serialize and deserializeRecordBatch plus VectorLoader, instead of ArrowStreamWriter and ArrowStreamReader.

• Register CometCachedBatch for Kryo. Done in feat: add experimental native support for in-memory cache, disabled by default #5051 as org.apache.comet.CometKryoRegistrator. The scoping in the original wording above was wrong, as @viirya pointed out: this is not limited to the explicitly serialized storage levels. The _SER levels, replication, cross-executor fetches and the disk half of the default MEMORY_AND_DISK all serialize a CachedBatch, so a plain .cache() that spills is enough to hit it. It also turned out not to be cache-specific -- CometBroadcastExchangeExec broadcasts an Array[ChunkedByteBuffer], which Spark does not register either, so native broadcast failed under registrationRequired=true independently of this feature (filed separately as Comet native broadcast fails under spark.kryo.registrationRequired=true #5510, since it is a pre-existing bug that need not wait on the cache feature). Unlike spark.sql.cache.serializer, Comet cannot install spark.kryo.registrator itself, because KryoSerializer reads it when SparkEnv builds it and that is before any plugin runs; it is documented on the config and CometDriverPlugin warns at startup instead.

• Typed readers for the row path. Spark builds typed ArrowColumnReaders once and writes straight into an UnsafeRowWriter, avoiding per-row pattern matching, with an explicit fallback to a columnar-to-row path for complex types (array, struct, map, UDT, variant, geometry, nanosecond timestamps). Comet's convertCachedBatchToInternalRow decodes to a ColumnarBatch and then does batch.rowIterator().map(row => toUnsafe(row).copy()). This is a concrete candidate for the gap in Cached reads that feed Spark operators are slower than Spark's own cache format #5485, which currently has no established cause.

• Evaluate projection by buffer selection instead of per-column streams. This is the significant one. Rather than splitting the payload, Spark keeps one RecordBatch per cached batch and its readProjectedRecordBatch parses the IPC message flatbuffer, which lists every buffer's offset and length within the body, then copies only the byte ranges belonging to the selected columns into a single off-heap buffer, so VectorLoader.load decompresses only those. It depends on using Arrow's native per-buffer compression (Lz4CompressionCodec, ZstdCompressionCodec from arrow-compression) rather than wrapping the whole stream in a Spark CompressionCodec as Comet does.

  That reaches the same projection-proportional decode Comet now has, with no per-column framing overhead at all, so it dominates the current design on footprint. The cost is roughly 120 lines of intricate code: field node counts, buffer span arithmetic, variadic buffer counts and 8-byte alignment, with correctness that is not locally obvious. Worth its own PR rather than an amendment to an existing one. Note this subsumes the first item above.

• Optional background prefetch of the next batch, decompressing and deserializing off the consumer thread, config-gated and off by default.

• Prune on collated string columns. Spark compares string bounds with UTF8String.semanticCompare(min, collationId). Comet's tracksBounds matches case StringType, which a collated StringType does not equal, so collated columns get null bounds and buildFilter declines to push predicates on them. That is correct and covered by a test, but it means no pruning where Spark manages it.

• Documentation and committed benchmark results. Spark ships a sql-arrow-cache-format.md page linked from the SQL docs menu, and commits ArrowCacheBenchmark-jdk{17,21,25}-results.txt generated by a benchmark workflow.

Additional context

Longer term, once Comet supports a Spark version that ships ArrowCachedBatchSerializer, it is worth asking whether Comet should consume Spark's ArrowCachedBatch directly instead of installing its own serializer. That would work with vanilla Spark's cache and drop a format from Comet's maintenance surface. Out of scope until Comet supports Spark 4.3.

Follow-up from #5051. Related: #4781, #5245, #5484, #5485.

[viirya]

Yea, I also thought that maybe we can incorporate the Arrow cache with Comet to have better query execution.