perf: compare string 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, VarCharVector}
 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)
@@ -141,16 +146,10 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
         case d: DecimalType => gatherDecimalStats(col, numRows, d)
         // Every StringType, collated ones included, and compared with the ordering resolved for
         // that column rather than with byte order. See boundsOrderings.
-        case _: StringType => gatherStringStats(col, numRows, orderings(c))
+        case dt: StringType => gatherStringStats(col, numRows, dt == StringType, 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) = {
@@ -313,6 +379,96 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
   }
 
   private def gatherStringStats(
+      col: ColumnVector,
+      numRows: Int,
+      defaultCollation: Boolean,
+      ordering: Ordering[Any]): (UTF8String, UTF8String, Int) = col match {
+    // The default StringType is equal only to itself, so no collated type takes this path.
+    case v: CometPlainVector
+        if defaultCollation && v.getValueVector.isInstanceOf[VarCharVector] =>
+      gatherStringBytesStats(v, numRows)
+    case _ => gatherStringStatsByValue(col, numRows, ordering)
+  }
+
+  // Compares the bytes of the strings in Arrow's data buffer instead of building a UTF8String for
+  // every row. The default collation orders strings by their bytes, unsigned and lexicographic,
+  // which is the order UTF8String.compareTo implements and the one this follows. A row is compared
+  // with the lower bound first, and with the upper one only when it is not below it. Only the two
+  // bounds become UTF8Strings, copied because the vector's buffers are released once the batch is
+  // written.
+  private def gatherStringBytesStats(
+      col: CometPlainVector,
+      numRows: Int): (UTF8String, UTF8String, Int) = {
+    val vector = col.getValueVector.asInstanceOf[VarCharVector]
+    val data = vector.getDataBuffer.memoryAddress()
+    val offsets = vector.getOffsetBuffer.memoryAddress()
+    var minRow = -1
+    var maxRow = -1
+    var minAddress = 0L
+    var minLength = 0
+    var maxAddress = 0L
+    var maxLength = 0
+    var nullCount = 0
+    var r = 0
+    while (r < numRows) {
+      if (col.isNullAt(r)) {
+        nullCount += 1
+      } else {
+        val start = Platform.getInt(null, offsets + r * 4L)
+        val length = Platform.getInt(null, offsets + (r + 1L) * 4L) - start
+        val address = data + start
+        if (minRow < 0) {
+          minRow = r
+          maxRow = r
+          minAddress = address
+          minLength = length
+          maxAddress = address
+          maxLength = length
+        } else if (compareBytes(address, length, minAddress, minLength) < 0) {
+          minRow = r
+          minAddress = address
+          minLength = length
+        } else if (compareBytes(address, length, maxAddress, maxLength) > 0) {
+          maxRow = r
+          maxAddress = address
+          maxLength = length
+        }
+      }
+      r += 1
+    }
+    if (minRow < 0) {
+      (null, null, nullCount)
+    } else {
+      (col.getUTF8String(minRow).copy(), col.getUTF8String(maxRow).copy(), nullCount)
+    }
+  }
+
+  // The comparison of UTF8String.compareTo: eight bytes at a time as big-endian unsigned longs,
+  // then byte by byte, and the shorter of two strings with the same prefix first.
+  private def compareBytes(left: Long, leftLength: Int, right: Long, rightLength: Int): Int = {
+    val length = math.min(leftLength, rightLength)
+    val words = length & ~7
+    var i = 0
+    while (i < words) {
+      val a = Platform.getLong(null, left + i)
+      val b = Platform.getLong(null, right + i)
+      if (a != b) {
+        return JLong.compareUnsigned(JLong.reverseBytes(a), JLong.reverseBytes(b))
+      }
+      i += 8
+    }
+    while (i < length) {
+      val a = Platform.getByte(null, left + i) & 0xff
+      val b = Platform.getByte(null, right + i) & 0xff
+      if (a != b) return a - b
+      i += 1
+    }
+    leftLength - rightLength
+  }
+
+  // The fallback for a collated string, a large-offset vector, or a column that is not a plain
+  // Arrow vector.
+  private def gatherStringStatsByValue(
       col: ColumnVector,
       numRows: Int,
       ordering: Ordering[Any]): (UTF8String, UTF8String, Int) = {
@@ -405,9 +561,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 +569,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 +680,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)
     }
   }