fix: size JVM shuffle pointer array growth from the array, not the data pages

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

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();

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)
+    }
+  }
+
 }