feat: new settings fuse_parquet_read_batch_size

dantengsky · dantengsky · commit d32bc162dd24 · 2025-04-21T16:59:59.000+08:00
Which controls the bach size during deserializing of fuse parquet data
block. The default value of this setting is 8192.
diff --git a/src/query/settings/src/settings_default.rs b/src/query/settings/src/settings_default.rs
@@ -1264,6 +1264,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 {
diff --git a/src/query/settings/src/settings_getter_setter.rs b/src/query/settings/src/settings_getter_setter.rs
@@ -932,4 +932,8 @@ impl Settings {
     pub fn get_enable_block_stream_write(&self) -> Result<bool> {
         Ok(self.try_get_u64("enable_block_stream_write")? == 1)
     }
+
+    pub fn get_fuse_parquet_read_batch_size(&self) -> Result<usize> {
+        Ok(self.try_get_u64("fuse_parquet_read_batch_size")? as usize)
+    }
 }
diff --git a/src/query/storages/fuse/src/io/read/agg_index/agg_index_reader.rs b/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<_>>()
+    }
 }
diff --git a/src/query/storages/fuse/src/io/read/agg_index/agg_index_reader_native.rs b/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)
     }
 }
diff --git a/src/query/storages/fuse/src/io/read/agg_index/agg_index_reader_parquet.rs b/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)
diff --git a/src/query/storages/fuse/src/io/read/block/parquet/deserialize.rs b/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)
 }
diff --git a/src/query/storages/fuse/src/io/read/block/parquet/mod.rs b/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)
     }
 }
 
diff --git a/src/query/storages/fuse/src/io/read/virtual_column/virtual_column_reader_parquet.rs b/src/query/storages/fuse/src/io/read/virtual_column/virtual_column_reader_parquet.rs
@@ -15,13 +15,16 @@
 use std::collections::HashSet;
 use std::sync::Arc;
 
+use arrow_array::RecordBatch;
+use databend_common_catalog::plan::VirtualColumnField;
 use databend_common_exception::Result;
 use databend_common_expression::eval_function;
 use databend_common_expression::types::DataType;
 use databend_common_expression::BlockEntry;
 use databend_common_expression::Column;
 use databend_common_expression::ColumnId;
 use databend_common_expression::DataBlock;
+use databend_common_expression::FunctionContext;
 use databend_common_expression::TableSchema;
 use databend_common_expression::TableSchemaRef;
 use databend_common_expression::Value;
diff --git a/src/query/storages/fuse/src/operations/read/native_data_source_deserializer.rs b/src/query/storages/fuse/src/operations/read/native_data_source_deserializer.rs
@@ -191,7 +191,7 @@ pub struct NativeDeserializeDataTransform {
     // Structures for driving the pipeline:
     input: Arc<InputPort>,
     output: Arc<OutputPort>,
-    output_data: Option<DataBlock>,
+    output_data: Vec<DataBlock>,
     parts: VecDeque<PartInfoPtr>,
     columns: VecDeque<NativeDataSource>,
     scan_progress: Arc<Progress>,
@@ -309,7 +309,7 @@ impl NativeDeserializeDataTransform {
                 block_reader,
                 input,
                 output,
-                output_data: None,
+                output_data: vec![],
                 parts: VecDeque::new(),
                 columns: VecDeque::new(),
                 prewhere_columns,
@@ -353,7 +353,7 @@ impl NativeDeserializeDataTransform {
         };
         self.scan_progress.incr(&progress_values);
         Profile::record_usize_profile(ProfileStatisticsName::ScanBytes, data_block.memory_size());
-        self.output_data = Some(data_block);
+        self.output_data = vec![data_block];
     }
 
     /// Check if can skip the whole block by default values.
@@ -846,7 +846,7 @@ impl Processor for NativeDeserializeDataTransform {
             return Ok(Event::NeedConsume);
         }
 
-        if let Some(data_block) = self.output_data.take() {
+        if let Some(data_block) = self.output_data.pop() {
             self.output.push_data(Ok(data_block));
             return Ok(Event::NeedConsume);
         }
@@ -891,8 +891,8 @@ impl Processor for NativeDeserializeDataTransform {
                 let columns = match columns {
                     NativeDataSource::AggIndex(data) => {
                         let agg_index_reader = self.index_reader.as_ref().as_ref().unwrap();
-                        let block = agg_index_reader.deserialize_native_data(data)?;
-                        self.output_data = Some(block);
+                        let blocks = agg_index_reader.deserialize_native_data(data)?;
+                        self.output_data = blocks;
                         self.finish_partition();
                         return Ok(());
                     }
diff --git a/src/query/storages/fuse/src/operations/read/parquet_data_source_deserializer.rs b/src/query/storages/fuse/src/operations/read/parquet_data_source_deserializer.rs

Original file line number	Diff line number	Diff line change
`@@ -932,4 +932,8 @@ impl Settings {`
`932`	`932`	`pub fn get_enable_block_stream_write(&self) -> Result<bool> {`
`933`	`933`	`Ok(self.try_get_u64("enable_block_stream_write")? == 1)`
`934`	`934`	`}`
	`935`	`+`
	`936`	`+ pub fn get_fuse_parquet_read_batch_size(&self) -> Result<usize> {`
	`937`	`+ Ok(self.try_get_u64("fuse_parquet_read_batch_size")? as usize)`
	`938`	`+ }`
`935`	`939`	`}`
Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,7 @@ impl AggIndexReader {`
`97`	`97`	`self.index_id`
`98`	`98`	`}`
`99`	`99`
`100`		`- pub(super) fn apply_agg_info(&self, block: DataBlock) -> Result<DataBlock> {`
	`100`	`+ pub(super) fn apply_agg_info_to_block(&self, block: DataBlock) -> Result<DataBlock> {`
`101`	`101`	`let evaluator = Evaluator::new(&block, &self.func_ctx, &BUILTIN_FUNCTIONS);`
`102`	`102`
`103`	`103`	`// 1. Filter the block if there is a filter.`
`@@ -145,4 +145,11 @@ impl AggIndexReader {`
`145`	`145`	`)),`
`146`	`146`	`))`
`147`	`147`	`}`
	`148`	`+`
	`149`	`+ pub(super) fn apply_agg_info(&self, block: Vec<DataBlock>) -> Result<Vec<DataBlock>> {`
	`150`	`+ block`
	`151`	`+ .into_iter()`
	`152`	`+ .map(\|block\| self.apply_agg_info_to_block(block))`
	`153`	`+ .collect::<Result<_>>()`
	`154`	`+ }`
`148`	`155`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@`
`13`	`13`	`// limitations under the License.`
`14`	`14`
`15`	`15`	`use std::sync::Arc;`
	`16`	`+use std::vec;`
`16`	`17`
`17`	`18`	`use databend_common_exception::Result;`
`18`	`19`	`use databend_common_expression::DataBlock;`
`@@ -138,7 +139,7 @@ impl AggIndexReader {`
`138`	`139`	`}`
`139`	`140`	`}`
`140`	`141`
`141`		`- pub fn deserialize_native_data(&self, data: &mut NativeSourceData) -> Result<DataBlock> {`
	`142`	`+ pub fn deserialize_native_data(&self, data: &mut NativeSourceData) -> Result<Vec<DataBlock>> {`
`142`	`143`	`let mut all_columns_arrays = vec![];`
`143`	`144`
`144`	`145`	`for (index, column_node) in self.reader.project_column_nodes.iter().enumerate() {`
`@@ -148,9 +149,9 @@ impl AggIndexReader {`
`148`	`149`	`all_columns_arrays.push(arrays);`
`149`	`150`	`}`
`150`	`151`	`if all_columns_arrays.is_empty() {`
`151`		`- return Ok(DataBlock::empty_with_schema(Arc::new(`
	`152`	`+ return Ok(vec![DataBlock::empty_with_schema(Arc::new(`
`152`	`153`	`self.reader.data_schema(),`
`153`		`- )));`
	`154`	`+ ))]);`
`154`	`155`	`}`
`155`	`156`	`debug_assert!(all_columns_arrays`
`156`	`157`	`.iter()`
`@@ -166,7 +167,6 @@ impl AggIndexReader {`
`166`	`167`	`let block = DataBlock::new_from_columns(columns);`
`167`	`168`	`blocks.push(block);`
`168`	`169`	`}`
`169`		`- let block = DataBlock::concat(&blocks)?;`
`170`		`- self.apply_agg_info(block)`
	`170`	`+ self.apply_agg_info(blocks)`
`171`	`171`	`}`
`172`	`172`	`}`