feat(cube): Add and use split_row_group_reads in ParquetOptions

Author: srh

Cargo.lock

@@ -452,6 +452,9 @@ config_namespace! {
         /// BLOB instead.
         pub binary_as_string: bool, default = false
 
+        /// (reading) Allows multiple I/O's per row group, if the batch size is small enough.
+        pub split_row_group_reads: bool, default = false
+
         // The following options affect writing to parquet files
         // and map to parquet::file::properties::WriterProperties
 

datafusion/common/src/config.rs

@@ -272,6 +272,7 @@ impl ParquetOptions {
             bloom_filter_on_read: _, // reads not used for writer props
             schema_force_view_types: _,
             binary_as_string: _, // not used for writer props
+            split_row_group_reads: _,
             skip_arrow_metadata: _,
         } = self;
 
@@ -548,6 +549,7 @@ mod tests {
             bloom_filter_on_read: defaults.bloom_filter_on_read,
             schema_force_view_types: defaults.schema_force_view_types,
             binary_as_string: defaults.binary_as_string,
+            split_row_group_reads: defaults.split_row_group_reads,
             skip_arrow_metadata: defaults.skip_arrow_metadata,
         }
     }
@@ -654,6 +656,7 @@ mod tests {
                 bloom_filter_on_read: global_options_defaults.bloom_filter_on_read,
                 schema_force_view_types: global_options_defaults.schema_force_view_types,
                 binary_as_string: global_options_defaults.binary_as_string,
+                split_row_group_reads: global_options_defaults.split_row_group_reads,
                 skip_arrow_metadata: global_options_defaults.skip_arrow_metadata,
             },
             column_specific_options,

datafusion/common/src/file_options/parquet_writer.rs

@@ -55,6 +55,8 @@ pub(super) struct ParquetOpener {
     pub projection: Arc<[usize]>,
     /// Target number of rows in each output RecordBatch
     pub batch_size: usize,
+    /// If true, perform multiple I/O's per row group (if batch size is smaller than row group size)
+    pub split_row_group_reads: bool,
     /// Optional limit on the number of rows to read
     pub limit: Option<usize>,
     /// Optional predicate to apply during the scan
@@ -107,6 +109,7 @@ impl FileOpener for ParquetOpener {
             )?;
 
         let batch_size = self.batch_size;
+        let split_row_group_reads = self.split_row_group_reads;
 
         let projected_schema =
             SchemaRef::from(self.table_schema.project(&self.projection)?);
@@ -262,6 +265,7 @@ impl FileOpener for ParquetOpener {
             let stream = builder
                 .with_projection(mask)
                 .with_batch_size(batch_size)
+                .with_split_row_group_reads(split_row_group_reads)
                 .with_row_groups(row_group_indexes)
                 .build()?;
 

datafusion/core/src/datasource/physical_plan/parquet/opener.rs

@@ -515,6 +515,7 @@ impl FileSource for ParquetSource {
             batch_size: self
                 .batch_size
                 .expect("Batch size must set before creating ParquetOpener"),
+            split_row_group_reads: self.table_parquet_options.global.split_row_group_reads,
             limit: base_config.limit,
             predicate: self.predicate.clone(),
             pruning_predicate: self.pruning_predicate.clone(),

datafusion/core/src/datasource/physical_plan/parquet/source.rs

@@ -497,6 +497,7 @@ message ParquetOptions {
   bool schema_force_view_types = 28; // default = false
   bool binary_as_string = 29; // default = false
   bool skip_arrow_metadata = 30; // default = false
+  bool split_row_group_reads = 32; // default = false
 
   oneof metadata_size_hint_opt {
     uint64 metadata_size_hint = 4;

datafusion/proto-common/proto/datafusion_common.proto

@@ -984,6 +984,7 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions {
             maximum_buffered_record_batches_per_stream: value.maximum_buffered_record_batches_per_stream as usize,
             schema_force_view_types: value.schema_force_view_types,
             binary_as_string: value.binary_as_string,
+            split_row_group_reads: value.split_row_group_reads,
             skip_arrow_metadata: value.skip_arrow_metadata,
         })
     }

datafusion/proto-common/src/from_proto/mod.rs

@@ -4939,6 +4939,9 @@ impl serde::Serialize for ParquetOptions {
         if self.skip_arrow_metadata {
             len += 1;
         }
+        if self.split_row_group_reads {
+            len += 1;
+        }
         if self.dictionary_page_size_limit != 0 {
             len += 1;
         }
@@ -5038,6 +5041,9 @@ impl serde::Serialize for ParquetOptions {
         if self.skip_arrow_metadata {
             struct_ser.serialize_field("skipArrowMetadata", &self.skip_arrow_metadata)?;
         }
+        if self.split_row_group_reads {
+            struct_ser.serialize_field("splitRowGroupReads", &self.split_row_group_reads)?;
+        }
         if self.dictionary_page_size_limit != 0 {
             #[allow(clippy::needless_borrow)]
             #[allow(clippy::needless_borrows_for_generic_args)]
@@ -5177,6 +5183,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions {
             "binaryAsString",
             "skip_arrow_metadata",
             "skipArrowMetadata",
+            "split_row_group_reads",
+            "splitRowGroupReads",
             "dictionary_page_size_limit",
             "dictionaryPageSizeLimit",
             "data_page_row_count_limit",
@@ -5223,6 +5231,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions {
             SchemaForceViewTypes,
             BinaryAsString,
             SkipArrowMetadata,
+            SplitRowGroupReads,
             DictionaryPageSizeLimit,
             DataPageRowCountLimit,
             MaxRowGroupSize,
@@ -5274,6 +5283,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions {
                             "schemaForceViewTypes" | "schema_force_view_types" => Ok(GeneratedField::SchemaForceViewTypes),
                             "binaryAsString" | "binary_as_string" => Ok(GeneratedField::BinaryAsString),
                             "skipArrowMetadata" | "skip_arrow_metadata" => Ok(GeneratedField::SkipArrowMetadata),
+                            "splitRowGroupReads" | "split_row_group_reads" => Ok(GeneratedField::SplitRowGroupReads),
                             "dictionaryPageSizeLimit" | "dictionary_page_size_limit" => Ok(GeneratedField::DictionaryPageSizeLimit),
                             "dataPageRowCountLimit" | "data_page_row_count_limit" => Ok(GeneratedField::DataPageRowCountLimit),
                             "maxRowGroupSize" | "max_row_group_size" => Ok(GeneratedField::MaxRowGroupSize),
@@ -5323,6 +5333,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions {
                 let mut schema_force_view_types__ = None;
                 let mut binary_as_string__ = None;
                 let mut skip_arrow_metadata__ = None;
+                let mut split_row_group_reads__ = None;
                 let mut dictionary_page_size_limit__ = None;
                 let mut data_page_row_count_limit__ = None;
                 let mut max_row_group_size__ = None;
@@ -5443,6 +5454,12 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions {
                             }
                             skip_arrow_metadata__ = Some(map_.next_value()?);
                         }
+                        GeneratedField::SplitRowGroupReads => {
+                            if split_row_group_reads__.is_some() {
+                                return Err(serde::de::Error::duplicate_field("splitRowGroupReads"));
+                            }
+                            split_row_group_reads__ = Some(map_.next_value()?);
+                        }
                         GeneratedField::DictionaryPageSizeLimit => {
                             if dictionary_page_size_limit__.is_some() {
                                 return Err(serde::de::Error::duplicate_field("dictionaryPageSizeLimit"));
@@ -5552,6 +5569,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions {
                     schema_force_view_types: schema_force_view_types__.unwrap_or_default(),
                     binary_as_string: binary_as_string__.unwrap_or_default(),
                     skip_arrow_metadata: skip_arrow_metadata__.unwrap_or_default(),
+                    split_row_group_reads: split_row_group_reads__.unwrap_or_default(),
                     dictionary_page_size_limit: dictionary_page_size_limit__.unwrap_or_default(),
                     data_page_row_count_limit: data_page_row_count_limit__.unwrap_or_default(),
                     max_row_group_size: max_row_group_size__.unwrap_or_default(),

datafusion/proto-common/src/generated/pbjson.rs

@@ -764,6 +764,9 @@ pub struct ParquetOptions {
     /// default = false
     #[prost(bool, tag = "30")]
     pub skip_arrow_metadata: bool,
+    /// default = false
+    #[prost(bool, tag = "32")]
+    pub split_row_group_reads: bool,
     #[prost(uint64, tag = "12")]
     pub dictionary_page_size_limit: u64,
     #[prost(uint64, tag = "18")]

datafusion/proto-common/src/generated/prost.rs

@@ -835,6 +835,7 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions {
             maximum_buffered_record_batches_per_stream: value.maximum_buffered_record_batches_per_stream as u64,
             schema_force_view_types: value.schema_force_view_types,
             binary_as_string: value.binary_as_string,
+            split_row_group_reads: value.split_row_group_reads,
             skip_arrow_metadata: value.skip_arrow_metadata,
         })
     }