clean up

Author: dantengsky

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

@@ -74,7 +74,11 @@ pub fn column_chunks_to_record_batch(
         batch_size,
         None,
     )?;
-    let records = record_reader.try_collect()?;
-    // TODO assert the row numbers?
+
+    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/virtual_column/virtual_column_reader_parquet.rs

@@ -15,7 +15,7 @@
 use std::collections::HashSet;
 
 use arrow_array::RecordBatch;
-use databend_common_catalog::plan::VirtualColumnInfo;
+use databend_common_catalog::plan::VirtualColumnField;
 use databend_common_exception::Result;
 use databend_common_expression::eval_function;
 use databend_common_expression::types::DataType;
@@ -132,12 +132,13 @@ impl VirtualColumnReader {
         ))
     }
 
+    /// Deserialize virtual column data into record batches, according to the `batch_size`.
     pub fn try_create_paster(
         &self,
         virtual_data: Option<VirtualBlockReadResult>,
         batch_size: usize,
-    ) -> Result<VirtualColumnDataModifier> {
-        let chunks = if let Some(virtual_data) = virtual_data {
+    ) -> Result<VirtualColumnDataPaster> {
+        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,
@@ -151,27 +152,31 @@ impl VirtualColumnReader {
             None
         };
 
-        let func_ctx = self.ctx.get_function_context()?;
+        let function_context = self.ctx.get_function_context()?;
 
-        Ok(VirtualColumnDataModifier {
-            record_batches: chunks,
-            function_context: func_ctx,
+        // 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_info: self.virtual_column_info.clone(),
+            virtual_column_fields: self.virtual_column_info.virtual_column_fields.clone(),
             source_schema: self.source_schema.clone(),
         })
     }
 }
 
-pub struct VirtualColumnDataModifier {
+pub struct VirtualColumnDataPaster {
     record_batches: Option<Vec<RecordBatch>>,
     next_record_batch_index: usize,
     function_context: FunctionContext,
-    virtual_column_info: VirtualColumnInfo,
+    virtual_column_fields: Vec<VirtualColumnField>,
     source_schema: TableSchemaRef,
 }
 
-impl VirtualColumnDataModifier {
+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);
@@ -183,7 +188,7 @@ impl VirtualColumnDataModifier {
         self.next_record_batch_index += 1;
 
         let func_ctx = &self.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() {
             if let Some(arrow_array) =
                 record_batch.and_then(|r| r.column_by_name(&virtual_column_field.name).cloned())
             {
@@ -207,7 +212,7 @@ impl VirtualColumnDataModifier {
                 None,
                 "get_by_keypath",
                 [src_arg, path_arg],
-                &func_ctx,
+                func_ctx,
                 data_block.num_rows(),
                 &BUILTIN_FUNCTIONS,
             )?;
@@ -217,7 +222,7 @@ impl VirtualColumnDataModifier {
                     None,
                     cast_func_name,
                     [(value, data_type)],
-                    &func_ctx,
+                    func_ctx,
                     data_block.num_rows(),
                     &BUILTIN_FUNCTIONS,
                 )?;