clamp offset spill priorities

NVIDIA · Jan 13, 2022 · 2d60a3f · 2d60a3f
1 parent 04c0a03
commit 2d60a3f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsHostMemoryStore.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsHostMemoryStore.scala
@@ -17,7 +17,7 @@
 package com.nvidia.spark.rapids
 
 import ai.rapids.cudf.{Cuda, DeviceMemoryBuffer, HostMemoryBuffer, MemoryBuffer, PinnedMemoryPool}
-import com.nvidia.spark.rapids.SpillPriorities.{HOST_MEMORY_BUFFER_DIRECT_OFFSET, HOST_MEMORY_BUFFER_PAGEABLE_OFFSET, HOST_MEMORY_BUFFER_PINNED_OFFSET}
+import com.nvidia.spark.rapids.SpillPriorities.{getHostMemoryBufferSpillPriority, HOST_MEMORY_BUFFER_DIRECT_OFFSET, HOST_MEMORY_BUFFER_PAGEABLE_OFFSET, HOST_MEMORY_BUFFER_PINNED_OFFSET}
 import com.nvidia.spark.rapids.StorageTier.StorageTier
 import com.nvidia.spark.rapids.format.TableMeta
 
@@ -98,7 +98,8 @@ class RapidsHostMemoryStore(
         other.id,
         other.size,
         other.meta,
-        other.getSpillPriority + allocationMode.spillPriorityOffset,
+        getHostMemoryBufferSpillPriority(
+          other.getSpillPriority, allocationMode.spillPriorityOffset),
         hostBuffer,
         allocationMode,
         other.spillCallback,

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/SpillPriorities.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/SpillPriorities.scala
@@ -78,4 +78,21 @@ object SpillPriorities {
    * priorities, thus freeing up memory pools first.
    */
   val HOST_MEMORY_BUFFER_DIRECT_OFFSET: Long = 100
+
+  /**
+   * Calculate the host memory buffer spill priority based on an offset, clamping it to avoid
+   * wraparound.
+   *
+   * @param originalSpillPriority the original spill priority
+   * @param offset the desired offset
+   * @return the offset spill priority, with clamping if needed
+   */
+  def getHostMemoryBufferSpillPriority(originalSpillPriority: Long, offset: Long): Long = {
+    if (offset < 0 && originalSpillPriority < Long.MinValue - offset ||
+        offset > 0 && originalSpillPriority > Long.MaxValue - offset) {
+      originalSpillPriority
+    } else {
+      originalSpillPriority + offset
+    }
+  }
 }