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