Document .split better

NVIDIA · Mar 7, 2023 · dd1bad3 · dd1bad3
1 parent 3a5c1d0
commit dd1bad3
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RmmRapidsRetryIterator.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RmmRapidsRetryIterator.scala
@@ -240,6 +240,14 @@ object RmmRapidsRetryIterator extends Arm {
   trait Spliterator[K] extends Iterator[K] with AutoCloseable {
     override def hasNext: Boolean
 
+    /**
+     * Split is a function that is invoked by `RmmRapidsRetryIterator` when `SplitAndRetryOOM`
+     * is thrown. This function is implemented by `Spliterator` classes to attempt to handle
+     * this exception by reducing the size of attempts (the thing that `.next` is
+     * using as an input), usually by splitting a batch in half by number of rows, or
+     * splitting a collection of batches into smaller collections to be attempted separately,
+     * likely reducing GPU memory that needs to be manifested while calling `.next`.
+     */
     def split(): Unit
 
     override def next(): K