From d87014acb9c5134fcee7fd8016fee5d44e17d3aa Mon Sep 17 00:00:00 2001 From: pchintar <89355405+pchintar@users.noreply.github.com> Date: Thu, 4 Jun 2026 18:52:39 +0530 Subject: [PATCH 1/8] Add Comet-native in-memory cache scan support --- .../scala/org/apache/comet/CometConf.scala | 10 + .../apache/comet/rules/CometExecRule.scala | 44 +++ .../main/scala/org/apache/spark/Plugins.scala | 16 +- .../comet/CometInMemoryTableScanExec.scala | 124 +++++++++ .../arrow/ArrowCachedBatchSerializer.scala | 253 +++++++++++++++++ .../apache/spark/sql/comet/operators.scala | 3 +- .../comet/exec/CometInMemoryCacheSuite.scala | 255 ++++++++++++++++++ 7 files changed, 703 insertions(+), 2 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 78ea0f0168..dee67a044f 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -275,6 +275,16 @@ object CometConf extends ShimCometConf { val COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("localTableScan", defaultValue = false) + 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. " + + "When disabled or when spark.comet.enabled=false, Spark's default cache " + + "serializer and execution path will be used.") + .booleanConf + .createWithDefault(false) + val COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index d116d2f407..114304492e 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/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.CometInMemoryTableScanExec 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 @@ -85,6 +87,7 @@ object CometExecRule { classOf[SortMergeJoinExec] -> CometSortMergeJoinExec, classOf[SortExec] -> CometSortExec, classOf[LocalTableScanExec] -> CometLocalTableScanExec, + classOf[InMemoryTableScanExec] -> CometInMemoryTableScanExec, classOf[WindowExec] -> CometWindowExec) /** @@ -282,6 +285,47 @@ case class CometExecRule(session: SparkSession) case op if isCometScan(op) => convertToComet(op, CometScanWrapper).getOrElse(op) + case scan: InMemoryTableScanExec => + val cachedBuffers = scan.relation.cacheBuilder.cachedColumnBuffers + val firstBatchOpt = cachedBuffers.take(1).headOption + val expectedBatchClass = + "org.apache.spark.sql.comet.execution.arrow.CometCachedBatch" + + if (CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf)) { + firstBatchOpt match { + case Some(firstBatch) if firstBatch.getClass.getName == expectedBatchClass => + convertToComet(scan, CometInMemoryTableScanExec).getOrElse(scan) + + case Some(firstBatch) => + withFallbackReason( + scan, + s"Comet in-memory cache requires $expectedBatchClass, " + + s"but found ${firstBatch.getClass.getName}") + scan + + case None => + withFallbackReason( + scan, + "Comet in-memory cache rewrite skipped because cached buffer is empty") + scan + } + } else { + firstBatchOpt match { + case Some(firstBatch) if firstBatch.getClass.getName == expectedBatchClass => + withFallbackReason( + scan, + s"Native support for operator InMemoryTableScanExec is disabled. " + + s"Set ${CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key}=true to enable it.") + case _ => + } + + if (shouldApplySparkToColumnar(conf, scan)) { + convertToComet(scan, CometSparkToColumnarExec).getOrElse(scan) + } else { + scan + } + } + case op if shouldApplySparkToColumnar(conf, op) => convertToComet(op, CometSparkToColumnarExec).getOrElse(op) diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 7290ab436a..82235d380c 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -28,6 +28,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.internal.config.{EXECUTOR_MEMORY, EXECUTOR_MEMORY_OVERHEAD, EXECUTOR_MEMORY_OVERHEAD_FACTOR} import org.apache.spark.sql.internal.StaticSQLConf +import org.apache.comet.CometConf import org.apache.comet.CometConf.{COMET_METRICS_ENABLED, COMET_ONHEAP_ENABLED} import org.apache.comet.CometSparkSessionExtensions @@ -54,6 +55,19 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl return Collections.emptyMap[String, String] } + val extraConfs = new ju.HashMap[String, String]() + + // Always register Comet's cache serializer class. + // The serializer itself decides at runtime whether to use Comet cache format + // or delegate to DefaultCachedBatchSerializer based on + // spark.comet.exec.inMemoryCache.enabled. + val serializerKey = "spark.sql.cache.serializer" + val serializerValue = + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer" + extraConfs.put(serializerKey, serializerValue) + sc.conf.set(serializerKey, serializerValue) + logInfo(s"Auto-set $serializerKey=$serializerValue") + // register CometSparkSessionExtensions if it isn't already registered CometDriverPlugin.registerCometSessionExtension(sc.conf) @@ -87,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) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala new file mode 100644 index 0000000000..cb12f2a3ac --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -0,0 +1,124 @@ +/* + * 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 + +import scala.collection.JavaConverters._ + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer} +import org.apache.spark.sql.execution.LeafExecNode +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.CometConf +import org.apache.comet.serde.CometOperatorSerde +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.QueryPlanSerde.serializeDataType + +/** + * Reads Spark cached table data when the cache was written by Comet's cache serializer. + * + * Spark stores cached data through `CachedBatchSerializer`. This node keeps the scan inside Comet + * by asking the serializer to decode cached batches directly into `ColumnarBatch` output, + * avoiding the extra Spark columnar-to-Comet columnar conversion used by the default path. + * + * `relationOutput` is the full schema stored in the cache. `scanOutput` is the subset requested + * by this scan after pruning. + */ +case class CometInMemoryTableScanExec( + originalPlan: InMemoryTableScanExec, + serializer: CachedBatchSerializer, + cachedBuffers: RDD[CachedBatch], + relationOutput: Seq[Attribute], + scanOutput: Seq[Attribute]) + extends CometExec + with LeafExecNode { + + override lazy val metrics: Map[String, SQLMetric] = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + + override def output: Seq[Attribute] = originalPlan.output + + // Use the serializer's vector types because the cached batch layout is owned by the serializer. + override def vectorTypes: Option[Seq[String]] = + serializer.vectorTypes(scanOutput, conf) + + // Decode only the requested columns from the cached batches and update scan output metrics. + override def doExecuteColumnar(): RDD[ColumnarBatch] = { + val numOutputRows = longMetric("numOutputRows") + + serializer + .convertCachedBatchToColumnarBatch(cachedBuffers, relationOutput, scanOutput, conf) + .map { cb => + numOutputRows += cb.numRows() + cb + } + } +} + +object CometInMemoryTableScanExec extends CometOperatorSerde[InMemoryTableScanExec] { + + override def enabledConfig: Option[org.apache.comet.ConfigEntry[Boolean]] = + Some(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED) + + override def convert( + op: InMemoryTableScanExec, + builder: OperatorOuterClass.Operator.Builder, + childOp: Operator*): Option[Operator] = { + + // Empty-output scans still need a schema for native planning, so fall back to the cache schema. + val actualOutput = + if (op.output.nonEmpty) op.output + else op.relation.output + + val scanTypes = actualOutput.flatMap(attr => serializeDataType(attr.dataType)) + + val scanBuilder = OperatorOuterClass.Scan + .newBuilder() + .setSource(op.getClass.getSimpleName) + .addAllFields(scanTypes.asJava) + // Cached batches are decoded on the JVM side; the native scan only receives Spark batches. + .setArrowFfiSafe(false) + + Some(builder.setScan(scanBuilder).build()) + } + + // Reuse Spark's InMemoryRelation metadata so cache materialization, pruning, and storage + // behavior remain controlled by Spark's cache manager. + override def createExec(nativeOp: Operator, op: InMemoryTableScanExec): CometNativeExec = { + val relation = op.relation + + val actualOutput = + if (op.output.nonEmpty) op.output + else relation.output + + CometScanWrapper( + nativeOp, + CometInMemoryTableScanExec( + op, + relation.cacheBuilder.serializer, + relation.cacheBuilder.cachedColumnBuffers, + relation.output, + actualOutput)) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala new file mode 100644 index 0000000000..5d54930141 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -0,0 +1,253 @@ +/* + * 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 scala.collection.JavaConverters._ + +import org.apache.spark.TaskContext +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, UnsafeProjection} +import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer} +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.execution.columnar.{DefaultCachedBatch, DefaultCachedBatchSerializer} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{StructField, StructType} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.storage.StorageLevel +import org.apache.spark.util.io.ChunkedByteBuffer + +import org.apache.comet.CometConf + +/** + * Cached batch format used when Comet writes Spark in-memory cache data. + * + * `bytes` contains compressed Arrow stream data produced by `Utils.serializeBatches`. The cache + * manager still owns storage and eviction; this class only changes the cached payload. + */ +private case class CometCachedBatch( + override val numRows: Int, + override val sizeInBytes: Long, + stats: InternalRow, + bytes: ChunkedByteBuffer) + extends CachedBatch + +/** + * Cache serializer that stores Comet-compatible Arrow batches in Spark's in-memory cache. + * + * When Comet cache support is disabled, row-based cache writes and default cache reads are + * delegated to Spark's `DefaultCachedBatchSerializer`. + */ +class ArrowCachedBatchSerializer extends CachedBatchSerializer { + + private val fallback = new DefaultCachedBatchSerializer() + + // Cache writes use Comet format only when both Comet and the in-memory cache scan are enabled. + private def enabled(conf: SQLConf): Boolean = { + CometConf.COMET_ENABLED.get(conf) && + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf) + } + + // Row-to-Arrow conversion needs a StructType, while cache APIs pass attributes. + private def toStructType(schema: Seq[Attribute]): StructType = { + StructType(schema.map { attr => + StructField(attr.name, attr.dataType, attr.nullable, attr.metadata) + }) + } + + override def supportsColumnarInput(schema: Seq[Attribute]): Boolean = { + val activeConf = SQLConf.get + activeConf != null && enabled(activeConf) + } + override def supportsColumnarOutput(schema: StructType): Boolean = true + + // Columnar Comet output is stored as compressed Arrow stream bytes. + override def convertColumnarBatchToCachedBatch( + input: RDD[ColumnarBatch], + schema: Seq[Attribute], + storageLevel: StorageLevel, + conf: SQLConf): RDD[CachedBatch] = { + + input.mapPartitions { batches => + Utils.serializeBatches(batches).map { case (rows, buffer) => + CometCachedBatch( + numRows = rows.toInt, + sizeInBytes = buffer.size, + stats = InternalRow.empty, + bytes = buffer) + } + } + } + + override def convertCachedBatchToColumnarBatch( + input: RDD[CachedBatch], + cacheAttributes: Seq[Attribute], + selectedAttributes: Seq[Attribute], + conf: SQLConf): RDD[ColumnarBatch] = { + + // Resolve requested columns by exprId, not by name, because aliases may reuse names. + val selectedIndices = + if (selectedAttributes.isEmpty) { + cacheAttributes.indices.toArray + } else { + val byExprId = cacheAttributes.zipWithIndex.map { case (attr, idx) => + attr.exprId -> idx + }.toMap + + selectedAttributes.map { attr => + byExprId.getOrElse( + attr.exprId, + throw new IllegalStateException( + s"Could not resolve selected attribute ${attr.name} from cache attributes")) + }.toArray + } + + val batchTypes = input.map(_.getClass.getName).distinct().collect() + + if (batchTypes.isEmpty) { + input.sparkContext.emptyRDD[ColumnarBatch] + } else if (batchTypes.length > 1) { + throw new IllegalStateException( + s"Mixed cached batch types are not supported: ${batchTypes.mkString(", ")}") + } else if (batchTypes.head == classOf[CometCachedBatch].getName) { + input.mapPartitions { it => + it.flatMap { + case cb: CometCachedBatch => + Utils.decodeBatches(cb.bytes, "CometCache").map { batch => + if (selectedIndices.length == batch.numCols()) { + batch + } else { + val cols = + selectedIndices.map(i => batch.column(i).asInstanceOf[ColumnVector]) + new ColumnarBatch(cols, batch.numRows()) + } + } + + case other => + throw new IllegalStateException( + s"Expected CometCachedBatch, got ${other.getClass.getName}") + } + } + } else if (batchTypes.head == classOf[DefaultCachedBatch].getName) { + fallback.convertCachedBatchToColumnarBatch(input, cacheAttributes, selectedAttributes, conf) + } else { + throw new IllegalStateException(s"Unsupported cached batch type: ${batchTypes.head}") + } + } + + // Row input can still be cached in Comet format by converting rows to Arrow batches first. + override def convertInternalRowToCachedBatch( + input: RDD[InternalRow], + schema: Seq[Attribute], + storageLevel: StorageLevel, + conf: SQLConf): RDD[CachedBatch] = { + + if (!enabled(conf)) { + fallback.convertInternalRowToCachedBatch(input, schema, storageLevel, conf) + } else { + val batchSize = conf.columnBatchSize + val sessionTz = conf.sessionLocalTimeZone + + input.mapPartitions { rows => + val iter = CometArrowConverters.rowToArrowBatchIter( + rows, + toStructType(schema), + batchSize, + sessionTz, + TaskContext.get()) + + Utils.serializeBatches(iter).map { case (rows, buffer) => + CometCachedBatch( + numRows = rows.toInt, + sizeInBytes = buffer.size, + stats = InternalRow.empty, + bytes = buffer) + } + } + } + } + + override def convertCachedBatchToInternalRow( + input: RDD[CachedBatch], + cacheAttributes: Seq[Attribute], + selectedAttributes: Seq[Attribute], + conf: SQLConf): RDD[InternalRow] = { + + // Resolve requested columns by exprId, not by name, because aliases may reuse names. + val selectedIndices = + if (selectedAttributes.isEmpty) { + cacheAttributes.indices.toArray + } else { + val byExprId = cacheAttributes.zipWithIndex.map { case (attr, idx) => + attr.exprId -> idx + }.toMap + + selectedAttributes.map { attr => + byExprId.getOrElse( + attr.exprId, + throw new IllegalStateException( + s"Could not resolve selected attribute ${attr.name} from cache attributes")) + }.toArray + } + + val batchTypes = input.map(_.getClass.getName).distinct().collect() + + if (batchTypes.isEmpty) { + input.sparkContext.emptyRDD[InternalRow] + } else if (batchTypes.length > 1) { + throw new IllegalStateException( + s"Mixed cached batch types are not supported: ${batchTypes.mkString(", ")}") + } else if (batchTypes.head == classOf[DefaultCachedBatch].getName) { + fallback.convertCachedBatchToInternalRow(input, cacheAttributes, selectedAttributes, conf) + } else if (batchTypes.head == classOf[CometCachedBatch].getName) { + input.mapPartitions { it => + it.flatMap { + case cb: CometCachedBatch => + Utils.decodeBatches(cb.bytes, "CometCache").flatMap { batch => + val projectedBatch = + if (selectedIndices.length == batch.numCols()) { + batch + } else { + val cols = + selectedIndices.map(i => batch.column(i).asInstanceOf[ColumnVector]) + new ColumnarBatch(cols, batch.numRows()) + } + + // Spark's row collect path expects UnsafeRow, not ColumnarBatchRow wrappers. + val toUnsafe = UnsafeProjection.create(selectedAttributes, selectedAttributes) + projectedBatch.rowIterator().asScala.map(row => toUnsafe(row).copy()) + } + + case other => + throw new IllegalStateException( + s"Expected CometCachedBatch, got ${other.getClass.getName}") + } + } + } else { + throw new IllegalStateException(s"Unsupported cached batch type: ${batchTypes.head}") + } + } + + override def buildFilter( + predicates: Seq[Expression], + cachedAttributes: Seq[Attribute]): (Int, Iterator[CachedBatch]) => Iterator[CachedBatch] = { + (partitionIndex: Int, it: Iterator[CachedBatch]) => it + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 8cbf7c9189..eb506ddc6e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -612,7 +612,8 @@ abstract class CometNativeExec extends CometExec { _: AQEShuffleReadExec | _: CometShuffleExchangeExec | _: CometUnionExec | _: CometTakeOrderedAndProjectExec | _: CometCoalesceExec | _: ReusedExchangeExec | _: CometBroadcastExchangeExec | _: BroadcastQueryStageExec | - _: CometSparkToColumnarExec | _: CometLocalTableScanExec => + _: CometSparkToColumnarExec | _: CometLocalTableScanExec | + _: CometInMemoryTableScanExec => func(plan) case _: CometPlan => // Other Comet operators, continue to traverse the tree. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala new file mode 100644 index 0000000000..11174d9bde --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -0,0 +1,255 @@ +/* + * 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.exec + +import org.apache.spark.SparkConf +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +class CometInMemoryCacheSuite extends CometTestBase { + override protected def sparkConf: SparkConf = { + val conf = new SparkConf() + conf.set("spark.driver.memory", "1G") + conf.set("spark.executor.memory", "1G") + conf.set("spark.executor.memoryOverhead", "2G") + conf.set("spark.plugins", "org.apache.spark.CometPlugin") + conf.set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + conf.set("spark.comet.enabled", "true") + conf.set("spark.comet.exec.enabled", "true") + conf.set("spark.comet.exec.onHeap.enabled", "true") + conf.set("spark.comet.metrics.enabled", "true") + conf.set( + "spark.sql.cache.serializer", + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + conf + } + + private def cachedBatchTypes(table: String): Array[String] = { + val ds = spark.table(table).asInstanceOf[org.apache.spark.sql.classic.Dataset[_]] + val cached = spark.sharedState.cacheManager.lookupCachedData(ds).get + cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers + .map(_.getClass.getName) + .distinct() + .collect() + } + + test("CometInMemoryTableScan over CometCachedBatch") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + spark + .range(1000) + .selectExpr("id as key", "id % 8 as value") + .createOrReplaceTempView("abc") + + spark.catalog.cacheTable("abc") + spark.table("abc").count() + + assert( + cachedBatchTypes("abc").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val df = spark.sql("SELECT key, count(*) FROM abc GROUP BY key") + checkSparkAnswer(df) + + val plan = df.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + assert(!plan.contains("CometSparkColumnarToColumnar")) + + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache disabled keeps SparkToColumnar fallback path") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + spark + .range(1000) + .selectExpr("id as key", "id % 8 as value") + .createOrReplaceTempView("comet_cache_disabled") + + spark.catalog.cacheTable("comet_cache_disabled") + spark.table("comet_cache_disabled").count() + + assert( + cachedBatchTypes("comet_cache_disabled").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + } + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + val df = spark.sql("SELECT key, count(*) FROM comet_cache_disabled GROUP BY key") + checkSparkAnswer(df) + + val plan = df.queryExecution.executedPlan.toString() + assert(!plan.contains("CometInMemoryTableScan")) + assert(plan.contains("CometSparkColumnarToColumnar")) + + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache handles multi-partition cache") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + val multiPartition = + spark.range(0, 1000, 1, 5).toDF("id").cache() + multiPartition.createOrReplaceTempView("multi_partition_cache") + multiPartition.count() + + assert( + cachedBatchTypes("multi_partition_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val grouped = spark.sql(""" + SELECT id % 100, count(*) + FROM multi_partition_cache + GROUP BY id % 100 + """) + checkSparkAnswer(grouped) + + val groupedPlan = grouped.queryExecution.executedPlan.toString() + assert(groupedPlan.contains("CometInMemoryTableScan")) + + multiPartition.unpersist() + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache handles empty cache") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + val empty = spark.range(0).toDF("id").cache() + empty.createOrReplaceTempView("empty_cache") + empty.count() + + val emptyDf = spark.sql("SELECT * FROM empty_cache") + checkSparkAnswer(emptyDf) + + val emptyPlan = emptyDf.queryExecution.executedPlan.toString() + assert(!emptyPlan.contains("CometInMemoryTableScan")) + + empty.unpersist() + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache supports projection-only read") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + spark + .range(1000) + .selectExpr("id as key", "id % 8 as value", "id + 1 as key_plus_1") + .createOrReplaceTempView("project_cache") + + spark.catalog.cacheTable("project_cache") + spark.table("project_cache").count() + + assert( + cachedBatchTypes("project_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val df = spark.sql("SELECT key FROM project_cache") + checkSparkAnswer(df) + + val plan = df.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + assert(plan.contains("CometNativeColumnarToRow")) + + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache supports shuffle after cache read") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + spark + .range(1000) + .selectExpr("id as key", "id % 100 as group") + .createOrReplaceTempView("shuffle_cache") + + spark.catalog.cacheTable("shuffle_cache") + spark.table("shuffle_cache").count() + + assert( + cachedBatchTypes("shuffle_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val df = spark.sql("SELECT group, count(*) FROM shuffle_cache GROUP BY group") + checkSparkAnswer(df) + + val plan = df.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + assert(plan.contains("CometHashAggregate")) + + spark.catalog.clearCache() + } + } +} From 86a97cc3fe1365926d758e3396602579772d355d Mon Sep 17 00:00:00 2001 From: pchintar <89355405+pchintar@users.noreply.github.com> Date: Thu, 4 Jun 2026 18:52:39 +0530 Subject: [PATCH 2/8] Add Comet-native in-memory cache scan support --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../apache/comet/rules/CometExecRule.scala | 43 +-- .../main/scala/org/apache/spark/Plugins.scala | 34 +- .../comet/CometInMemoryTableScanExec.scala | 14 +- .../arrow/ArrowCachedBatchSerializer.scala | 350 +++++++++++------- .../comet/exec/CometInMemoryCacheSuite.scala | 253 ++++++++++++- .../CometInMemoryCacheBenchmark.scala | 152 ++++++++ 8 files changed, 680 insertions(+), 168 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 3fbe052aff..78ba754a70 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -335,6 +335,7 @@ 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.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index dbdd325848..eb451cd8a0 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -151,6 +151,7 @@ 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.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 114304492e..cc80304165 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -286,37 +286,26 @@ case class CometExecRule(session: SparkSession) convertToComet(op, CometScanWrapper).getOrElse(op) case scan: InMemoryTableScanExec => - val cachedBuffers = scan.relation.cacheBuilder.cachedColumnBuffers - val firstBatchOpt = cachedBuffers.take(1).headOption - val expectedBatchClass = - "org.apache.spark.sql.comet.execution.arrow.CometCachedBatch" + val usesCometCacheSerializer = + scan.relation.cacheBuilder.serializer + .isInstanceOf[org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer] if (CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf)) { - firstBatchOpt match { - case Some(firstBatch) if firstBatch.getClass.getName == expectedBatchClass => - convertToComet(scan, CometInMemoryTableScanExec).getOrElse(scan) - - case Some(firstBatch) => - withFallbackReason( - scan, - s"Comet in-memory cache requires $expectedBatchClass, " + - s"but found ${firstBatch.getClass.getName}") - scan - - case None => - withFallbackReason( - scan, - "Comet in-memory cache rewrite skipped because cached buffer is empty") - scan + if (usesCometCacheSerializer) { + convertToComet(scan, CometInMemoryTableScanExec).getOrElse(scan) + } else { + withFallbackReason( + scan, + "Comet in-memory cache requires ArrowCachedBatchSerializer, " + + s"but found ${scan.relation.cacheBuilder.serializer.getClass.getName}") + scan } } else { - firstBatchOpt match { - case Some(firstBatch) if firstBatch.getClass.getName == expectedBatchClass => - withFallbackReason( - scan, - s"Native support for operator InMemoryTableScanExec is disabled. " + - s"Set ${CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key}=true to enable it.") - case _ => + if (usesCometCacheSerializer) { + withFallbackReason( + scan, + s"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)) { diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 82235d380c..b849a481d9 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -57,16 +57,7 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl val extraConfs = new ju.HashMap[String, String]() - // Always register Comet's cache serializer class. - // The serializer itself decides at runtime whether to use Comet cache format - // or delegate to DefaultCachedBatchSerializer based on - // spark.comet.exec.inMemoryCache.enabled. - val serializerKey = "spark.sql.cache.serializer" - val serializerValue = - "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer" - extraConfs.put(serializerKey, serializerValue) - sc.conf.set(serializerKey, serializerValue) - logInfo(s"Auto-set $serializerKey=$serializerValue") + CometDriverPlugin.maybeSetCacheSerializer(sc.conf, extraConfs) // register CometSparkSessionExtensions if it isn't already registered CometDriverPlugin.registerCometSessionExtension(sc.conf) @@ -118,6 +109,29 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl } object CometDriverPlugin extends Logging { + // 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") + } + } + } + def registerCometMetrics(sc: SparkContext): Unit = { if (sc.getConf.getBoolean( COMET_METRICS_ENABLED.key, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala index cb12f2a3ac..3da81d1f54 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -63,12 +63,22 @@ case class CometInMemoryTableScanExec( override def vectorTypes: Option[Seq[String]] = serializer.vectorTypes(scanOutput, conf) - // Decode only the requested columns from the cached batches and update scan output metrics. + // Apply Spark's cache batch filter before decoding. Spark's InMemoryTableScanExec does this in + // filteredCachedBatches(), but that method is private. Reusing the serializer's buildFilter here + // keeps Comet on the same stats-based pruning path instead of decoding every cached batch. override def doExecuteColumnar(): RDD[ColumnarBatch] = { val numOutputRows = longMetric("numOutputRows") + val filteredBuffers = + if (originalPlan.predicates.nonEmpty) { + val filter = serializer.buildFilter(originalPlan.predicates, relationOutput) + cachedBuffers.mapPartitionsWithIndex(filter) + } else { + cachedBuffers + } + serializer - .convertCachedBatchToColumnarBatch(cachedBuffers, relationOutput, scanOutput, conf) + .convertCachedBatchToColumnarBatch(filteredBuffers, relationOutput, scanOutput, conf) .map { cb => numOutputRows += cb.numRows() cb diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 5d54930141..bfae724c59 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -24,14 +24,17 @@ import scala.collection.JavaConverters._ import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, UnsafeProjection} -import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer} +import org.apache.spark.sql.catalyst.expressions.{Attribute, GenericInternalRow, UnsafeProjection} +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch, SimpleMetricsCachedBatchSerializer} import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.columnar.{DefaultCachedBatch, DefaultCachedBatchSerializer} +import org.apache.spark.sql.execution.columnar.{ColumnAccessor, DefaultCachedBatch, DefaultCachedBatchSerializer} +import org.apache.spark.sql.execution.vectorized.{OnHeapColumnVector, WritableColumnVector} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{StructField, StructType} +import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.storage.StorageLevel +import org.apache.spark.unsafe.types.{ByteArray, UTF8String} import org.apache.spark.util.io.ChunkedByteBuffer import org.apache.comet.CometConf @@ -45,17 +48,19 @@ import org.apache.comet.CometConf private case class CometCachedBatch( override val numRows: Int, override val sizeInBytes: Long, - stats: InternalRow, + override val stats: InternalRow, bytes: ChunkedByteBuffer) - extends CachedBatch + extends SimpleMetricsCachedBatch /** * Cache serializer that stores Comet-compatible Arrow batches in Spark's in-memory cache. * - * When Comet cache support is disabled, row-based cache writes and default cache reads are - * delegated to Spark's `DefaultCachedBatchSerializer`. + * Writes use Comet's Arrow cache format only when Comet and the native in-memory cache path are + * enabled. Reads of CometCachedBatch are still supported even if the native scan is disabled + * later, because Spark may then read the same cached data through the SparkToColumnar fallback + * path. */ -class ArrowCachedBatchSerializer extends CachedBatchSerializer { +class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { private val fallback = new DefaultCachedBatchSerializer() @@ -72,10 +77,164 @@ class ArrowCachedBatchSerializer extends CachedBatchSerializer { }) } + // Build the statistics row expected by SimpleMetricsCachedBatchSerializer. + // For each cached column Spark expects five values in this order: + // lower bound, upper bound, null count, row count, and size in bytes. + private def computeStats(batch: ColumnarBatch, attrs: Seq[Attribute]): InternalRow = { + val numCols = attrs.length + val lower = new Array[Any](numCols) + val upper = new Array[Any](numCols) + val nulls = Array.fill[Int](numCols)(0) + val numRows = batch.numRows() + + var c = 0 + while (c < numCols) { + val dt = attrs(c).dataType + val col = batch.column(c) + var r = 0 + while (r < numRows) { + if (col.isNullAt(r)) { + nulls(c) += 1 + } else if (tracksBounds(dt)) { + val value = readValue(col, dt, r) + if (lower(c) == null || compare(dt, value, lower(c)) < 0) { + lower(c) = value + } + if (upper(c) == null || compare(dt, value, upper(c)) > 0) { + upper(c) = value + } + } + r += 1 + } + c += 1 + } + + val values = new Array[Any](numCols * 5) + c = 0 + while (c < numCols) { + val base = c * 5 + values(base) = lower(c) + values(base + 1) = upper(c) + values(base + 2) = nulls(c) + values(base + 3) = numRows + // Spark reserves the fifth field for per-column size. Comet stores the whole + // Arrow stream as one compressed buffer, so per-column size is not tracked here. + // Cache pruning uses bounds/null-count/row-count, not this size field. + values(base + 4) = 0L + c += 1 + } + + new GenericInternalRow(values) + } + + // Spark can prune cache batches only for types whose bounds can be compared. + // Other types still report null count and row count but leave bounds as null. + private def tracksBounds(dt: DataType): Boolean = dt match { + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + _: DecimalType | StringType | DateType | TimestampType | TimestampNTZType => + true + case _ => false + } + + // Read a non-null value from a ColumnVector using Spark's internal value type + // for the corresponding DataType. + private def readValue(col: ColumnVector, dt: DataType, rowId: Int): Any = dt match { + case BooleanType => col.getBoolean(rowId) + case ByteType => col.getByte(rowId) + case ShortType => col.getShort(rowId) + case IntegerType | DateType => col.getInt(rowId) + case LongType | TimestampType | TimestampNTZType => col.getLong(rowId) + case FloatType => col.getFloat(rowId) + case DoubleType => col.getDouble(rowId) + case d: DecimalType => col.getDecimal(rowId, d.precision, d.scale) + case StringType => col.getUTF8String(rowId).copy() + case _ => null + } + + // Compare values using the same physical representation used in the stats row. + private def compare(dt: DataType, left: Any, right: Any): Int = dt match { + case BooleanType => + java.lang.Boolean.compare(left.asInstanceOf[Boolean], right.asInstanceOf[Boolean]) + case ByteType => + java.lang.Byte.compare(left.asInstanceOf[Byte], right.asInstanceOf[Byte]) + case ShortType => + java.lang.Short.compare(left.asInstanceOf[Short], right.asInstanceOf[Short]) + case IntegerType | DateType => + java.lang.Integer.compare(left.asInstanceOf[Int], right.asInstanceOf[Int]) + case LongType | TimestampType | TimestampNTZType => + java.lang.Long.compare(left.asInstanceOf[Long], right.asInstanceOf[Long]) + case FloatType => + java.lang.Float.compare(left.asInstanceOf[Float], right.asInstanceOf[Float]) + case DoubleType => + java.lang.Double.compare(left.asInstanceOf[Double], right.asInstanceOf[Double]) + case _: DecimalType => + left.asInstanceOf[Decimal].compare(right.asInstanceOf[Decimal]) + case StringType => + ByteArray.compareBinary( + left.asInstanceOf[UTF8String].getBytes, + right.asInstanceOf[UTF8String].getBytes) + case other => + throw new IllegalStateException(s"compare called for unsupported type $other") + } + + // Compute Spark-compatible cache stats before serializing each batch to Arrow. + // The stats are stored beside the Arrow bytes so Spark's cache filter can prune + // CometCachedBatch without decoding the batch first. + private def encodeBatches( + batches: Iterator[ColumnarBatch], + attrs: Seq[Attribute]): Iterator[CachedBatch] = { + batches.flatMap { batch => + val stats = computeStats(batch, attrs) + + Utils.serializeBatches(Iterator.single(batch)).map { case (rows, buffer) => + CometCachedBatch( + numRows = rows.toInt, + sizeInBytes = buffer.size, + stats = stats, + bytes = buffer) + } + } + } + + // Resolve requested columns by exprId, not by name, because aliases may reuse names. + private def selectedIndices( + cacheAttributes: Seq[Attribute], + selectedAttributes: Seq[Attribute]): Array[Int] = { + if (selectedAttributes.isEmpty) { + cacheAttributes.indices.toArray + } else { + val byExprId = cacheAttributes.zipWithIndex.map { case (attr, idx) => + attr.exprId -> idx + }.toMap + + selectedAttributes.map { attr => + byExprId.getOrElse( + attr.exprId, + throw new IllegalStateException( + s"Could not resolve selected attribute ${attr.name} from cache attributes")) + }.toArray + } + } + + // A full-width projection is only an identity projection if every selected index + // is already in column order. For example, [1, 0] must still be projected. + private def isIdentityProjection(indices: Array[Int], numCols: Int): Boolean = + indices.length == numCols && indices.indices.forall(i => indices(i) == i) + + private def projectBatch(batch: ColumnarBatch, indices: Array[Int]): ColumnarBatch = { + if (isIdentityProjection(indices, batch.numCols())) { + batch + } else { + val cols = indices.map(i => batch.column(i).asInstanceOf[ColumnVector]) + new ColumnarBatch(cols, batch.numRows()) + } + } + override def supportsColumnarInput(schema: Seq[Attribute]): Boolean = { val activeConf = SQLConf.get activeConf != null && enabled(activeConf) } + override def supportsColumnarOutput(schema: StructType): Boolean = true // Columnar Comet output is stored as compressed Arrow stream bytes. @@ -86,69 +245,68 @@ class ArrowCachedBatchSerializer extends CachedBatchSerializer { conf: SQLConf): RDD[CachedBatch] = { input.mapPartitions { batches => - Utils.serializeBatches(batches).map { case (rows, buffer) => - CometCachedBatch( - numRows = rows.toInt, - sizeInBytes = buffer.size, - stats = InternalRow.empty, - bytes = buffer) - } + encodeBatches(batches, schema) } } + // A cached relation can contain DefaultCachedBatch when this serializer is installed + // but spark.comet.exec.inMemoryCache.enabled was disabled while the table was cached. + // Decode Spark's default cache format here so the read path stays symmetric with the + // fallback write path without launching another Spark job from inside a task. + private def decodeDefaultCachedBatch( + batch: DefaultCachedBatch, + cacheAttributes: Seq[Attribute], + selectedAttributes: Seq[Attribute], + conf: SQLConf): ColumnarBatch = { + val schema = DataTypeUtils.fromAttributes(selectedAttributes) + val indices = selectedIndices(cacheAttributes, selectedAttributes) + val numRows = batch.numRows + + // This fallback path is used only for Spark's DefaultCachedBatch format. Use on-heap + // vectors here to avoid reading SQLConf inside executor-side cache decode code. + val vectors = OnHeapColumnVector.allocateColumns(numRows, schema) + + val columnarBatch = new ColumnarBatch(vectors.asInstanceOf[Array[ColumnVector]]) + columnarBatch.setNumRows(numRows) + + var i = 0 + while (i < selectedAttributes.length) { + ColumnAccessor.decompress( + batch.buffers(indices(i)), + columnarBatch.column(i).asInstanceOf[WritableColumnVector], + schema.fields(i).dataType, + numRows) + i += 1 + } + + Option(TaskContext.get()).foreach { taskContext => + taskContext.addTaskCompletionListener[Unit](_ => columnarBatch.close()) + } + + columnarBatch + } + override def convertCachedBatchToColumnarBatch( input: RDD[CachedBatch], cacheAttributes: Seq[Attribute], selectedAttributes: Seq[Attribute], conf: SQLConf): RDD[ColumnarBatch] = { + val indices = selectedIndices(cacheAttributes, selectedAttributes) - // Resolve requested columns by exprId, not by name, because aliases may reuse names. - val selectedIndices = - if (selectedAttributes.isEmpty) { - cacheAttributes.indices.toArray - } else { - val byExprId = cacheAttributes.zipWithIndex.map { case (attr, idx) => - attr.exprId -> idx - }.toMap - - selectedAttributes.map { attr => - byExprId.getOrElse( - attr.exprId, - throw new IllegalStateException( - s"Could not resolve selected attribute ${attr.name} from cache attributes")) - }.toArray - } + input.mapPartitions { it => + it.flatMap { + case cb: CometCachedBatch => + Utils.decodeBatches(cb.bytes, "CometCache").map { batch => + projectBatch(batch, indices) + } - val batchTypes = input.map(_.getClass.getName).distinct().collect() - - if (batchTypes.isEmpty) { - input.sparkContext.emptyRDD[ColumnarBatch] - } else if (batchTypes.length > 1) { - throw new IllegalStateException( - s"Mixed cached batch types are not supported: ${batchTypes.mkString(", ")}") - } else if (batchTypes.head == classOf[CometCachedBatch].getName) { - input.mapPartitions { it => - it.flatMap { - case cb: CometCachedBatch => - Utils.decodeBatches(cb.bytes, "CometCache").map { batch => - if (selectedIndices.length == batch.numCols()) { - batch - } else { - val cols = - selectedIndices.map(i => batch.column(i).asInstanceOf[ColumnVector]) - new ColumnarBatch(cols, batch.numRows()) - } - } - - case other => - throw new IllegalStateException( - s"Expected CometCachedBatch, got ${other.getClass.getName}") - } + case cb: DefaultCachedBatch => + Iterator(decodeDefaultCachedBatch(cb, cacheAttributes, selectedAttributes, conf)) + + case other => + throw new IllegalStateException( + s"Unsupported cached batch type ${other.getClass.getName}") } - } else if (batchTypes.head == classOf[DefaultCachedBatch].getName) { - fallback.convertCachedBatchToColumnarBatch(input, cacheAttributes, selectedAttributes, conf) - } else { - throw new IllegalStateException(s"Unsupported cached batch type: ${batchTypes.head}") } } @@ -173,13 +331,7 @@ class ArrowCachedBatchSerializer extends CachedBatchSerializer { sessionTz, TaskContext.get()) - Utils.serializeBatches(iter).map { case (rows, buffer) => - CometCachedBatch( - numRows = rows.toInt, - sizeInBytes = buffer.size, - stats = InternalRow.empty, - bytes = buffer) - } + encodeBatches(iter, schema) } } } @@ -189,65 +341,13 @@ class ArrowCachedBatchSerializer extends CachedBatchSerializer { cacheAttributes: Seq[Attribute], selectedAttributes: Seq[Attribute], conf: SQLConf): RDD[InternalRow] = { + convertCachedBatchToColumnarBatch(input, cacheAttributes, selectedAttributes, conf) + .mapPartitions { batches => + val toUnsafe = UnsafeProjection.create(selectedAttributes, selectedAttributes) - // Resolve requested columns by exprId, not by name, because aliases may reuse names. - val selectedIndices = - if (selectedAttributes.isEmpty) { - cacheAttributes.indices.toArray - } else { - val byExprId = cacheAttributes.zipWithIndex.map { case (attr, idx) => - attr.exprId -> idx - }.toMap - - selectedAttributes.map { attr => - byExprId.getOrElse( - attr.exprId, - throw new IllegalStateException( - s"Could not resolve selected attribute ${attr.name} from cache attributes")) - }.toArray - } - - val batchTypes = input.map(_.getClass.getName).distinct().collect() - - if (batchTypes.isEmpty) { - input.sparkContext.emptyRDD[InternalRow] - } else if (batchTypes.length > 1) { - throw new IllegalStateException( - s"Mixed cached batch types are not supported: ${batchTypes.mkString(", ")}") - } else if (batchTypes.head == classOf[DefaultCachedBatch].getName) { - fallback.convertCachedBatchToInternalRow(input, cacheAttributes, selectedAttributes, conf) - } else if (batchTypes.head == classOf[CometCachedBatch].getName) { - input.mapPartitions { it => - it.flatMap { - case cb: CometCachedBatch => - Utils.decodeBatches(cb.bytes, "CometCache").flatMap { batch => - val projectedBatch = - if (selectedIndices.length == batch.numCols()) { - batch - } else { - val cols = - selectedIndices.map(i => batch.column(i).asInstanceOf[ColumnVector]) - new ColumnarBatch(cols, batch.numRows()) - } - - // Spark's row collect path expects UnsafeRow, not ColumnarBatchRow wrappers. - val toUnsafe = UnsafeProjection.create(selectedAttributes, selectedAttributes) - projectedBatch.rowIterator().asScala.map(row => toUnsafe(row).copy()) - } - - case other => - throw new IllegalStateException( - s"Expected CometCachedBatch, got ${other.getClass.getName}") + batches.flatMap { batch => + batch.rowIterator().asScala.map(row => toUnsafe(row).copy()) } } - } else { - throw new IllegalStateException(s"Unsupported cached batch type: ${batchTypes.head}") - } - } - - override def buildFilter( - predicates: Seq[Expression], - cachedAttributes: Seq[Attribute]): (Int, Iterator[CachedBatch]) => Iterator[CachedBatch] = { - (partitionIndex: Int, it: Iterator[CachedBatch]) => it } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 11174d9bde..68e22df97d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -19,9 +19,14 @@ package org.apache.comet.exec +import java.{util => ju} + +import org.apache.spark.CometDriverPlugin import org.apache.spark.SparkConf import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.catalyst.expressions.{And, Expression, GreaterThanOrEqual, LessThan, Literal} +import org.apache.spark.sql.columnar.SimpleMetricsCachedBatch +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.comet.CometConf @@ -46,8 +51,7 @@ class CometInMemoryCacheSuite extends CometTestBase { } private def cachedBatchTypes(table: String): Array[String] = { - val ds = spark.table(table).asInstanceOf[org.apache.spark.sql.classic.Dataset[_]] - val cached = spark.sharedState.cacheManager.lookupCachedData(ds).get + val cached = spark.sharedState.cacheManager.lookupCachedData(spark.table(table)).get cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers .map(_.getClass.getName) .distinct() @@ -128,6 +132,59 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet cache serializer can read DefaultCachedBatch fallback data") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + spark + .range(1000) + .selectExpr("id as key", "id % 8 as value", "id + 1 as key_plus_1") + .createOrReplaceTempView("default_cached_batch") + + spark.catalog.cacheTable("default_cached_batch") + spark.table("default_cached_batch").count() + + assert( + cachedBatchTypes("default_cached_batch").sameElements( + Array("org.apache.spark.sql.execution.columnar.DefaultCachedBatch"))) + + // Columnar read path: reads DefaultCachedBatch through + // convertCachedBatchToColumnarBatch instead of throwing. + val columnarDf = spark.sql(""" + SELECT key, value + FROM default_cached_batch + WHERE key >= 10 AND key < 20 + """) + checkSparkAnswer(columnarDf) + + val columnarPlan = columnarDf.queryExecution.executedPlan.toString() + assert(!columnarPlan.contains("CometInMemoryTableScan")) + + // Row read path: disabling the vectorized cache reader makes Spark use + // convertCachedBatchToInternalRow. This verifies DefaultCachedBatch is + // readable there as well. + withSQLConf(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "false") { + val rowDf = spark.sql(""" + SELECT key_plus_1 + FROM default_cached_batch + WHERE key >= 10 AND key < 20 + """) + checkSparkAnswer(rowDf) + + val rowPlan = rowDf.queryExecution.executedPlan.toString() + assert(!rowPlan.contains("CometInMemoryTableScan")) + } + + spark.catalog.clearCache() + } + } + test("Comet in-memory cache handles multi-partition cache") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", @@ -180,7 +237,8 @@ class CometInMemoryCacheSuite extends CometTestBase { checkSparkAnswer(emptyDf) val emptyPlan = emptyDf.queryExecution.executedPlan.toString() - assert(!emptyPlan.contains("CometInMemoryTableScan")) + assert(emptyPlan.contains("CometInMemoryTableScan")) + assert(!emptyPlan.contains("CometSparkColumnarToColumnar")) empty.unpersist() spark.catalog.clearCache() @@ -252,4 +310,191 @@ class CometInMemoryCacheSuite extends CometTestBase { spark.catalog.clearCache() } } + + test("Comet in-memory cache supports stats-based batch pruning") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "100") { + + spark.catalog.clearCache() + + spark + .range(0, 1000, 1, 10) + .selectExpr("id as key", "id % 7 as value") + .createOrReplaceTempView("prune_cache") + + spark.catalog.cacheTable("prune_cache") + spark.table("prune_cache").count() + + assert( + cachedBatchTypes("prune_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val cached = spark.sharedState.cacheManager.lookupCachedData(spark.table("prune_cache")).get + val relation = cached.cachedRepresentation + val cachedBuffers = relation.cacheBuilder.cachedColumnBuffers + + // Spark's cache pruning reads statistics through SimpleMetricsCachedBatch. + // CometCachedBatch must expose the same five statistics per column: + // lower bound, upper bound, null count, row count, and size in bytes. + val firstBatch = cachedBuffers.take(1).head + assert(firstBatch.isInstanceOf[SimpleMetricsCachedBatch]) + assert( + firstBatch.asInstanceOf[SimpleMetricsCachedBatch].stats.numFields == + relation.output.length * 5) + + val keyAttr = relation.output.find(_.name == "key").get + + // Call the serializer filter directly so the test fails if buildFilter is + // accidentally changed back to a no-op. + def prunedCount(predicate: Expression): Long = { + val filter = relation.cacheBuilder.serializer.buildFilter(Seq(predicate), relation.output) + cachedBuffers.mapPartitionsWithIndex(filter).count() + } + + val totalBatches = cachedBuffers.count() + assert(totalBatches > 1) + + val targetPredicate = + And(GreaterThanOrEqual(keyAttr, Literal(900L)), LessThan(keyAttr, Literal(905L))) + assert(prunedCount(targetPredicate) == 1) + + val outsidePredicate = LessThan(keyAttr, Literal(0L)) + assert(prunedCount(outsidePredicate) == 0) + + val allPredicate = + And(GreaterThanOrEqual(keyAttr, Literal(0L)), LessThan(keyAttr, Literal(1000L))) + assert(prunedCount(allPredicate) == totalBatches) + + val df = spark.sql(""" + SELECT key, value + FROM prune_cache + WHERE key >= 900 AND key < 905 + """) + checkSparkAnswer(df) + + val plan = df.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + assert(!plan.contains("CometSparkColumnarToColumnar")) + + spark.catalog.clearCache() + } + } + + test("Comet plugin respects user-provided cache serializer") { + val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key + val cometSerializer = + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer" + val userSerializer = "com.example.CustomCachedBatchSerializer" + + val defaultConf = new SparkConf() + .set(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") + val defaultExtraConfs = new ju.HashMap[String, String]() + + // With no user serializer configured, the plugin should install Comet's + // serializer and also return it through extraConfs for executors. + CometDriverPlugin.maybeSetCacheSerializer(defaultConf, defaultExtraConfs) + + assert(defaultConf.get(serializerKey) == cometSerializer) + assert(defaultExtraConfs.get(serializerKey) == cometSerializer) + + val userConf = new SparkConf() + .set(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") + .set(serializerKey, userSerializer) + val userExtraConfs = new ju.HashMap[String, String]() + + // If the user already configured a cache serializer, keep it and do not + // send a replacement serializer through extraConfs. + CometDriverPlugin.maybeSetCacheSerializer(userConf, userExtraConfs) + + assert(userConf.get(serializerKey) == userSerializer) + assert(!userExtraConfs.containsKey(serializerKey)) + } + + test("Comet in-memory cache supports empty projection scan") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + spark + .range(1000) + .selectExpr("id as key", "id % 8 as value") + .createOrReplaceTempView("count_cache") + + spark.catalog.cacheTable("count_cache") + spark.table("count_cache").count() + + val df = spark.sql("SELECT count(*) FROM count_cache") + checkSparkAnswer(df) + + val plan = df.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache pruning handles NaN floating-point values") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "2") { + + spark.catalog.clearCache() + + spark + .sql(""" + SELECT * + FROM VALUES + (0, CAST('NaN' AS DOUBLE), CAST('NaN' AS FLOAT)), + (1, 1.0D, CAST(1.0 AS FLOAT)), + (2, -0.0D, CAST(-0.0 AS FLOAT)), + (3, 0.0D, CAST(0.0 AS FLOAT)) + AS t(id, d, f) + """) + .createOrReplaceTempView("nan_prune_cache") + + spark.catalog.cacheTable("nan_prune_cache") + spark.table("nan_prune_cache").count() + + val doubleDf = spark.sql(""" + SELECT id + FROM nan_prune_cache + WHERE isnan(d) + """) + checkSparkAnswer(doubleDf) + + val floatDf = spark.sql(""" + SELECT id + FROM nan_prune_cache + WHERE isnan(f) + """) + checkSparkAnswer(floatDf) + + val zeroDf = spark.sql(""" + SELECT id + FROM nan_prune_cache + WHERE d = 0.0D OR f = CAST(0.0 AS FLOAT) + """) + checkSparkAnswer(zeroDf) + + val plan = doubleDf.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + assert(!plan.contains("CometSparkColumnarToColumnar")) + + spark.catalog.clearCache() + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala new file mode 100644 index 0000000000..f3ceb831ce --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -0,0 +1,152 @@ +/* + * 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.benchmark + +import org.apache.spark.SparkConf +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +object CometInMemoryCacheBenchmark extends CometBenchmarkBase { + private val numRows = 5 * 1000 * 1000 + private val cacheTable = "comet_cache_bench" + private val sourceTable = "comet_cache_bench_src" + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName("CometInMemoryCacheBenchmark") + .set("spark.master", "local[1]") + .setIfMissing("spark.driver.memory", "3g") + .setIfMissing("spark.executor.memory", "3g") + .set("spark.plugins", "org.apache.spark.CometPlugin") + .set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set( + "spark.sql.cache.serializer", + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + + val sparkSession = SparkSession + .builder() + .config(conf) + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + + sparkSession.conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") + sparkSession.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") + sparkSession.conf.set(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, "true") + sparkSession.conf.set(SQLConf.ANSI_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") + sparkSession + } + + override def runCometBenchmark(args: Array[String]): Unit = { + withTempTable(sourceTable, cacheTable) { + spark + .range(0, numRows, 1, 16) + .selectExpr("id", "id % 1000 AS k", "id + 1 AS v") + .createOrReplaceTempView(sourceTable) + + runCacheBenchmark( + "in-memory cache repeated scan", + s"SELECT sum(id), sum(k), sum(v) FROM $cacheTable") + + runCacheBenchmark( + "in-memory cache selective filter", + s""" + |SELECT sum(id), sum(k), sum(v) + |FROM $cacheTable + |WHERE id >= 4500000 AND id < 4750000 + """.stripMargin) + } + } + + private def runCacheBenchmark(name: String, query: String): Unit = { + withCachedTable { + withSQLConf(cacheConf(nativeCacheEnabled = false): _*) { + verifyPlan(query, nativeCacheEnabled = false) + } + withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { + verifyPlan(query, nativeCacheEnabled = true) + } + + val benchmark = new Benchmark(name, numRows, output = output) + + benchmark.addCase("Comet cache disabled") { _ => + withSQLConf(cacheConf(nativeCacheEnabled = false): _*) { + spark.sql(query).noop() + } + } + + benchmark.addCase("Comet cache enabled") { _ => + withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { + spark.sql(query).noop() + } + } + + benchmark.run() + } + } + + private def withCachedTable(f: => Unit): Unit = { + spark.catalog.clearCache() + + // Materialize the cache once using Comet's cache serializer. + // The benchmark measures repeated cache reads by comparing the + // fallback read path against CometInMemoryTableScan. + withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { + spark.sql(s"SELECT id, k, v FROM $sourceTable").createOrReplaceTempView(cacheTable) + spark.catalog.cacheTable(cacheTable) + spark.table(cacheTable).count() + } + + try f + finally { + spark.catalog.uncacheTable(cacheTable) + spark.catalog.clearCache() + } + } + + private def verifyPlan(query: String, nativeCacheEnabled: Boolean): Unit = { + val plan = spark.sql(query).queryExecution.executedPlan.toString() + + if (nativeCacheEnabled) { + assert(plan.contains("CometInMemoryTableScan"), s"Expected native cache scan:\n$plan") + assert(!plan.contains("CometSparkColumnarToColumnar"), s"Unexpected conversion:\n$plan") + } else { + assert( + !plan.contains("CometInMemoryTableScan"), + s"Native cache scan should be disabled:\n$plan") + } + } + + private def cacheConf(nativeCacheEnabled: Boolean): Seq[(String, String)] = { + Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> nativeCacheEnabled.toString, + "spark.comet.sparkToColumnar.enabled" -> "true", + "spark.comet.exec.onHeap.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + } +} From 64be140701cf26b605a3d201df71b8beb2df3734 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 27 Jul 2026 09:03:02 -0600 Subject: [PATCH 3/8] fix: reset memoized cache serializer in in-memory cache tests InMemoryRelation resolves spark.sql.cache.serializer once per JVM and memoizes the instance in a static field. Test suites share a forked JVM, so whichever suite caches a table first pins the serializer for every suite that follows. CometInMemoryCacheSuite runs after CometExecSuite in the exec group, so its configured serializer was ignored, every cached batch was a DefaultCachedBatch, and 9 of its 11 tests failed on all Spark profiles. Reset the memoized serializer around the suite, via a test-only shim for the private[columnar] clearSerializer. Also address review feedback: - Fall through to the SparkToColumnar path when the native cache is enabled but the relation was cached by a foreign serializer, so enabling the feature is never worse for a scan than leaving it off. Covered by a new CometExecSuite test. - Explain on CometInMemoryTableScanExec.output why the declared output can be narrower than the emitted batch width for an empty projection. - Drop the import made redundant by the org.apache.spark.sql.comet._ wildcard. --- .../apache/comet/rules/CometExecRule.scala | 28 ++++++++-------- .../comet/CometInMemoryTableScanExec.scala | 5 +++ .../apache/comet/exec/CometExecSuite.scala | 29 ++++++++++++++++ .../comet/exec/CometInMemoryCacheSuite.scala | 21 ++++++++++++ .../CometInMemoryRelationHelper.scala | 33 +++++++++++++++++++ 5 files changed, 102 insertions(+), 14 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/execution/columnar/CometInMemoryRelationHelper.scala diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index d1ed04cd5b..41cfcfb226 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -29,7 +29,7 @@ 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.CometInMemoryTableScanExec +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._ @@ -287,22 +287,22 @@ case class CometExecRule(session: SparkSession) convertToComet(op, CometScanWrapper).getOrElse(op) case scan: InMemoryTableScanExec => - val usesCometCacheSerializer = - scan.relation.cacheBuilder.serializer - .isInstanceOf[org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer] + val serializer = scan.relation.cacheBuilder.serializer + val usesCometCacheSerializer = serializer.isInstanceOf[ArrowCachedBatchSerializer] + val nativeCacheEnabled = CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf) - if (CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf)) { - if (usesCometCacheSerializer) { - convertToComet(scan, CometInMemoryTableScanExec).getOrElse(scan) - } else { + if (nativeCacheEnabled && usesCometCacheSerializer) { + 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) { withFallbackReason( scan, - "Comet in-memory cache requires ArrowCachedBatchSerializer, " + - s"but found ${scan.relation.cacheBuilder.serializer.getClass.getName}") - scan - } - } else { - if (usesCometCacheSerializer) { + s"Comet in-memory cache requires ${classOf[ArrowCachedBatchSerializer].getName} " + + s"but this relation was cached with ${serializer.getClass.getName}") + } else if (usesCometCacheSerializer) { withFallbackReason( scan, "Native support for operator InMemoryTableScanExec is disabled. " + diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala index e1e42bc137..c7032579df 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -57,6 +57,11 @@ case class CometInMemoryTableScanExec( override lazy val metrics: Map[String, SQLMetric] = Map( "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + // For an empty-projection scan (`SELECT count(*)`) this is empty while `scanOutput` holds the + // full cache schema, so the emitted batches are wider than the declared output. That is safe + // because the only consumer of an empty-output scan is a count-style aggregate, which reads + // the row count rather than any column; `convert` and `createExec` deliberately fall back to + // the cache schema in that case because the native plan still needs a non-empty scan schema. override def output: Seq[Attribute] = originalPlan.output // Use the serializer's vector types because the cached batch layout is owned by the serializer. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index ce6071bf1a..bf60d5de3f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -38,6 +38,7 @@ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, Comet import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec} +import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec, SortMergeJoinExec} @@ -3700,6 +3701,34 @@ class CometExecSuite extends CometTestBase { }) } + test("SparkToColumnar over InMemoryTableScanExec with a non-Comet cache serializer") { + // Enabling the native in-memory cache must never leave a cached scan worse off than having + // the feature disabled. When the relation was cached by a serializer Comet cannot decode, the + // native scan is unavailable, but the scan should still take the SparkToColumnar fallback + // rather than staying entirely on Spark. + // + // This session does not configure spark.sql.cache.serializer, so it uses Spark's default. + // The reset makes that deterministic regardless of which suite ran first in this JVM, since + // InMemoryRelation memoizes the serializer per JVM. + CometInMemoryRelationHelper.clearSerializer() + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + spark + .range(1000) + .selectExpr("id as key", "id % 8 as value") + .createOrReplaceTempView("foreign_cache_serializer") + spark.catalog.cacheTable("foreign_cache_serializer") + try { + val df = spark.sql("SELECT * FROM foreign_cache_serializer").groupBy("key").count() + checkSparkAnswerAndOperator(df, includeClasses = Seq(classOf[CometSparkToColumnarExec])) + } finally { + spark.catalog.uncacheTable("foreign_cache_serializer") + } + } + } + test("SparkToColumnar eliminate redundant in AQE") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 68e22df97d..24c0f6dbbd 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -26,11 +26,32 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.catalyst.expressions.{And, Expression, GreaterThanOrEqual, LessThan, Literal} import org.apache.spark.sql.columnar.SimpleMetricsCachedBatch +import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.comet.CometConf class CometInMemoryCacheSuite extends CometTestBase { + + // `InMemoryRelation` resolves `spark.sql.cache.serializer` once per JVM and memoizes the + // instance in a static field. Test suites share a forked JVM, so whichever suite caches a + // table first pins the serializer for everything that follows: without this reset the + // serializer configured below is ignored and every cached batch here is a `DefaultCachedBatch`. + // Clear it on the way out as well so this suite does not pin Comet's serializer for the rest + // of the JVM. + override protected def beforeAll(): Unit = { + CometInMemoryRelationHelper.clearSerializer() + super.beforeAll() + } + + override protected def afterAll(): Unit = { + try { + super.afterAll() + } finally { + CometInMemoryRelationHelper.clearSerializer() + } + } + override protected def sparkConf: SparkConf = { val conf = new SparkConf() conf.set("spark.driver.memory", "1G") diff --git a/spark/src/test/scala/org/apache/spark/sql/execution/columnar/CometInMemoryRelationHelper.scala b/spark/src/test/scala/org/apache/spark/sql/execution/columnar/CometInMemoryRelationHelper.scala new file mode 100644 index 0000000000..ff5240af2a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/execution/columnar/CometInMemoryRelationHelper.scala @@ -0,0 +1,33 @@ +/* + * 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.execution.columnar + +/** + * Test-only access to `InMemoryRelation`'s JVM-wide cached `CachedBatchSerializer`. + * + * `InMemoryRelation` resolves `spark.sql.cache.serializer` once per JVM and memoizes the instance + * in a static field, so the first suite in a forked JVM that caches a table pins the serializer + * for every suite that follows. A suite that needs a specific cache serializer must reset that + * state around itself. `InMemoryRelation.clearSerializer` is `private[columnar]`, hence this + * shim. + */ +object CometInMemoryRelationHelper { + def clearSerializer(): Unit = InMemoryRelation.clearSerializer() +} From 63c5a916b2e5626b040dc1f097384ab1a443bf91 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 27 Jul 2026 09:43:45 -0600 Subject: [PATCH 4/8] fix: correct cache format selection and stats pruning in Comet cache serializer Three defects found while reviewing type coverage, each with a regression test that fails without the corresponding fix. Pruning dropped every batch for columns without bounds. `tracksBounds` only computes bounds for the types it lists, so a collated `StringType` on Spark 4.x left them null. Spark still builds a partition filter for such a column because a collated string literal is an `AtomicType`, and comparing against null bounds yields null, which the generated predicate treats as false. A filtered query over a collated string column returned zero rows instead of the matching ones, with no error. `buildFilter` now drops predicates over columns that have no bounds, keeping `IsNull` and `IsNotNull` which only read null and row counts. Interval columns broke `cache()`. `Utils.getFieldVector` has no case for `DurationVector` or `IntervalYearVector`, so materializing a cached relation containing a `DayTimeIntervalType` or `YearMonthIntervalType` column threw `Unsupported Arrow Vector for serialize`. The serializer now reports the schemas its Arrow writer supports and delegates the rest to Spark's default cache serializer. Reading a `DefaultCachedBatch` failed for any non-primitive column. `supportsColumnarOutput` returned true unconditionally while `decodeDefaultCachedBatch` decoded through `ColumnAccessor.decompress`, whose `PassThrough` decoder only handles the seven primitive physical types. Spark's own serializer gates `supportsColumnarOutput` on exactly those types to avoid this. Reading a cached string column threw `scala.MatchError: PhysicalStringType`. The cached format is now decided by schema alone rather than by a runtime config, so a relation is either entirely Comet format or entirely delegated to Spark, and the hand-rolled `DefaultCachedBatch` decoder is gone. `spark.sql.cache.serializer` is a static conf, so a format that could flip mid-session was never readable back reliably. `CometExecRule` checks the schema before choosing the native scan. Also add coverage for all supported types, the row read path over `CometCachedBatch`, and a reordered full-width projection. --- .../scala/org/apache/comet/CometConf.scala | 10 +- .../apache/comet/rules/CometExecRule.scala | 16 +- .../arrow/ArrowCachedBatchSerializer.scala | 157 +++++++++------ .../comet/exec/CometInMemoryCacheSuite.scala | 190 +++++++++++++++++- 4 files changed, 296 insertions(+), 77 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index ded79c544a..f520cc756f 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -238,9 +238,13 @@ object CometConf extends ShimCometConf { conf("spark.comet.exec.inMemoryCache.enabled") .category(CATEGORY_EXEC) .doc( - "Whether to enable Comet native execution for in-memory cached tables. " + - "When disabled or when spark.comet.enabled=false, Spark's default cache " + - "serializer and execution path will be used.") + "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.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 41cfcfb226..780c7d0cb3 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -289,19 +289,31 @@ case class CometExecRule(session: SparkSession) 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 && usesCometCacheSerializer) { + 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) { + 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, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 071934c0cd..eaadc90f0a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -21,14 +21,12 @@ package org.apache.spark.sql.comet.execution.arrow import scala.collection.JavaConverters._ -import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, GenericInternalRow, UnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull, UnsafeProjection} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch, SimpleMetricsCachedBatchSerializer} import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.columnar.{ColumnAccessor, DefaultCachedBatch, DefaultCachedBatchSerializer} -import org.apache.spark.sql.execution.vectorized.{OnHeapColumnVector, WritableColumnVector} +import org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} @@ -36,7 +34,7 @@ import org.apache.spark.storage.StorageLevel import org.apache.spark.unsafe.types.{ByteArray, UTF8String} import org.apache.spark.util.io.ChunkedByteBuffer -import org.apache.comet.{CometArrowAllocator, CometConf} +import org.apache.comet.CometArrowAllocator /** * Cached batch format used when Comet writes Spark in-memory cache data. @@ -54,20 +52,23 @@ private case class CometCachedBatch( /** * Cache serializer that stores Comet-compatible Arrow batches in Spark's in-memory cache. * - * Writes use Comet's Arrow cache format only when Comet and the native in-memory cache path are - * enabled. Reads of CometCachedBatch are still supported even if the native scan is disabled - * later, because Spark may then read the same cached data through the SparkToColumnar fallback - * path. + * The cached payload format is decided by the schema alone. A relation whose schema Comet's Arrow + * writer supports is stored as `CometCachedBatch`, and every other relation is delegated in full + * to Spark's `DefaultCachedBatchSerializer`. The format deliberately does not depend on any + * runtime config: `spark.sql.cache.serializer` is a static conf, so installing this serializer is + * already a per-application decision, and a relation whose format could flip mid-session cannot + * be read back reliably. `spark.comet.exec.inMemoryCache.enabled` still governs whether a scan + * over the cache runs natively, and its value at startup is what makes `CometDriverPlugin` + * install this serializer in the first place. + * + * Reads of `CometCachedBatch` keep working when the native scan is disabled, because Spark then + * reads the same cached data through the SparkToColumnar fallback path. */ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { - private val fallback = new DefaultCachedBatchSerializer() + import ArrowCachedBatchSerializer.supportsSchema - // Cache writes use Comet format only when both Comet and the in-memory cache scan are enabled. - private def enabled(conf: SQLConf): Boolean = { - CometConf.COMET_ENABLED.get(conf) && - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf) - } + private val fallback = new DefaultCachedBatchSerializer() // Row-to-Arrow conversion needs a StructType, while cache APIs pass attributes. private def toStructType(schema: Seq[Attribute]): StructType = { @@ -229,14 +230,45 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } } - override def supportsColumnarInput(schema: Seq[Attribute]): Boolean = { - val activeConf = SQLConf.get - activeConf != null && enabled(activeConf) + // Spark's SimpleMetricsCachedBatchSerializer prunes a batch when the generated partition filter + // does not evaluate to true against the stats row. Bounds are only computed for the types + // tracksBounds accepts, and for every other column the lower and upper bounds stay null, which + // makes a comparison against them evaluate to null and therefore prune the batch. That would + // silently drop rows, so predicates over columns without bounds are not pushed down at all. + // Null counts and row counts are recorded for every column, so IsNull and IsNotNull stay safe. + override def buildFilter( + predicates: Seq[Expression], + cachedAttributes: Seq[Attribute]): (Int, Iterator[CachedBatch]) => Iterator[CachedBatch] = { + val prunable = cachedAttributes.collect { + case a if tracksBounds(a.dataType) => a.exprId + }.toSet + + val prunablePredicates = predicates.filter { + case _: IsNull | _: IsNotNull => true + case p => p.references.forall(a => prunable.contains(a.exprId)) + } + + super.buildFilter(prunablePredicates, cachedAttributes) } - override def supportsColumnarOutput(schema: StructType): Boolean = true + // Comet's Arrow writer only handles the types listed in supportsSchema. Reporting false here + // sends the relation down the row path, where it is delegated to Spark's default serializer, + // instead of failing at cache materialization inside Utils.serializeBatches. + override def supportsColumnarInput(schema: Seq[Attribute]): Boolean = supportsSchema(schema) - // Columnar Comet output is stored as compressed Arrow stream bytes. + // A relation Comet stores is always readable as columnar Arrow. Anything else holds + // DefaultCachedBatch, so defer to Spark, which only claims columnar output for the primitive + // types its ColumnAccessor.decompress path can actually decode. + override def supportsColumnarOutput(schema: StructType): Boolean = { + if (schema.fields.forall(f => ArrowCachedBatchSerializer.supportsType(f.dataType))) { + true + } else { + fallback.supportsColumnarOutput(schema) + } + } + + // Columnar Comet output is stored as compressed Arrow stream bytes. Spark only calls this when + // supportsColumnarInput returned true, so the schema is known to be Comet-writable here. override def convertColumnarBatchToCachedBatch( input: RDD[ColumnarBatch], schema: Seq[Attribute], @@ -248,48 +280,19 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } } - // A cached relation can contain DefaultCachedBatch when this serializer is installed - // but spark.comet.exec.inMemoryCache.enabled was disabled while the table was cached. - // Decode Spark's default cache format here so the read path stays symmetric with the - // fallback write path without launching another Spark job from inside a task. - private def decodeDefaultCachedBatch( - batch: DefaultCachedBatch, - cacheAttributes: Seq[Attribute], - selectedAttributes: Seq[Attribute], - conf: SQLConf): ColumnarBatch = { - val schema = toStructType(selectedAttributes) - val indices = selectedIndices(cacheAttributes, selectedAttributes) - val numRows = batch.numRows - - // This fallback path is used only for Spark's DefaultCachedBatch format. Use on-heap - // vectors here to avoid reading SQLConf inside executor-side cache decode code. - val vectors = OnHeapColumnVector.allocateColumns(numRows, schema) - - val columnarBatch = new ColumnarBatch(vectors.asInstanceOf[Array[ColumnVector]]) - columnarBatch.setNumRows(numRows) - - var i = 0 - while (i < selectedAttributes.length) { - ColumnAccessor.decompress( - batch.buffers(indices(i)), - columnarBatch.column(i).asInstanceOf[WritableColumnVector], - schema.fields(i).dataType, - numRows) - i += 1 - } - - Option(TaskContext.get()).foreach { taskContext => - taskContext.addTaskCompletionListener[Unit](_ => columnarBatch.close()) - } - - columnarBatch - } - override def convertCachedBatchToColumnarBatch( input: RDD[CachedBatch], cacheAttributes: Seq[Attribute], selectedAttributes: Seq[Attribute], conf: SQLConf): RDD[ColumnarBatch] = { + if (!supportsSchema(cacheAttributes)) { + return fallback.convertCachedBatchToColumnarBatch( + input, + cacheAttributes, + selectedAttributes, + conf) + } + val indices = selectedIndices(cacheAttributes, selectedAttributes) input.mapPartitions { it => @@ -299,9 +302,6 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { projectBatch(batch, indices) } - case cb: DefaultCachedBatch => - Iterator(decodeDefaultCachedBatch(cb, cacheAttributes, selectedAttributes, conf)) - case other => throw new IllegalStateException( s"Unsupported cached batch type ${other.getClass.getName}") @@ -309,14 +309,14 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } } - // Row input can still be cached in Comet format by converting rows to Arrow batches first. + // Row input is cached in Comet format by converting rows to Arrow batches first. override def convertInternalRowToCachedBatch( input: RDD[InternalRow], schema: Seq[Attribute], storageLevel: StorageLevel, conf: SQLConf): RDD[CachedBatch] = { - if (!enabled(conf)) { + if (!supportsSchema(schema)) { fallback.convertInternalRowToCachedBatch(input, schema, storageLevel, conf) } else { val batchSize = conf.columnBatchSize @@ -340,6 +340,14 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { cacheAttributes: Seq[Attribute], selectedAttributes: Seq[Attribute], conf: SQLConf): RDD[InternalRow] = { + if (!supportsSchema(cacheAttributes)) { + return fallback.convertCachedBatchToInternalRow( + input, + cacheAttributes, + selectedAttributes, + conf) + } + convertCachedBatchToColumnarBatch(input, cacheAttributes, selectedAttributes, conf) .mapPartitions { batches => val toUnsafe = UnsafeProjection.create(selectedAttributes, selectedAttributes) @@ -350,3 +358,28 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } } } + +object ArrowCachedBatchSerializer { + + /** + * Whether Comet's Arrow cache format can store this type. + * + * This mirrors the vectors `Utils.getFieldVector` accepts. A type missing from that list throws + * during cache materialization, so it has to be delegated to Spark's default cache format + * instead. Interval types are the notable omission. + */ + def supportsType(dt: DataType): Boolean = dt match { + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + DateType | TimestampType | TimestampNTZType | BinaryType | NullType => + true + case _: DecimalType => true + case _: StringType => true + case ArrayType(elementType, _) => supportsType(elementType) + case MapType(keyType, valueType, _) => supportsType(keyType) && supportsType(valueType) + case StructType(fields) => fields.forall(f => supportsType(f.dataType)) + case _ => false + } + + def supportsSchema(schema: Seq[Attribute]): Boolean = + schema.forall(a => supportsType(a.dataType)) +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 24c0f6dbbd..d5dc8c282d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -30,6 +30,7 @@ import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus class CometInMemoryCacheSuite extends CometTestBase { @@ -153,19 +154,24 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Comet cache serializer can read DefaultCachedBatch fallback data") { + test("Comet cache serializer delegates unsupported types to Spark's cache format") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_SHUFFLE_MODE.key -> "jvm", SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", "spark.comet.sparkToColumnar.enabled" -> "true") { spark.catalog.clearCache() + // Interval types have no Arrow vector in Utils.getFieldVector. Without the schema check in + // the serializer, caching this relation fails outright with "Unsupported Arrow Vector for + // serialize: class org.apache.arrow.vector.DurationVector". spark - .range(1000) - .selectExpr("id as key", "id % 8 as value", "id + 1 as key_plus_1") + .sql(""" + SELECT id AS key, make_dt_interval(0, 0, 0, id) AS dt + FROM range(1000) + """) .createOrReplaceTempView("default_cached_batch") spark.catalog.cacheTable("default_cached_batch") @@ -175,27 +181,27 @@ class CometInMemoryCacheSuite extends CometTestBase { cachedBatchTypes("default_cached_batch").sameElements( Array("org.apache.spark.sql.execution.columnar.DefaultCachedBatch"))) - // Columnar read path: reads DefaultCachedBatch through - // convertCachedBatchToColumnarBatch instead of throwing. + // Columnar read path, delegated to Spark's serializer. val columnarDf = spark.sql(""" - SELECT key, value + SELECT key, dt FROM default_cached_batch WHERE key >= 10 AND key < 20 """) + assert(columnarDf.collect().length == 10) checkSparkAnswer(columnarDf) val columnarPlan = columnarDf.queryExecution.executedPlan.toString() assert(!columnarPlan.contains("CometInMemoryTableScan")) // Row read path: disabling the vectorized cache reader makes Spark use - // convertCachedBatchToInternalRow. This verifies DefaultCachedBatch is - // readable there as well. + // convertCachedBatchToInternalRow. withSQLConf(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "false") { val rowDf = spark.sql(""" - SELECT key_plus_1 + SELECT dt FROM default_cached_batch WHERE key >= 10 AND key < 20 """) + assert(rowDf.collect().length == 10) checkSparkAnswer(rowDf) val rowPlan = rowDf.queryExecution.executedPlan.toString() @@ -464,6 +470,170 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + private def withNativeCache(f: => Unit): Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + spark.catalog.clearCache() + try f + finally spark.catalog.clearCache() + } + } + + test("Comet in-memory cache round-trips all supported types") { + withNativeCache { + val query = + """ + SELECT + id AS l, + CAST(id AS INT) AS i, + CAST(id AS SMALLINT) AS sh, + CAST(id AS TINYINT) AS ti, + CAST(id % 2 AS BOOLEAN) AS bo, + CAST(id AS FLOAT) AS fl, + CAST(id AS DOUBLE) AS db, + CAST(id AS DECIMAL(20,4)) AS de, + CAST(id AS STRING) AS st, + CAST(CAST(id AS STRING) AS BINARY) AS bi, + DATE_ADD(DATE'2020-01-01', CAST(id AS INT)) AS da, + TIMESTAMP'2020-01-01 00:00:00' + make_dt_interval(0, 0, 0, id) AS ts, + CAST(TIMESTAMP'2020-01-01 00:00:00' + make_dt_interval(0, 0, 0, id) AS TIMESTAMP_NTZ) + AS tsntz, + struct(id AS a, CAST(id AS STRING) AS b) AS sc, + array(id, id + 1) AS ar, + map('k', id) AS mp + FROM range(100) + """ + + // Expected values come from the uncached query so a wrong-but-consistent cached answer + // cannot make this pass. + val expected = spark.sql(query).orderBy("l").collect() + + spark.sql(query).createOrReplaceTempView("all_types_cache") + spark.catalog.cacheTable("all_types_cache") + spark.table("all_types_cache").count() + + assert( + cachedBatchTypes("all_types_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val df = spark.sql("SELECT * FROM all_types_cache").orderBy("l") + assert(df.collect() === expected) + assert(df.queryExecution.executedPlan.toString().contains("CometInMemoryTableScan")) + } + } + + test("Comet in-memory cache prunes only on columns that have bounds") { + assume(isSpark40Plus, "collated string types require Spark 4.0+") + withNativeCache { + // A collated StringType does not match `case StringType` in the serializer's bounds + // tracking, so its lower and upper bounds stay null. Spark still builds a partition filter + // for it because a collated string literal is an AtomicType, and comparing against null + // bounds prunes every batch. Without the buildFilter guard this query returns no rows. + spark + .sql("SELECT id, CAST(id AS STRING) COLLATE UTF8_LCASE AS s FROM range(100)") + .createOrReplaceTempView("collated_cache") + spark.catalog.cacheTable("collated_cache") + spark.table("collated_cache").count() + + assert( + cachedBatchTypes("collated_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val expected = + spark.sql("SELECT id FROM range(100) WHERE CAST(id AS STRING) >= '5'").collect().length + assert(expected > 0) + assert( + spark.sql("SELECT id FROM collated_cache WHERE s >= '5'").collect().length == expected) + + // Null-count based pruning stays available for columns without bounds. + assert( + spark.sql("SELECT id FROM collated_cache WHERE s IS NOT NULL").collect().length == 100) + } + } + + test("Comet in-memory cache is readable when Comet is disabled") { + // spark.sql.cache.serializer is static, so the cached format cannot depend on a runtime + // config. Disabling Comet must still leave the cached relation readable, including for + // string columns, which Spark's DefaultCachedBatch columnar decoder cannot handle. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + spark.catalog.clearCache() + spark + .sql("SELECT id, CAST(id AS STRING) AS s FROM range(100)") + .createOrReplaceTempView("comet_off_cache") + spark.catalog.cacheTable("comet_off_cache") + spark.table("comet_off_cache").count() + + val rows = spark.sql("SELECT s FROM comet_off_cache WHERE id >= 90").collect() + assert(rows.length == 10) + assert(rows.map(_.getString(0)).toSet == (90 until 100).map(_.toString).toSet) + + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache supports the row read path over CometCachedBatch") { + withNativeCache { + spark + .sql("SELECT id AS key, CAST(id AS STRING) AS s FROM range(100)") + .createOrReplaceTempView("row_path_cache") + spark.catalog.cacheTable("row_path_cache") + spark.table("row_path_cache").count() + + assert( + cachedBatchTypes("row_path_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + // Turning off the vectorized cache reader routes the scan through + // convertCachedBatchToInternalRow rather than convertCachedBatchToColumnarBatch. + withSQLConf(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "false") { + val rows = spark.sql("SELECT s FROM row_path_cache WHERE key >= 90").collect() + assert(rows.length == 10) + assert(rows.map(_.getString(0)).toSet == (90 until 100).map(_.toString).toSet) + } + } + } + + test("Comet in-memory cache projects a reordered full-width selection") { + withNativeCache { + spark + .sql("SELECT id AS key, CAST(id * 10 AS STRING) AS value FROM range(10)") + .createOrReplaceTempView("reorder_cache") + spark.catalog.cacheTable("reorder_cache") + spark.table("reorder_cache").count() + + val relation = + spark.sharedState.cacheManager + .lookupCachedData(spark.table("reorder_cache")) + .get + .cachedRepresentation + val serializer = relation.cacheBuilder.serializer + + // A full-width but reordered projection has the same length as the cache schema, so an + // identity check based on length alone would return the columns in the wrong order. + val reordered = Seq(relation.output(1), relation.output(0)) + val rows = serializer + .convertCachedBatchToInternalRow( + relation.cacheBuilder.cachedColumnBuffers, + relation.output, + reordered, + spark.sessionState.conf) + .map(row => (row.getString(0).toString, row.getLong(1))) + .collect() + .sortBy(_._2) + + assert(rows.length == 10) + assert(rows === (0 until 10).map(i => ((i * 10).toString, i.toLong)).toArray) + } + } + test("Comet in-memory cache pruning handles NaN floating-point values") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", From 9dce0d66a3488153014d6fe17bc7a97054e4016f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 29 Jul 2026 14:19:38 -0600 Subject: [PATCH 5/8] fix: convert non-Arrow columnar input in Comet cache serializer `ArrowCachedBatchSerializer.supportsColumnarInput` decides the cache input format from the schema alone, which is all Spark gives it. Returning true makes `InMemoryRelation` strip the `ColumnarToRow` above the cached plan and feed `cachedPlan.executeColumnar()` directly to `convertColumnarBatchToCachedBatch`. Any Spark or third-party columnar leaf under that transition (Spark's vectorized Parquet/ORC reader, a connector's own vectors) therefore reaches the serializer with `On/OffHeapColumnVector`-style columns, and `Utils.serializeBatches` fails with "Comet execution only takes Arrow Arrays". The serializer cannot detect this from the schema, so copy non-Arrow batches into Arrow at write time. Comet-native cached plans keep the existing zero-copy path. --- .../arrow/ArrowCachedBatchSerializer.scala | 43 +++++++-- .../arrow/CometArrowConverters.scala | 45 +++++++++- .../comet/exec/CometInMemoryCacheSuite.scala | 90 +++++++++++++++++++ 3 files changed, 170 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index eaadc90f0a..22069bd1d0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -180,18 +180,49 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // Compute Spark-compatible cache stats before serializing each batch to Arrow. // The stats are stored beside the Arrow bytes so Spark's cache filter can prune // CometCachedBatch without decoding the batch first. + // + // Spark decides the input path from `supportsColumnarInput`, which only sees the schema, so a + // columnar input batch is not guaranteed to be Arrow-backed: any Spark or third-party columnar + // leaf (Spark's vectorized Parquet/ORC reader, a connector's own vectors) reaches + // `convertColumnarBatchToCachedBatch` with `On/OffHeapColumnVector`-style columns, which + // `Utils.serializeBatches` cannot write. Copy those into Arrow first. private def encodeBatches( batches: Iterator[ColumnarBatch], attrs: Seq[Attribute]): Iterator[CachedBatch] = { + lazy val structType = toStructType(attrs) + batches.flatMap { batch => val stats = computeStats(batch, attrs) - Utils.serializeBatches(Iterator.single(batch)).map { case (rows, buffer) => - CometCachedBatch( - numRows = rows.toInt, - sizeInBytes = buffer.size, - stats = stats, - bytes = buffer) + val (arrowBatch, ownsBatch) = + if (CometArrowConverters.isArrowBacked(batch)) { + (batch, false) + } else { + val converted = CometArrowConverters.columnarBatchToArrowBatch( + batch, + structType, + CometArrowStream.NATIVE_TIMEZONE, + CometArrowAllocator) + (converted, true) + } + + try { + // `Utils.serializeBatches` is lazy and clears the batch's vectors as it writes, so the + // result has to be materialized before a converted batch is closed. + Utils + .serializeBatches(Iterator.single(arrowBatch)) + .map { case (rows, buffer) => + CometCachedBatch( + numRows = rows.toInt, + sizeInBytes = buffer.size, + stats = stats, + bytes = buffer) + } + .toList + } finally { + if (ownsBatch) { + arrowBatch.close() + } } } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala index 2d4fd71376..8eb97b2eb0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala @@ -26,9 +26,9 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch} -import org.apache.comet.vector.NativeUtil +import org.apache.comet.vector.{CometVector, NativeUtil} /** * Convert a stream of Spark `InternalRow`s to a stream of independently-owned Arrow @@ -74,4 +74,45 @@ object CometArrowConverters extends Logging { } } } + + /** + * Copy a Spark `ColumnarBatch` whose columns are not Arrow-backed (e.g. + * `On/OffHeapColumnVector` from Spark's vectorized Parquet reader, or a third-party connector's + * vectors) into a freshly allocated Arrow `ColumnarBatch` of `CometVector`s. + * + * The input batch is not consumed or closed; the caller owns the returned batch and must close + * it. Values are copied element-wise, since Spark's `ColumnVector` implementations do not + * expose Arrow buffers. + */ + def columnarBatchToArrowBatch( + batch: ColumnarBatch, + schema: StructType, + timeZoneId: String, + allocator: BufferAllocator): ColumnarBatch = { + val arrowSchema: Schema = Utils.toArrowSchema(schema, timeZoneId) + val root = VectorSchemaRoot.create(arrowSchema, allocator) + val writer = ArrowWriter.create(root) + val numRows = batch.numRows() + var col = 0 + while (col < batch.numCols()) { + val column = batch.column(col) + val columnArray = new ColumnarArray(column, 0, numRows) + if (column.hasNull) { + writer.writeCol(columnArray, col) + } else { + writer.writeColNoNull(columnArray, col) + } + col += 1 + } + writer.finish() + // ArrowWriter derives the root row count from its per-column writes, so a zero-column input + // batch (Spark's count-from-metadata scan: numRows > 0, numCols == 0) would otherwise produce + // a root with rowCount == 0 and silently drop the rows. + root.setRowCount(numRows) + NativeUtil.rootAsBatch(root) + } + + /** Whether every column in `batch` is already an Arrow-backed `CometVector`. */ + def isArrowBacked(batch: ColumnarBatch): Boolean = + (0 until batch.numCols()).forall(i => batch.column(i).isInstanceOf[CometVector]) } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index d5dc8c282d..a57f74f956 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -688,4 +688,94 @@ class CometInMemoryCacheSuite extends CometTestBase { spark.catalog.clearCache() } } + + test("cache a Spark columnar plan whose vectors are not Arrow-backed") { + withTempPath { path => + spark + .range(1000) + .selectExpr( + "id as key", + "id % 8 as value", + "cast(id as string) as s", + "cast(id as double) as d", + "cast(null as int) as n", + "cast(id as decimal(20,3)) as dec", + "date_add(date'2020-01-01', cast(id as int)) as dt", + "timestamp_micros(id * 1000000) as ts") + .write + .parquet(path.toString) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + // Force the cached plan to be Spark's own vectorized Parquet reader, which produces + // On/OffHeapColumnVector rather than CometVector. Spark's InMemoryRelation strips the + // ColumnarToRow above it because supportsColumnarInput is true for this schema, so the + // serializer receives non-Arrow columnar batches. + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Denver") { + + spark.catalog.clearCache() + + spark.read.parquet(path.toString).createOrReplaceTempView("spark_columnar_cache") + + spark.catalog.cacheTable("spark_columnar_cache") + assert(spark.table("spark_columnar_cache").count() == 1000) + + assert( + cachedBatchTypes("spark_columnar_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + checkSparkAnswer( + spark.sql( + "SELECT * FROM spark_columnar_cache WHERE key >= 10 AND key < 20 ORDER BY key")) + checkSparkAnswer( + spark.sql("SELECT sum(key), sum(d), sum(dec), count(s), count(n), max(dt), max(ts) " + + "FROM spark_columnar_cache")) + + spark.catalog.clearCache() + } + } + } + + test("cache a non-Arrow-backed Spark columnar plan with complex types") { + withTempPath { path => + spark + .range(200) + .selectExpr( + "id as key", + "if(id % 5 = 0, null, array(id, id + 1)) as a", + "named_struct('x', id, 'y', cast(id as string)) as st", + "if(id % 7 = 0, null, map(cast(id as string), id)) as m", + "cast(id as binary) as b") + .write + .parquet(path.toString) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true", + SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true") { + + spark.catalog.clearCache() + + spark.read.parquet(path.toString).createOrReplaceTempView("spark_columnar_complex") + + spark.catalog.cacheTable("spark_columnar_complex") + assert(spark.table("spark_columnar_complex").count() == 200) + + checkSparkAnswer(spark.sql("SELECT key, a, st, m, b FROM spark_columnar_complex")) + + spark.catalog.clearCache() + } + } + } } From c7bc49e52ed381cfe4482e798025ce3d0ca6e6e3 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 29 Jul 2026 14:40:56 -0600 Subject: [PATCH 6/8] refactor: dedup Spark-columnar-to-Arrow copy and tidy cache serializer Extract the element-wise Spark `ColumnVector` -> Arrow copy loop that `columnarBatchToArrowBatch` had duplicated from `SparkColumnarArrowReader` into a shared `CometArrowConverters.writeColumns`, so the zero-column rowCount workaround lives in one place. Also: - hoist the Arrow schema out of the per-batch path - replace the `(batch, ownsBatch)` tuple and `.toList` with a `serializeBatch` helper that takes the single element eagerly - move `isArrowBacked` to `Utils`, next to the `getBatchFieldVectors` precondition it describes - drop `toStructType` in favour of the existing `Utils.fromAttributes` - document at `supportsColumnarInput` why its schema-only answer means the conversion in `encodeBatches` is load bearing - route the new tests through the existing `withNativeCache` helper and assert both of them really cached in Comet's format --- .../arrow/ArrowCachedBatchSerializer.scala | 76 +++++------- .../arrow/CometArrowConverters.scala | 61 ++++++---- .../arrow/SparkColumnarArrowReader.scala | 29 ++--- .../apache/spark/sql/comet/util/Utils.scala | 9 ++ .../comet/exec/CometInMemoryCacheSuite.scala | 113 ++++++++---------- 5 files changed, 137 insertions(+), 151 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 22069bd1d0..c2cdbe4cd1 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -70,13 +70,6 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { private val fallback = new DefaultCachedBatchSerializer() - // Row-to-Arrow conversion needs a StructType, while cache APIs pass attributes. - private def toStructType(schema: Seq[Attribute]): StructType = { - StructType(schema.map { attr => - StructField(attr.name, attr.dataType, attr.nullable, attr.metadata) - }) - } - // Build the statistics row expected by SimpleMetricsCachedBatchSerializer. // For each cached column Spark expects five values in this order: // lower bound, upper bound, null count, row count, and size in bytes. @@ -181,52 +174,40 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // The stats are stored beside the Arrow bytes so Spark's cache filter can prune // CometCachedBatch without decoding the batch first. // - // Spark decides the input path from `supportsColumnarInput`, which only sees the schema, so a - // columnar input batch is not guaranteed to be Arrow-backed: any Spark or third-party columnar - // leaf (Spark's vectorized Parquet/ORC reader, a connector's own vectors) reaches - // `convertColumnarBatchToCachedBatch` with `On/OffHeapColumnVector`-style columns, which - // `Utils.serializeBatches` cannot write. Copy those into Arrow first. + // A columnar input batch is not guaranteed to be Arrow-backed; see supportsColumnarInput for + // why. Batches that are not get copied into Arrow first, since Utils.serializeBatches only + // writes CometVector columns. private def encodeBatches( batches: Iterator[ColumnarBatch], attrs: Seq[Attribute]): Iterator[CachedBatch] = { - lazy val structType = toStructType(attrs) + val arrowSchema = + Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE) - batches.flatMap { batch => + batches.map { batch => val stats = computeStats(batch, attrs) - val (arrowBatch, ownsBatch) = - if (CometArrowConverters.isArrowBacked(batch)) { - (batch, false) - } else { - val converted = CometArrowConverters.columnarBatchToArrowBatch( - batch, - structType, - CometArrowStream.NATIVE_TIMEZONE, - CometArrowAllocator) - (converted, true) - } - - try { - // `Utils.serializeBatches` is lazy and clears the batch's vectors as it writes, so the - // result has to be materialized before a converted batch is closed. - Utils - .serializeBatches(Iterator.single(arrowBatch)) - .map { case (rows, buffer) => - CometCachedBatch( - numRows = rows.toInt, - sizeInBytes = buffer.size, - stats = stats, - bytes = buffer) - } - .toList - } finally { - if (ownsBatch) { - arrowBatch.close() - } + if (Utils.isArrowBacked(batch)) { + serializeBatch(batch, stats) + } else { + val arrowBatch = + CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) + try serializeBatch(arrowBatch, stats) + finally arrowBatch.close() } } } + // Utils.serializeBatches is one-in/one-out, so take the single element eagerly: the write has to + // happen before a converted batch is closed. + private def serializeBatch(batch: ColumnarBatch, stats: InternalRow): CachedBatch = { + val (rows, buffer) = Utils.serializeBatches(Iterator.single(batch)).next() + CometCachedBatch( + numRows = rows.toInt, + sizeInBytes = buffer.size, + stats = stats, + bytes = buffer) + } + // Resolve requested columns by exprId, not by name, because aliases may reuse names. private def selectedIndices( cacheAttributes: Seq[Attribute], @@ -285,6 +266,13 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // Comet's Arrow writer only handles the types listed in supportsSchema. Reporting false here // sends the relation down the row path, where it is delegated to Spark's default serializer, // instead of failing at cache materialization inside Utils.serializeBatches. + // + // This answer is schema-only, because attributes are all Spark gives us; it says nothing about + // the vectors. Returning true also makes InMemoryRelation strip the ColumnarToRow above the + // cached plan, so convertColumnarBatchToCachedBatch then receives whatever that plan produces: + // a Comet scan's CometVectors, but equally Spark's vectorized Parquet/ORC reader or a + // connector's own vectors. encodeBatches converts the non-Arrow ones; that conversion is load + // bearing, not defensive. override def supportsColumnarInput(schema: Seq[Attribute]): Boolean = supportsSchema(schema) // A relation Comet stores is always readable as columnar Arrow. Anything else holds @@ -356,7 +344,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { input.mapPartitions { rows => val iter = CometArrowConverters.rowToArrowBatchIter( rows, - toStructType(schema), + Utils.fromAttributes(schema), batchSize, sessionTz, CometArrowAllocator) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala index 8eb97b2eb0..eb6cc5cc35 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala @@ -28,17 +28,19 @@ import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch} -import org.apache.comet.vector.{CometVector, NativeUtil} +import org.apache.comet.vector.NativeUtil /** - * Convert a stream of Spark `InternalRow`s to a stream of independently-owned Arrow - * `ColumnarBatch`es: each emitted batch owns a fresh `VectorSchemaRoot` with newly allocated - * buffers and the consumer is responsible for closing it. + * Convert Spark data that is not Arrow-backed (`InternalRow`s, or `ColumnarBatch`es whose columns + * are Spark/third-party `ColumnVector`s) into independently-owned Arrow `ColumnarBatch`es: each + * emitted batch owns a fresh `VectorSchemaRoot` with newly allocated buffers and the consumer is + * responsible for closing it. * - * This differs from [[RowArrowReader]], which reuses one stable `VectorSchemaRoot` - * (release-and-replace) so only one batch is valid at a time. Use this when multiple emitted - * batches must be alive simultaneously (e.g. tests that buffer several batches before consuming). - * Buffers come from the caller-provided `BufferAllocator`, whose lifecycle the caller owns. + * This differs from [[RowArrowReader]] and [[SparkColumnarArrowReader]], which reuse one stable + * `VectorSchemaRoot` (release-and-replace) so only one batch is valid at a time. Use this when + * multiple emitted batches must be alive simultaneously (e.g. tests that buffer several batches + * before consuming). Buffers come from the caller-provided `BufferAllocator`, whose lifecycle the + * caller owns. */ object CometArrowConverters extends Logging { @@ -76,27 +78,22 @@ object CometArrowConverters extends Logging { } /** - * Copy a Spark `ColumnarBatch` whose columns are not Arrow-backed (e.g. - * `On/OffHeapColumnVector` from Spark's vectorized Parquet reader, or a third-party connector's - * vectors) into a freshly allocated Arrow `ColumnarBatch` of `CometVector`s. + * Copy `numRows` rows starting at `startRow` from a Spark `ColumnarBatch` into `root`. * - * The input batch is not consumed or closed; the caller owns the returned batch and must close - * it. Values are copied element-wise, since Spark's `ColumnVector` implementations do not - * expose Arrow buffers. + * Spark's `ColumnVector` implementations do not expose Arrow buffers, so values are necessarily + * copied element-wise. Shared by [[SparkColumnarArrowReader]], which slices into a stable root, + * and [[columnarBatchToArrowBatch]], which fills a fresh one. */ - def columnarBatchToArrowBatch( + private[arrow] def writeColumns( + root: VectorSchemaRoot, batch: ColumnarBatch, - schema: StructType, - timeZoneId: String, - allocator: BufferAllocator): ColumnarBatch = { - val arrowSchema: Schema = Utils.toArrowSchema(schema, timeZoneId) - val root = VectorSchemaRoot.create(arrowSchema, allocator) + startRow: Int, + numRows: Int): Unit = { val writer = ArrowWriter.create(root) - val numRows = batch.numRows() var col = 0 while (col < batch.numCols()) { val column = batch.column(col) - val columnArray = new ColumnarArray(column, 0, numRows) + val columnArray = new ColumnarArray(column, startRow, numRows) if (column.hasNull) { writer.writeCol(columnArray, col) } else { @@ -109,10 +106,22 @@ object CometArrowConverters extends Logging { // batch (Spark's count-from-metadata scan: numRows > 0, numCols == 0) would otherwise produce // a root with rowCount == 0 and silently drop the rows. root.setRowCount(numRows) - NativeUtil.rootAsBatch(root) } - /** Whether every column in `batch` is already an Arrow-backed `CometVector`. */ - def isArrowBacked(batch: ColumnarBatch): Boolean = - (0 until batch.numCols()).forall(i => batch.column(i).isInstanceOf[CometVector]) + /** + * Copy a Spark `ColumnarBatch` whose columns are not Arrow-backed (e.g. + * `On/OffHeapColumnVector` from Spark's vectorized Parquet reader, or a third-party connector's + * vectors) into a freshly allocated Arrow `ColumnarBatch` of `CometVector`s. + * + * The input batch is not consumed or closed; the caller owns the returned batch and must close + * it. + */ + def columnarBatchToArrowBatch( + batch: ColumnarBatch, + arrowSchema: Schema, + allocator: BufferAllocator): ColumnarBatch = { + val root = VectorSchemaRoot.create(arrowSchema, allocator) + writeColumns(root, batch, 0, batch.numRows()) + NativeUtil.rootAsBatch(root) + } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala index 157aca7423..b2a6e048d7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala @@ -22,13 +22,13 @@ package org.apache.spark.sql.comet.execution.arrow import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.Schema -import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch} +import org.apache.spark.sql.vectorized.ColumnarBatch /** * `ArrowReader` over an iterator of Spark-side `ColumnarBatch`es (not Arrow-backed). Slices up to * `maxRecordsPerBatch` rows per `loadNextBatch` from the current Spark batch into the reader's - * stable VSR via `ArrowWriter.writeCol`. Spark's `ColumnVector` implementations aren't Arrow - * buffers, so this reader necessarily copies element values into Arrow format. + * stable VSR via `CometArrowConverters.writeColumns`. Spark's `ColumnVector` implementations + * aren't Arrow buffers, so this reader necessarily copies element values into Arrow format. */ private[comet] class SparkColumnarArrowReader( allocator: BufferAllocator, @@ -76,26 +76,13 @@ private[comet] class SparkColumnarArrowReader( if (maxRecordsPerBatch <= 0) rowsRemaining else math.min(maxRecordsPerBatch, rowsRemaining) - val writer = ArrowWriter.create(getVectorSchemaRoot) - var col = 0 - while (col < current.numCols()) { - val column = current.column(col) - val columnArray = new ColumnarArray(column, rowsConsumedInCurrent, rowsToProduce) - if (column.hasNull) { - writer.writeCol(columnArray, col) - } else { - writer.writeColNoNull(columnArray, col) - } - col += 1 - } + CometArrowConverters.writeColumns( + getVectorSchemaRoot, + current, + rowsConsumedInCurrent, + rowsToProduce) rowsConsumedInCurrent += rowsToProduce - writer.finish() - // ArrowWriter derives the root row count from its per-column writes, so a zero-column - // input batch (Spark's count-from-metadata scan: numRows > 0, numCols == 0) would otherwise - // produce a root with rowCount == 0 and silently drop the rows. Set rowCount explicitly so - // downstream aggregations (e.g. df.count()) see the correct value. - getVectorSchemaRoot.setRowCount(rowsToProduce) onConversionNs(System.nanoTime() - startNs) true } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 164d53fe19..8844605505 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -397,6 +397,15 @@ object Utils extends CometTypeShim with Logging { } } + /** + * Whether every column in `batch` satisfies [[getBatchFieldVectors]]'s precondition, i.e. is an + * Arrow-backed `CometVector`. Callers that may receive batches from a plan they did not build + * (e.g. Comet's cache serializer, which Spark hands the cached plan's columnar output) use this + * to convert foreign vectors to Arrow instead of tripping the exception below. + */ + def isArrowBacked(batch: ColumnarBatch): Boolean = + (0 until batch.numCols()).forall(i => batch.column(i).isInstanceOf[CometVector]) + def getBatchFieldVectors( batch: ColumnarBatch): (Seq[FieldVector], Option[DictionaryProvider]) = { var provider: Option[DictionaryProvider] = None diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index a57f74f956..023fc97548 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -689,8 +689,44 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("cache a Spark columnar plan whose vectors are not Arrow-backed") { + /** + * Cache `view` over a Parquet file written by `write`, with the cached plan forced to be + * Spark's own vectorized Parquet reader: its columns are On/OffHeapColumnVector rather than + * CometVector. Spark's InMemoryRelation strips the ColumnarToRow above that scan because + * supportsColumnarInput is true for the schema, so the serializer receives non-Arrow columnar + * batches. Asserts the relation really was stored in Comet's format before handing control to + * `f`. + */ + private def withSparkColumnarCache(view: String, extraConfs: (String, String)*)( + write: String => Unit)(f: => Unit): Unit = { withTempPath { path => + write(path.toString) + + withNativeCache { + withSQLConf( + Seq( + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true") ++ extraConfs: _*) { + + spark.read.parquet(path.toString).createOrReplaceTempView(view) + spark.catalog.cacheTable(view) + spark.table(view).count() + + assert( + cachedBatchTypes(view).sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + f + } + } + } + } + + test("cache a Spark columnar plan whose vectors are not Arrow-backed") { + withSparkColumnarCache( + "spark_columnar_cache", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Denver") { path => spark .range(1000) .selectExpr( @@ -703,47 +739,22 @@ class CometInMemoryCacheSuite extends CometTestBase { "date_add(date'2020-01-01', cast(id as int)) as dt", "timestamp_micros(id * 1000000) as ts") .write - .parquet(path.toString) - - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - CometConf.COMET_SHUFFLE_MODE.key -> "jvm", - SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", - // Force the cached plan to be Spark's own vectorized Parquet reader, which produces - // On/OffHeapColumnVector rather than CometVector. Spark's InMemoryRelation strips the - // ColumnarToRow above it because supportsColumnarInput is true for this schema, so the - // serializer receives non-Arrow columnar batches. - CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", - CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false", - SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true", - SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Denver") { - - spark.catalog.clearCache() - - spark.read.parquet(path.toString).createOrReplaceTempView("spark_columnar_cache") - - spark.catalog.cacheTable("spark_columnar_cache") - assert(spark.table("spark_columnar_cache").count() == 1000) - - assert( - cachedBatchTypes("spark_columnar_cache").sameElements( - Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) - - checkSparkAnswer( - spark.sql( - "SELECT * FROM spark_columnar_cache WHERE key >= 10 AND key < 20 ORDER BY key")) - checkSparkAnswer( - spark.sql("SELECT sum(key), sum(d), sum(dec), count(s), count(n), max(dt), max(ts) " + - "FROM spark_columnar_cache")) - - spark.catalog.clearCache() - } + .parquet(path) + } { + assert(spark.table("spark_columnar_cache").count() == 1000) + + checkSparkAnswer( + spark.sql("SELECT * FROM spark_columnar_cache WHERE key >= 10 AND key < 20 ORDER BY key")) + checkSparkAnswer( + spark.sql("SELECT sum(key), sum(d), sum(dec), count(s), count(n), max(dt), max(ts) " + + "FROM spark_columnar_cache")) } } test("cache a non-Arrow-backed Spark columnar plan with complex types") { - withTempPath { path => + withSparkColumnarCache( + "spark_columnar_complex", + SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true") { path => spark .range(200) .selectExpr( @@ -753,29 +764,11 @@ class CometInMemoryCacheSuite extends CometTestBase { "if(id % 7 = 0, null, map(cast(id as string), id)) as m", "cast(id as binary) as b") .write - .parquet(path.toString) - - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - CometConf.COMET_SHUFFLE_MODE.key -> "jvm", - SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", - CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", - CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false", - SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true", - SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true") { - - spark.catalog.clearCache() + .parquet(path) + } { + assert(spark.table("spark_columnar_complex").count() == 200) - spark.read.parquet(path.toString).createOrReplaceTempView("spark_columnar_complex") - - spark.catalog.cacheTable("spark_columnar_complex") - assert(spark.table("spark_columnar_complex").count() == 200) - - checkSparkAnswer(spark.sql("SELECT key, a, st, m, b FROM spark_columnar_complex")) - - spark.catalog.clearCache() - } + checkSparkAnswer(spark.sql("SELECT key, a, st, m, b FROM spark_columnar_complex")) } } } From 1367665bd3679bf8cb4d4b369d51efe1a24f3442 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 31 Jul 2026 14:30:49 -0600 Subject: [PATCH 7/8] fix: honor cache pruning config, close readers on early termination Three review fixes on the cache path, plus tests for two of them. Honor spark.sql.inMemoryColumnarStorage.partitionPruning. CometInMemoryTableScanExec applied the serializer's stats filter unconditionally, so setting the config to false still pruned. Spark's InMemoryTableScanExec.filteredCachedBatches gates on it, and a user reaching for that knob -- to rule out a stats bug, say -- is specifically trying to stop pruning, so silently ignoring it makes Comet diverge on the one thing they are controlling. Close Arrow readers when a consumer stops early. ArrowReaderIterator closes its reader only on exhaustion, so LIMIT, take() or a cancelled task left the reader it was part-way through open. convertCachedBatchToColumnarBatch now registers a TaskCompletionListener, as Spark's own ArrowCachedBatchSerializer does. flatMap consumes each inner iterator fully before building the next, so at most one reader is open at a time and tracking the current one suffices; close() is already idempotent and synchronized, so closing an exhausted one is a no-op. Release the VectorSchemaRoot if conversion throws. columnarBatchToArrowBatch allocated a root before NativeUtil.rootAsBatch wrapped it, and the caller only owns the returned batch, so a throw from writeColumns leaked the allocation. rowToArrowBatchIter had the same shape and gets the same guard. Tests: - "honors inMemoryColumnarStorage.partitionPruning=false" observes the scan's numOutputRows, which counts rows in the batches actually decoded: 100 with pruning on (one batch), 1000 with it off. Note the metric has to be read after forcing this df to run, since checkSparkAnswer executes its own copies and leaves the collected plan instance at zero. - "supports DISK_ONLY storage level" asserts memSize == 0, diskSize > 0, all partitions cached, the payload is still CometCachedBatch, and the cache reads back through CometInMemoryTableScan. CometCachedBatch holds a ChunkedByteBuffer, which is Externalizable, so BlockManager can spill it like any other block. All 19 tests in CometInMemoryCacheSuite pass on spark-3.5. --- .../comet/CometInMemoryTableScanExec.scala | 8 +- .../arrow/ArrowCachedBatchSerializer.scala | 28 ++++- .../arrow/CometArrowConverters.scala | 38 ++++-- .../comet/exec/CometInMemoryCacheSuite.scala | 109 ++++++++++++++++++ 4 files changed, 170 insertions(+), 13 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala index c7032579df..7cc1d74ead 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -71,11 +71,17 @@ case class CometInMemoryTableScanExec( // Apply Spark's cache batch filter before decoding. Spark's InMemoryTableScanExec does this in // filteredCachedBatches(), but that method is private. Reusing the serializer's buildFilter here // keeps Comet on the same stats-based pruning path instead of decoding every cached batch. + // + // Gated on conf.inMemoryPartitionPruning the same way Spark's filteredCachedBatches is, so + // spark.sql.inMemoryColumnarStorage.partitionPruning=false disables pruning here too. Pruning is + // normally a win, but the config exists to be able to turn it off -- for debugging a suspected + // stats bug, for instance -- and silently ignoring it would make Comet diverge from Spark on a + // knob a user reaching for it is specifically trying to control. override def doExecuteColumnar(): RDD[ColumnarBatch] = { val numOutputRows = longMetric("numOutputRows") val filteredBuffers = - if (originalPlan.predicates.nonEmpty) { + if (originalPlan.predicates.nonEmpty && conf.inMemoryPartitionPruning) { val filter = serializer.buildFilter(originalPlan.predicates, relationOutput) cachedBuffers.mapPartitionsWithIndex(filter) } else { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index c2cdbe4cd1..06b70f28da 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -21,6 +21,7 @@ package org.apache.spark.sql.comet.execution.arrow import scala.collection.JavaConverters._ +import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull, UnsafeProjection} @@ -315,10 +316,33 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val indices = selectedIndices(cacheAttributes, selectedAttributes) input.mapPartitions { it => + // An ArrowReaderIterator closes its reader (releasing the batch it is holding) only when it + // runs to exhaustion. A consumer that stops early -- LIMIT, take(), or a cancelled task -- + // leaves the reader it was part-way through open, so close it on task completion. Spark's + // own ArrowCachedBatchSerializer registers a listener for the same reason. + // + // flatMap consumes each inner iterator fully before building the next, so at most one reader + // is open at a time and tracking the current one is enough. close() is idempotent, so closing + // one that already exhausted itself is a no-op. + @volatile var current: ArrowReaderIterator = null + Option(TaskContext.get()).foreach { tc => + tc.addTaskCompletionListener[Unit] { _ => + val reader = current + current = null + if (reader != null) { + reader.close() + } + } + } + it.flatMap { case cb: CometCachedBatch => - Utils.decodeBatches(cb.bytes, "CometCache").map { batch => - projectBatch(batch, indices) + Utils.decodeBatches(cb.bytes, "CometCache") match { + case reader: ArrowReaderIterator => + current = reader + reader.map(batch => projectBatch(batch, indices)) + case empty => + empty.map(batch => projectBatch(batch, indices)) } case other => diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala index eb6cc5cc35..b492fbafbe 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.comet.execution.arrow +import scala.util.control.NonFatal + import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.types.pojo.Schema @@ -64,15 +66,23 @@ object CometArrowConverters extends Logging { override def next(): ColumnarBatch = { val root = VectorSchemaRoot.create(arrowSchema, allocator) - val writer = ArrowWriter.create(root) - var rowCount = 0L - while (rowIter.hasNext && - (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch)) { - writer.write(rowIter.next()) - rowCount += 1 + // Same ownership rule as columnarBatchToArrowBatch: the caller only owns the batch that + // rootAsBatch returns, so a throw from writing a row has to release the root here. + try { + val writer = ArrowWriter.create(root) + var rowCount = 0L + while (rowIter.hasNext && + (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch)) { + writer.write(rowIter.next()) + rowCount += 1 + } + writer.finish() + NativeUtil.rootAsBatch(root) + } catch { + case NonFatal(e) => + root.close() + throw e } - writer.finish() - NativeUtil.rootAsBatch(root) } } } @@ -121,7 +131,15 @@ object CometArrowConverters extends Logging { arrowSchema: Schema, allocator: BufferAllocator): ColumnarBatch = { val root = VectorSchemaRoot.create(arrowSchema, allocator) - writeColumns(root, batch, 0, batch.numRows()) - NativeUtil.rootAsBatch(root) + // The caller only owns the returned batch, so anything that throws before `rootAsBatch` wraps + // the root has to release it here or the allocation leaks. + try { + writeColumns(root, batch, 0, batch.numRows()) + NativeUtil.rootAsBatch(root) + } catch { + case NonFatal(e) => + root.close() + throw e + } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 023fc97548..5679588474 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Expression, GreaterThanOr import org.apache.spark.sql.columnar.SimpleMetricsCachedBatch import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.storage.StorageLevel import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus @@ -412,6 +413,114 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache honors inMemoryColumnarStorage.partitionPruning=false") { + // CometInMemoryTableScanExec applies the serializer's stats filter before decoding, the same + // way Spark's InMemoryTableScanExec.filteredCachedBatches does. Spark gates that on + // spark.sql.inMemoryColumnarStorage.partitionPruning, so Comet must too. + // + // Pruning is transparent in the results, so it is observed through the scan's numOutputRows: + // that counts the rows in the batches actually decoded, so pruning fewer batches means fewer + // rows. With pruning off, every cached row must be decoded. + def scanRowsFor(pruning: Boolean): (Long, Long) = { + var result: (Long, Long) = (0L, 0L) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "100", + SQLConf.IN_MEMORY_PARTITION_PRUNING.key -> pruning.toString) { + + spark.catalog.clearCache() + spark + .range(0, 1000, 1, 10) + .selectExpr("id as key", "id % 7 as value") + .createOrReplaceTempView("prune_conf_cache") + spark.catalog.cacheTable("prune_conf_cache") + val totalRows = spark.table("prune_conf_cache").count() + + val df = + spark.sql("SELECT key, value FROM prune_conf_cache WHERE key >= 900 AND key < 905") + checkSparkAnswer(df) + + val scans = df.queryExecution.executedPlan.collect { + case s: org.apache.spark.sql.comet.CometInMemoryTableScanExec => s + } + assert(scans.length == 1, s"expected one CometInMemoryTableScan, got ${scans.length}") + // scalastyle:off println + println( + "DIAG rows=" + df.collect().length + " metrics=" + scans.head.metrics + .map { case (k, v) => k + "=" + v.value } + .mkString(",")) + println("DIAG plan=" + df.queryExecution.executedPlan.getClass.getName) + // scalastyle:on println + result = (scans.head.metrics("numOutputRows").value, totalRows) + spark.catalog.clearCache() + } + result + } + + val (prunedRows, total) = scanRowsFor(pruning = true) + val (unprunedRows, total2) = scanRowsFor(pruning = false) + assert(total == total2) + // With pruning on, only the batch holding keys 900-904 is decoded. + assert(prunedRows < total, s"expected pruning to decode fewer than $total rows") + // With pruning off, every cached batch is decoded. + assert( + unprunedRows == total, + s"expected all $total rows to be decoded with pruning disabled, got $unprunedRows") + } + + test("Comet in-memory cache supports DISK_ONLY storage level") { + // CometCachedBatch holds a ChunkedByteBuffer, which is Externalizable, so BlockManager can + // spill it to the DiskStore like any other cached block. Pins that: nothing is held in memory, + // the bytes really do land on disk, every partition is cached, and the cache still reads back + // through the native scan. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + spark + .range(0, 1000, 1, 4) + .selectExpr("id as key", "id % 7 as value") + .createOrReplaceTempView("disk_cache") + + spark.catalog.cacheTable("disk_cache", StorageLevel.DISK_ONLY) + val total = spark.table("disk_cache").count() + assert(total == 1000) + + assert( + cachedBatchTypes("disk_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")), + "DISK_ONLY must still use Comet's cached batch format") + + val cached = + spark.sharedState.cacheManager.lookupCachedData(spark.table("disk_cache")).get + val rddId = cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers.id + val info = spark.sparkContext.getRDDStorageInfo + .find(_.id == rddId) + .getOrElse(fail(s"no storage info for cached RDD $rddId")) + + assert(info.memSize == 0, s"expected nothing in memory, got ${info.memSize} bytes") + assert(info.diskSize > 0, "expected the cached bytes to be on disk") + assert( + info.numCachedPartitions == info.numPartitions, + s"expected all ${info.numPartitions} partitions cached, got ${info.numCachedPartitions}") + + val df = spark.sql("SELECT key, value FROM disk_cache WHERE key >= 900 AND key < 905") + checkSparkAnswer(df) + val plan = df.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + + spark.catalog.clearCache() + } + } + test("Comet plugin respects user-provided cache serializer") { val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key val cometSerializer = From c25c7b32f99604fd16b1243aeb4a6fda1f8c3621 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 31 Jul 2026 15:02:44 -0600 Subject: [PATCH 8/8] fix: label cached timestamps UTC on the row write path too The row write path passed conf.sessionLocalTimeZone into rowToArrowBatchIter while the columnar path already encoded with CometArrowStream.NATIVE_TIMEZONE, so the same logical cache had two physical formats depending on which path filled it, and the row one persisted the writing session's timezone into cached data. Unlike Spark's Arrow cache, whose RecordBatch is deliberately schema-less and reconstructs the timezone on read, CometCachedBatch stores a full IPC stream including the schema, so that label really is written down. Standardise on NATIVE_TIMEZONE ("UTC") for TimestampType, per the analysis on the review thread. This is a label only. Spark's internal timestamp representation is micros since the Unix epoch regardless of session timezone, so no values are converted, and matching Comet's native schema also avoids a cast at the native boundary. TimestampNTZType already had no timezone: Utils.toArrowType maps it to Timestamp(MICROSECOND, null) whatever is passed in. The closure no longer needs the session timezone at all, so it no longer captures anything derived from `conf`. Test: "stores timestamps with a UTC schema label" caches a row-based (local Seq) plan under two non-UTC session timezones, decodes the cached batches back through the serializer, and asserts the Arrow field carries "UTC" -- then checks the values and their string rendering still match Spark. Reverting the one-line change fails it with `got [America/Los_Angeles]`, so it pins the format rather than passing vacuously. All 20 tests in CometInMemoryCacheSuite pass on spark-3.5. --- .../arrow/ArrowCachedBatchSerializer.scala | 11 ++- .../comet/exec/CometInMemoryCacheSuite.scala | 83 +++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 06b70f28da..e6d5aa06cf 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -363,14 +363,21 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { fallback.convertInternalRowToCachedBatch(input, schema, storageLevel, conf) } else { val batchSize = conf.columnBatchSize - val sessionTz = conf.sessionLocalTimeZone input.mapPartitions { rows => val iter = CometArrowConverters.rowToArrowBatchIter( rows, Utils.fromAttributes(schema), batchSize, - sessionTz, + // NATIVE_TIMEZONE ("UTC"), not conf.sessionLocalTimeZone, so both write paths produce + // the same physical format: the columnar path above already encodes with + // NATIVE_TIMEZONE. Unlike Spark's Arrow cache, whose RecordBatch is deliberately + // schema-less, CometCachedBatch stores a full IPC stream including the schema, so a + // session-local label would persist the writing session's mutable timezone into cached + // data. This is a label only: Spark's internal timestamp representation is micros since + // the Unix epoch regardless of session timezone, so no values are converted. It also + // matches Comet's native schema, avoiding a cast at the native boundary. + CometArrowStream.NATIVE_TIMEZONE, CometArrowAllocator) encodeBatches(iter, schema) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 5679588474..839ede5281 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -21,6 +21,7 @@ package org.apache.comet.exec import java.{util => ju} +import org.apache.arrow.vector.types.pojo.ArrowType import org.apache.spark.CometDriverPlugin import org.apache.spark.SparkConf import org.apache.spark.sql.CometTestBase @@ -32,9 +33,12 @@ import org.apache.spark.storage.StorageLevel import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.vector.CometVector class CometInMemoryCacheSuite extends CometTestBase { + import testImplicits._ + // `InMemoryRelation` resolves `spark.sql.cache.serializer` once per JVM and memoizes the // instance in a static field. Test suites share a forked JVM, so whichever suite caches a // table first pins the serializer for everything that follows: without this reset the @@ -521,6 +525,85 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache stores timestamps with a UTC schema label") { + // Unlike Spark's Arrow cache, whose RecordBatch is deliberately schema-less, CometCachedBatch + // stores a full IPC stream including the schema. Labelling TimestampType with the writing + // session's timezone would persist a mutable session value into cached data and would make the + // row write path disagree with the columnar one, which already encodes with NATIVE_TIMEZONE. + // So both paths must write "UTC". This is a label only -- Spark stores timestamps as micros + // since the Unix epoch regardless of session timezone -- so values must be unaffected. + Seq("America/Los_Angeles", "Asia/Kolkata").foreach { sessionTz => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> sessionTz) { + + spark.catalog.clearCache() + + // A local Seq gives a row-based plan, so this exercises + // convertInternalRowToCachedBatch rather than the columnar path. + val rows = Seq( + (1, java.sql.Timestamp.valueOf("2024-01-31 12:34:56.789")), + (2, java.sql.Timestamp.valueOf("1970-01-01 00:00:00")), + (3, null)) + rows.toDF("id", "ts").createOrReplaceTempView("ts_cache") + + spark.catalog.cacheTable("ts_cache") + assert(spark.table("ts_cache").count() == 3) + + assert( + cachedBatchTypes("ts_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")), + s"expected Comet cache format for sessionTz=$sessionTz") + + // Decode the cached bytes through the serializer and read the Arrow field metadata back + // out. The timezone has to be extracted inside the closure: ColumnarBatch is not + // serializable. + val relation = + spark.sharedState.cacheManager + .lookupCachedData(spark.table("ts_cache")) + .get + .cachedRepresentation + val tsIndex = relation.output.indexWhere(_.name == "ts") + val labels = relation.cacheBuilder.serializer + .convertCachedBatchToColumnarBatch( + relation.cacheBuilder.cachedColumnBuffers, + relation.output, + relation.output, + spark.sessionState.conf) + .mapPartitions { batches => + batches.take(1).map { batch => + batch.column(tsIndex) match { + case v: CometVector => + v.getValueVector.getField.getType match { + case t: ArrowType.Timestamp => String.valueOf(t.getTimezone) + case other => s"unexpected arrow type $other" + } + case other => s"unexpected vector ${other.getClass.getName}" + } + } + } + .collect() + .distinct + + assert( + labels.sameElements(Array("UTC")), + s"expected the cached timestamp schema to be labelled UTC for sessionTz=$sessionTz, " + + s"got ${labels.mkString("[", ",", "]")}") + + // The label change must not move any values. + checkSparkAnswer(spark.sql("SELECT id, ts FROM ts_cache ORDER BY id")) + checkSparkAnswer( + spark.sql("SELECT id, CAST(ts AS STRING) AS s FROM ts_cache ORDER BY id")) + + spark.catalog.clearCache() + } + } + } + test("Comet plugin respects user-provided cache serializer") { val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key val cometSerializer =