perf: compare decimal bounds in the Arrow data buffer when gathering cache statistics

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

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)
     }
   }