Add withRetry to GpuBatchedBoundedWindowIterator (#14139)

Contributes to #13672

### Description

This PR:
Add retry support to GpuBatchedBoundedWindowIterator to handle OOM:
- Protect the following 3 operations with OOM retry support.
  - Window computation (by `computeWindowWithRetry`)
- Input batch concatenation with the cache (by
`getNextInputBatchWithRetry`)
  - Batch trim (by `trimWithRetry`)
- Add unit tests for retry and split-and-retry OOM scenarios

NDS numbers (:Seconds) with 10k data size shows no perf regressions. 

'rapids-4-spark_2.12-26.02.0-20260102.073925-32-cuda12.jar' was used for
nightly runs, not sure why it is a little slower. Will try to run this
more.

|ID|with PR| Nitghtly|
|--|--|--| 
|1| 1302 | 1369|
|2| 1315 | 1389|
|avg| 1308.5|1379|


### Checklists

- [x] This PR has added documentation for new or modified features or
behaviors.
- [x] This PR has added new tests or modified existing tests to cover
new code paths.
(Please explain in the PR description how the new code paths are tested,
such as names of the new/existing tests that cover them.)
- [x] Performance testing has been performed and its results are added
in the PR description. Or, an issue has been filed with a link in the PR
description.

---------

Signed-off-by: Firestarman <firestarmanllc@gmail.com>
Co-authored-by: Firestarman <firestarmanllc@gmail.com>

Author: thirtiseven

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2026, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,9 +16,11 @@
 
 package com.nvidia.spark.rapids.window
 
-import ai.rapids.cudf.{ColumnVector => CudfColumnVector, Table => CudfTable}
+import ai.rapids.cudf.Table
 import com.nvidia.spark.rapids._
-import com.nvidia.spark.rapids.Arm.withResource
+import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
+import com.nvidia.spark.rapids.RapidsPluginImplicits._
+import com.nvidia.spark.rapids.RmmRapidsRetryIterator.withRetryNoSplit
 import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion
 
 import org.apache.spark.TaskContext
@@ -46,7 +48,8 @@ class GpuBatchedBoundedWindowIterator(
 
   override def hasNext: Boolean = numUnprocessedInCache > 0 || input.hasNext
 
-  var cached: Option[Array[CudfColumnVector]] = None  // For processing with the next batch.
+  // For processing with the next batch.
+  private var cached: Option[SpillableColumnarBatch] = None
 
   private var numUnprocessedInCache: Int = 0  // numRows at the bottom not processed completely.
   private var numPrecedingRowsAdded: Int = 0  // numRows at the top, added for preceding context.
@@ -59,151 +62,140 @@ class GpuBatchedBoundedWindowIterator(
   }
 
   // Caches input column schema on first read.
-  var inputTypes: Option[Array[DataType]] = None
+  private var inputTypes: Option[Array[DataType]] = None
 
-  // Clears cached column vectors, after consumption.
-  private def clearCached(): Unit = {
-    cached.foreach(_.foreach(_.close))
+  // Clears the cache after consumption.
+  private[rapids] def clearCached(): Unit = {
+    cached.foreach(_.close())
     cached = None
   }
 
-  private def getNextInputBatch: SpillableColumnarBatch = {
-    // Sets column batch types using the types cached from the
-    // first input column read.
-    def optionallySetInputTypes(inputCB: ColumnarBatch): Unit = {
-      if (inputTypes.isEmpty) {
-        inputTypes = Some(GpuColumnVector.extractTypes(inputCB))
-      }
-    }
-
-    // Reads fresh batch from iterator, initializes input data-types if necessary.
-    def getFreshInputBatch: ColumnarBatch = {
-      val fresh_batch = input.next()
-      optionallySetInputTypes(fresh_batch)
-      fresh_batch
-    }
-
-    def concatenateColumns(cached: Array[CudfColumnVector],
-                           freshBatchTable: CudfTable)
-    : Array[CudfColumnVector] = {
-
-      if (cached.length != freshBatchTable.getNumberOfColumns) {
-        throw new IllegalArgumentException("Expected the same number of columns " +
-          "in input batch and cached batch.")
-      }
-      cached.zipWithIndex.map { case (cachedCol, idx) =>
-        CudfColumnVector.concatenate(cachedCol, freshBatchTable.getColumn(idx))
-      }
-    }
+  protected final def hasCache: Boolean = cached.isDefined // for unit test
 
+  protected def getNextInputBatchWithRetry: SpillableColumnarBatch = {
     // Either cached has unprocessed rows, or input.hasNext().
     if (input.hasNext) {
-      if (cached.isDefined) {
+      val freshSCB = closeOnExcept(input.next()) { freshCB =>
+        if (inputTypes.isEmpty) {
+          // initializes input data-types if necessary.
+          inputTypes = Some(GpuColumnVector.extractTypes(freshCB))
+        }
+        SpillableColumnarBatch(freshCB, SpillPriorities.ACTIVE_BATCHING_PRIORITY)
+      }
+      if (hasCache) {
         // Cached input AND new input rows exist.  Return concat-ed rows.
-        withResource(getFreshInputBatch) { freshBatchCB =>
-          withResource(GpuColumnVector.from(freshBatchCB)) { freshBatchTable =>
-            withResource(concatenateColumns(cached.get, freshBatchTable)) { concat =>
-              clearCached()
-              SpillableColumnarBatch(convertToBatch(inputTypes.get, concat),
-                SpillPriorities.ACTIVE_BATCHING_PRIORITY)
-            }
+        // The two input batches will be closed by `withRetryNoSplit`.
+        val concatedTbl = withRetryNoSplit(Seq(cached.get, freshSCB)) { toConcat =>
+          val cbs = toConcat.safeMap(_.getColumnarBatch())
+          val tbls = withResource(cbs)(_ => cbs.safeMap(GpuColumnVector.from))
+          withResource(tbls)(_ => Table.concatenate(tbls: _*))
+        }
+        withResource(concatedTbl) { _ =>
+          cached = None
+          closeOnExcept(GpuColumnVector.from(concatedTbl, inputTypes.get)) { concatedCB =>
+            SpillableColumnarBatch(concatedCB, SpillPriorities.ACTIVE_BATCHING_PRIORITY)
           }
         }
-      } else {
-          // No cached input available. Return fresh input rows, only.
-          SpillableColumnarBatch(getFreshInputBatch,
-                                 SpillPriorities.ACTIVE_BATCHING_PRIORITY)
+      } else { // no cache, return the input batch directly
+        freshSCB
       }
-    }
-    else {
-      // No fresh input available. Return cached input.
-      val cachedCB = convertToBatch(inputTypes.get, cached.get)
-      clearCached()
-      SpillableColumnarBatch(cachedCB,
-                             SpillPriorities.ACTIVE_BATCHING_PRIORITY)
+    } else { // No fresh input available. Return cached input.
+      val cachedSCB = cached.get
+      cached = None
+      cachedSCB
     }
   }
 
   /**
    * Helper to trim specified number of rows off the top and bottom,
    * of all specified columns.
    */
-  private def trim(columns: Array[CudfColumnVector],
-                    offTheTop: Int,
-                    offTheBottom: Int): Array[CudfColumnVector] = {
+  protected def trimWithRetry(scb: SpillableColumnarBatch, offTheTop: Int,
+      offTheBottom: Int): ColumnarBatch = {
+    val batchRowsNum = scb.numRows()
+    if ((offTheTop + offTheBottom) > batchRowsNum) {
+      throw new IllegalArgumentException(s"Cannot trim batch of size ${batchRowsNum} by " +
+        s"$offTheTop rows at the top, and $offTheBottom rows at the bottom.")
+    }
 
-    def checkValidSizes(col: CudfColumnVector): Unit =
-      if ((offTheTop + offTheBottom) > col.getRowCount) {
-        throw new IllegalArgumentException(s"Cannot trim column of size ${col.getRowCount} by " +
-          s"$offTheTop rows at the top, and $offTheBottom rows at the bottom.")
+    withRetryNoSplit[ColumnarBatch] {
+      withResource(scb.getColumnarBatch()) { cb =>
+        val types = GpuColumnVector.extractTypes(cb)
+        val baseCols = GpuColumnVector.extractBases(cb)
+        withResource(baseCols.safeMap(_.subVector(offTheTop, batchRowsNum - offTheBottom))) {
+          convertToBatch(types, _)
+        }
       }
+    }
+  }
 
-    columns.map{ col =>
-      checkValidSizes(col)
-      col.subVector(offTheTop, col.getRowCount.toInt - offTheBottom)
+  /**
+   * Compute the window aggregations on the input batch with retry support.
+   * It does not close the input batch.
+   */
+  protected def computeWindowWithRetry(
+      scb: SpillableColumnarBatch): SpillableColumnarBatch = {
+    val outputCB = withRetryNoSplit[ColumnarBatch] {
+      val outputCols = withResource(scb.getColumnarBatch())(computeBasicWindow)
+      withResource(outputCols)(convertToBatch(outputTypes, _))
     }
+    SpillableColumnarBatch(outputCB, SpillPriorities.ACTIVE_BATCHING_PRIORITY)
   }
 
-  private def resetInputCache(newCache: Option[Array[CudfColumnVector]],
+  private def resetInputCache(newCache: Option[ColumnarBatch],
                               newPrecedingAdded: Int): Unit= {
-    cached.foreach(_.foreach(_.close))
-    cached = newCache
+    clearCached()
+    cached = newCache.map(
+      SpillableColumnarBatch(_, SpillPriorities.ACTIVE_BATCHING_PRIORITY)
+    )
     numPrecedingRowsAdded = newPrecedingAdded
   }
 
   override def next(): ColumnarBatch = {
     var outputBatch: ColumnarBatch = null
-    while (outputBatch == null  &&  hasNext) {
-      withResource(getNextInputBatch) { inputCbSpillable =>
-        withResource(inputCbSpillable.getColumnarBatch()) { inputCB =>
-
-          val inputRowCount = inputCB.numRows()
-          val noMoreInput = !input.hasNext
-          numUnprocessedInCache = if (noMoreInput) {
-            // If there are no more input rows expected,
-            // this is the last output batch.
-            // Consider all rows in the batch as processed.
+    while (outputBatch == null && hasNext) {
+      withResource(getNextInputBatchWithRetry) { inputSCB =>
+        val inputRowCount = inputSCB.numRows()
+        val noMoreInput = !input.hasNext
+        numUnprocessedInCache = if (noMoreInput) {
+          // If there are no more input rows expected,
+          // this is the last output batch.
+          // Consider all rows in the batch as processed.
+          0
+        } else {
+          // More input rows expected. The last `maxFollowing` rows can't be finalized.
+          // Cannot exceed `inputRowCount`.
+          if (maxFollowing < 0) { // E.g. LAG(3) => [ preceding=-3, following=-3 ]
+            // -ve following => No need to wait for more following rows.
+            // All "following" context is already available in the current batch.
             0
           } else {
-            // More input rows expected. The last `maxFollowing` rows can't be finalized.
-            // Cannot exceed `inputRowCount`.
-            if (maxFollowing < 0) { // E.g. LAG(3) => [ preceding=-3, following=-3 ]
-              // -ve following => No need to wait for more following rows.
-              // All "following" context is already available in the current batch.
-              0
-            } else {
-              maxFollowing min inputRowCount
-            }
+            maxFollowing min inputRowCount
           }
+        }
 
-          if (numPrecedingRowsAdded + numUnprocessedInCache >= inputRowCount) {
-            // No point calling windowing kernel: the results will simply be ignored.
-            logWarning("Not enough rows! Cannot output a batch.")
-          } else {
-            NvtxIdWithMetrics(NvtxRegistry.WINDOW_EXEC, opTime) {
-              withResource(computeBasicWindow(inputCB)) { outputCols =>
-                outputBatch =  withResource(
-                                  trim(outputCols,
-                                    numPrecedingRowsAdded, numUnprocessedInCache)) { trimmed =>
-                                  convertToBatch(outputTypes, trimmed)
-                }
-              }
+        if (numPrecedingRowsAdded + numUnprocessedInCache >= inputRowCount) {
+          // No point calling windowing kernel: the results will simply be ignored.
+          logWarning("Not enough rows! Cannot output a batch.")
+        } else {
+          NvtxIdWithMetrics(NvtxRegistry.WINDOW_EXEC, opTime) {
+            outputBatch = withResource(computeWindowWithRetry(inputSCB)) {
+              trimWithRetry(_, numPrecedingRowsAdded, numUnprocessedInCache)
             }
           }
+        }
 
-          // Compute new cache using current input.
-          numPrecedingRowsAdded = if (minPreceding > 0) { // E.g. LEAD(3) => [prec=3, foll=3]
-            // preceding > 0 => No "preceding" rows need be carried forward.
-            // Only the rows that need to be recomputed.
-            0
-          } else {
-            Math.abs(minPreceding) min (inputRowCount - numUnprocessedInCache)
-          }
-          val inputCols = Range(0, inputCB.numCols()).map {
-            inputCB.column(_).asInstanceOf[GpuColumnVector].getBase
-          }.toArray
+        // Compute new cache using current input.
+        numPrecedingRowsAdded = if (minPreceding > 0) { // E.g. LEAD(3) => [prec=3, foll=3]
+          // preceding > 0 => No "preceding" rows need be carried forward.
+          // Only the rows that need to be recomputed.
+          0
+        } else {
+          Math.abs(minPreceding) min (inputRowCount - numUnprocessedInCache)
+        }
 
-          val newCached = trim(inputCols,
+        if (!noMoreInput) { // cache is needed only when have more input data
+          val newCached = trimWithRetry(inputSCB,
             inputRowCount - (numPrecedingRowsAdded + numUnprocessedInCache),
             0)
           resetInputCache(Some(newCached), numPrecedingRowsAdded)