apache / apache/datafusion-comet
Adopt techniques from Spark's ArrowCachedBatchSerializer (SPARK-57268) in Comet's cache format
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
### What is the problem the feature request solves?
Spark added its own Arrow in-memory cache format in [SPARK-57268](https://github.com/apache/spark/pull/56334) (`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](https://github.com/apache/spark/commit/073c8d8f4), 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 #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`.
- [x] **Register `CometCachedBatch` for Kryo.** Done in #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 #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 `ArrowColumnReader`s 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 #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.
Contributor guide
Assessment
This issue has not been assessed yet.