diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index b7c88254c93..eea569501ad 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -339,6 +339,8 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometInMemoryCacheSuite + org.apache.comet.exec.CometInMemoryCacheKryoSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 8210f91b7f9..146872793f5 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -155,6 +155,8 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometInMemoryCacheSuite + org.apache.comet.exec.CometInMemoryCacheKryoSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index d8fe5b69890..507bccf6138 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -258,6 +258,26 @@ object CometConf extends ShimCometConf { val COMET_EXEC_SAMPLE_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("sample", defaultValue = true) + val COMET_EXEC_IN_MEMORY_CACHE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.inMemoryCache.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether to enable Comet native execution for in-memory cached tables. Its value at " + + "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " + + "which stores cached data in Arrow format. Because spark.sql.cache.serializer is a " + + "static config, the cached format is fixed for the application, and disabling this " + + "at runtime only sends cached scans back to Spark's execution path. Relations whose " + + "schema Comet's Arrow writer does not support are always cached in Spark's default " + + "format. Each cached column is stored as its own compressed Arrow IPC stream, so a " + + "scan decodes only the columns it projected. Reads that feed Spark operators rather " + + "than Comet ones still pay a row conversion the default format avoids, and can be " + + "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " + + "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " + + "SparkContext, otherwise caching fails as soon as a block is serialized, including " + + "the disk half of the default MEMORY_AND_DISK storage level.") + .booleanConf + .createWithDefault(false) + val COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.enabled") .category(CATEGORY_EXEC) 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/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index c69602fc80b..17df4b0b422 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.execution.arrow.ArrowCachedBatchSerializer import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} import org.apache.spark.sql.execution.datasources.WriteFilesExec import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat @@ -86,6 +88,7 @@ object CometExecRule { classOf[SortMergeJoinExec] -> CometSortMergeJoinExec, classOf[SortExec] -> CometSortExec, classOf[LocalTableScanExec] -> CometLocalTableScanExec, + classOf[InMemoryTableScanExec] -> CometInMemoryTableScanExec, classOf[SampleExec] -> CometSampleExec, classOf[WindowExec] -> CometWindowExec) ++ // WindowGroupLimitExec exists only on Spark 3.5+; the shim returns None on 3.4. @@ -295,6 +298,48 @@ case class CometExecRule(session: SparkSession) case op if isCometScan(op) => convertToComet(op, CometScanWrapper).getOrElse(op) + case scan: InMemoryTableScanExec => + val serializer = scan.relation.cacheBuilder.serializer + val usesCometCacheSerializer = serializer.isInstanceOf[ArrowCachedBatchSerializer] + // The serializer only stores Comet's Arrow format for schemas it supports and delegates + // everything else to Spark's default cache format, which the native scan cannot read. + val cometCacheFormat = usesCometCacheSerializer && + ArrowCachedBatchSerializer.supportsSchema(scan.relation.output) + val nativeCacheEnabled = CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf) + + if (nativeCacheEnabled && cometCacheFormat) { + convertToComet(scan, CometInMemoryTableScanExec).getOrElse(scan) + } else { + // The native cache scan is not available for this relation. Record why, then take the + // same SparkToColumnar fallback that any other unsupported operator would take, so + // that turning the feature on is never worse for a scan than leaving it off. + if (nativeCacheEnabled && !usesCometCacheSerializer) { + withFallbackReason( + scan, + s"Comet in-memory cache requires ${classOf[ArrowCachedBatchSerializer].getName} " + + s"but this relation was cached with ${serializer.getClass.getName}") + } else if (nativeCacheEnabled) { + val unsupported = scan.relation.output + .filterNot(a => ArrowCachedBatchSerializer.supportsType(a.dataType)) + .map(a => s"${a.name}: ${a.dataType.simpleString}") + withFallbackReason( + scan, + "Comet in-memory cache does not support the type of these cached columns, so the " + + s"relation was cached in Spark's default format: ${unsupported.mkString(", ")}") + } else if (usesCometCacheSerializer) { + withFallbackReason( + scan, + "Native support for operator InMemoryTableScanExec is disabled. " + + s"Set ${CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key}=true to enable it.") + } + + if (shouldApplySparkToColumnar(conf, scan)) { + convertToComet(scan, CometSparkToColumnarExec).getOrElse(scan) + } else { + scan + } + } + case op if shouldApplySparkToColumnar(conf, op) => convertToComet(op, CometSparkToColumnarExec).getOrElse(op) 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/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 321d530e9ee..eaeac316655 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -29,7 +29,9 @@ import org.apache.spark.internal.config.{EXECUTOR_MEMORY, EXECUTOR_MEMORY_OVERHE import org.apache.spark.sql.internal.StaticSQLConf import org.apache.comet.{COMET_VERSION, CometSparkSessionExtensions, NativeBase} +import org.apache.comet.CometConf import org.apache.comet.CometConf.{COMET_METRICS_ENABLED, COMET_ONHEAP_ENABLED} +import org.apache.comet.CometKryoRegistrator import org.apache.comet.annotation.Public /** @@ -61,6 +63,11 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl return Collections.emptyMap[String, String] } + val extraConfs = new ju.HashMap[String, String]() + + CometDriverPlugin.maybeSetCacheSerializer(sc.conf, extraConfs) + CometDriverPlugin.warnIfKryoRegistratorMissing(sc.conf) + // register CometSparkSessionExtensions if it isn't already registered CometDriverPlugin.registerCometSessionExtension(sc.conf) @@ -94,7 +101,7 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl logInfo("Comet is running in unified memory mode and sharing off-heap memory with Spark") } - Collections.emptyMap[String, String] + extraConfs } override def receive(message: Any): AnyRef = super.receive(message) @@ -117,6 +124,58 @@ object CometDriverPlugin extends Logging { /** Spark config key under which the loaded Comet version is exposed at runtime. */ val COMET_VERSION_CONFIG = "spark.comet.version" + // Use Comet's cache serializer only for the native in-memory cache path. + // If the application already set spark.sql.cache.serializer, leave that value + // unchanged so Comet does not replace a user-selected cache format. + private[apache] def maybeSetCacheSerializer( + conf: SparkConf, + extraConfs: ju.HashMap[String, String]): Unit = { + if (conf.getBoolean(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, false)) { + val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key + val serializerValue = + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer" + val defaultSerializer = StaticSQLConf.SPARK_CACHE_SERIALIZER.defaultValueString + val currentSerializer = conf.get(serializerKey, defaultSerializer) + + if (currentSerializer == defaultSerializer) { + extraConfs.put(serializerKey, serializerValue) + conf.set(serializerKey, serializerValue) + logInfo(s"Auto-set $serializerKey=$serializerValue") + } else { + logInfo(s"Not overriding user-provided $serializerKey=$currentSerializer") + } + } + } + + // Comet hands Spark's serializer classes that Kryo has not been told about, so with + // spark.kryo.registrationRequired=true it rejects them with "Class is not registered", which + // names neither Comet nor the operation that failed. Two paths reach it: a native broadcast, + // which broadcasts an Array[ChunkedByteBuffer], and any cached block Spark serializes -- the + // disk half of MEMORY_AND_DISK, the _SER levels, replication, a cross-executor fetch. + // CometKryoRegistrator covers both, but spark.kryo.registrator is read when SparkEnv builds the + // serializer, before any plugin runs, so it cannot be set from here. Say so while the + // application is still starting up rather than leaving the user to attribute the failure later. + private[apache] def warnIfKryoRegistratorMissing(conf: SparkConf): Unit = { + val usingKryo = + conf.get("spark.serializer", "") == "org.apache.spark.serializer.KryoSerializer" + val registrationRequired = conf.getBoolean("spark.kryo.registrationRequired", false) + val registered = conf + .get("spark.kryo.registrator", "") + .split(',') + .map(_.trim) + .contains(CometKryoRegistrator.CLASS_NAME) + + if (usingKryo && registrationRequired && !registered) { + logWarning( + "spark.kryo.registrationRequired=true but spark.kryo.registrator does not include " + + s"${CometKryoRegistrator.CLASS_NAME}. Comet's native broadcast and its in-memory " + + "cache format will fail with Kryo's \"Class is not registered\" as soon as their " + + "payloads are serialized. Add " + + s"spark.kryo.registrator=${CometKryoRegistrator.CLASS_NAME} before creating the " + + "SparkContext; it cannot be set later.") + } + } + def registerCometMetrics(sc: SparkContext): Unit = { if (sc.getConf.getBoolean( COMET_METRICS_ENABLED.key, 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..261102bc72a --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -0,0 +1,158 @@ +/* + * 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.CachedBatchSerializer +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 + +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, + cacheBuilder: CachedRDDBuilder, + 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")) + + // `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 + + // `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) + + // 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") + + // 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) + cachedBuffers.mapPartitionsWithIndex(filter) + } else { + cachedBuffers + } + + serializer + .convertCachedBatchToColumnarBatch(filteredBuffers, 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] = { + + val scanTypes = op.output.flatMap(attr => serializeDataType(attr.dataType)) + + val scanBuilder = OperatorOuterClass.Scan + .newBuilder() + .setSource(op.getClass.getSimpleName) + .addAllFields(scanTypes.asJava) + + 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 + + CometScanWrapper( + nativeOp, + CometInMemoryTableScanExec( + op, + relation.cacheBuilder.serializer, + relation.cacheBuilder, + relation.output, + op.output)) + } + +} 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..46a51ad8775 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -0,0 +1,570 @@ +/* + * 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 scala.util.control.NonFatal + +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} +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.internal.SQLConf +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.CometArrowAllocator + +/** + * Cached batch format used when Comet writes Spark in-memory cache data. + * + * `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, + columns: Array[ChunkedByteBuffer]) + extends SimpleMetricsCachedBatch + +/** + * Cache serializer that stores Comet-compatible Arrow batches in Spark's in-memory cache. + * + * 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 { + + import ArrowCachedBatchSerializer.supportsSchema + + private val fallback = new DefaultCachedBatchSerializer() + + // 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) + 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 + } + + (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) + 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 + // 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 + } + + 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. + // + // 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] = { + val arrowSchema = + Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE) + + batches.map { batch => + // 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 Utils.serializeBatchColumns(arrowBatch) + finally arrowBatch.close() + } + + 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] = { + 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 + // 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) + } + + // 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 + // 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], + storageLevel: StorageLevel, + conf: SQLConf): RDD[CachedBatch] = { + + input.mapPartitions { batches => + encodeBatches(batches, schema) + } + } + + 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 => + // 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 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 readers = current + current = null + if (readers != null) { + readers.close() + } + } + } + + it.flatMap { + case cb: CometCachedBatch => + 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 => + throw new IllegalStateException( + s"Unsupported cached batch type ${other.getClass.getName}") + } + } + } + + // 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) { + // 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 { + 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], + schema: Seq[Attribute], + storageLevel: StorageLevel, + conf: SQLConf): RDD[CachedBatch] = { + + if (!supportsSchema(schema)) { + fallback.convertInternalRowToCachedBatch(input, schema, storageLevel, conf) + } else { + val batchSize = conf.columnBatchSize + + input.mapPartitions { rows => + val iter = CometArrowConverters.rowToArrowBatchIter( + rows, + Utils.fromAttributes(schema), + batchSize, + // 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) + } + } + } + + override def convertCachedBatchToInternalRow( + input: RDD[CachedBatch], + 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) + + batches.flatMap { batch => + batch.rowIterator().asScala.map(row => toUnsafe(row).copy()) + } + } + } +} + +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)) + + /** + * 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/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/execution/arrow/CometArrowConverters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala index 0b099a04373..e68eee6b79b 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 @@ -31,14 +33,16 @@ import org.apache.spark.sql.vectorized.ColumnarBatch 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 { @@ -63,15 +67,67 @@ object CometArrowConverters extends Logging { override def next(): ColumnarBatch = { val root = VectorSchemaRoot.create(arrowSchema, allocator) - val writer = ArrowWriter.create(root, maxRecordsPerBatch) - var rowCount = 0 - while (rowIter.hasNext && 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. + closingRootOnFailure(root) { + val writer = ArrowWriter.create(root, maxRecordsPerBatch) + var rowCount = 0 + while (rowIter.hasNext && rowCount < maxRecordsPerBatch) { + writer.write(rowIter.next()) + rowCount += 1 + } + writer.finish() + NativeUtil.rootAsBatch(root) } - writer.finish() - NativeUtil.rootAsBatch(root) } } } + + /** + * 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 numRows = batch.numRows() + 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. + closingRootOnFailure(root) { + val writer = ArrowWriter.create(root, numRows) + writer.writeColumns(batch, 0, numRows) + writer.finish() + 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) => + try { + root.close() + } catch { + case NonFatal(closeError) => e.addSuppressed(closeError) + } + throw e + } + } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala index 4697c172601..e2454f51322 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala @@ -158,7 +158,20 @@ object CometArrowStream extends Logging { val expectedFields = expected.getFields val actualFields = (0 until first.numCols()).map { i => val col = first.column(i).asInstanceOf[CometVector] - actualFieldOf(col, expectedFields.get(i)) + + // The Arrow C Stream exports the producer's actual Arrow schema. If the + // consumer-declared schema has fewer fields, fall back to the producer's + // field and synthesize a non-null field name, since Arrow Field names cannot + // be null. + val expectedField = + if (i < expectedFields.size()) { + expectedFields.get(i) + } else { + val raw = col.getValueVector.getField + new Field(s"_c$i", raw.getFieldType, raw.getChildren) + } + + actualFieldOf(col, expectedField) } val mismatches = actualFields.zip(expectedFields.asScala).zipWithIndex.collect { case ((actual, exp), idx) if actual.getType != exp.getType => 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 216f9f1e4bd..3700e97642b 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 @@ -862,7 +862,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/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 9d4b0bce881..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 @@ -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} @@ -268,6 +269,54 @@ 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 codec = CompressionCodec.createCodec(SparkEnv.get.conf) + + // 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)) + + 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 + } + + /** + * 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. * @@ -398,33 +447,109 @@ object Utils extends CometTypeShim with Logging { } } + /** + * 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) 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), 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]] 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])] = { 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( @@ -438,19 +563,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/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 4eb6d001787..89e9d69bc50 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} @@ -3786,6 +3787,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/CometInMemoryCacheKryoSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala new file mode 100644 index 00000000000..13d0623a355 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala @@ -0,0 +1,189 @@ +/* + * 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 { + + 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/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala new file mode 100644 index 00000000000..a7a9a6c8590 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -0,0 +1,1407 @@ +/* + * 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 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 +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, InMemoryRelation} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.storage.StorageLevel + +import org.apache.comet.{CometArrowAllocator, 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 + // 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") + 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 cached = spark.sharedState.cacheManager.lookupCachedData(spark.table(table)).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 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 -> "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 + .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") + spark.table("default_cached_batch").count() + + assert( + cachedBatchTypes("default_cached_batch").sameElements( + Array("org.apache.spark.sql.execution.columnar.DefaultCachedBatch"))) + + // Columnar read path, delegated to Spark's serializer. + val columnarDf = spark.sql(""" + 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. + withSQLConf(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "false") { + val rowDf = spark.sql(""" + 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() + assert(!rowPlan.contains("CometInMemoryTableScan")) + } + + 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")) + assert(!emptyPlan.contains("CometSparkColumnarToColumnar")) + + 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")) + // 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() + } + } + + 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() + } + } + + 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 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) + + // 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}") + 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 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 = + "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() + } + } + + 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", + 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() + } + } + + /** + * 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( + "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) + } { + 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") { + withSparkColumnarCache( + "spark_columnar_complex", + SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true") { 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", + // 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) + } { + assert(spark.table("spark_columnar_complex").count() == 200) + + 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 { (_, 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 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", + "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.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", + // 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") { + + 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") + } + } + + /** + * 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/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala new file mode 100644 index 00000000000..959b5590b34 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -0,0 +1,192 @@ +/* + * 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", + "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( + "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) + + // 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") + + 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") + } + } + + 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("Spark cache scan + CometSparkColumnarToColumnar") { _ => + withSQLConf(cacheConf(nativeCacheEnabled = false): _*) { + spark.sql(query).noop() + } + } + + benchmark.addCase("CometInMemoryTableScan") { _ => + 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, then read it both ways. + // + // 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") + .createOrReplaceTempView(cacheTable) + spark.catalog.cacheTable(cacheTable) + spark.table(cacheTable).count() + } + + try f + finally { + spark.catalog.uncacheTable(cacheTable) + spark.catalog.clearCache() + } + } + + // 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() + + 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") + assert( + plan.contains("CometSparkColumnarToColumnar"), + s"Expected the fallback read to bridge into Comet operators:\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") + } +} 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..5548f6dadbd --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala @@ -0,0 +1,119 @@ +/* + * 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.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, ChunkedByteBufferOutputStream} + +import org.apache.comet.CometArrowAllocator + +/** + * 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)))) + } + + /** 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) +} 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() + } + } + } } 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() +}