Start to use explicit memory limits in the parquet chunked reader

docs/additional-functionality/advanced_configs.md

@@ -133,6 +133,7 @@ Name | Description | Default Value | Applicable at
 <a name="sql.optimizer.joinReorder.enabled"></a>spark.rapids.sql.optimizer.joinReorder.enabled|When enabled, joins may be reordered for improved query performance|true|Runtime
 <a name="sql.python.gpu.enabled"></a>spark.rapids.sql.python.gpu.enabled|This is an experimental feature and is likely to change in the future. Enable (true) or disable (false) support for scheduling Python Pandas UDFs with GPU resources. When enabled, pandas UDFs are assumed to share the same GPU that the RAPIDs accelerator uses and will honor the python GPU configs|false|Runtime
 <a name="sql.reader.chunked"></a>spark.rapids.sql.reader.chunked|Enable a chunked reader where possible. A chunked reader allows reading highly compressed data that could not be read otherwise, but at the expense of more GPU memory, and in some cases more GPU computation.|true|Runtime
+<a name="sql.reader.chunked.subPage"></a>spark.rapids.sql.reader.chunked.subPage|Enable a chunked reader where possible for reading data that is smaller than the typical row group/page limit. Currently this only works for parquet.|false|Runtime
 <a name="sql.reader.multithreaded.combine.sizeBytes"></a>spark.rapids.sql.reader.multithreaded.combine.sizeBytes|The target size in bytes to combine multiple small files together when using the MULTITHREADED parquet or orc reader. With combine disabled, the MULTITHREADED reader reads the files in parallel and sends individual files down to the GPU, but that can be inefficient for small files. When combine is enabled, files that are ready within spark.rapids.sql.reader.multithreaded.combine.waitTime together, up to this threshold size, are combined before sending down to GPU. This can be disabled by setting it to 0. Note that combine also will not go over the spark.rapids.sql.reader.batchSizeRows or spark.rapids.sql.reader.batchSizeBytes limits.|67108864|Runtime
 <a name="sql.reader.multithreaded.combine.waitTime"></a>spark.rapids.sql.reader.multithreaded.combine.waitTime|When using the multithreaded parquet or orc reader with combine mode, how long to wait, in milliseconds, for more files to finish if haven't met the size threshold. Note that this will wait this amount of time from when the last file was available, so total wait time could be larger then this.|200|Runtime
 <a name="sql.reader.multithreaded.read.keepOrder"></a>spark.rapids.sql.reader.multithreaded.read.keepOrder|When using the MULTITHREADED reader, if this is set to true we read the files in the same order Spark does, otherwise the order may not be the same. Now it is supported only for parquet and orc.|true|Runtime

integration_tests/src/main/python/parquet_test.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2023, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -77,6 +77,12 @@ def read_parquet_sql(data_path):
 coalesce_parquet_file_reader_multithread_filter_chunked_conf = {'spark.rapids.sql.format.parquet.reader.type': 'COALESCING',
         'spark.rapids.sql.coalescing.reader.numFilterParallel': '2',
         'spark.rapids.sql.reader.chunked': True}
+coalesce_parquet_file_reader_multithread_filter_sub_chunked_conf = {'spark.rapids.sql.format.parquet.reader.type': 'COALESCING',
+        'spark.rapids.sql.coalescing.reader.numFilterParallel': '2',
+        'spark.rapids.sql.reader.chunked': True,
+        'spark.rapids.sql.reader.chunked.subPage': True,
+        #'spark.rapids.sql.parquet.debug.dumpPrefix': '/data/tmp/debug/'
+        }
 coalesce_parquet_file_reader_multithread_filter_conf = {'spark.rapids.sql.format.parquet.reader.type': 'COALESCING',
         'spark.rapids.sql.coalescing.reader.numFilterParallel': '2',
         'spark.rapids.sql.reader.chunked': False}
@@ -91,6 +97,12 @@ def read_parquet_sql(data_path):
 native_coalesce_parquet_file_reader_chunked_conf = {'spark.rapids.sql.format.parquet.reader.type': 'COALESCING',
         'spark.rapids.sql.format.parquet.reader.footer.type': 'NATIVE',
         'spark.rapids.sql.reader.chunked': True}
+native_coalesce_parquet_file_reader_sub_chunked_conf = {'spark.rapids.sql.format.parquet.reader.type': 'COALESCING',
+        'spark.rapids.sql.format.parquet.reader.footer.type': 'NATIVE',
+        'spark.rapids.sql.reader.chunked': True,
+        'spark.rapids.sql.reader.chunked.subPage': True,
+        #'spark.rapids.sql.parquet.debug.dumpPrefix': '/data/tmp/debug/'
+        }
 combining_multithreaded_parquet_file_reader_conf_ordered = {'spark.rapids.sql.format.parquet.reader.type': 'MULTITHREADED',
         'spark.rapids.sql.reader.multithreaded.combine.sizeBytes': '64m',
         'spark.rapids.sql.reader.multithreaded.read.keepOrder': True}
@@ -108,7 +120,9 @@ def read_parquet_sql(data_path):
 reader_opt_confs_native = [native_parquet_file_reader_conf, native_multithreaded_parquet_file_reader_conf,
                     native_coalesce_parquet_file_reader_conf,
                     coalesce_parquet_file_reader_multithread_filter_chunked_conf,
-                    native_coalesce_parquet_file_reader_chunked_conf]
+                    coalesce_parquet_file_reader_multithread_filter_sub_chunked_conf,
+                    native_coalesce_parquet_file_reader_chunked_conf,
+                    native_coalesce_parquet_file_reader_sub_chunked_conf]
 
 reader_opt_confs_no_native = [original_parquet_file_reader_conf, multithreaded_parquet_file_reader_conf,
                     coalesce_parquet_file_reader_conf, coalesce_parquet_file_reader_multithread_filter_conf,
@@ -121,7 +135,8 @@ def read_parquet_sql(data_path):
 @pytest.mark.parametrize('parquet_gens', [[byte_gen, short_gen, int_gen, long_gen]], ids=idfn)
 @pytest.mark.parametrize('read_func', [read_parquet_df])
 @pytest.mark.parametrize('reader_confs', [coalesce_parquet_file_reader_multithread_filter_conf,
-    coalesce_parquet_file_reader_multithread_filter_chunked_conf])
+    coalesce_parquet_file_reader_multithread_filter_chunked_conf,
+    coalesce_parquet_file_reader_multithread_filter_sub_chunked_conf])
 @pytest.mark.parametrize('v1_enabled_list', ["", "parquet"])
 def test_parquet_read_coalescing_multiple_files(spark_tmp_path, parquet_gens, read_func, reader_confs, v1_enabled_list):
     gen_list = [('_c' + str(i), gen) for i, gen in enumerate(parquet_gens)]

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -69,6 +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 scala.Option<String> debugDumpPrefix = null;
     private boolean debugDumpAlways = false;
     private scala.collection.immutable.Map<String, GpuMetric> metrics = null;
@@ -140,8 +141,9 @@ public ReadBuilder withTargetBatchSizeBytes(long targetBatchSizeBytes) {
       return this;
     }
 
-    public ReadBuilder withUseChunkedReader(boolean useChunkedReader) {
+    public ReadBuilder withUseChunkedReader(boolean useChunkedReader, boolean useSubPageChunked) {
       this.useChunkedReader = useChunkedReader;
+      this.useSubPageChunked = useSubPageChunked;
       return this;
     }
 
@@ -162,7 +164,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, debugDumpPrefix, debugDumpAlways, metrics);
+          targetBatchSizeBytes, useChunkedReader, useSubPageChunked, debugDumpPrefix,
+          debugDumpAlways, metrics);
     }
   }
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -87,6 +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 scala.Option<String> debugDumpPrefix;
   private final boolean debugDumpAlways;
   private final scala.collection.immutable.Map<String, GpuMetric> metrics;
@@ -97,6 +98,7 @@ public GpuParquetReader(
       Map<Integer, ?> idToConstant, GpuDeleteFilter deleteFilter,
       PartitionedFile partFile, Configuration conf, int maxBatchSizeRows,
       long maxBatchSizeBytes, long targetBatchSizeBytes, boolean useChunkedReader,
+      boolean useSubPageChunked,
       scala.Option<String> debugDumpPrefix, boolean debugDumpAlways,
       scala.collection.immutable.Map<String, GpuMetric> metrics) {
     this.input = input;
@@ -113,6 +115,7 @@ public GpuParquetReader(
     this.maxBatchSizeBytes = maxBatchSizeBytes;
     this.targetBatchSizeBytes = targetBatchSizeBytes;
     this.useChunkedReader = useChunkedReader;
+    this.useSubPageChunked = useSubPageChunked;
     this.debugDumpPrefix = debugDumpPrefix;
     this.debugDumpAlways = debugDumpAlways;
     this.metrics = metrics;
@@ -139,7 +142,9 @@ public org.apache.iceberg.io.CloseableIterator<ColumnarBatch> iterator() {
       ParquetPartitionReader parquetPartReader = new ParquetPartitionReader(conf, partFile,
           new Path(input.location()), clippedBlocks, fileReadSchema, caseSensitive,
           partReaderSparkSchema, debugDumpPrefix, debugDumpAlways,
-          maxBatchSizeRows, maxBatchSizeBytes, targetBatchSizeBytes, useChunkedReader, metrics,
+          maxBatchSizeRows, maxBatchSizeBytes, targetBatchSizeBytes, useChunkedReader,
+          useSubPageChunked,
+          metrics,
           DateTimeRebaseCorrected$.MODULE$, // dateRebaseMode
           DateTimeRebaseCorrected$.MODULE$, // timestampRebaseMode
           true, // hasInt96Timestamps

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -48,13 +48,14 @@ class GpuBatchDataReader extends BaseDataReader<ColumnarBatch> {
   private final long maxBatchSizeBytes;
   private final long targetBatchSizeBytes;
   private final boolean useChunkedReader;
+  private final boolean useSubPageChunked;
   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,
+                     long targetBatchSizeBytes, boolean useChunkedReader, boolean useSubPageChunked,
                      scala.Option<String> parquetDebugDumpPrefix, boolean parquetDebugDumpAlways,
                      scala.collection.immutable.Map<String, GpuMetric> metrics) {
     super(table, task);
@@ -66,6 +67,7 @@ class GpuBatchDataReader extends BaseDataReader<ColumnarBatch> {
     this.maxBatchSizeBytes = maxBatchSizeBytes;
     this.targetBatchSizeBytes = targetBatchSizeBytes;
     this.useChunkedReader = useChunkedReader;
+    this.useSubPageChunked = useSubPageChunked;
     this.parquetDebugDumpPrefix = parquetDebugDumpPrefix;
     this.parquetDebugDumpAlways = parquetDebugDumpAlways;
     this.metrics = metrics;
@@ -100,7 +102,7 @@ CloseableIterator<ColumnarBatch> open(FileScanTask task) {
           .withMaxBatchSizeRows(maxBatchSizeRows)
           .withMaxBatchSizeBytes(maxBatchSizeBytes)
           .withTargetBatchSizeBytes(targetBatchSizeBytes)
-          .withUseChunkedReader(useChunkedReader)
+          .withUseChunkedReader(useChunkedReader, useSubPageChunked)
           .withDebugDump(parquetDebugDumpPrefix, parquetDebugDumpAlways)
           .withMetrics(metrics);
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -68,6 +68,7 @@ class GpuMultiFileBatchReader extends BaseDataReader<ColumnarBatch> {
   private final long maxGpuColumnSizeBytes;
 
   private final boolean useChunkedReader;
+  private final boolean useSubPageChunked;
   private final scala.Option<String> parquetDebugDumpPrefix;
   private final boolean parquetDebugDumpAlways;
   private final scala.collection.immutable.Map<String, GpuMetric> metrics;
@@ -86,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 useChunkedReader, boolean useSubPageChunked,
       scala.Option<String> parquetDebugDumpPrefix, boolean parquetDebugDumpAlways,
       int numThreads, int maxNumFileProcessed,
       boolean useMultiThread, FileFormat fileFormat,
@@ -101,6 +102,7 @@ class GpuMultiFileBatchReader extends BaseDataReader<ColumnarBatch> {
     this.targetBatchSizeBytes = targetBatchSizeBytes;
     this.maxGpuColumnSizeBytes = maxGpuColumnSizeBytes;
     this.useChunkedReader = useChunkedReader;
+    this.useSubPageChunked = useSubPageChunked;
     this.parquetDebugDumpPrefix = parquetDebugDumpPrefix;
     this.parquetDebugDumpAlways = parquetDebugDumpAlways;
     this.useMultiThread = useMultiThread;
@@ -350,7 +352,7 @@ protected FilePartitionReaderBase createRapidsReader(PartitionedFile[] pFiles,
       return new MultiFileCloudParquetPartitionReader(conf, pFiles,
           this::filterParquetBlocks, caseSensitive, parquetDebugDumpPrefix, parquetDebugDumpAlways,
           maxBatchSizeRows, maxBatchSizeBytes, targetBatchSizeBytes, maxGpuColumnSizeBytes,
-          useChunkedReader, metrics, partitionSchema,
+          useChunkedReader, useSubPageChunked, metrics, partitionSchema,
           numThreads, maxNumFileProcessed,
           false, // ignoreMissingFiles
           false, // ignoreCorruptFiles
@@ -427,6 +429,7 @@ protected FilePartitionReaderBase createRapidsReader(PartitionedFile[] pFiles,
       return new MultiFileParquetPartitionReader(conf, pFiles,
           JavaConverters.asScalaBuffer(clippedBlocks).toSeq(),
           caseSensitive, parquetDebugDumpPrefix, parquetDebugDumpAlways, useChunkedReader,
+          useSubPageChunked,
           maxBatchSizeRows, maxBatchSizeBytes, targetBatchSizeBytes, maxGpuColumnSizeBytes,
           metrics, partitionSchema, numThreads,
           false, // ignoreMissingFiles

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -283,6 +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.getParquetDebugDumpPrefix(), task.getParquetDebugDumpAlways(),
           task.getNumThreads(), task.getMaxNumFileProcessed(),
           useMultiThread, ff, metrics, queryUsesInputFile);
@@ -293,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.getTargetBatchSizeBytes(), task.useChunkedReader(), task.useSubPageChunked(),
           task.getParquetDebugDumpPrefix(), task.getParquetDebugDumpAlways(), metrics);
     }
   }
@@ -304,6 +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 Broadcast<SerializableConfiguration> confBroadcast;
     private final int maxBatchSizeRows;
     private final long maxBatchSizeBytes;
@@ -341,6 +343,7 @@ static class ReadTask implements InputPartition, Serializable {
       this.numThreads = rapidsConf.multiThreadReadNumThreads();
       this.maxNumFileProcessed = rapidsConf.maxNumParquetFilesParallel();
       this.useChunkedReader = rapidsConf.chunkedReaderEnabled();
+      this.useSubPageChunked = rapidsConf.chunkedSubPageReaderEnabled();
     }
 
     @Override
@@ -406,5 +409,9 @@ private Schema expectedSchema() {
     public boolean useChunkedReader() {
       return useChunkedReader;
     }
+
+    public boolean useSubPageChunked() {
+      return useSubPageChunked;
+    }
   }
 }