feat: add experimental native support for in-memory cache, disabled by default

.github/workflows/pr_build_linux.yml

@@ -339,6 +339,8 @@ jobs:
               org.apache.comet.exec.CometAggregateSuite
               org.apache.comet.exec.CometExec3_4PlusSuite
               org.apache.comet.exec.CometExecSuite
+              org.apache.comet.exec.CometInMemoryCacheSuite
+              org.apache.comet.exec.CometInMemoryCacheKryoSuite
               org.apache.comet.exec.CometGenerateExecSuite
               org.apache.comet.exec.CometWindowExecSuite
               org.apache.comet.exec.CometJoinSuite

.github/workflows/pr_build_macos.yml

@@ -155,6 +155,8 @@ jobs:
               org.apache.comet.exec.CometAggregateSuite
               org.apache.comet.exec.CometExec3_4PlusSuite
               org.apache.comet.exec.CometExecSuite
+              org.apache.comet.exec.CometInMemoryCacheSuite
+              org.apache.comet.exec.CometInMemoryCacheKryoSuite
               org.apache.comet.exec.CometGenerateExecSuite
               org.apache.comet.exec.CometWindowExecSuite
               org.apache.comet.exec.CometJoinSuite

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

@@ -258,6 +258,26 @@ object CometConf extends ShimCometConf {
   val COMET_EXEC_SAMPLE_ENABLED: ConfigEntry[Boolean] =
     createExecEnabledConfig("sample", defaultValue = true)
 
+  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.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] =
     conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.enabled")
       .category(CATEGORY_EXEC)

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

@@ -0,0 +1,58 @@
+/*
+ * 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.comet
+
+import org.apache.spark.serializer.KryoRegistrator
+import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer
+import org.apache.spark.sql.comet.util.Utils
+
+import com.esotericsoftware.kryo.Kryo
+
+/**
+ * Registers the classes Comet hands to Spark's serializer with Kryo.
+ *
+ * This is only needed when `spark.kryo.registrationRequired=true`, which makes Kryo reject any
+ * unregistered class rather than writing its name. Set it alongside Comet's own configuration:
+ *
+ * {{{
+ *   spark.serializer          org.apache.spark.serializer.KryoSerializer
+ *   spark.kryo.registrator    org.apache.comet.CometKryoRegistrator
+ * }}}
+ *
+ * `spark.kryo.registrator` has to be set before the `SparkContext` is created, because
+ * `KryoSerializer` reads it when `SparkEnv` builds it. That is earlier than `CometDriverPlugin`
+ * runs, so Comet cannot add this for you the way it can add `spark.sql.cache.serializer`;
+ * `CometDriverPlugin` logs a warning instead when the combination looks unsafe.
+ *
+ * Two payloads need it: the `Array[ChunkedByteBuffer]` a native broadcast broadcasts, and
+ * `CometCachedBatch`. The first applies whether or not the in-memory cache feature is enabled.
+ */
+class CometKryoRegistrator extends KryoRegistrator {
+  override def registerClasses(kryo: Kryo): Unit = {
+    CometKryoRegistrator.classes.foreach(kryo.register)
+  }
+}
+
+object CometKryoRegistrator {
+  val CLASS_NAME: String = classOf[CometKryoRegistrator].getName
+
+  def classes: Seq[Class[_]] =
+    Utils.arrowBytesKryoClasses ++ ArrowCachedBatchSerializer.kryoClasses
+}

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -29,11 +29,13 @@ import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 import org.apache.spark.sql.catalyst.util.sideBySide
 import org.apache.spark.sql.comet._
+import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec}
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
 import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec}
 import org.apache.spark.sql.execution.datasources.WriteFilesExec
 import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat
@@ -86,6 +88,7 @@ object CometExecRule {
       classOf[SortMergeJoinExec] -> CometSortMergeJoinExec,
       classOf[SortExec] -> CometSortExec,
       classOf[LocalTableScanExec] -> CometLocalTableScanExec,
+      classOf[InMemoryTableScanExec] -> CometInMemoryTableScanExec,
       classOf[SampleExec] -> CometSampleExec,
       classOf[WindowExec] -> CometWindowExec) ++
       // WindowGroupLimitExec exists only on Spark 3.5+; the shim returns None on 3.4.
@@ -295,6 +298,48 @@ case class CometExecRule(session: SparkSession)
       case op if isCometScan(op) =>
         convertToComet(op, CometScanWrapper).getOrElse(op)
 
+      case scan: InMemoryTableScanExec =>
+        val serializer = scan.relation.cacheBuilder.serializer
+        val usesCometCacheSerializer = serializer.isInstanceOf[ArrowCachedBatchSerializer]
+        // The serializer only stores Comet's Arrow format for schemas it supports and delegates
+        // everything else to Spark's default cache format, which the native scan cannot read.
+        val cometCacheFormat = usesCometCacheSerializer &&
+          ArrowCachedBatchSerializer.supportsSchema(scan.relation.output)
+        val nativeCacheEnabled = CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf)
+
+        if (nativeCacheEnabled && cometCacheFormat) {
+          convertToComet(scan, CometInMemoryTableScanExec).getOrElse(scan)
+        } else {
+          // The native cache scan is not available for this relation. Record why, then take the
+          // same SparkToColumnar fallback that any other unsupported operator would take, so
+          // that turning the feature on is never worse for a scan than leaving it off.
+          if (nativeCacheEnabled && !usesCometCacheSerializer) {
+            withFallbackReason(
+              scan,
+              s"Comet in-memory cache requires ${classOf[ArrowCachedBatchSerializer].getName} " +
+                s"but this relation was cached with ${serializer.getClass.getName}")
+          } else if (nativeCacheEnabled) {
+            val unsupported = scan.relation.output
+              .filterNot(a => ArrowCachedBatchSerializer.supportsType(a.dataType))
+              .map(a => s"${a.name}: ${a.dataType.simpleString}")
+            withFallbackReason(
+              scan,
+              "Comet in-memory cache does not support the type of these cached columns, so the " +
+                s"relation was cached in Spark's default format: ${unsupported.mkString(", ")}")
+          } else if (usesCometCacheSerializer) {
+            withFallbackReason(
+              scan,
+              "Native support for operator InMemoryTableScanExec is disabled. " +
+                s"Set ${CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key}=true to enable it.")
+          }
+
+          if (shouldApplySparkToColumnar(conf, scan)) {
+            convertToComet(scan, CometSparkToColumnarExec).getOrElse(scan)
+          } else {
+            scan
+          }
+        }
+
       case op if shouldApplySparkToColumnar(conf, op) =>
         convertToComet(op, CometSparkToColumnarExec).getOrElse(op)
 

spark/src/main/scala/org/apache/comet/vector/StreamReader.scala

@@ -21,6 +21,8 @@ package org.apache.comet.vector
 
 import java.nio.channels.ReadableByteChannel
 
+import scala.util.control.NonFatal
+
 import org.apache.arrow.vector.VectorSchemaRoot
 import org.apache.arrow.vector.ipc.{ArrowStreamReader, ReadChannel}
 import org.apache.arrow.vector.ipc.message.MessageChannelReader
@@ -35,7 +37,17 @@ case class StreamReader(channel: ReadableByteChannel, source: String) extends Au
   private val channelReader =
     new MessageChannelReader(new ReadChannel(channel), CometArrowAllocator)
   private var arrowReader = new ArrowStreamReader(channelReader, CometArrowAllocator)
-  private var root = arrowReader.getVectorSchemaRoot
+
+  // Reading the schema allocates the root's vectors, so it can fail with buffers already taken.
+  // No caller holds this reader until its constructor returns, so close it here or nothing will.
+  private var root =
+    try arrowReader.getVectorSchemaRoot
+    catch {
+      case NonFatal(e) =>
+        try arrowReader.close()
+        catch { case NonFatal(closeError) => e.addSuppressed(closeError) }
+        throw e
+    }
 
   def nextBatch(): Option[ColumnarBatch] = {
     if (arrowReader.loadNextBatch()) {

spark/src/main/scala/org/apache/spark/Plugins.scala

@@ -29,7 +29,9 @@ import org.apache.spark.internal.config.{EXECUTOR_MEMORY, EXECUTOR_MEMORY_OVERHE
 import org.apache.spark.sql.internal.StaticSQLConf
 
 import org.apache.comet.{COMET_VERSION, CometSparkSessionExtensions, NativeBase}
+import org.apache.comet.CometConf
 import org.apache.comet.CometConf.{COMET_METRICS_ENABLED, COMET_ONHEAP_ENABLED}
+import org.apache.comet.CometKryoRegistrator
 import org.apache.comet.annotation.Public
 
 /**
@@ -61,6 +63,11 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
       return Collections.emptyMap[String, String]
     }
 
+    val extraConfs = new ju.HashMap[String, String]()
+
+    CometDriverPlugin.maybeSetCacheSerializer(sc.conf, extraConfs)
+    CometDriverPlugin.warnIfKryoRegistratorMissing(sc.conf)
+
     // register CometSparkSessionExtensions if it isn't already registered
     CometDriverPlugin.registerCometSessionExtension(sc.conf)
 
@@ -94,7 +101,7 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
       logInfo("Comet is running in unified memory mode and sharing off-heap memory with Spark")
     }
 
-    Collections.emptyMap[String, String]
+    extraConfs
   }
 
   override def receive(message: Any): AnyRef = super.receive(message)
@@ -117,6 +124,58 @@ object CometDriverPlugin extends Logging {
   /** Spark config key under which the loaded Comet version is exposed at runtime. */
   val COMET_VERSION_CONFIG = "spark.comet.version"
 
+  // Use Comet's cache serializer only for the native in-memory cache path.
+  // If the application already set spark.sql.cache.serializer, leave that value
+  // unchanged so Comet does not replace a user-selected cache format.
+  private[apache] def maybeSetCacheSerializer(
+      conf: SparkConf,
+      extraConfs: ju.HashMap[String, String]): Unit = {
+    if (conf.getBoolean(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, false)) {
+      val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key
+      val serializerValue =
+        "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer"
+      val defaultSerializer = StaticSQLConf.SPARK_CACHE_SERIALIZER.defaultValueString
+      val currentSerializer = conf.get(serializerKey, defaultSerializer)
+
+      if (currentSerializer == defaultSerializer) {
+        extraConfs.put(serializerKey, serializerValue)
+        conf.set(serializerKey, serializerValue)
+        logInfo(s"Auto-set $serializerKey=$serializerValue")
+      } else {
+        logInfo(s"Not overriding user-provided $serializerKey=$currentSerializer")
+      }
+    }
+  }
+
+  // Comet hands Spark's serializer classes that Kryo has not been told about, so with
+  // spark.kryo.registrationRequired=true it rejects them with "Class is not registered", which
+  // names neither Comet nor the operation that failed. Two paths reach it: a native broadcast,
+  // which broadcasts an Array[ChunkedByteBuffer], and any cached block Spark serializes -- the
+  // disk half of MEMORY_AND_DISK, the _SER levels, replication, a cross-executor fetch.
+  // CometKryoRegistrator covers both, but spark.kryo.registrator is read when SparkEnv builds the
+  // serializer, before any plugin runs, so it cannot be set from here. Say so while the
+  // application is still starting up rather than leaving the user to attribute the failure later.
+  private[apache] def warnIfKryoRegistratorMissing(conf: SparkConf): Unit = {
+    val usingKryo =
+      conf.get("spark.serializer", "") == "org.apache.spark.serializer.KryoSerializer"
+    val registrationRequired = conf.getBoolean("spark.kryo.registrationRequired", false)
+    val registered = conf
+      .get("spark.kryo.registrator", "")
+      .split(',')
+      .map(_.trim)
+      .contains(CometKryoRegistrator.CLASS_NAME)
+
+    if (usingKryo && registrationRequired && !registered) {
+      logWarning(
+        "spark.kryo.registrationRequired=true but spark.kryo.registrator does not include " +
+          s"${CometKryoRegistrator.CLASS_NAME}. Comet's native broadcast and its in-memory " +
+          "cache format will fail with Kryo's \"Class is not registered\" as soon as their " +
+          "payloads are serialized. Add " +
+          s"spark.kryo.registrator=${CometKryoRegistrator.CLASS_NAME} before creating the " +
+          "SparkContext; it cannot be set later.")
+    }
+  }
+
   def registerCometMetrics(sc: SparkContext): Unit = {
     if (sc.getConf.getBoolean(
         COMET_METRICS_ENABLED.key,