feat: new settings fuse_parquet_read_batch_size

Which controls the bach size during deserializing of fuse parquet data
block. The default value of this setting is 8192.

Author: dantengsky

src/query/settings/src/settings_default.rs

@@ -1242,6 +1242,13 @@ impl DefaultSettings {
                     scope: SettingScope::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
                 }),
+                ("fuse_parquet_read_batch_size", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(8192),
+                    desc: "The batch size while deserializing fuse table with parquet storage format",
+                    mode: SettingMode::Both,
+                    scope: SettingScope::Both,
+                    range: Some(SettingRange::Numeric(0..=1_000_000)),
+                }),
             ]);
 
             Ok(Arc::new(DefaultSettings {

src/query/settings/src/settings_getter_setter.rs

@@ -924,4 +924,8 @@ impl Settings {
     pub fn get_enable_use_vacuum2_to_purge_transient_table_data(&self) -> Result<bool> {
         Ok(self.try_get_u64("use_vacuum2_to_purge_transient_table_data")? == 1)
     }
+
+    pub fn get_fuse_parquet_read_batch_size(&self) -> Result<usize> {
+        Ok(self.try_get_u64("fuse_parquet_read_batch_size")? as usize)
+    }
 }

src/query/storages/fuse/src/io/read/agg_index/agg_index_reader.rs

@@ -97,7 +97,7 @@ impl AggIndexReader {
         self.index_id
     }
 
-    pub(super) fn apply_agg_info(&self, block: DataBlock) -> Result<DataBlock> {
+    pub(super) fn apply_agg_info_to_block(&self, block: DataBlock) -> Result<DataBlock> {
         let evaluator = Evaluator::new(&block, &self.func_ctx, &BUILTIN_FUNCTIONS);
 
         // 1. Filter the block if there is a filter.
@@ -145,4 +145,11 @@ impl AggIndexReader {
             )),
         ))
     }
+
+    pub(super) fn apply_agg_info(&self, block: Vec<DataBlock>) -> Result<Vec<DataBlock>> {
+        block
+            .into_iter()
+            .map(|block| self.apply_agg_info_to_block(block))
+            .collect::<Result<_>>()
+    }
 }

src/query/storages/fuse/src/io/read/agg_index/agg_index_reader_native.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::sync::Arc;
+use std::vec;
 
 use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
@@ -138,7 +139,7 @@ impl AggIndexReader {
         }
     }
 
-    pub fn deserialize_native_data(&self, data: &mut NativeSourceData) -> Result<DataBlock> {
+    pub fn deserialize_native_data(&self, data: &mut NativeSourceData) -> Result<Vec<DataBlock>> {
         let mut all_columns_arrays = vec![];
 
         for (index, column_node) in self.reader.project_column_nodes.iter().enumerate() {
@@ -148,9 +149,9 @@ impl AggIndexReader {
             all_columns_arrays.push(arrays);
         }
         if all_columns_arrays.is_empty() {
-            return Ok(DataBlock::empty_with_schema(Arc::new(
+            return Ok(vec![DataBlock::empty_with_schema(Arc::new(
                 self.reader.data_schema(),
-            )));
+            ))]);
         }
         debug_assert!(all_columns_arrays
             .iter()
@@ -166,7 +167,6 @@ impl AggIndexReader {
             let block = DataBlock::new_from_columns(columns);
             blocks.push(block);
         }
-        let block = DataBlock::concat(&blocks)?;
-        self.apply_agg_info(block)
+        self.apply_agg_info(blocks)
     }
 }

src/query/storages/fuse/src/io/read/agg_index/agg_index_reader_parquet.rs

@@ -113,15 +113,17 @@ impl AggIndexReader {
         &self,
         part: PartInfoPtr,
         data: BlockReadResult,
-    ) -> Result<DataBlock> {
+        batch_size: usize,
+    ) -> Result<Vec<DataBlock>> {
         let columns_chunks = data.columns_chunks()?;
         let part = FuseBlockPartInfo::from_part(&part)?;
-        let block = self.reader.deserialize_parquet_chunks(
+        let block = self.reader.deserialize_parquet_to_blocks(
             part.nums_rows,
             &part.columns_meta,
             columns_chunks,
             &part.compression,
             &part.location,
+            batch_size,
         )?;
 
         self.apply_agg_info(block)

src/query/storages/fuse/src/io/read/block/parquet/deserialize.rs

@@ -19,6 +19,7 @@ use arrow_schema::Schema;
 use databend_common_expression::ColumnId;
 use databend_common_expression::TableSchema;
 use databend_storages_common_table_meta::meta::Compression;
+use itertools::Itertools;
 use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
 use parquet::arrow::parquet_to_arrow_field_levels;
 use parquet::arrow::ArrowSchemaConverter;
@@ -34,7 +35,8 @@ pub fn column_chunks_to_record_batch(
     num_rows: usize,
     column_chunks: &HashMap<ColumnId, DataItem>,
     compression: &Compression,
-) -> databend_common_exception::Result<RecordBatch> {
+    batch_size: usize,
+) -> databend_common_exception::Result<Vec<RecordBatch>> {
     let arrow_schema = Schema::from(original_schema);
     let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema)?;
 
@@ -66,13 +68,17 @@ pub fn column_chunks_to_record_batch(
         ProjectionMask::leaves(&parquet_schema, projection_mask),
         Some(arrow_schema.fields()),
     )?;
-    let mut record_reader = ParquetRecordBatchReader::try_new_with_row_groups(
+    let record_reader = ParquetRecordBatchReader::try_new_with_row_groups(
         &field_levels,
         row_group.as_ref(),
-        num_rows,
+        batch_size,
         None,
     )?;
-    let record = record_reader.next().unwrap()?;
-    assert!(record_reader.next().is_none());
-    Ok(record)
+
+    let records: Vec<_> = record_reader.try_collect()?;
+    assert_eq!(
+        num_rows,
+        records.iter().map(|r| r.num_rows()).sum::<usize>()
+    );
+    Ok(records)
 }

src/query/storages/fuse/src/io/read/block/parquet/mod.rs

@@ -35,6 +35,7 @@ mod adapter;
 mod deserialize;
 
 pub use adapter::RowGroupImplBuilder;
+use databend_common_exception::Result;
 pub use deserialize::column_chunks_to_record_batch;
 
 use crate::io::read::block::block_reader_merge_io::DataItem;
@@ -48,17 +49,41 @@ impl BlockReader {
         column_chunks: HashMap<ColumnId, DataItem>,
         compression: &Compression,
         block_path: &str,
-    ) -> databend_common_exception::Result<DataBlock> {
+    ) -> Result<DataBlock> {
+        let mut blocks = self.deserialize_parquet_to_blocks(
+            num_rows,
+            column_metas,
+            column_chunks,
+            compression,
+            block_path,
+            num_rows,
+        )?;
+        // Defensive check: using `num_rows` as batch_size, expects only one block
+        assert_eq!(blocks.len(), 1);
+        Ok(blocks.pop().unwrap())
+    }
+
+    pub(crate) fn deserialize_parquet_to_blocks(
+        &self,
+        num_rows: usize,
+        column_metas: &HashMap<ColumnId, ColumnMeta>,
+        column_chunks: HashMap<ColumnId, DataItem>,
+        compression: &Compression,
+        block_path: &str,
+        batch_size: usize,
+    ) -> Result<Vec<DataBlock>> {
         if column_chunks.is_empty() {
-            return self.build_default_values_block(num_rows);
+            return Ok(vec![self.build_default_values_block(num_rows)?]);
         }
-        let record_batch = column_chunks_to_record_batch(
+
+        let record_batches = column_chunks_to_record_batch(
             &self.original_schema,
             num_rows,
             &column_chunks,
             compression,
+            batch_size,
         )?;
-        let mut columns = Vec::with_capacity(self.projected_schema.fields.len());
+
         let name_paths = column_name_paths(&self.projection, &self.original_schema);
 
         let array_cache = if self.put_cache {
@@ -67,58 +92,71 @@ impl BlockReader {
             None
         };
 
-        for ((i, field), column_node) in self
-            .projected_schema
-            .fields
-            .iter()
-            .enumerate()
-            .zip(self.project_column_nodes.iter())
-        {
-            let data_type = field.data_type().into();
-
-            // NOTE, there is something tricky here:
-            // - `column_chunks` always contains data of leaf columns
-            // - here we may processing a nested type field
-            // - But, even if the field being processed is a field with multiple leaf columns
-            //    `column_chunks.get(&field.column_id)` will still return Some(DataItem::_)[^1],
-            //    even if we are getting data from `column_chunks` using a non-leaf
-            //    `column_id` of `projected_schema.fields`
-            //
-            //   [^1]: Except in the current block, there is no data stored for the
-            //         corresponding field, and a default value has been declared for
-            //         the corresponding field.
-            //
-            //  Yes, it is too obscure, we need to polish it later.
-
-            let value = match column_chunks.get(&field.column_id) {
-                Some(DataItem::RawData(data)) => {
-                    // get the deserialized arrow array, which may be a nested array
-                    let arrow_array = column_by_name(&record_batch, &name_paths[i]);
-                    if !column_node.is_nested {
-                        if let Some(cache) = &array_cache {
-                            let meta = column_metas.get(&field.column_id).unwrap();
-                            let (offset, len) = meta.offset_length();
-                            let key =
-                                TableDataCacheKey::new(block_path, field.column_id, offset, len);
-                            cache.insert(key.into(), (arrow_array.clone(), data.len()));
+        let mut blocks = Vec::with_capacity(record_batches.len());
+
+        for record_batch in record_batches {
+            let num_rows_record_batch = record_batch.num_rows();
+            let mut columns = Vec::with_capacity(self.projected_schema.fields.len());
+            for ((i, field), column_node) in self
+                .projected_schema
+                .fields
+                .iter()
+                .enumerate()
+                .zip(self.project_column_nodes.iter())
+            {
+                let data_type = field.data_type().into();
+
+                // NOTE, there is something tricky here:
+                // - `column_chunks` always contains data of leaf columns
+                // - here we may processing a nested type field
+                // - But, even if the field being processed is a field with multiple leaf columns
+                //    `column_chunks.get(&field.column_id)` will still return Some(DataItem::_)[^1],
+                //    even if we are getting data from `column_chunks` using a non-leaf
+                //    `column_id` of `projected_schema.fields`
+                //
+                //   [^1]: Except in the current block, there is no data stored for the
+                //         corresponding field, and a default value has been declared for
+                //         the corresponding field.
+                //
+                //  Yes, it is too obscure, we need to polish it later.
+
+                let value = match column_chunks.get(&field.column_id) {
+                    Some(DataItem::RawData(data)) => {
+                        // get the deserialized arrow array, which may be a nested array
+                        let arrow_array = column_by_name(&record_batch, &name_paths[i]);
+                        if !column_node.is_nested {
+                            if let Some(cache) = &array_cache {
+                                let meta = column_metas.get(&field.column_id).unwrap();
+                                let (offset, len) = meta.offset_length();
+                                let key = TableDataCacheKey::new(
+                                    block_path,
+                                    field.column_id,
+                                    offset,
+                                    len,
+                                );
+                                cache.insert(key.into(), (arrow_array.clone(), data.len()));
+                            }
                         }
+                        Value::from_arrow_rs(arrow_array, &data_type)?
                     }
-                    Value::from_arrow_rs(arrow_array, &data_type)?
-                }
-                Some(DataItem::ColumnArray(cached)) => {
-                    if column_node.is_nested {
-                        // a defensive check, should never happen
-                        return Err(ErrorCode::StorageOther(
-                            "unexpected nested field: nested leaf field hits cached",
-                        ));
+                    Some(DataItem::ColumnArray(cached)) => {
+                        // TODO this is NOT correct!
+                        if column_node.is_nested {
+                            // a defensive check, should never happen
+                            return Err(ErrorCode::StorageOther(
+                                "unexpected nested field: nested leaf field hits cached",
+                            ));
+                        }
+                        Value::from_arrow_rs(cached.0.clone(), &data_type)?
                     }
-                    Value::from_arrow_rs(cached.0.clone(), &data_type)?
-                }
-                None => Value::Scalar(self.default_vals[i].clone()),
-            };
-            columns.push(BlockEntry::new(data_type, value));
+                    None => Value::Scalar(self.default_vals[i].clone()),
+                };
+                columns.push(BlockEntry::new(data_type, value));
+            }
+            blocks.push(DataBlock::new(columns, num_rows_record_batch));
         }
-        Ok(DataBlock::new(columns, num_rows))
+
+        Ok(blocks)
     }
 }