Fix SpillablePartialFileHandle read overflow when written more than 2GB (#15327)

Fixes #15325

### Description

This is a bug where the length of the `SpillablePartialFileHandle`
(actually written) was larger than Int.MaxValue, and we overflowed while
reading, due to casting we were doing too early. I refactored it a bit
to be able to test the functionality.

I believe this to be small/targeted enough that I don't see a need to
re-run performance runs when we have spilled partial file handles.

### Checklists

Documentation
- [ ] Updated for new or modified user-facing features or behaviors
- [x] No user-facing change

Testing
- [x] Added or modified tests to cover new code paths
- [ ] Covered by existing tests
(Please provide the names of the existing tests in the PR description.)
- [ ] Not required

Performance
- [ ] Tests ran and results are added in the PR description
- [ ] Issue filed with a link in the PR description
- [x] Not required

Signed-off-by: Alessandro Bellina <abellina@nvidia.com>

Author: abellina

sql-plugin/src/main/scala/com/nvidia/spark/rapids/spill/SpillablePartialFileHandle.scala

@@ -59,6 +59,7 @@ object PartialFileStorageMode extends Enumeration {
  * @param memoryThreshold Host memory usage threshold for buffer expansion decisions
  * @param priority Spill priority for memory-based mode
  * @param syncWrites Whether to force outstanding writes to disk
+ * @param bufferedOutputStreamFactory Creates buffered streams for file writes
  * @param capacityHintProvider Optional function that provides capacity hints based on
  *                             current bytes written and required capacity. When provided,
  *                             buffer expansion will use this hint instead of simple doubling.
@@ -73,6 +74,7 @@ class SpillablePartialFileHandle private (
     memoryThreshold: Double,
     priority: Long,
     syncWrites: Boolean,
+    bufferedOutputStreamFactory: FileOutputStream => BufferedOutputStream,
     capacityHintProvider: Option[(Long, Long) => Long])
   extends HostSpillableHandle[ai.rapids.cudf.HostMemoryBuffer] with Logging {
 
@@ -407,7 +409,8 @@ class SpillablePartialFileHandle private (
       return -1  // EOF
     }
 
-    val actualLength = math.min(length, (totalBytesWritten - readPosition).toInt)
+    val actualLength = SpillablePartialFileHandle.boundedReadLengthAsInt(
+      length, totalBytesWritten - readPosition)
 
     def readFromFile(): Int = {
       ensureFileInputStreamOpen()
@@ -475,7 +478,8 @@ class SpillablePartialFileHandle private (
       return -1
     }
 
-    val actualLength = math.min(length, (totalBytesWritten - position).toInt)
+    val actualLength = SpillablePartialFileHandle.boundedReadLengthAsInt(
+      length, totalBytesWritten - position)
     if (actualLength <= 0) {
       return -1
     }
@@ -698,7 +702,7 @@ class SpillablePartialFileHandle private (
     if (fileOutputStream.isEmpty) {
       val fos = new FileOutputStream(file, true)  // append mode
       fileOutputStream = Some(fos)
-      bufferedOutputStream = Some(new BufferedOutputStream(fos, 64 * 1024))
+      bufferedOutputStream = Some(bufferedOutputStreamFactory(fos))
     }
   }
 
@@ -800,6 +804,14 @@ class SpillablePartialFileHandle private (
 
 object SpillablePartialFileHandle extends Logging {
 
+  private val DEFAULT_FILE_BUFFER_SIZE = 64 * 1024
+
+  private[spill] def boundedReadLengthAsInt(requestedLength: Int, remainingBytes: Long): Int = {
+    require(requestedLength >= 0, s"requestedLength must be non-negative: $requestedLength")
+    require(remainingBytes >= 0, s"remainingBytes must be non-negative: $remainingBytes")
+    math.min(requestedLength.toLong, remainingBytes).toInt
+  }
+
   /**
    * Create a file-only handle.
    * Data is written directly to disk without using host memory.
@@ -817,6 +829,24 @@ object SpillablePartialFileHandle extends Logging {
       memoryThreshold = 0.0,
       priority = Long.MinValue,
       syncWrites = syncWrites,
+      bufferedOutputStreamFactory = new BufferedOutputStream(_, DEFAULT_FILE_BUFFER_SIZE),
+      capacityHintProvider = None)
+  }
+
+  private[spill] def createFileOnly(
+      file: File,
+      syncWrites: Boolean,
+      bufferedOutputStreamFactory: FileOutputStream => BufferedOutputStream):
+  SpillablePartialFileHandle = {
+    new SpillablePartialFileHandle(
+      storageMode = PartialFileStorageMode.FILE_ONLY,
+      file = file,
+      initialCapacity = 0L,
+      maxBufferSize = 0L,
+      memoryThreshold = 0.0,
+      priority = Long.MinValue,
+      syncWrites = syncWrites,
+      bufferedOutputStreamFactory = bufferedOutputStreamFactory,
       capacityHintProvider = None)
   }
 
@@ -857,6 +887,7 @@ object SpillablePartialFileHandle extends Logging {
       memoryThreshold = memoryThreshold,
       priority = priority,
       syncWrites = syncWrites,
+      bufferedOutputStreamFactory = new BufferedOutputStream(_, DEFAULT_FILE_BUFFER_SIZE),
       capacityHintProvider = capacityHintProvider)
   }
 }

tests/src/test/scala/com/nvidia/spark/rapids/spill/SpillablePartialFileHandleSuite.scala

@@ -16,15 +16,17 @@
 
 package com.nvidia.spark.rapids.spill
 
-import java.io.File
+import java.io.{BufferedOutputStream, File, FileOutputStream}
 import java.util.Arrays
 
 import com.nvidia.spark.rapids.Arm.withResource
 import com.nvidia.spark.rapids.RapidsConf
 import org.scalatest.BeforeAndAfterEach
 import org.scalatest.funsuite.AnyFunSuite
+import org.scalatestplus.mockito.MockitoSugar
 
-class SpillablePartialFileHandleSuite extends AnyFunSuite with BeforeAndAfterEach {
+class SpillablePartialFileHandleSuite
+    extends AnyFunSuite with BeforeAndAfterEach with MockitoSugar {
 
   // Use 1GB max buffer size for tests to avoid memory issues on test machines
   private val testMaxBufferSize = 1L * 1024 * 1024 * 1024
@@ -286,6 +288,37 @@ class SpillablePartialFileHandleSuite extends AnyFunSuite with BeforeAndAfterEac
     }
   }
 
+  test("FILE_ONLY mode: read with logical size larger than Int.MaxValue") {
+    val testData = "dummy shuffle bytes".getBytes("UTF-8")
+    val tempFile = createTempFileWithData("test-file-only-large-logical-", testData)
+
+    withLargeLogicalFileOnlyHandle(tempFile) { handle =>
+      val readBuffer = new Array[Byte](testData.length)
+      assert(handle.read(readBuffer, 0, readBuffer.length) == readBuffer.length)
+      assert(readBuffer.sameElements(testData))
+    }
+  }
+
+  test("FILE_ONLY mode: readAt with logical size larger than Int.MaxValue") {
+    val testData = "dummy shuffle bytes".getBytes("UTF-8")
+    val tempFile = createTempFileWithData("test-file-only-large-logical-random-", testData)
+
+    withLargeLogicalFileOnlyHandle(tempFile) { handle =>
+      val readBuffer = new Array[Byte](testData.length)
+      assert(handle.readAt(0, readBuffer, 0, readBuffer.length) == readBuffer.length)
+      assert(readBuffer.sameElements(testData))
+    }
+  }
+
+  test("boundedReadLength requires non-negative inputs") {
+    assertThrows[IllegalArgumentException] {
+      SpillablePartialFileHandle.boundedReadLengthAsInt(-1, 1L)
+    }
+    assertThrows[IllegalArgumentException] {
+      SpillablePartialFileHandle.boundedReadLengthAsInt(1, -1L)
+    }
+  }
+
   test("Error handling: write after finish should fail") {
     val tempFile = File.createTempFile("test-error-", ".tmp")
 
@@ -572,5 +605,34 @@ class SpillablePartialFileHandleSuite extends AnyFunSuite with BeforeAndAfterEac
       assert(handle.isSpilled)
     }
   }
+
+  private def withLargeLogicalFileOnlyHandle(tempFile: File)(
+      testBody: SpillablePartialFileHandle => Unit): Unit = {
+    val outputStream = mock[BufferedOutputStream]
+    withResource(SpillablePartialFileHandle.createFileOnly(
+      file = tempFile,
+      syncWrites = false,
+      bufferedOutputStreamFactory = _ => outputStream)) { handle =>
+      val ignored = new Array[Byte](1)
+      val expectedLogicalSize = Int.MaxValue.toLong + 4096L
+      handle.write(ignored, 0, Int.MaxValue)
+      handle.write(ignored, 0, 4096)
+      handle.finishWrite()
+      assert(handle.getTotalBytesWritten == expectedLogicalSize)
+      testBody(handle)
+    }
+  }
+
+  private def createTempFileWithData(prefix: String, data: Array[Byte]): File = {
+    val tempFile = File.createTempFile(prefix, ".tmp")
+    val out = new FileOutputStream(tempFile)
+    try {
+      out.write(data)
+    } finally {
+      out.close()
+    }
+    tempFile
+  }
+
 }