(datafusion-comet) branch branch-1.0 updated: fix: size JVM shuffle pointer array growth from the array, not the data pages (#5907) (#6227)

[ag...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

andygrove pushed a commit to branch branch-1.0
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git


The following commit(s) were added to refs/heads/branch-1.0 by this push:
     new 9dadf442dd fix: size JVM shuffle pointer array growth from the array, not the data pages (#5907) (#6227)
9dadf442dd is described below

commit 9dadf442dd0bc2e29f4fc12d24d61c7c09f3dabb
Author: Andy Grove <ag...@apache.org>
AuthorDate: Fri Sep 25 14:06:11 2026 -0600

    fix: size JVM shuffle pointer array growth from the array, not the data pages (#5907) (#6227)
    
    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.
    
    (cherry picked from commit 2d1aab3a6abc3461c2320e8a18d203711b1957cf)
---
 .../shuffle/sort/CometShuffleExternalSorter.java   |  5 ++-
 .../org/apache/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 13ea96dfcb..6375419c80 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
@@ -253,7 +253,10 @@ public final class CometShuffleExternalSorter implements CometShuffleChecksumSup
   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 36b50e6202..992ce0f1c1 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
@@ -165,6 +165,21 @@ public class SpillSorter extends SpillWriter {
     }
   }
 
+  /**
+   * 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 dfbe38b648..c6330fb28e 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)
+    }
+  }
+
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org