Prefetch Iceberg Parquet footers with a suffix read

Signed-off-by: Zach Puller <zpuller@nvidia.com>

Author: zpuller

iceberg/common/src/main/java/com/nvidia/spark/rapids/fileio/iceberg/IcebergInputFile.java

@@ -16,6 +16,7 @@
 
 package com.nvidia.spark.rapids.fileio.iceberg;
 
+import ai.rapids.cudf.HostMemoryBuffer;
 import com.nvidia.spark.rapids.jni.fileio.RapidsInputFile;
 import com.nvidia.spark.rapids.jni.fileio.SeekableInputStream;
 import org.apache.iceberg.io.InputFile;
@@ -31,10 +32,16 @@
  */
 public class IcebergInputFile implements RapidsInputFile {
   private final InputFile delegate;
+  private final SuffixReader suffixReader;
 
   public IcebergInputFile(InputFile delegate) {
+    this(delegate, null);
+  }
+
+  public IcebergInputFile(InputFile delegate, SuffixReader suffixReader) {
     Objects.requireNonNull(delegate, "delegate can't be null");
     this.delegate = delegate;
+    this.suffixReader = suffixReader;
   }
 
   @Override
@@ -60,4 +67,21 @@ public SeekableInputStream open() throws IOException {
   public InputFile getDelegate() {
     return delegate;
   }
+
+  public boolean supportsSuffixReads() {
+    return suffixReader != null;
+  }
+
+  public long readSuffix(long length, HostMemoryBuffer output, long outputOffset)
+      throws IOException {
+    if (suffixReader == null) {
+      throw new UnsupportedOperationException("Suffix reads are not supported");
+    }
+    return suffixReader.read(length, output, outputOffset);
+  }
+
+  @FunctionalInterface
+  public interface SuffixReader {
+    long read(long length, HostMemoryBuffer output, long outputOffset) throws IOException;
+  }
 }

iceberg/common/src/main/java/org/apache/iceberg/aws/s3/IcebergS3InputFile.java

@@ -85,6 +85,16 @@ public static RapidsInputFile maybeCreate(InputFile inputFile, FileIO fileIO) {
     return new IcebergS3InputFile(delegate, s3Uri, icebergS3Client);
   }
 
+  public static IcebergInputFile maybeCreateIcebergInputFile(
+      InputFile inputFile, FileIO fileIO) {
+    RapidsInputFile optimizedFile = maybeCreate(inputFile, fileIO);
+    if (optimizedFile instanceof IcebergS3InputFile) {
+      IcebergS3InputFile s3File = (IcebergS3InputFile) optimizedFile;
+      return new IcebergInputFile(inputFile, s3File::readTail);
+    }
+    return (IcebergInputFile) optimizedFile;
+  }
+
   @Override
   public String path() {
     return delegate.path();
@@ -127,12 +137,18 @@ public void readVectored(HostMemoryBuffer output, List<CopyRange> copyRanges)
    */
   @Override
   public void readTail(long length, HostMemoryBuffer output) throws IOException {
+    readTail(length, output, 0L);
+  }
+
+  public long readTail(long length, HostMemoryBuffer output, long outputOffset)
+      throws IOException {
     if (length == 0) {
-      return;
+      return 0;
     }
     if (length < 0) {
       throw new IllegalArgumentException("length must be non-negative");
     }
-    IcebergS3RangeCopier.copyTailToHMB(icebergS3Client, output, s3Uri, length, /*dstOffset*/ 0L);
+    return IcebergS3RangeCopier.copyTailToHMB(
+        icebergS3Client, output, s3Uri, length, outputOffset);
   }
 }

iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuIcebergPartitionReader.scala

@@ -20,12 +20,13 @@ import scala.collection.JavaConverters._
 
 import com.nvidia.spark.rapids.GpuMetric
 import com.nvidia.spark.rapids.MapUtil.toMapStrict
-import com.nvidia.spark.rapids.fileio.iceberg.{IcebergFileIO, IcebergInputFile}
+import com.nvidia.spark.rapids.fileio.iceberg.IcebergFileIO
 import com.nvidia.spark.rapids.iceberg.ShimUtils
 import com.nvidia.spark.rapids.iceberg.ShimUtils.locationOf
 import com.nvidia.spark.rapids.iceberg.data.GpuDeleteFilter
 import com.nvidia.spark.rapids.iceberg.parquet._
 import org.apache.iceberg._
+import org.apache.iceberg.aws.s3.IcebergS3InputFile
 import org.apache.iceberg.encryption.EncryptedFiles
 import org.apache.iceberg.mapping.NameMappingParser
 
@@ -110,7 +111,7 @@ class GpuIcebergPartitionReader(private val task: GpuSparkInputPartition,
     val inputFiles = table.encryption()
       .decrypt(encryptedFiles.asJava)
       .asScala
-      .map(f => f.location() -> new IcebergInputFile(f))
+      .map(f => f.location() -> IcebergS3InputFile.maybeCreateIcebergInputFile(f, fileIO))
       .toMap
 
     val taskMap = toMapStrict(tasks.map(t => {

iceberg/iceberg-1-10-x/src/main/scala/com/nvidia/spark/rapids/iceberg/iceberg110x/GpuParquetIOShim.scala

@@ -17,7 +17,7 @@
 package com.nvidia.spark.rapids.iceberg.iceberg110x
 
 import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
-import com.nvidia.spark.rapids.GpuMetric
+import com.nvidia.spark.rapids.{GpuMetric, PerfIO}
 import com.nvidia.spark.rapids.fileio.iceberg.IcebergInputFile
 import com.nvidia.spark.rapids.iceberg.parquet.converter.ToIcebergShaded
 import com.nvidia.spark.rapids.parquet.{HMBInputFile, ParquetFooterUtils}
@@ -41,7 +41,15 @@ object GpuParquetIOShim {
       metrics: Map[String, GpuMetric]): ParquetFileReader = {
     val metadata = withResource(ParquetFooterUtils.getFooterBuffer(
         inputFile, metrics,
-        ParquetFooterUtils.readFooterBufferFromInputFile(inputFile, filePath))) { hmb =>
+        if (inputFile.supportsSuffixReads) {
+          PerfIO.readParquetFooterBufferFromTail(
+            filePath,
+            (length, output, outputOffset) =>
+              inputFile.readSuffix(length, output, outputOffset),
+            ParquetFooterUtils.verifyParquetMagic)
+        } else {
+          ParquetFooterUtils.readFooterBufferFromInputFile(inputFile, filePath)
+        })) { hmb =>
       val shadedHmbFile = ToIcebergShaded.shade(new HMBInputFile(hmb))
       withResource(shadedHmbFile.newStream()) { hmbStream =>
         ParquetFileReader.readFooter(shadedHmbFile, options, hmbStream)

iceberg/iceberg-1-11-x/src/main/scala/com/nvidia/spark/rapids/iceberg/iceberg111x/GpuParquetIOShim.scala

@@ -17,7 +17,7 @@
 package com.nvidia.spark.rapids.iceberg.iceberg111x
 
 import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
-import com.nvidia.spark.rapids.GpuMetric
+import com.nvidia.spark.rapids.{GpuMetric, PerfIO}
 import com.nvidia.spark.rapids.fileio.iceberg.IcebergInputFile
 import com.nvidia.spark.rapids.iceberg.parquet.converter.ToIcebergShaded
 import com.nvidia.spark.rapids.parquet.{HMBInputFile, ParquetFooterUtils}
@@ -41,7 +41,15 @@ object GpuParquetIOShim {
       metrics: Map[String, GpuMetric]): ParquetFileReader = {
     val metadata = withResource(ParquetFooterUtils.getFooterBuffer(
         inputFile, metrics,
-        ParquetFooterUtils.readFooterBufferFromInputFile(inputFile, filePath))) { hmb =>
+        if (inputFile.supportsSuffixReads) {
+          PerfIO.readParquetFooterBufferFromTail(
+            filePath,
+            (length, output, outputOffset) =>
+              inputFile.readSuffix(length, output, outputOffset),
+            ParquetFooterUtils.verifyParquetMagic)
+        } else {
+          ParquetFooterUtils.readFooterBufferFromInputFile(inputFile, filePath)
+        })) { hmb =>
       val shadedHmbFile = ToIcebergShaded.shade(new HMBInputFile(hmb))
       withResource(shadedHmbFile.newStream()) { hmbStream =>
         ParquetFileReader.readFooter(shadedHmbFile, options, hmbStream)