chore(query): Refresh virtual column support limit and selection (#19001)

Author: b41sh

Cargo.lock

@@ -19,6 +19,7 @@ use derive_visitor::Drive;
 use derive_visitor::DriveMut;
 
 use crate::ast::write_dot_separated_list;
+use crate::ast::Expr;
 use crate::ast::Identifier;
 use crate::ast::ShowLimit;
 
@@ -27,19 +28,30 @@ pub struct RefreshVirtualColumnStmt {
     pub catalog: Option<Identifier>,
     pub database: Option<Identifier>,
     pub table: Identifier,
+    pub selection: Option<Box<Expr>>,
+    pub limit: Option<u64>,
+    pub overwrite: bool,
 }
 
 impl Display for RefreshVirtualColumnStmt {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        write!(f, "REFRESH VIRTUAL COLUMN FOR ")?;
+        write!(f, "REFRESH VIRTUAL COLUMN ON ")?;
         write_dot_separated_list(
             f,
             self.catalog
                 .iter()
                 .chain(&self.database)
                 .chain(Some(&self.table)),
         )?;
-
+        if let Some(selection) = &self.selection {
+            write!(f, " WHERE {selection}")?;
+        }
+        if let Some(limit) = self.limit {
+            write!(f, " LIMIT {limit}")?;
+        }
+        if self.overwrite {
+            write!(f, " OVERWRITE")?;
+        }
         Ok(())
     }
 }

src/query/ast/src/ast/statements/virtual_column.rs

@@ -1554,13 +1554,16 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
 
     let refresh_virtual_column = map(
         rule! {
-            REFRESH ~ VIRTUAL ~ COLUMN ~ FOR ~ #dot_separated_idents_1_to_3
+            REFRESH ~ VIRTUAL ~ ^COLUMN ~ ^( FOR | ON ) ~ ^#dot_separated_idents_1_to_3 ~ ( WHERE ~ ^#expr )? ~ ( LIMIT ~ ^#literal_u64 )? ~ OVERWRITE?
         },
-        |(_, _, _, _, (catalog, database, table))| {
+        |(_, _, _, _, (catalog, database, table), opt_selection, opt_limit, opt_overwrite)| {
             Statement::RefreshVirtualColumn(RefreshVirtualColumnStmt {
                 catalog,
                 database,
                 table,
+                selection: opt_selection.map(|(_, selection)| Box::new(selection)),
+                limit: opt_limit.map(|(_, limit)| limit),
+                overwrite: opt_overwrite.is_some(),
             })
         },
     );

src/query/ast/src/parser/statement.rs

@@ -22511,7 +22511,7 @@ DropDatamaskPolicy(
 ---------- Input ----------
 REFRESH VIRTUAL COLUMN FOR t
 ---------- Output ---------
-REFRESH VIRTUAL COLUMN FOR t
+REFRESH VIRTUAL COLUMN ON t
 ---------- AST ------------
 RefreshVirtualColumn(
     RefreshVirtualColumnStmt {
@@ -22525,6 +22525,9 @@ RefreshVirtualColumn(
             quote: None,
             ident_type: None,
         },
+        selection: None,
+        limit: None,
+        overwrite: false,
     },
 )
 

src/query/ast/tests/it/testdata/stmt.txt

@@ -39,6 +39,7 @@ 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_sql::plans::RefreshSelection;
 use databend_common_storages_fuse::io::read::read_segment_stats;
 use databend_common_storages_fuse::io::write_data;
 use databend_common_storages_fuse::io::BlockReader;
@@ -60,6 +61,8 @@ use databend_storages_common_table_meta::meta::BlockMeta;
 use databend_storages_common_table_meta::meta::ExtendedBlockMeta;
 use databend_storages_common_table_meta::meta::RawBlockHLL;
 use databend_storages_common_table_meta::meta::Statistics;
+use log::debug;
+use log::info;
 use opendal::Operator;
 
 // The big picture of refresh virtual column into pipeline:
@@ -81,6 +84,9 @@ pub async fn do_refresh_virtual_column(
     ctx: Arc<dyn TableContext>,
     fuse_table: &FuseTable,
     pipeline: &mut Pipeline,
+    limit: Option<u64>,
+    overwrite: bool,
+    selection: Option<RefreshSelection>,
 ) -> Result<()> {
     let Some(snapshot) = fuse_table.read_table_snapshot().await? else {
         // no snapshot
@@ -127,9 +133,29 @@ pub async fn do_refresh_virtual_column(
 
     let operator = fuse_table.get_operator_ref();
 
+    let limit = limit.unwrap_or_default() as usize;
+    let segment_filter = selection.as_ref().and_then(|sel| match sel {
+        RefreshSelection::SegmentLocation(loc) => Some(loc.clone()),
+        _ => None,
+    });
+    let block_filter = selection.as_ref().and_then(|sel| match sel {
+        RefreshSelection::BlockLocation(loc) => Some(loc.clone()),
+        _ => None,
+    });
+    let mut matched_selection = false;
+    let mut reached_limit = false;
     // Iterates through all segments and collect blocks don't have virtual block meta.
     let mut virtual_column_metas = VecDeque::new();
     for (segment_idx, (location, ver)) in snapshot.segments.iter().enumerate() {
+        if reached_limit {
+            break;
+        }
+        if let Some(target) = segment_filter.as_ref() {
+            if location != target {
+                continue;
+            }
+            matched_selection = true;
+        }
         let segment_info = segment_reader
             .read(&LoadParams {
                 location: location.to_string(),
@@ -144,9 +170,23 @@ pub async fn do_refresh_virtual_column(
         };
 
         for (block_idx, block_meta) in segment_info.block_metas()?.into_iter().enumerate() {
-            if block_meta.virtual_block_meta.is_some() {
+            let mut matched_block_filter = false;
+            if let Some(target) = block_filter.as_ref() {
+                if &block_meta.location.0 != target {
+                    continue;
+                }
+                matched_selection = true;
+                matched_block_filter = true;
+            }
+
+            if !overwrite && block_meta.virtual_block_meta.is_some() {
+                if matched_block_filter {
+                    reached_limit = true;
+                    break;
+                }
                 continue;
             }
+
             virtual_column_metas.push_back(VirtualColumnMeta {
                 index: BlockMetaIndex {
                     segment_idx,
@@ -158,13 +198,40 @@ pub async fn do_refresh_virtual_column(
                     .and_then(|v| v.block_hlls.get(block_idx))
                     .cloned(),
             });
+
+            if limit > 0 && virtual_column_metas.len() >= limit {
+                reached_limit = true;
+                break;
+            }
+            if matched_block_filter {
+                reached_limit = true;
+                break;
+            }
         }
     }
 
+    if let (Some(sel), false) = (selection.as_ref(), matched_selection) {
+        let message = match sel {
+            RefreshSelection::SegmentLocation(loc) => {
+                format!("segment_location '{loc}' not found")
+            }
+            RefreshSelection::BlockLocation(loc) => {
+                format!("block_location '{loc}' not found")
+            }
+        };
+        return Err(ErrorCode::VirtualColumnError(message));
+    }
+
     if virtual_column_metas.is_empty() {
         return Ok(());
     }
 
+    let block_nums = virtual_column_metas.len();
+    info!(
+        "Prepared {} blocks for virtual column refresh (limit={}, overwrite={})",
+        block_nums, limit, overwrite
+    );
+
     // Read source blocks.
     let settings = ReadSettings::from_ctx(&ctx)?;
     pipeline.add_source(
@@ -181,9 +248,12 @@ pub async fn do_refresh_virtual_column(
     )?;
 
     // Extract inner fields as virtual columns and write virtual block data.
-    let block_nums = virtual_column_metas.len();
     let max_threads = ctx.get_settings().get_max_threads()? as usize;
     let max_threads = std::cmp::min(block_nums, max_threads);
+    info!(
+        "Virtual column pipeline will process {} blocks with {} async workers",
+        block_nums, max_threads
+    );
     pipeline.try_resize(max_threads)?;
     pipeline.add_async_transformer(|| {
         VirtualColumnTransform::new(
@@ -229,12 +299,16 @@ pub async fn do_refresh_virtual_column(
     Ok(())
 }
 
+const VIRTUAL_COLUMN_PROGRESS_LOG_STEP: usize = 10;
+
 /// `VirtualColumnSource` is used to read data blocks that need generate virtual columns.
 pub struct VirtualColumnSource {
     settings: ReadSettings,
     storage_format: FuseStorageFormat,
     block_reader: Arc<BlockReader>,
     virtual_column_metas: VecDeque<VirtualColumnMeta>,
+    total_blocks: usize,
+    processed_blocks: usize,
     is_finished: bool,
 }
 
@@ -245,11 +319,14 @@ impl VirtualColumnSource {
         block_reader: Arc<BlockReader>,
         virtual_column_metas: VecDeque<VirtualColumnMeta>,
     ) -> Self {
+        let total_blocks = virtual_column_metas.len();
         Self {
             settings,
             storage_format,
             block_reader,
             virtual_column_metas,
+            total_blocks,
+            processed_blocks: 0,
             is_finished: false,
         }
     }
@@ -267,6 +344,21 @@ impl AsyncSource for VirtualColumnSource {
 
         match self.virtual_column_metas.pop_front() {
             Some(meta) => {
+                self.processed_blocks += 1;
+                if self.processed_blocks == 1
+                    || self.processed_blocks == self.total_blocks
+                    || self.processed_blocks % VIRTUAL_COLUMN_PROGRESS_LOG_STEP == 0
+                {
+                    info!(
+                        "Virtual column source progress: {}/{}",
+                        self.processed_blocks, self.total_blocks
+                    );
+                } else {
+                    debug!(
+                        "Virtual column source progress: {}/{}",
+                        self.processed_blocks, self.total_blocks
+                    );
+                }
                 let block = self
                     .block_reader
                     .read_by_meta(&self.settings, &meta.block_meta, &self.storage_format)
@@ -275,6 +367,12 @@ impl AsyncSource for VirtualColumnSource {
                 Ok(Some(block))
             }
             None => {
+                if !self.is_finished {
+                    info!(
+                        "Virtual column source finished reading {} blocks",
+                        self.processed_blocks
+                    );
+                }
                 self.is_finished = true;
                 Ok(None)
             }
@@ -348,6 +446,15 @@ impl AsyncTransform for VirtualColumnTransform {
                     start.elapsed().as_millis() as u64
                 );
             }
+            info!(
+                "Virtual column written for segment {} block {} at {} ({} bytes)",
+                index.segment_idx, index.block_idx, location, virtual_column_size
+            );
+        } else {
+            info!(
+                "No virtual column data produced for segment {} block {}",
+                index.segment_idx, index.block_idx
+            );
         }
 
         let extended_block_meta = ExtendedBlockMeta {

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

@@ -18,6 +18,7 @@ use databend_common_base::base::GlobalInstance;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_pipeline::core::Pipeline;
+use databend_common_sql::plans::RefreshSelection;
 use databend_common_storages_fuse::FuseTable;
 use databend_enterprise_virtual_column::VirtualColumnHandler;
 use databend_enterprise_virtual_column::VirtualColumnHandlerWrapper;
@@ -33,8 +34,11 @@ impl VirtualColumnHandler for RealVirtualColumnHandler {
         ctx: Arc<dyn TableContext>,
         fuse_table: &FuseTable,
         pipeline: &mut Pipeline,
+        limit: Option<u64>,
+        overwrite: bool,
+        selection: Option<RefreshSelection>,
     ) -> Result<()> {
-        do_refresh_virtual_column(ctx, fuse_table, pipeline).await
+        do_refresh_virtual_column(ctx, fuse_table, pipeline, limit, overwrite, selection).await
     }
 }
 

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

@@ -51,7 +51,15 @@ async fn test_fuse_do_refresh_virtual_column() -> Result<()> {
     let snapshot = snapshot_opt.unwrap();
 
     let mut build_res = PipelineBuildResult::create();
-    do_refresh_virtual_column(table_ctx.clone(), fuse_table, &mut build_res.main_pipeline).await?;
+    do_refresh_virtual_column(
+        table_ctx.clone(),
+        fuse_table,
+        &mut build_res.main_pipeline,
+        None,
+        false,
+        None,
+    )
+    .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

@@ -15,6 +15,7 @@ databend-common-base = { workspace = true }
 databend-common-catalog = { workspace = true }
 databend-common-exception = { workspace = true }
 databend-common-pipeline = { workspace = true }
+databend-common-sql = { workspace = true }
 databend-common-storages-fuse = { workspace = true }
 
 [build-dependencies]

src/query/ee_features/virtual_column/Cargo.toml

@@ -18,6 +18,7 @@ use databend_common_base::base::GlobalInstance;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_pipeline::core::Pipeline;
+use databend_common_sql::plans::RefreshSelection;
 use databend_common_storages_fuse::FuseTable;
 
 #[async_trait::async_trait]
@@ -27,6 +28,9 @@ pub trait VirtualColumnHandler: Sync + Send {
         ctx: Arc<dyn TableContext>,
         fuse_table: &FuseTable,
         pipeline: &mut Pipeline,
+        limit: Option<u64>,
+        overwrite: bool,
+        selection: Option<RefreshSelection>,
     ) -> Result<()>;
 }
 
@@ -45,9 +49,12 @@ impl VirtualColumnHandlerWrapper {
         ctx: Arc<dyn TableContext>,
         fuse_table: &FuseTable,
         pipeline: &mut Pipeline,
+        limit: Option<u64>,
+        overwrite: bool,
+        selection: Option<RefreshSelection>,
     ) -> Result<()> {
         self.handler
-            .do_refresh_virtual_column(ctx, fuse_table, pipeline)
+            .do_refresh_virtual_column(ctx, fuse_table, pipeline, limit, overwrite, selection)
             .await
     }
 }

src/query/ee_features/virtual_column/src/virtual_column.rs

@@ -57,8 +57,8 @@ pub const SEARCH_SCORE_COLUMN_ID: u32 = u32::MAX - 8;
 pub const VECTOR_SCORE_COLUMN_ID: u32 = u32::MAX - 9;
 
 pub const VIRTUAL_COLUMN_ID_START: u32 = 3_000_000_000;
-pub const VIRTUAL_COLUMNS_ID_UPPER: u32 = 3_000_001_000;
-pub const VIRTUAL_COLUMNS_LIMIT: usize = 1000;
+pub const VIRTUAL_COLUMNS_ID_UPPER: u32 = 3_000_002_000;
+pub const VIRTUAL_COLUMNS_LIMIT: usize = 2000;
 
 // internal column name.
 pub const ROW_ID_COL_NAME: &str = "_row_id";