bug: no column projection should still persist row count (#4444)

Author: coderfender

.github/workflows/pr_build_linux.yml

@@ -383,6 +383,7 @@ jobs:
               org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite
               org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite
               org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite
+              org.apache.spark.sql.comet.util.UtilsSuite
               org.apache.comet.objectstore.NativeConfigSuite
               org.apache.spark.sql.CometToPrettyStringSuite
               org.apache.spark.sql.CometCollationSuite

.github/workflows/pr_build_macos.yml

@@ -223,6 +223,7 @@ jobs:
               org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite
               org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite
               org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite
+              org.apache.spark.sql.comet.util.UtilsSuite
               org.apache.comet.objectstore.NativeConfigSuite
               org.apache.spark.sql.CometToPrettyStringSuite
               org.apache.spark.sql.CometCollationSuite

spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala

@@ -224,6 +224,10 @@ object Utils extends CometTypeShim with Logging {
 
       val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch)
       val root = new VectorSchemaRoot(fieldVectors.asJava)
+      if (fieldVectors.isEmpty) {
+        // VSR cannot infer rowCount without field vectors
+        root.setRowCount(batch.numRows())
+      }
       val provider = batchProviderOpt.getOrElse(dictionaryProvider)
 
       val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out))
@@ -336,6 +340,11 @@ object Utils extends CometTypeShim with Logging {
           return (Array.empty, 0L, 0L)
         }
 
+        if (targetRoot.getSchema.getFields.isEmpty) {
+          // VSRAppender does not update rowCount with no columns
+          targetRoot.setRowCount(totalRows.toInt)
+        }
+
         assert(
           targetRoot.getRowCount.toLong == totalRows,
           s"Row count mismatch after coalesce: ${targetRoot.getRowCount} != $totalRows")

spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala

@@ -0,0 +1,54 @@
+/*
+ * 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.util
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+class UtilsSuite extends CometTestBase {
+
+  test("serializeBatches preserves row count for a zero-column batch") {
+    val numRows = 5
+    val batch = new ColumnarBatch(Array.empty[ColumnVector], numRows)
+
+    val (rowCount, buf) = Utils.serializeBatches(Iterator(batch)).next()
+    assert(rowCount == numRows)
+
+    val decoded = Utils.decodeBatches(buf, "test").toSeq
+    assert(decoded.map(_.numRows()).sum == numRows)
+  }
+
+  test("coalesceBroadcastBatches preserves row count across zero-column inputs") {
+    val numRows = 5
+    val numBatches = 3
+    val batches =
+      (0 until numBatches).map(_ => new ColumnarBatch(Array.empty[ColumnVector], numRows))
+
+    val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq.iterator
+    val (coalesced, batchCount, totalRows) = Utils.coalesceBroadcastBatches(bufs)
+
+    val expected = numRows.toLong * numBatches
+    assert(batchCount == numBatches)
+    assert(totalRows == expected)
+
+    val decoded = coalesced.iterator.flatMap(b => Utils.decodeBatches(b, "test")).toSeq
+    assert(decoded.map(_.numRows()).sum == expected)
+  }
+}