From 320b384a8f3a0bddbd187250df65c1cbbe8f4605 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Sat, 10 Oct 2026 07:35:31 -0600
Subject: [PATCH 1/9] perf: read cache statistics null counts from Comet's
 Arrow vectors

gatherColumnStats visited every row to count nulls, including for binary,
array, map and struct columns, which record a null count and no bounds.
A CometDecodedVector already holds that count, taken from Arrow when the
vector is built. Use it when the vector holds exactly the batch's rows: a
column without bounds records it without a pass over its rows, and an
all-null column of any type skips its bounds loop.

Spark's own vectors are still counted row by row. Its vectorized Parquet
reader marks a column missing from a file all-null without counting any
nulls, so numNulls is 0 there while every row isNullAt.

The statistics benchmark now measures Comet's Arrow vectors as well as
Spark's, and binary and nested columns.

Part of #6834.
---
 .../arrow/ArrowCachedBatchSerializer.scala    |  36 ++++-
 .../comet/exec/CometInMemoryCacheSuite.scala  | 141 +++++++++++++++++-
 .../CometInMemoryCacheBenchmark.scala         |  94 +++++++++---
 .../arrow/CometCachedBatchHelper.scala        |  13 ++
 4 files changed, 256 insertions(+), 28 deletions(-)

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..7aec0ad8128 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
@@ -39,7 +39,7 @@ 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, NativeUtil}
 
 /**
  * Cached batch format used when Comet writes Spark in-memory cache data.
@@ -130,7 +130,10 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
     var c = 0
     while (c < numCols) {
       val col = batch.column(c)
+      val knownNulls = arrowNullCount(col, numRows)
       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 +147,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, if (knownNulls >= 0) knownNulls else countNulls(col, numRows))
       }
       if (nullCount < numRows) {
         lower(c) = min
@@ -163,6 +160,29 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
     (lower, upper, nulls)
   }
 
+  // The column's null count if its Arrow vector already holds one, or -1 if only a pass over the
+  // rows can tell.
+  //
+  // A CometDecodedVector takes its null count from Arrow once, when it is built, and its isNullAt
+  // answers from the same validity bitmap, so the two agree whenever the vector holds exactly the
+  // batch's rows. Spark's own vectors do not qualify: its vectorized Parquet reader marks a column
+  // missing from the file all-null, and that vector's numNulls is then 0 while isNullAt is true
+  // for every row.
+  private def arrowNullCount(col: ColumnVector, numRows: Int): Int = col match {
+    case v: CometDecodedVector if v.numValues() == numRows => v.numNulls()
+    case _ => -1
+  }
+
+  private def countNulls(col: ColumnVector, numRows: Int): Int = {
+    var nullCount = 0
+    var r = 0
+    while (r < numRows) {
+      if (col.isNullAt(r)) nullCount += 1
+      r += 1
+    }
+    nullCount
+  }
+
   // 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) = {
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 aba97d80219..161fc830b13 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -37,11 +37,11 @@ import org.apache.spark.SparkConf
 import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row}
 import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal}
 import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch}
-import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec}
+import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, 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, SparkPlanInfo}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FileSourceScanExec, FilterExec, FormattedMode, RowToColumnarExec, SortExec, 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}
@@ -1141,6 +1141,143 @@ class CometInMemoryCacheSuite extends CometTestBase {
     }
   }
 
+  // Each column's null count summed over the cached batches, asserting along the way that no
+  // batch records bounds for the column.
+  private def nullCountsWithoutBounds(view: String): Seq[Long] = {
+    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. Comet's Arrow vectors give the count without a pass over the rows,
+    // while Spark's vectors are counted row by row, so it is checked from the row write path,
+    // Comet's native scan and Spark's vectorized reader.
+    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 numNulls 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 statistics count nulls only within a batch's rows") {
+    // A batch can report fewer rows than its vectors hold, and the nulls here are all past the
+    // batch's rows. A count taken from the whole vector would say every row of this batch is null,
+    // and Spark would then prune the batch for IsNotNull.
+    val ints = new IntVector("i", CometArrowAllocator)
+    val binaries = new VarBinaryVector("b", CometArrowAllocator)
+    try {
+      ints.allocateNew(4)
+      binaries.allocateNew(4)
+      Seq(0, 1).foreach { r =>
+        ints.set(r, r + 1)
+        binaries.set(r, Array[Byte](r.toByte))
+      }
+      Seq(2, 3).foreach { r =>
+        ints.setNull(r)
+        binaries.setNull(r)
+      }
+      ints.setValueCount(4)
+      binaries.setValueCount(4)
+      val batch = new ColumnarBatch(
+        Array[ColumnVector](new CometPlainVector(ints), new CometPlainVector(binaries)),
+        2)
+      val attrs =
+        Seq(AttributeReference("i", IntegerType)(), AttributeReference("b", BinaryType)())
+      val (lower, upper, nulls) = CometCachedBatchHelper.columnStats(batch, attrs)
+      assert(nulls.toSeq == Seq(0, 0))
+      assert(lower(0) == 1 && upper(0) == 2)
+    } finally {
+      ints.close()
+      binaries.close()
+    }
+  }
+
   test("Comet in-memory cache supports stats-based batch pruning") {
     withConversions(
       SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
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 230126a7025..87f74c7f7f1 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
@@ -24,19 +24,23 @@ import java.nio.charset.StandardCharsets
 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, 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
@@ -747,8 +751,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))
@@ -764,22 +788,56 @@ 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. Spark's vectors are what the collector reads
+      // from Spark's vectorized readers, and Arrow vectors are what it reads from Comet operators
+      // and from the row write path.
+      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, "Spark vectors", batch, attrs)
+        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()
   }
 
   private def runCacheBenchmark(
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala
index 47e876b5f86..ba4b616b57b 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala
@@ -33,6 +33,7 @@ import org.apache.arrow.vector.ipc.{ReadChannel, WriteChannel}
 import org.apache.arrow.vector.ipc.message.{ArrowBodyCompression, ArrowRecordBatch, MessageMetadataResult, MessageSerializer}
 import org.apache.arrow.vector.types.pojo.Field
 import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.columnar.CachedBatch
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.types.StructType
@@ -129,6 +130,18 @@ object CometCachedBatchHelper {
   def cachedBatch(payload: ChunkedByteBuffer, numRows: Int): CachedBatch =
     CometCachedBatch(numRows, InternalRow.empty, payload)
 
+  /**
+   * The lower bounds, upper bounds and null counts the writer records for `batch`, one entry per
+   * column. Another thin shim, for tests that hand the statistics collector a batch no plan
+   * produces.
+   */
+  def columnStats(
+      batch: ColumnarBatch,
+      attrs: Seq[Attribute]): (Array[Any], Array[Any], Array[Int]) = {
+    val serializer = new ArrowCachedBatchSerializer
+    serializer.gatherColumnStats(batch, attrs, serializer.boundsOrderings(attrs))
+  }
+
   /**
    * Decode the `selected` columns of a cached batch the way a scan does, into a root the caller
    * owns. Another thin shim, for tests about what the read path leaves allocated.

From 1b3990d111a166574fdbd9524437c768fbe99a47 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Tue, 13 Oct 2026 17:03:32 -0600
Subject: [PATCH 2/9] test: assert the no-bounds null count test caches in
 Comet's format

Spark's DefaultCachedBatch records the same null counts and no bounds
for binary, nested and void columns. nullCountsWithoutBounds now checks
that the relation is stored as CometCachedBatch before reading its
statistics, so the row write path and native scan legs check the format
as the Spark vectorized reader leg already did.
---
 .../apache/comet/exec/CometInMemoryCacheSuite.scala   | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

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 161fc830b13..161620103b7 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -1141,9 +1141,16 @@ class CometInMemoryCacheSuite extends CometTestBase {
     }
   }
 
-  // Each column's null count summed over the cached batches, asserting along the way that no
-  // batch records bounds for the column.
+  // 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()

From b979e1b07e709043c4746d573dbd3ecb5a28103d Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 14 Oct 2026 08:16:23 -0600
Subject: [PATCH 3/9] perf: gather cache statistics from the Arrow batch that
 is written

encodeBatches now converts a batch that is not written directly before it
gathers the statistics, so every column the collector reads is a
CometDecodedVector holding the null count Arrow took when it was built. The
row-by-row null count for Spark's vectors and the check that a vector holds
exactly the batch's rows go, and so do the test of that check and its helper.
The statistics benchmark measures only Arrow vectors, which is all the
collector reads now.
---
 .../arrow/ArrowCachedBatchSerializer.scala    | 76 +++++++++----------
 .../comet/exec/CometInMemoryCacheSuite.scala  | 43 ++---------
 .../CometInMemoryCacheBenchmark.scala         |  7 +-
 .../arrow/CometCachedBatchHelper.scala        | 13 ----
 4 files changed, 43 insertions(+), 96 deletions(-)

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 7aec0ad8128..d34e07582ce 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
@@ -114,9 +114,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 +130,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
     var c = 0
     while (c < numCols) {
       val col = batch.column(c)
-      val knownNulls = arrowNullCount(col, numRows)
+      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)
@@ -147,7 +147,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
         case _: StringType => gatherStringStats(col, numRows, orderings(c))
         case other =>
           assert(!tracksBounds(other), s"Missing cache bounds implementation for $other")
-          (null, null, if (knownNulls >= 0) knownNulls else countNulls(col, numRows))
+          (null, null, knownNulls)
       }
       if (nullCount < numRows) {
         lower(c) = min
@@ -160,27 +160,13 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
     (lower, upper, nulls)
   }
 
-  // The column's null count if its Arrow vector already holds one, or -1 if only a pass over the
-  // rows can tell.
-  //
-  // A CometDecodedVector takes its null count from Arrow once, when it is built, and its isNullAt
-  // answers from the same validity bitmap, so the two agree whenever the vector holds exactly the
-  // batch's rows. Spark's own vectors do not qualify: its vectorized Parquet reader marks a column
-  // missing from the file all-null, and that vector's numNulls is then 0 while isNullAt is true
-  // for every row.
-  private def arrowNullCount(col: ColumnVector, numRows: Int): Int = col match {
-    case v: CometDecodedVector if v.numValues() == numRows => v.numNulls()
-    case _ => -1
-  }
-
-  private def countNulls(col: ColumnVector, numRows: Int): Int = {
-    var nullCount = 0
-    var r = 0
-    while (r < numRows) {
-      if (col.isNullAt(r)) nullCount += 1
-      r += 1
-    }
-    nullCount
+  // 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
@@ -425,9 +411,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
@@ -436,20 +419,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()
+        }
+      }
     }
   }
 
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 161620103b7..6844bcb072a 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -1163,9 +1163,9 @@ class CometInMemoryCacheSuite extends CometTestBase {
 
   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. Comet's Arrow vectors give the count without a pass over the rows,
-    // while Spark's vectors are counted row by row, so it is checked from the row write path,
-    // Comet's native scan and Spark's vectorized reader.
+    // 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")
@@ -1223,8 +1223,8 @@ class CometInMemoryCacheSuite extends CometTestBase {
   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 numNulls would record no nulls, and the IsNull
-    // filter below would prune every batch read from the file that lacks the column.
+    // 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 =>
@@ -1252,39 +1252,6 @@ class CometInMemoryCacheSuite extends CometTestBase {
     }
   }
 
-  test("Comet in-memory cache statistics count nulls only within a batch's rows") {
-    // A batch can report fewer rows than its vectors hold, and the nulls here are all past the
-    // batch's rows. A count taken from the whole vector would say every row of this batch is null,
-    // and Spark would then prune the batch for IsNotNull.
-    val ints = new IntVector("i", CometArrowAllocator)
-    val binaries = new VarBinaryVector("b", CometArrowAllocator)
-    try {
-      ints.allocateNew(4)
-      binaries.allocateNew(4)
-      Seq(0, 1).foreach { r =>
-        ints.set(r, r + 1)
-        binaries.set(r, Array[Byte](r.toByte))
-      }
-      Seq(2, 3).foreach { r =>
-        ints.setNull(r)
-        binaries.setNull(r)
-      }
-      ints.setValueCount(4)
-      binaries.setValueCount(4)
-      val batch = new ColumnarBatch(
-        Array[ColumnVector](new CometPlainVector(ints), new CometPlainVector(binaries)),
-        2)
-      val attrs =
-        Seq(AttributeReference("i", IntegerType)(), AttributeReference("b", BinaryType)())
-      val (lower, upper, nulls) = CometCachedBatchHelper.columnStats(batch, attrs)
-      assert(nulls.toSeq == Seq(0, 0))
-      assert(lower(0) == 1 && upper(0) == 2)
-    } finally {
-      ints.close()
-      binaries.close()
-    }
-  }
-
   test("Comet in-memory cache supports stats-based batch pruning") {
     withConversions(
       SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
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 87f74c7f7f1..89841d69715 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
@@ -788,9 +788,9 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase {
         columns(5).putByteArray(r, s"str_c_$r".getBytes(StandardCharsets.UTF_8))
         r += 1
       }
-      // The same values in Comet's Arrow vectors. Spark's vectors are what the collector reads
-      // from Spark's vectorized readers, and Arrow vectors are what it reads from Comet operators
-      // and from the row write path.
+      // 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),
@@ -798,7 +798,6 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase {
       try {
         val benchmark =
           new Benchmark("in-memory cache statistics", numRows.toLong, output = output)
-        addStatsCase(benchmark, "Spark vectors", batch, attrs)
         addStatsCase(benchmark, "Arrow vectors", arrowBatch, attrs)
         benchmark.run()
       } finally arrowBatch.close()
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala
index ba4b616b57b..47e876b5f86 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala
@@ -33,7 +33,6 @@ import org.apache.arrow.vector.ipc.{ReadChannel, WriteChannel}
 import org.apache.arrow.vector.ipc.message.{ArrowBodyCompression, ArrowRecordBatch, MessageMetadataResult, MessageSerializer}
 import org.apache.arrow.vector.types.pojo.Field
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.columnar.CachedBatch
 import org.apache.spark.sql.comet.util.Utils
 import org.apache.spark.sql.types.StructType
@@ -130,18 +129,6 @@ object CometCachedBatchHelper {
   def cachedBatch(payload: ChunkedByteBuffer, numRows: Int): CachedBatch =
     CometCachedBatch(numRows, InternalRow.empty, payload)
 
-  /**
-   * The lower bounds, upper bounds and null counts the writer records for `batch`, one entry per
-   * column. Another thin shim, for tests that hand the statistics collector a batch no plan
-   * produces.
-   */
-  def columnStats(
-      batch: ColumnarBatch,
-      attrs: Seq[Attribute]): (Array[Any], Array[Any], Array[Int]) = {
-    val serializer = new ArrowCachedBatchSerializer
-    serializer.gatherColumnStats(batch, attrs, serializer.boundsOrderings(attrs))
-  }
-
   /**
    * Decode the `selected` columns of a cached batch the way a scan does, into a root the caller
    * owns. Another thin shim, for tests about what the read path leaves allocated.

From f1f57aeef84a9242848fadef250b47bcbc1792f6 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 14 Oct 2026 09:04:22 -0600
Subject: [PATCH 4/9] perf: combine small columnar batches before caching them

The columnar cache write stored each input batch as one cached batch,
and nothing bounds how small those batches are. A shuffle read returns
at least one batch per map output block, so a relation cached straight
from a shuffle over 200 map tasks into 50 partitions held 10000 batches
of about 500 rows, each with its own message metadata, compression
framing and statistics row.

CacheBatchCoalescer now sits in front of the columnar write. An
Arrow-backed batch laid out the way the reader reads it is passed
through when it holds at least half of
spark.sql.inMemoryColumnarStorage.batchSize, so full native batches are
still written without a copy. Every other batch is copied through
ArrowWriter.writeColumns into batches of batchSize rows, the size the
row path cuts at: a small batch is worth the copy, and a batch of Spark
column vectors was already copied into Arrow. Row order is kept, and
each batch is copied before the input advances, because a native
producer releases the batch it returned last when asked for the next.

Closes #6950.
---
 .../user-guide/latest/in-memory-cache.md      |   8 +
 .../arrow/ArrowCachedBatchSerializer.scala    |   9 +-
 .../execution/arrow/CacheBatchCoalescer.scala | 216 ++++++++++++++++++
 .../comet/exec/CometInMemoryCacheSuite.scala  | 119 ++++++++++
 .../apache/comet/exec/ReleasingBatches.scala  |  89 ++++++++
 5 files changed, 440 insertions(+), 1 deletion(-)
 create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala
 create mode 100644 spark/src/test/scala/org/apache/comet/exec/ReleasingBatches.scala

diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md
index f0c2cd6a31e..f944d1f9d76 100644
--- a/docs/source/user-guide/latest/in-memory-cache.md
+++ b/docs/source/user-guide/latest/in-memory-cache.md
@@ -78,6 +78,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..a2beac6003e 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
@@ -516,16 +516,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/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
index b657cf1526f..9a45b46c62d 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -2036,6 +2036,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)
+  }
+}

From 93d8f7dbe62f5d5545094ce699ea59cb6e6fb6da Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 26 Oct 2026 05:58:48 -0600
Subject: [PATCH 5/9] fix: convert a cache batch whose vectors hold more than
 its rows

The cache statistics now come from the Arrow batch that is written, so a
vector longer than its batch would count nulls past the batch's last row
and could record the batch all null. Require each vector to hold exactly
the batch's rows to take the direct path, and convert otherwise.
---
 .../execution/arrow/CachedBatchIpc.scala      |  8 +++++-
 .../comet/exec/CometInMemoryCacheSuite.scala  | 25 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

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 16dc6813899..3c985d19210 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 90cf8298604..0f61d060e77 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -2697,6 +2697,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 =

From e30bc831b9585ba47c33530e4e4b9072ee1c9d1d Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 26 Oct 2026 06:24:46 -0600
Subject: [PATCH 6/9] bench: add decimal columns to the cache statistics
 benchmark

---
 .../CometInMemoryCacheBenchmark.scala         | 38 ++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

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(

From 8a9394fdaa7968cae27c260c10f6c3434602f4d8 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 26 Oct 2026 06:24:46 -0600
Subject: [PATCH 7/9] perf: compare decimal bounds in the Arrow data buffer
 when gathering cache statistics

---
 .../arrow/ArrowCachedBatchSerializer.scala    | 62 ++++++++++++++++++-
 1 file changed, 61 insertions(+), 1 deletion(-)

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 ddc6b31993d..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) = {

From f20e6e123e739de60d789d584f26601afc727295 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 26 Oct 2026 06:31:46 -0600
Subject: [PATCH 8/9] test: check decimal bounds across the word boundaries of
 a 128-bit value

---
 .../arrow/ArrowCachedBatchStatsSuite.scala    | 159 ++++++++++++++++++
 1 file changed, 159 insertions(+)
 create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchStatsSuite.scala

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..7ae5ee4f406
--- /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), nullEvery = 0)
+        check(precision, scale, randomValues(precision, 500, seed), 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)
+  }
+}

From 158e105a719a4bb2c08e2335a8a915c1e1ae0349 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 26 Oct 2026 06:41:50 -0600
Subject: [PATCH 9/9] test: widen the seed of the random decimal bounds
 explicitly

---
 .../comet/execution/arrow/ArrowCachedBatchStatsSuite.scala    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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
index 7ae5ee4f406..55eaacfc223 100644
--- 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
@@ -146,8 +146,8 @@ class ArrowCachedBatchStatsSuite extends AnyFunSuite {
 
     test(s"decimal($precision, $scale) bounds of random values and nulls") {
       (0 until 5).foreach { seed =>
-        check(precision, scale, randomValues(precision, 500, seed), nullEvery = 0)
-        check(precision, scale, randomValues(precision, 500, seed), nullEvery = 7)
+        check(precision, scale, randomValues(precision, 500, seed.toLong), nullEvery = 0)
+        check(precision, scale, randomValues(precision, 500, seed.toLong), nullEvery = 7)
       }
     }
   }