Refactor GpuSorter further to keep it in SortUtils and just shim out …

…the Stack implementation Signed-off-by: Navin Kumar <[email protected]>
NVIDIA · Oct 26, 2023 · eb2dbd3 · eb2dbd3
1 parent 5db5520
commit eb2dbd3
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 311 deletions.
diff --git a/sql-plugin/src/main/scala-2.12/com/nvidia/spark/rapids/GpuSorter.scala b/sql-plugin/src/main/scala-2.12/com/nvidia/spark/rapids/GpuSorter.scala
diff --git a/...vidia/spark/rapids/CollectionHelper.scala → .../com/nvidia/spark/rapids/ScalaStack.scala b/...vidia/spark/rapids/CollectionHelper.scala → .../com/nvidia/spark/rapids/ScalaStack.scala
@@ -18,8 +18,4 @@ package com.nvidia.spark.rapids
 
 import scala.collection.mutable.ArrayStack
 
-object CollectionHelper {
-  def createStack[T](): ArrayStack[T] = {
-    new ArrayStack[T]()
-  } 
-}
+class ScalaStack[T] extends ArrayStack[T]
diff --git a/sql-plugin/src/main/scala-2.13/com/nvidia/spark/rapids/GpuSorter.scala b/sql-plugin/src/main/scala-2.13/com/nvidia/spark/rapids/GpuSorter.scala
diff --git a/...vidia/spark/rapids/CollectionHelper.scala → .../com/nvidia/spark/rapids/ScalaStack.scala b/...vidia/spark/rapids/CollectionHelper.scala → .../com/nvidia/spark/rapids/ScalaStack.scala
@@ -18,8 +18,4 @@ package com.nvidia.spark.rapids
 
 import scala.collection.mutable.Stack
 
-object CollectionHelper {
-  def createStack[T](): Stack[T] = {
-    new Stack[T]()
-  } 
-}
+class ScalaStack[T] extends Stack[T]
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuSortExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuSortExec.scala
@@ -450,7 +450,7 @@ case class GpuOutOfCoreSortIterator(
     while (!pending.isEmpty && sortedSize < targetSize) {
       // Keep going until we have enough data to return
       var bytesLeftToFetch = targetSize
-      val pendingSort = CollectionHelper.createStack[SpillableColumnarBatch]()
+      val pendingSort = new ScalaStack[SpillableColumnarBatch]()
       closeOnExcept(pendingSort) { _ =>
         while (!pending.isEmpty &&
             (bytesLeftToFetch - pending.peek().buffer.sizeInBytes >= 0 || pendingSort.isEmpty)) {