[SPARK-56511][CORE] Fix NPE in ShuffleInMemorySorter.getMemoryUsage after failed reset

### What changes were proposed in this pull request?

Null-check `array` in `ShuffleInMemorySorter.getMemoryUsage()`.

When `ShuffleExternalSorter.spill()` → `ShuffleInMemorySorter.reset()` → `MemoryConsumer.allocateArray()` throws OOM, `ShuffleInMemorySorter.array` is left null. The OOM propagates to
`UnsafeShuffleWriter.write()`'s finally block, which calls `ShuffleExternalSorter.cleanupResources()` → `freeMemory()` → `updatePeakMemoryUsed()` → `ShuffleInMemorySorter.getMemoryUsage()` → NPE on `array.size()`.

Example stack trace we see in prod:
```
java.lang.NullPointerException
	at org.apache.spark.shuffle.sort.ShuffleInMemorySorter.getMemoryUsage(ShuffleInMemorySorter.java:131)
	at org.apache.spark.shuffle.sort.ShuffleExternalSorter.getMemoryUsage(ShuffleExternalSorter.java:349)
	at org.apache.spark.shuffle.sort.ShuffleExternalSorter.insertRecord(ShuffleExternalSorter.java:472)
	at org.apache.spark.shuffle.sort.UnsafeShuffleWriter.insertRecordIntoSorter(UnsafeShuffleWriter.java:297)
	at org.apache.spark.shuffle.sort.UnsafeShuffleWriter.write(UnsafeShuffleWriter.java:213)
	at org.apache.spark.shuffle.ShuffleWriteProcessor.write(ShuffleWriteProcessor.scala:58)
	at org.apache.spark.scheduler.ShuffleMapTask.$anonfun$runTask$3(ShuffleMapTask.scala:87)
	at org.apache.spark.scheduler.ShuffleMapTask.$anonfun$runTask$1(ShuffleMapTask.scala:82)
	at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:58)
	at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:39)
...
```

Returning 0 for `ShuffleInMemorySorter.getMemoryUsage()` is correct: when `array` is null the pointer array was already freed by `ShuffleInMemorySorter.reset()` and never reallocated — the actual memory usage IS zero. The value is only consumed by
`ShuffleExternalSorter.updatePeakMemoryUsed()` for bookkeeping.

### Why are the changes needed?

Without the fix, the NPE in `cleanupResources()` prevents `ShuffleInMemorySorter.free()` and page cleanup from running, causing a memory leak on top of the original OOM.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

- Unit test (`testGetMemoryUsageAfterFree`) verifying `ShuffleInMemorySorter.getMemoryUsage()` returns 0 after `free()`
- Integration test in `ShuffleExternalSorterSuite`: constrains memory so `ShuffleInMemorySorter.reset()` → `allocateArray()` fails with OOM, then verifies `ShuffleExternalSorter.cleanupResources()` does not throw

Before fix: `NullPointerException: Cannot invoke "LongArray.size()" because "this.array" is null at ShuffleExternalSorter.cleanupResources`

After fix: test passes.

### Was this patch authored or co-authored using generative AI tooling?

Yes.

Closes #55373 from timlee0119/fix-shuffle-sorter-npe.

Authored-by: Tim Lee <tim.lee@databricks.com>
Signed-off-by: Xingbo Jiang <xingbo.jiang@databricks.com>

Author: timlee0119

core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java

@@ -128,6 +128,9 @@ public boolean hasSpaceForAnotherRecord() {
   }
 
   public long getMemoryUsage() {
+    if (array == null) {
+      return 0L;
+    }
     return array.size() * 8;
   }
 

core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java

@@ -120,6 +120,15 @@ public void testBasicSorting() {
     Assertions.assertFalse(iter.hasNext());
   }
 
+  @Test
+  public void testGetMemoryUsageAfterFree() {
+    final ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(
+      consumer, 100, shouldUseRadixSort());
+    Assertions.assertTrue(sorter.getMemoryUsage() > 0);
+    sorter.free();
+    Assertions.assertEquals(0, sorter.getMemoryUsage());
+  }
+
   @Test
   public void testSortingManyNumbers() {
     ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(consumer, 4, shouldUseRadixSort());

core/src/test/scala/org/apache/spark/shuffle/sort/ShuffleExternalSorterSuite.scala

@@ -114,4 +114,93 @@ class ShuffleExternalSorterSuite extends SparkFunSuite with LocalSparkContext wi
       condition = "UNABLE_TO_ACQUIRE_MEMORY",
       parameters = Map("requestedBytes" -> "800", "receivedBytes" -> "400"))
   }
+
+  test("cleanupResources should not NPE when reset fails to reallocate array") {
+    // Reproduces a bug where:
+    //   1. insertRecord() triggers spill -> reset() -> array = null -> allocateArray() throws OOM
+    //   2. OOM propagates out of insertRecord()
+    //   3. UnsafeShuffleWriter's finally block calls cleanupResources()
+    //   4. cleanupResources() -> freeMemory() -> updatePeakMemoryUsed() -> getMemoryUsage()
+    //      -> inMemSorter.getMemoryUsage() -> NPE because inMemSorter.array is still null
+    //
+    // The root cause: reset() sets array = null, then allocateArray() fails. The sorter is left
+    // with inMemSorter != null but inMemSorter.array == null. cleanupResources() calls
+    // freeMemory() which calls getMemoryUsage() before reaching inMemSorter.free().
+    val conf = new SparkConf()
+      .setMaster("local[1]")
+      .setAppName("ShuffleExternalSorterSuite")
+      .set(IS_TESTING, true)
+      .set(TEST_MEMORY, 1600L)
+      .set(MEMORY_FRACTION, 0.9999)
+
+    sc = new SparkContext(conf)
+    val memoryManager = UnifiedMemoryManager(conf, 1)
+
+    var shouldStealMemory = false
+
+    // Override acquireExecutionMemory to steal freed memory during reset()'s allocateArray(),
+    // forcing the allocation to fail with OOM.
+    val taskMemoryManager = new TaskMemoryManager(memoryManager, 0) {
+      override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = {
+        if (shouldStealMemory &&
+          memoryManager.maxHeapMemory - memoryManager.executionMemoryUsed > 400) {
+          val acquireExecutionMemoryMethod =
+            memoryManager.getClass.getMethods.filter(_.getName == "acquireExecutionMemory").head
+          acquireExecutionMemoryMethod.invoke(
+            memoryManager,
+            JLong.valueOf(
+              memoryManager.maxHeapMemory - memoryManager.executionMemoryUsed - 400),
+            JLong.valueOf(1L),
+            MemoryMode.ON_HEAP
+          ).asInstanceOf[java.lang.Long]
+        }
+        super.acquireExecutionMemory(required, consumer)
+      }
+    }
+    val taskContext = mock[TaskContext]
+    val taskMetrics = new TaskMetrics
+    when(taskContext.taskMetrics()).thenReturn(taskMetrics)
+    val sorter = new ShuffleExternalSorter(
+      taskMemoryManager,
+      sc.env.blockManager,
+      taskContext,
+      100, // initialSize: ShuffleInMemorySorter needs 800 bytes (100 * 8)
+      1,
+      conf,
+      new ShuffleWriteMetrics)
+    val inMemSorter = {
+      val field = sorter.getClass.getDeclaredField("inMemSorter")
+      field.setAccessible(true)
+      field.get(sorter).asInstanceOf[ShuffleInMemorySorter]
+    }
+
+    // Fill the pointer array until there's no space for another record.
+    val bytes = new Array[Byte](1)
+    while (inMemSorter.hasSpaceForAnotherRecord) {
+      sorter.insertRecord(bytes, Platform.BYTE_ARRAY_OFFSET, 1, 0)
+    }
+
+    // Enable memory stealing so that when spill -> reset() -> allocateArray() runs, the freed
+    // memory is consumed before allocateArray can use it, causing OOM.
+    shouldStealMemory = true
+
+    // insertRecord triggers spill -> reset() -> array = null -> allocateArray fails -> OOM
+    intercept[SparkOutOfMemoryError] {
+      sorter.insertRecord(bytes, Platform.BYTE_ARRAY_OFFSET, 1, 0)
+    }
+
+    // Verify the broken state: inMemSorter != null but inMemSorter.array == null
+    val inMemSorterField = sorter.getClass.getDeclaredField("inMemSorter")
+    inMemSorterField.setAccessible(true)
+    assert(inMemSorterField.get(sorter) != null, "inMemSorter should still be non-null")
+    val arrayField = classOf[ShuffleInMemorySorter].getDeclaredField("array")
+    arrayField.setAccessible(true)
+    assert(arrayField.get(inMemSorterField.get(sorter)) == null,
+      "inMemSorter.array should be null (reset freed it, allocateArray failed)")
+
+    // Without the fix, this NPEs in:
+    //   cleanupResources -> freeMemory -> updatePeakMemoryUsed -> getMemoryUsage
+    //     -> inMemSorter.getMemoryUsage -> array.size() -> NPE
+    sorter.cleanupResources()
+  }
 }