fix virtual column builder

Author: zhyass

Cargo.lock

@@ -50,6 +50,7 @@ backoff = { workspace = true, features = ["futures", "tokio"] }
 bytes = { workspace = true }
 chrono = { workspace = true }
 enum-as-inner = { workspace = true }
+enum_dispatch = { workspace = true }
 fastrace = { workspace = true }
 futures = { workspace = true }
 futures-util = { workspace = true }

src/query/storages/fuse/Cargo.toml

@@ -35,6 +35,7 @@ use databend_common_expression::TableSchemaRef;
 use databend_common_expression::ORIGIN_BLOCK_ROW_NUM_COLUMN_ID;
 use databend_common_io::constants::DEFAULT_BLOCK_BUFFER_SIZE;
 use databend_common_native::write::NativeWriter;
+use databend_common_native::write::WriteOptions;
 use databend_common_sql::executor::physical_plans::MutationKind;
 use databend_storages_common_index::BloomIndex;
 use databend_storages_common_index::BloomIndexBuilder;
@@ -54,12 +55,14 @@ use crate::io::create_inverted_index_builders;
 use crate::io::write::stream::cluster_statistics::ClusterStatisticsBuilder;
 use crate::io::write::stream::cluster_statistics::ClusterStatisticsState;
 use crate::io::write::stream::ColumnStatisticsState;
+use crate::io::write::stream::VirtualColumnWriter;
 use crate::io::write::InvertedIndexState;
 use crate::io::BlockSerialization;
 use crate::io::BloomIndexState;
 use crate::io::InvertedIndexBuilder;
 use crate::io::InvertedIndexWriter;
 use crate::io::TableMetaLocationGenerator;
+use crate::io::VirtualColumnBuilder;
 use crate::io::WriteSettings;
 use crate::operations::column_parquet_metas;
 use crate::FuseStorageFormat;
@@ -149,6 +152,7 @@ pub struct StreamBlockBuilder {
     block_writer: BlockWriterImpl,
     inverted_index_writers: Vec<InvertedIndexWriter>,
     bloom_index_builder: BloomIndexBuilder,
+    virtual_column_writer: Option<VirtualColumnWriter>,
 
     cluster_stats_state: ClusterStatisticsState,
     column_stats_state: ColumnStatisticsState,
@@ -187,7 +191,7 @@ impl StreamBlockBuilder {
                 let writer = NativeWriter::new(
                     buffer,
                     properties.source_schema.as_ref().clone(),
-                    databend_common_native::write::WriteOptions {
+                    WriteOptions {
                         default_compression: properties.write_settings.table_compression.into(),
                         max_page_size: Some(properties.write_settings.max_page_size),
                         default_compress_ratio,
@@ -216,11 +220,18 @@ impl StreamBlockBuilder {
             ClusterStatisticsState::new(properties.cluster_stats_builder.clone());
         let column_stats_state =
             ColumnStatisticsState::new(&properties.stats_columns, &properties.distinct_columns);
+        let virtual_column_writer = properties.virtual_column_builder.as_ref().map(|builder| {
+            VirtualColumnWriter::create(
+                builder.clone(),
+                properties.write_settings.table_compression,
+            )
+        });
 
         Ok(StreamBlockBuilder {
             properties,
             block_writer,
             inverted_index_writers,
+            virtual_column_writer,
             bloom_index_builder,
             row_count: 0,
             block_size: 0,
@@ -257,6 +268,9 @@ impl StreamBlockBuilder {
         for writer in self.inverted_index_writers.iter_mut() {
             writer.add_block(&self.properties.source_schema, &block)?;
         }
+        if let Some(writer) = self.virtual_column_writer.as_mut() {
+            writer.add_block(&block)?;
+        }
 
         self.row_count += block.num_rows();
         self.block_size += block.estimate_block_size();
@@ -303,6 +317,12 @@ impl StreamBlockBuilder {
             inverted_index_states.push(inverted_index_state);
         }
 
+        let virtual_column_state = if let Some(writer) = self.virtual_column_writer.take() {
+            Some(writer.finalize(&block_location)?)
+        } else {
+            None
+        };
+
         let col_metas = self.block_writer.finish(&self.properties.source_schema)?;
         let block_raw_data = mem::take(self.block_writer.inner_mut());
 
@@ -333,15 +353,18 @@ impl StreamBlockBuilder {
             compression: self.properties.write_settings.table_compression.into(),
             inverted_index_size,
             create_on: Some(Utc::now()),
-            ngram_filter_index_size: None,
+            ngram_filter_index_size: bloom_index_state
+                .as_ref()
+                .map(|v| v.ngram_size)
+                .unwrap_or_default(),
             virtual_block_meta: None,
         };
         let serialized = BlockSerialization {
             block_raw_data,
             block_meta,
             bloom_index_state,
             inverted_index_states,
-            virtual_column_state: None,
+            virtual_column_state,
         };
         Ok(serialized)
     }
@@ -361,6 +384,7 @@ pub struct StreamBlockProperties {
     bloom_columns_map: BTreeMap<FieldIndex, TableField>,
     ngram_args: Vec<NgramArgs>,
     inverted_index_builders: Vec<InvertedIndexBuilder>,
+    virtual_column_builder: Option<VirtualColumnBuilder>,
     table_meta_timestamps: TableMetaTimestamps,
 }
 
@@ -406,6 +430,15 @@ impl StreamBlockProperties {
             .collect::<HashSet<_>>();
 
         let inverted_index_builders = create_inverted_index_builders(&table.table_info.meta);
+        let virtual_column_builder = if ctx
+            .get_settings()
+            .get_enable_refresh_virtual_column_after_write()
+            .unwrap_or_default()
+        {
+            VirtualColumnBuilder::try_create(ctx.clone(), table, source_schema.clone()).ok()
+        } else {
+            None
+        };
 
         let cluster_stats_builder =
             ClusterStatisticsBuilder::try_create(table, ctx.clone(), &source_schema)?;
@@ -432,6 +465,7 @@ impl StreamBlockProperties {
             source_schema,
             write_settings,
             cluster_stats_builder,
+            virtual_column_builder,
             stats_columns,
             distinct_columns,
             bloom_columns_map,

src/query/storages/fuse/src/io/write/stream/block_builder.rs

@@ -24,56 +24,135 @@ use databend_common_expression::types::Decimal;
 use databend_common_expression::types::Decimal128Type;
 use databend_common_expression::types::Decimal256Type;
 use databend_common_expression::types::Decimal64Type;
+use databend_common_expression::types::Float32Type;
+use databend_common_expression::types::Float64Type;
+use databend_common_expression::types::Int16Type;
+use databend_common_expression::types::Int32Type;
+use databend_common_expression::types::Int64Type;
+use databend_common_expression::types::Int8Type;
 use databend_common_expression::types::NumberDataType;
-use databend_common_expression::types::NumberType;
 use databend_common_expression::types::StringType;
 use databend_common_expression::types::TimestampType;
+use databend_common_expression::types::UInt16Type;
+use databend_common_expression::types::UInt32Type;
+use databend_common_expression::types::UInt64Type;
+use databend_common_expression::types::UInt8Type;
 use databend_common_expression::types::ValueType;
-use databend_common_expression::with_number_mapped_type;
 use databend_common_expression::Column;
 use databend_common_expression::Scalar;
 use databend_common_expression::ScalarRef;
 use databend_common_expression::SELECTIVITY_THRESHOLD;
 use databend_storages_common_table_meta::meta::ColumnStatistics;
+use enum_dispatch::enum_dispatch;
 
 use crate::statistics::Trim;
 
-pub trait ColumnStatisticsBuilder: Send + Sync {
-    fn update_column(&mut self, column: &Column);
+pub type CommonBuilder<T> = GenericColumnStatisticsBuilder<T, CommonAdapter>;
+pub type DecimalBuilder<T> = GenericColumnStatisticsBuilder<T, DecimalAdapter>;
+
+#[enum_dispatch(ColumnStatsOps)]
+pub enum ColumnStatisticsBuilder {
+    Int8(CommonBuilder<Int8Type>),
+    Int16(CommonBuilder<Int16Type>),
+    Int32(CommonBuilder<Int32Type>),
+    Int64(CommonBuilder<Int64Type>),
+    UInt8(CommonBuilder<UInt8Type>),
+    UInt16(CommonBuilder<UInt16Type>),
+    UInt32(CommonBuilder<UInt32Type>),
+    UInt64(CommonBuilder<UInt64Type>),
+    Float32(CommonBuilder<Float32Type>),
+    Float64(CommonBuilder<Float64Type>),
+    String(CommonBuilder<StringType>),
+    Date(CommonBuilder<DateType>),
+    Timestamp(CommonBuilder<TimestampType>),
+    Decimal64(DecimalBuilder<Decimal64Type>),
+    Decimal128(DecimalBuilder<Decimal128Type>),
+    Decimal256(DecimalBuilder<Decimal256Type>),
+}
 
+#[enum_dispatch]
+pub trait ColumnStatsOps {
+    fn update_column(&mut self, column: &Column);
     fn update_scalar(&mut self, scalar: &ScalarRef, num_rows: usize, data_type: &DataType);
+    fn finalize(self) -> Result<ColumnStatistics>;
+}
+
+impl<T, A> ColumnStatsOps for GenericColumnStatisticsBuilder<T, A>
+where
+    T: ValueType + Send + Sync,
+    T::Scalar: Send + Sync,
+    A: ColumnStatisticsAdapter<T> + 'static,
+    for<'a, 'b> T::ScalarRef<'a>: PartialOrd<T::ScalarRef<'b>>,
+{
+    fn update_column(&mut self, column: &Column) {
+        GenericColumnStatisticsBuilder::update_column(self, column);
+    }
+
+    fn update_scalar(&mut self, scalar: &ScalarRef, num_rows: usize, data_type: &DataType) {
+        GenericColumnStatisticsBuilder::update_scalar(self, scalar, num_rows, data_type);
+    }
+
+    fn finalize(self) -> Result<ColumnStatistics> {
+        GenericColumnStatisticsBuilder::finalize(self)
+    }
+}
 
-    fn finalize(self: Box<Self>) -> Result<ColumnStatistics>;
+macro_rules! create_builder_for_type {
+    ($data_type:expr, $variant:ident, $type:ty) => {
+        ColumnStatisticsBuilder::$variant(CommonBuilder::<$type>::create($data_type))
+    };
+    ($data_type:expr, $variant:ident, $type:ty, decimal) => {
+        ColumnStatisticsBuilder::$variant(DecimalBuilder::<$type>::create($data_type))
+    };
 }
 
-pub fn create_column_stats_builder(data_type: &DataType) -> Box<dyn ColumnStatisticsBuilder> {
+pub fn create_column_stats_builder(data_type: &DataType) -> ColumnStatisticsBuilder {
     let inner_type = data_type.remove_nullable();
-    with_number_mapped_type!(|NUM_TYPE| match inner_type {
-        DataType::Number(NumberDataType::NUM_TYPE) => {
-            GenericColumnStatisticsBuilder::<NumberType<NUM_TYPE>, CommonAdapter>::create(
-                inner_type,
-            )
+    match inner_type {
+        DataType::Number(NumberDataType::Int8) => {
+            create_builder_for_type!(inner_type, Int8, Int8Type)
+        }
+        DataType::Number(NumberDataType::Int16) => {
+            create_builder_for_type!(inner_type, Int16, Int16Type)
+        }
+        DataType::Number(NumberDataType::Int32) => {
+            create_builder_for_type!(inner_type, Int32, Int32Type)
+        }
+        DataType::Number(NumberDataType::Int64) => {
+            create_builder_for_type!(inner_type, Int64, Int64Type)
         }
-        DataType::String => {
-            GenericColumnStatisticsBuilder::<StringType, CommonAdapter>::create(inner_type)
+        DataType::Number(NumberDataType::UInt8) => {
+            create_builder_for_type!(inner_type, UInt8, UInt8Type)
         }
-        DataType::Date => {
-            GenericColumnStatisticsBuilder::<DateType, CommonAdapter>::create(inner_type)
+        DataType::Number(NumberDataType::UInt16) => {
+            create_builder_for_type!(inner_type, UInt16, UInt16Type)
         }
-        DataType::Timestamp => {
-            GenericColumnStatisticsBuilder::<TimestampType, CommonAdapter>::create(inner_type)
+        DataType::Number(NumberDataType::UInt32) => {
+            create_builder_for_type!(inner_type, UInt32, UInt32Type)
         }
+        DataType::Number(NumberDataType::UInt64) => {
+            create_builder_for_type!(inner_type, UInt64, UInt64Type)
+        }
+        DataType::Number(NumberDataType::Float32) => {
+            create_builder_for_type!(inner_type, Float32, Float32Type)
+        }
+        DataType::Number(NumberDataType::Float64) => {
+            create_builder_for_type!(inner_type, Float64, Float64Type)
+        }
+        DataType::String => create_builder_for_type!(inner_type, String, StringType),
+        DataType::Date => create_builder_for_type!(inner_type, Date, DateType),
+        DataType::Timestamp => create_builder_for_type!(inner_type, Timestamp, TimestampType),
         DataType::Decimal(size) => {
             if size.can_carried_by_64() {
-                GenericColumnStatisticsBuilder::<Decimal64Type, DecimalAdapter>::create(inner_type)
+                create_builder_for_type!(inner_type, Decimal64, Decimal64Type, decimal)
             } else if size.can_carried_by_128() {
-                GenericColumnStatisticsBuilder::<Decimal128Type, DecimalAdapter>::create(inner_type)
+                create_builder_for_type!(inner_type, Decimal128, Decimal128Type, decimal)
             } else {
-                GenericColumnStatisticsBuilder::<Decimal256Type, DecimalAdapter>::create(inner_type)
+                create_builder_for_type!(inner_type, Decimal256, Decimal256Type, decimal)
             }
         }
         _ => unreachable!("Unsupported data type: {:?}", data_type),
-    })
+    }
 }
 
 pub trait ColumnStatisticsAdapter<T: ValueType>: Send + Sync {
@@ -86,7 +165,7 @@ pub trait ColumnStatisticsAdapter<T: ValueType>: Send + Sync {
     fn update_value(value: &mut Self::Value, scalar: T::ScalarRef<'_>, ordering: Ordering);
 }
 
-struct CommonAdapter;
+pub struct CommonAdapter;
 
 impl<T> ColumnStatisticsAdapter<T> for CommonAdapter
 where
@@ -111,7 +190,7 @@ where
     }
 }
 
-struct DecimalAdapter;
+pub struct DecimalAdapter;
 
 impl<T> ColumnStatisticsAdapter<T> for DecimalAdapter
 where
@@ -137,7 +216,7 @@ where
     }
 }
 
-struct GenericColumnStatisticsBuilder<T, A>
+pub struct GenericColumnStatisticsBuilder<T, A>
 where
     T: ValueType,
     A: ColumnStatisticsAdapter<T>,
@@ -158,15 +237,15 @@ where
     A: ColumnStatisticsAdapter<T> + 'static,
     for<'a, 'b> T::ScalarRef<'a>: PartialOrd<T::ScalarRef<'b>>,
 {
-    fn create(data_type: DataType) -> Box<dyn ColumnStatisticsBuilder> {
-        Box::new(Self {
+    fn create(data_type: DataType) -> Self {
+        Self {
             min: None,
             max: None,
             null_count: 0,
             in_memory_size: 0,
             data_type,
             _phantom: PhantomData,
-        })
+        }
     }
 
     fn add_batch<'a, I>(&mut self, mut iter: I)
@@ -201,15 +280,7 @@ where
             self.max = Some(A::scalar_to_value(max));
         }
     }
-}
 
-impl<T, A> ColumnStatisticsBuilder for GenericColumnStatisticsBuilder<T, A>
-where
-    T: ValueType + Send + Sync,
-    T::Scalar: Send + Sync,
-    A: ColumnStatisticsAdapter<T> + 'static,
-    for<'a, 'b> T::ScalarRef<'a>: PartialOrd<T::ScalarRef<'b>>,
-{
     fn update_column(&mut self, column: &Column) {
         self.in_memory_size += column.memory_size();
         if column.len() == 0 {
@@ -265,7 +336,7 @@ where
         self.add(val.clone(), val);
     }
 
-    fn finalize(self: Box<Self>) -> Result<ColumnStatistics> {
+    fn finalize(self) -> Result<ColumnStatistics> {
         let min = if let Some(v) = self.min {
             let v = A::value_to_scalar(v);
             // safe upwrap.

src/query/storages/fuse/src/io/write/stream/column_statistics_builder.rs

@@ -26,10 +26,11 @@ use crate::io::write::stream::create_column_ndv_estimator;
 use crate::io::write::stream::create_column_stats_builder;
 use crate::io::write::stream::ColumnNDVEstimator;
 use crate::io::write::stream::ColumnStatisticsBuilder;
+use crate::io::write::stream::ColumnStatsOps;
 use crate::statistics::traverse_values_dfs;
 
 pub struct ColumnStatisticsState {
-    col_stats: HashMap<ColumnId, Box<dyn ColumnStatisticsBuilder>>,
+    col_stats: HashMap<ColumnId, ColumnStatisticsBuilder>,
     distinct_columns: HashMap<ColumnId, Box<dyn ColumnNDVEstimator>>,
 }