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

[andygrove]

Which issue does this PR close?

Part of #5905 (finding J2). Does not close it.

Rationale for this change

CometShuffleExternalSorter.growPointerArrayIfNecessary sizes the replacement pointer array from SpillSorter.getMemoryUsage(), which Comet overrides to return the pointer array bytes plus all allocated data pages. Spark's ShuffleExternalSorter sizes it from inMemSorter.getMemoryUsage(), the array alone.

The result is that the first growth, which happens after initialSize / 2 records because half the array is reserved for radix sort, requests 2 x (pageBytes + arrayBytes) / 8 entries instead of 2 x arrayBytes / 8. With the default page size that is tens of megabytes of pointer array for a couple of thousand rows, and every later growth compounds on top of the pages allocated since. Under the bounded allocator this either wastes memory that should have gone to data pages or fails the allocation and forces a spill after very few rows.

What changes are included in this PR?

• SpillSorter gains getPointerArrayMemoryUsage(), returning only the in-memory sorter's array size under the same lock as getMemoryUsage().
• CometShuffleExternalSorter.growPointerArrayIfNecessary uses it to size the new array, matching Spark's behaviour of doubling the pointer array.

getMemoryUsage() itself is unchanged, so peak memory reporting and spill sizing still include the data pages.

How are these changes tested?

New test in SpillSorterSuite that inserts enough records to trigger the first pointer-array growth against a private off-heap allocator and asserts that memory in use afterwards equals one data page plus twice the initial array. Before this change the assertion fails because the array grows to more than the page size.

[sunchao]

Correctness

Summary

Reviewed a82936abdc63aca1d504fe32a4d0d2590097c518 against 451c99963206fa6bf0387239aa12887a16255516. No verified P1/P2 findings remain.

Previously, pointer-array growth used SpillSorter.getMemoryUsage(), which includes both the pointer array and buffered data pages. Consequently, larger data pages caused a larger pointer array even when the number of buffered record pointers was unchanged. The new getter supplies array bytes only; total memory accounting still includes both components.

Compatibility and memory ownership

The units are consistent: Spark's in-memory sorter reports array.size() * 8 bytes, while allocateArray accepts a count of longs. Therefore, used / 8 * 2 doubles the array. With radix sorting, half the array holds records and half is scratch space; the 1,024-entry initial allocation grows before record 513. This matches the growth calculation in the maintained Spark 3.5 and 4.0 sources.

The growth path preserves allocation-before-copy/free ordering and the post-allocation capacity check that releases an unused candidate after a spill. The adopted array still becomes the array passed to native sorting. The new getter uses the existing monitor and guards freed/null state. Both allocator choices use byte-sized blocks for long arrays; configuring Spark on-heap still selects Comet's separate bounded off-heap allocator. Doubling a valid array stays within long arithmetic; existing page-size checks remain responsible for rejecting oversized allocations.

Validation

The regression directly constructs the JVM external sorter with a private unified allocator and checks the first expansion plus cleanup. Its 513 small records stay within one data page, so its assertion isolates pointer-array growth. It intentionally never spills; it does not test native sorting or spill-file encoding.

CI ran all ten SpillSorterSuite tests, including the new regression, successfully on Spark 3.4, 3.5, 4.0, 4.1, and 4.2. Those jobs checked out b05165f5, whose entire source tree equals the reviewed head. The shuffle jobs passed 443 tests with six cancellations on 3.4/3.5 and 492 with no cancellations on 4.0/4.1/4.2.

A local Java component probe compiled the exact base/head growth methods with controlled allocator and spill doubles. It reproduced the regression assertion failing for the base calculation and passing for the head calculation, and checked full/not-full, allocation failure, spill/reset, candidate release, and size-limit branches. This was method-level validation, not a local execution of the Scala suite or native writer. Maintained Spark 3.4 and 4.1 source branches were unavailable; the source comparison is limited to 3.5 and 4.0.

Performance

For array bytes A and buffered page bytes D, the allocation changes from 2(A + D) to 2A. In the regression's first-growth case, the array becomes 16 KiB instead of 8,404,992 bytes. After releasing the old array, total retained memory is 4,210,688 bytes instead of 12,599,296 bytes: exactly 8 MiB less.

The getter also avoids summing data pages when selecting the new array size. Geometric growth can require more expansions than the previous oversized allocation, but it restores Spark's ordinary capacity policy. The direct allocation regression is sufficient for this bounded fix; no throughput benchmark or end-to-end speedup is claimed. The test's peak counter assertion reflects the sorter's sampled accounting, not the transient interval when both old and new arrays are allocated.

Design

The change separates capacity sizing from total memory accounting at the owning sorter. That keeps peak/spill accounting meaningful while removing data-page size from the capacity decision. The allocator interfaces, spill/reset sequence, record-pointer representation, and native output format remain unchanged.

Keeping the existing total-memory getter and adding an explicitly named array-memory getter is a clear fit for these two different callers. The regression exercises the external sorter's real growth trigger rather than asserting only the helper's return value.

Abstraction & complexity

The implementation adds one small accessor and changes one call site. It introduces no additional state or allocation policy. Delegating the value to Spark's in-memory sorter avoids duplicating its capacity calculation or exposing the array.

The private test allocator makes the memory assertions independent of the shared bounded allocator's state, and cleanup checks that no task-managed memory remains. I found no simplification or additional test requirement that should block this change.

[andygrove]

Thanks for the review @sunchao