perf: read cache statistics null counts from Comet's Arrow vectors

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.
@@ -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,10 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
     var c = 0
     while (c < numCols) {
       val col = batch.column(c)
+      val knownNulls = arrowNullCount(col)
       val (min, max, nullCount) = attrs(c).dataType match {
+        // Every row is null, so there are no bounds and nothing to read.
+        case _ if knownNulls == numRows => (null, null, numRows)
         case BooleanType => gatherBooleanStats(col, numRows)
         case ByteType => gatherByteStats(col, numRows)
         case ShortType => gatherShortStats(col, numRows)
@@ -144,13 +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, knownNulls)
       }
       if (nullCount < numRows) {
         lower(c) = min
@@ -163,6 +160,15 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
     (lower, upper, nulls)
   }
 
+  // Every column of a batch encodeBatches writes is a CometDecodedVector, which takes its null
+  // count from Arrow once, when it is built, and answers isNullAt from the same validity bitmap.
+  private def arrowNullCount(col: ColumnVector): Int = col match {
+    case v: CometDecodedVector => v.numNulls()
+    case other =>
+      throw new IllegalStateException(
+        s"Cache statistics read from ${other.getClass.getName}, not a Comet Arrow vector")
+  }
+
   // Keep each loop specialized and box only its final bounds. r == nullCount identifies
   // the first non-null value, including when a column starts with nulls.
   private def gatherBooleanStats(col: ColumnVector, numRows: Int): (Boolean, Boolean, Int) = {
@@ -405,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
@@ -416,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()
+        }
+      }
     }
   }
 

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

spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala

@@ -40,11 +40,11 @@ import org.apache.spark.sql.catalyst.expressions.{And, Ascending, Attribute, Att
 import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
 import org.apache.spark.sql.catalyst.plans.{LeftOuter, RightOuter}
 import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch}
-import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoin, CometHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec}
+import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoin, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec}
 import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper}
 import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager
 import org.apache.spark.sql.comet.util.Utils
-import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FileSourceScanExec, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo}
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec}
 import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike}
@@ -1147,6 +1147,117 @@ class CometInMemoryCacheSuite extends CometTestBase {
     }
   }
 
+  // Each column's null count summed over the cached batches. Asserts first that the relation is
+  // stored in Comet's format, because Spark's DefaultCachedBatch records the same null counts and
+  // no bounds for binary, nested and void columns, and then that no batch records bounds for a
+  // column.
+  private def nullCountsWithoutBounds(view: String): Seq[Long] = {
+    val batchTypes = cachedBatchTypes(view)
+    assert(
+      batchTypes.sameElements(
+        Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")),
+      s"$view was cached as ${batchTypes.mkString(", ")}")
+    val relation = spark.sharedState.cacheManager.lookupCachedData(spark.table(view)).get
+    val stats = relation.cachedRepresentation.cacheBuilder.cachedColumnBuffers
+      .collect()
+      .map(_.asInstanceOf[SimpleMetricsCachedBatch].stats)
+    relation.cachedRepresentation.output.indices.map { i =>
+      stats.foreach(s => assert(s.isNullAt(i * 5) && s.isNullAt(i * 5 + 1), s"column $i"))
+      stats.map(_.getInt(i * 5 + 2).toLong).sum
+    }
+  }
+
+  test("Comet in-memory cache statistics count nulls in columns that record no bounds") {
+    // These columns record a null count and no bounds: binary and nested types, and a bigint
+    // column with no values. The count is the one the written Arrow vectors hold, so it is checked
+    // from the row write path and Comet's native scan, which produce those vectors, and from
+    // Spark's vectorized reader, whose batches are converted to them first.
+    val df = spark
+      .range(1000)
+      .selectExpr("if(id % 7 = 0, NULL, id) AS v")
+      .selectExpr(
+        "cast(concat('b', v) AS binary) AS bin",
+        "if(v IS NULL, NULL, array(v, v + 1)) AS arr",
+        "if(v IS NULL, NULL, map(concat('k', v), v)) AS mp",
+        "if(v IS NULL, NULL, named_struct('a', v, 'b', cast(v AS string))) AS st",
+        "cast(NULL AS bigint) AS no_values")
+    // 143 of the ids 0 to 999 are multiples of 7.
+    val expected = Seq(143L, 143L, 143L, 143L, 1000L)
+
+    withSQLConf(
+      CometConf.COMET_ENABLED.key -> "false",
+      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+      // Only the row write path can take a void column; Parquet cannot store one.
+      val rows = df.withColumn("nt", lit(null))
+      rows.createOrReplaceTempView("no_bounds_rows")
+      rows.cache()
+      try {
+        rows.count()
+        assert(nullCountsWithoutBounds("no_bounds_rows") == expected :+ 1000L)
+      } finally {
+        rows.unpersist(blocking = true)
+        spark.catalog.dropTempView("no_bounds_rows")
+      }
+    }
+
+    withTempPath { path =>
+      df.write.parquet(path.toString)
+      withNativeCache {
+        spark.read.parquet(path.toString).createOrReplaceTempView("no_bounds_native")
+        spark.catalog.cacheTable("no_bounds_native")
+        spark.table("no_bounds_native").count()
+        val cachedPlan = spark.sharedState.cacheManager
+          .lookupCachedData(spark.table("no_bounds_native"))
+          .get
+          .cachedRepresentation
+          .cacheBuilder
+          .cachedPlan
+        assert(cachedPlan.isInstanceOf[CometNativeScanExec], cachedPlan)
+        assert(nullCountsWithoutBounds("no_bounds_native") == expected)
+      }
+    }
+
+    withSparkColumnarCache(
+      "no_bounds_columnar",
+      SQLConf.PARQUET_VECTORIZED_READER_NESTED_COLUMN_ENABLED.key -> "true") { path =>
+      df.write.parquet(path)
+    } { _ =>
+      assert(nullCountsWithoutBounds("no_bounds_columnar") == expected)
+    }
+  }
+
+  test("Comet in-memory cache statistics count the nulls of a column missing from a file") {
+    // Spark's vectorized Parquet reader marks a column that a file lacks all-null without counting
+    // any nulls, so that vector's numNulls is 0 while every row isNullAt. A binary column records
+    // that count and no bounds, so taking it from that vector would record no nulls, and the
+    // IsNull filter below would prune every batch read from the file that lacks the column.
+    withSparkColumnarCache(
+      "missing_column_columnar",
+      SQLConf.PARQUET_SCHEMA_MERGING_ENABLED.key -> "true") { path =>
+      spark.range(0, 400).selectExpr("id").write.parquet(path)
+      spark
+        .range(400, 1000)
+        .selectExpr("id", "cast(cast(id AS string) AS binary) AS b")
+        .write
+        .mode("append")
+        .parquet(path)
+    } { _ =>
+      val relation = spark.sharedState.cacheManager
+        .lookupCachedData(spark.table("missing_column_columnar"))
+        .get
+        .cachedRepresentation
+      assert(relation.cacheBuilder.cachedPlan.isInstanceOf[FileSourceScanExec])
+      checkAnswer(
+        spark.sql("SELECT count(*) FROM missing_column_columnar WHERE b IS NULL"),
+        Row(400L))
+      val b = relation.output.indexWhere(_.name == "b")
+      val stats = relation.cacheBuilder.cachedColumnBuffers
+        .collect()
+        .map(_.asInstanceOf[SimpleMetricsCachedBatch].stats)
+      assert(stats.map(_.getInt(b * 5 + 2)).sum == 400)
+    }
+  }
+
   test("Comet in-memory cache supports stats-based batch pruning") {
     withConversions(
       SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
@@ -2588,6 +2699,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 =