read batch size for native format deser

Author: dantengsky

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

@@ -139,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() {
@@ -149,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()
@@ -167,8 +167,6 @@ impl AggIndexReader {
             let block = DataBlock::new_from_columns(columns);
             blocks.push(block);
         }
-        let blocks = self.apply_agg_info(blocks)?;
-
-        DataBlock::concat(blocks.as_slice())
+        self.apply_agg_info(blocks)
     }
 }

src/query/storages/fuse/src/operations/read/native_data_source_deserializer.rs

@@ -194,7 +194,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>,
@@ -369,7 +369,7 @@ impl NativeDeserializeDataTransform {
                 block_reader,
                 input,
                 output,
-                output_data: None,
+                output_data: vec![],
                 parts: VecDeque::new(),
                 columns: VecDeque::new(),
                 prewhere_columns,
@@ -417,7 +417,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];
     }
 
     /// If the virtual column has been already generated, add it directly,
@@ -1049,7 +1049,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);
         }
@@ -1094,8 +1094,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(());
                     }

src/query/storages/fuse/src/operations/read/parquet_data_source_deserializer.rs

@@ -225,11 +225,8 @@ impl Processor for DeserializeDataTransform {
             return Ok(Event::NeedConsume);
         }
 
-        while let Some(data_block) = self.output_data.pop() {
+        if let Some(data_block) = self.output_data.pop() {
             self.output.push_data(Ok(data_block));
-            if !self.output.can_push() {
-                break;
-            }
             return Ok(Event::NeedConsume);
         }