Reuse Hadoop vectored read copy buffer in parquet reader

Signed-off-by: Hongbin Ma (Mahone) <mahongbin@apache.org>

Author: binmahone

sql-plugin/src/main/java/com/nvidia/spark/rapids/fileio/CopyBufferReadableInputFile.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.nvidia.spark.rapids.fileio;
+
+import ai.rapids.cudf.HostMemoryBuffer;
+import com.nvidia.spark.rapids.jni.fileio.RapidsInputFile;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Capability interface for input files that can use a caller-supplied copy buffer.
+ */
+public interface CopyBufferReadableInputFile extends RapidsInputFile {
+    /**
+     * Reads ranges into {@code output} using {@code copyBuffer} as caller-owned scratch space.
+     * Implementations must not retain a reference to {@code copyBuffer}.
+     *
+     * @param output the destination buffer
+     * @param copyRanges input ranges and output offsets
+     * @param copyBuffer caller-owned reusable copy buffer
+     * @throws IOException if an I/O error occurs during reading
+     */
+    void readVectored(
+            HostMemoryBuffer output,
+            List<RapidsInputFile.CopyRange> copyRanges,
+            byte[] copyBuffer) throws IOException;
+}

sql-plugin/src/main/java/com/nvidia/spark/rapids/fileio/hadoop/HadoopInputFile.java

@@ -17,6 +17,7 @@
 package com.nvidia.spark.rapids.fileio.hadoop;
 
 import ai.rapids.cudf.HostMemoryBuffer;
+import com.nvidia.spark.rapids.fileio.CopyBufferReadableInputFile;
 import com.nvidia.spark.rapids.jni.fileio.RapidsInputFile;
 import com.nvidia.spark.rapids.jni.fileio.SeekableInputStream;
 import org.apache.hadoop.conf.Configuration;
@@ -35,31 +36,22 @@
  * This class provides methods to get the length of the file and to open a seekable input stream
  * for reading the file.
  */
-public class HadoopInputFile implements RapidsInputFile {
-    private static final String PARQUET_READ_ALLOCATION_SIZE = "parquet.read.allocation.size";
-
+public class HadoopInputFile implements CopyBufferReadableInputFile {
     private final Path filePath;
     private final FileSystem fs;
-    private final int copyBufferSize;
 
     public static HadoopInputFile create(Path filePath, Configuration conf) throws IOException {
         Objects.requireNonNull(filePath, "filePath can't be null!");
         Objects.requireNonNull(conf, "Hadoop conf can't be null");
         FileSystem fs = filePath.getFileSystem(conf);
-        int copyBufferSize = conf.getInt(PARQUET_READ_ALLOCATION_SIZE,
-                RapidsInputFile.DEFAULT_READ_VECTORED_COPY_BUFFER_SIZE);
-        return new HadoopInputFile(filePath, fs, copyBufferSize);
+        return new HadoopInputFile(filePath, fs);
     }
 
-    private HadoopInputFile(Path filePath, FileSystem fs, int copyBufferSize) {
+    private HadoopInputFile(Path filePath, FileSystem fs) {
         Objects.requireNonNull(filePath, "filePath can't be null!");
         Objects.requireNonNull(fs, "FileSystem can't be null");
-        if (copyBufferSize <= 0) {
-            throw new IllegalArgumentException(PARQUET_READ_ALLOCATION_SIZE + " must be positive");
-        }
         this.filePath = filePath;
         this.fs = fs;
-        this.copyBufferSize = copyBufferSize;
     }
 
     @Override
@@ -83,8 +75,10 @@ public SeekableInputStream open() throws IOException {
     }
 
     @Override
-    public void readVectored(HostMemoryBuffer output, List<RapidsInputFile.CopyRange> copyRanges)
-            throws IOException {
-        RapidsInputFile.readVectoredUsingCopyBuffer(this, output, copyRanges, copyBufferSize);
+    public void readVectored(
+            HostMemoryBuffer output,
+            List<RapidsInputFile.CopyRange> copyRanges,
+            byte[] copyBuffer) throws IOException {
+        RapidsInputFile.readVectoredUsingCopyBuffer(this, output, copyRanges, copyBuffer);
     }
 }

sql-plugin/src/main/scala/com/nvidia/spark/rapids/parquet/GpuParquetScan.scala

@@ -38,6 +38,7 @@ import com.nvidia.spark.rapids.RapidsConf.ParquetFooterReaderType
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
 import com.nvidia.spark.rapids.RmmRapidsRetryIterator.withRetryNoSplit
 import com.nvidia.spark.rapids.filecache.FileCache
+import com.nvidia.spark.rapids.fileio.CopyBufferReadableInputFile
 import com.nvidia.spark.rapids.fileio.hadoop.HadoopFileIO
 import com.nvidia.spark.rapids.io.async._
 import com.nvidia.spark.rapids.jni.{DateTimeRebase, ParquetFooter, RmmSpark}
@@ -280,6 +281,22 @@ object GpuParquetScan {
     }
   }
 
+  private[parquet] def readRangesToHostMemory(
+      inputFile: CopyBufferReadableInputFile,
+      output: HostMemoryBuffer,
+      ranges: Seq[CopyRange],
+      metrics: Map[String, GpuMetric],
+      copyBuffer: Array[Byte]): Long = {
+    if (ranges.isEmpty) {
+      0L
+    } else {
+      metrics.getOrElse(READ_FS_TIME, NoopMetric).ns {
+        inputFile.readVectored(output, ranges.asJava, copyBuffer)
+      }
+      ranges.map(_.getLength).sum
+    }
+  }
+
   def throwIfRebaseNeededInExceptionMode(table: Table, dateRebaseMode: DateTimeRebaseMode,
       timestampRebaseMode: DateTimeRebaseMode): Unit = {
     (0 until table.getNumberOfColumns).foreach { i =>
@@ -1579,6 +1596,9 @@ trait ParquetPartitionReaderBase extends Logging with ScanWithMetrics
   def compressCfg: CpuCompressionConfig
 
   val copyBufferSize = conf.getInt("parquet.read.allocation.size", 8 * 1024 * 1024)
+  require(copyBufferSize > 0, "parquet.read.allocation.size must be positive")
+
+  private[this] lazy val readVectoredCopyBuffer: Array[Byte] = new Array[Byte](copyBufferSize)
 
   def checkIfNeedToSplitBlocks(currentDateRebaseMode: DateTimeRebaseMode,
       nextDateRebaseMode: DateTimeRebaseMode,
@@ -2029,8 +2049,14 @@ trait ParquetPartitionReaderBase extends Logging with ScanWithMetrics
     if (scheme != null && scheme.startsWith("s3")) {
       GpuTaskMetrics.get.recordPerfioS3BackendOnce()
     }
-    val totalBytesCopied = GpuParquetScan.readRangesToHostMemory(
-      inputFile, out.buffer, coalescedRanges, metrics)
+    val totalBytesCopied = inputFile match {
+      case copyBufferReadable: CopyBufferReadableInputFile =>
+        GpuParquetScan.readRangesToHostMemory(
+          copyBufferReadable, out.buffer, coalescedRanges, metrics,
+          readVectoredCopyBuffer)
+      case _ =>
+        GpuParquetScan.readRangesToHostMemory(inputFile, out.buffer, coalescedRanges, metrics)
+    }
 
     // try to cache the remote ranges that were copied
     remoteCopies.foreach { range =>