refactot build bloom info

JackTan25 · JackTan25 · commit 7a41e4e5707d · 2024-03-21T01:17:03.000+08:00
diff --git a/src/query/catalog/src/merge_into_join.rs b/src/query/catalog/src/merge_into_join.rs
@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use databend_common_expression::TableSchemaRef;
 use databend_common_meta_app::schema::CatalogInfo;
 use databend_common_meta_app::schema::TableInfo;
 use databend_storages_common_table_meta::meta::Location;
@@ -61,6 +62,7 @@ pub struct MergeIntoJoin {
     pub table_info: Option<TableInfo>,
     pub catalog_info: Option<CatalogInfo>,
     pub database_name: String,
+    pub table_schema: Option<TableSchemaRef>,
 }
 
 impl Default for MergeIntoJoin {
@@ -73,6 +75,7 @@ impl Default for MergeIntoJoin {
             table_info: None,
             catalog_info: None,
             database_name: Default::default(),
+            table_schema: None,
         }
     }
 }
diff --git a/src/query/catalog/src/runtime_filter_info.rs b/src/query/catalog/src/runtime_filter_info.rs
@@ -14,8 +14,6 @@
 
 use std::sync::Arc;
 
-use databend_common_arrow::arrow::bitmap::Bitmap;
-use databend_common_arrow::arrow::buffer::Buffer;
 use databend_common_expression::Expr;
 use xorf::BinaryFuse16;
 
diff --git a/src/query/catalog/src/table_context.rs b/src/query/catalog/src/table_context.rs
@@ -278,7 +278,9 @@ pub trait TableContext: Send + Sync {
     fn get_merge_into_source_build_siphashkeys_with_id(
         &self,
         id: usize,
-    ) -> Vec<(String, (Buffer<u64>, Option<Bitmap>))>;
+    ) -> Vec<(String, Arc<Vec<u64>>)>;
+
+    fn get_merge_into_source_build_bloom_probe_keys(&self, id: usize) -> Vec<String>;
 
     fn has_bloom_runtime_filters(&self, id: usize) -> bool;
     fn txn_mgr(&self) -> TxnManagerRef;
diff --git a/src/query/service/src/sessions/query_ctx.rs b/src/query/service/src/sessions/query_ctx.rs
@@ -1017,6 +1017,7 @@ impl TableContext for QueryContext {
             catalog_info: merge_into_join.catalog_info.clone(),
             table_info: merge_into_join.table_info.clone(),
             database_name: merge_into_join.database_name.clone(),
+            table_schema: merge_into_join.table_schema.clone(),
         }
     }
 
@@ -1031,14 +1032,26 @@ impl TableContext for QueryContext {
     fn get_merge_into_source_build_siphashkeys_with_id(
         &self,
         id: IndexType,
-    ) -> Vec<(String, (Buffer<u64>, Option<Bitmap>))> {
+    ) -> Vec<(String, Arc<Vec<u64>>)> {
         let runtime_filters = self.shared.runtime_filters.read();
         match runtime_filters.get(&id) {
             Some(v) => v.get_merge_into_source_build_siphashkeys(),
             None => vec![],
         }
     }
 
+    fn get_merge_into_source_build_bloom_probe_keys(&self, id: usize) -> Vec<String> {
+        let runtime_filters = self.shared.runtime_filters.read();
+        match runtime_filters.get(&id) {
+            Some(v) => v
+                .get_merge_into_source_build_siphashkeys()
+                .iter()
+                .map(|key| key.0)
+                .collect(),
+            None => vec![],
+        }
+    }
+
     fn get_inlist_runtime_filter_with_id(&self, id: IndexType) -> Vec<Expr<String>> {
         let runtime_filters = self.shared.runtime_filters.read();
         match runtime_filters.get(&id) {
diff --git a/src/query/service/tests/it/sql/exec/get_table_bind_test.rs b/src/query/service/tests/it/sql/exec/get_table_bind_test.rs
@@ -833,7 +833,11 @@ impl TableContext for CtxDelegation {
     fn get_merge_into_source_build_siphashkeys_with_id(
         &self,
         id: usize,
-    ) -> Vec<(String, (Buffer<u64>, Option<Bitmap>))> {
+    ) -> Vec<(String, Arc<Vec<u64>>)> {
+        todo!()
+    }
+
+    fn get_merge_into_source_build_bloom_probe_keys(&self, id: usize) -> Vec<String> {
         todo!()
     }
 
diff --git a/src/query/sql/src/planner/binder/merge_into.rs b/src/query/sql/src/planner/binder/merge_into.rs
@@ -488,6 +488,7 @@ impl Binder {
             row_id_index: column_binding.index,
             split_idx,
             can_try_update_column_only: self.can_try_update_column_only(&matched_clauses),
+            table_schema,
         })
     }
 
diff --git a/src/query/sql/src/planner/optimizer/optimizer.rs b/src/query/sql/src/planner/optimizer/optimizer.rs
@@ -427,6 +427,7 @@ fn optimize_merge_into(opt_ctx: OptimizerContext, plan: Box<MergeInto>) -> Resul
             is_distributed: false, // we will set it after later optimization.
             merge_into_join_type: MergeIntoJoinType::Right,
             database_name: plan.database.clone(),
+            table_schema: Some(plan.table_schema.clone()),
         })
     }
 
@@ -517,6 +518,7 @@ fn optimize_merge_into(opt_ctx: OptimizerContext, plan: Box<MergeInto>) -> Resul
                 catalog_info: merge_into_join.catalog_info.clone(),
                 table_info: merge_into_join.table_info.clone(),
                 database_name: merge_into_join.database_name.clone(),
+                table_schema: merge_into_join.table_schema.clone(),
             })
         }
 
@@ -547,6 +549,7 @@ fn optimize_merge_into(opt_ctx: OptimizerContext, plan: Box<MergeInto>) -> Resul
                 catalog_info: merge_into_join.catalog_info.clone(),
                 table_info: merge_into_join.table_info.clone(),
                 database_name: merge_into_join.database_name,
+                table_schema: merge_into_join.table_schema.clone(),
             })
         }
         Ok(Plan::MergeInto(Box::new(MergeInto {
diff --git a/src/query/sql/src/planner/plans/merge_into.rs b/src/query/sql/src/planner/plans/merge_into.rs
@@ -14,6 +14,7 @@
 
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::sync::Arc;
 
 use databend_common_ast::ast::TableAlias;
 use databend_common_exception::ErrorCode;
@@ -24,6 +25,7 @@ use databend_common_expression::DataField;
 use databend_common_expression::DataSchemaRef;
 use databend_common_expression::DataSchemaRefExt;
 use databend_common_expression::FieldIndex;
+use databend_common_expression::TableSchema;
 use databend_common_meta_types::MetaId;
 
 use crate::binder::MergeIntoType;
@@ -77,6 +79,8 @@ pub struct MergeInto {
     // `update *`` or `update set t1.a = t2.a ...`, the right expr on the `=` must be only a column,
     // we don't support complex expressions.
     pub can_try_update_column_only: bool,
+
+    pub table_schema: Arc<TableSchema>,
 }
 
 impl std::fmt::Debug for MergeInto {
diff --git a/src/query/storages/fuse/src/operations/read/native_data_source_reader.rs b/src/query/storages/fuse/src/operations/read/native_data_source_reader.rs
@@ -33,8 +33,8 @@ use databend_common_pipeline_sources::SyncSourcer;
 use databend_common_sql::IndexType;
 use log::debug;
 
-use super::can_merge_into_target_build_bloom_filter;
 use super::native_data_source::NativeDataSource;
+use super::util::build_merge_into_source_build_bloom_info;
 use super::util::MergeIntoSourceBuildBloomInfo;
 use crate::io::AggIndexReader;
 use crate::io::BlockReader;
@@ -91,16 +91,11 @@ impl ReadNativeDataSource<true> {
             virtual_reader,
             table_schema,
             table_index,
-            merge_into_source_build_bloom_info: MergeIntoSourceBuildBloomInfo {
-                can_do_merge_into_rumtime_filter_bloom: can_merge_into_target_build_bloom_filter(
-                    ctx.clone(),
-                    table_index,
-                )?,
-                segment_infos: Default::default(),
-                catalog_info: merge_into_join.catalog_info.clone(),
-                table_info: merge_into_join.table_info.clone(),
-                database_name: merge_into_join.database_name.clone(),
-            },
+            merge_into_source_build_bloom_info: build_merge_into_source_build_bloom_info(
+                ctx,
+                table_index,
+                merge_into_join,
+            )?,
         })
     }
 }
@@ -135,16 +130,11 @@ impl ReadNativeDataSource<false> {
             virtual_reader,
             table_schema,
             table_index,
-            merge_into_source_build_bloom_info: MergeIntoSourceBuildBloomInfo {
-                can_do_merge_into_rumtime_filter_bloom: can_merge_into_target_build_bloom_filter(
-                    ctx.clone(),
-                    table_index,
-                )?,
-                segment_infos: Default::default(),
-                catalog_info: merge_into_join.catalog_info.clone(),
-                table_info: merge_into_join.table_info.clone(),
-                database_name: merge_into_join.database_name.clone(),
-            },
+            merge_into_source_build_bloom_info: build_merge_into_source_build_bloom_info(
+                ctx,
+                table_index,
+                merge_into_join,
+            )?,
         })))
     }
 }
diff --git a/src/query/storages/fuse/src/operations/read/parquet_data_source_reader.rs b/src/query/storages/fuse/src/operations/read/parquet_data_source_reader.rs
@@ -32,8 +32,8 @@ use databend_common_pipeline_sources::SyncSourcer;
 use databend_common_sql::IndexType;
 use log::debug;
 
-use super::can_merge_into_target_build_bloom_filter;
 use super::parquet_data_source::ParquetDataSource;
+use super::util::build_merge_into_source_build_bloom_info;
 use super::util::MergeIntoSourceBuildBloomInfo;
 use crate::fuse_part::FusePartInfo;
 use crate::io::AggIndexReader;
@@ -95,14 +95,11 @@ impl<const BLOCKING_IO: bool> ReadParquetDataSource<BLOCKING_IO> {
                 index_reader,
                 virtual_reader,
                 table_schema,
-                merge_into_source_build_bloom_info: MergeIntoSourceBuildBloomInfo {
-                    can_do_merge_into_rumtime_filter_bloom:
-                        can_merge_into_target_build_bloom_filter(ctx.clone(), table_index)?,
-                    segment_infos: Default::default(),
-                    catalog_info: merge_into_join.catalog_info.clone(),
-                    table_info: merge_into_join.table_info.clone(),
-                    database_name: merge_into_join.database_name.clone(),
-                },
+                merge_into_source_build_bloom_info: build_merge_into_source_build_bloom_info(
+                    ctx,
+                    table_index,
+                    merge_into_join,
+                )?,
             })
         } else {
             let merge_into_join = ctx.get_merge_into_join();
@@ -121,14 +118,11 @@ impl<const BLOCKING_IO: bool> ReadParquetDataSource<BLOCKING_IO> {
                 index_reader,
                 virtual_reader,
                 table_schema,
-                merge_into_source_build_bloom_info: MergeIntoSourceBuildBloomInfo {
-                    can_do_merge_into_rumtime_filter_bloom:
-                        can_merge_into_target_build_bloom_filter(ctx.clone(), table_index)?,
-                    segment_infos: Default::default(),
-                    catalog_info: merge_into_join.catalog_info.clone(),
-                    table_info: merge_into_join.table_info.clone(),
-                    database_name: merge_into_join.database_name.clone(),
-                },
+                merge_into_source_build_bloom_info: build_merge_into_source_build_bloom_info(
+                    ctx,
+                    table_index,
+                    merge_into_join,
+                )?,
             })))
         }
     }
diff --git a/src/query/storages/fuse/src/operations/read/runtime_filter_prunner.rs b/src/query/storages/fuse/src/operations/read/runtime_filter_prunner.rs
@@ -182,6 +182,7 @@ pub fn runtime_filter_pruner(
             segment_idx, block_idx
         );
         let block_meta = segment_info.blocks[block_idx].clone();
+        // iterate every probe key name to do prune
     }
 
     Ok(false)
diff --git a/src/query/storages/fuse/src/operations/read/util.rs b/src/query/storages/fuse/src/operations/read/util.rs
@@ -15,13 +15,16 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
+use databend_common_catalog::merge_into_join::MergeIntoJoin;
 use databend_common_catalog::merge_into_join::MergeIntoJoinType;
 use databend_common_catalog::plan::gen_mutation_stream_meta;
 use databend_common_catalog::plan::InternalColumnMeta;
 use databend_common_catalog::table_context::TableContext;
+use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::BlockMetaInfoPtr;
 use databend_common_expression::DataBlock;
+use databend_common_expression::FieldIndex;
 use databend_common_expression::Scalar;
 use databend_common_meta_app::schema::CatalogInfo;
 use databend_common_meta_app::schema::TableInfo;
@@ -37,6 +40,7 @@ pub struct MergeIntoSourceBuildBloomInfo {
     pub table_info: Option<TableInfo>,
     pub catalog_info: Option<CatalogInfo>,
     pub database_name: String,
+    pub bloom_indexes: Vec<FieldIndex>,
 }
 
 pub fn need_reserve_block_info(ctx: Arc<dyn TableContext>, table_idx: usize) -> (bool, bool) {
@@ -68,6 +72,7 @@ pub fn can_merge_into_target_build_bloom_filter(
         assert!(merge_into_join.database_name.as_str() != "");
         assert!(matches!(merge_into_join.catalog_info, Some(_)));
         assert!(matches!(merge_into_join.table_info, Some(_)));
+        assert!(matches!(merge_into_join.table_schema, Some(_)));
     }
     Ok(enabled)
 }
@@ -116,3 +121,36 @@ pub(crate) fn add_data_block_meta(
     }
     block.add_meta(meta)
 }
+
+pub fn build_merge_into_source_build_bloom_info(
+    ctx: Arc<dyn TableContext>,
+    table_index: usize,
+    merge_into_join: MergeIntoJoin,
+) -> Result<MergeIntoSourceBuildBloomInfo> {
+    let enabled_bloom_filter = can_merge_into_target_build_bloom_filter(ctx.clone(), table_index)?;
+    let bloom_indexes = if enabled_bloom_filter {
+        ctx.get_merge_into_source_build_bloom_probe_keys(table_index)
+            .iter()
+            .try_fold(Vec::new(), |mut acc, probe_key_name| {
+                let table_schema = merge_into_join.table_schema.as_ref().ok_or_else(|| {
+                    ErrorCode::Internal(
+                        "can't get merge into target table schema when build bloom info, it's a bug",
+                    )
+                })?;
+                let index = table_schema.index_of(probe_key_name)?;
+                acc.push(index);
+                Ok::<_, ErrorCode>(acc)
+            })?
+    } else {
+        vec![]
+    };
+
+    Ok(MergeIntoSourceBuildBloomInfo {
+        can_do_merge_into_rumtime_filter_bloom: enabled_bloom_filter,
+        segment_infos: Default::default(),
+        catalog_info: merge_into_join.catalog_info.clone(),
+        table_info: merge_into_join.table_info.clone(),
+        database_name: merge_into_join.database_name.clone(),
+        bloom_indexes,
+    })
+}

Original file line number	Diff line number	Diff line change
`@@ -14,6 +14,7 @@`
`14`	`14`
`15`	`15`	`use std::sync::Arc;`
`16`	`16`
	`17`	`+use databend_common_expression::TableSchemaRef;`
`17`	`18`	`use databend_common_meta_app::schema::CatalogInfo;`
`18`	`19`	`use databend_common_meta_app::schema::TableInfo;`
`19`	`20`	`use databend_storages_common_table_meta::meta::Location;`
`@@ -61,6 +62,7 @@ pub struct MergeIntoJoin {`
`61`	`62`	`pub table_info: Option<TableInfo>,`
`62`	`63`	`pub catalog_info: Option<CatalogInfo>,`
`63`	`64`	`pub database_name: String,`
	`65`	`+ pub table_schema: Option<TableSchemaRef>,`
`64`	`66`	`}`
`65`	`67`
`66`	`68`	`impl Default for MergeIntoJoin {`
`@@ -73,6 +75,7 @@ impl Default for MergeIntoJoin {`
`73`	`75`	`table_info: None,`
`74`	`76`	`catalog_info: None,`
`75`	`77`	`database_name: Default::default(),`
	`78`	`+ table_schema: None,`
`76`	`79`	`}`
`77`	`80`	`}`
`78`	`81`	`}`
Original file line number	Diff line number	Diff line change
`@@ -488,6 +488,7 @@ impl Binder {`
`488`	`488`	`row_id_index: column_binding.index,`
`489`	`489`	`split_idx,`
`490`	`490`	`can_try_update_column_only: self.can_try_update_column_only(&matched_clauses),`
	`491`	`+ table_schema,`
`491`	`492`	`})`
`492`	`493`	`}`
`493`	`494`