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

docs/source/user-guide/latest/in-memory-cache.md

@@ -0,0 +1,148 @@
+<!---
+  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**.
+
+```scala
+spark.conf.set("spark.comet.exec.inMemoryCache.enabled", "true")
+```
+
+## What changes when it is enabled
+
+`spark.comet.exec.inMemoryCache.enabled` is read at startup, and its value then decides whether
+Comet installs its cache serializer as `spark.sql.cache.serializer`. When it is installed:
+
+- 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`, which is faster than storing cached batches uncompressed: the
+bytes it saves cost more to copy and store than compressing them costs. Measured 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 |
+
+Arrow's other IPC codec, LZ4, is deliberately not offered. It is commons-compress's pure-Java
+implementation and is unrelated to the JNI-accelerated lz4-java behind `spark.io.compression.codec`;
+it measured three orders of magnitude slower to write than `zstd` while also producing larger
+output, so no workload prefers it.
+
+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` on a 5M-row, six-column relation (Apple M3 Ultra,
+JDK 17, Spark 4.1, release build). Regenerate with:
+
+```sh
+SPARK_GENERATE_BENCHMARK_FILES=1 \
+  make benchmark-org.apache.spark.sql.benchmark.CometInMemoryCacheBenchmark
+```
+
+| Query shape                    | Spark cache scan + convert | `CometInMemoryTableScan` | Relative |
+| ------------------------------ | -------------------------: | -----------------------: | -------: |
+| Repeated scan (3 of 6 columns) |                     157 ms |                   116 ms |     1.4x |
+| Selective filter               |                      44 ms |                    38 ms |     1.1x |
+| Row count only (0 of 6)        |                      30 ms |                    28 ms |     1.1x |
+| Narrow projection (1 of 6)     |                      49 ms |                    39 ms |     1.3x |
+| Full projection (6 of 6)       |                     299 ms |                   135 ms |     2.2x |
+
+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.
+
+Comet's serializer exists because Spark's own Arrow cache format
+([SPARK-57268](https://issues.apache.org/jira/browse/SPARK-57268)) is only available from Spark
+4.3, which Comet does not yet support.

docs/source/user-guide/latest/index.rst

@@ -74,6 +74,7 @@ to read more.
    Understanding Comet Plans <understanding-comet-plans>
    Tuning Guide <tuning>
    Metrics Guide <metrics>
+   In-Memory Cache <in-memory-cache>
    PyArrow UDF Acceleration <pyarrow-udfs>
 
 .. toctree::

pom.xml

@@ -225,6 +225,32 @@ under the License.
         <artifactId>arrow-c-data</artifactId>
         <version>${arrow.version}</version>
       </dependency>
+      <dependency>
+        <groupId>org.apache.arrow</groupId>
+        <artifactId>arrow-compression</artifactId>
+        <version>${arrow.version}</version>
+        <!-- Arrow's per-buffer IPC compression codecs. Their backing libraries ship with Spark,
+             so take them from there rather than bundling a second copy: commons-compress (LZ4
+             frame) and zstd-jni are both on every Spark version Comet supports. -->
+        <exclusions>
+          <exclusion>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-compress</artifactId>
+          </exclusion>
+          <exclusion>
+            <groupId>com.github.luben</groupId>
+            <artifactId>zstd-jni</artifactId>
+          </exclusion>
+          <exclusion>
+            <groupId>io.netty</groupId>
+            <artifactId>netty-common</artifactId>
+          </exclusion>
+          <exclusion>
+            <groupId>com.google.code.findbugs</groupId>
+            <artifactId>jsr305</artifactId>
+          </exclusion>
+        </exclusions>
+      </dependency>
 
       <!-- Parquet dependencies -->
       <dependency>

spark/pom.xml

@@ -60,6 +60,10 @@ under the License.
       <groupId>org.apache.arrow</groupId>
       <artifactId>arrow-vector</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.arrow</groupId>
+      <artifactId>arrow-compression</artifactId>
+    </dependency>
     <dependency>
       <groupId>org.scala-lang.modules</groupId>
       <artifactId>scala-collection-compat_${scala.binary.version}</artifactId>
@@ -630,6 +634,16 @@ under the License.
                   <shadedPattern>${comet.shade.packageName}.guava.thirdparty</shadedPattern>
                 </relocation>
               </relocations>
+              <transformers>
+                <!-- arrow-compression ships META-INF/services/org.apache.arrow.vector.compression.CompressionCodec$Factory.
+                     Copied verbatim, that file names Spark's own unshaded Arrow interface while
+                     pointing at a provider class that only exists here under the relocated name, so
+                     Spark's ServiceLoader lookup fails and takes CompressionCodec.Factory's static
+                     initializer down with it. This transformer relocates both the service file name
+                     and its contents. -->
+                <transformer
+                  implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer" />
+              </transformers>
             </configuration>
           </execution>
         </executions>

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -261,23 +261,50 @@ object CometConf extends ShimCometConf {
   val COMET_EXEC_IN_MEMORY_CACHE_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.exec.inMemoryCache.enabled")
       .category(CATEGORY_EXEC)
-      .doc(
-        "Whether to enable Comet native execution for in-memory cached tables. Its value at " +
-          "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " +
-          "which stores cached data in Arrow format. Because spark.sql.cache.serializer is a " +
-          "static config, the cached format is fixed for the application, and disabling this " +
-          "at runtime only sends cached scans back to Spark's execution path. Relations whose " +
-          "schema Comet's Arrow writer does not support are always cached in Spark's default " +
-          "format. Each cached column is stored as its own compressed Arrow IPC stream, so a " +
-          "scan decodes only the columns it projected. Reads that feed Spark operators rather " +
-          "than Comet ones still pay a row conversion the default format avoids, and can be " +
-          "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " +
-          "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " +
-          "SparkContext, otherwise caching fails as soon as a block is serialized, including " +
-          "the disk half of the default MEMORY_AND_DISK storage level.")
+      .doc("Whether to enable Comet native execution for in-memory cached tables. Its value at " +
+        "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " +
+        "which stores cached data in Arrow format. Because spark.sql.cache.serializer is a " +
+        "static config, the cached format is fixed for the application, and disabling this " +
+        "at runtime only sends cached scans back to Spark's execution path. Relations whose " +
+        "schema Comet's Arrow writer does not support are always cached in Spark's default " +
+        "format. Each cached batch is stored as one Arrow IPC record batch with per-buffer " +
+        "zstd compression, and a scan copies out only the buffers of the columns it projected, " +
+        "so the unselected ones are never decompressed. Reads that feed Spark operators rather " +
+        "than Comet ones still pay a row conversion the default format avoids, and can be " +
+        "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " +
+        "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " +
+        "SparkContext, otherwise caching fails as soon as a block is serialized, including " +
+        "the disk half of the default MEMORY_AND_DISK storage level.")
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC: ConfigEntry[String] =
+    conf("spark.comet.exec.inMemoryCache.compression.codec")
+      .category(CATEGORY_EXEC)
+      .doc(
+        "The Arrow IPC compression codec used when Comet's cache serializer writes cached " +
+          "data. Unlike spark.io.compression.codec, this compresses each Arrow buffer " +
+          "separately rather than the batch as a whole, which is what lets a projected scan " +
+          "decompress only the columns it selected. Set to none to store cached batches " +
+          "uncompressed, which is both slower to write and larger than zstd because the extra " +
+          "bytes cost more to move and store than compressing them costs. Only affects newly " +
+          "cached data; the codec a batch was written with is recorded in the batch itself and " +
+          "is what the read path uses. Arrow's lz4 is deliberately not offered: it is a " +
+          "pure-Java implementation, unrelated to the JNI-accelerated lz4 behind " +
+          "spark.io.compression.codec, and is orders of magnitude slower to write than zstd " +
+          "while also producing larger output.")
+      .stringConf
+      .checkValues(Set("none", "zstd"))
+      .createWithDefault("zstd")
+
+  val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL: ConfigEntry[Int] =
+    conf("spark.comet.exec.inMemoryCache.compression.zstd.level")
+      .category(CATEGORY_EXEC)
+      .doc("The compression level to use when Comet's cache serializer compresses cached data " +
+        "with zstd. Ignored for other codecs.")
+      .intConf
+      .createWithDefault(1)
+
   val COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] =
     conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.enabled")
       .category(CATEGORY_EXEC)