refactor: Enables the new setting for Query only

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<Vec<DataBlock>> {
+    pub fn deserialize_native_data(&self, data: &mut NativeSourceData) -> Result<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(vec![DataBlock::empty_with_schema(Arc::new(
+            return Ok(DataBlock::empty_with_schema(Arc::new(
                 self.reader.data_schema(),
-            ))]);
+            )));
         }
         debug_assert!(all_columns_arrays
             .iter()
@@ -167,6 +167,7 @@ impl AggIndexReader {
             let block = DataBlock::new_from_columns(columns);
             blocks.push(block);
         }
-        self.apply_agg_info(blocks)
+        let block = DataBlock::concat(&blocks)?;
+        self.apply_agg_info_to_block(block)
     }
 }

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

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

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

@@ -35,7 +35,7 @@ pub fn column_chunks_to_record_batch(
     num_rows: usize,
     column_chunks: &HashMap<ColumnId, DataItem>,
     compression: &Compression,
-    batch_size: usize,
+    batch_size: Option<usize>,
 ) -> databend_common_exception::Result<Vec<RecordBatch>> {
     let arrow_schema = Schema::from(original_schema);
     let parquet_schema = ArrowSchemaConverter::new().convert(&arrow_schema)?;
@@ -68,6 +68,8 @@ pub fn column_chunks_to_record_batch(
         ProjectionMask::leaves(&parquet_schema, projection_mask),
         Some(arrow_schema.fields()),
     )?;
+
+    let batch_size = batch_size.unwrap_or(num_rows);
     let record_reader = ParquetRecordBatchReader::try_new_with_row_groups(
         &field_levels,
         row_group.as_ref(),

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

@@ -56,7 +56,7 @@ impl BlockReader {
             column_chunks,
             compression,
             block_path,
-            num_rows,
+            None,
         )?;
         // Defensive check: using `num_rows` as batch_size, expects only one block
         assert_eq!(blocks.len(), 1);
@@ -70,7 +70,7 @@ impl BlockReader {
         column_chunks: HashMap<ColumnId, DataItem>,
         compression: &Compression,
         block_path: &str,
-        batch_size: usize,
+        batch_size_hint: Option<usize>,
     ) -> Result<Vec<DataBlock>> {
         if column_chunks.is_empty() {
             return Ok(vec![self.build_default_values_block(num_rows)?]);
@@ -81,7 +81,7 @@ impl BlockReader {
             num_rows,
             &column_chunks,
             compression,
-            batch_size,
+            batch_size_hint,
         )?;
 
         let name_paths = column_name_paths(&self.projection, &self.original_schema);

src/query/storages/fuse/src/io/read/virtual_column/virtual_column_reader_parquet.rs

@@ -136,7 +136,7 @@ impl VirtualColumnReader {
     pub fn try_create_paster(
         &self,
         virtual_data: Option<VirtualBlockReadResult>,
-        batch_size: usize,
+        batch_size_hint: Option<usize>,
     ) -> Result<VirtualColumnDataPaster> {
         let record_batches = if let Some(virtual_data) = virtual_data {
             let columns_chunks = virtual_data.data.columns_chunks()?;
@@ -145,7 +145,7 @@ impl VirtualColumnReader {
                 virtual_data.num_rows,
                 &columns_chunks,
                 &virtual_data.compression,
-                batch_size,
+                batch_size_hint,
             )?;
             Some(chunks)
         } else {

src/query/storages/fuse/src/operations/merge_into/mutator/matched_mutator.rs

@@ -201,7 +201,7 @@ impl MatchedAggregator {
                         .insert(offset as usize)
                     {
                         return Err(ErrorCode::UnresolvableConflict(
-                            "multi rows from source match one and the same row in the target_table multi times",
+                            "1 multi rows from source match one and the same row in the target_table multi times",
                         ));
                     }
                 }
@@ -335,7 +335,7 @@ impl MatchedAggregator {
                 < update_modified_offsets.len() + delete_modified_offsets.len()
             {
                 return Err(ErrorCode::UnresolvableConflict(
-                    "multi rows from source match one and the same row in the target_table multi times",
+                    "2 multi rows from source match one and the same row in the target_table multi times",
                 ));
             }
 

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: Vec<DataBlock>,
+    output_data: Option<DataBlock>,
     parts: VecDeque<PartInfoPtr>,
     columns: VecDeque<NativeDataSource>,
     scan_progress: Arc<Progress>,
@@ -369,7 +369,7 @@ impl NativeDeserializeDataTransform {
                 block_reader,
                 input,
                 output,
-                output_data: vec![],
+                output_data: None,
                 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 = vec![data_block];
+        self.output_data = Some(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.pop() {
+        if let Some(data_block) = self.output_data.take() {
             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 blocks = agg_index_reader.deserialize_native_data(data)?;
-                        self.output_data = blocks;
+                        let block = agg_index_reader.deserialize_native_data(data)?;
+                        self.output_data = Some(block);
                         self.finish_partition();
                         return Ok(());
                     }

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

@@ -23,6 +23,7 @@ use databend_common_base::runtime::profile::Profile;
 use databend_common_base::runtime::profile::ProfileStatisticsName;
 use databend_common_catalog::plan::DataSourcePlan;
 use databend_common_catalog::plan::PartInfoPtr;
+use databend_common_catalog::query_kind::QueryKind;
 use databend_common_catalog::runtime_filter_info::RuntimeFilterReady;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
@@ -80,7 +81,7 @@ pub struct DeserializeDataTransform {
     need_wait_runtime_filter: bool,
     runtime_filter_ready: Option<Arc<RuntimeFilterReady>>,
 
-    batch_size: usize,
+    batch_size_hint: Option<usize>,
 }
 
 unsafe impl Send for DeserializeDataTransform {}
@@ -99,6 +100,12 @@ impl DeserializeDataTransform {
         let need_wait_runtime_filter =
             !ctx.get_cluster().is_empty() && ctx.get_wait_runtime_filter(plan.scan_id);
 
+        // Unfortunately, batch size is hint is only safe for Query now.
+        let batch_size_hint = match ctx.get_query_kind() {
+            QueryKind::Query => Some(ctx.get_settings().get_fuse_parquet_read_batch_size()?),
+            _ => None,
+        };
+
         let mut src_schema: DataSchema = (block_reader.schema().as_ref()).into();
         if let Some(virtual_reader) = virtual_reader.as_ref() {
             let mut fields = src_schema.fields().clone();
@@ -117,7 +124,6 @@ impl DeserializeDataTransform {
         let output_schema: DataSchema = (&output_schema).into();
         let (need_reserve_block_info, _) = need_reserve_block_info(ctx.clone(), plan.table_index);
 
-        let batch_size = ctx.get_settings().get_fuse_parquet_read_batch_size()?;
         Ok(ProcessorPtr::create(Box::new(DeserializeDataTransform {
             ctx,
             table_index: plan.table_index,
@@ -138,7 +144,7 @@ impl DeserializeDataTransform {
             need_reserve_block_info,
             need_wait_runtime_filter,
             runtime_filter_ready: None,
-            batch_size,
+            batch_size_hint,
         })))
     }
 
@@ -270,7 +276,7 @@ impl Processor for DeserializeDataTransform {
                     let blocks = agg_index_reader.deserialize_parquet_data(
                         actual_part,
                         data,
-                        self.batch_size,
+                        self.batch_size_hint,
                     )?;
 
                     self.update_scan_metrics(blocks.as_slice());
@@ -289,13 +295,13 @@ impl Processor for DeserializeDataTransform {
                         columns_chunks,
                         &part.compression,
                         &part.location,
-                        self.batch_size,
+                        self.batch_size_hint,
                     )?;
 
                     let mut virtual_columns_paster =
                         if let Some(virtual_column_reader) = self.virtual_reader.as_ref() {
                             let record_batches = virtual_column_reader
-                                .try_create_paster(virtual_data, self.batch_size)?;
+                                .try_create_paster(virtual_data, self.batch_size_hint)?;
                             Some(record_batches)
                         } else {
                             None