Add metrics, spotless.

Author: mbutrovich

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

@@ -26,7 +26,7 @@ import java.nio.channels.Channels
 import scala.jdk.CollectionConverters._
 
 import org.apache.arrow.c.CDataDictionaryProvider
-import org.apache.arrow.vector.{BigIntVector, BitVector, DateDayVector, DecimalVector, FieldVector, FixedSizeBinaryVector, Float4Vector, Float8Vector, IntVector, NullVector, SmallIntVector, TimeStampMicroTZVector, TimeStampMicroVector, TinyIntVector, ValueVector, VarBinaryVector, VarCharVector, VectorSchemaRoot}
+import org.apache.arrow.vector._
 import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
 import org.apache.arrow.vector.dictionary.DictionaryProvider
 import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
@@ -234,6 +234,7 @@ object Utils extends CometTypeShim with Logging {
 
   /**
    * Decodes the byte arrays back to ColumnarBatchs and put them into buffer.
+   *
    * @param bytes
    *   the serialized batches
    * @param source
@@ -264,10 +265,11 @@ object Utils extends CometTypeShim with Logging {
    * re-serialize once via ArrowStreamWriter. This is done on the driver (not per-task) so the
    * cost is paid once rather than once per consumer partition.
    */
-  def coalesceBroadcastBatches(input: Iterator[ChunkedByteBuffer]): Array[ChunkedByteBuffer] = {
+  def coalesceBroadcastBatches(
+      input: Iterator[ChunkedByteBuffer]): (Array[ChunkedByteBuffer], Long, Long) = {
     val buffers = input.filterNot(_.size == 0).toArray
     if (buffers.isEmpty) {
-      return Array.empty
+      return (Array.empty, 0L, 0L)
     }
 
     val allocator = org.apache.comet.CometArrowAllocator
@@ -308,7 +310,7 @@ object Utils extends CometTypeShim with Logging {
         }
 
         if (targetRoot == null) {
-          return Array.empty
+          return (Array.empty, 0L, 0L)
         }
 
         assert(
@@ -320,7 +322,7 @@ object Utils extends CometTypeShim with Logging {
         val outCodec = CompressionCodec.createCodec(SparkEnv.get.conf)
         val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate)
         val out = new DataOutputStream(outCodec.compressedOutputStream(cbbos))
-        // null provider is safe here — we assert no dictionary-encoded columns above
+        // null provider is safe here because we assert no dictionary-encoded columns above
         val writer = new ArrowStreamWriter(targetRoot, null, Channels.newChannel(out))
         try {
           writer.start()
@@ -329,7 +331,7 @@ object Utils extends CometTypeShim with Logging {
           writer.close()
         }
 
-        Array(cbbos.toChunkedByteBuffer)
+        (Array(cbbos.toChunkedByteBuffer), batchCount.toLong, totalRows)
       } finally {
         if (targetRoot != null) {
           targetRoot.close()

spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala

@@ -77,7 +77,13 @@ case class CometBroadcastExchangeExec(
     "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
     "collectTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to collect"),
     "buildTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to build"),
-    "broadcastTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to broadcast"))
+    "broadcastTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to broadcast"),
+    "numCoalescedBatches" -> SQLMetrics.createMetric(
+      sparkContext,
+      "number of coalesced batches for broadcast"),
+    "numCoalescedRows" -> SQLMetrics.createMetric(
+      sparkContext,
+      "number of coalesced rows for broadcast"))
 
   override def doCanonicalize(): SparkPlan = {
     CometBroadcastExchangeExec(null, null, mode, child.canonicalized)
@@ -157,7 +163,9 @@ case class CometBroadcastExchangeExec(
 
         // Coalesce many small per-partition buffers into a single buffer so each
         // consumer partition only deserializes one Arrow IPC stream.
-        val batches = Utils.coalesceBroadcastBatches(input)
+        val (batches, coalescedBatches, coalescedRows) = Utils.coalesceBroadcastBatches(input)
+        longMetric("numCoalescedBatches") += coalescedBatches
+        longMetric("numCoalescedRows") += coalescedRows
 
         val dataSize = batches.map(_.size).sum
 

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

@@ -427,13 +427,24 @@ class CometJoinSuite extends CometTestBase {
           // Without coalescing, build_input_batches would be ~numPartitions per task,
           // totaling ~numPartitions * numPartitions across all tasks.
           // With coalescing, each task gets 1 batch, so total ≈ numPartitions.
-          // scalastyle:off println
-          println(s"Build-side metrics: batches=$buildBatches, rows=$buildRows")
-          // scalastyle:on println
           assert(
             buildBatches <= numPartitions,
             s"Expected at most $numPartitions build batches (1 per task), got $buildBatches. " +
               "Broadcast batch coalescing may not be working.")
+
+          val broadcasts = collect(df2.queryExecution.executedPlan) {
+            case b: CometBroadcastExchangeExec => b
+          }
+          assert(broadcasts.nonEmpty, "Expected CometBroadcastExchangeExec in plan")
+
+          val broadcast = broadcasts.head
+          val coalescedBatches = broadcast.metrics("numCoalescedBatches").value
+          val coalescedRows = broadcast.metrics("numCoalescedRows").value
+
+          assert(
+            coalescedBatches >= numPartitions,
+            s"Expected at least $numPartitions coalesced batches, got $coalescedBatches")
+          assert(coalescedRows == 10000, s"Expected 10000 coalesced rows, got $coalescedRows")
         }
       }
     }