diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 58c04f406a7..f30fedc5cab 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -460,6 +460,7 @@ jobs: org.apache.spark.sql.CometCollationSuite org.apache.comet.CometFuzzAggregateSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite + org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite - name: "expressions" value: | diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 5731b2f5be7..28d768270cf 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -214,6 +214,7 @@ jobs: org.apache.spark.sql.CometCollationSuite org.apache.comet.CometFuzzAggregateSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite + org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite - name: "expressions" value: | diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 5a67d5eef80..e7a8f950dcf 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf._ import org.apache.comet.iceberg.IcebergWriteStrategy -import org.apache.comet.rules.{CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages} +import org.apache.comet.rules.{CometCacheColumnarRule, CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages} import org.apache.comet.shims.ShimCometSparkSessionExtensions /** @@ -54,7 +54,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * CometSubqueryBroadcastExec for exchange reuse with Comet broadcasts * b. insertTransitions: ColumnarToRow/RowToColumnar added * c. postColumnarTransitions: RevertNativeForTransitionHeavyStages, - * EliminateRedundantTransitions + * EliminateRedundantTransitions, CometCacheColumnarRule * 5. ReuseExchangeAndSubquery -- Spark deduplicates subqueries (sees Comet nodes) * }}} * @@ -78,7 +78,7 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * a. preColumnarTransitions: CometScanRule, CometExecRule (no-ops, already converted) * b. insertTransitions * c. postColumnarTransitions: RevertNativeForTransitionHeavyStages, - * EliminateRedundantTransitions + * EliminateRedundantTransitions, CometCacheColumnarRule * }}} * * On Spark 3.4, injectQueryStageOptimizerRule is unavailable. CometExecRule does not wrap SABs, @@ -113,7 +113,10 @@ class CometSparkSessionExtensions override def postColumnarTransitions: Rule[SparkPlan] = { val rules = - Seq(RevertNativeForTransitionHeavyStages(session), EliminateRedundantTransitions(session)) + Seq( + RevertNativeForTransitionHeavyStages(session), + EliminateRedundantTransitions(session), + CometCacheColumnarRule) plan => rules.foldLeft(plan) { case (p, rule) => rule(p) } } } diff --git a/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala new file mode 100644 index 00000000000..7d8ba194aed --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometCacheColumnarRule.scala @@ -0,0 +1,90 @@ +/* + * 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.rules + +import org.apache.spark.sql.catalyst.expressions.LeafExpression +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.{CodegenSupport, ColumnarToRowExec, ColumnarToRowTransition, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + +/** + * Lets Spark's generated consumers read cached Arrow vectors without an intermediate UnsafeRow. + * + * Data flows upward. Spark's InputAdapter/whole-stage wrappers and an optional AQE cache stage + * are omitted: + * {{{ + * Before After + * +------------------------+ +------------------------+ + * | Spark codegen consumer | | Spark codegen consumer | + * +------------------------+ +------------------------+ + * ^ ^ + * | UnsafeRow | column values + * +------------------------+ +------------------------+ + * | InMemoryTableScanExec | | ColumnarToRowExec | + * | row iterator | | fused with consumer | + * +------------------------+ +------------------------+ + * ^ + * | ColumnarBatch + * +------------------------+ + * | InMemoryTableScanExec | + * | Arrow vectors | + * +------------------------+ + * }}} + */ +object CometCacheColumnarRule extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.wholeStageEnabled) return plan + + plan.transformUp { + case parent: CodegenSupport + if parent.supportCodegen && !parent.supportsColumnar && + !parent.isInstanceOf[ColumnarToRowTransition] && + !WholeStageCodegenExec.isTooManyFields(conf, parent.schema) && + !parent.children.exists(p => WholeStageCodegenExec.isTooManyFields(conf, p.schema)) && + !parent.expressions.exists(_.exists { + case _: LeafExpression => false + case _: CodegenFallback => true + case _ => false + }) => + // Match the consuming edge rather than every scan: an existing columnar consumer (or a + // cache stage being materialized by AQE) must keep receiving batches. Spark inserts an + // InputAdapter around the scan later, while this transition fuses with the row consumer. + parent.withNewChildren(parent.children.map { + case child if isColumnarCometCache(child) => ColumnarToRowExec(child) + case child => child + }) + } + } + + private def isColumnarCometCache(plan: SparkPlan): Boolean = { + plan.supportsColumnar && (plan match { + case scan: InMemoryTableScanExec => + // The materialized format is fixed even when Comet execution is later disabled. The + // serializer delegates unsupported schemas to Spark, whose cache keeps its own reader. + scan.relation.cacheBuilder.serializer.isInstanceOf[ArrowCachedBatchSerializer] && + ArrowCachedBatchSerializer.supportsSchema(scan.relation.output) + case stage: QueryStageExec => isColumnarCometCache(stage.plan) + case _ => false + }) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 821f84e0c2b..59013ff9ead 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -21,13 +21,13 @@ package org.apache.spark.sql.comet.execution.arrow import java.lang.{Boolean => JBoolean, Byte => JByte, Double => JDouble, Float => JFloat, Integer => JInteger, Long => JLong, Short => JShort} -import scala.collection.JavaConverters._ import scala.util.control.NonFatal -import org.apache.spark.TaskContext +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.io.CompressionCodec 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.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull} import org.apache.spark.sql.catalyst.util.TypeUtils import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch, SimpleMetricsCachedBatchSerializer} import org.apache.spark.sql.comet.util.Utils @@ -40,6 +40,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.io.ChunkedByteBuffer import org.apache.comet.CometArrowAllocator +import org.apache.comet.vector.CometVector /** * Cached batch format used when Comet writes Spark in-memory cache data. @@ -48,13 +49,16 @@ import org.apache.comet.CometArrowAllocator * 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. + * this class only changes the cached payload. `deltaEncoded` marks numeric streams whose values + * need a prefix sum after decoding; the validity bits and logical schema remain the same as the + * source column. */ private case class CometCachedBatch( override val numRows: Int, override val sizeInBytes: Long, override val stats: InternalRow, - columns: Array[ChunkedByteBuffer]) + columns: Array[ChunkedByteBuffer], + deltaEncoded: Array[Boolean]) extends SimpleMetricsCachedBatch /** @@ -352,7 +356,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val (lower, upper, nulls) = gatherColumnStats(batch, attrs) val numRows = batch.numRows() - val columns = if (Utils.isArrowBacked(batch)) { + val (columns, deltaEncoded) = if (Utils.isArrowBacked(batch)) { Utils.serializeBatchColumns(batch) } else { val arrowBatch = @@ -366,7 +370,8 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { numRows = numRows, sizeInBytes = columnSizes.sum, stats = statsRow(lower, upper, nulls, numRows, columnSizes), - columns = columns) + columns = columns, + deltaEncoded = deltaEncoded) } } @@ -464,6 +469,8 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val indices = selectedIndices(cacheAttributes, selectedAttributes) input.mapPartitions { it => + // Codec factories are reusable; each selected column still gets its own input stream. + lazy val codec = CompressionCodec.createCodec(SparkEnv.get.conf) // 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 @@ -490,7 +497,11 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // 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) + val readers = new ColumnReaders( + indices.map(i => cb.columns(i)), + indices.map(i => cb.deltaEncoded(i)), + cb.numRows, + codec) current = readers readers.batches } @@ -508,7 +519,11 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // decoded vectors stay owned by their readers: closing them releases the batch, which is why // this yields a single-element iterator that closes on exhaustion, matching what // ArrowReaderIterator did when the payload was one stream. - private class ColumnReaders(buffers: Array[ChunkedByteBuffer], numRows: Int) { + private class ColumnReaders( + buffers: Array[ChunkedByteBuffer], + deltaEncoded: Array[Boolean], + numRows: Int, + codec: CompressionCodec) { // 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 @@ -518,7 +533,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { var i = 0 try { while (i < buffers.length) { - opened(i) = Utils.decodeBatches(buffers(i), "CometCache") + opened(i) = Utils.decodeBatches(buffers(i), "CometCache", codec) i += 1 } } catch { @@ -570,7 +585,11 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { throw new IllegalStateException( s"Cached column stream $i decoded ${decoded.numRows()} rows, expected $numRows") } - columns(i) = decoded.column(0) + val column = decoded.column(0) + if (deltaEncoded(i)) { + Utils.decodeDeltaLongs(column.asInstanceOf[CometVector].getValueVector) + } + columns(i) = column i += 1 } new ColumnarBatch(columns, numRows) @@ -646,11 +665,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { convertCachedBatchToColumnarBatch(input, cacheAttributes, selectedAttributes, conf) .mapPartitions { batches => - val toUnsafe = UnsafeProjection.create(selectedAttributes, selectedAttributes) - - batches.flatMap { batch => - batch.rowIterator().asScala.map(row => toUnsafe(row).copy()) - } + new CachedBatchRowIterator(selectedAttributes).createObject(batches) } } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala new file mode 100644 index 00000000000..907236999aa --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIterator.scala @@ -0,0 +1,128 @@ +/* + * 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 org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, CodeGeneratorWithInterpretedFallback, InterpretedUnsafeProjection} +import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.catalyst.expressions.codegen.Block._ +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +/** + * Reads vectors directly into Spark's reusable UnsafeRow buffer. The input iterator owns the + * batches and releases them on advancement or task completion. As with Spark's cache reader, + * callers must copy rows they retain across next(), but the returned row owns its variable-width + * values and remains valid when hasNext() releases the batch that supplied them. + */ +private[arrow] class CachedBatchRowIterator(attributes: Seq[Attribute]) + extends CodeGeneratorWithInterpretedFallback[Iterator[ColumnarBatch], Iterator[InternalRow]] { + + private def fields: Seq[BoundReference] = attributes.zipWithIndex.map { case (attr, i) => + BoundReference(i, attr.dataType, attr.nullable) + } + + override protected def createCodeGeneratedObject( + batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = { + val ctx = new CodegenContext + val columns = attributes.indices.map { i => + ctx.addMutableState(classOf[ColumnVector].getName, s"column$i") + } + ctx.currentVars = attributes.zip(columns).map { case (attr, column) => + val value = JavaCode.variable(ctx.freshName("value"), attr.dataType) + val getter = CodeGenerator.getValueFromVector(column, attr.dataType, "rowId") + val javaType = CodeGenerator.javaType(attr.dataType) + if (attr.nullable) { + val isNull = JavaCode.isNullVariable(ctx.freshName("isNull")) + ExprCode( + code""" + boolean $isNull = $column.isNullAt(rowId); + $javaType $value = $isNull ? ${CodeGenerator.defaultValue(attr.dataType)} : ($getter); + """, + isNull, + value) + } else { + ExprCode(code"$javaType $value = $getter;", FalseLiteral, value) + } + } + val projection = GenerateUnsafeProjection.createCode(ctx, fields) + val bindColumns = columns.zipWithIndex + .map { case (column, i) => + s"$column = batch.column($i);" + } + .mkString("\n") + val code = s""" + public Object generate(Object[] references) { + return new SpecificCachedBatchRowIterator((scala.collection.Iterator) references[0]); + } + + class SpecificCachedBatchRowIterator extends scala.collection.AbstractIterator { + private final scala.collection.Iterator batches; + private int rowId = 0; + private int numRows = 0; + ${ctx.declareMutableStates()} + + public SpecificCachedBatchRowIterator(scala.collection.Iterator batches) { + this.batches = batches; + ${ctx.initMutableStates()} + } + + public boolean hasNext() { + while (rowId >= numRows && batches.hasNext()) { + ${classOf[ColumnarBatch].getName} batch = + (${classOf[ColumnarBatch].getName}) batches.next(); + numRows = batch.numRows(); + rowId = 0; + $bindColumns + } + return rowId < numRows; + } + + public InternalRow next() { + if (!hasNext()) throw new java.util.NoSuchElementException(); + ${projection.code} + rowId++; + return ${projection.value}; + } + + ${ctx.declareAddedFunctions()} + } + """ + val (compiled, _) = + CodeGenerator.compile(new CodeAndComment(code, ctx.getPlaceHolderToComments())) + compiled.generate(Array[Any](batches)).asInstanceOf[Iterator[InternalRow]] + } + + override protected def createInterpretedObject( + batches: Iterator[ColumnarBatch]): Iterator[InternalRow] = { + val toUnsafe = InterpretedUnsafeProjection.createProjection(fields) + batches.flatMap { batch => + new Iterator[InternalRow] { + private var rowId = 0 + override def hasNext: Boolean = rowId < batch.numRows() + override def next(): InternalRow = { + if (!hasNext) throw new NoSuchElementException + val row = toUnsafe(batch.getRow(rowId)) + rowId += 1 + row + } + } + } + } +} 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 d70fdab35e6..bd8a1523ccb 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 @@ -42,6 +42,7 @@ import org.apache.spark.sql.comet.execution.arrow.{ArrowReaderIterator, Constant import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.unsafe.Platform import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStream} import org.apache.comet.Constants.COMET_CONF_DIR_ENV @@ -284,30 +285,89 @@ object Utils extends CometTypeShim with Logging { * [[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. + * and any dictionaries the column needs. Returns the streams and flags identifying long columns + * stored as deltas. Their readers restore the original values after Arrow decoding, before + * exposing the vectors to consumers. * * 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] = { + def serializeBatchColumns(batch: ColumnarBatch): (Array[ChunkedByteBuffer], Array[Boolean]) = { 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)) + getBatchFieldVectorsWithProviders(batch) + .map { case (fieldVector, providerOpt) => + val provider = providerOpt.getOrElse(new CDataDictionaryProvider) + def writeColumn(vector: FieldVector): ChunkedByteBuffer = { + val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) + val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) + val root = new VectorSchemaRoot(Seq(vector).asJava) + val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) + writer.start() + writer.writeBatch() + root.clear() + writer.close() + cbbos.toChunkedByteBuffer + } - val root = new VectorSchemaRoot(Seq(fieldVector).asJava) - val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) - writer.start() - writer.writeBatch() - root.clear() - writer.close() + fieldVector match { + case longs: BigIntVector + if longs.getField.getDictionary == null && longs.getValueCount > 0 => + // Copy before writing: a column may be borrowed from a cache scan, and writeColumn + // clears its input. The validity bits and logical schema stay unchanged. + val deltas = new BigIntVector(longs.getField, longs.getAllocator) + try { + val count = longs.getValueCount + deltas.allocateNew(count) + deltas.setValueCount(count) + deltas.getValidityBuffer.setBytes( + 0, + longs.getValidityBuffer, + 0, + BitVectorHelper.getValidityBufferSize(count)) + val source = longs.getDataBuffer.memoryAddress() + val target = deltas.getDataBuffer.memoryAddress() + var previous = 0L + var smallDeltas = 0 + var i = 0 + while (i < count) { + val value = Platform.getLong(null, source + i * 8L) + val delta = value - previous + Platform.putLong(null, target + i * 8L, delta) + if (delta == delta.toInt.toLong) smallDeltas += 1 + previous = value + i += 1 + } + val plain = writeColumn(longs) + // ponytail: this cheap filter skips full-width random longs; the size comparison + // below still rejects poorly compressing deltas from narrower distributions. + if (smallDeltas.toLong * 2 < count) (plain, false) + else { + val encoded = writeColumn(deltas) + // Require a substantial size reduction to pay for reconstructing each read. + if (encoded.size < plain.size * 3 / 4) (encoded, true) else (plain, false) + } + } finally deltas.close() + case _ => (writeColumn(fieldVector), false) + } + } + .toArray + .unzip + } - cbbos.toChunkedByteBuffer - }.toArray + private[sql] def decodeDeltaLongs(vector: ValueVector): Unit = { + require(vector.isInstanceOf[BigIntVector], "Delta-encoded cache column must contain longs") + val address = vector.getDataBuffer.memoryAddress() + var previous = 0L + var i = 0 + while (i < vector.getValueCount) { + val value = Platform.getLong(null, address + i * 8L) + previous + Platform.putLong(null, address + i * 8L, value) + previous = value + i += 1 + } } /** @@ -337,12 +397,19 @@ object Utils extends CometTypeShim with Logging { * an iterator of ColumnarBatch */ def decodeBatches(bytes: ChunkedByteBuffer, source: String): Iterator[ColumnarBatch] = { + if (bytes.size == 0) Iterator.empty + else decodeBatches(bytes, source, CompressionCodec.createCodec(SparkEnv.get.conf)) + } + + def decodeBatches( + bytes: ChunkedByteBuffer, + source: String, + codec: CompressionCodec): Iterator[ColumnarBatch] = { if (bytes.size == 0) { return Iterator.empty } // use Spark's compression codec (LZ4 by default) and not Comet's compression - val codec = CompressionCodec.createCodec(SparkEnv.get.conf) val cbbis = bytes.toInputStream() val ins = new DataInputStream(codec.compressedInputStream(cbbis)) // batches are in Arrow IPC format diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index e72f4d10f74..52be0409527 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -24,14 +24,14 @@ 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, Row} +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} 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.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.CometCachedBatchHelper -import org.apache.spark.sql.execution.SortExec +import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SortExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} -import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation} +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.max @@ -41,6 +41,7 @@ import org.apache.spark.storage.StorageLevel import org.apache.comet.{CometArrowAllocator, CometConf} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} +import org.apache.comet.rules.CometCacheColumnarRule import org.apache.comet.vector.CometVector class CometInMemoryCacheSuite extends CometTestBase { @@ -295,6 +296,149 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Spark row consumers of Comet cache preserve values across batches") { + for { + mode <- Seq("CODEGEN_ONLY", "NO_CODEGEN") + vectorized <- Seq(false, true) + } { + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> vectorized.toString, + SQLConf.COLUMN_BATCH_SIZE.key -> "7", + SQLConf.CODEGEN_FACTORY_MODE.key -> mode, + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> (mode == "CODEGEN_ONLY").toString, + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val scalars = Seq( + "boolean", + "tinyint", + "smallint", + "int", + "bigint", + "float", + "double", + "decimal(10,2)", + "decimal(38,2)", + "date", + "timestamp", + "timestamp_ntz").zipWithIndex.map { case (dt, i) => + val value = dt match { + case "date" | "timestamp" | "timestamp_ntz" => + s"cast(date_add(DATE '2000-01-01', cast(id AS INT)) AS $dt)" + case _ => s"cast(id AS $dt)" + } + s"if(id % 3 = 0, null, $value) AS c$i" + } + val source = spark + .range(0, 41, 1, 2) + .selectExpr((Seq("id AS key") ++ scalars ++ Seq( + "if(id % 3 = 0, null, repeat(concat('字', id), cast(id + 1 AS INT))) AS s", + "if(id % 3 = 0, null, cast(concat('binary', id) AS BINARY)) AS b", + "if(id % 3 = 0, null, array(cast(id AS STRING), null)) AS a", + "if(id % 3 = 0, null, named_struct('x', id, 'a', array(cast(id AS STRING)))) AS st", + "if(id % 3 = 0, null, map('k', array(cast(id AS STRING), null))) AS m", + "null AS n")): _*) + + def queries(df: DataFrame): Seq[DataFrame] = Seq( + df.select("*"), + df.selectExpr("s AS renamed", "key", "b", "a", "st", "m"), + df.orderBy($"s".desc, $"key"), + df.join(spark.range(41).toDF("join_key"), $"key" === $"join_key").select(df("*")), + df.selectExpr("count(*)"), + df.limit(1)) + + val expected = queries(source).map(_.collect().toSeq) + source.cache() + try { + assert(source.count() == 41) + val relation = + spark.sharedState.cacheManager.lookupCachedData(source).get.cachedRepresentation + val buffers = relation.cacheBuilder.cachedColumnBuffers.collect() + assert(buffers.length > 2) + assert(buffers.forall(_.getClass.getSimpleName == "CometCachedBatch")) + queries(source).zip(expected).foreach { case (df, answer) => + val scans = + df.queryExecution.executedPlan.collect { case scan: InMemoryTableScanExec => + scan + } + assert( + scans.nonEmpty && scans.forall(_.supportsColumnar == vectorized), + df.queryExecution.executedPlan.toString) + if (!vectorized || mode == "NO_CODEGEN") { + assert( + !df.queryExecution.executedPlan.exists(_.isInstanceOf[ColumnarToRowExec]), + df.queryExecution.executedPlan.toString) + } + checkAnswer(df, answer) + } + } finally source.unpersist(blocking = true) + } + } + } + + test("Spark generated consumers read cold and warm Comet caches columnarly") { + for { + adaptive <- Seq(false, true) + cometEnabled <- Seq(false, true) + } { + withSQLConf( + CometConf.COMET_ENABLED.key -> cometEnabled.toString, + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.COLUMN_BATCH_SIZE.key -> "7", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val source = spark + .range(0, 41, 1, 2) + .selectExpr("id AS key", "if(id % 3 = 0, null, concat('字', id)) AS s") + def query = source + .filter("key >= 7") + .selectExpr("sum(key)", "sum(length(s))", "count(*)") + val expected = query.collect().toSeq + source.cache() + try { + val builder = spark.sharedState.cacheManager + .lookupCachedData(source) + .get + .cachedRepresentation + .cacheBuilder + Seq(true, false).foreach { cold => + val df = query + val plan = df.queryExecution.executedPlan + // Planning must not materialize the cache or replace AQE's cache-stage metadata. + assert(builder.isCachedColumnBuffersLoaded != cold, plan.toString) + checkAnswer(df, expected) + assert(builder.isCachedColumnBuffersLoaded) + val transitions = collect(plan) { + case c: ColumnarToRowExec if collect(c.child) { case s: InMemoryTableScanExec => + s + }.nonEmpty => + c + } + assert(transitions.size == 1, plan.toString) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.isEmpty) + if (adaptive && isSpark35Plus) { + assert(collect(plan) { + case s: QueryStageExec + if s.getClass.getSimpleName == "TableCacheQueryStageExec" => + s + }.size == 1) + } + val scan = collect(plan) { case s: InMemoryTableScanExec => s }.head + // A cache scan can also be the root of a columnar request or already have a + // transition. Applying the rule again must preserve those input/output contracts. + Seq(scan, ColumnarToRowExec(scan), RowToColumnarExec(scan)).foreach { boundary => + assert(CometCacheColumnarRule(boundary).fastEquals(boundary)) + } + } + } finally source.unpersist(blocking = true) + } + } + } + test("Comet cache serializer delegates unsupported types to Spark's cache format") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", @@ -1531,6 +1675,64 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache preserves delta-encoded longs across every reader") { + val random = new java.util.Random(5485) + val rows = (0 until 4096).map { i => + // Cross the signed-long boundary; delta reconstruction must preserve wrapping arithmetic. + val value = Long.MaxValue - 2048 + i + Row(value, if (i % 7 == 0) null else value, random.nextLong(), s"value_${i % 11}") + } + val schema = new StructType() + .add("seq", LongType, nullable = false) + .add("nullable", LongType, nullable = true) + .add("random", LongType, nullable = false) + .add("text", StringType, nullable = false) + 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 cached = spark.createDataFrame(spark.sparkContext.parallelize(rows, 2), schema).cache() + try { + assert(cached.count() == rows.length) + val relation = + spark.sharedState.cacheManager.lookupCachedData(cached).get.cachedRepresentation + val batches = relation.cacheBuilder.cachedColumnBuffers.collect() + assert(batches.forall(b => CometCachedBatchHelper.columnsAreDeltaEncoded(b)(0))) + assert(batches.forall(b => !CometCachedBatchHelper.columnsAreDeltaEncoded(b)(2))) + + val expected = + rows.map(r => Row(r.getLong(2), r.getLong(0), r.get(1), r.getLong(0), r.get(3))) + for ((native, vectorized) <- Seq((true, true), (false, true), (false, false))) { + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> native.toString, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> native.toString, + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> vectorized.toString) { + val projected = + cached.selectExpr("random", "seq", "nullable", "seq AS repeated", "text") + checkAnswer(projected, expected) + val plan = projected.queryExecution.executedPlan + if (native) assert(plan.exists(_.isInstanceOf[CometInMemoryTableScanExec])) + else { + assert(plan.exists(_.isInstanceOf[InMemoryTableScanExec])) + assert(plan.exists(_.isInstanceOf[ColumnarToRowExec]) == vectorized) + } + } + } + + // Re-caching decoded vectors must neither encode them twice nor change the original cache. + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false") { + val recached = cached.union(cached).cache() + try checkAnswer(recached, rows ++ rows) + finally recached.unpersist() + checkAnswer(cached, rows) + } + } finally cached.unpersist() + } + } + 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 diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala new file mode 100644 index 00000000000..db0ce00853e --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometCacheRowReaderBenchmark.scala @@ -0,0 +1,207 @@ +/* + * 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 java.nio.charset.StandardCharsets + +import org.apache.spark.benchmark.BenchmarkBase +import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.ColumnarToRowExec +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.storage.StorageLevel + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +/** + * Compare Spark consumers of Comet and Spark caches (issue #5485). + * + * Arguments: [spark|comet|comet-row|all] [rows] [iterations] [all|mixed|numeric] [codec]. Run one + * format per JVM in alternating order on main and the patch. comet-row disables vectorized cache + * reading to isolate the row iterator. Cache creation and validation are outside timing. + */ +object CometCacheRowReaderBenchmark extends BenchmarkBase { + private val warmups = 5 + + override def runBenchmarkSuite(args: Array[String]): Unit = { + require(args.length <= 5, "Expected format, rows, iterations, schema, compression codec") + val format = args.headOption.getOrElse("all") + val rows = args.lift(1).map(_.toLong).getOrElse(5000000L) + val iterations = args.lift(2).map(_.toInt).getOrElse(15) + val schema = args.lift(3).getOrElse("all") + val codec = args.lift(4).getOrElse("lz4") + require(Set("all", "spark", "comet", "comet-row").contains(format)) + require(Set("all", "mixed", "numeric").contains(schema)) + require(rows > 0 && iterations > 0) + + emit("CACHE_SAMPLE,format,schema,query,rows,iteration,elapsed_ns") + val formats = + if (format == "all") Seq("spark", "comet", "comet-row") else Seq(format) + val schemas = if (schema == "all") Seq("mixed", "numeric") else Seq(schema) + formats.foreach { name => + CometInMemoryRelationHelper.clearSerializer() + SparkSession.clearActiveSession() + SparkSession.clearDefaultSession() + val serializer = if (name == "spark") { + classOf[DefaultCachedBatchSerializer].getName + } else { + classOf[ArrowCachedBatchSerializer].getName + } + val spark = SparkSession + .builder() + .master("local[1]") + .appName(getClass.getSimpleName) + .config("spark.ui.enabled", "false") + .config("spark.sql.cache.serializer", serializer) + .config("spark.sql.shuffle.partitions", "1") + .config("spark.sql.inMemoryColumnarStorage.batchSize", "10000") + .config("spark.sql.inMemoryColumnarStorage.compressed", "true") + .config("spark.io.compression.codec", codec) + .config(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") + .config(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") + .config(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key, (name != "comet-row").toString) + .config(SQLConf.CODEGEN_FACTORY_MODE.key, "CODEGEN_ONLY") + .config(CometConf.COMET_ENABLED.key, "false") + .config(CometConf.COMET_EXEC_ENABLED.key, "false") + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + spark.sparkContext.setLogLevel("WARN") + try { + emit( + s"CACHE_ENV,$name,Spark=${spark.version}," + + s"Java=${System.getProperty("java.version")},codec=$codec") + schemas.foreach(runSchema(spark, name, _, rows, iterations, serializer)) + } finally { + spark.stop() + SparkSession.clearActiveSession() + SparkSession.clearDefaultSession() + CometInMemoryRelationHelper.clearSerializer() + } + } + } + + private def runSchema( + spark: SparkSession, + format: String, + schema: String, + rows: Long, + iterations: Int, + serializer: String): Unit = { + val mixed = schema == "mixed" + val first = Seq("id", "id % 1000 AS k", "id + 1 AS v") + val rest = if (mixed) { + Seq( + "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") + } else { + Seq("id % 100000 AS n1", "id % 7919 AS n2", "id * 3 AS n3") + } + val source = spark.range(0, rows, 1, 16).selectExpr((first ++ rest): _*) + val columns = source.columns.toSeq + val three = if (mixed) Seq("id", "s1", "s2") else columns.take(3) + val projections = Seq("count" -> Seq.empty[String], "long" -> Seq("id")) ++ + (if (mixed) Seq("string" -> Seq("s1")) else Seq.empty) ++ + Seq("three" -> three, "all" -> columns) + def expressions(selected: Seq[String]): Seq[String] = { + if (selected.isEmpty) Seq("count(*)") + else + selected.map { name => + if (name.startsWith("s")) s"sum(length($name))" else s"sum($name)" + } + } + // Obtain the expected values before the relation is cached, using Spark's ordinary row plan. + val expected = projections.map { case (_, selected) => + source.selectExpr(expressions(selected): _*).collect() + } + val cached = source.persist(StorageLevel.MEMORY_ONLY) + try { + val buildStart = System.nanoTime() + assert(cached.count() == rows) + emit(s"CACHE_BUILD,$format,$schema,$rows,${System.nanoTime() - buildStart}") + val relation = cached.queryExecution.withCachedData.collectFirst { + case relation: InMemoryRelation => relation + }.get + val builder = relation.cacheBuilder + assert(builder.serializer.getClass.getName == serializer) + val batches = builder.cachedColumnBuffers + val batchSummary = batches + .map { batch => + // Spark's sizeInBytes comes from statistics; measure its encoded column buffers. + val bytes = batch match { + case b: DefaultCachedBatch => b.buffers.map(_.length.toLong).sum + case _ => batch.sizeInBytes + } + (batch.getClass.getSimpleName, batch.numRows.toLong, bytes) + } + .collect() + val expectedClass = if (format == "spark") "DefaultCachedBatch" else "CometCachedBatch" + assert(batchSummary.forall(_._1 == expectedClass), "Wrong cached payload format") + assert(batchSummary.map(_._2).sum == rows) + val storage = spark.sparkContext.getRDDStorageInfo.find(_.id == batches.id).get + assert(storage.numCachedPartitions == batches.getNumPartitions && storage.diskSize == 0) + emit( + s"CACHE_STORAGE,$format,$schema,${batchSummary.length}," + + s"${batchSummary.map(_._3).sum},${storage.memSize}") + + projections.zip(expected).foreach { case ((name, selected), answer) => + val query = cached.selectExpr(expressions(selected): _*) + val plan = query.queryExecution.executedPlan + val scans = plan.collect { case scan: InMemoryTableScanExec => scan } + assert(scans.size == 1, s"Expected one Spark cache scan:\n$plan") + val scan = scans.head + assert(scan.attributes.map(_.name).toSet == selected.toSet, s"Wrong projection:\n$plan") + val columnar = plan.exists(_.isInstanceOf[ColumnarToRowExec]) + if (format != "comet") { + assert(!columnar, s"Expected the cache row reader:\n$plan") + } + assert(!plan.exists(_.getClass.getName.startsWith("org.apache.spark.sql.comet."))) + val reader = if (columnar) "columnar" else "row" + emit(s"CACHE_PLAN,$format,$schema,$name,columns=${selected.size},reader=$reader\n$plan") + runQuery(query, answer, format, schema, name, rows, iterations) + } + } finally cached.unpersist(blocking = true) + } + + private def runQuery( + query: DataFrame, + expected: Array[Row], + format: String, + schema: String, + name: String, + rows: Long, + iterations: Int): Unit = { + (0 until warmups).foreach { _ => assert(query.collect().sameElements(expected)) } + (0 until iterations).foreach { i => + val start = System.nanoTime() + val actual = query.collect() + val elapsed = System.nanoTime() - start + assert(actual.sameElements(expected), s"Wrong result for $format/$schema/$name") + emit(s"CACHE_SAMPLE,$format,$schema,$name,$rows,$i,$elapsed") + } + emit(s"CACHE_RESULT,$format,$schema,$name,${expected.mkString(";")}") + } + + private def emit(line: String): Unit = { + println(line) + output.foreach(_.write((line + "\n").getBytes(StandardCharsets.UTF_8))) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala new file mode 100644 index 00000000000..1bea5c5e3e1 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchRowIteratorSuite.scala @@ -0,0 +1,137 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import java.nio.charset.StandardCharsets.UTF_8 + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.VarCharVector +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, UnsafeRow} +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StringType} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.comet.vector.CometPlainVector + +class CachedBatchRowIteratorSuite extends AnyFunSuite { + Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { mode => + def withMode(f: => Unit): Unit = { + val conf = new SQLConf + conf.setConfString(SQLConf.CODEGEN_FACTORY_MODE.key, mode) + SQLConf.withExistingConf(conf)(f) + } + + test(s"$mode: rows own Arrow values across batch release and reuse the output buffer") { + withMode { + val allocator = new RootAllocator(Long.MaxValue) + val vectors = Seq(Seq("first", null), Seq("字" * 1000, "last")).map { values => + val vector = new VarCharVector("s", allocator) + values.zipWithIndex.foreach { case (value, i) => + if (value == null) vector.setNull(i) else vector.setSafe(i, value.getBytes(UTF_8)) + } + vector.setValueCount(values.size) + vector + } + try { + // Match the cache decoder: hasNext releases a consumed batch before the next is read. + val batches = vectors.iterator.flatMap { vector => + new Iterator[ColumnarBatch] { + private var emitted = false + override def hasNext: Boolean = { + if (emitted) vector.close() + !emitted + } + override def next(): ColumnarBatch = { + emitted = true + new ColumnarBatch(Array(new CometPlainVector(vector, false)), 2) + } + } + } + val attributes = Seq(AttributeReference("s", StringType, nullable = true)()) + val rows = new CachedBatchRowIterator(attributes).createObject(batches) + assert(rows.hasNext && rows.hasNext) + val first = rows.next().asInstanceOf[UnsafeRow] + val saved = first.copy() + assert(first.getUTF8String(0).toString == "first") + assert(rows.next() eq first) + assert(first.isNullAt(0)) + assert(rows.hasNext && rows.hasNext) + assert(first.isNullAt(0)) + assert(rows.next().getUTF8String(0).toString == "字" * 1000) + val last = rows.next() + assert(!rows.hasNext && !rows.hasNext) + assert(allocator.getAllocatedMemory == 0) + assert(last.getUTF8String(0).toString == "last") + assert(saved.getUTF8String(0).toString == "first") + intercept[NoSuchElementException](rows.next()) + } finally { + vectors.foreach(_.close()) + allocator.close() + } + } + } + + test(s"$mode: empty input, empty batches, and zero-column rows") { + withMode { + val factory = new CachedBatchRowIterator(Seq.empty) + val empty = factory.createObject(Iterator.empty) + assert(!empty.hasNext) + intercept[NoSuchElementException](empty.next()) + val batches = Seq(0, 2, 0, 3, 0).map { n => + new ColumnarBatch(Array.empty[ColumnVector], n) + } + val rows = factory.createObject(batches.iterator) + assert(rows.map { row => + assert(row.isInstanceOf[UnsafeRow] && row.numFields == 0) + 1 + }.sum == 5) + intercept[NoSuchElementException](rows.next()) + } + } + + test(s"$mode: wide projections preserve nullable and required columns") { + withMode { + val attributes = (0 until 150).map { i => + AttributeReference(s"c$i", IntegerType, nullable = i % 2 == 0)() + } + val columns = attributes.indices.map { i => + val column = new OnHeapColumnVector(2, IntegerType) + column.putInt(0, i) + if (i % 2 == 0) column.putNull(1) else column.putInt(1, -i) + column + } + val batch = new ColumnarBatch(columns.toArray[ColumnVector], 2) + try { + val rows = new CachedBatchRowIterator(attributes).createObject(Iterator.single(batch)) + val first = rows.next().copy() + val second = rows.next() + attributes.indices.foreach { i => + assert(first.getInt(i) == i) + if (i % 2 == 0) assert(second.isNullAt(i)) else assert(second.getInt(i) == -i) + } + assert(!rows.hasNext) + } finally batch.close() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala index 5548f6dadbd..24dfedf92c7 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala @@ -48,6 +48,9 @@ object CometCachedBatchHelper { def columnStreamSizes(batch: CachedBatch): Seq[Long] = batch.asInstanceOf[CometCachedBatch].columns.map(_.size).toSeq + def columnsAreDeltaEncoded(batch: CachedBatch): Seq[Boolean] = + batch.asInstanceOf[CometCachedBatch].deltaEncoded.toSeq + /** * Replace one column's stream with bytes that cannot be decoded, in place. *