refactor(query): Refactor and optimize virtual column refresh (#17950)

* fix(query): fix refresh virtual column

* fix

* add table lock

Author: b41sh

Cargo.lock

@@ -271,8 +271,8 @@ build_exceptions! {
     LicenceDenied(1112),
     /// Unknown datamask
     UnknownDatamask(1113),
-    /// Virtual column not found
-    VirtualColumnNotFound(1115),
+    /// Virtual column error
+    VirtualColumnError(1115),
     /// Virtual column already exists
     VirtualColumnAlreadyExists(1116),
     /// Column referenced by computed column

src/common/exception/src/exception_code.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::collections::VecDeque;
 use std::sync::Arc;
 use std::time::Instant;
@@ -21,8 +22,10 @@ use databend_common_catalog::table::Table;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_expression::BlockMetaInfoDowncast;
+use databend_common_expression::ComputedExpr;
 use databend_common_expression::DataBlock;
 use databend_common_expression::TableDataType;
+use databend_common_expression::TableSchema;
 use databend_common_metrics::storage::metrics_inc_block_virtual_column_write_bytes;
 use databend_common_metrics::storage::metrics_inc_block_virtual_column_write_milliseconds;
 use databend_common_metrics::storage::metrics_inc_block_virtual_column_write_nums;
@@ -31,18 +34,26 @@ use databend_common_pipeline_sources::AsyncSource;
 use databend_common_pipeline_sources::AsyncSourcer;
 use databend_common_pipeline_transforms::processors::AsyncTransform;
 use databend_common_pipeline_transforms::processors::TransformPipelineHelper;
+use databend_common_sql::executor::physical_plans::MutationKind;
 use databend_common_storages_fuse::io::write_data;
 use databend_common_storages_fuse::io::BlockReader;
 use databend_common_storages_fuse::io::MetaReaders;
 use databend_common_storages_fuse::io::VirtualColumnBuilder;
 use databend_common_storages_fuse::io::WriteSettings;
+use databend_common_storages_fuse::operations::BlockMetaIndex;
+use databend_common_storages_fuse::operations::CommitSink;
+use databend_common_storages_fuse::operations::MutationGenerator;
+use databend_common_storages_fuse::operations::MutationLogEntry;
+use databend_common_storages_fuse::operations::MutationLogs;
+use databend_common_storages_fuse::operations::TableMutationAggregator;
 use databend_common_storages_fuse::FuseStorageFormat;
 use databend_common_storages_fuse::FuseTable;
 use databend_storages_common_cache::LoadParams;
 use databend_storages_common_io::ReadSettings;
 use databend_storages_common_table_meta::meta::BlockMeta;
 use databend_storages_common_table_meta::meta::ExtendedBlockMeta;
 use databend_storages_common_table_meta::meta::Location;
+use databend_storages_common_table_meta::meta::Statistics;
 use opendal::Operator;
 
 // The big picture of refresh virtual column into pipeline:
@@ -63,38 +74,33 @@ use opendal::Operator;
 pub async fn do_refresh_virtual_column(
     ctx: Arc<dyn TableContext>,
     fuse_table: &FuseTable,
-    segment_locs: Option<Vec<Location>>,
     pipeline: &mut Pipeline,
 ) -> Result<()> {
-    let snapshot_opt = fuse_table.read_table_snapshot().await?;
-    let snapshot = if let Some(val) = snapshot_opt {
-        val
-    } else {
+    let Some(snapshot) = fuse_table.read_table_snapshot().await? else {
         // no snapshot
         return Ok(());
     };
-
     let table_schema = &fuse_table.get_table_info().meta.schema;
 
     // Collect source fields used by virtual columns.
+    let mut fields = Vec::new();
     let mut field_indices = Vec::new();
     for (i, f) in table_schema.fields().iter().enumerate() {
-        if f.data_type().remove_nullable() != TableDataType::Variant {
+        if f.data_type().remove_nullable() != TableDataType::Variant
+            || matches!(f.computed_expr(), Some(ComputedExpr::Virtual(_)))
+        {
             continue;
         }
+        fields.push(f.clone());
         field_indices.push(i);
     }
 
-    if field_indices.is_empty() {
-        // no source variant column
-        return Ok(());
-    }
-
-    let table_info = &fuse_table.get_table_info();
-    let Some(virtual_column_builder) = VirtualColumnBuilder::try_create(ctx.clone(), table_info)
-    else {
-        return Ok(());
-    };
+    let source_schema = Arc::new(TableSchema {
+        fields,
+        ..fuse_table.schema().as_ref().clone()
+    });
+    let virtual_column_builder =
+        VirtualColumnBuilder::try_create(ctx.clone(), fuse_table, source_schema)?;
 
     let projection = Projection::Columns(field_indices);
     let block_reader =
@@ -108,16 +114,11 @@ pub async fn do_refresh_virtual_column(
 
     let operator = fuse_table.get_operator_ref();
 
-    // If no segment locations are specified, iterates through all segments
-    let segment_locs = if let Some(segment_locs) = segment_locs {
-        segment_locs
-    } else {
-        snapshot.segments.clone()
-    };
-
-    // Read source variant columns and extract inner fields as virtual columns.
+    // Iterates through all segments and collect blocks don't have virtual block meta.
+    let segment_locs = snapshot.segments.clone();
     let mut block_metas = VecDeque::new();
-    for (location, ver) in segment_locs {
+    let mut block_meta_index_map = HashMap::new();
+    for (segment_idx, (location, ver)) in segment_locs.into_iter().enumerate() {
         let segment_info = segment_reader
             .read(&LoadParams {
                 location: location.to_string(),
@@ -127,10 +128,15 @@ pub async fn do_refresh_virtual_column(
             })
             .await?;
 
-        for block_meta in segment_info.block_metas()? {
+        for (block_idx, block_meta) in segment_info.block_metas()?.into_iter().enumerate() {
             if block_meta.virtual_block_meta.is_some() {
                 continue;
             }
+            let index = BlockMetaIndex {
+                segment_idx,
+                block_idx,
+            };
+            block_meta_index_map.insert(block_meta.location.clone(), index);
             block_metas.push_back(block_meta);
         }
     }
@@ -139,6 +145,7 @@ pub async fn do_refresh_virtual_column(
         return Ok(());
     }
 
+    // Read source blocks.
     let settings = ReadSettings::from_ctx(&ctx)?;
     pipeline.add_source(
         |output| {
@@ -153,31 +160,52 @@ pub async fn do_refresh_virtual_column(
         1,
     )?;
 
+    // Extract inner fields as virtual columns and write virtual block data.
     let block_nums = block_metas.len();
     let max_threads = ctx.get_settings().get_max_threads()? as usize;
     let max_threads = std::cmp::min(block_nums, max_threads);
     pipeline.try_resize(max_threads)?;
     pipeline.add_async_transformer(|| {
         VirtualColumnTransform::new(
+            operator.clone(),
             write_settings.clone(),
+            block_meta_index_map.clone(),
             virtual_column_builder.clone(),
-            operator.clone(),
         )
     });
 
-    let base_snapshot = fuse_table.read_table_snapshot().await?;
-    let table_meta_timestamps = ctx.get_table_meta_timestamps(fuse_table, base_snapshot.clone())?;
-
-    fuse_table.do_commit(
-        ctx,
-        pipeline,
-        None,
-        vec![],
-        false,
-        None,
-        None,
-        table_meta_timestamps,
-    )?;
+    pipeline.try_resize(1)?;
+    let table_meta_timestamps =
+        ctx.get_table_meta_timestamps(fuse_table, Some(snapshot.clone()))?;
+    pipeline.add_async_accumulating_transformer(|| {
+        TableMutationAggregator::create(
+            fuse_table,
+            ctx.clone(),
+            vec![],
+            vec![],
+            vec![],
+            Statistics::default(),
+            MutationKind::Update,
+            table_meta_timestamps,
+        )
+    });
+
+    let prev_snapshot_id = snapshot.snapshot_id;
+    let snapshot_gen = MutationGenerator::new(Some(snapshot), MutationKind::Update);
+    pipeline.add_sink(|input| {
+        CommitSink::try_create(
+            fuse_table,
+            ctx.clone(),
+            None,
+            vec![],
+            snapshot_gen.clone(),
+            input,
+            None,
+            Some(prev_snapshot_id),
+            None,
+            table_meta_timestamps,
+        )
+    })?;
 
     Ok(())
 }
@@ -238,21 +266,24 @@ impl AsyncSource for VirtualColumnSource {
 
 /// `VirtualColumnTransform` is used to generate virtual columns for each blocks.
 pub struct VirtualColumnTransform {
+    operator: Operator,
     write_settings: WriteSettings,
+    block_meta_index_map: HashMap<Location, BlockMetaIndex>,
     virtual_column_builder: VirtualColumnBuilder,
-    operator: Operator,
 }
 
 impl VirtualColumnTransform {
     pub fn new(
+        operator: Operator,
         write_settings: WriteSettings,
+        block_meta_index_map: HashMap<Location, BlockMetaIndex>,
         virtual_column_builder: VirtualColumnBuilder,
-        operator: Operator,
     ) -> Self {
         Self {
+            operator,
             write_settings,
+            block_meta_index_map,
             virtual_column_builder,
-            operator,
         }
     }
 }
@@ -301,13 +332,23 @@ impl AsyncTransform for VirtualColumnTransform {
             }
         }
 
+        let block_meta_index = self
+            .block_meta_index_map
+            .remove(&block_meta.location)
+            .unwrap();
         let extended_block_meta = ExtendedBlockMeta {
             block_meta: block_meta.clone(),
             draft_virtual_block_meta: Some(virtual_column_state.draft_virtual_block_meta),
         };
 
-        let new_block = DataBlock::new(vec![], 0);
-        let new_block = new_block.add_meta(Some(extended_block_meta.boxed()))?;
+        let entry = MutationLogEntry::ReplacedBlock {
+            index: block_meta_index,
+            block_meta: Arc::new(extended_block_meta),
+        };
+        let meta = MutationLogs {
+            entries: vec![entry],
+        };
+        let new_block = DataBlock::empty_with_meta(Box::new(meta));
         Ok(new_block)
     }
 }

src/query/ee/src/storages/fuse/operations/virtual_columns.rs

@@ -21,7 +21,6 @@ use databend_common_pipeline_core::Pipeline;
 use databend_common_storages_fuse::FuseTable;
 use databend_enterprise_virtual_column::VirtualColumnHandler;
 use databend_enterprise_virtual_column::VirtualColumnHandlerWrapper;
-use databend_storages_common_table_meta::meta::Location;
 
 use crate::storages::fuse::do_refresh_virtual_column;
 
@@ -33,10 +32,9 @@ impl VirtualColumnHandler for RealVirtualColumnHandler {
         &self,
         ctx: Arc<dyn TableContext>,
         fuse_table: &FuseTable,
-        segment_locs: Option<Vec<Location>>,
         pipeline: &mut Pipeline,
     ) -> Result<()> {
-        do_refresh_virtual_column(ctx, fuse_table, segment_locs, pipeline).await
+        do_refresh_virtual_column(ctx, fuse_table, pipeline).await
     }
 }
 

src/query/ee/src/virtual_column/virtual_column_handler.rs

@@ -55,14 +55,7 @@ async fn test_fuse_do_refresh_virtual_column() -> Result<()> {
     let snapshot = snapshot_opt.unwrap();
 
     let mut build_res = PipelineBuildResult::create();
-    let segment_locs = Some(snapshot.segments.clone());
-    do_refresh_virtual_column(
-        table_ctx.clone(),
-        fuse_table,
-        segment_locs,
-        &mut build_res.main_pipeline,
-    )
-    .await?;
+    do_refresh_virtual_column(table_ctx.clone(), fuse_table, &mut build_res.main_pipeline).await?;
 
     let settings = table_ctx.get_settings();
     build_res.set_max_threads(settings.get_max_threads()? as usize);

src/query/ee/tests/it/storages/fuse/operations/virtual_columns.rs

@@ -48,6 +48,7 @@ async fn test_virtual_column_builder() -> Result<()> {
 
     let table = fixture.latest_default_table().await?;
     let table_info = table.get_table_info();
+    let schema = table_info.meta.schema.clone();
 
     let fuse_table = FuseTable::try_from_table(table.as_ref())?;
 
@@ -57,7 +58,7 @@ async fn test_virtual_column_builder() -> Result<()> {
         0,
     ); // Dummy location
 
-    let builder = VirtualColumnBuilder::try_create(ctx, table_info).unwrap();
+    let builder = VirtualColumnBuilder::try_create(ctx, fuse_table, schema).unwrap();
 
     let block = DataBlock::new(
         vec![

src/query/ee/tests/it/storages/fuse/operations/virtual_columns_builder.rs

@@ -16,7 +16,6 @@ databend-common-catalog = { workspace = true }
 databend-common-exception = { workspace = true }
 databend-common-pipeline-core = { workspace = true }
 databend-common-storages-fuse = { workspace = true }
-databend-storages-common-table-meta = { workspace = true }
 
 [build-dependencies]
 

src/query/ee_features/virtual_column/Cargo.toml

@@ -19,15 +19,13 @@ use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_pipeline_core::Pipeline;
 use databend_common_storages_fuse::FuseTable;
-use databend_storages_common_table_meta::meta::Location;
 
 #[async_trait::async_trait]
 pub trait VirtualColumnHandler: Sync + Send {
     async fn do_refresh_virtual_column(
         &self,
         ctx: Arc<dyn TableContext>,
         fuse_table: &FuseTable,
-        segment_locs: Option<Vec<Location>>,
         pipeline: &mut Pipeline,
     ) -> Result<()>;
 }
@@ -46,11 +44,10 @@ impl VirtualColumnHandlerWrapper {
         &self,
         ctx: Arc<dyn TableContext>,
         fuse_table: &FuseTable,
-        segment_locs: Option<Vec<Location>>,
         pipeline: &mut Pipeline,
     ) -> Result<()> {
         self.handler
-            .do_refresh_virtual_column(ctx, fuse_table, segment_locs, pipeline)
+            .do_refresh_virtual_column(ctx, fuse_table, pipeline)
             .await
     }
 }