rapids shuffle manager V2 phase 1: writer use as much memory as allowed and pipelined write

integration_tests/run_pyspark_from_build.sh

@@ -57,6 +57,7 @@
 # ============================================================================
 
 set -ex
+PS4='+${LINENO}: '
 
 SCRIPTPATH="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
 cd "$SCRIPTPATH"
@@ -271,7 +272,7 @@ else
 
     REPORT_CHARS=${REPORT_CHARS:="fE"} # default as (f)ailed, (E)rror
     STD_INPUT_PATH="$INPUT_PATH"/src/test/resources
-    TEST_COMMON_OPTS=(-v
+    TEST_COMMON_OPTS=(-v -x
           -r"$REPORT_CHARS"
           "$TEST_TAGS"
           --std_input_path="$STD_INPUT_PATH"
@@ -627,6 +628,11 @@ PY
         unset PYSP_TEST_spark_jars_repositories
         unset PYSP_TEST_spark_rapids_memory_gpu_allocSize
 
+
+        # Comment this out if you want to run remote debug this local mode spark process
+        # Don't forget to set TEST_PARALLEL=1 to ensure local mode spark 
+        # export SPARK_SUBMIT_OPTS="-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=5005"
+
         exec "$SPARK_HOME"/bin/spark-submit "${jarOpts[@]}" \
             --driver-java-options "$driverJavaOpts" \
             $SPARK_SUBMIT_FLAGS \

sql-plugin/src/main/java/com/nvidia/spark/rapids/InternalRowToColumnarBatchIterator.java

@@ -228,7 +228,7 @@ private HostMemoryBuffer[] getHostBuffersWithRetry(
     try {
       hBuf = HostAlloc$.MODULE$.alloc((dataBytes + offsetBytes),true);
       SpillableHostBuffer sBuf = SpillableHostBuffer$.MODULE$.apply(hBuf, hBuf.getLength(),
-          SpillPriorities$.MODULE$.ACTIVE_ON_DECK_PRIORITY());
+          Option.empty());
       hBuf = null;  // taken over by spillable host buffer
       return Tuple2.apply(sBuf, numRowsWrapper);
     } finally {

sql-plugin/src/main/java/com/nvidia/spark/rapids/RapidsHostColumnVectorCore.java

@@ -1,4 +1,5 @@
 
+
 /*
  * Copyright (c) 2020-2021, NVIDIA CORPORATION.
  *

sql-plugin/src/main/java/com/nvidia/spark/rapids/SlicedGpuColumnVector.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2025, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

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

@@ -312,7 +312,7 @@ trait GpuExec extends SparkPlan with Logging {
 
   // For GpuShuffleExchangeExecBase and GpuCustomShuffleReaderExec,
   // we want the op time metrics to be called:
-  // - "op time (shuffle write partition & serial)" for shuffle write, and
+  // - "op time (shuffle write partition)" for shuffle write, and
   // - "op time (shuffle read)" for shuffle read.
   // That's why we have this separate method to get the metric.
   def getOpTimeNewMetric: Option[GpuMetric] = allMetrics.get(OP_TIME_NEW)

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

@@ -153,7 +153,7 @@ object GpuMetric extends Logging {
   val DESCRIPTION_NUM_PARTITIONS = "partitions"
   val DESCRIPTION_OP_TIME_LEGACY = "op time (legacy)"
   val DESCRIPTION_OP_TIME_NEW = "op time"
-  val DESCRIPTION_OP_TIME_NEW_SHUFFLE_WRITE = "op time (shuffle write partition & serial)"
+  val DESCRIPTION_OP_TIME_NEW_SHUFFLE_WRITE = "op time (shuffle write partitioning)"
   val DESCRIPTION_OP_TIME_NEW_SHUFFLE_READ = "op time (shuffle read)"
   val DESCRIPTION_COLLECT_TIME = "collect batch time"
   val DESCRIPTION_CONCAT_TIME = "concat batch time"

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

@@ -1406,8 +1406,7 @@ abstract class MultiFileCoalescingPartitionReaderBase(
       logDebug(s"$getFileFormatShortName Coalescing reading estimates the initTotalSize:" +
         s" $initTotalSize, and the true size: $finalBufferSize")
       SpillableHostBuffer(finalBuffer,
-        finalBufferSize,
-        SpillPriorities.ACTIVE_BATCHING_PRIORITY)
+        finalBufferSize)
     }
   }
 

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

@@ -1096,7 +1096,7 @@ trait OrcPartitionReaderBase extends OrcCommonFunctions with Logging
       closeOnExcept(hostBuffer) { hmb =>
         withResource(new HostMemoryOutputStream(hmb)) { out =>
           writeOrcOutputFile(ctx, out, stripes)
-          (SpillableHostBuffer(hmb, out.getPos, SpillPriorities.ACTIVE_BATCHING_PRIORITY),
+          (SpillableHostBuffer(hmb, out.getPos),
             out.getPos)
         }
       }
@@ -2393,8 +2393,7 @@ class MultiFileCloudOrcPartitionReader(
 
           // e: Create the new meta for the combined buffer
           val numRows = combinedMeta.allPartValues.map(_._1).sum
-          val finalBuf = SpillableHostBuffer(maybeNewBuf, maybeNewBuf.getLength,
-            SpillPriorities.ACTIVE_BATCHING_PRIORITY)
+          val finalBuf = SpillableHostBuffer(maybeNewBuf, maybeNewBuf.getLength)
           val combinedRet = SingleHMBAndMeta(Array(finalBuf), outStream.getPos, numRows,
             blockMetas)
           val newHmbWithMeta = metaToUse.copy(

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

@@ -1133,7 +1133,7 @@ sealed trait SpillableHostConcatResult extends AutoCloseable {
   def getDataLen: Long
 
   protected var spillableBuffer = {
-    SpillableHostBuffer(hmb, hmb.getLength, SpillPriorities.ACTIVE_BATCHING_PRIORITY)
+    SpillableHostBuffer(hmb, hmb.getLength)
   }
 
   override def close(): Unit = if (spillableBuffer != null) {

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

@@ -42,6 +42,16 @@ private class HostAlloc(nonPinnedLimit: Long) extends HostMemoryAllocator with L
   private val isUnlimited = nonPinnedLimit < 0
   private val isPinnedOnly = nonPinnedLimit == 0
 
+  // Expose for usage ratio calculation
+  def getCurrentAllocated: Long = synchronized {
+    currentNonPinnedAllocated + currentPinnedAllocated
+  }
+
+  def getTotalLimit: Long = {
+    if (isUnlimited) Long.MaxValue
+    else pinnedLimit + nonPinnedLimit
+  }
+
   /**
    * A callback class so we know when a non-pinned host buffer was released
    */
@@ -295,6 +305,28 @@ object HostAlloc extends Logging {
     getSingleton.alloc(amount, preferPinned)
   }
 
+  /**
+   * Get current host memory usage ratio (0.0 to 1.0).
+   * Returns current allocated / limit.
+   */
+  def getUsageRatio(): Double = {
+    val alloc = getSingleton
+    val currentAllocated = alloc.getCurrentAllocated
+    val totalLimit = alloc.getTotalLimit
+    if (totalLimit == Long.MaxValue) {
+      0.0  // Unlimited, consider as 0% used
+    } else {
+      currentAllocated.toDouble / totalLimit.toDouble
+    }
+  }
+
+  /**
+   * Check if host memory usage is below the given threshold (0.0 to 1.0).
+   */
+  def isUsageBelowThreshold(threshold: Double): Boolean = {
+    getUsageRatio() < threshold
+  }
+
   def addEventHandler(buff: HostMemoryBuffer,
                       handler: MemoryBuffer.EventHandler): HostMemoryBuffer = {
     buff.synchronized {

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

@@ -546,6 +546,41 @@ val GPU_COREDUMP_PIPE_PATTERN = conf("spark.rapids.gpu.coreDump.pipePattern")
     .bytesConf(ByteUnit.BYTE)
     .createWithDefault(-1)
 
+  val PARTIAL_FILE_BUFFER_INITIAL_SIZE = 
+    conf("spark.rapids.memory.host.partialFileBufferInitialSize")
+    .doc("The initial size in bytes for a host memory buffer used by " +
+        "SpillablePartialFileHandle. The buffer can expand dynamically up to " +
+        "partialFileBufferMaxSize. A smaller initial size reduces upfront memory " +
+        "allocation but may require more expansions.")
+    .startupOnly()
+    .internal()
+    .bytesConf(ByteUnit.BYTE)
+    .createWithDefault(1024L * 1024 * 1024)  // 1GB default
+
+  val PARTIAL_FILE_BUFFER_MAX_SIZE = 
+    conf("spark.rapids.memory.host.partialFileBufferMaxSize")
+    .doc("The maximum size in bytes for a single host memory buffer used by " +
+        "SpillablePartialFileHandle. When a buffer needs to expand beyond this limit, " +
+        "it will be spilled to disk instead. This prevents excessive memory usage " +
+        "for large shuffle partitions.")
+    .startupOnly()
+    .internal()
+    .bytesConf(ByteUnit.BYTE)
+    .createWithDefault(8L * 1024 * 1024 * 1024)  // 8GB
+
+  val PARTIAL_FILE_BUFFER_MEMORY_THRESHOLD =
+    conf("spark.rapids.memory.host.partialFileBufferMemoryThreshold")
+    .doc("The host memory usage threshold (as a fraction from 0.0 to 1.0) for deciding " +
+        "whether to use memory-based buffering for partial files. When host memory usage " +
+        "exceeds this threshold, file-based storage will be used directly. This threshold " +
+        "also applies when expanding buffers dynamically.")
+    .startupOnly()
+    .internal()
+    .doubleConf
+    .checkValue(v => v > 0.0 && v <= 1.0,
+      "The memory threshold must be in the range (0.0, 1.0]")
+    .createWithDefault(0.5)
+
   val UNSPILL = conf("spark.rapids.memory.gpu.unspill.enabled")
     .doc("When a spilled GPU buffer is needed again, should it be unspilled, or only copied " +
         "back into GPU memory temporarily. Unspilling may be useful for GPU buffers that are " +
@@ -3281,6 +3316,12 @@ class RapidsConf(conf: Map[String, String]) extends Logging {
 
   lazy val hostSpillStorageSize: Long = get(HOST_SPILL_STORAGE_SIZE)
 
+  lazy val partialFileBufferInitialSize: Long = get(PARTIAL_FILE_BUFFER_INITIAL_SIZE)
+
+  lazy val partialFileBufferMaxSize: Long = get(PARTIAL_FILE_BUFFER_MAX_SIZE)
+
+  lazy val partialFileBufferMemoryThreshold: Double = get(PARTIAL_FILE_BUFFER_MEMORY_THRESHOLD)
+
   lazy val isUnspillEnabled: Boolean = get(UNSPILL)
 
   lazy val needDecimalGuarantees: Boolean = get(NEED_DECIMAL_OVERFLOW_GUARANTEES)

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

@@ -526,17 +526,19 @@ object SpillableHostBuffer {
    *               must be <= than buffer.getLength, otherwise this function throws
    *               and closes `buffer`
    * @param buffer the buffer to make spillable
-   * @param priority the initial spill priority of this buffer
+   * @param overrideTaskPriority optional task priority to override the default task priority
    */
   def apply(buffer: HostMemoryBuffer,
             length: Long,
-            priority: Long): SpillableHostBuffer = {
+            overrideTaskPriority: Option[Long] = None): SpillableHostBuffer = {
     closeOnExcept(buffer) { _ =>
       require(length <= buffer.getLength,
         s"Attempted to add a host spillable with a length ${length} B which is " +
           s"greater than the backing host buffer length ${buffer.getLength} B")
     }
-    new SpillableHostBuffer(SpillableHostBufferHandle(buffer), length)
+    val handle = SpillableHostBufferHandle(buffer)
+    overrideTaskPriority.foreach(priority => handle.taskPriority = priority)
+    new SpillableHostBuffer(handle, length)
   }
 
   def sliceWithRetry(shb: SpillableHostBuffer, start: Long, len: Long): HostMemoryBuffer = {

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

@@ -51,8 +51,7 @@ object SpillableKudoTable {
         buffer.getLength,
         SpillableHostBuffer.apply(
           buffer,
-          buffer.getLength,
-          SpillPriorities.ACTIVE_BATCHING_PRIORITY)
+          buffer.getLength)
       )
     }
   }

sql-plugin/src/main/scala/com/nvidia/spark/rapids/jni/kudo/KudoHostMergeResultWrapper.scala

@@ -17,7 +17,7 @@
 package com.nvidia.spark.rapids.jni.kudo
 
 import ai.rapids.cudf.Schema
-import com.nvidia.spark.rapids.{CoalescedHostResult, GpuColumnVector, RmmRapidsRetryIterator, SpillableHostBuffer, SpillPriorities}
+import com.nvidia.spark.rapids.{CoalescedHostResult, GpuColumnVector, RmmRapidsRetryIterator, SpillableHostBuffer}
 import com.nvidia.spark.rapids.Arm.withResource
 
 import org.apache.spark.sql.types.DataType
@@ -57,7 +57,7 @@ object KudoHostMergeResultWrapper {
   def apply(inner: KudoHostMergeResult): KudoHostMergeResultWrapper = {
     KudoHostMergeResultWrapper(inner.getSchema, inner.getColumnInfoList,
       SpillableHostBuffer(inner.getHostBuf,
-        inner.getHostBuf.getLength, SpillPriorities.ACTIVE_BATCHING_PRIORITY
+        inner.getHostBuf.getLength
       )
     )
   }

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

@@ -2022,7 +2022,7 @@ trait ParquetPartitionReaderBase extends Logging with ScanWithMetrics
           throw new QueryExecutionException(s"Calculated buffer size $estTotalSize is too " +
               s"small, actual written: ${out.getPos}")
         }
-        (SpillableHostBuffer(hmb, out.getPos, SpillPriorities.ACTIVE_BATCHING_PRIORITY),
+        (SpillableHostBuffer(hmb, out.getPos),
           outputBlocks)
       }
     }
@@ -2552,7 +2552,7 @@ class MultiFileCloudParquetPartitionReader(
         offset += footerOut.getPos
       }
       val spHostBufs = buffers.map(b =>
-        SpillableHostBuffer(b, b.getLength, SpillPriorities.ACTIVE_BATCHING_PRIORITY))
+        SpillableHostBuffer(b, b.getLength))
       val newHmbBufferInfo = SingleHMBAndMeta(spHostBufs.toArray, offset,
         combinedMeta.allPartValues.map(_._1).sum, Seq.empty)
       val newHmbMeta = HostMemoryBuffersWithMetaData(

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

@@ -349,6 +349,7 @@ trait HostSpillableHandle[T <: AutoCloseable] extends SpillableHandle {
   }
 }
 
+
 object SpillableHostBufferHandle extends Logging {
   def apply(hmb: HostMemoryBuffer): SpillableHostBufferHandle = {
     val handle = new SpillableHostBufferHandle(hmb.getLength, host = Some(hmb))
@@ -388,6 +389,7 @@ object SpillableHostBufferHandle extends Logging {
   }
 }
 
+
 class SpillableHostBufferHandle private (
     val sizeInBytes: Long,
     private[spill] override var host: Option[HostMemoryBuffer] = None,
@@ -1867,7 +1869,7 @@ object SpillFramework extends Logging {
   // public for tests. Some tests not in the `spill` package require setting this
   // because they need fine control over allocations.
   var storesInternal: SpillableStores = _
-  
+
   def stores: SpillableStores = {
     if (storesInternal == null) {
       throw new IllegalStateException(