mbutrovich commented on code in PR #5543:
URL: https://github.com/apache/datafusion-comet/pull/5543#discussion_r4085089070


##########
docs/source/user-guide/latest/in-memory-cache.md:
##########
@@ -0,0 +1,184 @@
+<!---
+  Licensed to the Apache Software Foundation (ASF) under one
+  or more contributor license agreements.  See the NOTICE file
+  distributed with this work for additional information
+  regarding copyright ownership.  The ASF licenses this file
+  to you under the Apache License, Version 2.0 (the
+  "License"); you may not use this file except in compliance
+  with the License.  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing,
+  software distributed under the License is distributed on an
+  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  KIND, either express or implied.  See the License for the
+  specific language governing permissions and limitations
+  under the License.
+-->
+
+# In-Memory Cache
+
+Comet can store Spark's in-memory cache (`CACHE TABLE`, `df.cache()`, 
`df.persist()`) in an Arrow
+format that Comet operators read directly. Without it, a cached table is 
stored in Spark's own
+format and every scan of it has to convert each batch before Comet can 
continue, which shows up in
+the plan as a `CometSparkColumnarToColumnar` above the cache scan.
+
+This feature is **experimental and disabled by default**. Turn it on at 
startup, alongside the rest
+of Comet's configuration:
+
+```shell
+$SPARK_HOME/bin/spark-shell \
+    ... \
+    --conf spark.comet.exec.inMemoryCache.enabled=true
+```
+
+It has to be set before the `SparkContext` starts. Comet's driver plugin 
chooses
+`spark.sql.cache.serializer` once, while the context is initializing, so a 
session that started
+with the default goes on using Spark's cache format however the config is set 
afterwards.
+
+## What changes when it is enabled
+
+With Comet's serializer installed as `spark.sql.cache.serializer`:
+
+- Cached data is stored as `CometCachedBatch` rather than Spark's 
`DefaultCachedBatch`.
+- Cached tables are scanned by `CometInMemoryTableScan`, which feeds Comet 
operators directly.
+- Per-batch column statistics are recorded in the layout Spark's 
`SimpleMetricsCachedBatchSerializer`
+  expects, so Spark can prune whole cached batches on a predicate before any 
of them is decoded.
+
+Relations whose schema Comet's Arrow writer cannot store — interval types, 
most notably — are
+delegated in full to Spark's default cache format, per relation. Nothing about 
the format depends
+on a runtime config, because `spark.sql.cache.serializer` is a static setting 
and a relation whose
+format could change mid-session could not be read back reliably. Turning
+`spark.comet.exec.inMemoryCache.enabled` off at runtime only sends cached 
scans back to Spark's
+execution path; the cached data stays readable either way.
+
+## Storage format
+
+Each cached batch is stored as a single Arrow IPC record batch message and its 
body.
+
+The message carries **no Arrow schema**. The reader already has one: 
`InMemoryRelation` knows the
+cached relation's attributes, and Comet maps them to exactly the Arrow fields 
the writer produced.
+Storing a schema in every batch would repeat the same bytes once per cached 
batch — for a wide
+relation cached in many batches, a large share of a payload that is not data.
+
+Compression is applied by Arrow to **each buffer separately**, rather than by 
wrapping the whole
+payload in a Spark compression codec. That is what makes a projected read 
cheap: the message
+metadata records every buffer's offset and length within the body, so a scan 
copies out only the
+byte ranges belonging to the columns it selected, and only those are 
decompressed. A read of one
+column out of six does roughly a sixth of the decompression work, and a 
`SELECT count(*)`, which
+selects no columns at all, answers from the row count stored beside the 
payload without touching
+it.
+
+Compression defaults to `zstd`, for footprint rather than for speed. Over the 
same 5M-row,
+six-column relation the tables under [Performance](#performance) use — and 
measured by the same
+benchmark — it holds the data in a sixth of the memory and pays for that on 
both sides: about 40%
+longer to materialize, and, on a read wide enough to inflate everything, close 
to five times
+longer. A narrow projection pays far less, because it only inflates the 
columns it asked for.
+
+| Codec  | Materialize | Footprint | Read 1 of 6 | Read 6 of 6 |
+| ------ | ----------: | --------: | ----------: | ----------: |
+| `zstd` |     1503 ms |    55 MiB |       47 ms |      296 ms |
+| `none` |     1091 ms |   315 MiB |       36 ms |       63 ms |
+
+`none` is the better setting for a relation that fits in memory uncompressed 
and is read at close
+to full width. The default is the other way round because a cache that does 
not fit costs more than
+one that is slower to read, and Spark's own cache format compresses by default 
too.
+
+Arrow's other IPC codec, LZ4, is deliberately not offered and the config 
rejects it. It is
+commons-compress's pure-Java implementation, unrelated to the JNI-accelerated 
lz4-java behind
+`spark.io.compression.codec`, and is orders of magnitude slower to write than 
`zstd` while also
+producing larger output.
+
+Dictionary-encoded columns are decoded before they are stored. A payload with 
no schema message has
+nowhere to record either that a column is dictionary encoded or the dictionary 
itself.
+
+## Configuration
+
+| Config                                                  | Default | 
Description                                                                     
                                                               |
+| ------------------------------------------------------- | ------- | 
----------------------------------------------------------------------------------------------------------------------------------------------
 |
+| `spark.comet.exec.inMemoryCache.enabled`                | `false` | Whether 
to store and scan Spark's in-memory cache in Comet's format. Read at startup.   
                                                       |
+| `spark.comet.exec.inMemoryCache.compression.codec`      | `zstd`  | Arrow 
IPC compression codec for cached data: `zstd` or `none`. Affects newly cached 
data only — a batch records the codec it was written with. |
+| `spark.comet.exec.inMemoryCache.compression.zstd.level` | `1`     | 
Compression level when the codec is `zstd`. Ignored otherwise.                  
                                                               |
+
+## Performance
+
+Measured with `CometInMemoryCacheBenchmark` (Apple M3 Max, JDK 17, Spark 4.1, 
release build).
+Regenerate with:
+
+```sh
+SPARK_GENERATE_BENCHMARK_FILES=1 \
+  make benchmark-org.apache.spark.sql.benchmark.CometInMemoryCacheBenchmark
+```
+
+On a 5M-row relation of six flat columns:
+
+| Query shape                    | Spark cache scan + convert | 
`CometInMemoryTableScan` | Relative |
+| ------------------------------ | -------------------------: | 
-----------------------: | -------: |
+| Repeated scan (3 of 6 columns) |                     180 ms |                
   147 ms |     1.2x |
+| Selective filter               |                      58 ms |                
    50 ms |     1.2x |
+| Row count only (0 of 6)        |                      37 ms |                
    34 ms |     1.1x |
+| Narrow projection (1 of 6)     |                      52 ms |                
    48 ms |     1.1x |
+| Full projection (6 of 6)       |                     547 ms |                
   291 ms |     1.9x |
+
+And on a 1M-row relation of six columns whose middle three are structs, one of 
them nested two
+levels deep:
+
+| Query shape                | Spark cache scan + convert | 
`CometInMemoryTableScan` | Relative |
+| -------------------------- | -------------------------: | 
-----------------------: | -------: |
+| Row count only (0 of 6)    |                      30 ms |                    
29 ms |     1.0x |
+| Narrow projection (1 of 6) |                      98 ms |                    
53 ms |     1.8x |
+| Full projection (6 of 6)   |                     250 ms |                   
125 ms |     2.0x |
+
+Both columns read the cache at the default codec, `zstd`. The codec table 
above shows what `none`
+changes, and it is the full projection that moves most: nothing has to be 
inflated, so it runs
+several times faster, at six times the memory.
+
+The two relations are not comparable to each other — different row counts, and 
a struct column
+carries several values per row.
+
+Array and map columns are deliberately absent from the benchmark, not from the 
format — the cache
+stores and projects them, and `CometInMemoryCacheSuite` covers them. They 
cannot be measured _here_
+because the left column would not exist: it needs Spark's cache scan to bridge 
into Comet operators,
+and `CometSparkToColumnarExec` declines `ArrayType` and `MapType`, so a query 
projecting one falls
+back to Spark row execution above the scan and the two columns stop measuring 
the same boundary.
+
+Read what this compares carefully. Comet execution is on in both columns, so 
the aggregation runs
+on Comet either way and only the cache-scan boundary moves: on the left, 
Spark's
+`InMemoryTableScanExec` feeds those same Comet operators through a 
`CometSparkColumnarToColumnar`
+bridge; on the right, `CometInMemoryTableScan` feeds them directly. Both 
columns read the same
+Comet-written `CometCachedBatch` — `spark.sql.cache.serializer` is static, so 
one session cannot
+also materialize Spark's format to compare against. These numbers are 
therefore "keep the cached
+scan native" against "fall back to a Spark cache scan and convert", not Comet 
against Spark
+execution, and not a comparison with Spark's own cache format.
+
+## Kryo
+
+Spark serializes a cached batch with `spark.serializer` whenever the block 
leaves the heap: the
+`_SER` storage levels, replication, cross-executor fetches, and the disk half 
of the default
+`MEMORY_AND_DISK`. So an ordinary `df.cache()` that spills is enough to reach 
it.
+
+If you run with `spark.kryo.registrationRequired=true`, register Comet's 
classes:
+
+```
+spark.serializer=org.apache.spark.serializer.KryoSerializer
+spark.kryo.registrationRequired=true
+spark.kryo.registrator=org.apache.comet.CometKryoRegistrator
+```
+
+Comet cannot set `spark.kryo.registrator` for you the way it sets 
`spark.sql.cache.serializer`:
+`KryoSerializer` reads it when `SparkEnv` builds the serializer, which happens 
before any plugin
+runs. Without it, caching fails with a "Class is not registered" error that 
does not name this
+feature. Comet's driver plugin warns at startup when it sees Kryo, 
`registrationRequired`, and no
+registrator.
+
+## Limitations
+
+Reads that feed **Spark** operators rather than Comet ones are still slower 
than Spark's own cache
+format, by roughly 1.7x to 2.5x depending on how wide the projection is. Those 
reads pay a row
+conversion that Spark's format avoids with generated code over its own layout. 
This is why the
+feature is off by default.

Review Comment:
   Where does "roughly 1.7x to 2.5x" come from? It matches the table in #5485, 
which was measured on the per-column format with Spark's whole-stream codec. 
This PR changes both, and the new codec table shows a full-width zstd read at 
296 ms against 63 ms for `none`, so the Spark-operator read path has probably 
moved too. Could you re-measure it on this format, or link #5485 and say which 
format the number describes? This is the reason the page gives for the feature 
being off by default, so it should describe the format that ships.



##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala:
##########
@@ -0,0 +1,633 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.arrow
+
+import java.nio.ByteBuffer
+import java.nio.channels.Channels
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
+
+import org.apache.arrow.compression.{CommonsCompressionFactory, 
ZstdCompressionCodec}
+import org.apache.arrow.flatbuf.{RecordBatch => FlatBufRecordBatch}
+import org.apache.arrow.memory.{ArrowBuf, BufferAllocator}
+import org.apache.arrow.vector.{FieldVector, TypeLayout, ValueVector, 
VectorLoader, VectorSchemaRoot, VectorUnloader}
+import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, 
NoCompressionCodec}
+import org.apache.arrow.vector.dictionary.DictionaryEncoder
+import org.apache.arrow.vector.ipc.{ReadChannel, WriteChannel}
+import org.apache.arrow.vector.ipc.message.{ArrowBodyCompression, 
ArrowFieldNode, ArrowRecordBatch, MessageSerializer}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema}
+import org.apache.arrow.vector.util.DataSizeRoundingUtil
+import org.apache.spark.SparkException
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.util.io.{ChunkedByteBuffer, 
ChunkedByteBufferOutputStream}
+
+import org.apache.comet.vector.CometVector
+
+/**
+ * The on-disk shape of a `CometCachedBatch` payload, and the two operations 
over it.
+ *
+ * A cached batch is one encapsulated Arrow IPC RecordBatch message followed 
by its body, with no
+ * Schema message and no end-of-stream marker. The schema is not stored 
because the reader already
+ * has it: `InMemoryRelation` knows the cached relation's attributes, and 
`Utils.toArrowSchema`
+ * maps them to exactly the fields the writer unloaded. Leaving it out saves a 
schema message per
+ * cached batch, which for a wide relation cached in many batches is a large 
share of the payload
+ * that is not data.
+ *
+ * Compression is applied by Arrow per buffer rather than by wrapping the 
whole payload in a Spark
+ * `CompressionCodec`. That is what makes projection cheap: the message 
metadata records every
+ * buffer's offset and length within the body, so [[Projection.load]] can copy 
out only the
+ * buffers of the columns a scan selected and let `VectorLoader` decompress 
just those. A
+ * whole-payload codec would have to inflate everything before any column 
could be read.
+ */
+private[comet] object CachedBatchIpc {
+
+  /**
+   * The Arrow compression codec named by 
`spark.comet.exec.inMemoryCache.compression.codec`.
+   *
+   * Only the write path consults the config. A batch records which codec 
compressed it, so the
+   * read path looks the codec up from the batch itself and keeps reading data 
cached before the
+   * config changed.
+   */
+  def compressionCodec(codecName: String, zstdLevel: Int): CompressionCodec = 
codecName match {
+    case "none" => NoCompressionCodec.INSTANCE
+    // Constructed directly rather than through CompressionCodec.Factory, 
which ignores the level
+    // and always builds a codec at zstd's default.
+    case "zstd" => new ZstdCompressionCodec(zstdLevel)
+    // Arrow's other codec, LZ4_FRAME, is not offered. It is 
commons-compress's pure-Java LZ4 --
+    // no relation to the JNI-accelerated lz4-java behind 
spark.io.compression.codec -- and
+    // measures three orders of magnitude slower to write than zstd while also 
producing larger
+    // output, so nothing prefers it. Reads still accept it, since the factory 
the read path uses
+    // handles whatever codec a batch records.
+    case other =>
+      throw new SparkException(
+        s"Unsupported Arrow compression codec for Comet's cache: $other. " +
+          "Supported values: none, zstd")
+  }
+
+  // Decompressors are stateless and shared. Resolving one per cached batch 
would allocate a codec
+  // per batch on every scan, and the enum lookup walks the CodecType values 
each time.
+  private val readCodecs: Map[CompressionUtil.CodecType, CompressionCodec] =
+    CompressionUtil.CodecType
+      .values()
+      .filter(_ != CompressionUtil.CodecType.NO_COMPRESSION)
+      .map(t => t -> CommonsCompressionFactory.INSTANCE.createCodec(t))
+      .toMap
+
+  /**
+   * The decompressor for a body-compression byte, or None when the batch is 
stored plain.
+   *
+   * A byte this build does not recognize is rejected rather than read as 
plain bytes.
+   * `CodecType.fromCompressionType` answers `NO_COMPRESSION` for anything 
outside its enum, so
+   * taking its word for it would turn a corrupt payload into garbage values 
instead of an error.
+   */
+  private def readCodec(compressionType: Byte): Option[CompressionCodec] =
+    if (compressionType == NoCompressionCodec.COMPRESSION_TYPE) {
+      None
+    } else {
+      val codecType = 
CompressionUtil.CodecType.fromCompressionType(compressionType)
+      if (codecType == CompressionUtil.CodecType.NO_COMPRESSION) {
+        throw new SparkException(
+          s"Comet cached batch records an unknown Arrow compression codec: 
$compressionType")
+      }
+      Some(readCodecs(codecType))
+    }
+
+  /**
+   * Whether `batch`'s vectors can be unloaded as they stand, or have to be 
converted first.
+   *
+   * The payload records no schema, so [[Projection]] rebuilds the fields from 
the cached
+   * relation's Spark attributes and reads the body against them. The direct 
write path unloads
+   * whatever vectors the cached plan produced, and one Spark type can arrive 
as more than one
+   * Arrow type: `BinaryType` is a `VarBinaryVector` from Comet's own scans 
but a
+   * `FixedSizeBinaryVector` from an accelerated `mapInArrow` or an Iceberg 
`fixed[N]` read, and
+   * those occupy three buffers and two. Writing one and reading the other 
shifts every buffer
+   * from that column on, which is wrong values rather than an error, so a 
batch that does not
+   * already carry the reader's types is converted instead.
+   *
+   * The same holds inside a nested column, which `Utils.isArrowBacked` does 
not look at: it
+   * answers for the top-level vector only, so a struct of large strings 
passes it while its child
+   * is stored with 64-bit offsets and read with 32-bit ones.
+   *
+   * Names, nullability and a timestamp's timezone are not compared. None of 
them changes how the
+   * reader interprets the body, and the writer's legitimately differ -- a 
Comet scan labels
+   * timestamps with the session's zone where the reader rebuilds them as UTC, 
which is a label
+   * only: Spark's representation is micros since the epoch either way.
+   */
+  def matchesReaderLayout(batch: ColumnarBatch, readerFields: Seq[Field]): 
Boolean =
+    batch.numCols() == readerFields.length &&
+      (0 until batch.numCols()).forall { i =>
+        batch.column(i) match {
+          case v: CometVector => sameLayout(writtenField(v), readerFields(i))
+          case _ => false
+        }
+      }
+
+  /**
+   * The field a column reaches the body as.
+   *
+   * A dictionary-encoded vector's own field carries the index type, not the 
values', because
+   * [[decodeDictionaries]] replaces it with the decoded form before anything 
is unloaded.
+   * Resolved through the same `lookupDictionary` the write path uses, so a 
batch missing its
+   * dictionary fails here exactly as it would there.
+   */
+  private def writtenField(column: CometVector): Field = {
+    val vector = column.getValueVector
+    if (vector.getField.getDictionary == null) {
+      vector.getField
+    } else {
+      Utils
+        .lookupDictionary(vector.asInstanceOf[FieldVector], 
Option(column.getDictionaryProvider))
+        .getVector
+        .getField
+    }
+  }
+
+  private def sameLayout(written: Field, read: Field): Boolean =
+    layoutType(written.getType) == layoutType(read.getType) && {
+      val writtenChildren = written.getChildren
+      val readChildren = read.getChildren
+      writtenChildren.size == readChildren.size &&
+      (0 until writtenChildren.size).forall(i =>
+        sameLayout(writtenChildren.get(i), readChildren.get(i)))
+    }
+
+  private def layoutType(t: ArrowType): ArrowType = t match {
+    case ts: ArrowType.Timestamp if ts.getTimezone != null =>
+      new ArrowType.Timestamp(ts.getUnit, "UTC")
+    case other => other
+  }
+
+  /**
+   * Serialize `batch` into one encapsulated IPC RecordBatch message.
+   *
+   * The message is written into heap chunks of `chunkSize` bytes rather than 
one array, so a
+   * batch is not bounded by the 2 GiB a single JVM array can hold. Nothing 
bounds a cached
+   * batch's bytes upstream: the row path cuts batches by row count alone, and 
a columnar input
+   * batch arrives at whatever size the plan above produced. Chunks are 
appended rather than grown
+   * and recopied, so the write also never holds the payload twice.
+   *
+   * Returns the message and the on-body compressed size of each top-level 
column, which the
+   * caller records in the statistics row. The sizes come from the message's 
own buffer layout, so
+   * they are the real stored sizes rather than an estimate.
+   *
+   * Dictionary-encoded columns are decoded to their plain form first. A 
payload with no Schema
+   * message cannot describe a dictionary encoding, and the schema the reader 
rebuilds from Spark
+   * attributes never carries one, so a dictionary-encoded column has nowhere 
to record either its
+   * index type or the dictionary itself. Comet's native scans do produce such 
columns, so this is
+   * a real path, not a defensive one.
+   *
+   * As in `Utils.serializeBatches`, `batch`'s vectors are cleared once 
written, so callers gather
+   * anything they need from the batch (statistics, for instance) before 
calling this.
+   */
+  def serialize(
+      batch: ColumnarBatch,
+      codec: CompressionCodec,
+      allocator: BufferAllocator,
+      chunkSize: Int): (ChunkedByteBuffer, Array[Long]) = {
+    val (vectors, decoded) = decodeDictionaries(batch, allocator)
+    try {
+      val root = new VectorSchemaRoot(vectors.asJava)
+      // A batch of zero columns carries only a row count, which a 
VectorSchemaRoot cannot infer
+      // without vectors to measure.
+      if (vectors.isEmpty) {
+        root.setRowCount(batch.numRows())
+      }
+
+      // Unloaded plain and compressed afterwards rather than by handing the 
codec to the unloader;
+      // see compressed for why.
+      val unloader = new VectorUnloader(root, true, 
NoCompressionCodec.INSTANCE, true)
+      val plainBatch = unloader.getRecordBatch
+      val recordBatch =
+        try compressed(plainBatch, codec, allocator)
+        finally plainBatch.close()
+      try {
+        val fields = vectors.map(_.getField)
+        // Leaves the batch in the state serializeBatches leaves one. The 
record batch holds its
+        // own buffers by now, so this does not touch it, and getField still 
answers afterwards:
+        // clearing releases buffers, not the schema.
+        root.clear()
+
+        val out = new ChunkedByteBufferOutputStream(chunkSize, 
ByteBuffer.allocate)
+        try {
+          val channel = new WriteChannel(Channels.newChannel(out))
+          MessageSerializer.serialize(channel, recordBatch)
+        } finally {
+          out.close()
+        }
+        (out.toChunkedByteBuffer, columnSizes(fields, recordBatch))
+      } finally {
+        recordBatch.close()
+      }
+    } finally {
+      // Only the vectors this method allocated. The rest belong to the input 
batch.
+      decoded.foreach(v =>
+        try v.close()
+        catch { case NonFatal(_) => () })
+    }
+  }
+
+  /**
+   * Everything about reading one projection of this format that does not 
change between batches.
+   * A scan builds one of these per partition.
+   *
+   * The index arithmetic walks every field of the cached relation rather than 
just the projected
+   * ones, so recomputing 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.
+   *
+   * Holding the projected `Schema` here too is what keeps it consistent with 
the buffers:
+   * [[Projection.load]] packs field nodes and buffers by walking 
`selectedIndices` in order, and
+   * the schema is built from the same walk, so the two cannot drift apart.
+   */
+  final class Projection(arrowFields: IndexedSeq[Field], selectedIndices: 
Array[Int]) {
+
+    private val schema = new 
Schema(selectedIndices.map(arrowFields).toSeq.asJava)
+
+    // A record batch body is a flat, depth-first sequence of buffers in 
schema order, so each
+    // top-level column owns a contiguous run of it; field nodes run in the 
same order. The totals
+    // are what a payload is checked against in load.
+    private val (nodeIndices, totalNodes) =
+      selectedRange(arrowFields, selectedIndices, fieldNodeCount)
+    private val (bufferIndices, totalBuffers) =
+      selectedRange(arrowFields, selectedIndices, fieldBufferCount)
+
+    /**
+     * Decode the projected columns of one cached payload into a fresh root 
the caller owns.
+     *
+     * A buffer's recorded (offset, length) covers its on-body bytes including 
the
+     * uncompressed-length prefix, so a window copied out of the payload is 
exactly what the
+     * writer emitted, 8-byte aligned as Arrow's IPC body lays it out. The 
columns that were not
+     * selected are never read, let alone inflated. The windows are then 
decompressed in one pass;
+     * see [[decompressed]] for why that is not left to `VectorLoader`.
+     */
+    def load(data: ChunkedByteBuffer, allocator: BufferAllocator): 
VectorSchemaRoot = {
+      val readChannel = new 
ReadChannel(Channels.newChannel(data.toInputStream()))
+      // Reads the message metadata only. The body stays in `data` and is 
copied selectively.
+      val metadata = MessageSerializer.readMessage(readChannel)
+      if (metadata == null) {
+        throw new SparkException("Unexpected end of input reading a Comet 
cached batch")
+      }
+      val batch =
+        metadata.getMessage.header(new 
FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch]
+
+      // The payload carries no schema, so nothing in it says the writer laid 
the body out the way
+      // these windows read it. batch.buffers(j) is an unchecked flatbuffer 
accessor, so a
+      // disagreement would otherwise surface as wrong values, or as an 
out-of-range read from
+      // inside the copy below, rather than as an error naming the cause. See 
matchesReaderLayout
+      // for how the write path avoids producing one.
+      if (batch.nodesLength() != totalNodes || batch.buffersLength() != 
totalBuffers) {
+        throw new SparkException(
+          "Comet cached batch does not match the cached schema: the payload 
holds " +
+            s"${batch.nodesLength()} field nodes and ${batch.buffersLength()} 
buffers, but the " +
+            s"schema describes $totalNodes and $totalBuffers")
+      }
+
+      // serialize writes exactly [encapsulated message][body] and nothing 
after it, so the body is
+      // the tail of `data`.
+      val bodyStart = data.size - metadata.getMessageBodyLength
+      val chunks = new PayloadChunks(data)
+
+      val compression =
+        if (batch.compression() == null) 
NoCompressionCodec.DEFAULT_BODY_COMPRESSION
+        else new ArrowBodyCompression(batch.compression().codec(), 
batch.compression().method())
+
+      val nodes = new java.util.ArrayList[ArrowFieldNode](nodeIndices.length)
+      nodeIndices.foreach { j =>
+        val node = batch.nodes(j)
+        nodes.add(new ArrowFieldNode(node.length(), node.nullCount()))
+      }
+
+      val offsets = new Array[Long](bufferIndices.length)
+      val lengths = new Array[Long](bufferIndices.length)
+      var total = 0L
+      var k = 0
+      while (k < bufferIndices.length) {
+        val buffer = batch.buffers(bufferIndices(k))
+        offsets(k) = buffer.offset()
+        lengths(k) = buffer.length()
+        total += DataSizeRoundingUtil.roundUpTo8Multiple(lengths(k))
+        k += 1
+      }
+
+      // allocator.buffer(0) is legal but yields a buffer no window can be 
sliced from, and an
+      // all-empty projection (every selected column a NullVector, say) would 
ask for exactly that.

Review Comment:
   Could you add a test for the case this comment describes? I couldn't find a 
test that caches a `NullType` column, so a projection of only NullVector 
columns (`total == 0`) never reaches `load`. A zero-row batch is similar: every 
buffer is empty, so it goes through zstd's empty-buffer framing on write and 
the zero-length windows on read. Something like caching `SELECT id, NULL AS n` 
and reading `SELECT n` under both codecs, plus a zero-row batch through 
`CometCachedBatchHelper.serialize` and `Projection.load`, would cover both 
paths.



##########
docs/source/user-guide/latest/in-memory-cache.md:
##########
@@ -0,0 +1,184 @@
+<!---
+  Licensed to the Apache Software Foundation (ASF) under one
+  or more contributor license agreements.  See the NOTICE file
+  distributed with this work for additional information
+  regarding copyright ownership.  The ASF licenses this file
+  to you under the Apache License, Version 2.0 (the
+  "License"); you may not use this file except in compliance
+  with the License.  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing,
+  software distributed under the License is distributed on an
+  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  KIND, either express or implied.  See the License for the
+  specific language governing permissions and limitations
+  under the License.
+-->
+
+# In-Memory Cache
+
+Comet can store Spark's in-memory cache (`CACHE TABLE`, `df.cache()`, 
`df.persist()`) in an Arrow
+format that Comet operators read directly. Without it, a cached table is 
stored in Spark's own
+format and every scan of it has to convert each batch before Comet can 
continue, which shows up in
+the plan as a `CometSparkColumnarToColumnar` above the cache scan.
+
+This feature is **experimental and disabled by default**. Turn it on at 
startup, alongside the rest
+of Comet's configuration:
+
+```shell
+$SPARK_HOME/bin/spark-shell \
+    ... \
+    --conf spark.comet.exec.inMemoryCache.enabled=true
+```
+
+It has to be set before the `SparkContext` starts. Comet's driver plugin 
chooses
+`spark.sql.cache.serializer` once, while the context is initializing, so a 
session that started
+with the default goes on using Spark's cache format however the config is set 
afterwards.
+
+## What changes when it is enabled
+
+With Comet's serializer installed as `spark.sql.cache.serializer`:
+
+- Cached data is stored as `CometCachedBatch` rather than Spark's 
`DefaultCachedBatch`.
+- Cached tables are scanned by `CometInMemoryTableScan`, which feeds Comet 
operators directly.
+- Per-batch column statistics are recorded in the layout Spark's 
`SimpleMetricsCachedBatchSerializer`
+  expects, so Spark can prune whole cached batches on a predicate before any 
of them is decoded.
+
+Relations whose schema Comet's Arrow writer cannot store — interval types, 
most notably — are
+delegated in full to Spark's default cache format, per relation. Nothing about 
the format depends
+on a runtime config, because `spark.sql.cache.serializer` is a static setting 
and a relation whose
+format could change mid-session could not be read back reliably. Turning
+`spark.comet.exec.inMemoryCache.enabled` off at runtime only sends cached 
scans back to Spark's
+execution path; the cached data stays readable either way.

Review Comment:
   "Nothing about the format depends on a runtime config" isn't true anymore. 
`spark.comet.exec.inMemoryCache.compression.codec` is a runtime config, and it 
decides the codec byte each batch is written with. The data stays readable 
because the reader takes the codec from the batch, so how about saying that 
instead?
   
   ```suggestion
   Relations whose schema Comet's Arrow writer cannot store — interval types, 
most notably — are
   delegated in full to Spark's default cache format, per relation. Which 
format a relation uses does
   not depend on a runtime config, because `spark.sql.cache.serializer` is a 
static setting and a
   relation whose format could change mid-session could not be read back 
reliably. The compression
   codec is a runtime config, but each batch records the codec it was written 
with, so data cached
   under one setting stays readable after the setting changes. Turning
   `spark.comet.exec.inMemoryCache.enabled` off at runtime only sends cached 
scans back to Spark's
   execution path; the cached data stays readable either way.
   ```



-- 
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]

Reply via email to