Address review comments

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

@@ -68,7 +68,7 @@ class RapidsBufferCatalog extends AutoCloseable with Arm {
       override val id: RapidsBufferId,
       var priority: Long,
       spillCallback: SpillCallback)
-    extends RapidsBufferHandle with Arm {
+    extends RapidsBufferHandle {
 
     override def setSpillPriority(newPriority: Long): Unit = {
       priority = newPriority
@@ -142,7 +142,9 @@ class RapidsBufferCatalog extends AutoCloseable with Arm {
    * returns true, otherwise it returns false.
    *
    * @param handle handle to stop tracking
-   * @return
+   * @return true: if this was the last `RapidsBufferHandle` associated with the
+   *         underlying buffer.
+   *         false: if there are remaining live handles
    */
   private def stopTrackingHandle(handle: RapidsBufferHandle): Boolean = {
     withResource(acquireBuffer(handle)) { buffer =>
@@ -159,9 +161,7 @@ class RapidsBufferCatalog extends AutoCloseable with Arm {
           null
         } else {
           val newHandles = handles.filter(h => h != handle).map { h =>
-            if (h.getSpillPriority > maxPriority) {
-              maxPriority = h.getSpillPriority
-            }
+            maxPriority = maxPriority.max(h.getSpillPriority)
             h
           }
           if (newHandles.isEmpty) {
@@ -308,6 +308,10 @@ class RapidsBufferCatalog extends AutoCloseable with Arm {
   /**
    * Remove a buffer handle from the catalog and, if it this was the final handle,
    * release the resources of the registered buffers.
+   *
+   * @return true: if the buffer for this handle was removed from the spill framework
+   *               (`handle` was the last handle)
+   *         false: if buffer was not removed due to other live handles.
    */
   def removeBuffer(handle: RapidsBufferHandle): Boolean = {
     // if this is the last handle, remove the buffer
@@ -324,7 +328,7 @@ class RapidsBufferCatalog extends AutoCloseable with Arm {
   def numBuffers: Int = bufferMap.size()
 
   override def close(): Unit = {
-    bufferIdToHandles.forEach { case (_, handles) =>
+    bufferIdToHandles.values.forEach { handles =>
       handles.foreach(removeBuffer)
     }
     bufferIdToHandles.clear()
@@ -396,9 +400,7 @@ object RapidsBufferCatalog extends Logging with Arm {
   }
 
   private def closeImpl(): Unit = {
-    if (singleton != null) {
-      singleton.close()
-    }
+    singleton.close()
 
     if (memoryEventHandler != null) {
       // Workaround for shutdown ordering problems where device buffers allocated with this handler

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

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

@@ -194,10 +194,8 @@ class ShuffleBufferCatalog(
           tableMap.remove(id.tableId)
           val didRemove = catalog.removeBuffer(bufferIdToHandle.get(id))
           if (!didRemove) {
-            logWarning(s"Unable to remove from underlying storage ${id} when cleaning " +
+            logWarning(s"Unable to remove $id from underlying storage when cleaning " +
               s"shuffle blocks.")
-          } else {
-            logWarning(s"Did remove ${id}")
           }
         }
       }

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

@@ -58,16 +58,13 @@ class JustRowsColumnarBatch(numRows: Int, semWait: GpuMetric)
   override def numRows(): Int = numRows
   override def setSpillPriority(priority: Long): Unit = () // NOOP nothing to spill
 
-  private def makeJustRowsBatch(): ColumnarBatch = {
+  def getColumnarBatch(): ColumnarBatch = {
     GpuSemaphore.acquireIfNecessary(TaskContext.get(), semWait)
     new ColumnarBatch(Array.empty, numRows)
   }
+
   override def close(): Unit = () // NOOP nothing to close
   override val sizeInBytes: Long = 0L
-
-  def getColumnarBatch(): ColumnarBatch = {
-    makeJustRowsBatch()
-  }
 }
 
 /**

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

tests/src/test/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleClientSuite.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

tests/src/test/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleTestHelper.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.