Migrate CollectTimeIterator to use new NvtxIdWithMetrics class

docs/dev/nvtx_ranges.md

@@ -20,13 +20,22 @@ See [nvtx_profiling.md](https://nvidia.github.io/spark-rapids/docs/dev/nvtx_prof
 Name | Description
 -----|-------------
 getMapSizesByExecId|Call to internal Spark API for retrieving size and location of shuffle map output blocks
-Release GPU|Releasing the GPU semaphore
+probe right|Probing the right side of a join input iterator to get the data size for preparing the join
+fetch join stream|stream iterator time for GpuShuffleSizeHashJoinExec
+shuffled join stream|GpuShuffledHashJoinExec op is preparing build batches for join
 Acquire GPU|Time waiting for GPU semaphore to be acquired
-ParallelDeserializerIterator.next|Calling next on the MT shuffle reader iterator
+Release GPU|Releasing the GPU semaphore
+GpuCoalesceBatches: collect|GPU combining of small batches post-kernel processing
+broadcast join stream|GpuBroadcastHashJoinExec.getBroadcastBuiltBatchAndStreamIter -  Gets the ColumnarBatch for the build side and the stream iterator by acquiring the GPU only after first stream batch has been streamed to GPU.
 CommitShuffle|After all temporary shuffle writes are done, produce a single file (shuffle_[map_id]_0) in the commit phase
+ParallelDeserializerIterator.next|Calling next on the MT shuffle reader iterator
 queueFetched|MT shuffle manager is using the RapidsShuffleBlockFetcherIterator to queue the next set of fetched results
 WaitingForWrites|Rapids Shuffle Manager (multi threaded) is waiting for any queued writes to finish before finalizing the map output writer
-ThreadedReader.read|Rapids Shuffle Manager (multi threaded) reading
+AbstractGpuCoalesceIterator|Default range for a code path in the AbstractGpuCoalesceIterator for an op which is not explicitly documented in its own range
 ThreadedWriter.write|Rapids Shuffle Manager (multi threaded) writing
-RapidsCachingWriter.write|Rapids Shuffle Manager (ucx) writing
+ThreadedReader.read|Rapids Shuffle Manager (multi threaded) reading
+hash join build|
+probe left|Probing the left side of a join input iterator to get the data size for preparing the join
+build batch: collect|Perform a join where the build side fits in a single GPU batch
 BatchWait|Rapids Shuffle Manager (multi threaded) reader blocked waiting for batches to finish decoding
+RapidsCachingWriter.write|Rapids Shuffle Manager (ucx) writing

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -245,6 +245,15 @@ case class BatchedByKey(gpuOrder: Seq[SortOrder])(val cpuOrder: Seq[SortOrder])
   override def children: Seq[Expression] = gpuOrder
 }
 
+object OpNameNvtxMap {
+  private val map = Map(
+    "GpuCoalesceBatches: collect" -> NvtxRegistry.GPU_COALESCE_BATCHES_COLLECT,
+    "build batch: collect" -> NvtxRegistry.BUILD_BATCH_COLLECT
+  )
+
+  def get(opName: String): Option[NvtxId] = map.get(opName)
+}
+
 abstract class AbstractGpuCoalesceIterator(
     inputIter: Iterator[ColumnarBatch],
     goal: CoalesceSizeGoal,
@@ -257,7 +266,9 @@ abstract class AbstractGpuCoalesceIterator(
     opTime: GpuMetric,
     opName: String) extends Iterator[ColumnarBatch] with Logging {
 
-  private val iter = new CollectTimeIterator(s"$opName: collect", inputIter, streamTime)
+  private val iter = new CollectTimeIterator(
+    OpNameNvtxMap.get(s"$opName: collect").getOrElse(NvtxRegistry.GPU_COALESCE_ITERATOR),
+    inputIter, streamTime)
 
   private var batchInitialized: Boolean = false
 

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuMetrics.scala

@@ -18,9 +18,6 @@ package com.nvidia.spark.rapids
 
 import scala.collection.immutable.TreeMap
 
-import ai.rapids.cudf.NvtxColor
-import com.nvidia.spark.rapids.Arm.withResource
-
 import org.apache.spark.{SparkContext, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.SparkSession
@@ -356,17 +353,17 @@ final class LocalGpuMetric extends GpuMetric {
 }
 
 class CollectTimeIterator[T](
-    nvtxName: String,
+    nvtxId: NvtxId,
     it: Iterator[T],
     collectTime: GpuMetric) extends Iterator[T] {
   override def hasNext: Boolean = {
-    withResource(new NvtxWithMetrics(nvtxName, NvtxColor.BLUE, collectTime)) { _ =>
+    NvtxIdWithMetrics(nvtxId, collectTime) {
       it.hasNext
     }
   }
 
   override def next(): T = {
-    withResource(new NvtxWithMetrics(nvtxName, NvtxColor.BLUE, collectTime)) { _ =>
+    NvtxIdWithMetrics(nvtxId, collectTime) {
       it.next
     }
   }

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuShuffledHashJoinExec.scala

@@ -241,7 +241,7 @@ case class GpuShuffledHashJoinExec(
       (streamIter, buildIter) => {
         val (buildData, maybeBufferedStreamIter) =
           GpuShuffledHashJoinExec.prepareBuildBatchesForJoin(buildIter,
-            new CollectTimeIterator("shuffled join stream", streamIter, streamTime),
+            new CollectTimeIterator(NvtxRegistry.SHUFFLED_JOIN_STREAM, streamIter, streamTime),
             realTarget, localBuildOutput, buildGoal, subPartConf, coalesceMetrics, readOption)
 
         buildData match {
@@ -371,7 +371,7 @@ object GpuShuffledHashJoinExec extends Logging {
             } else {
               logDebug("Return multiple batches as the build side data for the following " +
                 "sub-partitioning join")
-              Right(new CollectTimeIterator("hash join build", gpuBuildIter, buildTime))
+              Right(new CollectTimeIterator(NvtxRegistry.HASH_JOIN_BUILD, gpuBuildIter, buildTime))
             }
           }
           buildTime += System.nanoTime() - startTime
@@ -421,7 +421,7 @@ object GpuShuffledHashJoinExec extends Logging {
         val safeIter = GpuSubPartitionHashJoin.safeIteratorFromSeq(spillBuf.toSeq).map { sp =>
           withRetryNoSplit(sp)(_.getColumnarBatch())
         } ++ filteredIter
-        Right(new CollectTimeIterator("hash join build", safeIter, buildTime))
+        Right(new CollectTimeIterator(NvtxRegistry.HASH_JOIN_BUILD, safeIter, buildTime))
       } else {
         // The size after filtering is within the target size or sub-partitioning is disabled.
         while(filteredIter.hasNext) {

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuShuffledSizedHashJoinExec.scala

@@ -668,8 +668,10 @@ object GpuShuffledSymmetricHashJoinExec {
       val rightTime = new LocalGpuMetric
       val buildTime = metrics(BUILD_TIME)
       val streamTime = metrics(STREAM_TIME)
-      val leftIter = new CollectTimeIterator("probe left", setupForProbe(rawLeftIter), leftTime)
-      val rightIter = new CollectTimeIterator("probe right", setupForProbe(rawRightIter), rightTime)
+      val leftIter = new CollectTimeIterator(NvtxRegistry.PROBE_LEFT,
+        setupForProbe(rawLeftIter), leftTime)
+      val rightIter = new CollectTimeIterator(NvtxRegistry.PROBE_RIGHT,
+        setupForProbe(rawRightIter), rightTime)
       closeOnExcept(mutable.Queue.empty[T]) { leftQueue =>
         closeOnExcept(mutable.Queue.empty[T]) { rightQueue =>
           var leftSize = 0L
@@ -718,7 +720,7 @@ object GpuShuffledSymmetricHashJoinExec {
           } else {
             baseBuildIter
           }
-          val streamIter = new CollectTimeIterator("fetch join stream",
+          val streamIter = new CollectTimeIterator(NvtxRegistry.FETCH_JOIN_STREAM,
             setupForJoin(streamQueue, rawStreamIter, exprs.streamTypes, gpuBatchSizeBytes, metrics),
             streamTime)
           JoinInfo(joinType, buildSide, buildIter, buildSize, None, streamIter, exprs)

sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxRangeWithDoc.scala

@@ -86,6 +86,36 @@ object NvtxRegistry {
   val GET_MAP_SIZES_BY_EXEC_ID: NvtxId = NvtxId("getMapSizesByExecId", NvtxColor.CYAN,
     "Call to internal Spark API for retrieving size and location of shuffle map output blocks")
 
+  val GPU_COALESCE_BATCHES_COLLECT: NvtxId = NvtxId("GpuCoalesceBatches: collect", NvtxColor.BLUE,
+    "GPU combining of small batches post-kernel processing")
+
+  val BUILD_BATCH_COLLECT: NvtxId = NvtxId("build batch: collect", NvtxColor.BLUE,
+    "Perform a join where the build side fits in a single GPU batch")
+
+  val GPU_COALESCE_ITERATOR: NvtxId = NvtxId("AbstractGpuCoalesceIterator", NvtxColor.BLUE,
+    "Default range for a code path in the AbstractGpuCoalesceIterator for an op which " +
+      "is not explicitly documented in its own range")
+
+  val SHUFFLED_JOIN_STREAM: NvtxId = NvtxId("shuffled join stream", NvtxColor.BLUE,
+    "GpuShuffledHashJoinExec op is preparing build batches for join")
+
+  val HASH_JOIN_BUILD: NvtxId = NvtxId("hash join build", NvtxColor.BLUE,
+    "")
+
+  val PROBE_LEFT: NvtxId = NvtxId("probe left", NvtxColor.BLUE,
+    "Probing the left side of a join input iterator to get the data size for preparing the join")
+
+  val PROBE_RIGHT: NvtxId = NvtxId("probe right", NvtxColor.BLUE,
+    "Probing the right side of a join input iterator to get the data size for preparing the join")
+
+  val FETCH_JOIN_STREAM: NvtxId = NvtxId("fetch join stream", NvtxColor.BLUE,
+    "stream iterator time for GpuShuffleSizeHashJoinExec")
+
+  val BROADCAST_JOIN_STREAM: NvtxId = NvtxId("broadcast join stream", NvtxColor.BLUE,
+    "GpuBroadcastHashJoinExec.getBroadcastBuiltBatchAndStreamIter -  Gets the ColumnarBatch for " +
+      "the build side and the stream iterator by acquiring the GPU only after first stream batch " +
+      "has been streamed to GPU.")
+
   def init(): Unit = {
     register(ACQUIRE_GPU)
     register(RELEASE_GPU)
@@ -98,6 +128,15 @@ object NvtxRegistry {
     register(QUEUE_FETCHED)
     register(RAPIDS_CACHING_WRITER_WRITE)
     register(GET_MAP_SIZES_BY_EXEC_ID)
+    register(GPU_COALESCE_BATCHES_COLLECT)
+    register(BUILD_BATCH_COLLECT)
+    register(GPU_COALESCE_ITERATOR)
+    register(SHUFFLED_JOIN_STREAM)
+    register(HASH_JOIN_BUILD)
+    register(PROBE_LEFT)
+    register(PROBE_RIGHT)
+    register(FETCH_JOIN_STREAM)
+    register(BROADCAST_JOIN_STREAM)
   }
 }
 

sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxWithMetrics.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -48,6 +48,29 @@ class NvtxWithMetrics(name: String, color: NvtxColor, val metrics: GpuMetric*)
   }
 }
 
+/**
+ *  NvtxId with option to pass one or more nano timing metric(s) that are updated upon close
+ *  by the amount of time spent in the range
+ */
+object NvtxIdWithMetrics {
+
+  def apply[V](nvtxId: NvtxId, metrics: GpuMetric*)(block: => V): V = {
+    val needTracks = metrics.map(_.tryActivateTimer())
+    val start = System.nanoTime()
+
+    try {
+      nvtxId.apply(block)
+    } finally {
+      val time = System.nanoTime() - start
+      metrics.toSeq.zip(needTracks).foreach { pair =>
+        if (pair._2) {
+          pair._1.deactivateTimer(time)
+        }
+      }
+    }
+  }
+}
+
 class MetricRange(val metrics: GpuMetric*) extends AutoCloseable {
   val needTracks = metrics.map(_.tryActivateTimer())
   private val start = System.nanoTime()

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastHashJoinExecBase.scala

@@ -151,7 +151,7 @@ abstract class GpuBroadcastHashJoinExecBase(
         GpuBroadcastHelper.getBroadcastBuiltBatchAndStreamIter(
           broadcastRelation,
           buildSchema,
-          new CollectTimeIterator("broadcast join stream", it, streamTime))
+          new CollectTimeIterator(NvtxRegistry.BROADCAST_JOIN_STREAM, it, streamTime))
       // builtBatch will be closed in doJoin
       doJoin(builtBatch, streamIter, targetSize, numOutputRows, numOutputBatches, opTime, joinTime)
     }