Use per-batch retry block in R2C with lightweight OOM recovery (#14428)

thirtiseven · claude · web-flow · commit d7fd09353fc8 · 2026-04-07T17:01:15.000+08:00
Fixes #14368 ### Description PR #13842 wrapped each row in `RowToColumnarIterator.buildBatch` with `withRetryNoSplit` + `withRestoreOnRetry`, which introduced significant per-row overhead (~19x slower) from JNI calls, synchronized blocks, and object allocations on every row — even in the common no-OOM path. This PR replaces the per-row retry framework usage with a single per-batch retry block and a lightweight catch-based OOM recovery: **Design: per-batch retry block + inline OOM handling** The conversion loop runs inside `RmmRapidsRetryIterator.withRetryBlock`, which enters the RMM retry block once per batch (not per row). A per-row `captureState()` / `restoreState()` on the column builders enables rollback on OOM — these are lightweight offset snapshots with negligible overhead. On OOM, a single catch arm handles all cases: | Scenario | Action | |----------|--------| | RetryOOM/SplitAndRetryOOM with rows already converted (non-RequireSingleBatch) | Emit partial batch, save failed row as `pendingRow` for next batch | | RetryOOM with no rows yet (or RequireSingleBatch) | `blockUntilMemoryFreed()` — wait for spill, then while-loop retries naturally | | SplitAndRetryOOM with no rows yet | Propagate — can't split a single row | This is an **optimistic strategy**: pay minimal overhead (one `captureState` per row) in the common case, and only block/wait when absolutely necessary. **New retry framework utilities** Two general-purpose methods are added to `RmmRapidsRetryIterator` for incremental operations where partial progress is valuable (doesn't fit the standard atomic-retry model of `withRetryNoSplit`): - `withRetryBlock[T](fn: => T): T` — manages the retry block lifecycle without automatic retry - `blockUntilMemoryFreed()` — follows the standard protocol (exit retry block → `blockThreadUntilReady` → re-enter) for use within `withRetryBlock` **Other changes** - `ENABLE_R2C_RETRY` config default flipped from `false` to `true` — retry is now on by default with negligible overhead. The config is retained as an internal kill-switch to disable retry if needed. - Removed `RetryableRowConverter` class — no longer needed without per-row `withRetryNoSplit` **Performance** [Benchmark script](#14368 (comment)) | Configuration | Median | vs no-retry | |---------------|--------|-------------| | No retry | 829 ms | baseline | | Per-row retry (old) | 15,526 ms | **~19x slower** | | Per-batch retry (this PR) | 876 ms | **~5% slower** | ### 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. - `test simple GPU/CPU OOM retry` — OOM during conversion with RequireSingleBatch - `test CPU OOM retry preserves all rows for non-RequireSingleBatch` — emit-early path - `test first-row CPU OOM with TargetSize/RequireSingleBatch falls back to retry` — blockUntilMemoryFreed path - `test CPU SplitAndRetryOOM emit-early for non-RequireSingleBatch` — SplitAndRetryOOM emit-early path - `test simple OOM split and retry` — SplitAndRetryOOM propagation - [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: Haoyang Li <haoyangl@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuRowToColumnarExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuRowToColumnarExec.scala
@@ -16,10 +16,9 @@
 
 package com.nvidia.spark.rapids
 
-import com.nvidia.spark.Retryable
 import com.nvidia.spark.rapids.Arm.withResource
 import com.nvidia.spark.rapids.GpuColumnVector.GpuColumnarBatchBuilder
-import com.nvidia.spark.rapids.RmmRapidsRetryIterator.{withRestoreOnRetry, withRetryNoSplit}
+import com.nvidia.spark.rapids.jni.{CpuRetryOOM, CpuSplitAndRetryOOM, GpuRetryOOM, GpuSplitAndRetryOOM}
 import com.nvidia.spark.rapids.shims.{CudfUnsafeRow, GpuTypeShims, ShimUnaryExecNode}
 
 import org.apache.spark.TaskContext
@@ -578,7 +577,7 @@ class RowToColumnarIterator(
     localGoal: CoalesceSizeGoal,
     batchSizeBytes: Long,
     converters: GpuRowToColumnConverter,
-    enableRetry: Boolean = false,
+    enableRetry: Boolean = true,
     numInputRows: GpuMetric = NoopMetric,
     numOutputRows: GpuMetric = NoopMetric,
     numOutputBatches: GpuMetric = NoopMetric,
@@ -592,11 +591,18 @@ class RowToColumnarIterator(
   private var totalOutputRows: Long = 0
   private lazy val rowCopyProjection: UnsafeProjection = UnsafeProjection.create(localSchema)
 
+  // Carries a failed row across batch boundaries when OOM interrupted conversion.
+  private var pendingRow: InternalRow = _
+
   override def hasNext: Boolean = {
-    val start = System.nanoTime()
-    val result = rowIter.hasNext
-    streamTime += System.nanoTime() - start
-    result
+    if (pendingRow != null) {
+      true
+    } else {
+      val start = System.nanoTime()
+      val result = rowIter.hasNext
+      streamTime += System.nanoTime() - start
+      result
+    }
   }
 
   override def next(): ColumnarBatch = {
@@ -606,56 +612,110 @@ class RowToColumnarIterator(
     buildBatch()
   }
 
-  private def buildBatch(): ColumnarBatch = {
-    NvtxRegistry.ROW_TO_COLUMNAR {
-      val streamStart = System.nanoTime()
-      // estimate the size of the first batch based on the schema
-      if (targetRows == 0) {
-        if (localSchema.fields.isEmpty) {
-          // if there are no columns then we just default to a small number
-          // of rows for the first batch
-          targetRows = 1024
-          initialRows = targetRows
-        } else {
-          val sampleRows = GpuBatchUtils.VALIDITY_BUFFER_BOUNDARY_ROWS
-          val sampleBytes = GpuBatchUtils.estimateGpuMemory(localSchema, sampleRows)
-          targetRows = GpuBatchUtils.estimateRowCount(targetSizeBytes, sampleBytes, sampleRows)
-          initialRows = GpuBatchUtils.estimateRowCount(batchSizeBytes, sampleBytes, sampleRows)
-        }
-      }
+  private def copyRow(row: InternalRow): InternalRow = row match {
+    case unsafe: UnsafeRow => unsafe.copy()
+    case other => rowCopyProjection.apply(other).copy()
+  }
 
-      withResource(new GpuColumnarBatchBuilder(localSchema, initialRows)) { builders =>
-        var rowCount = 0
-        // Double because validity can be < 1 byte, and this is just an estimate anyways
-        var byteCount: Double = 0
-
-        if (enableRetry) {
-          val converter = new RetryableRowConverter(builders, rowCopyProjection)
-          // read at least one row
-          while (rowIter.hasNext &&
-              (rowCount == 0 || rowCount < targetRows && byteCount < targetSizeBytes)) {
-            converter.attempt(rowIter.next())
-            val bytesWritten = withRetryNoSplit {
-              withRestoreOnRetry(converter) {
-                converters.convert(converter.currentRow, builders)
-              }
-            }
-            byteCount += bytesWritten
-            rowCount += 1
-          }
-        } else {
-          // Disabling R2C retry only removes the per-row retry wrapper around convert().
-          // The final builders.tryBuild(rowCount) call below still uses withRetryNoSplit.
-          while (rowIter.hasNext &&
-              (rowCount == 0 || rowCount < targetRows && byteCount < targetSizeBytes)) {
-            val row = rowIter.next()
+  /** Consume the pending row if available, otherwise advance the iterator. */
+  private def nextRow(): InternalRow = {
+    if (pendingRow != null) {
+      val r = pendingRow
+      pendingRow = null
+      r
+    } else {
+      rowIter.next()
+    }
+  }
+
+  /**
+   * Core conversion loop. When retry is enabled, runs inside a single RMM retry block
+   * so host allocations throw retryable OOMs instead of fatal OutOfMemoryError.
+   *
+   * OOM strategy (retry enabled):
+   *  - Has rows to emit (and not RequireSingleBatch) → emit partial batch, save failed row
+   *  - No rows yet + RetryOOM → block until memory freed, then while loop retries
+   *  - No rows yet + SplitAndRetryOOM → propagate (can't split a single row)
+   */
+  private def convertRows(builders: GpuColumnarBatchBuilder): (Int, Double) = {
+    var rowCount = 0
+    var byteCount: Double = 0
+
+    if (enableRetry) {
+      var batchDone = false
+      RmmRapidsRetryIterator.withRetryBlock {
+        while (!batchDone && hasNext &&
+            (rowCount == 0 || rowCount < targetRows && byteCount < targetSizeBytes)) {
+          val snapshots = builders.captureState()
+          var row: InternalRow = null
+          try {
+            row = nextRow()
             byteCount += converters.convert(row, builders)
             rowCount += 1
+          } catch {
+            case oom @ (_: CpuRetryOOM | _: CpuSplitAndRetryOOM |
+                        _: GpuRetryOOM | _: GpuSplitAndRetryOOM) =>
+              builders.restoreState(snapshots)
+              if (rowCount > 0 && !localGoal.isInstanceOf[RequireSingleBatchLike]) {
+                // Emit partial batch. This also handles SplitAndRetryOOM: emitting
+                // a smaller batch IS the right response to memory pressure, and
+                // tryBuild() has its own withRetryNoSplit to handle GPU OOM.
+                if (row != null) {
+                  pendingRow = copyRow(row)
+                }
+                batchDone = true
+              } else {
+                // No rows to emit — must retry or fail.
+                oom match {
+                  case _: CpuSplitAndRetryOOM | _: GpuSplitAndRetryOOM => throw oom
+                  case _ =>
+                    if (row != null) {
+                      pendingRow = copyRow(row)
+                    }
+                    RmmRapidsRetryIterator.blockUntilMemoryFreed()
+                }
+              }
           }
         }
+      }
+    } else {
+      while (hasNext && (rowCount == 0 || rowCount < targetRows && byteCount < targetSizeBytes)) {
+        val row = nextRow()
+        byteCount += converters.convert(row, builders)
+        rowCount += 1
+      }
+    }
+    (rowCount, byteCount)
+  }
+
+  private def estimateInitialTargetRows(): Unit = {
+    // estimate the size of the first batch based on the schema
+    if (targetRows == 0) {
+      if (localSchema.fields.isEmpty) {
+        // if there are no columns then we just default to a small number
+        // of rows for the first batch
+        targetRows = 1024
+        initialRows = targetRows
+      } else {
+        val sampleRows = GpuBatchUtils.VALIDITY_BUFFER_BOUNDARY_ROWS
+        val sampleBytes = GpuBatchUtils.estimateGpuMemory(localSchema, sampleRows)
+        targetRows = GpuBatchUtils.estimateRowCount(targetSizeBytes, sampleBytes, sampleRows)
+        initialRows = GpuBatchUtils.estimateRowCount(batchSizeBytes, sampleBytes, sampleRows)
+      }
+    }
+  }
+
+  private def buildBatch(): ColumnarBatch = {
+    NvtxRegistry.ROW_TO_COLUMNAR {
+      val streamStart = System.nanoTime()
+      estimateInitialTargetRows()
+
+      withResource(new GpuColumnarBatchBuilder(localSchema, initialRows)) { builders =>
+        val (rowCount, _) = convertRows(builders)
 
         // enforce RequireSingleBatch limit
-        if (rowIter.hasNext && localGoal.isInstanceOf[RequireSingleBatchLike]) {
+        if ((pendingRow != null || rowIter.hasNext) &&
+            localGoal.isInstanceOf[RequireSingleBatchLike]) {
           throw new IllegalStateException("A single batch is required for this operation." +
               " Please try increasing your partition count.")
         }
@@ -693,84 +753,6 @@ class RowToColumnarIterator(
 
 }
 
-/**
- * A retryable row converter that integrates with the retry framework via `withRestoreOnRetry`.
- *
- * The design defers expensive operations to when OOM actually occurs:
- * - Builder state is captured lazily on first `currentRow` access (before convert runs)
- * - Row copying is deferred to `restore()` - only happens if OOM occurs
- * - `checkpoint()` is a no-op since we capture state lazily
- *
- * This class is designed to be reused across multiple rows to minimize object allocation
- * overhead. Call `attempt()` with a new row before each conversion attempt.
- *
- * This minimizes overhead in the common case (no OOM) while still enabling proper
- * rollback when OOM does occur.
- */
-private class RetryableRowConverter(
-    builders: GpuColumnarBatchBuilder,
-    projection: UnsafeProjection)
-  extends Retryable {
-
-  // The current row being converted - set via attempt()
-  private var initialRow: InternalRow = _
-
-  // Builder state - captured once on first call to ensureSnapshotCaptured()
-  private var builderSnapshots: Array[RapidsHostColumnBuilder.BuilderSnapshot] = _
-  private var snapshotCaptured: Boolean = false
-
-  // Row snapshot - only created when restore() is called (i.e., OOM happened)
-  private var rowSnapshot: Option[UnsafeRow] = None
-
-  /**
-   * Attempt the conversion for a new row. This must be called before each row conversion.
-   */
-  def attempt(row: InternalRow): Unit = {
-    initialRow = row
-    snapshotCaptured = false
-    rowSnapshot = None
-  }
-
-  /**
-   * Ensures builder state is captured exactly once, before conversion starts.
-   * This must be called before convert() to enable rollback on OOM.
-   */
-  private def ensureSnapshotCaptured(): Unit = {
-    if (!snapshotCaptured) {
-      builderSnapshots = builders.captureState()
-      snapshotCaptured = true
-    }
-  }
-
-  /**
-   * The row to convert. Captures builder state on first access.
-   * After restore(), returns the copied row for retry.
-   */
-  def currentRow: InternalRow = {
-    ensureSnapshotCaptured()
-    rowSnapshot.getOrElse(initialRow)
-  }
-
-  override def checkpoint(): Unit = ()
-
-  /**
-   * Called by withRestoreOnRetry when an OOM occurs.
-   * Copies the row (for retry) and rolls back the builders to pre-conversion state.
-   */
-  override def restore(): Unit = {
-    // Snapshot the row for the retry attempt (only on first restore)
-    if (rowSnapshot.isEmpty) {
-      rowSnapshot = Some(initialRow match {
-        case unsafe: UnsafeRow => unsafe.copy()
-        case other => projection.apply(other).copy()
-      })
-    }
-    // Roll back builders to state before this row's conversion started
-    builders.restoreState(builderSnapshots)
-  }
-
-}
-
 object GeneratedInternalRowToCudfRowIterator extends Logging {
   def apply(input: Iterator[InternalRow],
       schema: Array[Attribute],
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala
@@ -439,13 +439,12 @@ object RapidsConf extends Logging {
       .createWithDefault(2)
 
   val ENABLE_R2C_RETRY = conf("spark.rapids.sql.rowToColumnar.retry.enabled")
-    .doc("When true, the row-to-columnar conversion wraps each row's conversion with " +
-      "retry logic so that host OOM during conversion can be recovered. This adds a small " +
-      "per-row overhead. When false (default), the retry is disabled to avoid that overhead, " +
-      "at the risk of failing the task on host OOM during R2C conversion.")
+    .doc("When true (default), the row-to-columnar conversion uses a per-batch retry block " +
+      "so that host OOM during conversion can be recovered with negligible overhead. " +
+      "Set to false to disable retry and let host OOM fail the task immediately.")
     .internal()
     .booleanConf
-    .createWithDefault(false)
+    .createWithDefault(true)
 
   val GPU_COREDUMP_DIR = conf("spark.rapids.gpu.coreDump.dir")
     .doc("The URI to a directory where a GPU core dump will be created if the GPU encounters " +
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
@@ -188,6 +188,44 @@ object RmmRapidsRetryIterator extends Logging {
       new RmmRapidsRetryAutoCloseableIterator(attemptIter))
   }
 
+  /**
+   * Execute `fn` inside a retry block where host/GPU allocations throw retryable
+   * OOM exceptions instead of fatal errors. Unlike `withRetryNoSplit`, this does
+   * NOT automatically retry — the caller manages its own retry logic within `fn`.
+   *
+   * Use this for incremental operations (e.g. row-by-row accumulation) where
+   * partial progress is valuable and the standard atomic-retry model doesn't fit.
+   *
+   * @param fn the work to perform inside the retry block
+   * @tparam T result type
+   * @return the result of `fn`
+   */
+  def withRetryBlock[T](fn: => T): T = {
+    RmmSpark.currentThreadStartRetryBlock()
+    try {
+      fn
+    } finally {
+      RmmSpark.currentThreadEndRetryBlock()
+    }
+  }
+
+  /**
+   * Block the current thread until memory is freed, following the standard
+   * protocol of exiting and re-entering the retry block around the blocking call.
+   *
+   * Must be called from within an active retry block (i.e. inside `withRetryBlock`).
+   * After this call returns, the retry block is re-entered and the caller can retry
+   * the failed operation.
+   */
+  def blockUntilMemoryFreed(): Unit = {
+    RmmSpark.currentThreadEndRetryBlock()
+    try {
+      RmmSpark.blockThreadUntilReady()
+    } finally {
+      RmmSpark.currentThreadStartRetryBlock()
+    }
+  }
+
   /**
    * Returns a tuple of (shouldRetry, shouldSplit, isFromGpuOom) depending the exception
    * passed
diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/RowToColumnarIteratorRetrySuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/RowToColumnarIteratorRetrySuite.scala
