[SPARK-49217][CORE] Support separate buffer size configuration in UnsafeShuffleWriter

[cxzl25]

What changes were proposed in this pull request?

This PR aims to support separate buffer size configuration in UnsafeShuffleWriter.

Introduce spark.shuffle.file.merge.buffer configuration.

Why are the changes needed?

UnsafeShuffleWriter#mergeSpillsWithFileStream uses spark.shuffle.file.buffer as the buffer for reading spill files, and this buffer is an off-heap buffer.

In the spill process, we hope that the buffer size is larger, but once there are too many files in the spill, UnsafeShuffleWriter#mergeSpillsWithFileStream needs to create a lot of off-heap memory, which makes the executor easily killed by YARN.

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

Lines 372 to 375 in e72d21c

	for (int i = 0; i < spills.length; i++) {
	spillInputStreams[i] = new NioBufferedFileInputStream(
	spills[i].file,
	inputBufferSizeInBytes);

Does this PR introduce any user-facing change?

No

How was this patch tested?

Production environment verification

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

No

[mridulm]

The change looks reasonable.
+CC @JoshRosen as well.

As a side note, should we be using Platform.allocateDirectBuffer for NioBufferedFileInputStream as well @JoshRosen ?

[JoshRosen]

LGTM, as this looks reasonable to me as well: this is a bit of a low-level configuration but it seems fine to allow it to be tuned.

---

As a side note, should we be using Platform.allocateDirectBuffer for NioBufferedFileInputStream as well @JoshRosen ?

The difference between ByteBuffer.allocateDirect and Platform.allocateDirectBuffer is that the latter bypasses / ignores the JVM's -XX:MaxDirectMemorySize limit.

Given that almost all other Spark-initiated allocations use the Platform version, we probably should make that change.

That said, I also spot another ByteBuffer.allocateDirect usage at

spark/core/src/main/scala/org/apache/spark/storage/DiskStore.scala

Lines 327 to 328 in 9b9a7a7

	private val buffer = ByteBuffer.allocateDirect(64 * 1024)
	buffer.flip()

plus a potential need for additional StorageUtils.dispose() calls for that other call in ReadableChannelFileRegion (which is only used by EncryptedBlockData, as far as I know, though), so perhaps it would be better to update both of those in a separate PR instead of doing it here.

[cxzl25]

I also found another one, which uses Platform.allocateDirectBuffer when initializing, but uses ByteBuffer.allocateDirect when growing.

We can do this in another PR.

spark/core/src/main/scala/org/apache/spark/util/DirectByteBufferOutputStream.scala

Lines 31 to 32 in 4ceacbe

	private[spark] class DirectByteBufferOutputStream(capacity: Int) extends OutputStream {
	private var buffer = Platform.allocateDirectBuffer(capacity)

spark/core/src/main/scala/org/apache/spark/util/DirectByteBufferOutputStream.scala

Lines 60 to 63 in 4ceacbe

	val newBuffer = ByteBuffer.allocateDirect(newCapacity)
	newBuffer.put(oldBuffer)
	StorageUtils.dispose(oldBuffer)
	buffer = newBuffer

[mridulm]

Merged to master.
Thanks for fixing this @cxzl25 !
Thanks for the review @JoshRosen :-)

[mridulm]

@cxzl25, please do verify if the jira has been updated correctly - thanks !