Change parameters for memory limit in Parquet chunked reader

sql-plugin/src/main/java/com/nvidia/spark/rapids/iceberg/parquet/GpuParquet.java

@@ -69,7 +69,7 @@ public static class ReadBuilder {
     private long maxBatchSizeBytes = Integer.MAX_VALUE;
     private long targetBatchSizeBytes = Integer.MAX_VALUE;
     private boolean useChunkedReader = false;
-    private boolean useSubPageChunked = false;
+    private long maxChunkedReaderMemoryUsageSizeBytes = 0;
     private scala.Option<String> debugDumpPrefix = null;
     private boolean debugDumpAlways = false;
     private scala.collection.immutable.Map<String, GpuMetric> metrics = null;
@@ -141,9 +141,10 @@ public ReadBuilder withTargetBatchSizeBytes(long targetBatchSizeBytes) {
       return this;
     }
 
-    public ReadBuilder withUseChunkedReader(boolean useChunkedReader, boolean useSubPageChunked) {
+    public ReadBuilder withUseChunkedReader(boolean useChunkedReader,
+        long maxChunkedReaderMemoryUsageSizeBytes) {
       this.useChunkedReader = useChunkedReader;
-      this.useSubPageChunked = useSubPageChunked;
+      this.maxChunkedReaderMemoryUsageSizeBytes = maxChunkedReaderMemoryUsageSizeBytes;
       return this;
     }
 
@@ -164,8 +165,8 @@ public CloseableIterable<ColumnarBatch> build() {
           InternalRow.empty(), file.location(), start, length);
       return new GpuParquetReader(file, projectSchema, options, nameMapping, filter, caseSensitive,
           idToConstant, deleteFilter, partFile, conf, maxBatchSizeRows, maxBatchSizeBytes,
-          targetBatchSizeBytes, useChunkedReader, useSubPageChunked, debugDumpPrefix,
-          debugDumpAlways, metrics);
+          targetBatchSizeBytes, useChunkedReader, maxChunkedReaderMemoryUsageSizeBytes,
+          debugDumpPrefix, debugDumpAlways, metrics);
     }
   }
 

sql-plugin/src/main/java/com/nvidia/spark/rapids/iceberg/parquet/GpuParquetReader.java

@@ -87,7 +87,7 @@ public class GpuParquetReader extends CloseableGroup implements CloseableIterabl
   private final long maxBatchSizeBytes;
   private final long targetBatchSizeBytes;
   private final boolean useChunkedReader;
-  private final boolean useSubPageChunked;
+  private final long maxChunkedReaderMemoryUsageSizeBytes;
   private final scala.Option<String> debugDumpPrefix;
   private final boolean debugDumpAlways;
   private final scala.collection.immutable.Map<String, GpuMetric> metrics;
@@ -98,7 +98,7 @@ public GpuParquetReader(
       Map<Integer, ?> idToConstant, GpuDeleteFilter deleteFilter,
       PartitionedFile partFile, Configuration conf, int maxBatchSizeRows,
       long maxBatchSizeBytes, long targetBatchSizeBytes, boolean useChunkedReader,
-      boolean useSubPageChunked,
+      long maxChunkedReaderMemoryUsageSizeBytes,
       scala.Option<String> debugDumpPrefix, boolean debugDumpAlways,
       scala.collection.immutable.Map<String, GpuMetric> metrics) {
     this.input = input;
@@ -115,7 +115,7 @@ public GpuParquetReader(
     this.maxBatchSizeBytes = maxBatchSizeBytes;
     this.targetBatchSizeBytes = targetBatchSizeBytes;
     this.useChunkedReader = useChunkedReader;
-    this.useSubPageChunked = useSubPageChunked;
+    this.maxChunkedReaderMemoryUsageSizeBytes = maxChunkedReaderMemoryUsageSizeBytes;
     this.debugDumpPrefix = debugDumpPrefix;
     this.debugDumpAlways = debugDumpAlways;
     this.metrics = metrics;
@@ -143,7 +143,7 @@ public org.apache.iceberg.io.CloseableIterator<ColumnarBatch> iterator() {
           new Path(input.location()), clippedBlocks, fileReadSchema, caseSensitive,
           partReaderSparkSchema, debugDumpPrefix, debugDumpAlways,
           maxBatchSizeRows, maxBatchSizeBytes, targetBatchSizeBytes, useChunkedReader,
-          useSubPageChunked,
+          maxChunkedReaderMemoryUsageSizeBytes,
           metrics,
           DateTimeRebaseCorrected$.MODULE$, // dateRebaseMode
           DateTimeRebaseCorrected$.MODULE$, // timestampRebaseMode

sql-plugin/src/main/java/com/nvidia/spark/rapids/iceberg/spark/source/GpuBatchDataReader.java

@@ -48,14 +48,15 @@ class GpuBatchDataReader extends BaseDataReader<ColumnarBatch> {
   private final long maxBatchSizeBytes;
   private final long targetBatchSizeBytes;
   private final boolean useChunkedReader;
-  private final boolean useSubPageChunked;
+  private final long maxChunkedReaderMemoryUsageSizeBytes;
   private final scala.Option<String> parquetDebugDumpPrefix;
   private final boolean parquetDebugDumpAlways;
   private final scala.collection.immutable.Map<String, GpuMetric> metrics;
 
   GpuBatchDataReader(CombinedScanTask task, Table table, Schema expectedSchema, boolean caseSensitive,
                      Configuration conf, int maxBatchSizeRows, long maxBatchSizeBytes,
-                     long targetBatchSizeBytes, boolean useChunkedReader, boolean useSubPageChunked,
+                     long targetBatchSizeBytes,
+                     boolean useChunkedReader, long maxChunkedReaderMemoryUsageSizeBytes,
                      scala.Option<String> parquetDebugDumpPrefix, boolean parquetDebugDumpAlways,
                      scala.collection.immutable.Map<String, GpuMetric> metrics) {
     super(table, task);
@@ -67,7 +68,7 @@ class GpuBatchDataReader extends BaseDataReader<ColumnarBatch> {
     this.maxBatchSizeBytes = maxBatchSizeBytes;
     this.targetBatchSizeBytes = targetBatchSizeBytes;
     this.useChunkedReader = useChunkedReader;
-    this.useSubPageChunked = useSubPageChunked;
+    this.maxChunkedReaderMemoryUsageSizeBytes = maxChunkedReaderMemoryUsageSizeBytes;
     this.parquetDebugDumpPrefix = parquetDebugDumpPrefix;
     this.parquetDebugDumpAlways = parquetDebugDumpAlways;
     this.metrics = metrics;
@@ -102,7 +103,7 @@ CloseableIterator<ColumnarBatch> open(FileScanTask task) {
           .withMaxBatchSizeRows(maxBatchSizeRows)
           .withMaxBatchSizeBytes(maxBatchSizeBytes)
           .withTargetBatchSizeBytes(targetBatchSizeBytes)
-          .withUseChunkedReader(useChunkedReader, useSubPageChunked)
+          .withUseChunkedReader(useChunkedReader, maxChunkedReaderMemoryUsageSizeBytes)
           .withDebugDump(parquetDebugDumpPrefix, parquetDebugDumpAlways)
           .withMetrics(metrics);
 

sql-plugin/src/main/java/com/nvidia/spark/rapids/iceberg/spark/source/GpuMultiFileBatchReader.java

@@ -68,7 +68,7 @@ class GpuMultiFileBatchReader extends BaseDataReader<ColumnarBatch> {
   private final long maxGpuColumnSizeBytes;
 
   private final boolean useChunkedReader;
-  private final boolean useSubPageChunked;
+  private final long maxChunkedReaderMemoryUsageSizeBytes;
   private final scala.Option<String> parquetDebugDumpPrefix;
   private final boolean parquetDebugDumpAlways;
   private final scala.collection.immutable.Map<String, GpuMetric> metrics;
@@ -87,7 +87,7 @@ class GpuMultiFileBatchReader extends BaseDataReader<ColumnarBatch> {
   GpuMultiFileBatchReader(CombinedScanTask task, Table table, Schema expectedSchema,
       boolean caseSensitive, Configuration conf, int maxBatchSizeRows, long maxBatchSizeBytes,
       long targetBatchSizeBytes, long maxGpuColumnSizeBytes,
-      boolean useChunkedReader, boolean useSubPageChunked,
+      boolean useChunkedReader, long maxChunkedReaderMemoryUsageSizeBytes,
       scala.Option<String> parquetDebugDumpPrefix, boolean parquetDebugDumpAlways,
       int numThreads, int maxNumFileProcessed,
       boolean useMultiThread, FileFormat fileFormat,
@@ -102,7 +102,7 @@ class GpuMultiFileBatchReader extends BaseDataReader<ColumnarBatch> {
     this.targetBatchSizeBytes = targetBatchSizeBytes;
     this.maxGpuColumnSizeBytes = maxGpuColumnSizeBytes;
     this.useChunkedReader = useChunkedReader;
-    this.useSubPageChunked = useSubPageChunked;
+    this.maxChunkedReaderMemoryUsageSizeBytes = maxChunkedReaderMemoryUsageSizeBytes;
     this.parquetDebugDumpPrefix = parquetDebugDumpPrefix;
     this.parquetDebugDumpAlways = parquetDebugDumpAlways;
     this.useMultiThread = useMultiThread;
@@ -352,7 +352,7 @@ protected FilePartitionReaderBase createRapidsReader(PartitionedFile[] pFiles,
       return new MultiFileCloudParquetPartitionReader(conf, pFiles,
           this::filterParquetBlocks, caseSensitive, parquetDebugDumpPrefix, parquetDebugDumpAlways,
           maxBatchSizeRows, maxBatchSizeBytes, targetBatchSizeBytes, maxGpuColumnSizeBytes,
-          useChunkedReader, useSubPageChunked, metrics, partitionSchema,
+          useChunkedReader, maxChunkedReaderMemoryUsageSizeBytes, metrics, partitionSchema,
           numThreads, maxNumFileProcessed,
           false, // ignoreMissingFiles
           false, // ignoreCorruptFiles
@@ -428,9 +428,9 @@ protected FilePartitionReaderBase createRapidsReader(PartitionedFile[] pFiles,
 
       return new MultiFileParquetPartitionReader(conf, pFiles,
           JavaConverters.asScalaBuffer(clippedBlocks).toSeq(),
-          caseSensitive, parquetDebugDumpPrefix, parquetDebugDumpAlways, useChunkedReader,
-          useSubPageChunked,
+          caseSensitive, parquetDebugDumpPrefix, parquetDebugDumpAlways,
           maxBatchSizeRows, maxBatchSizeBytes, targetBatchSizeBytes, maxGpuColumnSizeBytes,
+          useChunkedReader, maxChunkedReaderMemoryUsageSizeBytes,
           metrics, partitionSchema, numThreads,
           false, // ignoreMissingFiles
           false, // ignoreCorruptFiles

sql-plugin/src/main/java/com/nvidia/spark/rapids/iceberg/spark/source/GpuSparkScan.java

@@ -283,7 +283,7 @@ private static class MultiFileBatchReader
       super(task.task, task.table(), task.expectedSchema(), task.isCaseSensitive(),
           task.getConfiguration(), task.getMaxBatchSizeRows(), task.getMaxBatchSizeBytes(),
           task.getTargetBatchSizeBytes(), task.getMaxGpuColumnSizeBytes(), task.useChunkedReader(),
-          task.useSubPageChunked(),
+          task.maxChunkedReaderMemoryUsageSizeBytes(),
           task.getParquetDebugDumpPrefix(), task.getParquetDebugDumpAlways(),
           task.getNumThreads(), task.getMaxNumFileProcessed(),
           useMultiThread, ff, metrics, queryUsesInputFile);
@@ -294,7 +294,7 @@ private static class BatchReader extends GpuBatchDataReader implements Partition
     BatchReader(ReadTask task, scala.collection.immutable.Map<String, GpuMetric> metrics) {
       super(task.task, task.table(), task.expectedSchema(), task.isCaseSensitive(),
           task.getConfiguration(), task.getMaxBatchSizeRows(), task.getMaxBatchSizeBytes(),
-          task.getTargetBatchSizeBytes(), task.useChunkedReader(), task.useSubPageChunked(),
+          task.getTargetBatchSizeBytes(), task.useChunkedReader(), task.maxChunkedReaderMemoryUsageSizeBytes(),
           task.getParquetDebugDumpPrefix(), task.getParquetDebugDumpAlways(), metrics);
     }
   }
@@ -305,7 +305,7 @@ static class ReadTask implements InputPartition, Serializable {
     private final String expectedSchemaString;
     private final boolean caseSensitive;
     private final boolean useChunkedReader;
-    private final boolean useSubPageChunked;
+    private final long maxChunkedReaderMemoryUsageSizeBytes;
     private final Broadcast<SerializableConfiguration> confBroadcast;
     private final int maxBatchSizeRows;
     private final long maxBatchSizeBytes;
@@ -343,7 +343,12 @@ static class ReadTask implements InputPartition, Serializable {
       this.numThreads = rapidsConf.multiThreadReadNumThreads();
       this.maxNumFileProcessed = rapidsConf.maxNumParquetFilesParallel();
       this.useChunkedReader = rapidsConf.chunkedReaderEnabled();
-      this.useSubPageChunked = rapidsConf.chunkedSubPageReaderEnabled();
+      if(rapidsConf.limitChunkedReaderMemoryUsage()) {
+        double limitRatio = rapidsConf.chunkedReaderMemoryUsageRatio();
+        this.maxChunkedReaderMemoryUsageSizeBytes = (long)(limitRatio * this.targetBatchSizeBytes);
+      } else {
+        this.maxChunkedReaderMemoryUsageSizeBytes = 0L;
+      }
     }
 
     @Override
@@ -410,8 +415,8 @@ public boolean useChunkedReader() {
       return useChunkedReader;
     }
 
-    public boolean useSubPageChunked() {
-      return useSubPageChunked;
+    public long maxChunkedReaderMemoryUsageSizeBytes() {
+      return maxChunkedReaderMemoryUsageSizeBytes;
     }
   }
 }

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

@@ -214,8 +214,14 @@ abstract class MultiFilePartitionReaderFactoryBase(
   protected val maxReadBatchSizeRows: Int = rapidsConf.maxReadBatchSizeRows
   protected val maxReadBatchSizeBytes: Long = rapidsConf.maxReadBatchSizeBytes
   protected val targetBatchSizeBytes: Long = rapidsConf.gpuTargetBatchSizeBytes
-  protected val subPageChunked: Boolean = rapidsConf.chunkedSubPageReaderEnabled
   protected val maxGpuColumnSizeBytes: Long = rapidsConf.maxGpuColumnSizeBytes
+  protected val useChunkedReader: Boolean = rapidsConf.chunkedReaderEnabled
+  protected val maxChunkedReaderMemoryUsageSizeBytes: Long =
+    if(rapidsConf.limitChunkedReaderMemoryUsage) {
+      (rapidsConf.chunkedReaderMemoryUsageRatio * targetBatchSizeBytes).toLong
+    } else {
+      0L
+    }
   private val allCloudSchemes = rapidsConf.getCloudSchemes.toSet
 
   override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {