diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index 9a15ee39c3a..e4988df5d09 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 2771dc8d2cb..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.NativeUtil +import org.apache.comet.vector.{CometDecodedVector, CometPlainVector, NativeUtil} /** * Cached batch format used when Comet writes Spark in-memory cache data. @@ -114,9 +116,9 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { if (tracksBounds(attr.dataType)) TypeUtils.getInterpretedOrdering(attr.dataType) else null }.toArray - // Bounds and null counts per column, gathered before the batch is serialized: serializing - // clears the batch's vectors, and the per-column byte sizes that complete the statistics row - // are only known afterwards. See statsRow. + // Bounds and null counts per column of the Arrow batch encodeBatches writes, gathered before + // it is serialized: serializing clears the batch's vectors, and the per-column byte sizes that + // complete the statistics row are only known afterwards. See statsRow. private[sql] def gatherColumnStats( batch: ColumnarBatch, attrs: Seq[Attribute], @@ -130,7 +132,10 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { var c = 0 while (c < numCols) { val col = batch.column(c) + val knownNulls = arrowNullCount(col) val (min, max, nullCount) = attrs(c).dataType match { + // Every row is null, so there are no bounds and nothing to read. + case _ if knownNulls == numRows => (null, null, numRows) case BooleanType => gatherBooleanStats(col, numRows) case ByteType => gatherByteStats(col, numRows) case ShortType => gatherShortStats(col, numRows) @@ -144,13 +149,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { case _: StringType => gatherStringStats(col, numRows, orderings(c)) case other => assert(!tracksBounds(other), s"Missing cache bounds implementation for $other") - var nullCount = 0 - var r = 0 - while (r < numRows) { - if (col.isNullAt(r)) nullCount += 1 - r += 1 - } - (null, null, nullCount) + (null, null, knownNulls) } if (nullCount < numRows) { lower(c) = min @@ -163,6 +162,15 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { (lower, upper, nulls) } + // Every column of a batch encodeBatches writes is a CometDecodedVector, which takes its null + // count from Arrow once, when it is built, and answers isNullAt from the same validity bitmap. + private def arrowNullCount(col: ColumnVector): Int = col match { + case v: CometDecodedVector => v.numNulls() + case other => + throw new IllegalStateException( + s"Cache statistics read from ${other.getClass.getName}, not a Comet Arrow vector") + } + // Keep each loop specialized and box only its final bounds. r == nullCount identifies // the first non-null value, including when a column starts with nulls. private def gatherBooleanStats(col: ColumnVector, numRows: Int): (Boolean, Boolean, Int) = { @@ -292,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) = { @@ -405,9 +471,6 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val orderings = boundsOrderings(attrs) batches.map { batch => - // Bounds and null counts are read from the input batch before it is serialized, and the row - // is only assembled once serialize has measured the per-column sizes. - val (lower, upper, nulls) = gatherColumnStats(batch, attrs, orderings) val numRows = batch.numRows() // Written as it stands only if its vectors are ones the writer accepts and are already laid @@ -416,20 +479,31 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // rather than merely usual. val writeDirectly = Utils.isArrowBacked(batch) && CachedBatchIpc.matchesReaderLayout(batch, readerFields) - - val (bytes, columnSizes) = if (writeDirectly) { - CachedBatchIpc.serialize(batch, codec, CometArrowAllocator, settings.chunkSize) - } else { - val arrowBatch = + val arrowBatch = + if (writeDirectly) { + batch + } else { CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) - try CachedBatchIpc.serialize(arrowBatch, codec, CometArrowAllocator, settings.chunkSize) - finally arrowBatch.close() - } + } - CometCachedBatch( - numRows = numRows, - stats = statsRow(lower, upper, nulls, numRows, columnSizes), - bytes = bytes) + try { + // Bounds and null counts are read from the batch that is written, so every null count is + // the one Arrow holds. Spark's own numNulls can be wrong: its vectorized Parquet reader + // marks a column missing from a file all-null without counting any nulls. The row is + // only assembled once serialize has measured the per-column sizes. + val (lower, upper, nulls) = gatherColumnStats(arrowBatch, attrs, orderings) + val (bytes, columnSizes) = + CachedBatchIpc.serialize(arrowBatch, codec, CometArrowAllocator, settings.chunkSize) + + CometCachedBatch( + numRows = numRows, + stats = statsRow(lower, upper, nulls, numRows, columnSizes), + bytes = bytes) + } finally { + if (!writeDirectly) { + arrowBatch.close() + } + } } } @@ -516,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..17bbbd53b1e --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala @@ -0,0 +1,216 @@ +/* + * 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. Each + * batch is copied before the input is advanced, 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 + + // 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 + held = null + } + } + + private def advance(): Unit = { + while (filled.isEmpty && held == null && input.hasNext) { + val batch = input.next() + if (passesThrough(batch)) { + finishPending() + held = batch + } else { + append(batch) + } + } + 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 every row of `batch`, filling batches of targetRows rows. + private def append(batch: ColumnarBatch): Unit = { + val numRows = batch.numRows() + var start = 0 + while (start < numRows) { + if (pending == null) { + startPending() + } + val rows = math.min(numRows - start, targetRows - pendingRows) + pendingWriter.writeColumns(batch, start, rows) + pendingRows += rows + start += rows + 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/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala index 2526d6b6511..81f8e56918a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -132,6 +132,11 @@ private[comet] object CachedBatchIpc { * answers for the top-level vector only, so a struct of large strings passes it while its child * is stored with 64-bit offsets and read with 32-bit ones. * + * A vector must also hold exactly the batch's rows. The cache records the null counts of the + * vectors it writes, so a vector longer than the batch it sits in would count nulls past the + * batch's last row, and could be recorded all null and pruned by `IS NOT NULL` without being + * decoded. Converting copies exactly the batch's rows. + * * Names, nullability and a timestamp's timezone are not compared. None of them changes how the * reader interprets the body, and the writer's legitimately differ -- a Comet scan labels * timestamps with the session's zone where the reader rebuilds them as UTC, which is a label @@ -141,7 +146,8 @@ private[comet] object CachedBatchIpc { batch.numCols() == readerFields.length && (0 until batch.numCols()).forall { i => batch.column(i) match { - case v: CometVector => sameLayout(writtenField(v), readerFields(i)) + case v: CometVector => + v.numValues() == batch.numRows() && sameLayout(writtenField(v), readerFields(i)) case _ => false } } 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 b0edda42338..e7370ccdcfe 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -40,11 +40,11 @@ import org.apache.spark.sql.catalyst.expressions.{And, Ascending, Attribute, Att import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{LeftOuter, RightOuter} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} -import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoin, CometHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoin, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo} +import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FileSourceScanExec, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} @@ -1147,6 +1147,117 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + // Each column's null count summed over the cached batches. Asserts first that the relation is + // stored in Comet's format, because Spark's DefaultCachedBatch records the same null counts and + // no bounds for binary, nested and void columns, and then that no batch records bounds for a + // column. + private def nullCountsWithoutBounds(view: String): Seq[Long] = { + val batchTypes = cachedBatchTypes(view) + assert( + batchTypes.sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")), + s"$view was cached as ${batchTypes.mkString(", ")}") + val relation = spark.sharedState.cacheManager.lookupCachedData(spark.table(view)).get + val stats = relation.cachedRepresentation.cacheBuilder.cachedColumnBuffers + .collect() + .map(_.asInstanceOf[SimpleMetricsCachedBatch].stats) + relation.cachedRepresentation.output.indices.map { i => + stats.foreach(s => assert(s.isNullAt(i * 5) && s.isNullAt(i * 5 + 1), s"column $i")) + stats.map(_.getInt(i * 5 + 2).toLong).sum + } + } + + test("Comet in-memory cache statistics count nulls in columns that record no bounds") { + // These columns record a null count and no bounds: binary and nested types, and a bigint + // column with no values. The count is the one the written Arrow vectors hold, so it is checked + // from the row write path and Comet's native scan, which produce those vectors, and from + // Spark's vectorized reader, whose batches are converted to them first. + val df = spark + .range(1000) + .selectExpr("if(id % 7 = 0, NULL, id) AS v") + .selectExpr( + "cast(concat('b', v) AS binary) AS bin", + "if(v IS NULL, NULL, array(v, v + 1)) AS arr", + "if(v IS NULL, NULL, map(concat('k', v), v)) AS mp", + "if(v IS NULL, NULL, named_struct('a', v, 'b', cast(v AS string))) AS st", + "cast(NULL AS bigint) AS no_values") + // 143 of the ids 0 to 999 are multiples of 7. + val expected = Seq(143L, 143L, 143L, 143L, 1000L) + + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + // Only the row write path can take a void column; Parquet cannot store one. + val rows = df.withColumn("nt", lit(null)) + rows.createOrReplaceTempView("no_bounds_rows") + rows.cache() + try { + rows.count() + assert(nullCountsWithoutBounds("no_bounds_rows") == expected :+ 1000L) + } finally { + rows.unpersist(blocking = true) + spark.catalog.dropTempView("no_bounds_rows") + } + } + + withTempPath { path => + df.write.parquet(path.toString) + withNativeCache { + spark.read.parquet(path.toString).createOrReplaceTempView("no_bounds_native") + spark.catalog.cacheTable("no_bounds_native") + spark.table("no_bounds_native").count() + val cachedPlan = spark.sharedState.cacheManager + .lookupCachedData(spark.table("no_bounds_native")) + .get + .cachedRepresentation + .cacheBuilder + .cachedPlan + assert(cachedPlan.isInstanceOf[CometNativeScanExec], cachedPlan) + assert(nullCountsWithoutBounds("no_bounds_native") == expected) + } + } + + withSparkColumnarCache( + "no_bounds_columnar", + SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true") { path => + df.write.parquet(path) + } { _ => + assert(nullCountsWithoutBounds("no_bounds_columnar") == expected) + } + } + + test("Comet in-memory cache statistics count the nulls of a column missing from a file") { + // Spark's vectorized Parquet reader marks a column that a file lacks all-null without counting + // any nulls, so that vector's numNulls is 0 while every row isNullAt. A binary column records + // that count and no bounds, so taking it from that vector would record no nulls, and the + // IsNull filter below would prune every batch read from the file that lacks the column. + withSparkColumnarCache( + "missing_column_columnar", + SQLConf.PARQUET_SCHEMA_MERGING_ENABLED.key -> "true") { path => + spark.range(0, 400).selectExpr("id").write.parquet(path) + spark + .range(400, 1000) + .selectExpr("id", "cast(cast(id AS string) AS binary) AS b") + .write + .mode("append") + .parquet(path) + } { _ => + val relation = spark.sharedState.cacheManager + .lookupCachedData(spark.table("missing_column_columnar")) + .get + .cachedRepresentation + assert(relation.cacheBuilder.cachedPlan.isInstanceOf[FileSourceScanExec]) + checkAnswer( + spark.sql("SELECT count(*) FROM missing_column_columnar WHERE b IS NULL"), + Row(400L)) + val b = relation.output.indexWhere(_.name == "b") + val stats = relation.cacheBuilder.cachedColumnBuffers + .collect() + .map(_.asInstanceOf[SimpleMetricsCachedBatch].stats) + assert(stats.map(_.getInt(b * 5 + 2)).sum == 400) + } + } + test("Comet in-memory cache supports stats-based batch pruning") { withConversions( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", @@ -2038,6 +2149,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 @@ -2588,6 +2818,31 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache converts a batch whose vectors hold more than its rows") { + // The statistics are gathered from the batch that is written, and a vector's null count covers + // all of its values. Here the batch reports 2 rows over vectors of 4 whose last two values are + // null, so a count taken from the vectors would be 2, which equals the batch's row count: the + // batch would be recorded all null and pruned by IsNotNull without being decoded. Converting + // copies only the batch's rows. + val cacheSchema = StructType(Seq(StructField("b", BinaryType))) + val binaries = new VarBinaryVector("b", CometArrowAllocator) + try { + binaries.allocateNew(4) + Seq(0, 1).foreach(r => binaries.set(r, Array[Byte](r.toByte))) + Seq(2, 3).foreach(r => binaries.setNull(r)) + binaries.setValueCount(4) + val longer = new ColumnarBatch(Array[ColumnVector](new CometPlainVector(binaries)), 2) + assert(Utils.isArrowBacked(longer)) + assert( + !CometCachedBatchHelper.writesDirectly(longer, cacheSchema), + "a vector longer than its batch must be converted, not written as it stands") + val exact = new ColumnarBatch(Array[ColumnVector](new CometPlainVector(binaries)), 4) + assert(CometCachedBatchHelper.writesDirectly(exact, cacheSchema)) + } finally { + binaries.close() + } + } + // Ordering by the JSON form rather than by the columns themselves, since a projection that // excludes `id` has no orderable key of its own and ORDER BY over a map is not allowed. private def orderedByJson(cols: Seq[String], from: String): String = 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 a7b8412eb72..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 @@ -27,19 +27,23 @@ import scala.collection.mutable import org.apache.spark.SparkConf import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, GenericInternalRow} +import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.comet.CometInMemoryTableScanExec -import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} +import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometArrowConverters, CometArrowStream, CometCachedBatchHelper} +import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.ColumnarToRowExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec 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.{DataType, LongType, StringType} +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 -import org.apache.comet.{CometConf, CometSparkSessionExtensions} +import org.apache.comet.{CometArrowAllocator, CometConf, CometSparkSessionExtensions} object CometInMemoryCacheBenchmark extends CometBenchmarkBase { private val numRows = 5 * 1000 * 1000 @@ -786,8 +790,28 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { .collectFirst { case r: InMemoryRelation => r } .getOrElse(sys.error(s"$view is not cached")) + // Each case measures this collector across commits; Spark's default cache has its own collector. private def runStatsBenchmark(): Unit = { val batchSize = 10000 + val serializer = new ArrowCachedBatchSerializer + + def addStatsCase( + benchmark: Benchmark, + name: String, + batch: ColumnarBatch, + attrs: Seq[Attribute]): Unit = { + // Resolved outside the timed loop because that is where the serializer resolves it: once + // per partition, not once per batch. + val orderings = serializer.boundsOrderings(attrs) + benchmark.addCase(name) { _ => + var i = 0 + while (i < numRows / batchSize) { + statsResult = serializer.gatherColumnStats(batch, attrs, orderings) + i += 1 + } + } + } + val types: Seq[DataType] = Seq.fill(3)(LongType) ++ Seq.fill(3)(StringType) val attrs = types.zipWithIndex.map { case (dt, i) => AttributeReference(s"c$i", dt)() } val columns = types.map(dt => new OnHeapColumnVector(batchSize, dt)) @@ -803,22 +827,91 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { columns(5).putByteArray(r, s"str_c_$r".getBytes(StandardCharsets.UTF_8)) r += 1 } - val serializer = new ArrowCachedBatchSerializer - // Resolved outside the timed loop because that is where the serializer resolves it: once - // per partition, not once per batch. - val orderings = serializer.boundsOrderings(attrs) - val benchmark = - new Benchmark("in-memory cache statistics", numRows.toLong, output = output) - // One case measures this collector across commits; Spark's default cache has its own collector. - benchmark.addCase("Comet statistics collector") { _ => - var i = 0 - while (i < numRows / batchSize) { - statsResult = serializer.gatherColumnStats(batch, attrs, orderings) - i += 1 - } + // The same values in Comet's Arrow vectors, which are what the collector reads: the cache + // writes the batches of Comet operators and of the row write path as they are, and those + // of Spark's vectorized readers after converting them to Arrow. + val arrowBatch = CometArrowConverters.columnarBatchToArrowBatch( + batch, + Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE), + CometArrowAllocator) + try { + val benchmark = + new Benchmark("in-memory cache statistics", numRows.toLong, output = output) + addStatsCase(benchmark, "Arrow vectors", arrowBatch, attrs) + benchmark.run() + } finally arrowBatch.close() + } finally batch.close() + + // Columns that record only a null count, with one row in ten null. + val nestedTypes: Seq[DataType] = Seq( + new StructType().add("a", LongType).add("b", StringType), + ArrayType(IntegerType), + BinaryType) + val nestedAttrs = nestedTypes.zipWithIndex.map { case (dt, i) => + AttributeReference(s"n$i", dt)() + } + val rows = Iterator.tabulate(batchSize) { r => + if (r % 10 == 0) { + new GenericInternalRow(nestedTypes.length) + } else { + InternalRow( + InternalRow(r.toLong, UTF8String.fromString(s"str_$r")), + new GenericArrayData(Array[Any](r, r + 1)), + s"bin_$r".getBytes(StandardCharsets.UTF_8)) } + } + val nestedBatch = CometArrowConverters + .rowToArrowBatchIter( + rows, + Utils.fromAttributes(nestedAttrs), + batchSize, + CometArrowStream.NATIVE_TIMEZONE, + CometArrowAllocator) + .next() + try { + val benchmark = new Benchmark( + "in-memory cache statistics, nested and binary columns", + numRows.toLong, + output = output) + addStatsCase(benchmark, "Arrow vectors", nestedBatch, nestedAttrs) benchmark.run() - } finally batch.close() + } 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) + } +}