diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 019ae27722..5e01a713ec 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -336,6 +336,7 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 0c02fdf8ce..a3ca4be684 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -152,6 +152,7 @@ jobs: org.apache.comet.exec.CometAggregateSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 958a22cc2b..f520cc756f 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -234,6 +234,20 @@ object CometConf extends ShimCometConf { val COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("localTableScan", defaultValue = false) + val COMET_EXEC_IN_MEMORY_CACHE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.inMemoryCache.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether to enable Comet native execution for in-memory cached tables. Its value at " + + "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " + + "which stores cached data in Arrow format. Because spark.sql.cache.serializer is a " + + "static config, the cached format is fixed for the application, and disabling this " + + "at runtime only sends cached scans back to Spark's execution path. Relations whose " + + "schema Comet's Arrow writer does not support are always cached in Spark's default " + + "format.") + .booleanConf + .createWithDefault(false) + val COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index bd3b92f2cc..780c7d0cb3 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[WindowExec] -> CometWindowExec) /** @@ -283,6 +286,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/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index 43945b97fd..1ef3cd4f12 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -29,6 +29,7 @@ import org.apache.spark.internal.config.{EXECUTOR_MEMORY, EXECUTOR_MEMORY_OVERHE import org.apache.spark.sql.internal.StaticSQLConf import org.apache.comet.{CometSparkSessionExtensions, NativeBase} +import org.apache.comet.CometConf import org.apache.comet.CometConf.{COMET_METRICS_ENABLED, COMET_ONHEAP_ENABLED} /** @@ -54,6 +55,10 @@ 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) + // register CometSparkSessionExtensions if it isn't already registered CometDriverPlugin.registerCometSessionExtension(sc.conf) @@ -87,7 +92,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) @@ -106,6 +111,29 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl } object CometDriverPlugin extends Logging { + // Use Comet's cache serializer only for the native in-memory cache path. + // If the application already set spark.sql.cache.serializer, leave that value + // unchanged so Comet does not replace a user-selected cache format. + private[apache] def maybeSetCacheSerializer( + conf: SparkConf, + extraConfs: ju.HashMap[String, String]): Unit = { + if (conf.getBoolean(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, false)) { + val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key + val serializerValue = + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer" + val defaultSerializer = StaticSQLConf.SPARK_CACHE_SERIALIZER.defaultValueString + val currentSerializer = conf.get(serializerKey, defaultSerializer) + + if (currentSerializer == defaultSerializer) { + extraConfs.put(serializerKey, serializerValue) + conf.set(serializerKey, serializerValue) + logInfo(s"Auto-set $serializerKey=$serializerValue") + } else { + logInfo(s"Not overriding user-provided $serializerKey=$currentSerializer") + } + } + } + def registerCometMetrics(sc: SparkContext): Unit = { if (sc.getConf.getBoolean( COMET_METRICS_ENABLED.key, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala new file mode 100644 index 0000000000..7cc1d74ead --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -0,0 +1,143 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet + +import scala.collection.JavaConverters._ + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer} +import org.apache.spark.sql.execution.LeafExecNode +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.CometConf +import org.apache.comet.serde.CometOperatorSerde +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.QueryPlanSerde.serializeDataType + +/** + * Reads Spark cached table data when the cache was written by Comet's cache serializer. + * + * Spark stores cached data through `CachedBatchSerializer`. This node keeps the scan inside Comet + * by asking the serializer to decode cached batches directly into `ColumnarBatch` output, + * avoiding the extra Spark columnar-to-Comet columnar conversion used by the default path. + * + * `relationOutput` is the full schema stored in the cache. `scanOutput` is the subset requested + * by this scan after pruning. + */ +case class CometInMemoryTableScanExec( + originalPlan: InMemoryTableScanExec, + serializer: CachedBatchSerializer, + cachedBuffers: RDD[CachedBatch], + relationOutput: Seq[Attribute], + scanOutput: Seq[Attribute]) + extends CometExec + with LeafExecNode { + + override lazy val metrics: Map[String, SQLMetric] = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + + // For an empty-projection scan (`SELECT count(*)`) this is empty while `scanOutput` holds the + // full cache schema, so the emitted batches are wider than the declared output. That is safe + // because the only consumer of an empty-output scan is a count-style aggregate, which reads + // the row count rather than any column; `convert` and `createExec` deliberately fall back to + // the cache schema in that case because the native plan still needs a non-empty scan schema. + override def output: Seq[Attribute] = originalPlan.output + + // Use the serializer's vector types because the cached batch layout is owned by the serializer. + 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") + + 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] = { + + // Empty-output scans still need a schema for native planning, so fall back to the cache schema. + val actualOutput = + if (op.output.nonEmpty) op.output + else op.relation.output + + val scanTypes = actualOutput.flatMap(attr => serializeDataType(attr.dataType)) + + val scanBuilder = OperatorOuterClass.Scan + .newBuilder() + .setSource(op.getClass.getSimpleName) + .addAllFields(scanTypes.asJava) + + Some(builder.setScan(scanBuilder).build()) + } + + // Reuse Spark's InMemoryRelation metadata so cache materialization, pruning, and storage + // behavior remain controlled by Spark's cache manager. + override def createExec(nativeOp: Operator, op: InMemoryTableScanExec): CometNativeExec = { + val relation = op.relation + + val actualOutput = + if (op.output.nonEmpty) op.output + else relation.output + + CometScanWrapper( + nativeOp, + CometInMemoryTableScanExec( + op, + relation.cacheBuilder.serializer, + relation.cacheBuilder.cachedColumnBuffers, + relation.output, + actualOutput)) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala new file mode 100644 index 0000000000..e6d5aa06cf --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -0,0 +1,435 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import scala.collection.JavaConverters._ + +import org.apache.spark.TaskContext +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull, UnsafeProjection} +import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch, SimpleMetricsCachedBatchSerializer} +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer +import org.apache.spark.sql.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. + * + * `bytes` contains compressed Arrow stream data produced by `Utils.serializeBatches`. The cache + * manager still owns storage and eviction; this class only changes the cached payload. + */ +private case class CometCachedBatch( + override val numRows: Int, + override val sizeInBytes: Long, + override val stats: InternalRow, + bytes: 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() + + // Build the statistics row expected by SimpleMetricsCachedBatchSerializer. + // For each cached column Spark expects five values in this order: + // lower bound, upper bound, null count, row count, and size in bytes. + private def computeStats(batch: ColumnarBatch, attrs: Seq[Attribute]): InternalRow = { + val numCols = attrs.length + val lower = new Array[Any](numCols) + val upper = new Array[Any](numCols) + val nulls = Array.fill[Int](numCols)(0) + val numRows = batch.numRows() + + var c = 0 + while (c < numCols) { + val dt = attrs(c).dataType + val col = batch.column(c) + var r = 0 + while (r < numRows) { + if (col.isNullAt(r)) { + nulls(c) += 1 + } else if (tracksBounds(dt)) { + val value = readValue(col, dt, r) + if (lower(c) == null || compare(dt, value, lower(c)) < 0) { + lower(c) = value + } + if (upper(c) == null || compare(dt, value, upper(c)) > 0) { + upper(c) = value + } + } + r += 1 + } + c += 1 + } + + val values = new Array[Any](numCols * 5) + c = 0 + while (c < numCols) { + val base = c * 5 + values(base) = lower(c) + values(base + 1) = upper(c) + values(base + 2) = nulls(c) + values(base + 3) = numRows + // Spark reserves the fifth field for per-column size. Comet stores the whole + // Arrow stream as one compressed buffer, so per-column size is not tracked here. + // Cache pruning uses bounds/null-count/row-count, not this size field. + values(base + 4) = 0L + c += 1 + } + + new GenericInternalRow(values) + } + + // Spark can prune cache batches only for types whose bounds can be compared. + // Other types still report null count and row count but leave bounds as null. + private def tracksBounds(dt: DataType): Boolean = dt match { + case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | + _: DecimalType | StringType | DateType | TimestampType | TimestampNTZType => + true + case _ => false + } + + // Read a non-null value from a ColumnVector using Spark's internal value type + // for the corresponding DataType. + private def readValue(col: ColumnVector, dt: DataType, rowId: Int): Any = dt match { + case BooleanType => col.getBoolean(rowId) + case ByteType => col.getByte(rowId) + case ShortType => col.getShort(rowId) + case IntegerType | DateType => col.getInt(rowId) + case LongType | TimestampType | TimestampNTZType => col.getLong(rowId) + case FloatType => col.getFloat(rowId) + case DoubleType => col.getDouble(rowId) + case d: DecimalType => col.getDecimal(rowId, d.precision, d.scale) + case StringType => col.getUTF8String(rowId).copy() + case _ => null + } + + // Compare values using the same physical representation used in the stats row. + private def compare(dt: DataType, left: Any, right: Any): Int = dt match { + case BooleanType => + java.lang.Boolean.compare(left.asInstanceOf[Boolean], right.asInstanceOf[Boolean]) + case ByteType => + java.lang.Byte.compare(left.asInstanceOf[Byte], right.asInstanceOf[Byte]) + case ShortType => + java.lang.Short.compare(left.asInstanceOf[Short], right.asInstanceOf[Short]) + case IntegerType | DateType => + java.lang.Integer.compare(left.asInstanceOf[Int], right.asInstanceOf[Int]) + case LongType | TimestampType | TimestampNTZType => + java.lang.Long.compare(left.asInstanceOf[Long], right.asInstanceOf[Long]) + case FloatType => + java.lang.Float.compare(left.asInstanceOf[Float], right.asInstanceOf[Float]) + case DoubleType => + java.lang.Double.compare(left.asInstanceOf[Double], right.asInstanceOf[Double]) + case _: DecimalType => + left.asInstanceOf[Decimal].compare(right.asInstanceOf[Decimal]) + case StringType => + ByteArray.compareBinary( + left.asInstanceOf[UTF8String].getBytes, + right.asInstanceOf[UTF8String].getBytes) + case other => + throw new IllegalStateException(s"compare called for unsupported type $other") + } + + // Compute Spark-compatible cache stats before serializing each batch to Arrow. + // The stats are stored beside the Arrow bytes so Spark's cache filter can prune + // CometCachedBatch without decoding the batch first. + // + // 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 => + val stats = computeStats(batch, attrs) + + if (Utils.isArrowBacked(batch)) { + serializeBatch(batch, stats) + } else { + val arrowBatch = + CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) + try serializeBatch(arrowBatch, stats) + finally arrowBatch.close() + } + } + } + + // Utils.serializeBatches is one-in/one-out, so take the single element eagerly: the write has to + // happen before a converted batch is closed. + private def serializeBatch(batch: ColumnarBatch, stats: InternalRow): CachedBatch = { + val (rows, buffer) = Utils.serializeBatches(Iterator.single(batch)).next() + CometCachedBatch( + numRows = rows.toInt, + sizeInBytes = buffer.size, + stats = stats, + bytes = buffer) + } + + // Resolve requested columns by exprId, not by name, because aliases may reuse names. + private def selectedIndices( + cacheAttributes: Seq[Attribute], + selectedAttributes: Seq[Attribute]): Array[Int] = { + if (selectedAttributes.isEmpty) { + cacheAttributes.indices.toArray + } else { + val byExprId = cacheAttributes.zipWithIndex.map { case (attr, idx) => + attr.exprId -> idx + }.toMap + + selectedAttributes.map { attr => + byExprId.getOrElse( + attr.exprId, + throw new IllegalStateException( + s"Could not resolve selected attribute ${attr.name} from cache attributes")) + }.toArray + } + } + + // A full-width projection is only an identity projection if every selected index + // is already in column order. For example, [1, 0] must still be projected. + private def isIdentityProjection(indices: Array[Int], numCols: Int): Boolean = + indices.length == numCols && indices.indices.forall(i => indices(i) == i) + + private def projectBatch(batch: ColumnarBatch, indices: Array[Int]): ColumnarBatch = { + if (isIdentityProjection(indices, batch.numCols())) { + batch + } else { + val cols = indices.map(i => batch.column(i).asInstanceOf[ColumnVector]) + new ColumnarBatch(cols, batch.numRows()) + } + } + + // 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 => + // An ArrowReaderIterator closes its reader (releasing the batch it is holding) only when it + // runs to exhaustion. A consumer that stops early -- LIMIT, take(), or a cancelled task -- + // leaves the reader it was part-way through open, so close it on task completion. Spark's + // own ArrowCachedBatchSerializer registers a listener for the same reason. + // + // flatMap consumes each inner iterator fully before building the next, so at most one reader + // is open at a time and tracking the current one is enough. close() is idempotent, so closing + // one that already exhausted itself is a no-op. + @volatile var current: ArrowReaderIterator = null + Option(TaskContext.get()).foreach { tc => + tc.addTaskCompletionListener[Unit] { _ => + val reader = current + current = null + if (reader != null) { + reader.close() + } + } + } + + it.flatMap { + case cb: CometCachedBatch => + Utils.decodeBatches(cb.bytes, "CometCache") match { + case reader: ArrowReaderIterator => + current = reader + reader.map(batch => projectBatch(batch, indices)) + case empty => + empty.map(batch => projectBatch(batch, indices)) + } + + case other => + throw new IllegalStateException( + s"Unsupported cached batch type ${other.getClass.getName}") + } + } + } + + // 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)) +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala index 2d4fd71376..b492fbafbe 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.comet.execution.arrow +import scala.util.control.NonFatal + import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.types.pojo.Schema @@ -26,19 +28,21 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch} import org.apache.comet.vector.NativeUtil /** - * 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 { @@ -62,16 +66,80 @@ object CometArrowConverters extends Logging { override def next(): ColumnarBatch = { val root = VectorSchemaRoot.create(arrowSchema, allocator) - val writer = ArrowWriter.create(root) - var rowCount = 0L - while (rowIter.hasNext && - (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch)) { - writer.write(rowIter.next()) - rowCount += 1 + // Same ownership rule as columnarBatchToArrowBatch: the caller only owns the batch that + // rootAsBatch returns, so a throw from writing a row has to release the root here. + try { + val writer = ArrowWriter.create(root) + var rowCount = 0L + while (rowIter.hasNext && + (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch)) { + writer.write(rowIter.next()) + rowCount += 1 + } + writer.finish() + NativeUtil.rootAsBatch(root) + } catch { + case NonFatal(e) => + root.close() + throw e } - writer.finish() - NativeUtil.rootAsBatch(root) } } } + + /** + * Copy `numRows` rows starting at `startRow` from a Spark `ColumnarBatch` into `root`. + * + * Spark's `ColumnVector` implementations do not expose Arrow buffers, so values are necessarily + * copied element-wise. Shared by [[SparkColumnarArrowReader]], which slices into a stable root, + * and [[columnarBatchToArrowBatch]], which fills a fresh one. + */ + private[arrow] def writeColumns( + root: VectorSchemaRoot, + batch: ColumnarBatch, + startRow: Int, + numRows: Int): Unit = { + val writer = ArrowWriter.create(root) + var col = 0 + while (col < batch.numCols()) { + val column = batch.column(col) + val columnArray = new ColumnarArray(column, startRow, numRows) + if (column.hasNull) { + writer.writeCol(columnArray, col) + } else { + writer.writeColNoNull(columnArray, col) + } + col += 1 + } + writer.finish() + // ArrowWriter derives the root row count from its per-column writes, so a zero-column input + // batch (Spark's count-from-metadata scan: numRows > 0, numCols == 0) would otherwise produce + // a root with rowCount == 0 and silently drop the rows. + root.setRowCount(numRows) + } + + /** + * Copy a Spark `ColumnarBatch` whose columns are not Arrow-backed (e.g. + * `On/OffHeapColumnVector` from Spark's vectorized Parquet reader, or a third-party connector's + * vectors) into a freshly allocated Arrow `ColumnarBatch` of `CometVector`s. + * + * The input batch is not consumed or closed; the caller owns the returned batch and must close + * it. + */ + def columnarBatchToArrowBatch( + batch: ColumnarBatch, + arrowSchema: Schema, + allocator: BufferAllocator): ColumnarBatch = { + val root = VectorSchemaRoot.create(arrowSchema, allocator) + // The caller only owns the returned batch, so anything that throws before `rootAsBatch` wraps + // the root has to release it here or the allocation leaks. + try { + writeColumns(root, batch, 0, batch.numRows()) + NativeUtil.rootAsBatch(root) + } catch { + case NonFatal(e) => + root.close() + throw e + } + } } diff --git a/spark/src/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 4697c17260..e2454f5132 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/execution/arrow/SparkColumnarArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala index 157aca7423..b2a6e048d7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala @@ -22,13 +22,13 @@ package org.apache.spark.sql.comet.execution.arrow import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.Schema -import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch} +import org.apache.spark.sql.vectorized.ColumnarBatch /** * `ArrowReader` over an iterator of Spark-side `ColumnarBatch`es (not Arrow-backed). Slices up to * `maxRecordsPerBatch` rows per `loadNextBatch` from the current Spark batch into the reader's - * stable VSR via `ArrowWriter.writeCol`. Spark's `ColumnVector` implementations aren't Arrow - * buffers, so this reader necessarily copies element values into Arrow format. + * stable VSR via `CometArrowConverters.writeColumns`. Spark's `ColumnVector` implementations + * aren't Arrow buffers, so this reader necessarily copies element values into Arrow format. */ private[comet] class SparkColumnarArrowReader( allocator: BufferAllocator, @@ -76,26 +76,13 @@ private[comet] class SparkColumnarArrowReader( if (maxRecordsPerBatch <= 0) rowsRemaining else math.min(maxRecordsPerBatch, rowsRemaining) - val writer = ArrowWriter.create(getVectorSchemaRoot) - var col = 0 - while (col < current.numCols()) { - val column = current.column(col) - val columnArray = new ColumnarArray(column, rowsConsumedInCurrent, rowsToProduce) - if (column.hasNull) { - writer.writeCol(columnArray, col) - } else { - writer.writeColNoNull(columnArray, col) - } - col += 1 - } + CometArrowConverters.writeColumns( + getVectorSchemaRoot, + current, + rowsConsumedInCurrent, + rowsToProduce) rowsConsumedInCurrent += rowsToProduce - writer.finish() - // ArrowWriter derives the root row count from its per-column writes, so a zero-column - // input batch (Spark's count-from-metadata scan: numRows > 0, numCols == 0) would otherwise - // produce a root with rowCount == 0 and silently drop the rows. Set rowCount explicitly so - // downstream aggregations (e.g. df.count()) see the correct value. - getVectorSchemaRoot.setRowCount(rowsToProduce) onConversionNs(System.nanoTime() - startNs) true } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index fcd48d69b1..8355707a5f 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 @@ -821,7 +821,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 164d53fe19..8844605505 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -397,6 +397,15 @@ object Utils extends CometTypeShim with Logging { } } + /** + * Whether every column in `batch` satisfies [[getBatchFieldVectors]]'s precondition, i.e. is an + * Arrow-backed `CometVector`. Callers that may receive batches from a plan they did not build + * (e.g. Comet's cache serializer, which Spark hands the cached plan's columnar output) use this + * to convert foreign vectors to Arrow instead of tripping the exception below. + */ + def isArrowBacked(batch: ColumnarBatch): Boolean = + (0 until batch.numCols()).forall(i => batch.column(i).isInstanceOf[CometVector]) + def getBatchFieldVectors( batch: ColumnarBatch): (Seq[FieldVector], Option[DictionaryProvider]) = { var provider: Option[DictionaryProvider] = None diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index ce6071bf1a..bf60d5de3f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -38,6 +38,7 @@ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, Comet import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec} +import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec, SortMergeJoinExec} @@ -3700,6 +3701,34 @@ class CometExecSuite extends CometTestBase { }) } + test("SparkToColumnar over InMemoryTableScanExec with a non-Comet cache serializer") { + // Enabling the native in-memory cache must never leave a cached scan worse off than having + // the feature disabled. When the relation was cached by a serializer Comet cannot decode, the + // native scan is unavailable, but the scan should still take the SparkToColumnar fallback + // rather than staying entirely on Spark. + // + // This session does not configure spark.sql.cache.serializer, so it uses Spark's default. + // The reset makes that deterministic regardless of which suite ran first in this JVM, since + // InMemoryRelation memoizes the serializer per JVM. + CometInMemoryRelationHelper.clearSerializer() + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + spark + .range(1000) + .selectExpr("id as key", "id % 8 as value") + .createOrReplaceTempView("foreign_cache_serializer") + spark.catalog.cacheTable("foreign_cache_serializer") + try { + val df = spark.sql("SELECT * FROM foreign_cache_serializer").groupBy("key").count() + checkSparkAnswerAndOperator(df, includeClasses = Seq(classOf[CometSparkToColumnarExec])) + } finally { + spark.catalog.uncacheTable("foreign_cache_serializer") + } + } + } + test("SparkToColumnar eliminate redundant in AQE") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala new file mode 100644 index 0000000000..839ede5281 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -0,0 +1,966 @@ +/* + * 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, Expression, GreaterThanOrEqual, LessThan, Literal} +import org.apache.spark.sql.columnar.SimpleMetricsCachedBatch +import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} +import org.apache.spark.storage.StorageLevel + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.vector.CometVector + +class CometInMemoryCacheSuite extends CometTestBase { + + import testImplicits._ + + // `InMemoryRelation` resolves `spark.sql.cache.serializer` once per JVM and memoizes the + // instance in a static field. Test suites share a forked JVM, so whichever suite caches a + // table first pins the serializer for everything that follows: without this reset the + // 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")) + assert(plan.contains("CometNativeColumnarToRow")) + + spark.catalog.clearCache() + } + } + + test("Comet in-memory cache supports shuffle after cache read") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + + spark + .range(1000) + .selectExpr("id as key", "id % 100 as group") + .createOrReplaceTempView("shuffle_cache") + + spark.catalog.cacheTable("shuffle_cache") + spark.table("shuffle_cache").count() + + assert( + cachedBatchTypes("shuffle_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + val df = spark.sql("SELECT group, count(*) FROM shuffle_cache GROUP BY group") + checkSparkAnswer(df) + + val plan = df.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + assert(plan.contains("CometHashAggregate")) + + spark.catalog.clearCache() + } + } + + 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) + + val scans = df.queryExecution.executedPlan.collect { + case s: org.apache.spark.sql.comet.CometInMemoryTableScanExec => s + } + assert(scans.length == 1, s"expected one CometInMemoryTableScan, got ${scans.length}") + // scalastyle:off println + println( + "DIAG rows=" + df.collect().length + " metrics=" + scans.head.metrics + .map { case (k, v) => k + "=" + v.value } + .mkString(",")) + println("DIAG plan=" + df.queryExecution.executedPlan.getClass.getName) + // scalastyle:on println + result = (scans.head.metrics("numOutputRows").value, totalRows) + spark.catalog.clearCache() + } + result + } + + val (prunedRows, total) = scanRowsFor(pruning = true) + val (unprunedRows, total2) = scanRowsFor(pruning = false) + assert(total == total2) + // With pruning on, only the batch holding keys 900-904 is decoded. + assert(prunedRows < total, s"expected pruning to decode fewer than $total rows") + // With pruning off, every cached batch is decoded. + assert( + unprunedRows == total, + s"expected all $total rows to be decoded with pruning disabled, got $unprunedRows") + } + + test("Comet in-memory cache supports DISK_ONLY storage level") { + // CometCachedBatch holds a ChunkedByteBuffer, which is Externalizable, so BlockManager can + // spill it to the DiskStore like any other cached block. Pins that: nothing is held in memory, + // the bytes really do land on disk, every partition is cached, and the cache still reads back + // through the native scan. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + spark + .range(0, 1000, 1, 4) + .selectExpr("id as key", "id % 7 as value") + .createOrReplaceTempView("disk_cache") + + spark.catalog.cacheTable("disk_cache", StorageLevel.DISK_ONLY) + val total = spark.table("disk_cache").count() + assert(total == 1000) + + assert( + cachedBatchTypes("disk_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")), + "DISK_ONLY must still use Comet's cached batch format") + + val cached = + spark.sharedState.cacheManager.lookupCachedData(spark.table("disk_cache")).get + val rddId = cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers.id + val info = spark.sparkContext.getRDDStorageInfo + .find(_.id == rddId) + .getOrElse(fail(s"no storage info for cached RDD $rddId")) + + assert(info.memSize == 0, s"expected nothing in memory, got ${info.memSize} bytes") + assert(info.diskSize > 0, "expected the cached bytes to be on disk") + assert( + info.numCachedPartitions == info.numPartitions, + s"expected all ${info.numPartitions} partitions cached, got ${info.numCachedPartitions}") + + val df = spark.sql("SELECT key, value FROM disk_cache WHERE key >= 900 AND key < 905") + checkSparkAnswer(df) + val plan = df.queryExecution.executedPlan.toString() + assert(plan.contains("CometInMemoryTableScan")) + + spark.catalog.clearCache() + } + } + + test("Comet 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", + "cast(id 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")) + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala new file mode 100644 index 0000000000..f3ceb831ce --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -0,0 +1,152 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.spark.SparkConf +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +object CometInMemoryCacheBenchmark extends CometBenchmarkBase { + private val numRows = 5 * 1000 * 1000 + private val cacheTable = "comet_cache_bench" + private val sourceTable = "comet_cache_bench_src" + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName("CometInMemoryCacheBenchmark") + .set("spark.master", "local[1]") + .setIfMissing("spark.driver.memory", "3g") + .setIfMissing("spark.executor.memory", "3g") + .set("spark.plugins", "org.apache.spark.CometPlugin") + .set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set( + "spark.sql.cache.serializer", + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + + val sparkSession = SparkSession + .builder() + .config(conf) + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + + sparkSession.conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") + sparkSession.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") + sparkSession.conf.set(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, "true") + sparkSession.conf.set(SQLConf.ANSI_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") + sparkSession + } + + override def runCometBenchmark(args: Array[String]): Unit = { + withTempTable(sourceTable, cacheTable) { + spark + .range(0, numRows, 1, 16) + .selectExpr("id", "id % 1000 AS k", "id + 1 AS v") + .createOrReplaceTempView(sourceTable) + + runCacheBenchmark( + "in-memory cache repeated scan", + s"SELECT sum(id), sum(k), sum(v) FROM $cacheTable") + + runCacheBenchmark( + "in-memory cache selective filter", + s""" + |SELECT sum(id), sum(k), sum(v) + |FROM $cacheTable + |WHERE id >= 4500000 AND id < 4750000 + """.stripMargin) + } + } + + private def runCacheBenchmark(name: String, query: String): Unit = { + withCachedTable { + withSQLConf(cacheConf(nativeCacheEnabled = false): _*) { + verifyPlan(query, nativeCacheEnabled = false) + } + withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { + verifyPlan(query, nativeCacheEnabled = true) + } + + val benchmark = new Benchmark(name, numRows, output = output) + + benchmark.addCase("Comet cache disabled") { _ => + withSQLConf(cacheConf(nativeCacheEnabled = false): _*) { + spark.sql(query).noop() + } + } + + benchmark.addCase("Comet cache enabled") { _ => + withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { + spark.sql(query).noop() + } + } + + benchmark.run() + } + } + + private def withCachedTable(f: => Unit): Unit = { + spark.catalog.clearCache() + + // Materialize the cache once using Comet's cache serializer. + // The benchmark measures repeated cache reads by comparing the + // fallback read path against CometInMemoryTableScan. + withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { + spark.sql(s"SELECT id, k, v FROM $sourceTable").createOrReplaceTempView(cacheTable) + spark.catalog.cacheTable(cacheTable) + spark.table(cacheTable).count() + } + + try f + finally { + spark.catalog.uncacheTable(cacheTable) + spark.catalog.clearCache() + } + } + + private def verifyPlan(query: String, nativeCacheEnabled: Boolean): Unit = { + val plan = spark.sql(query).queryExecution.executedPlan.toString() + + if (nativeCacheEnabled) { + assert(plan.contains("CometInMemoryTableScan"), s"Expected native cache scan:\n$plan") + assert(!plan.contains("CometSparkColumnarToColumnar"), s"Unexpected conversion:\n$plan") + } else { + assert( + !plan.contains("CometInMemoryTableScan"), + s"Native cache scan should be disabled:\n$plan") + } + } + + private def cacheConf(nativeCacheEnabled: Boolean): Seq[(String, String)] = { + Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> nativeCacheEnabled.toString, + "spark.comet.sparkToColumnar.enabled" -> "true", + "spark.comet.exec.onHeap.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/execution/columnar/CometInMemoryRelationHelper.scala b/spark/src/test/scala/org/apache/spark/sql/execution/columnar/CometInMemoryRelationHelper.scala new file mode 100644 index 0000000000..ff5240af2a --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/execution/columnar/CometInMemoryRelationHelper.scala @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.execution.columnar + +/** + * Test-only access to `InMemoryRelation`'s JVM-wide cached `CachedBatchSerializer`. + * + * `InMemoryRelation` resolves `spark.sql.cache.serializer` once per JVM and memoizes the instance + * in a static field, so the first suite in a forked JVM that caches a table pins the serializer + * for every suite that follows. A suite that needs a specific cache serializer must reset that + * state around itself. `InMemoryRelation.clearSerializer` is `private[columnar]`, hence this + * shim. + */ +object CometInMemoryRelationHelper { + def clearSerializer(): Unit = InMemoryRelation.clearSerializer() +}