fix

Author: zhyass

src/query/service/src/test_kits/block_writer.rs

@@ -21,12 +21,12 @@ use databend_common_io::constants::DEFAULT_BLOCK_BUFFER_SIZE;
 use databend_common_io::constants::DEFAULT_BLOCK_INDEX_BUFFER_SIZE;
 use databend_common_sql::BloomIndexColumns;
 use databend_common_storages_fuse::io::serialize_block;
-use databend_common_storages_fuse::io::BloomIndexBuilder;
 use databend_common_storages_fuse::io::TableMetaLocationGenerator;
 use databend_common_storages_fuse::io::WriteSettings;
 use databend_common_storages_fuse::FuseStorageFormat;
 use databend_storages_common_blocks::blocks_to_parquet;
 use databend_storages_common_index::BloomIndex;
+use databend_storages_common_index::BloomIndexBuilder;
 use databend_storages_common_table_meta::meta::BlockMeta;
 use databend_storages_common_table_meta::meta::ClusterStatistics;
 use databend_storages_common_table_meta::meta::Compression;

src/query/service/tests/it/storages/fuse/mod.rs

@@ -14,7 +14,6 @@
 
 #![allow(clippy::too_many_arguments)]
 mod bloom_index_meta_size;
-mod bloom_pruner;
 mod conflict;
 mod io;
 mod meta;

src/query/storages/common/index/src/bloom_index.rs

@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::collections::HashMap;
 use std::ops::ControlFlow;
+use std::ops::Deref;
 use std::sync::Arc;
 
 use databend_common_ast::Span;
@@ -24,9 +26,11 @@ use databend_common_expression::eval_function;
 use databend_common_expression::expr::*;
 use databend_common_expression::types::boolean::BooleanDomain;
 use databend_common_expression::types::nullable::NullableDomain;
+use databend_common_expression::types::AnyType;
 use databend_common_expression::types::Bitmap;
 use databend_common_expression::types::Buffer;
 use databend_common_expression::types::DataType;
+use databend_common_expression::types::MapType;
 use databend_common_expression::types::NullableType;
 use databend_common_expression::types::Number;
 use databend_common_expression::types::NumberDataType;
@@ -35,16 +39,20 @@ use databend_common_expression::types::ValueType;
 use databend_common_expression::visit_expr;
 use databend_common_expression::BlockEntry;
 use databend_common_expression::Column;
+use databend_common_expression::ColumnBuilder;
 use databend_common_expression::ColumnId;
 use databend_common_expression::ConstantFolder;
 use databend_common_expression::DataBlock;
 use databend_common_expression::Domain;
 use databend_common_expression::Expr;
 use databend_common_expression::ExprVisitor;
+use databend_common_expression::FieldIndex;
 use databend_common_expression::FunctionContext;
 use databend_common_expression::Scalar;
+use databend_common_expression::ScalarRef;
 use databend_common_expression::TableDataType;
 use databend_common_expression::TableField;
+use databend_common_expression::TableSchema;
 use databend_common_expression::TableSchemaRef;
 use databend_common_expression::Value;
 use databend_common_functions::BUILTIN_FUNCTIONS;
@@ -59,8 +67,11 @@ use serde::Serialize;
 use super::eliminate_cast::is_injective_cast;
 use crate::eliminate_cast::cast_const;
 use crate::filters::BlockBloomFilterIndexVersion;
+use crate::filters::BlockFilter;
 use crate::filters::Filter;
+use crate::filters::FilterBuilder;
 use crate::filters::V2BloomBlock;
+use crate::filters::Xor8Builder;
 use crate::filters::Xor8Filter;
 use crate::statistics_to_domain;
 use crate::Index;
@@ -439,6 +450,175 @@ impl BloomIndex {
     }
 }
 
+pub struct BloomIndexBuilder {
+    func_ctx: FunctionContext,
+    columns: Vec<ColumnXor8Builder>,
+}
+
+struct ColumnXor8Builder {
+    index: FieldIndex,
+    field: TableField,
+    builder: Xor8Builder,
+}
+
+impl BloomIndexBuilder {
+    pub fn create(
+        func_ctx: FunctionContext,
+        bloom_columns_map: BTreeMap<FieldIndex, TableField>,
+    ) -> Self {
+        let columns = bloom_columns_map
+            .iter()
+            .map(|(&index, field)| ColumnXor8Builder {
+                index,
+                field: field.clone(),
+                builder: Xor8Builder::create(),
+            })
+            .collect();
+        Self { func_ctx, columns }
+    }
+
+    pub fn add_block(&mut self, block: &DataBlock) -> Result<()> {
+        if block.is_empty() {
+            return Err(ErrorCode::BadArguments("block is empty"));
+        }
+        if block.num_columns() == 0 {
+            return Ok(());
+        }
+
+        let mut keys_to_remove = Vec::with_capacity(self.columns.len());
+        for (index, bloom_index_column) in self.columns.iter_mut().enumerate() {
+            let field_type = &block.get_by_offset(bloom_index_column.index).data_type;
+            if !Xor8Filter::supported_type(field_type) {
+                keys_to_remove.push(index);
+                continue;
+            }
+
+            let column = match &block.get_by_offset(bloom_index_column.index).value {
+                Value::Scalar(s) => {
+                    let builder = ColumnBuilder::repeat(&s.as_ref(), 1, field_type);
+                    builder.build()
+                }
+                Value::Column(c) => c.clone(),
+            };
+
+            let (column, data_type) = match field_type.remove_nullable() {
+                DataType::Map(box inner_ty) => {
+                    // Add bloom filter for the value of map type
+                    let map_column = if field_type.is_nullable() {
+                        let nullable_column =
+                            NullableType::<MapType<AnyType, AnyType>>::try_downcast_column(&column)
+                                .unwrap();
+                        nullable_column.column
+                    } else {
+                        MapType::<AnyType, AnyType>::try_downcast_column(&column).unwrap()
+                    };
+                    let column = map_column.underlying_column().values;
+
+                    let DataType::Tuple(kv_tys) = inner_ty else {
+                        unreachable!();
+                    };
+                    let val_type = kv_tys[1].clone();
+                    // Extract JSON value of string type to create bloom index,
+                    // other types of JSON value will be ignored.
+                    if val_type.remove_nullable() == DataType::Variant {
+                        let mut builder = ColumnBuilder::with_capacity(
+                            &DataType::Nullable(Box::new(DataType::String)),
+                            column.len(),
+                        );
+                        for val in column.iter() {
+                            if let ScalarRef::Variant(v) = val {
+                                let raw_jsonb = RawJsonb::new(v);
+                                if let Ok(Some(str_val)) = raw_jsonb.as_str() {
+                                    builder.push(ScalarRef::String(&str_val));
+                                    continue;
+                                }
+                            }
+                            builder.push_default();
+                        }
+                        let str_column = builder.build();
+                        if BloomIndex::check_large_string(&str_column) {
+                            keys_to_remove.push(index);
+                            continue;
+                        }
+                        let str_type = DataType::Nullable(Box::new(DataType::String));
+                        (str_column, str_type)
+                    } else {
+                        if BloomIndex::check_large_string(&column) {
+                            keys_to_remove.push(index);
+                            continue;
+                        }
+                        (column, val_type)
+                    }
+                }
+                _ => {
+                    if BloomIndex::check_large_string(&column) {
+                        keys_to_remove.push(index);
+                        continue;
+                    }
+                    (column, field_type.clone())
+                }
+            };
+
+            let (column, validity) =
+                BloomIndex::calculate_nullable_column_digest(&self.func_ctx, &column, &data_type)?;
+
+            // create filter per column
+            if validity.as_ref().map(|v| v.null_count()).unwrap_or(0) > 0 {
+                let validity = validity.unwrap();
+                let it = column.deref().iter().zip(validity.iter()).map(
+                    |(v, b)| {
+                        if !b {
+                            &0
+                        } else {
+                            v
+                        }
+                    },
+                );
+                bloom_index_column.builder.add_digests(it);
+            } else {
+                bloom_index_column.builder.add_digests(column.deref());
+            }
+        }
+        for k in keys_to_remove {
+            self.columns.remove(k);
+        }
+        Ok(())
+    }
+
+    pub fn finalize(mut self) -> Result<Option<BloomIndex>> {
+        let mut column_distinct_count = HashMap::with_capacity(self.columns.len());
+        let mut filters = Vec::with_capacity(self.columns.len());
+        let mut filter_fields = Vec::with_capacity(self.columns.len());
+        for column in self.columns.iter_mut() {
+            let filter = column.builder.build()?;
+            if let Some(len) = filter.len() {
+                if !matches!(
+                    column.field.data_type().remove_nullable(),
+                    TableDataType::Map(_) | TableDataType::Variant
+                ) {
+                    column_distinct_count.insert(column.field.column_id, len);
+                }
+            }
+            let filter_name =
+                BloomIndex::build_filter_column_name(BlockFilter::VERSION, &column.field)?;
+            filter_fields.push(TableField::new(&filter_name, TableDataType::Binary));
+            filters.push(Arc::new(filter));
+        }
+
+        if filter_fields.is_empty() {
+            return Ok(None);
+        }
+        let filter_schema = Arc::new(TableSchema::new(filter_fields));
+        Ok(Some(BloomIndex {
+            func_ctx: self.func_ctx,
+            version: BlockFilter::VERSION,
+            filter_schema,
+            filters,
+            column_distinct_count,
+        }))
+    }
+}
+
 struct Visitor<T: EqVisitor>(T);
 
 impl<T> ExprVisitor<String> for Visitor<T>

src/query/storages/common/index/src/lib.rs

@@ -25,6 +25,7 @@ mod page_index;
 mod range_index;
 
 pub use bloom_index::BloomIndex;
+pub use bloom_index::BloomIndexBuilder;
 pub use bloom_index::BloomIndexMeta;
 pub use bloom_index::FilterEvalResult;
 pub use eliminate_cast::eliminate_cast;

src/query/service/tests/it/storages/fuse/bloom_pruner.rs renamed to src/query/storages/common/index/tests/it/bloom_pruner.rs

@@ -1,4 +1,4 @@
-// Copyright 2021 Datafuse Labs
+// Copyright 2022 Datafuse Labs
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -17,13 +17,13 @@ use std::collections::HashMap;
 use std::io::Write;
 use std::sync::Arc;
 
-use databend_common_column::buffer::Buffer;
 use databend_common_expression::type_check;
 use databend_common_expression::type_check::check_function;
 use databend_common_expression::types::map::KvColumn;
 use databend_common_expression::types::map::KvPair;
 use databend_common_expression::types::AnyType;
 use databend_common_expression::types::ArrayColumn;
+use databend_common_expression::types::Buffer;
 use databend_common_expression::types::DataType;
 use databend_common_expression::types::DateType;
 use databend_common_expression::types::Int16Type;
@@ -51,17 +51,17 @@ use databend_common_expression::TableSchema;
 use databend_common_expression::Value;
 use databend_common_functions::test_utils::parse_raw_expr;
 use databend_common_functions::BUILTIN_FUNCTIONS;
-use databend_common_storages_fuse::io::BloomIndexBuilder;
 use databend_storages_common_index::filters::Xor8Filter;
 use databend_storages_common_index::BloomIndex;
+use databend_storages_common_index::BloomIndexBuilder;
 use databend_storages_common_index::FilterEvalResult;
 use databend_storages_common_index::Index;
 use databend_storages_common_table_meta::meta::ColumnStatistics;
 use goldenfile::Mint;
 
 #[test]
 fn test_bloom_filter() {
-    let mut mint = Mint::new("tests/it/storages/testdata");
+    let mut mint = Mint::new("tests/it/testdata");
     let file = &mut mint.new_goldenfile("test_bloom_filter.txt").unwrap();
 
     test_base(file);

src/query/storages/common/index/tests/it/main.rs

@@ -15,6 +15,7 @@
 #![feature(box_patterns)]
 #![allow(clippy::uninlined_format_args)]
 
+mod bloom_pruner;
 mod eliminate_cast;
 mod page_pruner;
 mod range_pruner;

src/query/service/tests/it/storages/testdata/test_bloom_filter.txt renamed to src/query/storages/common/index/tests/it/testdata/test_bloom_filter.txt

@@ -44,11 +44,12 @@ pub use write::write_data;
 pub use write::BlockBuilder;
 pub use write::BlockSerialization;
 pub use write::BlockWriter;
-pub use write::BloomIndexBuilder;
 pub use write::BloomIndexRebuilder;
 pub use write::BloomIndexState;
 pub use write::CachedMetaWriter;
 pub use write::InvertedIndexBuilder;
 pub use write::InvertedIndexWriter;
 pub use write::MetaWriter;
+pub(crate) use write::StreamBlockBuilder;
+pub(crate) use write::StreamBlockProperties;
 pub use write::WriteSettings;