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 01/19] 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 78ea0f01687..dee67a044f7 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 d116d2f4076..114304492ee 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 7290ab436af..82235d380cf 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 00000000000..cb12f2a3acd --- /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 00000000000..5d549301412 --- /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 8cbf7c9189c..eb506ddc6ee 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 00000000000..11174d9bded --- /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 02/19] 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 3fbe052aff0..78ba754a70b 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 dbdd325848d..eb451cd8a0f 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 114304492ee..cc803041653 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 82235d380cf..b849a481d95 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 cb12f2a3acd..3da81d1f54d 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 5d549301412..bfae724c59a 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 11174d9bded..68e22df97dd 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 00000000000..f3ceb831ce1 --- /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 03/19] 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 d1ed04cd5bc..41cfcfb2264 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 e1e42bc137e..c7032579df4 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 ce6071bf1a6..bf60d5de3f0 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 68e22df97dd..24c0f6dbbd2 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 00000000000..ff5240af2a7 --- /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 04/19] 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 ded79c544a5..f520cc756f4 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 41cfcfb2264..780c7d0cb36 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 071934c0cd3..eaadc90f0a8 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 24c0f6dbbd2..d5dc8c282d6 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 05/19] 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 eaadc90f0a8..22069bd1d0f 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 2d4fd713763..8eb97b2eb01 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 d5dc8c282d6..a57f74f9569 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 06/19] 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 22069bd1d0f..c2cdbe4cd1f 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 8eb97b2eb01..eb6cc5cc356 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 157aca74232..b2a6e048d7b 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 164d53fe19b..8844605505a 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 a57f74f9569..023fc975485 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 07/19] 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 c7032579df4..7cc1d74ead5 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 c2cdbe4cd1f..06b70f28da9 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 eb6cc5cc356..b492fbafbe8 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 023fc975485..56795884748 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 08/19] 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 06b70f28da9..e6d5aa06cfa 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 56795884748..839ede52817 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 = From 4b9a5167ae74d5c7f9e8fee20c63c79861c0d2e1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 21:26:37 -0600 Subject: [PATCH 09/19] refactor: do not let a failing root close mask the original error Both Arrow batch producers released the VectorSchemaRoot on failure with a bare root.close(), so an IllegalStateException from close (outstanding child allocations) would replace the failure that actually mattered. Attach it as a suppressed exception instead, matching SparkErrorUtils.tryWithSafeFinally, and share the one guard between the row and columnar paths. --- .../arrow/CometArrowConverters.scala | 30 ++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) 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 b492fbafbe8..43b219ac49a 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 @@ -68,7 +68,7 @@ object CometArrowConverters extends Logging { val root = VectorSchemaRoot.create(arrowSchema, allocator) // 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 { + closingRootOnFailure(root) { val writer = ArrowWriter.create(root) var rowCount = 0L while (rowIter.hasNext && @@ -78,10 +78,6 @@ object CometArrowConverters extends Logging { } writer.finish() NativeUtil.rootAsBatch(root) - } catch { - case NonFatal(e) => - root.close() - throw e } } } @@ -133,12 +129,32 @@ object CometArrowConverters extends Logging { val root = VectorSchemaRoot.create(arrowSchema, allocator) // 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 { + closingRootOnFailure(root) { writeColumns(root, batch, 0, batch.numRows()) NativeUtil.rootAsBatch(root) + } + } + + /** + * Run `body`, closing `root` if it throws. On success the returned batch takes ownership of + * `root`, so it is deliberately left open. + * + * A failing `close` is attached as a suppressed exception rather than replacing the original, + * following `SparkErrorUtils.tryWithSafeFinally`: releasing an Arrow root can itself throw + * (e.g. `IllegalStateException` for outstanding child allocations), and that is the less + * informative of the two failures. + */ + private def closingRootOnFailure(root: VectorSchemaRoot)( + body: => ColumnarBatch): ColumnarBatch = { + try { + body } catch { case NonFatal(e) => - root.close() + try { + root.close() + } catch { + case NonFatal(closeError) => e.addSuppressed(closeError) + } throw e } } From ebefc1cf49cd76cd0ef68448ebd9bd90a60ace09 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 21:26:37 -0600 Subject: [PATCH 10/19] test: drop debug output from pruning test, fix binary cast under ANSI The pruning test carried leftover println diagnostics whose df.collect() was load-bearing: checkSparkAnswer executes its own copies of the query, so the metric read below it is zero unless this plan instance is forced. Make that an explicit call with the reason stated. The complex-types cache test cast bigint directly to binary, which ANSI mode rejects, so it failed on the Spark 4.x profiles. Cast via string. --- .../comet/exec/CometInMemoryCacheSuite.scala | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) 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 839ede52817..15e58ef4bc4 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -448,17 +448,15 @@ class CometInMemoryCacheSuite extends CometTestBase { spark.sql("SELECT key, value FROM prune_conf_cache WHERE key >= 900 AND key < 905") checkSparkAnswer(df) + // checkSparkAnswer takes its argument by name and executes its own copies of the query, + // so this df's plan instance has not run and its metrics are all still zero. Force this + // exact plan before reading them, or the comparison below passes vacuously with 0 == 0. + df.collect() + 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() } @@ -954,7 +952,9 @@ class CometInMemoryCacheSuite extends CometTestBase { "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") + // via string: ANSI mode (on by default in Spark 4.x) rejects a direct bigint -> binary + // cast. + "cast(cast(id as string) as binary) as b") .write .parquet(path) } { From 78ce24b6b0f12a5ad8f6e41cb86b3e94239360b8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 3 Aug 2026 21:33:52 -0600 Subject: [PATCH 11/19] test: update cache assertions for main The projection-only test asserted CometNativeColumnarToRow, but #5114 made the JVM converter the default. Accept either Comet converter, since the point of the assertion is that rows do not come from Spark's ColumnarToRow. Also correct the isArrowBacked docstring: #4532 taught getBatchFieldVectors to materialize a ConstantColumnVector, so isArrowBacked is now stricter than that method's precondition rather than equal to it. --- .../org/apache/spark/sql/comet/util/Utils.scala | 12 ++++++++---- .../apache/comet/exec/CometInMemoryCacheSuite.scala | 7 ++++++- 2 files changed, 14 insertions(+), 5 deletions(-) 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 f0977f20b83..89ad6cc0490 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 @@ -399,10 +399,14 @@ 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. + * Whether every column in `batch` is an Arrow-backed `CometVector`, so [[getBatchFieldVectors]] + * can hand out its vectors directly. 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. + * + * Stricter than what [[getBatchFieldVectors]] accepts: a `ConstantColumnVector` is rejected + * here even though that method materializes one, so such a batch takes the conversion path + * rather than being materialized column by column. */ 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 15e58ef4bc4..e96d8555716 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -304,7 +304,12 @@ class CometInMemoryCacheSuite extends CometTestBase { val plan = df.queryExecution.executedPlan.toString() assert(plan.contains("CometInMemoryTableScan")) - assert(plan.contains("CometNativeColumnarToRow")) + // Rows come out of a Comet converter rather than Spark's ColumnarToRow. Either variant + // satisfies that; which one is used depends on the default of + // spark.comet.exec.columnarToRow.native.enabled. + assert( + plan.contains("CometColumnarToRow") || plan.contains("CometNativeColumnarToRow"), + s"expected a Comet columnar-to-row above the cache scan, got:\n$plan") spark.catalog.clearCache() } From 3c35f7ac017d6dadf278b3a657fdd2285c0c7dc4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 26 Aug 2026 15:24:14 -0600 Subject: [PATCH 12/19] bench: expose cache read cost as projection width narrows A CometCachedBatch is one compressed Arrow IPC stream covering every cached column, so a scan decodes all of them and projects afterwards. Read cost is flat in the width of the projection where Spark's per-column format falls away as it narrows, which the existing cases could not show: both read a Comet-written cache, so neither is a baseline for Spark's own format. Widen the cached relation with string columns and add narrow and full projection cases that bracket the effect, note the missing baseline at the materialization site, and document the read-path cost on spark.comet.exec.inMemoryCache.enabled. --- .../scala/org/apache/comet/CometConf.scala | 6 +++- .../CometInMemoryCacheBenchmark.scala | 29 +++++++++++++++++-- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index ffd1dc8648c..81568cf842f 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -268,7 +268,11 @@ object CometConf extends ShimCometConf { "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.") + "format. Cached batches are stored as one compressed Arrow IPC stream covering every " + + "cached column, so a scan decodes all of them and then projects, while Spark's " + + "default format decodes only the projected columns. Reads that project a few columns " + + "of a wide cached relation, or that feed Spark operators rather than Comet ones, can " + + "therefore be slower than Spark's cache even though materializing it is faster.") .booleanConf .createWithDefault(false) 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 index f3ceb831ce1..9c50e81107b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -64,7 +64,13 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { withTempTable(sourceTable, cacheTable) { spark .range(0, numRows, 1, 16) - .selectExpr("id", "id % 1000 AS k", "id + 1 AS v") + .selectExpr( + "id", + "id % 1000 AS k", + "id + 1 AS v", + "concat('str_a_', cast(id % 100000 as string)) AS s1", + "concat('str_b_', cast(id % 7919 as string)) AS s2", + "concat('str_c_', cast(id as string)) AS s3") .createOrReplaceTempView(sourceTable) runCacheBenchmark( @@ -78,6 +84,18 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { |FROM $cacheTable |WHERE id >= 4500000 AND id < 4750000 """.stripMargin) + + // A CometCachedBatch is one compressed Arrow IPC stream covering every cached column, so a + // scan decodes all of them and projects afterwards. Cost is therefore flat in the width of + // the projection, where Spark's per-column format falls away as it narrows. These two cases + // bracket that: the same cached relation read one column wide and six columns wide. + runCacheBenchmark( + "in-memory cache narrow projection (1 of 6 columns)", + s"SELECT count(k) FROM $cacheTable") + + runCacheBenchmark( + "in-memory cache full projection (6 of 6 columns)", + s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $cacheTable") } } @@ -114,8 +132,15 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { // Materialize the cache once using Comet's cache serializer. // The benchmark measures repeated cache reads by comparing the // fallback read path against CometInMemoryTableScan. + // + // Both cases therefore read a Comet-written cache: spark.sql.cache.serializer is a static + // conf, so a single session cannot also materialize a DefaultCachedBatch to compare against. + // "Comet cache disabled" here means Spark execution over CometCachedBatch, NOT Spark's own + // cache format, and these numbers are not a baseline for it. withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { - spark.sql(s"SELECT id, k, v FROM $sourceTable").createOrReplaceTempView(cacheTable) + spark + .sql(s"SELECT id, k, v, s1, s2, s3 FROM $sourceTable") + .createOrReplaceTempView(cacheTable) spark.catalog.cacheTable(cacheTable) spark.table(cacheTable).count() } From 2f79c43d09e2f1260251df9986b7d1dd1daac033 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 26 Aug 2026 17:05:49 -0600 Subject: [PATCH 13/19] perf: store cached columns as separate streams so reads decode only what they project A CometCachedBatch held one compressed Arrow IPC stream covering every cached column, so convertCachedBatchToColumnarBatch inflated all of them and projected afterwards. Read cost was flat in the width of the projection, where Spark's per-column DefaultCachedBatch falls away as it narrows, and a narrow read of a wide cached relation was several times slower than Spark's cache despite materializing faster. Store one stream per column and decode only the selected ones. An empty selection now stays empty rather than expanding to every column, and the scan asks for a single cheap column instead of the whole schema when a query needs only the row count, since the native plan still requires a non-empty scan schema. Per-column sizes are now known, so the statistics field Spark reserves for them holds the real value. Measured on 5M rows and 6 columns, against Spark's cache format: read shape before after Spark count(*) 241 ms 77 ms 62 ms 1 of 6 59 ms 57 ms 105 ms 3 of 6 205 ms 204 ms 339 ms 6 of 6 448 ms 448 ms 375 ms Framing costs a schema block and compression framing per column per batch, and loses cross-column compression: footprint grows 2.5% at 6 columns and 32% at 60, where the cached relation is still 22x smaller than Spark's format. --- .../scala/org/apache/comet/CometConf.scala | 9 +- .../comet/CometInMemoryTableScanExec.scala | 58 +++-- .../arrow/ArrowCachedBatchSerializer.scala | 223 ++++++++++++------ .../apache/spark/sql/comet/util/Utils.scala | 31 +++ .../comet/exec/CometInMemoryCacheSuite.scala | 175 +++++++++++++- .../CometInMemoryCacheBenchmark.scala | 11 +- .../arrow/CometCachedBatchHelper.scala | 56 +++++ 7 files changed, 463 insertions(+), 100 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 81568cf842f..cf5a88d2359 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -268,11 +268,10 @@ object CometConf extends ShimCometConf { "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. Cached batches are stored as one compressed Arrow IPC stream covering every " + - "cached column, so a scan decodes all of them and then projects, while Spark's " + - "default format decodes only the projected columns. Reads that project a few columns " + - "of a wide cached relation, or that feed Spark operators rather than Comet ones, can " + - "therefore be slower than Spark's cache even though materializing it is faster.") + "format. Each cached column is stored as its own compressed Arrow IPC stream, so a " + + "scan decodes only the columns it projected. Reads that feed Spark operators rather " + + "than Comet ones still pay a row conversion the default format avoids, and can be " + + "slower than Spark's cache.") .booleanConf .createWithDefault(false) 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 7cc1d74ead5..99a22d4d208 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 @@ -27,6 +27,7 @@ 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.types._ import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometConf @@ -57,11 +58,10 @@ 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. + // For an empty-projection scan (`SELECT count(*)`) this is empty while `scanOutput` holds one + // placeholder column, 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; see `scanOutputFor` for why the scan cannot simply be empty. override def output: Seq[Attribute] = originalPlan.output // Use the serializer's vector types because the cached batch layout is owned by the serializer. @@ -107,12 +107,7 @@ object CometInMemoryTableScanExec extends CometOperatorSerde[InMemoryTableScanEx 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 scanTypes = scanOutputFor(op).flatMap(attr => serializeDataType(attr.dataType)) val scanBuilder = OperatorOuterClass.Scan .newBuilder() @@ -127,10 +122,6 @@ object CometInMemoryTableScanExec extends CometOperatorSerde[InMemoryTableScanEx 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( @@ -138,6 +129,41 @@ object CometInMemoryTableScanExec extends CometOperatorSerde[InMemoryTableScanEx relation.cacheBuilder.serializer, relation.cacheBuilder.cachedColumnBuffers, relation.output, - actualOutput)) + scanOutputFor(op))) + } + + /** + * Columns the cache scan asks the serializer to decode. + * + * An empty-output scan (`SELECT count(*)`) still needs a non-empty schema for native planning, + * and the batches the node emits have to match that schema. Falling back to the whole cache + * schema satisfies both, but the serializer decodes exactly what it is asked for, so the + * cheapest query in the workload would decode every cached column. One column is enough: the + * aggregate above an empty-output scan reads the row count and never a value, so pick the + * cheapest to decode rather than all of them. + * + * `convert` and `createExec` must choose identically, or the native scan's declared schema and + * the batches fed to it disagree. + */ + private def scanOutputFor(op: InMemoryTableScanExec): Seq[Attribute] = { + if (op.output.nonEmpty) { + op.output + } else if (op.relation.output.isEmpty) { + Nil + } else { + Seq(op.relation.output.minBy(a => decodeCostRank(a.dataType))) + } + } + + // Rank by how much work decoding a column of this type costs, cheapest first. Fixed-width types + // decode to a flat buffer; variable-width and nested ones carry offsets, children and possibly + // dictionaries. + private def decodeCostRank(dt: DataType): Int = dt match { + case BooleanType | ByteType => 0 + case ShortType => 1 + case IntegerType | FloatType | DateType => 2 + case LongType | DoubleType | TimestampType | TimestampNTZType => 3 + case _: DecimalType => 4 + case _ => 5 } } 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 e6d5aa06cfa..b3c9ba240ee 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 @@ -40,14 +40,17 @@ import org.apache.comet.CometArrowAllocator /** * 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. + * `columns` holds one compressed Arrow stream per cached column, in cache-schema order, produced + * by `Utils.serializeBatchColumns`. Storing columns separately is what lets a scan decode only + * the ones it projected; a single stream covering the whole batch would have to be inflated in + * full before any projection could be applied. 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, override val stats: InternalRow, - bytes: ChunkedByteBuffer) + columns: Array[ChunkedByteBuffer]) extends SimpleMetricsCachedBatch /** @@ -71,10 +74,12 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { private val fallback = new DefaultCachedBatchSerializer() - // 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 = { + // Bounds and null counts per column, gathered before the batch is serialized: serializing + // clears the batch's vectors, and the per-column byte sizes that complete the statistics row + // are only known afterwards. See statsRow. + private def gatherColumnStats( + batch: ColumnarBatch, + attrs: Seq[Attribute]): (Array[Any], Array[Any], Array[Int]) = { val numCols = attrs.length val lower = new Array[Any](numCols) val upper = new Array[Any](numCols) @@ -103,18 +108,31 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { c += 1 } + (lower, upper, nulls) + } + + // 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 statsRow( + lower: Array[Any], + upper: Array[Any], + nulls: Array[Int], + numRows: Int, + columnSizes: Array[Long]): InternalRow = { + val numCols = lower.length val values = new Array[Any](numCols * 5) - c = 0 + var 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 + // Each column is its own compressed stream, so its size is known exactly. Cache pruning + // uses bounds/null-count/row-count rather than this field, but Spark reserves it and + // reports it, so record the real value. + values(base + 4) = columnSizes(c) c += 1 } @@ -185,62 +203,49 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE) batches.map { batch => - val stats = computeStats(batch, attrs) - - if (Utils.isArrowBacked(batch)) { - serializeBatch(batch, stats) + // Bounds and null counts are read from the input batch, which serializing then clears, so + // they have to be gathered first. The row is only assembled once the per-column sizes are + // known. + val (lower, upper, nulls) = gatherColumnStats(batch, attrs) + val numRows = batch.numRows() + + val columns = if (Utils.isArrowBacked(batch)) { + Utils.serializeBatchColumns(batch) } else { val arrowBatch = CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) - try serializeBatch(arrowBatch, stats) + try Utils.serializeBatchColumns(arrowBatch) 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) + val columnSizes = columns.map(_.size) + CometCachedBatch( + numRows = numRows, + sizeInBytes = columnSizes.sum, + stats = statsRow(lower, upper, nulls, numRows, columnSizes), + columns = columns) + } } // Resolve requested columns by exprId, not by name, because aliases may reuse names. + // + // An empty selection stays empty rather than expanding to every column. Spark asks for no + // columns when the query only needs the row count (SELECT count(*)), and since projection now + // decides what gets decoded, expanding it would turn the cheapest possible read into the most + // expensive one. 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()) - } + 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 } // Spark's SimpleMetricsCachedBatchSerializer prunes a batch when the generated partition filter @@ -316,33 +321,35 @@ 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. + // A ColumnReaders closes its readers (releasing the vectors they are holding) only when the + // batch it produced has been consumed. A consumer that stops early -- LIMIT, take(), or a + // cancelled task -- leaves the readers for the batch in flight open, so close them 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 + // flatMap consumes each inner iterator fully before building the next, so at most one batch + // is open at a time and tracking the current one is enough. close() is idempotent, so + // closing one that already released itself is a no-op. + @volatile var current: ColumnReaders = null Option(TaskContext.get()).foreach { tc => tc.addTaskCompletionListener[Unit] { _ => - val reader = current + val readers = current current = null - if (reader != null) { - reader.close() + if (readers != null) { + readers.close() } } } it.flatMap { case cb: CometCachedBatch => - 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)) + if (indices.isEmpty) { + // Nothing to decode: the row count is the whole answer, and it is already here. + Iterator.single(new ColumnarBatch(Array.empty[ColumnVector], cb.numRows)) + } else { + val readers = new ColumnReaders(indices.map(i => cb.columns(i)), cb.numRows) + current = readers + readers.batches } case other => @@ -352,6 +359,76 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } } + // Decodes one selected column stream apiece and stitches the results back into a single batch. + // + // Each stream is self-contained, so the columns a scan did not select are never inflated. The + // decoded vectors stay owned by their readers: closing them releases the batch, which is why + // this yields a single-element iterator that closes on exhaustion, matching what + // ArrowReaderIterator did when the payload was one stream. + private class ColumnReaders(buffers: Array[ChunkedByteBuffer], numRows: Int) { + private val readers: Array[Iterator[ColumnarBatch]] = + buffers.map(Utils.decodeBatches(_, "CometCache")) + private var closed = false + + def close(): Unit = synchronized { + if (!closed) { + closed = true + readers.foreach { + case reader: ArrowReaderIterator => reader.close() + case _ => () + } + } + } + + private def assemble(): ColumnarBatch = { + val columns = new Array[ColumnVector](readers.length) + var i = 0 + while (i < readers.length) { + val reader = readers(i) + if (!reader.hasNext) { + throw new IllegalStateException( + s"Cached column stream $i of ${readers.length} decoded to no batch") + } + val decoded = reader.next() + // Each stream holds exactly one single-column record batch, and every column of a cached + // batch covers the same rows. Check rather than trust: a mismatch would otherwise build a + // batch whose columns disagree on length, which reads as corrupt data far from here. + if (decoded.numCols() != 1) { + throw new IllegalStateException( + s"Cached column stream $i decoded to ${decoded.numCols()} columns, expected 1") + } + if (decoded.numRows() != numRows) { + throw new IllegalStateException( + s"Cached column stream $i decoded ${decoded.numRows()} rows, expected $numRows") + } + columns(i) = decoded.column(0) + i += 1 + } + new ColumnarBatch(columns, numRows) + } + + def batches: Iterator[ColumnarBatch] = new Iterator[ColumnarBatch] { + private var emitted = false + + override def hasNext: Boolean = { + if (emitted) { + close() + false + } else { + true + } + } + + override def next(): ColumnarBatch = { + if (emitted) { + throw new NoSuchElementException + } + emitted = true + assemble() + } + } + } + // Row input is cached in Comet format by converting rows to Arrow batches first. override def convertInternalRowToCachedBatch( input: RDD[InternalRow], 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 89ad6cc0490..33533903307 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 @@ -268,6 +268,37 @@ object Utils extends CometTypeShim with Logging { } } + /** + * Serializes each column of `batch` into its own compressed Arrow IPC stream, in column order. + * + * [[serializeBatches]] writes one stream covering every column, so a reader has to inflate all + * of them before it can project. Comet's in-memory cache stores columns separately instead, so + * a scan decodes only the ones it selected. Each stream is self-contained, including its schema + * and any dictionaries the column needs. + * + * The row count is not recoverable from the result when `batch` has no columns, so callers keep + * it alongside. As with [[serializeBatches]], the batch's vectors are cleared once written. + */ + def serializeBatchColumns(batch: ColumnarBatch): Array[ChunkedByteBuffer] = { + val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch) + val provider = batchProviderOpt.getOrElse(new CDataDictionaryProvider) + val codec = CompressionCodec.createCodec(SparkEnv.get.conf) + + fieldVectors.map { fieldVector => + val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) + val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) + + val root = new VectorSchemaRoot(Seq(fieldVector).asJava) + val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) + writer.start() + writer.writeBatch() + root.clear() + writer.close() + + cbbos.toChunkedByteBuffer + }.toArray + } + /** * Decodes the byte arrays back to ColumnarBatchs and put them into buffer. * 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 e96d8555716..bbeaeca15a8 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -25,10 +25,13 @@ import org.apache.arrow.vector.types.pojo.ArrowType import org.apache.spark.CometDriverPlugin 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.catalyst.expressions.{And, Attribute, Expression, GreaterThanOrEqual, LessThan, Literal} +import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} +import org.apache.spark.sql.comet.CometInMemoryTableScanExec +import org.apache.spark.sql.comet.execution.arrow.CometCachedBatchHelper import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.sql.types.BooleanType import org.apache.spark.storage.StorageLevel import org.apache.comet.CometConf @@ -968,4 +971,172 @@ class CometInMemoryCacheSuite extends CometTestBase { checkSparkAnswer(spark.sql("SELECT key, a, st, m, b FROM spark_columnar_complex")) } } + + /** + * Cache a six-column relation and hand the collected batches to `f` along with the relation, so + * a test can doctor the payload before decoding it again through the serializer. + */ + private def withProjectionCache( + f: (org.apache.spark.sql.execution.columnar.InMemoryRelation, Array[CachedBatch]) => Unit) + : Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true") { + + spark.catalog.clearCache() + spark + .range(0, 500, 1, 2) + .selectExpr( + "id", + "id % 100 AS k", + "cast(id as double) / 3 AS d", + "concat('a_', cast(id as string)) AS s1", + "concat('b_', cast(id % 17 as string)) AS s2", + "cast(id % 2 = 0 as boolean) AS flag") + .createOrReplaceTempView("projection_cache") + spark.catalog.cacheTable("projection_cache") + assert(spark.table("projection_cache").count() == 500) + assert( + cachedBatchTypes("projection_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val relation = spark.sharedState.cacheManager + .lookupCachedData(spark.table("projection_cache")) + .get + .cachedRepresentation + + try { + f(relation, relation.cacheBuilder.cachedColumnBuffers.collect()) + } finally { + spark.catalog.clearCache() + } + } + } + + /** Decode `batches` through the cache serializer, selecting `selected`, and total the rows. */ + private def decodedRowCount( + relation: org.apache.spark.sql.execution.columnar.InMemoryRelation, + batches: Array[CachedBatch], + selected: Seq[Attribute]): Long = { + relation.cacheBuilder.serializer + .convertCachedBatchToColumnarBatch( + spark.sparkContext.parallelize(batches.toSeq, 1), + relation.output, + selected, + spark.sessionState.conf) + // ColumnarBatch is not serializable, so reduce to a count inside the closure. + .mapPartitions(batches => Iterator.single(batches.map(_.numRows().toLong).sum)) + .collect() + .sum + } + + test("Comet in-memory cache stores one stream per column") { + withProjectionCache { (relation, batches) => + assert(batches.nonEmpty) + batches.foreach { batch => + assert( + CometCachedBatchHelper.numColumnStreams(batch) == relation.output.length, + "a cached batch must hold one independently decodable stream per cached column") + assert( + CometCachedBatchHelper.columnStreamSizes(batch).forall(_ > 0), + "every column stream must carry data") + } + } + } + + test("Comet in-memory cache decodes only the projected columns") { + // Timings would be a weak assertion here, so this corrupts the streams the read must not + // touch. Reading still has to succeed, which it only can if those streams were never + // inflated. The second half checks the corruption is detectable at all, so that the first + // half cannot pass just because the bad bytes decode silently to nothing. + withProjectionCache { (relation, batches) => + val selectedIdx = 1 + val selected = Seq(relation.output(selectedIdx)) + + relation.output.indices.filter(_ != selectedIdx).foreach { i => + batches.foreach(b => CometCachedBatchHelper.corruptColumnStream(b, i)) + } + + assert( + decodedRowCount(relation, batches, selected) == 500, + "reading one column must not decode the other five") + + batches.foreach(b => CometCachedBatchHelper.corruptColumnStream(b, selectedIdx)) + intercept[Exception] { + decodedRowCount(relation, batches, selected) + } + } + } + + test("Comet in-memory cache decodes no columns for a row-count-only read") { + // SELECT count(*) selects no columns. Every stream is corrupted, so the read can only succeed + // by decoding none of them and answering from the row count the cached batch already carries. + withProjectionCache { (relation, batches) => + relation.output.indices.foreach { i => + batches.foreach(b => CometCachedBatchHelper.corruptColumnStream(b, i)) + } + + assert(decodedRowCount(relation, batches, Seq.empty) == 500) + } + } + + test("Comet in-memory cache records per-column sizes in its statistics") { + // SimpleMetricsCachedBatch reserves a fifth field per column for its size. Each column is now + // its own stream, so the real size is known and must be reported rather than left at zero. + withProjectionCache { (relation, batches) => + batches.foreach { batch => + val sizes = CometCachedBatchHelper.columnStreamSizes(batch) + val stats = batch.asInstanceOf[SimpleMetricsCachedBatch].stats + sizes.zipWithIndex.foreach { case (size, i) => + assert( + stats.getLong(i * 5 + 4) == size, + s"column $i should report its own stream size in the statistics row") + } + } + } + } + + test("Comet in-memory cache scans one narrow column for a row-count-only query") { + // SELECT count(*) needs no columns, but the native plan needs a non-empty scan schema, so the + // scan has to ask for something. It asks for one cheap column: since the serializer decodes + // exactly what it is asked for, falling back to the whole cache schema would make the + // cheapest query in a workload decode every cached column. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true") { + + spark.catalog.clearCache() + spark + .range(0, 500, 1, 2) + .selectExpr( + "id", + "id % 100 AS k", + "concat('a_', cast(id as string)) AS s1", + "cast(id % 2 = 0 as boolean) AS flag") + .createOrReplaceTempView("count_only_cache") + spark.catalog.cacheTable("count_only_cache") + assert(spark.table("count_only_cache").count() == 500) + + val df = spark.sql("SELECT count(*) FROM count_only_cache") + val scan = df.queryExecution.executedPlan.collectFirst { + case s: CometInMemoryTableScanExec => s + } + + assert(scan.isDefined, "expected a native cache scan") + assert(scan.get.output.isEmpty, "a count-only scan declares no output") + assert( + scan.get.scanOutput.length == 1, + s"expected one scanned column, got ${scan.get.scanOutput.map(_.name).mkString(",")}") + assert( + scan.get.scanOutput.head.dataType == BooleanType, + "expected the cheapest column to decode, not a string or the first column") + + checkSparkAnswer(df) + 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 index 9c50e81107b..11099e78ab9 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -85,10 +85,13 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { |WHERE id >= 4500000 AND id < 4750000 """.stripMargin) - // A CometCachedBatch is one compressed Arrow IPC stream covering every cached column, so a - // scan decodes all of them and projects afterwards. Cost is therefore flat in the width of - // the projection, where Spark's per-column format falls away as it narrows. These two cases - // bracket that: the same cached relation read one column wide and six columns wide. + // A CometCachedBatch stores each column as its own stream, so a scan decodes only what it + // projected and cost tracks the width of the projection. These three cases span that range + // over one cached relation: no columns, one column, and all six. + runCacheBenchmark( + "in-memory cache row count only (0 of 6 columns)", + s"SELECT count(*) FROM $cacheTable") + runCacheBenchmark( "in-memory cache narrow projection (1 of 6 columns)", s"SELECT count(k) FROM $cacheTable") diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala new file mode 100644 index 00000000000..edbf3763713 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala @@ -0,0 +1,56 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.ByteBuffer + +import org.apache.spark.sql.columnar.CachedBatch +import org.apache.spark.util.io.ChunkedByteBuffer + +/** + * Test-only access to the internals of `CometCachedBatch`. + * + * A top-level `private` class in Scala is visible to its own package, so this shim needs no + * reflection; it exists so tests outside `org.apache.spark.sql.comet.execution.arrow` can assert + * on the cached payload's shape. + */ +object CometCachedBatchHelper { + + /** Number of independently decodable column streams in a cached batch. */ + def numColumnStreams(batch: CachedBatch): Int = + batch.asInstanceOf[CometCachedBatch].columns.length + + /** Serialized size of each column stream, in column order. */ + def columnStreamSizes(batch: CachedBatch): Seq[Long] = + batch.asInstanceOf[CometCachedBatch].columns.map(_.size).toSeq + + /** + * Replace one column's stream with bytes that cannot be decoded, in place. + * + * Reading a column this has corrupted fails; reading any other column only succeeds if that + * column's stream was never touched. That is the difference between decoding what was projected + * and decoding everything and projecting afterwards, so it is what the projection tests assert + * on rather than timings. + */ + def corruptColumnStream(batch: CachedBatch, index: Int): Unit = { + val columns = batch.asInstanceOf[CometCachedBatch].columns + columns(index) = new ChunkedByteBuffer(Array(ByteBuffer.wrap(Array[Byte](1, 2, 3, 4)))) + } +} From 86a91ea7dd0fbc9765ea74340bc460c5e74ccfc4 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 26 Aug 2026 17:26:35 -0600 Subject: [PATCH 14/19] fix: drop unused binding in cache statistics test scalafix RemoveUnused rejects the named relation parameter, which this test does not use. Spotless does not catch it, so it only surfaced in the Lint Java jobs. --- .../scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 bbeaeca15a8..48da648896d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1085,7 +1085,7 @@ class CometInMemoryCacheSuite extends CometTestBase { test("Comet in-memory cache records per-column sizes in its statistics") { // SimpleMetricsCachedBatch reserves a fifth field per column for its size. Each column is now // its own stream, so the real size is known and must be reported rather than left at zero. - withProjectionCache { (relation, batches) => + withProjectionCache { (_, batches) => batches.foreach { batch => val sizes = CometCachedBatchHelper.columnStreamSizes(batch) val stats = batch.asInstanceOf[SimpleMetricsCachedBatch].stats From cb07053e30ba084dd63787bf65ecb88a194f414a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 26 Aug 2026 18:13:43 -0600 Subject: [PATCH 15/19] fix: address review findings on the cache scan and serializer Five issues from review of 86a91ea7d. Emit zero-column batches for an empty-output cache scan. The scan widened an empty projection to a placeholder column so the native plan would have a non-empty scan schema, but that claim was untrue: a zero-column scan plans and runs. Meanwhile the emitted batches disagreed with the declared output, so a join over a cached relation read the wrong column and silently returned a wrong sum rather than failing. Removing the widening fixes the join and makes count(*) decode nothing at all. Serialize each column with the dictionary provider it was decoded with. Columns decoded from separate streams have independent dictionary ID namespaces, so re-encoding a decoded batch with the first column's provider could not resolve the rest. Release readers opened before a later column fails to decode. The holder is published to the task-completion listener only after its constructor returns, so a partial failure leaked off-heap for the life of the executor. Build the cached RDD in doExecuteColumnar rather than at planning time. CachedRDDBuilder.cachedColumnBuffers executes the cached plan, so planning a query over an adaptively cached relation ran a job and finalized that plan during EXPLAIN. Report large-offset Arrow vectors as not directly writable. isArrowBacked accepted any CometVector, including one wrapping LargeVarCharVector or LargeVarBinaryVector, which getFieldVector rejects; such batches now take the conversion path instead of failing at cache materialization. --- .../comet/CometInMemoryTableScanExec.scala | 63 +++----- .../arrow/ArrowCachedBatchSerializer.scala | 31 +++- .../apache/spark/sql/comet/util/Utils.scala | 75 ++++++--- .../comet/exec/CometInMemoryCacheSuite.scala | 148 ++++++++++++++++-- .../spark/sql/comet/util/UtilsSuite.scala | 38 +++++ 5 files changed, 272 insertions(+), 83 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 99a22d4d208..0bb93de9bef 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 @@ -23,11 +23,10 @@ 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.columnar.CachedBatchSerializer import org.apache.spark.sql.execution.LeafExecNode -import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.spark.sql.execution.columnar.{CachedRDDBuilder, InMemoryTableScanExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometConf @@ -49,7 +48,7 @@ import org.apache.comet.serde.QueryPlanSerde.serializeDataType case class CometInMemoryTableScanExec( originalPlan: InMemoryTableScanExec, serializer: CachedBatchSerializer, - cachedBuffers: RDD[CachedBatch], + cacheBuilder: CachedRDDBuilder, relationOutput: Seq[Attribute], scanOutput: Seq[Attribute]) extends CometExec @@ -58,10 +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 one - // placeholder column, 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; see `scanOutputFor` for why the scan cannot simply be empty. + // `scanOutput` always equals this, including when it is empty. An empty-output scan + // (`SELECT count(*)`) emits genuinely zero-column batches carrying only a row count: widening it + // to a placeholder column, or to the whole cache schema, makes the emitted batches disagree with + // the declared output, and a consumer that reads by ordinal rather than by row count -- a join, + // for instance -- then reads the wrong column. override def output: Seq[Attribute] = originalPlan.output // Use the serializer's vector types because the cached batch layout is owned by the serializer. @@ -80,6 +80,13 @@ case class CometInMemoryTableScanExec( override def doExecuteColumnar(): RDD[ColumnarBatch] = { val numOutputRows = longMetric("numOutputRows") + // Resolved here rather than at planning time. CachedRDDBuilder.cachedColumnBuffers is not a + // metadata lookup: it builds the RDD by calling execute/executeColumnar on the cached plan, + // so touching it while Comet is still planning the outer query runs jobs during planning -- + // visibly, an EXPLAIN of a query over an adaptively-cached relation would launch a job and + // finalize that plan. + val cachedBuffers = cacheBuilder.cachedColumnBuffers + val filteredBuffers = if (originalPlan.predicates.nonEmpty && conf.inMemoryPartitionPruning) { val filter = serializer.buildFilter(originalPlan.predicates, relationOutput) @@ -107,7 +114,7 @@ object CometInMemoryTableScanExec extends CometOperatorSerde[InMemoryTableScanEx builder: OperatorOuterClass.Operator.Builder, childOp: Operator*): Option[Operator] = { - val scanTypes = scanOutputFor(op).flatMap(attr => serializeDataType(attr.dataType)) + val scanTypes = op.output.flatMap(attr => serializeDataType(attr.dataType)) val scanBuilder = OperatorOuterClass.Scan .newBuilder() @@ -127,43 +134,9 @@ object CometInMemoryTableScanExec extends CometOperatorSerde[InMemoryTableScanEx CometInMemoryTableScanExec( op, relation.cacheBuilder.serializer, - relation.cacheBuilder.cachedColumnBuffers, + relation.cacheBuilder, relation.output, - scanOutputFor(op))) - } - - /** - * Columns the cache scan asks the serializer to decode. - * - * An empty-output scan (`SELECT count(*)`) still needs a non-empty schema for native planning, - * and the batches the node emits have to match that schema. Falling back to the whole cache - * schema satisfies both, but the serializer decodes exactly what it is asked for, so the - * cheapest query in the workload would decode every cached column. One column is enough: the - * aggregate above an empty-output scan reads the row count and never a value, so pick the - * cheapest to decode rather than all of them. - * - * `convert` and `createExec` must choose identically, or the native scan's declared schema and - * the batches fed to it disagree. - */ - private def scanOutputFor(op: InMemoryTableScanExec): Seq[Attribute] = { - if (op.output.nonEmpty) { - op.output - } else if (op.relation.output.isEmpty) { - Nil - } else { - Seq(op.relation.output.minBy(a => decodeCostRank(a.dataType))) - } + op.output)) } - // Rank by how much work decoding a column of this type costs, cheapest first. Fixed-width types - // decode to a flat buffer; variable-width and nested ones carry offsets, children and possibly - // dictionaries. - private def decodeCostRank(dt: DataType): Int = dt match { - case BooleanType | ByteType => 0 - case ShortType => 1 - case IntegerType | FloatType | DateType => 2 - case LongType | DoubleType | TimestampType | TimestampNTZType => 3 - case _: DecimalType => 4 - case _ => 5 - } } 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 b3c9ba240ee..a1ab0d672a1 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 @@ -20,6 +20,7 @@ package org.apache.spark.sql.comet.execution.arrow import scala.collection.JavaConverters._ +import scala.util.control.NonFatal import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD @@ -366,8 +367,34 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // this yields a single-element iterator that closes on exhaustion, matching what // ArrowReaderIterator did when the payload was one stream. private class ColumnReaders(buffers: Array[ChunkedByteBuffer], numRows: Int) { - private val readers: Array[Iterator[ColumnarBatch]] = - buffers.map(Utils.decodeBatches(_, "CometCache")) + // decodeBatches opens a reader and eagerly decodes its first batch, so it allocates. If a + // later column throws, the readers already opened here are unreachable: the task-completion + // listener cannot release them because `current` is only assigned once this constructor + // returns, so they would leak off-heap for the life of the executor. + private val readers: Array[Iterator[ColumnarBatch]] = { + val opened = new Array[Iterator[ColumnarBatch]](buffers.length) + var i = 0 + try { + while (i < buffers.length) { + opened(i) = Utils.decodeBatches(buffers(i), "CometCache") + i += 1 + } + } catch { + case NonFatal(e) => + var j = 0 + while (j < i) { + opened(j) match { + case reader: ArrowReaderIterator => + try reader.close() + catch { case NonFatal(closeError) => e.addSuppressed(closeError) } + case _ => () + } + j += 1 + } + throw e + } + opened + } private var closed = false def close(): Unit = synchronized { 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 33533903307..edf75173c92 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 @@ -280,11 +280,12 @@ object Utils extends CometTypeShim with Logging { * it alongside. As with [[serializeBatches]], the batch's vectors are cleared once written. */ def serializeBatchColumns(batch: ColumnarBatch): Array[ChunkedByteBuffer] = { - val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch) - val provider = batchProviderOpt.getOrElse(new CDataDictionaryProvider) val codec = CompressionCodec.createCodec(SparkEnv.get.conf) - fieldVectors.map { fieldVector => + // Each column is written with the provider it was decoded with, not the batch's first one: + // columns decoded from separate streams have independent dictionary ID namespaces. + getBatchFieldVectorsWithProviders(batch).map { case (fieldVector, providerOpt) => + val provider = providerOpt.getOrElse(new CDataDictionaryProvider) val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) @@ -440,35 +441,58 @@ object Utils extends CometTypeShim with Logging { * rather than being materialized column by column. */ def isArrowBacked(batch: ColumnarBatch): Boolean = - (0 until batch.numCols()).forall(i => batch.column(i).isInstanceOf[CometVector]) + (0 until batch.numCols()).forall { i => + batch.column(i) match { + // Not every CometVector can be handed to getFieldVector: a CometPlainVector can wrap a + // LargeVarCharVector or LargeVarBinaryVector (an accelerated mapInArrow returning + // pa.large_string(), for instance), which it rejects. Answering true for those would + // send a batch down the direct write path that then fails, so check the vector itself + // and let the caller convert instead. + case v: CometVector => isSupportedFieldVector(v.getValueVector) + case _ => false + } + } def getBatchFieldVectors( batch: ColumnarBatch): (Seq[FieldVector], Option[DictionaryProvider]) = { - var provider: Option[DictionaryProvider] = None + val columns = getBatchFieldVectorsWithProviders(batch) + (columns.map(_._1), columns.flatMap(_._2).headOption) + } + + /** + * Field vectors of `batch` paired with the dictionary provider each column was decoded with. + * + * [[getBatchFieldVectors]] collapses these to the first provider, which is right when every + * column came from the same reader. Comet's cache decodes each column from its own stream, so a + * batch's dictionary-backed columns can carry independent providers whose IDs collide; writing + * such a batch back out with one column's provider cannot resolve the others. Callers that + * serialize columns individually use this instead. + */ + def getBatchFieldVectorsWithProviders( + batch: ColumnarBatch): Seq[(FieldVector, Option[DictionaryProvider])] = { val rows = batch.numRows() - val fieldVectors = (0 until batch.numCols()).map { index => + (0 until batch.numCols()).map { index => batch.column(index) match { case a: CometVector => val valueVector = a.getValueVector - if (valueVector.getField.getDictionary != null) { - if (provider.isEmpty) { - provider = Some(a.getDictionaryProvider) - } - } + val provider = + if (valueVector.getField.getDictionary != null) Some(a.getDictionaryProvider) + else None - getFieldVector(valueVector, "serialize") + (getFieldVector(valueVector, "serialize"), provider) case cv: ConstantColumnVector => // Spark wraps file-source partition columns and other per-batch constants in // `ConstantColumnVector`. Materialise to an Arrow vector so the serialisation path // doesn't reject the batch. "UTC" is intentional -- see `ConstantColumnVectors`. - ConstantColumnVectors.materialize( + val materialized = ConstantColumnVectors.materialize( cv, cv.dataType(), rows, s"_const_$index", org.apache.comet.CometArrowAllocator, "UTC") + (materialized, None) case c => throw new SparkException( @@ -482,19 +506,24 @@ object Utils extends CometTypeShim with Logging { "data to Arrow format automatically.") } } - (fieldVectors, provider) + } + + /** Whether [[getFieldVector]] accepts this vector, without throwing to find out. */ + def isSupportedFieldVector(valueVector: ValueVector): Boolean = valueVector match { + case _: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector | _: BigIntVector | + _: Float4Vector | _: Float8Vector | _: VarCharVector | _: DecimalVector | + _: DateDayVector | _: TimeStampMicroTZVector | _: VarBinaryVector | + _: FixedSizeBinaryVector | _: TimeStampMicroVector | _: StructVector | _: ListVector | + _: MapVector | _: NullVector | _: TimeNanoVector => + true + case _ => false } def getFieldVector(valueVector: ValueVector, reason: String): FieldVector = { - valueVector match { - case v @ (_: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector | - _: BigIntVector | _: Float4Vector | _: Float8Vector | _: VarCharVector | - _: DecimalVector | _: DateDayVector | _: TimeStampMicroTZVector | _: VarBinaryVector | - _: FixedSizeBinaryVector | _: TimeStampMicroVector | _: StructVector | _: ListVector | - _: MapVector | _: NullVector | _: TimeNanoVector) => - v.asInstanceOf[FieldVector] - case _ => - throw new SparkException(s"Unsupported Arrow Vector for $reason: ${valueVector.getClass}") + if (isSupportedFieldVector(valueVector)) { + valueVector.asInstanceOf[FieldVector] + } else { + throw new SparkException(s"Unsupported Arrow Vector for $reason: ${valueVector.getClass}") } } } 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 48da648896d..f07b9c4b039 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -29,12 +29,11 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, Expression, Gr import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.comet.execution.arrow.CometCachedBatchHelper -import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} -import org.apache.spark.sql.types.BooleanType import org.apache.spark.storage.StorageLevel -import org.apache.comet.CometConf +import org.apache.comet.{CometArrowAllocator, CometConf} import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus import org.apache.comet.vector.CometVector @@ -1098,11 +1097,11 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Comet in-memory cache scans one narrow column for a row-count-only query") { - // SELECT count(*) needs no columns, but the native plan needs a non-empty scan schema, so the - // scan has to ask for something. It asks for one cheap column: since the serializer decodes - // exactly what it is asked for, falling back to the whole cache schema would make the - // cheapest query in a workload decode every cached column. + test("Comet in-memory cache scans no columns for a row-count-only query") { + // SELECT count(*) selects no columns, and the scan must keep it that way. Widening it -- to + // the whole cache schema, or to a single placeholder column -- makes the emitted batches + // disagree with the scan's declared output, which is wrong for any consumer that reads by + // ordinal instead of by row count. See the join regression below. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", @@ -1129,14 +1128,137 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(scan.isDefined, "expected a native cache scan") assert(scan.get.output.isEmpty, "a count-only scan declares no output") assert( - scan.get.scanOutput.length == 1, - s"expected one scanned column, got ${scan.get.scanOutput.map(_.name).mkString(",")}") - assert( - scan.get.scanOutput.head.dataType == BooleanType, - "expected the cheapest column to decode, not a string or the first column") + scan.get.scanOutput.isEmpty, + s"expected no scanned columns, got ${scan.get.scanOutput.map(_.name).mkString(",")}") checkSparkAnswer(df) spark.catalog.clearCache() } } + + test("Comet in-memory cache joins correctly over an empty-output cache scan") { + // An empty-output cache scan can feed a join, not only a count-style aggregate. A join reads + // its inputs by ordinal, so any column the scan emits beyond its declared output shifts the + // right side's positions and silently produces wrong results rather than failing. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + val left = spark.range(10L, 13L).cache() + left.collect() + left.createOrReplaceTempView("cached_left") + + // 3 left rows joined to 2 right rows, summing only the right side: 3 * (0 + 1) == 3. + // Leaking the left id column into the scan output made this read 10 + 11 + 12 twice. + checkSparkAnswer(spark.sql(""" + |SELECT /*+ BROADCAST(r) */ sum(r.id) + |FROM cached_left l JOIN range(2) r ON true + """.stripMargin)) + + checkSparkAnswer(spark.sql(""" + |SELECT /*+ BROADCAST(r) */ r.id + |FROM cached_left l JOIN range(2) r ON true + """.stripMargin)) + + spark.catalog.clearCache() + } + } + + test( + "Comet in-memory cache re-encodes a decoded batch whose columns have separate dictionaries") { + // Each cached column is decoded from its own stream, so dictionary-backed columns come back + // with independent providers whose IDs collide. Re-encoding such a batch with only the first + // column's provider cannot resolve the later columns' dictionary IDs. Spark's columnar Union + // hands decoded cached batches straight back to this serializer, so caching a union of a + // cached relation exercises exactly that. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + val first = spark + .range(0, 200, 1, 2) + .selectExpr( + "concat('a_', cast(id % 3 as string)) AS s1", + "concat('b_', cast(id % 4 as string)) AS s2") + .repartition(2) + .cache() + assert(first.count() == 200) + + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false") { + val second = first.union(first).cache() + assert(second.count() == 400) + second.unpersist() + } + + first.unpersist() + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache does not build the cached RDD while planning") { + // CachedRDDBuilder.cachedColumnBuffers builds its RDD by executing the cached plan, so + // touching it during planning runs jobs before the outer query is even submitted. With an + // adaptively-cached relation that also finalizes the cached plan. EXPLAIN must launch nothing. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + val cached = spark.range(100).repartition(2).cache() + cached.createOrReplaceTempView("cached_adaptive") + + val builder = spark + .sql("SELECT * FROM cached_adaptive") + .queryExecution + .optimizedPlan + .collectFirst { case r: InMemoryRelation => r.cacheBuilder } + .get + // The cached plan is adaptive and has not run, so AQE has not finalized it. Building the + // cached RDD executes that plan, which finalizes it; isCachedColumnBuffersLoaded is not the + // signal to use here, since it additionally requires the blocks to be populated. + assert( + builder.cachedPlan.toString.contains("isFinalPlan=false"), + "cached plan was already finalized before the test ran") + + spark.sql("SELECT * FROM cached_adaptive").explain() + + assert( + builder.cachedPlan.toString.contains("isFinalPlan=false"), + "planning must not build the cached RDD: doing so executes the cached plan") + + // It must still be built when the query actually runs. + assert(spark.sql("SELECT * FROM cached_adaptive").count() == 100) + + cached.unpersist() + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache releases opened readers when a later column fails to decode") { + // A cached batch is several independent Arrow streams and decodeBatches opens each eagerly. + // If a later column throws, the readers already opened are unreachable: the task-completion + // listener cannot release them, because the holder is only published once its constructor + // returns. The failure would then leak off-heap for the life of the executor. + withProjectionCache { (relation, batches) => + // Corrupt the second selected column, so the first is opened successfully first. + val selected = Seq(relation.output(0), relation.output(1)) + batches.foreach(b => CometCachedBatchHelper.corruptColumnStream(b, 1)) + + val before = CometArrowAllocator.getAllocatedMemory + intercept[Exception] { + decodedRowCount(relation, batches, selected) + } + assert( + CometArrowAllocator.getAllocatedMemory == before, + "readers opened before the failure must be released") + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index c3b00a2814c..4510f9d0ac1 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -19,11 +19,15 @@ package org.apache.spark.sql.comet.util +import org.apache.arrow.c.CDataDictionaryProvider import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.comet.CometArrowAllocator +import org.apache.comet.vector.CometVector + class UtilsSuite extends CometTestBase { test("serializeBatches preserves row count for a zero-column batch") { @@ -159,4 +163,38 @@ class UtilsSuite extends CometTestBase { assert(nameNulls.forall(identity), s"expected all name null, got $nameNulls") assert(structNulls.forall(identity), s"expected all struct null, got $structNulls") } + + test("isArrowBacked rejects large-offset Arrow vectors") { + // A CometPlainVector can wrap a LargeVarCharVector or LargeVarBinaryVector -- an accelerated + // mapInArrow returning pa.large_string() produces one -- but getFieldVector rejects both. If + // isArrowBacked accepted them, a caller would take the direct write path and then fail, so it + // must report false and let the caller convert the batch instead. + val numRows = 2 + Seq[org.apache.arrow.vector.FieldVector]( + { + val v = new org.apache.arrow.vector.LargeVarCharVector("s", CometArrowAllocator) + v.allocateNew(numRows) + v.setSafe(0, "hello".getBytes("UTF-8")) + v.setSafe(1, "world".getBytes("UTF-8")) + v.setValueCount(numRows) + v + }, { + val v = new org.apache.arrow.vector.LargeVarBinaryVector("b", CometArrowAllocator) + v.allocateNew(numRows) + v.setSafe(0, "hello".getBytes("UTF-8")) + v.setSafe(1, "world".getBytes("UTF-8")) + v.setValueCount(numRows) + v + }).foreach { vector => + try { + val col = CometVector.getVector(vector, new CDataDictionaryProvider) + val batch = new ColumnarBatch(Array[ColumnVector](col), numRows) + assert( + !Utils.isArrowBacked(batch), + s"${vector.getClass.getSimpleName} must not be reported as directly writable") + } finally { + vector.close() + } + } + } } From 70f046abf797820a4e2df306e0895d0dc14ea2fc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 27 Aug 2026 08:03:01 -0600 Subject: [PATCH 16/19] fix: address the latest review findings on the cache scan Three findings from review, each with a regression test that fails without the fix: - A broadcast of a cache scan re-serializes each decoded batch as one stream covering every column, but `getBatchFieldVectors` handed the writer only the first column's dictionary provider. The columns are decoded from separate streams, so they arrive with separate providers and the write failed with "Could not find dictionary with ID 1". Combine the providers instead, and refuse a genuine ID clash rather than resolve one column against another's dictionary. - Opening an Arrow reader decodes the column's first batch, and a dictionary encoded column loads its dictionary before the record batch that indexes into it. A failure in between left that dictionary owned by a reader nobody holds: the constructor never returned. Close the reader on the way out, in `ArrowReaderIterator` and in `StreamReader`'s own schema read. - `CometInMemoryTableScanExec` carries the wrapped Spark scan as a plan-typed field rather than a child, so canonicalization walked past it and left its expression IDs in place. Two scans of one cache compared unequal, and a UNION of two identical aggregates over it ran two shuffles where Spark's cache scan runs one and reuses it. Canonicalize the wrapped scan, which keeps scans with different pushed predicates distinct. Also fix the Spark 3.4 failure in "does not build the cached RDD while planning": 3.4 defaults `canChangeCachedPlanOutputPartitioning` to false, which force-disables AQE inside the cached plan, so there is no finalization for the test to observe. Set the conf so the relation stays adaptive on every version. --- .../apache/comet/vector/StreamReader.scala | 14 +- .../comet/CometInMemoryTableScanExec.scala | 18 ++- .../execution/arrow/ArrowReaderIterator.scala | 16 +- .../apache/spark/sql/comet/util/Utils.scala | 55 ++++++- .../comet/exec/CometInMemoryCacheSuite.scala | 143 ++++++++++++++++++ .../arrow/CometCachedBatchHelper.scala | 65 +++++++- 6 files changed, 300 insertions(+), 11 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala index b8106a96e04..805eae988e8 100644 --- a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala +++ b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala @@ -21,6 +21,8 @@ package org.apache.comet.vector import java.nio.channels.ReadableByteChannel +import scala.util.control.NonFatal + import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.ipc.{ArrowStreamReader, ReadChannel} import org.apache.arrow.vector.ipc.message.MessageChannelReader @@ -35,7 +37,17 @@ case class StreamReader(channel: ReadableByteChannel, source: String) extends Au private val channelReader = new MessageChannelReader(new ReadChannel(channel), CometArrowAllocator) private var arrowReader = new ArrowStreamReader(channelReader, CometArrowAllocator) - private var root = arrowReader.getVectorSchemaRoot + + // Reading the schema allocates the root's vectors, so it can fail with buffers already taken. + // No caller holds this reader until its constructor returns, so close it here or nothing will. + private var root = + try arrowReader.getVectorSchemaRoot + catch { + case NonFatal(e) => + try arrowReader.close() + catch { case NonFatal(closeError) => e.addSuppressed(closeError) } + throw e + } def nextBatch(): Option[ColumnarBatch] = { if (arrowReader.loadNextBatch()) { 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 0bb93de9bef..261102bc72a 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 @@ -24,7 +24,7 @@ import scala.collection.JavaConverters._ import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.columnar.CachedBatchSerializer -import org.apache.spark.sql.execution.LeafExecNode +import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan} import org.apache.spark.sql.execution.columnar.{CachedRDDBuilder, InMemoryTableScanExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -64,6 +64,22 @@ case class CometInMemoryTableScanExec( // for instance -- then reads the wrong column. override def output: Seq[Attribute] = originalPlan.output + // `originalPlan` is a plan-typed field rather than a child, so QueryPlan's canonicalization + // walks straight past it: its attributes and predicates keep the expression IDs of whichever + // occurrence of the cached relation produced them. Two scans of one cache then compare unequal, + // and since sameResult is what exchange and broadcast reuse are keyed on, a UNION of two + // identical aggregates over a cached table runs two shuffles where Spark's own cache scan runs + // one and reuses it. + // + // Defer to `InMemoryTableScanExec`, which normalizes its own attributes and predicates against + // the relation's output. Dropping the field instead would also make the scans compare equal, + // but it would equate scans carrying different pruning predicates along with them. + override protected def doCanonicalize(): SparkPlan = + super + .doCanonicalize() + .asInstanceOf[CometInMemoryTableScanExec] + .copy(originalPlan = originalPlan.canonicalized.asInstanceOf[InMemoryTableScanExec]) + // 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) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowReaderIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowReaderIterator.scala index 0d0093a107e..fa29f72bb7b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowReaderIterator.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowReaderIterator.scala @@ -21,6 +21,8 @@ package org.apache.spark.sql.comet.execution.arrow import java.nio.channels.ReadableByteChannel +import scala.util.control.NonFatal + import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.vector._ @@ -29,7 +31,19 @@ class ArrowReaderIterator(channel: ReadableByteChannel, source: String) extends Iterator[ColumnarBatch] { private val reader = StreamReader(channel, source) - private var batch = nextBatch() + + // Decoding eagerly here is what makes hasNext cheap, but it allocates: loading a batch first + // loads the dictionaries it references. A failure part way through leaves those allocations + // owned by the reader, and nothing else can release them -- this constructor never returns, so + // no caller ever holds the iterator it would close. Close the reader on the way out instead. + private var batch = + try nextBatch() + catch { + case NonFatal(e) => + try reader.close() + catch { case NonFatal(closeError) => e.addSuppressed(closeError) } + throw e + } private var currentBatch: ColumnarBatch = null private var isClosed: Boolean = false 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 edf75173c92..cd3227877ea 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 @@ -28,7 +28,8 @@ import scala.jdk.CollectionConverters._ import org.apache.arrow.c.CDataDictionaryProvider import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} -import org.apache.arrow.vector.dictionary.DictionaryProvider +import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} +import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} import org.apache.arrow.vector.types._ import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} @@ -456,17 +457,57 @@ object Utils extends CometTypeShim with Logging { def getBatchFieldVectors( batch: ColumnarBatch): (Seq[FieldVector], Option[DictionaryProvider]) = { val columns = getBatchFieldVectorsWithProviders(batch) - (columns.map(_._1), columns.flatMap(_._2).headOption) + (columns.map(_._1), combineDictionaryProviders(columns)) + } + + /** + * The dictionaries every dictionary-encoded column of `columns` refers to, as one provider. + * + * Columns of a batch need not share a provider. Comet's cache decodes each column from its own + * Arrow stream, so a dictionary-backed column arrives carrying the provider its reader built, + * and a batch that reaches [[serializeBatches]] -- a native broadcast of a cache scan, say -- + * can hold several. Writing the whole batch emits one schema covering every column and resolves + * each column's dictionary ID against the single provider the writer was given, so handing it + * any one column's provider fails with "Could not find dictionary with ID n" for the others. + */ + private def combineDictionaryProviders( + columns: Seq[(FieldVector, Option[DictionaryProvider])]): Option[DictionaryProvider] = { + val dictionaries = scala.collection.mutable.LinkedHashMap.empty[Long, Dictionary] + + columns.foreach { case (vector, providerOpt) => + val encoding = vector.getField.getDictionary + if (encoding != null) { + val id = encoding.getId + val dictionary = providerOpt.map(_.lookup(id)).orNull + if (dictionary == null) { + throw new SparkException( + s"Column ${vector.getField.getName} is dictionary encoded with ID $id, but no " + + "dictionary with that ID was provided") + } + dictionaries.get(id) match { + // Every provider seen here descends from one upstream reader, which numbers the + // dictionaries it hands out, so two columns sharing an ID share the dictionary itself. + // A genuine clash would need renumbering, which means rewriting each vector's field, + // so refuse rather than silently decode one column against another's dictionary. + case Some(existing) if existing.getVector ne dictionary.getVector => + throw new SparkException( + s"Columns of the same batch carry different dictionaries under ID $id") + case _ => dictionaries.put(id, dictionary) + } + } + } + + if (dictionaries.isEmpty) None + else Some(new MapDictionaryProvider(dictionaries.values.toSeq: _*)) } /** * Field vectors of `batch` paired with the dictionary provider each column was decoded with. * - * [[getBatchFieldVectors]] collapses these to the first provider, which is right when every - * column came from the same reader. Comet's cache decodes each column from its own stream, so a - * batch's dictionary-backed columns can carry independent providers whose IDs collide; writing - * such a batch back out with one column's provider cannot resolve the others. Callers that - * serialize columns individually use this instead. + * [[getBatchFieldVectors]] folds these into one provider covering the whole batch, which is + * what a single stream over every column needs. Comet's cache decodes each column from its own + * stream and writes it back the same way, so it keeps the pairing instead: each column is + * written with the provider it was decoded with. */ def getBatchFieldVectorsWithProviders( batch: ColumnarBatch): Seq[(FieldVector, Option[DictionaryProvider])] = { 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 f07b9c4b039..a7a9a6c8590 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.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.comet.execution.arrow.CometCachedBatchHelper import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.storage.StorageLevel @@ -1208,6 +1209,11 @@ class CometInMemoryCacheSuite extends CometTestBase { // adaptively-cached relation that also finalizes the cached plan. EXPLAIN must launch nothing. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + // Spark caches through a session with some configs forced off, and on 3.4 that list still + // includes AQE itself, so the cached plan comes back non-adaptive and there is nothing to + // finalize. This conf is what decides that list; 3.5 defaults it on, and 4.0 stopped + // disabling AQE either way. Setting it keeps the relation adaptive on every version. + SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", "spark.comet.sparkToColumnar.enabled" -> "true") { @@ -1261,4 +1267,141 @@ class CometInMemoryCacheSuite extends CometTestBase { "readers opened before the failure must be released") } } + + /** + * Cache two low-cardinality string columns and hand the test the cached payload. + * + * The shuffle is what makes this worth its own fixture: its reader hands the cache writer + * dictionary-encoded columns, so each cached column stream carries a dictionary of its own. + */ + private def withDictionaryCache(f: (InMemoryRelation, Array[CachedBatch]) => Unit): Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + spark + .range(0, 2000, 1, 2) + .selectExpr( + "concat('a_', cast(id % 3 as string)) AS s1", + "concat('b_', cast(id % 4 as string)) AS s2") + .repartition(2) + .createOrReplaceTempView("dictionary_cache") + spark.catalog.cacheTable("dictionary_cache") + assert(spark.table("dictionary_cache").count() == 2000) + assert( + cachedBatchTypes("dictionary_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val relation = spark.sharedState.cacheManager + .lookupCachedData(spark.table("dictionary_cache")) + .get + .cachedRepresentation + + try { + f(relation, relation.cacheBuilder.cachedColumnBuffers.collect()) + } finally { + spark.catalog.clearCache() + } + } + } + + test("Comet in-memory cache broadcasts a batch whose columns have separate dictionaries") { + // A broadcast of a cache scan re-serializes each decoded batch as one stream covering every + // column, and the writer resolves all of their dictionary IDs against the single provider it + // is handed. The columns were decoded from separate streams, so they arrive carrying separate + // providers: passing any one of them cannot resolve the others. + withDictionaryCache { (relation, batches) => + assert( + CometCachedBatchHelper.columnsAreDictionaryEncoded(batches.head).forall(identity), + "this test is only meaningful over dictionary-encoded cached columns") + assert(relation.output.length == 2) + + val df = spark.sql( + "SELECT /*+ BROADCAST(c) */ c.s1, c.s2 FROM range(1) r JOIN dictionary_cache c ON true") + checkSparkAnswer(df) + assert(df.count() == 2000) + } + } + + test("Comet in-memory cache releases a reader whose own first batch fails to decode") { + // A dictionary-encoded column loads its dictionary before the record batch that indexes into + // it, so a reader can allocate and then fail while opening. Nothing else can release what it + // took: its constructor never returns, so no caller holds the reader it would close, and the + // task-completion listener has not been told about it either. + withDictionaryCache { (relation, batches) => + assert( + CometCachedBatchHelper.columnsAreDictionaryEncoded(batches.head).forall(identity), + "this test is only meaningful over dictionary-encoded cached columns") + + // Enough to take out the end-of-stream marker and bite into the record batch body, so the + // read fails after the dictionary has been loaded rather than before. + batches.foreach(b => CometCachedBatchHelper.truncateColumnStream(b, 0, 64)) + + val before = CometArrowAllocator.getAllocatedMemory + intercept[Exception] { + decodedRowCount(relation, batches, Seq(relation.output.head)) + } + assert( + CometArrowAllocator.getAllocatedMemory == before, + "a reader that fails while opening must release what it already allocated") + } + } + + test("Comet in-memory cache scans of one cache canonicalize equal, so exchanges are reused") { + // The wrapped Spark scan is a plan-typed field rather than a child, so canonicalization walks + // past it and leaves in place the expression IDs of whichever occurrence of the relation + // produced it. sameResult is what exchange and broadcast reuse are keyed on, so two + // equivalent scans that compare unequal make a query shuffle and aggregate one cache twice. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + spark + .range(0, 400, 1, 2) + .selectExpr("id", "id % 10 AS k") + .createOrReplaceTempView("reuse_cache") + spark.catalog.cacheTable("reuse_cache") + assert(spark.table("reuse_cache").count() == 400) + + val df = spark.sql( + "SELECT k, count(*) AS c FROM reuse_cache GROUP BY k " + + "UNION ALL SELECT k, count(*) AS c FROM reuse_cache GROUP BY k") + checkSparkAnswer(df) + + val plan = df.queryExecution.executedPlan + val exchanges = plan.collect { case e: Exchange => e } + val reused = plan.collect { case r: ReusedExchangeExec => r } + assert( + exchanges.length == 1 && reused.length == 1, + s"expected one exchange and one reuse of it, got ${exchanges.length} exchanges and " + + s"${reused.length} reuses:\n$plan") + + // Canonicalization must not simply drop the wrapped scan: scans that differ only in the + // predicates pushed into them have to stay distinct. + def scanOf(query: String): CometInMemoryTableScanExec = + spark + .sql(query) + .queryExecution + .executedPlan + .collectFirst { case s: CometInMemoryTableScanExec => s } + .get + + val under100 = scanOf("SELECT k FROM reuse_cache WHERE id < 100") + val under200 = scanOf("SELECT k FROM reuse_cache WHERE id < 200") + assert( + under100.originalPlan.predicates.nonEmpty, + "expected the filter to be pushed into the cache scan") + assert( + under100.canonicalized != under200.canonicalized, + "cache scans with different pruning predicates must not compare equal") + + spark.catalog.clearCache() + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala index edbf3763713..5548f6dadbd 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala @@ -19,10 +19,17 @@ package org.apache.spark.sql.comet.execution.arrow +import java.io.{DataInputStream, DataOutputStream} import java.nio.ByteBuffer +import java.nio.channels.Channels +import org.apache.arrow.vector.ipc.ArrowStreamReader +import org.apache.spark.SparkEnv +import org.apache.spark.io.CompressionCodec import org.apache.spark.sql.columnar.CachedBatch -import org.apache.spark.util.io.ChunkedByteBuffer +import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStream} + +import org.apache.comet.CometArrowAllocator /** * Test-only access to the internals of `CometCachedBatch`. @@ -53,4 +60,60 @@ object CometCachedBatchHelper { val columns = batch.asInstanceOf[CometCachedBatch].columns columns(index) = new ChunkedByteBuffer(Array(ByteBuffer.wrap(Array[Byte](1, 2, 3, 4)))) } + + /** Whether each column's stream stores that column dictionary encoded, in column order. */ + def columnsAreDictionaryEncoded(batch: CachedBatch): Seq[Boolean] = + batch.asInstanceOf[CometCachedBatch].columns.toSeq.map { buffer => + val in = new DataInputStream(codec.compressedInputStream(buffer.toInputStream())) + val reader = new ArrowStreamReader(Channels.newChannel(in), CometArrowAllocator) + try { + reader.getVectorSchemaRoot.getSchema.getFields.get(0).getDictionary != null + } finally { + reader.close() + } + } + + /** + * Drop the last `dropBytes` of one column's decoded Arrow stream, in place. + * + * [[corruptColumnStream]] replaces the stream outright, so a reader over it fails on the very + * first message, before it has allocated anything. This keeps the stream genuine up to the cut: + * the reader parses the schema and loads the column's dictionary, and only then runs out of + * input part way through the record batch that indexes into it. The cut is made on the decoded + * bytes rather than the compressed ones because a small column compresses to a single block, + * and truncating that fails the decompressor before Arrow reads anything at all. + */ + def truncateColumnStream(batch: CachedBatch, index: Int, dropBytes: Int): Unit = { + val columns = batch.asInstanceOf[CometCachedBatch].columns + + val decodedStream = new DataInputStream( + codec.compressedInputStream(columns(index).toInputStream())) + val decoded = + try { + val buffer = new java.io.ByteArrayOutputStream() + val chunk = new Array[Byte](8192) + var read = decodedStream.read(chunk) + while (read >= 0) { + buffer.write(chunk, 0, read) + read = decodedStream.read(chunk) + } + buffer.toByteArray + } finally { + decodedStream.close() + } + require( + decoded.length > dropBytes, + s"column $index decodes to ${decoded.length} bytes, too few to drop $dropBytes") + + val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) + val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) + try { + out.write(decoded, 0, decoded.length - dropBytes) + } finally { + out.close() + } + columns(index) = cbbos.toChunkedByteBuffer + } + + private def codec: CompressionCodec = CompressionCodec.createCodec(SparkEnv.get.conf) } From 7e19dff202288330c74bd777ae7e63dd02db0ae7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 27 Aug 2026 12:12:21 -0600 Subject: [PATCH 17/19] fix: register Comet's serialized payloads with Kryo, and correct the cache benchmark framing Under spark.kryo.registrationRequired=true, Kryo rejects any class it has not been told about. Two Comet payloads reach it: * CometCachedBatch, whenever Spark serializes a cached block. That is not only DISK_ONLY: the _SER levels, replication, cross-executor fetches, and the disk half of the default MEMORY_AND_DISK all serialize, so a plain df.cache() that spills is enough. * The Array[ChunkedByteBuffer] a native broadcast broadcasts. Spark registers ChunkedByteBuffer but not an array of them, so CometBroadcastExchangeExec fails here regardless of which Comet features are enabled. That one predates the cache work. Spark registers its own ArrowCachedBatch in KryoSerializer.loadableSparkClasses; Comet cannot add to that list, and spark.kryo.registrator is read when SparkEnv builds the serializer, before any plugin runs, so Comet cannot set it either. So: provide CometKryoRegistrator, document it on the cache config, and have CometDriverPlugin warn at startup when the combination is unsafe. DefaultCachedBatch is registered too, because a schema this serializer cannot store is delegated to Spark's, and Spark only registers that class itself from 4.1 onwards. Separately, correct what CometInMemoryCacheBenchmark claims to measure. Both cases run the aggregation on Comet; only the cache-scan boundary moves. The case labels and comment said "Spark execution over CometCachedBatch", which reads as a Spark-execution baseline it never was. verifyPlan now asserts the CometSparkColumnarToColumnar bridge in the disabled case so the framing is pinned by the benchmark itself. --- .../scala/org/apache/comet/CometConf.scala | 5 +- .../apache/comet/CometKryoRegistrator.scala | 58 ++++++ .../main/scala/org/apache/spark/Plugins.scala | 31 +++ .../arrow/ArrowCachedBatchSerializer.scala | 33 ++- .../apache/spark/sql/comet/util/Utils.scala | 16 ++ .../exec/CometInMemoryCacheKryoSuite.scala | 191 ++++++++++++++++++ .../CometInMemoryCacheBenchmark.scala | 30 ++- 7 files changed, 353 insertions(+), 11 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/CometKryoRegistrator.scala create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index cf5a88d2359..507bccf6138 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -271,7 +271,10 @@ object CometConf extends ShimCometConf { "format. Each cached column is stored as its own compressed Arrow IPC stream, so a " + "scan decodes only the columns it projected. Reads that feed Spark operators rather " + "than Comet ones still pay a row conversion the default format avoids, and can be " + - "slower than Spark's cache.") + "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " + + "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " + + "SparkContext, otherwise caching fails as soon as a block is serialized, including " + + "the disk half of the default MEMORY_AND_DISK storage level.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/comet/CometKryoRegistrator.scala b/spark/src/main/scala/org/apache/comet/CometKryoRegistrator.scala new file mode 100644 index 00000000000..9bc0be891d2 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/CometKryoRegistrator.scala @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet + +import org.apache.spark.serializer.KryoRegistrator +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.comet.util.Utils + +import com.esotericsoftware.kryo.Kryo + +/** + * Registers the classes Comet hands to Spark's serializer with Kryo. + * + * This is only needed when `spark.kryo.registrationRequired=true`, which makes Kryo reject any + * unregistered class rather than writing its name. Set it alongside Comet's own configuration: + * + * {{{ + * spark.serializer org.apache.spark.serializer.KryoSerializer + * spark.kryo.registrator org.apache.comet.CometKryoRegistrator + * }}} + * + * `spark.kryo.registrator` has to be set before the `SparkContext` is created, because + * `KryoSerializer` reads it when `SparkEnv` builds it. That is earlier than `CometDriverPlugin` + * runs, so Comet cannot add this for you the way it can add `spark.sql.cache.serializer`; + * `CometDriverPlugin` logs a warning instead when the combination looks unsafe. + * + * Two payloads need it: the `Array[ChunkedByteBuffer]` a native broadcast broadcasts, and + * `CometCachedBatch`. The first applies whether or not the in-memory cache feature is enabled. + */ +class CometKryoRegistrator extends KryoRegistrator { + override def registerClasses(kryo: Kryo): Unit = { + CometKryoRegistrator.classes.foreach(kryo.register) + } +} + +object CometKryoRegistrator { + val CLASS_NAME: String = classOf[CometKryoRegistrator].getName + + def classes: Seq[Class[_]] = + Utils.arrowBytesKryoClasses ++ ArrowCachedBatchSerializer.kryoClasses +} diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 1aa87c65eda..eaeac316655 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.internal.StaticSQLConf import org.apache.comet.{COMET_VERSION, CometSparkSessionExtensions, NativeBase} import org.apache.comet.CometConf import org.apache.comet.CometConf.{COMET_METRICS_ENABLED, COMET_ONHEAP_ENABLED} +import org.apache.comet.CometKryoRegistrator import org.apache.comet.annotation.Public /** @@ -65,6 +66,7 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl val extraConfs = new ju.HashMap[String, String]() CometDriverPlugin.maybeSetCacheSerializer(sc.conf, extraConfs) + CometDriverPlugin.warnIfKryoRegistratorMissing(sc.conf) // register CometSparkSessionExtensions if it isn't already registered CometDriverPlugin.registerCometSessionExtension(sc.conf) @@ -145,6 +147,35 @@ object CometDriverPlugin extends Logging { } } + // Comet hands Spark's serializer classes that Kryo has not been told about, so with + // spark.kryo.registrationRequired=true it rejects them with "Class is not registered", which + // names neither Comet nor the operation that failed. Two paths reach it: a native broadcast, + // which broadcasts an Array[ChunkedByteBuffer], and any cached block Spark serializes -- the + // disk half of MEMORY_AND_DISK, the _SER levels, replication, a cross-executor fetch. + // CometKryoRegistrator covers both, but spark.kryo.registrator is read when SparkEnv builds the + // serializer, before any plugin runs, so it cannot be set from here. Say so while the + // application is still starting up rather than leaving the user to attribute the failure later. + private[apache] def warnIfKryoRegistratorMissing(conf: SparkConf): Unit = { + val usingKryo = + conf.get("spark.serializer", "") == "org.apache.spark.serializer.KryoSerializer" + val registrationRequired = conf.getBoolean("spark.kryo.registrationRequired", false) + val registered = conf + .get("spark.kryo.registrator", "") + .split(',') + .map(_.trim) + .contains(CometKryoRegistrator.CLASS_NAME) + + if (usingKryo && registrationRequired && !registered) { + logWarning( + "spark.kryo.registrationRequired=true but spark.kryo.registrator does not include " + + s"${CometKryoRegistrator.CLASS_NAME}. Comet's native broadcast and its in-memory " + + "cache format will fail with Kryo's \"Class is not registered\" as soon as their " + + "payloads are serialized. Add " + + s"spark.kryo.registrator=${CometKryoRegistrator.CLASS_NAME} before creating the " + + "SparkContext; it cannot be set later.") + } + } + def registerCometMetrics(sc: SparkContext): Unit = { if (sc.getConf.getBoolean( COMET_METRICS_ENABLED.key, 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 a1ab0d672a1..46a51ad8775 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 @@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.InternalRow 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.DefaultCachedBatchSerializer +import org.apache.spark.sql.execution.columnar.{DefaultCachedBatch, DefaultCachedBatchSerializer} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} @@ -536,4 +536,35 @@ object ArrowCachedBatchSerializer { def supportsSchema(schema: Seq[Attribute]): Boolean = schema.forall(a => supportsType(a.dataType)) + + /** + * The classes a `CometCachedBatch` adds on top of [[org.apache.comet.CometKryoRegistrator]]'s + * shared Arrow-bytes classes. + * + * Spark serializes a `CachedBatch` with `spark.serializer` whenever the block leaves the heap: + * the disk half of `MEMORY_AND_DISK`, the `_SER` levels, replication, and cross-executor + * fetches. Under `spark.kryo.registrationRequired=true` Kryo rejects any class it has not been + * told about, so an ordinary `df.cache()` that spills would fail with "Class is not registered" + * rather than anything naming this feature. Spark registers its own `ArrowCachedBatch` in + * `KryoSerializer.loadableSparkClasses` for the same reason; Comet cannot add to that list, so + * `CometKryoRegistrator` registers these instead. + */ + def kryoClasses: Seq[Class[_]] = Seq( + classOf[CometCachedBatch], + // The statistics row, whose values are bounds in Spark's internal representation: boxed + // primitives, which Kryo registers by default, plus UTF8String and Decimal, which it does not. + // A Decimal above Long precision holds a scala.math.BigDecimal, which Chill's Scala registrar + // already covers with a serializer that writes the java.math.BigDecimal inside it as a + // class-and-object, so that one has to be registered here. + classOf[GenericInternalRow], + classOf[Array[Any]], + classOf[UTF8String], + classOf[Decimal], + classOf[java.math.BigDecimal], + classOf[java.math.BigInteger], + // A relation whose schema this serializer cannot store is delegated to Spark's + // DefaultCachedBatchSerializer, so its payload has to survive Kryo too. Spark registers + // DefaultCachedBatch itself only from 4.1 onwards, so on 3.4, 3.5 and 4.0 the delegated path + // fails without this. Registering it twice on 4.1 is a no-op. + classOf[DefaultCachedBatch]) } 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 cd3227877ea..769d8058de5 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 @@ -301,6 +301,22 @@ object Utils extends CometTypeShim with Logging { }.toArray } + /** + * The classes that carry the output of [[serializeBatches]] and [[serializeBatchColumns]] out + * of Comet, for Kryo registration by [[org.apache.comet.CometKryoRegistrator]]. + * + * Spark registers `ChunkedByteBuffer` itself but not an array of them, and + * `CometBroadcastExchangeExec` broadcasts exactly that array, so a native broadcast fails under + * `spark.kryo.registrationRequired=true` whichever Comet features are enabled. Comet's cache + * format stores one buffer per column and so needs the same registrations. + */ + def arrowBytesKryoClasses: Seq[Class[_]] = Seq( + classOf[ChunkedByteBuffer], + classOf[Array[ChunkedByteBuffer]], + // A ChunkedByteBuffer's own chunks. ChunkedByteBufferOutputStream allocates them on heap. + classOf[Array[ByteBuffer]], + ByteBuffer.allocate(1).getClass) + /** * Decodes the byte arrays back to ColumnarBatchs and put them into buffer. * diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala new file mode 100644 index 00000000000..01c3b5ec185 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala @@ -0,0 +1,191 @@ +/* + * 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.execution.columnar.CometInMemoryRelationHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.storage.StorageLevel + +import org.apache.comet.{CometConf, CometKryoRegistrator} + +/** + * Covers Comet's cached batch format under `spark.kryo.registrationRequired=true`. + * + * Kryo then rejects any class it has not been told about, and Spark serializes a `CachedBatch` + * whenever a cached block leaves the heap: the disk half of the default `MEMORY_AND_DISK`, the + * `_SER` levels, replication, and cross-executor fetches. So this is not a `DISK_ONLY`-only + * concern -- a plain `df.cache()` that spills is enough to reach it. Spark registers its own + * `ArrowCachedBatch` in `KryoSerializer.loadableSparkClasses`; Comet cannot add to that list, so + * [[CometKryoRegistrator]] has to be set explicitly, and this suite is what proves it is + * sufficient. + * + * This needs its own suite because `spark.serializer` and `spark.kryo.registrator` are read when + * `SparkEnv` builds the serializer, so they cannot be changed per test. + */ +class CometInMemoryCacheKryoSuite extends CometTestBase { + + import testImplicits._ + + 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 = super.sparkConf + conf.set("spark.plugins", "org.apache.spark.CometPlugin") + conf.set( + "spark.sql.cache.serializer", + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") + conf.set("spark.kryo.registrationRequired", "true") + conf.set("spark.kryo.registrator", CometKryoRegistrator.CLASS_NAME) + conf + } + + private def cachedBatchTypes(table: String): Array[String] = { + val cached = spark.sharedState.cacheManager.lookupCachedData(spark.table(table)).get + cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers + .map(_.getClass.getName) + .distinct() + .collect() + } + + // Every type whose bounds gatherColumnStats records, so the statistics row carries one of each + // internal representation Kryo has to write: boxed primitives, UTF8String, and Decimal at both + // sides of the long/BigDecimal split. + private val statsColumns = Seq( + "id AS c_long", + "cast(id % 2 = 0 as boolean) AS c_bool", + "cast(id % 100 as byte) AS c_byte", + "cast(id % 100 as short) AS c_short", + "cast(id as int) AS c_int", + "cast(id as float) AS c_float", + "cast(id as double) AS c_double", + "cast(id as decimal(9,2)) AS c_dec_short", + "cast(id as decimal(30,4)) AS c_dec_long", + "cast(id as string) AS c_string", + "cast(date '2020-01-01' + cast(id as int) as date) AS c_date", + "timestamp '2020-01-01 00:00:00' + make_interval(0, 0, 0, 0, 0, 0, id) AS c_ts") + + // DISK_ONLY and MEMORY_AND_DISK_SER both serialize the block on put, so each one reaches Kryo + // deterministically in local mode. The default MEMORY_AND_DISK reaches it only once a partition + // spills, which is the case that makes this more than a DISK_ONLY concern but is not something a + // test can force cheaply. + Seq(StorageLevel.DISK_ONLY, StorageLevel.MEMORY_AND_DISK_SER) + .foreach { level => + test(s"Comet in-memory cache round-trips through Kryo at $level") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + try { + spark + .range(0, 200, 1, 4) + .selectExpr(statsColumns: _*) + .createOrReplaceTempView("kryo_cache") + + spark.catalog.cacheTable("kryo_cache", level) + assert(spark.table("kryo_cache").count() == 200) + + assert( + cachedBatchTypes("kryo_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")), + "the payload Kryo serialized must be Comet's cached batch format") + + // Read the payload back rather than only the row count, so a Kryo round trip that + // silently mangles the Arrow bytes fails too. The predicate also exercises the + // statistics row, which is what carries UTF8String and Decimal through Kryo. + checkSparkAnswer( + spark.sql("SELECT c_long, c_string, c_dec_long, c_ts FROM kryo_cache " + + "WHERE c_dec_short >= 100 AND c_string > '1'")) + } finally { + spark.catalog.clearCache() + } + } + } + } + + test("Comet broadcast exchange survives Kryo with registration required") { + // Not about the cache: CometBroadcastExchangeExec broadcasts an Array[ChunkedByteBuffer], and + // Spark registers ChunkedByteBuffer but not an array of them, so this fails on main today + // under registrationRequired=true. The registrator this suite installs covers it because the + // cache write path hands back the same type. Kept here rather than split out because that + // registration is the thing under test. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { + withParquetTable((0 until 100).map(i => (i, i.toString)), "kryo_bcast_a") { + withParquetTable((0 until 10).map(i => (i, i.toString)), "kryo_bcast_b") { + val df = spark.sql( + "SELECT /*+ BROADCAST(b) */ a._1, b._2 " + + "FROM kryo_bcast_a a JOIN kryo_bcast_b b ON a._1 = b._1") + assert( + df.queryExecution.executedPlan.toString().contains("CometBroadcastExchange"), + "the broadcast has to run through Comet for this to test anything") + checkSparkAnswer(df) + } + } + } + } + + test("Comet in-memory cache falls back to Spark's format under Kryo for unsupported types") { + // A relation Comet cannot store is delegated to DefaultCachedBatch, which Spark registers + // itself. Pins that the fallback path is not collateral damage of the registration work. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + + spark.catalog.clearCache() + try { + spark + .range(0, 100, 1, 2) + .selectExpr("id", "make_interval(0, 0, 0, 0, 0, 0, id) AS iv") + .createOrReplaceTempView("kryo_cache_fallback") + + spark.catalog.cacheTable("kryo_cache_fallback", StorageLevel.DISK_ONLY) + assert(spark.table("kryo_cache_fallback").count() == 100) + assert( + cachedBatchTypes("kryo_cache_fallback").sameElements( + Array("org.apache.spark.sql.execution.columnar.DefaultCachedBatch"))) + + checkSparkAnswer(spark.sql("SELECT id FROM kryo_cache_fallback WHERE id > 90")) + } finally { + 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 index 11099e78ab9..959b5590b34 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -113,13 +113,13 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { val benchmark = new Benchmark(name, numRows, output = output) - benchmark.addCase("Comet cache disabled") { _ => + benchmark.addCase("Spark cache scan + CometSparkColumnarToColumnar") { _ => withSQLConf(cacheConf(nativeCacheEnabled = false): _*) { spark.sql(query).noop() } } - benchmark.addCase("Comet cache enabled") { _ => + benchmark.addCase("CometInMemoryTableScan") { _ => withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { spark.sql(query).noop() } @@ -132,14 +132,19 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { 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. + // Materialize the cache once using Comet's cache serializer, then read it both ways. // - // Both cases therefore read a Comet-written cache: spark.sql.cache.serializer is a static - // conf, so a single session cannot also materialize a DefaultCachedBatch to compare against. - // "Comet cache disabled" here means Spark execution over CometCachedBatch, NOT Spark's own - // cache format, and these numbers are not a baseline for it. + // What the two cases isolate is the cache-scan boundary, not the execution engine above it. + // cacheConf turns Comet execution on for both, so the aggregation runs on Comet either way; + // the only flag that moves is COMET_EXEC_IN_MEMORY_CACHE_ENABLED. Disabled, Spark's + // InMemoryTableScanExec feeds those same Comet operators through a + // CometSparkColumnarToColumnar bridge; enabled, CometInMemoryTableScan feeds them directly. + // So the numbers measure "keep the cached scan native" against "fall back to a Spark cache + // scan and convert" -- which is the overhead this feature exists to remove. + // + // Neither case is a baseline for Spark's own cache format. spark.sql.cache.serializer is a + // static conf, so a single session cannot also materialize a DefaultCachedBatch to compare + // against; both cases read the same Comet-written CometCachedBatch. withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { spark .sql(s"SELECT id, k, v, s1, s2, s3 FROM $sourceTable") @@ -155,6 +160,10 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { } } + // Pins the shape the case labels claim: enabled reads the cache natively with no conversion, + // disabled reads it through Spark's cache scan and a CometSparkColumnarToColumnar bridge. The + // bridge is what makes the disabled case a scan-boundary comparison rather than a Spark-vs-Comet + // execution one, since a Spark-columnar-to-Arrow transition only exists to feed Comet operators. private def verifyPlan(query: String, nativeCacheEnabled: Boolean): Unit = { val plan = spark.sql(query).queryExecution.executedPlan.toString() @@ -165,6 +174,9 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { assert( !plan.contains("CometInMemoryTableScan"), s"Native cache scan should be disabled:\n$plan") + assert( + plan.contains("CometSparkColumnarToColumnar"), + s"Expected the fallback read to bridge into Comet operators:\n$plan") } } From 1abc7ca3676416bd0f8a7e96ebca3ceab5e477a2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 28 Aug 2026 09:55:31 -0600 Subject: [PATCH 18/19] ci: register CometInMemoryCacheKryoSuite in the PR build workflows dev/ci/check-suites.py requires every suite to be listed in both pr_build_linux.yml and pr_build_macos.yml, so a new suite fails Preflight until it is registered in both. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 224af8fa903..eea569501ad 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -340,6 +340,7 @@ jobs: org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite org.apache.comet.exec.CometInMemoryCacheSuite + org.apache.comet.exec.CometInMemoryCacheKryoSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index a9719922b37..146872793f5 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -156,6 +156,7 @@ jobs: org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite org.apache.comet.exec.CometInMemoryCacheSuite + org.apache.comet.exec.CometInMemoryCacheKryoSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite From 444f0523bea2d927d290c0ef7475aade57126623 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 28 Aug 2026 10:37:18 -0600 Subject: [PATCH 19/19] fix: drop unused testImplicits import from the Kryo cache suite The lint matrix runs scalafix with -Psemanticdb, which is skipped on the default spark-4.1 profile because semanticdb-scalac_2.13.17 is unpublished, so RemoveUnused findings only surface on the 3.4, 3.5 and 4.0 jobs. --- .../org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala | 2 -- 1 file changed, 2 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala index 01c3b5ec185..13d0623a355 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala @@ -43,8 +43,6 @@ import org.apache.comet.{CometConf, CometKryoRegistrator} */ class CometInMemoryCacheKryoSuite extends CometTestBase { - import testImplicits._ - override protected def beforeAll(): Unit = { CometInMemoryRelationHelper.clearSerializer() super.beforeAll()