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

.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

.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

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

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