From a82936abdc63aca1d504fe32a4d0d2590097c518 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Sun, 13 Sep 2026 14:13:45 -0600
Subject: [PATCH] fix: size JVM shuffle pointer array growth from the array,
 not the data pages

CometShuffleExternalSorter.growPointerArrayIfNecessary sized the replacement
pointer array from SpillSorter.getMemoryUsage(), which includes every data
page, so the first growth requested an array proportional to the page size
instead of doubling the array as Spark's ShuffleExternalSorter does. Add
SpillSorter.getPointerArrayMemoryUsage() and use it for growth.
---
 .../sort/CometShuffleExternalSorter.java      |  5 +-
 .../spark/shuffle/sort/SpillSorter.java       | 15 ++++++
 .../spark/shuffle/sort/SpillSorterSuite.scala | 47 +++++++++++++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
index 1a2ba0f96c5..4837cd63b3b 100644
--- a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
+++ b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
@@ -254,7 +254,10 @@ public void cleanupResources() {
   private void growPointerArrayIfNecessary() throws IOException {
     assert (activeSpillSorter != null);
     if (!activeSpillSorter.hasSpaceForAnotherRecord()) {
-      long used = activeSpillSorter.getMemoryUsage();
+      // Size the new array from the pointer array alone, as Spark's ShuffleExternalSorter does.
+      // SpillSorter.getMemoryUsage() also counts the data pages, which would make the array grow
+      // in proportion to the pages instead of doubling.
+      long used = activeSpillSorter.getPointerArrayMemoryUsage();
       LongArray array;
       try {
         // could trigger spilling
diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java
index 2bff92f9fa2..f695a235f41 100644
--- a/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java
+++ b/spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java
@@ -180,6 +180,21 @@ public long getMemoryUsage() {
     }
   }
 
+  /**
+   * Memory held by the in-memory sorter's pointer array alone, excluding the data pages. This is
+   * what the pointer array growth in {@link CometShuffleExternalSorter} must be sized from: sizing
+   * it from {@link #getMemoryUsage()} would grow the array in proportion to the data pages instead
+   * of doubling it.
+   */
+  public long getPointerArrayMemoryUsage() {
+    synchronized (this) {
+      if (freed || inMemSorter == null) {
+        return 0;
+      }
+      return inMemSorter.getMemoryUsage();
+    }
+  }
+
   @Override
   protected void spill(int required) throws IOException {
     spillCallback.onSpillRequired();
diff --git a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala
index dfbe38b6484..c6330fb28e5 100644
--- a/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala
@@ -259,4 +259,51 @@ class SpillSorterSuite extends AnyFunSuite with BeforeAndAfterEach {
     }
   }
 
+  test("pointer array growth is sized from the array, not from the data pages") {
+    // Use the unified (off-heap) allocator so that every allocation made by this sorter is
+    // visible through `getUsed` on an allocator that is private to this test.
+    val offHeapConf = new SparkConf(false)
+      .set("spark.memory.offHeap.enabled", "true")
+      .set("spark.memory.offHeap.size", "64m")
+    val offHeapMemoryManager = new TestMemoryManager(offHeapConf)
+    offHeapMemoryManager.limit(64L * 1024 * 1024)
+    val offHeapTaskMemoryManager = new TaskMemoryManager(offHeapMemoryManager, 0)
+    val allocator =
+      CometShuffleMemoryAllocator.getInstance(offHeapConf, offHeapTaskMemoryManager, PAGE_SIZE)
+    // The block manager is only touched when spilling, which this test never does.
+    val sorter = new CometShuffleExternalSorter(
+      allocator,
+      null,
+      TaskContext.empty(),
+      INITIAL_SIZE,
+      2,
+      offHeapConf,
+      new ShuffleWriteMetrics(),
+      createTestSchema())
+
+    try {
+      val recordData = new Array[Byte](16)
+      def insert(i: Int): Unit =
+        sorter.insertRecord(recordData, Platform.BYTE_ARRAY_OFFSET, recordData.length, i % 2)
+
+      val initialArrayBytes = INITIAL_SIZE * 8L
+      assert(allocator.getUsed === initialArrayBytes)
+
+      insert(0)
+      val pageBytes = allocator.getUsed - initialArrayBytes
+      assert(pageBytes >= PAGE_SIZE)
+
+      // With radix sort enabled the in-memory sorter uses half the array for records, so the
+      // pointer array is grown when the (INITIAL_SIZE / 2 + 1)th record arrives. Growth must
+      // double the pointer array, not request an array sized from the data pages as well.
+      val recordsToTriggerGrowth = INITIAL_SIZE / 2 + 1
+      (1 until recordsToTriggerGrowth).foreach(insert)
+      assert(sorter.getPeakMemoryUsedBytes === pageBytes + 2 * initialArrayBytes)
+      assert(allocator.getUsed === pageBytes + 2 * initialArrayBytes)
+    } finally {
+      sorter.cleanupResources()
+      assert(offHeapTaskMemoryManager.cleanUpAllAllocatedMemory() === 0L)
+    }
+  }
+
 }