perf: reduce per-batch broadcast serializer allocations

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

@@ -260,15 +260,25 @@ object Utils extends CometTypeShim with Logging {
    *   the output stream
    */
   def serializeBatches(batches: Iterator[ColumnarBatch]): Iterator[(Long, ChunkedByteBuffer)] = {
+    lazy val codec = CompressionCodec.createCodec(SparkEnv.get.conf)
     batches.map { batch =>
       val dictionaryProvider: CDataDictionaryProvider = new CDataDictionaryProvider
 
-      val codec = CompressionCodec.createCodec(SparkEnv.get.conf)
-      val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate)
-      val out = new DataOutputStream(codec.compressedOutputStream(cbbos))
-
       val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch)
       try {
+        // Tiny shuffle batches otherwise allocate and zero a full 1 MiB heap chunk each. Keep
+        // larger chunks for large inputs, where a small fixed chunk would fragment the output.
+        val dictionaryBytes = batchProviderOpt
+          .map { provider =>
+            provider.getDictionaryIds.asScala.foldLeft(0L) { (bytes, id) =>
+              bytes + provider.lookup(id).getVector.getBufferSize
+            }
+          }
+          .getOrElse(0L)
+        val batchBytes = fieldVectors.foldLeft(dictionaryBytes)(_ + _.getBufferSize)
+        val chunkSize = math.max(4096L, math.min(1024 * 1024L, batchBytes)).toInt
+        val cbbos = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate)
+        val out = new DataOutputStream(codec.compressedOutputStream(cbbos))
         val root = new VectorSchemaRoot(fieldVectors.asJava)
         if (fieldVectors.isEmpty) {
           // VSR cannot infer rowCount without field vectors
@@ -292,19 +302,19 @@ object Utils extends CometTypeShim with Logging {
           // the write fails, as it does when a cancelled broadcast interrupts the task.
           if (written ne root) written.close()
         }
+
+        if (out.size() > 0) {
+          (batch.numRows().toLong, cbbos.toChunkedByteBuffer)
+        } else {
+          (batch.numRows().toLong, new ChunkedByteBuffer(Array.empty[ByteBuffer]))
+        }
       } catch {
         case NonFatal(e) =>
           // The batch's own vectors stay with their owner. The ones made for this write are held
           // by nothing else, and a narrowed copy can be close to 2 GiB.
           closeAfterFailure(allocatedFor(batch, fieldVectors), e)
           throw e
       }
-
-      if (out.size() > 0) {
-        (batch.numRows().toLong, cbbos.toChunkedByteBuffer)
-      } else {
-        (batch.numRows().toLong, new ChunkedByteBuffer(Array.empty[ByteBuffer]))
-      }
     }
   }
 

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

@@ -43,6 +43,38 @@ import org.apache.comet.vector.CometVector
 
 class UtilsSuite extends CometTestBase {
 
+  test("serializeBatches round trips schema metadata larger than the output chunk estimate") {
+    val random = new scala.util.Random(42)
+    // Keep vector data below 4 KiB while the varied field names exceed it even after compression.
+    val vectors = (0 until 256).map { i =>
+      val vector = new IntVector(random.alphanumeric.take(64).mkString, CometArrowAllocator)
+      vector.allocateNew(1)
+      vector.set(0, i)
+      vector.setValueCount(1)
+      vector
+    }
+    val columns = vectors.map { vector =>
+      CometVector.getVector(vector, new CDataDictionaryProvider): ColumnVector
+    }
+    val batch = new ColumnarBatch(columns.toArray, 1)
+    try {
+      assert(vectors.map(_.getBufferSize).sum < 4096)
+      val (rows, buffer) = Utils.serializeBatches(Iterator(batch)).next()
+      try {
+        assert(rows == 1)
+        assert(buffer.size > 4096)
+        assert(buffer.getChunks.length > 1)
+        assert(buffer.getChunks.head.remaining() == 4096)
+        val decoded = Utils.decodeBatches(buffer, "schema-chunks")
+        val result = decoded.next()
+        assert(result.numRows() == 1)
+        assert(result.numCols() == vectors.length)
+        vectors.indices.foreach(i => assert(result.column(i).getInt(0) == i))
+        assert(!decoded.hasNext)
+      } finally buffer.dispose()
+    } finally batch.close()
+  }
+
   test("serializeBatches preserves row count for a zero-column batch") {
     val numRows = 5
     val batch = new ColumnarBatch(Array.empty[ColumnVector], numRows)