Avoid eager input materialization for GPU range shuffle

delta-lake/common/src/main/databricks/scala/org/apache/spark/sql/rapids/delta/GpuOptimizeWriteExchangeExec.scala

@@ -124,6 +124,7 @@ case class GpuOptimizeWriteExchangeExec(
       serializer,
       useGPUShuffle=partitioning.usesGPUShuffle,
       useMultiThreadedShuffle=partitioning.usesMultiThreadedShuffle,
+      rangeInputBatchingEnabled=false,
       metrics=allMetrics,
       writeMetrics=writeMetrics,
       additionalMetrics=additionalMetrics,

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/GpuDeltaParquetFileFormatBase2.scala

@@ -1292,13 +1292,13 @@ case class DeltaParquetTableReader(
 
   logDebug("Using DeltaParquetTableReader for reading Parquet with deletion vectors")
 
-  override protected val reader = DeltaParquetChunkedReader(
+  override protected def createReader(): ChunkedReader = DeltaParquetChunkedReader(
     DeletionVector.newParquetChunkedReader(chunkSizeByteLimit,
       maxChunkedReaderMemoryUsageSizeBytes, opts, buffers, dvInfos)
   )
 
-  override protected lazy val resources: Seq[AutoCloseable] =
-    Seq(reader) ++ buffers ++ dvInfos.map(_.serializedBitmap)
+  override protected def additionalResources: Seq[AutoCloseable] =
+    dvInfos.map(_.serializedBitmap)
 
   override protected def postProcessChunk(chunk: Table): Table = {
     // The cuDF reader prepends an extra index column in the output table.

delta-lake/common/src/main/delta-33x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala

@@ -150,6 +150,7 @@ case class GpuOptimizeWriteExchangeExec(
       serializer,
       useGPUShuffle=actualPartitioning.usesGPUShuffle,
       useMultiThreadedShuffle=actualPartitioning.usesMultiThreadedShuffle,
+      rangeInputBatchingEnabled=false,
       metrics=allMetrics,
       writeMetrics=writeMetrics,
       additionalMetrics=additionalMetrics,

delta-lake/common/src/main/delta-40x-42x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala

@@ -140,6 +140,7 @@ case class GpuOptimizeWriteExchangeExec(
       serializer,
       useGPUShuffle=actualPartitioning.usesGPUShuffle,
       useMultiThreadedShuffle=actualPartitioning.usesMultiThreadedShuffle,
+      rangeInputBatchingEnabled=false,
       metrics=allMetrics,
       writeMetrics=writeMetrics,
       additionalMetrics=additionalMetrics)

delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatNativeDV.scala

@@ -1432,13 +1432,13 @@ case class DeltaParquetTableReader(
 
   logDebug("Using DeltaParquetTableReader for reading Parquet with deletion vectors")
 
-  override protected val reader = DeltaParquetChunkedReader(
+  override protected def createReader(): ChunkedReader = DeltaParquetChunkedReader(
     DeletionVector.newParquetChunkedReader(chunkSizeByteLimit,
       maxChunkedReaderMemoryUsageSizeBytes, opts, buffers, dvInfos)
   )
 
-  override protected lazy val resources: Seq[AutoCloseable] =
-    Seq(reader) ++ buffers ++ dvInfos.map(_.serializedBitmap)
+  override protected def additionalResources: Seq[AutoCloseable] =
+    dvInfos.map(_.serializedBitmap)
 
   private lazy val deletionVectorSkipRowIndexes =
     MakeParquetTableWithDVProducer.deletionVectorSkipRowIndexes(readDataSchema)

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

@@ -257,6 +257,34 @@ object OpNameNvtxMap {
   def get(opName: String): Option[NvtxId] = map.get(opName)
 }
 
+/**
+ * Marks upstream iterator calls made while feeding a range shuffle. This is deliberately a
+ * small execution-scope marker rather than a SQL metric or plan-level setting: the same scan can
+ * be reused by other consumers, and only the range-shuffle consumer needs one-batch-at-a-time
+ * coalescing.
+ */
+object RangeInputBatching {
+  private val active = new ThreadLocal[java.lang.Boolean]()
+
+  def isActive: Boolean = active.get() == java.lang.Boolean.TRUE
+
+  def withRangeInput[T](enabled: Boolean)(body: => T): T = {
+    if (enabled) {
+      val previous = active.get()
+      active.set(java.lang.Boolean.TRUE)
+      try body finally {
+        if (previous == null) {
+          active.remove()
+        } else {
+          active.set(previous)
+        }
+      }
+    } else {
+      body
+    }
+  }
+}
+
 abstract class AbstractGpuCoalesceIterator(
     inputIter: Iterator[ColumnarBatch],
     goal: CoalesceSizeGoal,
@@ -470,7 +498,12 @@ abstract class AbstractGpuCoalesceIterator(
     }
 
     // there is a hard limit of 2^31 rows
-    while (numRows < filteringModeRowsThreshold && !hasOnDeck && iter.hasNext) {
+    // A range shuffle consumes every splittable, size-based input batch independently. Avoid
+    // reading and retaining the next wide batch while the current range-shuffle batch is still
+    // live. Single-batch goals must continue reading the complete partition.
+    while (numRows < filteringModeRowsThreshold && !hasOnDeck &&
+        !(RangeInputBatching.isActive && goal.isInstanceOf[SplittableGoal] && hasAnyToConcat) &&
+        iter.hasNext) {
       val cbFromIter = iter.next()
       numInputBatches += 1
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-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.
@@ -19,8 +19,10 @@ package com.nvidia.spark.rapids
 import scala.collection.mutable
 
 import ai.rapids.cudf.Table
+import com.nvidia.spark.Retryable
 import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
 
+import org.apache.spark.TaskContext
 import org.apache.spark.sql.types.DataType
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
@@ -59,6 +61,13 @@ trait GpuDataProducer[T] extends AutoCloseable {
   }
 }
 
+/**
+ * A table producer whose state can be checkpointed and restored across an RMM retry.
+ * Implementations must keep their inputs alive until close and reproduce the next table after
+ * restore without skipping or duplicating previously returned data.
+ */
+private[rapids] trait RetryableTableProducer extends GpuDataProducer[Table] with Retryable
+
 object GpuDataProducer {
   /**
    * Essentially the same as doing a map on a regular iterator, but the resulting GpuDataProducer
@@ -153,6 +162,16 @@ object CachedGpuBatchIterator {
 
   def apply(producer: GpuDataProducer[Table],
       dataTypes: Array[DataType]): GpuColumnarBatchIterator = {
+    producer match {
+      case retryable: RetryableTableProducer if RangeInputBatching.isActive =>
+        new RangeGpuDataProducerIterator(retryable, dataTypes)
+      case _ =>
+        cacheProducer(producer, dataTypes)
+    }
+  }
+
+  private def cacheProducer(producer: GpuDataProducer[Table],
+      dataTypes: Array[DataType]): GpuColumnarBatchIterator = {
     withResource(producer) { _ =>
       if (producer.hasNext) {
         // Special case for the first one.
@@ -177,3 +196,46 @@ object CachedGpuBatchIterator {
     }
   }
 }
+
+/**
+ * Streams a restartable GPU table producer one batch at a time into a range shuffle.
+ *
+ * CachedGpuBatchIterator normally drains a chunked file reader eagerly so the producer can be
+ * closed before the GPU semaphore is released. A range shuffle consumes its input synchronously,
+ * and draining a wide reader there materializes several decoded batches before any of them can be
+ * partitioned. The restartable producer keeps native progress retry-safe while this
+ * iterator bounds live decoded data to the batch currently being partitioned.
+ */
+private class RangeGpuDataProducerIterator(
+    producer: RetryableTableProducer,
+    dataTypes: Array[DataType]) extends GpuColumnarBatchIterator(true) {
+
+  private def retry[T](body: => T): T = {
+    producer.checkpoint()
+    RmmRapidsRetryIterator.withRetryNoSplit {
+      RmmRapidsRetryIterator.withRestoreOnRetry(producer)(body)
+    }
+  }
+
+  override def hasNext: Boolean = closeOnExcept(this) { _ =>
+    val more = retry {
+      GpuSemaphore.acquireIfNecessary(TaskContext.get())
+      producer.hasNext
+    }
+    if (!more) {
+      close()
+    }
+    more
+  }
+
+  override def next(): ColumnarBatch = closeOnExcept(this) { _ =>
+    retry {
+      GpuSemaphore.acquireIfNecessary(TaskContext.get())
+      withResource(producer.next) { table =>
+        GpuColumnVector.from(table, dataTypes)
+      }
+    }
+  }
+
+  override def doClose(): Unit = producer.close()
+}

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

@@ -770,6 +770,14 @@ val GPU_COREDUMP_PIPE_PATTERN = conf("spark.rapids.gpu.coreDump.pipePattern")
     .booleanConf
     .createWithDefault(false)
 
+  val RANGE_SHUFFLE_INPUT_BATCHING_ENABLED =
+    conf("spark.rapids.sql.rangeShuffle.inputBatching.enabled")
+      .doc("Enables experimental one-input-batch-at-a-time consumption for GPU range shuffles " +
+        "to bound the amount of decoded input retained before partitioning.")
+      .internal()
+      .booleanConf
+      .createWithDefault(false)
+
   val EXPORT_COLUMNAR_RDD = conf("spark.rapids.sql.exportColumnarRdd")
     .doc("Spark has no simply way to export columnar RDD data.  This turns on special " +
       "processing/tagging that allows the RDD to be picked back apart into a Columnar RDD.")

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

@@ -3566,22 +3566,22 @@ abstract class AbstractParquetTableReader(
     clippedParquetSchema: MessageType,
     splits: Array[PartitionedFile],
     debugDumpPrefix: Option[String],
-    debugDumpAlways: Boolean) extends GpuDataProducer[Table] with Logging {
+    debugDumpAlways: Boolean) extends RetryableTableProducer with Logging {
 
-  protected val reader: ChunkedReader
+  protected def createReader(): ChunkedReader
 
-  private[this] lazy val splitsString = splits.mkString("; ")
-
-  // Should be lazy since the reader is not defined. Otherwise in practise, a native
-  // chunk reader will be leaked.
-  protected lazy val resources: Seq[AutoCloseable] = Seq(reader) ++ buffers
+  protected def additionalResources: Seq[AutoCloseable] = Seq.empty
 
-  override def hasNext: Boolean = reader.hasNext
+  private[this] lazy val splitsString = splits.mkString("; ")
+  private var activeReader: ChunkedReader = _
+  private var completedChunks = 0
+  private var checkpointedChunks = 0
+  private var closed = false
 
   protected def postProcessChunk(chunk: Table): Table
 
-  override def next: Table = {
-    val table = NvtxIdWithMetrics(NvtxRegistry.PARQUET_DECODE, metrics(GPU_DECODE_TIME)) {
+  private def decodeNext(reader: ChunkedReader): Table = {
+    NvtxIdWithMetrics(NvtxRegistry.PARQUET_DECODE, metrics(GPU_DECODE_TIME)) {
       try {
         reader.next
       } catch {
@@ -3597,9 +3597,11 @@ abstract class AbstractParquetTableReader(
           throw new IOException(s"Error when processing $splitsString$dumpMsg", e)
       }
     }
+  }
 
-    val postProcessedTable = postProcessChunk(table)
-
+  private def readNext(reader: ChunkedReader): Table = {
+    val table = decodeNext(reader)
+    val postProcessedTable = closeOnExcept(table)(postProcessChunk)
     closeOnExcept(postProcessedTable) { _ =>
       GpuParquetScan.throwIfRebaseNeededInExceptionMode(postProcessedTable, dateRebaseMode,
         timestampRebaseMode)
@@ -3617,8 +3619,54 @@ abstract class AbstractParquetTableReader(
     outputTable
   }
 
+  private def closeReader(): Unit = {
+    val reader = activeReader
+    activeReader = null
+    if (reader != null) {
+      reader.close()
+    }
+  }
+
+  private def getReader: ChunkedReader = {
+    if (activeReader == null) {
+      require(!closed, "Parquet table reader is closed")
+      val reader = createReader()
+      closeOnExcept(reader) { _ =>
+        var replayed = 0
+        while (replayed < completedChunks) {
+          require(reader.hasNext,
+            s"Unable to restore Parquet reader to chunk $completedChunks")
+          withResource(decodeNext(reader))(_ => ())
+          replayed += 1
+        }
+        activeReader = reader
+      }
+    }
+    activeReader
+  }
+
+  override def hasNext: Boolean = getReader.hasNext
+
+  override def next: Table = {
+    val result = readNext(getReader)
+    completedChunks += 1
+    result
+  }
+
+  override def checkpoint(): Unit = checkpointedChunks = completedChunks
+
+  override def restore(): Unit = {
+    completedChunks = checkpointedChunks
+    closeReader()
+  }
+
   override def close(): Unit = {
-    resources.safeClose()
+    if (!closed) {
+      closed = true
+      val reader = Option(activeReader).toSeq
+      activeReader = null
+      (reader ++ buffers ++ additionalResources).safeClose()
+    }
   }
 }
 
@@ -3642,7 +3690,7 @@ case class ParquetTableReader(
   opts, buffers, metrics, dateRebaseMode, timestampRebaseMode, isSchemaCaseSensitive, useFieldId,
   readDataSchema, clippedParquetSchema, splits, debugDumpPrefix, debugDumpAlways) {
 
-  override protected val reader: ChunkedReader = ParquetChunkedReader(
+  override protected def createReader(): ChunkedReader = ParquetChunkedReader(
     new JniParquetChunkedReader(chunkSizeByteLimit, maxChunkedReaderMemoryUsageSizeBytes,
       opts, buffers:_*)
   )

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

@@ -182,6 +182,8 @@ abstract class GpuShuffleExchangeExecBase(
   private lazy val kudoBufferCopyMeasurementEnabled = RapidsConf
     .SHUFFLE_KUDO_SERIALIZER_MEASURE_BUFFER_COPY_ENABLED
     .get(child.conf)
+  private lazy val rangeInputBatchingEnabled = RapidsConf
+    .RANGE_SHUFFLE_INPUT_BATCHING_ENABLED.get(child.conf)
 
   private lazy val useGPUShuffle = {
     gpuOutputPartitioning match {
@@ -270,6 +272,7 @@ abstract class GpuShuffleExchangeExecBase(
       serializer,
       useGPUShuffle,
       useMultiThreadedShuffle,
+      rangeInputBatchingEnabled,
       allMetrics,
       writeMetrics,
       additionalMetrics,
@@ -400,6 +403,7 @@ object GpuShuffleExchangeExecBase {
       serializer: Serializer,
       useGPUShuffle: Boolean,
       useMultiThreadedShuffle: Boolean,
+      rangeInputBatchingEnabled: Boolean,
       metrics: Map[String, GpuMetric],
       writeMetrics: Map[String, SQLMetric],
       additionalMetrics: Map[String, GpuMetric],
@@ -435,6 +439,8 @@ object GpuShuffleExchangeExecBase {
     }
     val partitioner: GpuExpression = getPartitioner(newRdd, outputAttributes,
       newPartitioning, metrics)
+    val useRangeInputBatching = rangeInputBatchingEnabled &&
+      newPartitioning.isInstanceOf[GpuRangePartitioning]
     // Inject debugging subMetrics, such as D2HTime before SliceOnCpu
     // The injected metrics will be serialized as the members of GpuPartitioning
     partitioner match {
@@ -456,18 +462,20 @@ object GpuShuffleExchangeExecBase {
           private var partitioned: Array[(ColumnarBatch, Int)] = _
           private var at = 0
           private val mutablePair = new MutablePair[Int, ColumnarBatch]()
+          private def rangeInput[T](body: => T): T =
+            RangeInputBatching.withRangeInput(useRangeInputBatching)(body)
           private def partNextBatch(): Unit = {
             if (partitioned != null) {
               partitioned.map(_._1).safeClose()
               partitioned = null
               at = 0
             }
             // Try to fill partitionedIter from iter if it's empty
-            if (!partitionedIter.hasNext && iter.hasNext) {
-              var batch = iter.next()
-              while (batch.numRows == 0 && iter.hasNext) {
+            if (!partitionedIter.hasNext && rangeInput(iter.hasNext)) {
+              var batch = rangeInput(iter.next())
+              while (batch.numRows == 0 && rangeInput(iter.hasNext)) {
                 batch.close()
-                batch = iter.next()
+                batch = rangeInput(iter.next())
               }
               // Get a non-empty batch or the last batch. So still need to
               // check if it is empty for the later case.