diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 52e86f59ffd..55be2b93806 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -631,6 +631,8 @@ jobs: org.apache.comet.CometFuzzAggregateSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite + org.apache.spark.sql.comet.execution.arrow.CacheBatchCoalescerSuite + org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchStatsSuite org.apache.spark.sql.comet.execution.arrow.CometArrowWriterSuite org.apache.spark.sql.comet.execution.arrow.CometStringWriterSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index b1d89ec1240..29e0c34dd9b 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -297,6 +297,8 @@ jobs: org.apache.comet.CometFuzzAggregateSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite org.apache.spark.sql.comet.execution.arrow.CachedBatchRowIteratorSuite + org.apache.spark.sql.comet.execution.arrow.CacheBatchCoalescerSuite + org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchStatsSuite org.apache.spark.sql.comet.execution.arrow.CometArrowWriterSuite org.apache.spark.sql.comet.execution.arrow.CometStringWriterSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index c1bbb479c42..76cc565c8e1 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -77,6 +77,14 @@ metrics only through that scan. Each cached batch is stored as a single Arrow IPC record batch message and its body. +A cache built from rows cuts its batches at `spark.sql.inMemoryColumnarStorage.batchSize` rows, +10000 by default, as Spark's own format does. A cache built from the Arrow batches of Comet +operators stores a batch of at least half that size as it stands, without copying it, and first +combines smaller ones up to that size. A shuffle read, for example, returns at least one batch per +map output block, so a relation cached from a shuffle over many map tasks would otherwise hold +thousands of small batches, each with its own metadata and statistics. Batches of Spark's own +column vectors have to be copied into Arrow anyway, so they are combined whatever their size. + The message carries **no Arrow schema**. The reader already has one: `InMemoryRelation` knows the cached relation's attributes, and Comet maps them to exactly the Arrow fields the writer produced. Storing a schema in every batch would repeat the same bytes once per cached batch — for a wide 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 d34e07582ce..624215908ac 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 @@ -23,6 +23,7 @@ import java.lang.{Boolean => JBoolean, Byte => JByte, Double => JDouble, Float = import scala.collection.JavaConverters._ +import org.apache.arrow.vector.DecimalVector import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow @@ -35,11 +36,12 @@ 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.Platform import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.io.ChunkedByteBuffer import org.apache.comet.{CometArrowAllocator, CometConf, DataTypeSupport} -import org.apache.comet.vector.{CometDecodedVector, NativeUtil} +import org.apache.comet.vector.{CometDecodedVector, CometPlainVector, NativeUtil} /** * Cached batch format used when Comet writes Spark in-memory cache data. @@ -298,6 +300,64 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } private def gatherDecimalStats( + col: ColumnVector, + numRows: Int, + dt: DecimalType): (Decimal, Decimal, Int) = col match { + case v: CometPlainVector if v.getValueVector.isInstanceOf[DecimalVector] => + gatherDecimal128Stats(v, numRows, dt) + case _ => gatherDecimalStatsByValue(col, numRows, dt) + } + + // Compares the 128-bit unscaled values in Arrow's data buffer rather than building a Decimal for + // every row, which for a precision past 18 digits means a BigInteger and a BigDecimal each. All + // the values of a column share one scale, so their unscaled order is their numeric order. A + // value is a signed high word over an unsigned low word, in little-endian order, and a narrow + // decimal is the sign extension of its low word, so one comparison covers every precision. Only + // the two bounds become Decimals, through the same getDecimal the by-value loop calls. + private def gatherDecimal128Stats( + col: CometPlainVector, + numRows: Int, + dt: DecimalType): (Decimal, Decimal, Int) = { + val base = col.getValueVector.getDataBuffer.memoryAddress() + var minRow = -1 + var maxRow = -1 + var minHi = 0L + var minLo = 0L + var maxHi = 0L + var maxLo = 0L + var nullCount = 0 + var r = 0 + while (r < numRows) { + if (col.isNullAt(r)) { + nullCount += 1 + } else { + val lo = Platform.getLong(null, base + r * 16L) + val hi = Platform.getLong(null, base + r * 16L + 8L) + if (minRow < 0 || hi < minHi || (hi == minHi && JLong.compareUnsigned(lo, minLo) < 0)) { + minRow = r + minHi = hi + minLo = lo + } + if (maxRow < 0 || hi > maxHi || (hi == maxHi && JLong.compareUnsigned(lo, maxLo) > 0)) { + maxRow = r + maxHi = hi + maxLo = lo + } + } + r += 1 + } + if (minRow < 0) { + (null, null, nullCount) + } else { + ( + col.getDecimal(minRow, dt.precision, dt.scale), + col.getDecimal(maxRow, dt.precision, dt.scale), + nullCount) + } + } + + // The fallback for a decimal column that is not a plain Arrow vector. + private def gatherDecimalStatsByValue( col: ColumnVector, numRows: Int, dt: DecimalType): (Decimal, Decimal, Int) = { @@ -530,16 +590,23 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // Columnar Comet output is stored as one Arrow IPC record batch message per cached batch. Spark // only calls this when supportsColumnarInput returned true, so the schema is known to be // Comet-writable here. + // + // Small input batches are combined first, up to spark.sql.inMemoryColumnarStorage.batchSize + // rows, the size the row path cuts at; see CacheBatchCoalescer. override def convertColumnarBatchToCachedBatch( input: RDD[ColumnarBatch], schema: Seq[Attribute], storageLevel: StorageLevel, conf: SQLConf): RDD[CachedBatch] = { + val batchSize = conf.columnBatchSize val settings = writeSettings(conf) input.mapPartitions { batches => - encodeBatches(batches, schema, settings) + encodeBatches( + new CacheBatchCoalescer(batches, schema, batchSize, CometArrowAllocator), + schema, + settings) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala new file mode 100644 index 00000000000..584d65e257f --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala @@ -0,0 +1,227 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet.execution.arrow + +import scala.collection.JavaConverters._ +import scala.collection.mutable +import scala.util.control.NonFatal + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.spark.TaskContext +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.vector.NativeUtil + +/** + * The columnar input of a cache build, with its small batches combined. + * + * The columnar write stores each batch it is handed as one cached batch, and nothing upstream + * bounds how small a batch can be. A shuffle read returns at least one batch per map output + * block, so a relation cached from a shuffle over many map tasks would otherwise hold thousands + * of batches of a few hundred rows. Each carries its own message metadata, compression framing + * and statistics row, and every read of the cache pays for the batch count. + * + * A batch the writer can store as it stands, Arrow-backed and laid out the way the reader reads + * it, is passed through when it holds at least half of `targetRows`, so full native batches are + * still written without a copy. Every other batch is copied into batches of `targetRows` rows, + * the size the row path cuts at, and split where it crosses that size. A small batch is worth the + * copy, and a batch the writer cannot store as it stands has to be copied into Arrow anyway, so + * combining it costs nothing more. The copy reads each column through the `ColumnVector` + * interface, as `ArrowWriter.writeColumns` does for any input, so it handles dictionary-encoded + * and nested columns, and offsets that do not start at zero, as a native producer's sliced arrays + * can have. Arrow's `VectorAppender` assumes they start at zero. + * + * Row order is kept: the rows copied so far are emitted before a batch that passes through. A + * large input is copied one output batch at a time, as the consumer asks for them, so the batches + * open at once do not grow with its size. The input is not advanced until it has been copied in + * full, because a native producer releases the batch it returned last when it is asked for the + * next one. + * + * The batches built here are owned here. Each is released when the next batch is requested, and + * whatever is still open when the task completes is released then. A batch passed through still + * belongs to its producer. + */ +private[arrow] class CacheBatchCoalescer( + input: Iterator[ColumnarBatch], + attrs: Seq[Attribute], + targetRows: Int, + allocator: BufferAllocator) + extends Iterator[ColumnarBatch] { + + require(targetRows > 0, s"The target batch size must be positive, but was $targetRows") + + private val arrowSchema = + Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE) + private val readerFields = arrowSchema.getFields.asScala.toIndexedSeq + + // A batch at least this large that the writer can store as it stands is not copied. + private val passThroughRows = math.max(1, targetRows / 2) + + // The batch being filled. + private var pending: VectorSchemaRoot = null + private var pendingWriter: ArrowWriter = null + private var pendingRows = 0 + + // The input batch being copied, and the first of its rows that has not been copied yet. + private var current: ColumnarBatch = null + private var currentStart = 0 + + // Batches built here and filled, in input order, ahead of `held`. + private val filled = mutable.Queue.empty[VectorSchemaRoot] + + // A batch that passes through, waiting for the rows copied ahead of it to be emitted. + private var held: ColumnarBatch = null + + // The batch the next call to `next` returns, and the root behind it when it was built here. + private var ready: ColumnarBatch = null + private var readyRoot: VectorSchemaRoot = null + + // The root behind the batch `next` returned last, released once the consumer asks for more. + private var emitted: VectorSchemaRoot = null + + private var closed = false + + Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ => close())) + + override def hasNext: Boolean = { + if (ready == null && !closed) { + release(emitted) + emitted = null + advance() + } + ready != null + } + + override def next(): ColumnarBatch = { + if (!hasNext) { + throw new NoSuchElementException("No more cache input batches") + } + val batch = ready + emitted = readyRoot + ready = null + readyRoot = null + batch + } + + /** Releases every batch built here that is still open. Idempotent. */ + def close(): Unit = { + if (!closed) { + closed = true + release(emitted) + release(readyRoot) + release(pending) + filled.foreach(release) + filled.clear() + emitted = null + ready = null + readyRoot = null + pending = null + pendingWriter = null + current = null + held = null + } + } + + private def advance(): Unit = { + while (filled.isEmpty && held == null && (current != null || input.hasNext)) { + if (current != null) { + copyRows() + } else { + val batch = input.next() + if (passesThrough(batch)) { + finishPending() + held = batch + } else if (batch.numRows() > 0) { + current = batch + currentStart = 0 + } + } + } + if (filled.isEmpty && held == null) { + // The input is exhausted. + finishPending() + } + if (filled.nonEmpty) { + readyRoot = filled.dequeue() + ready = NativeUtil.rootAsBatch(readyRoot) + } else if (held != null) { + ready = held + held = null + } + } + + private def passesThrough(batch: ColumnarBatch): Boolean = + batch.numRows() >= passThroughRows && + Utils.isArrowBacked(batch) && + CachedBatchIpc.matchesReaderLayout(batch, readerFields) + + // Copies rows of `current` until the batch being filled is full or `current` is exhausted. + private def copyRows(): Unit = { + if (pending == null) { + startPending() + } + val rows = math.min(current.numRows() - currentStart, targetRows - pendingRows) + pendingWriter.writeColumns(current, currentStart, rows) + pendingRows += rows + currentStart += rows + if (currentStart == current.numRows()) { + current = null + } + if (pendingRows == targetRows) { + finishPending() + } + } + + private def startPending(): Unit = { + val root = VectorSchemaRoot.create(arrowSchema, allocator) + try { + pendingWriter = ArrowWriter.create(root, targetRows) + } catch { + case NonFatal(e) => + root.close() + throw e + } + pending = root + pendingRows = 0 + } + + private def finishPending(): Unit = { + if (pending != null) { + if (pendingRows > 0) { + pendingWriter.finish() + filled.enqueue(pending) + } else { + pending.close() + } + pending = null + pendingWriter = null + pendingRows = 0 + } + } + + private def release(root: VectorSchemaRoot): Unit = { + if (root != null) { + root.close() + } + } +} 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 2e1f4d56d33..e25e8b63d43 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -2248,6 +2248,125 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache combines small columnar batches in order") { + // At a batch size of 1000, an Arrow batch under 500 rows is combined with its neighbours, a + // larger one is stored as it stands after the rows combined ahead of it, and a batch of + // Spark's own vectors, which has to be copied into Arrow anyway, is combined at any size. + // Combined batches are cut at 1000 rows. The input releases each batch when it is asked for + // the next one, as a native producer does, so a batch held rather than copied would be read + // after it was released (#6950). + val input = Seq( + 300 -> true, + 0 -> true, + 300 -> true, + 9000 -> true, + 700 -> true, + 400 -> true, + 2000 -> false, + 400 -> true, + 300 -> true) + val attrs = Seq(AttributeReference("id", LongType)(), AttributeReference("s", StringType)()) + val serializer = new ArrowCachedBatchSerializer + withSQLConf(SQLConf.COLUMN_BATCH_SIZE.key -> "1000") { + val cached = serializer.convertColumnarBatchToCachedBatch( + spark.sparkContext.parallelize(Seq(0), 1).mapPartitions(_ => new ReleasingBatches(input)), + attrs, + StorageLevel.MEMORY_ONLY, + SQLConf.get) + assert(cached.map(_.numRows).collect().toSeq == Seq(600, 9000, 700, 1000, 1000, 1000, 100)) + + val rows = serializer + .convertCachedBatchToColumnarBatch(cached, attrs, attrs, SQLConf.get) + .mapPartitions(_.flatMap { batch => + (0 until batch.numRows()).map(i => + (batch.column(0).getLong(i), batch.column(1).getUTF8String(i).toString)) + }) + .collect() + assert(rows.toSeq == (0L until input.map(_._1).sum.toLong).map(i => (i, s"v$i"))) + } + } + + test("Comet in-memory cache combines the small batches of a shuffle read") { + // A shuffle read returns at least one batch per map output block, so a relation cached + // straight from a shuffle over 20 map tasks into 2 partitions arrives in 40 batches of about + // 500 rows. They are combined up to the default batch size of 10000 rows (#6950). + withTempPath { path => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0, 20000, 1, 20) + .selectExpr( + "id", + "cast(id % 1000 as int) AS k", + "if(id % 11 = 0, NULL, concat('s', id)) AS s", + "named_struct('a', id, 'b', cast(id AS string)) AS st", + "if(id % 13 = 0, NULL, array(id, id + 1)) AS arr", + "map(cast(id % 7 AS string), id) AS mp") + .write + .parquet(path.toString) + } + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + // One map task per file. + SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1024") { + val query = s"SELECT /*+ REPARTITION(2, k) */ * FROM parquet.`${path.toString}`" + val expected = uncachedSparkAnswer(query) + val df = spark.sql(query) + assert(df.rdd.getNumPartitions == 2) + df.cache() + try { + assert(df.count() == 20000) + val relation = spark.sharedState.cacheManager.lookupCachedData(df).get + val cached = relation.cachedRepresentation.cacheBuilder.cachedColumnBuffers + assert( + cached.map(_.getClass.getName).distinct().collect().toSeq == + Seq("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")) + val sizes = cached.map(_.numRows).collect() + assert(sizes.length <= 4 && sizes.forall(_ <= 10000), sizes.mkString(", ")) + assert(QueryTest.sameRows(expected.toSeq, df.collect().toSeq).isEmpty) + } finally { + df.unpersist(blocking = true) + } + } + } + } + + test("Comet in-memory cache keeps a sorted cache sorted when it combines its batches") { + // A cache reports the order of the plan it was built from, so the rows of the batches it + // combines have to keep their order. The native sort emits batches of 100 rows here. + withTempPath { path => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 20000, 1, 4).selectExpr("id", "id % 97 AS k").write.parquet(path.toString) + } + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_BATCH_SIZE.key -> "100") { + val df = spark.read + .parquet(path.toString) + .repartition(2, $"k") + .sortWithinPartitions($"id".desc) + df.cache() + try { + assert(df.count() == 20000) + val relation = spark.sharedState.cacheManager.lookupCachedData(df).get + val sizes = + relation.cachedRepresentation.cacheBuilder.cachedColumnBuffers + .map(_.numRows) + .collect() + assert(sizes.length <= 4, sizes.mkString(", ")) + val sorted = df.rdd + .mapPartitions { rows => + val ids = rows.map(_.getLong(0)).toArray + Iterator(ids.toSeq == ids.sorted.reverse.toSeq) + } + .collect() + assert(sorted.toSeq == Seq(true, true)) + } finally { + df.unpersist(blocking = true) + } + } + } + } + // Enough rows that every column's buffers are big enough for Arrow to actually compress them. // Arrow stores a buffer verbatim when compressing it would not make it smaller, and a boolean // column of a few hundred rows is a few dozen bytes, which takes that fallback -- leaving the diff --git a/spark/src/test/scala/org/apache/comet/exec/ReleasingBatches.scala b/spark/src/test/scala/org/apache/comet/exec/ReleasingBatches.scala new file mode 100644 index 00000000000..82095bed451 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/ReleasingBatches.scala @@ -0,0 +1,89 @@ +/* + * 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.nio.charset.StandardCharsets + +import org.apache.arrow.vector.{BigIntVector, VarCharVector} +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.types.{LongType, StringType} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.comet.CometArrowAllocator +import org.apache.comet.vector.CometPlainVector + +/** + * Batches of `id`, counting from 0, and `s`, "v" followed by the id, one batch per element of + * `sizes`: a row count, and whether the batch holds Arrow vectors or Spark's own on-heap ones. + * Like a native producer, it releases each batch when it is asked for the next one. + */ +private class ReleasingBatches(sizes: Seq[(Int, Boolean)]) extends Iterator[ColumnarBatch] { + private var remaining = sizes.toList + private var nextId = 0L + private var previous: ColumnarBatch = null + + override def hasNext: Boolean = { + if (previous != null) { + previous.close() + previous = null + } + remaining.nonEmpty + } + + override def next(): ColumnarBatch = { + if (!hasNext) { + throw new NoSuchElementException("No more batches") + } + val (rows, arrow) = remaining.head + remaining = remaining.tail + val batch = if (arrow) arrowBatch(rows) else sparkBatch(rows) + nextId += rows + previous = batch + batch + } + + private def value(i: Int): Array[Byte] = s"v${nextId + i}".getBytes(StandardCharsets.UTF_8) + + private def arrowBatch(rows: Int): ColumnarBatch = { + val ids = new BigIntVector("id", CometArrowAllocator) + val strings = new VarCharVector("s", CometArrowAllocator) + ids.allocateNew(rows) + strings.allocateNew(rows.toLong * 8, rows) + (0 until rows).foreach { i => + ids.set(i, nextId + i) + strings.setSafe(i, value(i)) + } + ids.setValueCount(rows) + strings.setValueCount(rows) + new ColumnarBatch( + Array[ColumnVector](new CometPlainVector(ids), new CometPlainVector(strings)), + rows) + } + + private def sparkBatch(rows: Int): ColumnarBatch = { + val ids = new OnHeapColumnVector(rows, LongType) + val strings = new OnHeapColumnVector(rows, StringType) + (0 until rows).foreach { i => + ids.putLong(i, nextId + i) + strings.putByteArray(i, value(i)) + } + new ColumnarBatch(Array[ColumnVector](ids, strings), rows) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala index 187572c853f..b4aa3ee7aeb 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.execution.aggregate.HashAggregateExec import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} -import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, IntegerType, LongType, StringType, StructType} +import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, Decimal, DecimalType, IntegerType, LongType, StringType, StructType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.unsafe.types.UTF8String @@ -876,6 +876,42 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { addStatsCase(benchmark, "Arrow vectors", nestedBatch, nestedAttrs) benchmark.run() } finally nestedBatch.close() + + // Decimals of one narrow and two wide precisions, one row in ten null and the values on both + // sides of zero. A decimal past 18 digits is the one the by-value loop paid a BigInteger and a + // BigDecimal for on every row. + val decimalTypes: Seq[DataType] = + Seq(DecimalType(10, 2), DecimalType(18, 2), DecimalType(38, 2)) + val decimalAttrs = decimalTypes.zipWithIndex.map { case (dt, i) => + AttributeReference(s"d$i", dt)() + } + val decimalRows = Iterator.tabulate(batchSize) { r => + if (r % 10 == 0) { + new GenericInternalRow(decimalTypes.length) + } else { + val unscaled = BigInt(r - batchSize / 2) * 7919 + InternalRow( + Decimal(unscaled.toLong, 10, 2), + Decimal(unscaled.toLong * 100000, 18, 2), + Decimal(BigDecimal(unscaled * BigInt(10).pow(20) + r, 2), 38, 2)) + } + } + val decimalBatch = CometArrowConverters + .rowToArrowBatchIter( + decimalRows, + Utils.fromAttributes(decimalAttrs), + batchSize, + CometArrowStream.NATIVE_TIMEZONE, + CometArrowAllocator) + .next() + try { + val benchmark = new Benchmark( + "in-memory cache statistics, decimal columns", + numRows.toLong, + output = output) + addStatsCase(benchmark, "Arrow vectors", decimalBatch, decimalAttrs) + benchmark.run() + } finally decimalBatch.close() } private def runCacheBenchmark( diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchStatsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchStatsSuite.scala new file mode 100644 index 00000000000..55eaacfc223 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchStatsSuite.scala @@ -0,0 +1,159 @@ +/* + * 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.math.{BigDecimal => JBigDecimal, BigInteger} + +import scala.util.Random + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.arrow.c.CDataDictionaryProvider +import org.apache.arrow.vector.DecimalVector +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.types.{Decimal, DecimalType} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.comet.CometArrowAllocator +import org.apache.comet.vector.CometVector + +/** The bounds and null counts `ArrowCachedBatchSerializer` records for decimal columns. */ +class ArrowCachedBatchStatsSuite extends AnyFunSuite { + + private val serializer = new ArrowCachedBatchSerializer + + // Values that sit on either side of the boundaries between the two 64-bit words of a 128-bit + // unscaled value: the low word's sign bit, a carry into the high word, and both signs of each. + private def boundaryValues(precision: Int): Seq[BigInteger] = { + val two63 = BigInteger.ONE.shiftLeft(63) + val two64 = BigInteger.ONE.shiftLeft(64) + val limit = BigInteger.TEN.pow(precision).subtract(BigInteger.ONE) + val candidates = Seq( + BigInteger.ZERO, + BigInteger.ONE, + two63.subtract(BigInteger.ONE), + two63, + two63.add(BigInteger.ONE), + two64.subtract(BigInteger.ONE), + two64, + two64.add(BigInteger.ONE), + two64.multiply(BigInteger.valueOf(3)).add(two63)) + val signed = candidates ++ candidates.map(_.negate()) + (signed :+ limit :+ limit.negate()).filter(_.abs.compareTo(limit) <= 0).distinct + } + + private def randomValues(precision: Int, count: Int, seed: Long): Seq[BigInteger] = { + val random = new Random(seed) + val limit = BigInteger.TEN.pow(precision) + Seq.fill(count) { + // A random width, so the values spread over the words instead of all filling the high one. + val bits = 1 + random.nextInt(limit.bitLength()) + val magnitude = new BigInteger(bits, new java.util.Random(random.nextLong())).mod(limit) + if (random.nextBoolean()) magnitude else magnitude.negate() + } + } + + // Gathers the statistics of one decimal column holding `values`, with a null before and after + // every `nullEvery`-th value when `nullEvery` is positive. + private def check(precision: Int, scale: Int, values: Seq[BigInteger], nullEvery: Int): Unit = { + val rows: Seq[Option[BigInteger]] = + if (nullEvery > 0) { + values.zipWithIndex.flatMap { case (v, i) => + if (i % nullEvery == 0) Seq(None, Some(v), None) else Seq(Some(v)) + } + } else { + values.map(Some(_)) + } + val allocator = CometArrowAllocator.newChildAllocator("decimal-stats", 0, Long.MaxValue) + val vector = new DecimalVector("d", allocator, precision, scale) + try { + vector.allocateNew(rows.length) + rows.zipWithIndex.foreach { + case (None, i) => vector.setNull(i) + case (Some(v), i) => vector.setSafe(i, new JBigDecimal(v, scale)) + } + vector.setValueCount(rows.length) + + val column = CometVector.getVector(vector, new CDataDictionaryProvider) + val batch = new ColumnarBatch(Array[ColumnVector](column), rows.length) + val attrs = Seq(AttributeReference("d", DecimalType(precision, scale))()) + val (lower, upper, nulls) = + serializer.gatherColumnStats(batch, attrs, serializer.boundsOrderings(attrs)) + + val present = rows.flatten + assert(nulls(0) == rows.length - present.length) + if (present.isEmpty) { + assert(lower(0) == null && upper(0) == null) + } else { + val expectedMin = new JBigDecimal(present.min, scale) + val expectedMax = new JBigDecimal(present.max, scale) + val actualMin = lower(0).asInstanceOf[Decimal] + val actualMax = upper(0).asInstanceOf[Decimal] + val context = s"decimal($precision, $scale), ${rows.length} rows" + assert(actualMin.toJavaBigDecimal.compareTo(expectedMin) == 0, s"min of $context") + assert(actualMax.toJavaBigDecimal.compareTo(expectedMax) == 0, s"max of $context") + assert( + actualMin.precision == precision && actualMin.scale == scale, + s"min type, $context") + assert( + actualMax.precision == precision && actualMax.scale == scale, + s"max type, $context") + } + } finally { + vector.close() + allocator.close() + } + } + + // Both a narrow precision, whose unscaled value fits the low word, and a wide one, for which the + // high word decides the order and the low word is compared as unsigned. + private val precisions = + Seq((5, 2), (10, 2), (18, 0), (18, 6), (19, 0), (20, 4), (30, 10), (38, 2)) + + precisions.foreach { case (precision, scale) => + test(s"decimal($precision, $scale) bounds across the word boundaries") { + val values = boundaryValues(precision) + check(precision, scale, values, nullEvery = 0) + // Every order of the same values must give the same bounds. + check(precision, scale, values.reverse, nullEvery = 0) + check(precision, scale, new Random(7).shuffle(values), nullEvery = 3) + + // Without the largest magnitudes, which decide the bounds by their high word alone, the + // values are all near the word boundaries. Of one sign, the high words are equal for many of + // them, so the low word decides, and it is unsigned. + val inner = values.filter(_.abs.bitLength <= 66) + check(precision, scale, inner.filter(_.signum >= 0), nullEvery = 0) + check(precision, scale, inner.filter(_.signum <= 0), nullEvery = 0) + check(precision, scale, inner, nullEvery = 2) + } + + test(s"decimal($precision, $scale) bounds of random values and nulls") { + (0 until 5).foreach { seed => + check(precision, scale, randomValues(precision, 500, seed.toLong), nullEvery = 0) + check(precision, scale, randomValues(precision, 500, seed.toLong), nullEvery = 7) + } + } + } + + test("decimal bounds of a column with a single value") { + check(38, 2, Seq(BigInteger.ONE.shiftLeft(100).negate()), nullEvery = 0) + check(10, 2, Seq(BigInteger.valueOf(-5)), nullEvery = 1) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescerSuite.scala new file mode 100644 index 00000000000..cf351fd42ce --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescerSuite.scala @@ -0,0 +1,116 @@ +/* + * 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.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector +import org.apache.spark.sql.types.StringType +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +class CacheBatchCoalescerSuite extends AnyFunSuite { + + private val numColumns = 4 + private val attrs = (0 until numColumns).map(i => AttributeReference(s"s$i", StringType)()) + + // A batch of Spark's own on-heap vectors, which the coalescer always copies. Row `i` of column + // `c` holds "cr", so the output order can be checked. + private def sparkBatch(first: Int, rows: Int): ColumnarBatch = { + val columns = (0 until numColumns).map { c => + val vector = new OnHeapColumnVector(rows, StringType) + (0 until rows).foreach(i => vector.putByteArray(i, s"c${c}r${first + i}".getBytes(UTF_8))) + vector: ColumnVector + } + new ColumnarBatch(columns.toArray, rows) + } + + // Runs `f` on a coalescer over `inputs`, checks that closing it returns every byte it took, and + // returns the most the allocator held at once. + private def withCoalescer(limit: Long, inputs: Seq[ColumnarBatch], target: Int)( + f: CacheBatchCoalescer => Unit): Long = { + val allocator = new RootAllocator(limit) + val coalescer = new CacheBatchCoalescer(inputs.iterator, attrs, target, allocator) + try { + f(coalescer) + } finally { + coalescer.close() + } + assert(allocator.getAllocatedMemory == 0L) + val peak = allocator.getPeakMemoryAllocation + allocator.close() + peak + } + + test("splitting one large input copies one output batch at a time") { + val inputRows = 4096 + val target = 10 + // The limit is far above what one output batch needs and far below what 410 of them need. + val limit = 4L * 1024 * 1024 + val peak = withCoalescer(limit, Seq(sparkBatch(0, inputRows)), target) { coalescer => + var rows = 0 + var outputs = 0 + while (coalescer.hasNext) { + val out = coalescer.next() + assert(out.numRows() <= target) + (0 until out.numRows()).foreach { i => + (0 until numColumns).foreach { c => + assert(out.column(c).getUTF8String(i).toString == s"c${c}r${rows + i}") + } + } + rows += out.numRows() + outputs += 1 + } + assert(rows == inputRows) + assert(outputs == (inputRows + target - 1) / target) + } + // Only the batch being emitted is open, so the allocation does not grow with the number of + // outputs. Eagerly splitting the input would hold all of them at once. + assert(peak < 1024 * 1024, peak) + } + + test("rows left over from one input are filled up from the next") { + val inputs = Seq(sparkBatch(0, 25), sparkBatch(25, 7), sparkBatch(32, 30)) + withCoalescer(Long.MaxValue, inputs, 10) { coalescer => + val sizes = Seq.newBuilder[Int] + var rows = 0 + while (coalescer.hasNext) { + val out = coalescer.next() + (0 until out.numRows()).foreach { i => + assert(out.column(0).getUTF8String(i).toString == s"c0r${rows + i}") + } + rows += out.numRows() + sizes += out.numRows() + } + assert(sizes.result() == Seq(10, 10, 10, 10, 10, 10, 2)) + } + } + + test("an empty input batch is skipped") { + val inputs = Seq(sparkBatch(0, 0), sparkBatch(0, 3), sparkBatch(3, 0)) + withCoalescer(Long.MaxValue, inputs, 10) { coalescer => + assert(coalescer.next().numRows() == 3) + assert(!coalescer.hasNext) + } + } +}