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

@@ -149,37 +149,171 @@ impl VirtualColumnReader {
         ))
     }
 
-    pub fn deserialize_virtual_columns(
+    // pub fn deserialize_virtual_columns(
+    //    &self,
+    //    mut data_block: DataBlock,
+    //    virtual_data: Option<VirtualBlockReadResult>,
+    //) -> Result<DataBlock> {
+    //    let orig_schema = virtual_data
+    //        .as_ref()
+    //        .map(|virtual_data| virtual_data.schema.clone())
+    //        .unwrap_or_default();
+    //    let record_batch = virtual_data
+    //        .map(|virtual_data| {
+    //            let columns_chunks = virtual_data.data.columns_chunks()?;
+    //            column_chunks_to_record_batch(
+    //                &virtual_data.schema,
+    //                virtual_data.num_rows,
+    //                &columns_chunks,
+    //                &virtual_data.compression,
+    //            )
+    //        })
+    //        .transpose()?;
+
+    //    // If the virtual column has already generated, add it directly,
+    //    // otherwise extract it from the source column
+    //    let func_ctx = self.ctx.get_function_context()?;
+    //    for virtual_column_field in self.virtual_column_info.virtual_column_fields.iter() {
+    //        let name = format!("{}", virtual_column_field.column_id);
+    //        if let Some(arrow_array) = record_batch
+    //            .as_ref()
+    //            .and_then(|r| r.column_by_name(&name).cloned())
+    //        {
+    //            let orig_field = orig_schema.field_with_name(&name).unwrap();
+    //            let orig_type: DataType = orig_field.data_type().into();
+    //            let value = Value::Column(Column::from_arrow_rs(arrow_array, &orig_type)?);
+    //            let data_type: DataType = virtual_column_field.data_type.as_ref().into();
+    //            let column = if orig_type != data_type {
+    //                let cast_func_name = format!(
+    //                    "to_{}",
+    //                    data_type.remove_nullable().to_string().to_lowercase()
+    //                );
+    //                let (cast_value, cast_data_type) = eval_function(
+    //                    None,
+    //                    &cast_func_name,
+    //                    [(value, orig_type)],
+    //                    &func_ctx,
+    //                    data_block.num_rows(),
+    //                    &BUILTIN_FUNCTIONS,
+    //                )?;
+    //                BlockEntry::new(cast_data_type, cast_value)
+    //            } else {
+    //                BlockEntry::new(data_type, value)
+    //            };
+    //            data_block.add_column(column);
+    //            continue;
+    //        }
+    //        let src_index = self
+    //            .source_schema
+    //            .index_of(&virtual_column_field.source_name)
+    //            .unwrap();
+    //        let source = data_block.get_by_offset(src_index);
+    //        let src_arg = (source.value.clone(), source.data_type.clone());
+    //        let path_arg = (
+    //            Value::Scalar(virtual_column_field.key_paths.clone()),
+    //            DataType::String,
+    //        );
+
+    //        let (value, data_type) = eval_function(
+    //            None,
+    //            "get_by_keypath",
+    //            [src_arg, path_arg],
+    //            &func_ctx,
+    //            data_block.num_rows(),
+    //            &BUILTIN_FUNCTIONS,
+    //        )?;
+
+    //        let column = if let Some(cast_func_name) = &virtual_column_field.cast_func_name {
+    //            let (cast_value, cast_data_type) = eval_function(
+    //                None,
+    //                cast_func_name,
+    //                [(value, data_type)],
+    //                &func_ctx,
+    //                data_block.num_rows(),
+    //                &BUILTIN_FUNCTIONS,
+    //            )?;
+    //            BlockEntry::new(cast_data_type, cast_value)
+    //        } else {
+    //            BlockEntry::new(data_type, value)
+    //        };
+    //        data_block.add_column(column);
+    //    }
+
+    //    Ok(data_block)
+    //}
+    /// Deserialize virtual column data into record batches, according to the `batch_size`.
+    pub fn try_create_paster(
         &self,
-        mut data_block: DataBlock,
         virtual_data: Option<VirtualBlockReadResult>,
-    ) -> Result<DataBlock> {
+        batch_size_hint: Option<usize>,
+    ) -> Result<VirtualColumnDataPaster> {
         let orig_schema = virtual_data
             .as_ref()
             .map(|virtual_data| virtual_data.schema.clone())
             .unwrap_or_default();
-        let record_batch = virtual_data
-            .map(|virtual_data| {
-                let columns_chunks = virtual_data.data.columns_chunks()?;
-                column_chunks_to_record_batch(
-                    &virtual_data.schema,
-                    virtual_data.num_rows,
-                    &columns_chunks,
-                    &virtual_data.compression,
-                )
-            })
-            .transpose()?;
+
+        let record_batches = if let Some(virtual_data) = virtual_data {
+            let columns_chunks = virtual_data.data.columns_chunks()?;
+            let chunks = column_chunks_to_record_batch(
+                &self.virtual_column_info.schema,
+                virtual_data.num_rows,
+                &columns_chunks,
+                &virtual_data.compression,
+                batch_size_hint,
+            )?;
+            Some(chunks)
+        } else {
+            None
+        };
+
+        let function_context = self.ctx.get_function_context()?;
+
+        // Unfortunately, Paster cannot hold references to the fields that being cloned,
+        // since the caller `DeserializeDataTransform` will take mutable reference of
+        // VirtualColumnReader indirectly.
+        Ok(VirtualColumnDataPaster {
+            record_batches,
+            function_context,
+            next_record_batch_index: 0,
+            virtual_column_fields: self.virtual_column_info.virtual_column_fields.clone(),
+            source_schema: self.source_schema.clone(),
+            orig_schema,
+        })
+    }
+}
+
+pub struct VirtualColumnDataPaster {
+    record_batches: Option<Vec<RecordBatch>>,
+    next_record_batch_index: usize,
+    function_context: FunctionContext,
+    virtual_column_fields: Vec<VirtualColumnField>,
+    source_schema: TableSchemaRef,
+    orig_schema: TableSchemaRef,
+}
+
+impl VirtualColumnDataPaster {
+    /// Paste virtual column to `data_block` if necessary
+    pub fn paste_virtual_column(&mut self, mut data_block: DataBlock) -> Result<DataBlock> {
+        let record_batch = if let Some(record_batches) = &self.record_batches {
+            assert!(record_batches.len() > self.next_record_batch_index);
+            Some(&record_batches[self.next_record_batch_index])
+        } else {
+            None
+        };
+
+        self.next_record_batch_index += 1;
+
+        let func_ctx = &self.function_context;
 
         // If the virtual column has already generated, add it directly,
         // otherwise extract it from the source column
-        let func_ctx = self.ctx.get_function_context()?;
-        for virtual_column_field in self.virtual_column_info.virtual_column_fields.iter() {
+        for virtual_column_field in self.virtual_column_fields.iter() {
             let name = format!("{}", virtual_column_field.column_id);
             if let Some(arrow_array) = record_batch
                 .as_ref()
                 .and_then(|r| r.column_by_name(&name).cloned())
             {
-                let orig_field = orig_schema.field_with_name(&name).unwrap();
+                let orig_field = self.orig_schema.field_with_name(&name).unwrap();
                 let orig_type: DataType = orig_field.data_type().into();
                 let value = Value::Column(Column::from_arrow_rs(arrow_array, &orig_type)?);
                 let data_type: DataType = virtual_column_field.data_type.as_ref().into();
@@ -240,5 +374,52 @@ impl VirtualColumnReader {
         }
 
         Ok(data_block)
+
+        // for virtual_column_field in self.virtual_column_fields.iter() {
+        //    if let Some(arrow_array) =
+        //        record_batch.and_then(|r| r.column_by_name(&virtual_column_field.name).cloned())
+        //    {
+        //        let data_type: DataType = virtual_column_field.data_type.as_ref().into();
+        //        let value = Value::Column(Column::from_arrow_rs(arrow_array, &data_type)?);
+        //        data_block.add_column(BlockEntry::new(data_type, value));
+        //        continue;
+        //    }
+        //    let src_index = self
+        //        .source_schema
+        //        .index_of(&virtual_column_field.source_name)
+        //        .unwrap();
+        //    let source = data_block.get_by_offset(src_index);
+        //    let src_arg = (source.value.clone(), source.data_type.clone());
+        //    let path_arg = (
+        //        Value::Scalar(virtual_column_field.key_paths.clone()),
+        //        DataType::String,
+        //    );
+
+        //    let (value, data_type) = eval_function(
+        //        None,
+        //        "get_by_keypath",
+        //        [src_arg, path_arg],
+        //        func_ctx,
+        //        data_block.num_rows(),
+        //        &BUILTIN_FUNCTIONS,
+        //    )?;
+
+        //    let column = if let Some(cast_func_name) = &virtual_column_field.cast_func_name {
+        //        let (cast_value, cast_data_type) = eval_function(
+        //            None,
+        //            cast_func_name,
+        //            [(value, data_type)],
+        //            func_ctx,
+        //            data_block.num_rows(),
+        //            &BUILTIN_FUNCTIONS,
+        //        )?;
+        //        BlockEntry::new(cast_data_type, cast_value)
+        //    } else {
+        //        BlockEntry::new(data_type, value)
+        //    };
+        //    data_block.add_column(column);
+        //}
+
+        // Ok(data_block)
     }
 }

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",
                 ));
             }