diff --git a/delta-lake/common/src/main/databricks/scala/org/apache/spark/sql/rapids/delta/GpuOptimizeWriteExchangeExec.scala b/delta-lake/common/src/main/databricks/scala/org/apache/spark/sql/rapids/delta/GpuOptimizeWriteExchangeExec.scala index c3c171ffc5d..2444a39e938 100644 --- a/delta-lake/common/src/main/databricks/scala/org/apache/spark/sql/rapids/delta/GpuOptimizeWriteExchangeExec.scala +++ b/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, diff --git a/delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/GpuDeltaParquetFileFormatBase2.scala b/delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/GpuDeltaParquetFileFormatBase2.scala index 8381b49fba2..62c89b79dbc 100644 --- a/delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/GpuDeltaParquetFileFormatBase2.scala +++ b/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. diff --git a/delta-lake/common/src/main/delta-33x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala b/delta-lake/common/src/main/delta-33x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala index 5640376de72..75dc50438cb 100644 --- a/delta-lake/common/src/main/delta-33x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala +++ b/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, diff --git a/delta-lake/common/src/main/delta-40x-42x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala b/delta-lake/common/src/main/delta-40x-42x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala index 58025a47b2c..aca001f2575 100644 --- a/delta-lake/common/src/main/delta-40x-42x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala +++ b/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) diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatNativeDV.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatNativeDV.scala index e34764fab5c..3a5b73cb5f0 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatNativeDV.scala +++ b/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) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala index ffa84c5d5d7..47084c1a06e 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCoalesceBatches.scala +++ b/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 diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuDataProducer.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuDataProducer.scala index f6fc86b870d..5781c74e048 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuDataProducer.scala +++ b/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() +} 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 index 301f7379f5d..4b1046fcee7 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala +++ b/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.") diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/parquet/GpuParquetScan.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/parquet/GpuParquetScan.scala index 2c080d49da3..29f2c71efbd 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/parquet/GpuParquetScan.scala +++ b/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:_*) ) diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuShuffleExchangeExecBase.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuShuffleExchangeExecBase.scala index 5714c8ef44d..5b4fc3deb1b 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuShuffleExchangeExecBase.scala +++ b/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,6 +462,8 @@ 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() @@ -463,11 +471,11 @@ object GpuShuffleExchangeExecBase { 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. diff --git a/sql-plugin/src/test/scala/com/nvidia/spark/rapids/RangeInputBatchingSuite.scala b/sql-plugin/src/test/scala/com/nvidia/spark/rapids/RangeInputBatchingSuite.scala new file mode 100644 index 00000000000..060fc02b20b --- /dev/null +++ b/sql-plugin/src/test/scala/com/nvidia/spark/rapids/RangeInputBatchingSuite.scala @@ -0,0 +1,198 @@ +/* + * Copyright (c) 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. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids + +import scala.collection.mutable.ArrayBuffer + +import ai.rapids.cudf.Table +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.vectorized.ColumnarBatch + +class RangeInputBatchingSuite extends AnyFunSuite { + test("range input batching config is disabled by default and can be enabled") { + val entry = RapidsConf.RANGE_SHUFFLE_INPUT_BATCHING_ENABLED + assert(!entry.defaultValue) + assert(!entry.get(Map.empty[String, String])) + assert(entry.get(Map(entry.key -> "true"))) + } + + private class TestCoalesceIterator( + input: Iterator[ColumnarBatch], + goal: CoalesceSizeGoal = TargetSize(Long.MaxValue)) + extends AbstractGpuCoalesceIterator( + input, + goal, + NoopMetric, + NoopMetric, + NoopMetric, + NoopMetric, + NoopMetric, + NoopMetric, + NoopMetric, + "test") { + private val candidates = new ArrayBuffer[ColumnarBatch] + private var onDeck: ColumnarBatch = _ + + def collectCandidates(): Boolean = populateCandidateBatches() + def candidateCount: Int = candidates.size + def closeCandidates(): Unit = candidates.foreach(_.close()) + + override protected def hasOnDeck: Boolean = onDeck != null + override protected def saveOnDeck(batch: ColumnarBatch): Unit = onDeck = batch + override protected def clearOnDeck(): Unit = { + if (onDeck != null) onDeck.close() + onDeck = null + } + override protected def popOnDeck(): ColumnarBatch = { + val result = onDeck + onDeck = null + result + } + override def initNewBatch(batch: ColumnarBatch): Unit = candidates.clear() + override def addBatchToConcat(batch: ColumnarBatch): Unit = candidates += batch + override def hasAnyToConcat: Boolean = candidates.nonEmpty + override def concatAllAndPutOnGPU(): ColumnarBatch = candidates.remove(0) + override protected val supportsRetryIterator: Boolean = false + override def getCoalesceRetryIterator: Iterator[ColumnarBatch] = Iterator.empty + override def cleanupConcatIsDone(): Unit = candidates.clear() + } + + test("range input does not read ahead beyond the first candidate batch") { + var reads = 0 + val input = Seq( + new ColumnarBatch(Array.empty, 1), + new ColumnarBatch(Array.empty, 1), + new ColumnarBatch(Array.empty, 1)).iterator.map { batch => + reads += 1 + batch + } + val coalesce = new TestCoalesceIterator(input) + + try { + val isLast = RangeInputBatching.withRangeInput(enabled = true) { + assert(coalesce.hasNext) + coalesce.collectCandidates() + } + assert(!isLast) + assert(reads == 1) + assert(coalesce.candidateCount == 1) + } finally { + coalesce.closeCandidates() + input.foreach(_.close()) + } + } + + test("normal input retains size-based coalescing behavior") { + var reads = 0 + val input = Seq( + new ColumnarBatch(Array.empty, 1), + new ColumnarBatch(Array.empty, 1), + new ColumnarBatch(Array.empty, 1)).iterator.map { batch => + reads += 1 + batch + } + val coalesce = new TestCoalesceIterator(input) + + try { + assert(coalesce.hasNext) + assert(coalesce.collectCandidates()) + assert(reads == 3) + assert(coalesce.candidateCount == 3) + } finally { + coalesce.closeCandidates() + input.foreach(_.close()) + } + } + + Seq[CoalesceSizeGoal](RequireSingleBatch, RequireSingleBatchWithFilter(null)).foreach { goal => + test(s"range input retains $goal behavior") { + var reads = 0 + val input = Seq( + new ColumnarBatch(Array.empty, 1), + new ColumnarBatch(Array.empty, 1), + new ColumnarBatch(Array.empty, 1)).iterator.map { batch => + reads += 1 + batch + } + val coalesce = new TestCoalesceIterator(input, goal) + + try { + val isLast = RangeInputBatching.withRangeInput(enabled = true) { + assert(coalesce.hasNext) + coalesce.collectCandidates() + } + assert(isLast) + assert(reads == 3) + assert(coalesce.candidateCount == 3) + } finally { + coalesce.closeCandidates() + input.foreach(_.close()) + } + } + } + + test("range input marker restores its scope after nesting and exceptions") { + assert(!RangeInputBatching.isActive) + RangeInputBatching.withRangeInput(enabled = false) { + assert(!RangeInputBatching.isActive) + } + intercept[RuntimeException] { + RangeInputBatching.withRangeInput(enabled = true) { + assert(RangeInputBatching.isActive) + RangeInputBatching.withRangeInput(enabled = true) { + assert(RangeInputBatching.isActive) + } + assert(RangeInputBatching.isActive) + throw new RuntimeException("test") + } + } + assert(!RangeInputBatching.isActive) + } + + test("range input selects and closes a restartable producer lazily") { + var closed = false + val producer = new RetryableTableProducer { + override def hasNext: Boolean = false + override def next: Table = throw new NoSuchElementException + override def checkpoint(): Unit = () + override def restore(): Unit = () + override def close(): Unit = closed = true + } + val iter = RangeInputBatching.withRangeInput(enabled = true) { + CachedGpuBatchIterator(producer, Array.empty[DataType]) + } + + assert(!closed) + iter.close() + assert(closed) + } + + test("range input eagerly caches a non-restartable producer") { + var closed = false + val producer = new EmptyGpuDataProducer[Table] { + override def close(): Unit = closed = true + } + val iter = RangeInputBatching.withRangeInput(enabled = true) { + CachedGpuBatchIterator(producer, Array.empty[DataType]) + } + + assert(closed) + assert(!iter.hasNext) + } +} diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuKudoWritePartitioningSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuKudoWritePartitioningSuite.scala index 0da34594ac0..135df366bb7 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuKudoWritePartitioningSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuKudoWritePartitioningSuite.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. + * Copyright (c) 2025-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. @@ -235,6 +235,7 @@ class GpuKudoWritePartitioningSuite extends AnyFunSuite with BeforeAndAfterEach serializer, useGPUShuffle = false, useMultiThreadedShuffle = false, + rangeInputBatchingEnabled = false, metrics, writeMetrics, Map.empty, diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/WithRetrySuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/WithRetrySuite.scala index f800166e526..c9659ac63cb 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/WithRetrySuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/WithRetrySuite.scala @@ -178,6 +178,53 @@ class WithRetrySuite } } + test("range producer restores native progress before retrying a deferred batch") { + var nextIndex = 0 + var checkpointIndex = 0 + var failedOnce = false + var restoreCount = 0 + var closed = false + val successfulValues = new scala.collection.mutable.ArrayBuffer[Long] + val producer = new RetryableTableProducer { + override def hasNext: Boolean = nextIndex < 2 + + override def next: Table = { + val value = nextIndex + nextIndex += 1 + if (!failedOnce) { + failedOnce = true + throw new GpuRetryOOM("in tests") + } + successfulValues += value + new Table.TestBuilder().column(Array[java.lang.Long](value.toLong): _*).build() + } + + override def checkpoint(): Unit = checkpointIndex = nextIndex + + override def restore(): Unit = { + restoreCount += 1 + nextIndex = checkpointIndex + } + + override def close(): Unit = closed = true + } + + val iter = RangeInputBatching.withRangeInput(enabled = true) { + CachedGpuBatchIterator(producer, Array[DataType](LongType)) + } + withResource(iter) { _ => + assert(iter.hasNext) + withResource(iter.next())(batch => assert(batch.numRows() == 1)) + assert(iter.hasNext) + withResource(iter.next())(batch => assert(batch.numRows() == 1)) + assert(!iter.hasNext) + } + + assert(restoreCount == 1) + assert(successfulValues == Seq(0L, 1L)) + assert(closed) + } + test("withRestoreOnRetry restores state on retry") { val initialValue = 5 val increment = 5