fix: incorrect index size during merge segments (#10817)

Author: zhyass

src/query/service/tests/it/storages/fuse/operations/commit.rs

@@ -33,6 +33,7 @@ use common_catalog::table_context::StageAttachment;
 use common_catalog::table_context::TableContext;
 use common_exception::ErrorCode;
 use common_exception::Result;
+use common_expression::BlockThresholds;
 use common_expression::DataBlock;
 use common_expression::FunctionContext;
 use common_io::prelude::FormatSettings;
@@ -71,15 +72,22 @@ use common_settings::ChangeValue;
 use common_settings::Settings;
 use common_storage::DataOperator;
 use common_storage::StageFileInfo;
+use common_storages_fuse::io::SegmentWriter;
+use common_storages_fuse::io::TableMetaLocationGenerator;
 use common_storages_fuse::operations::AppendOperationLogEntry;
+use common_storages_fuse::statistics::reducers::reduce_block_metas;
 use common_storages_fuse::FuseTable;
 use common_storages_fuse::FUSE_TBL_SNAPSHOT_PREFIX;
 use databend_query::sessions::QueryContext;
 use futures::TryStreamExt;
+use rand::thread_rng;
+use rand::Rng;
 use storages_common_table_meta::meta::SegmentInfo;
 use storages_common_table_meta::meta::Statistics;
 use walkdir::WalkDir;
 
+use crate::storages::fuse::block_writer::BlockWriter;
+use crate::storages::fuse::operations::mutation::CompactSegmentTestFixture;
 use crate::storages::fuse::table_test_fixture::execute_query;
 use crate::storages::fuse::table_test_fixture::TestFixture;
 
@@ -301,6 +309,48 @@ async fn test_abort_on_error() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_merge_segments() -> common_exception::Result<()> {
+    let fixture = TestFixture::new().await;
+    let ctx = fixture.ctx();
+
+    let operator = ctx.get_data_operator()?.operator();
+    let data_accessor = operator.clone();
+    let location_gen = TableMetaLocationGenerator::with_prefix("test/".to_owned());
+    let block_writer = BlockWriter::new(&data_accessor, &location_gen);
+    let segment_writer = SegmentWriter::new(&data_accessor, &location_gen);
+
+    let mut rand = thread_rng();
+    let number_of_segments: usize = rand.gen_range(1..10);
+    let mut block_number_of_segments = Vec::with_capacity(number_of_segments);
+    let mut rows_per_blocks = Vec::with_capacity(number_of_segments);
+    for _ in 0..number_of_segments {
+        block_number_of_segments.push(rand.gen_range(10..30));
+        rows_per_blocks.push(rand.gen_range(1..8));
+    }
+
+    let threshold = BlockThresholds {
+        max_rows_per_block: 5,
+        min_rows_per_block: 4,
+        max_bytes_per_block: 1024,
+    };
+
+    let (locations, block_metas, segment_infos) = CompactSegmentTestFixture::gen_segments(
+        &block_writer,
+        &segment_writer,
+        &block_number_of_segments,
+        &rows_per_blocks,
+        threshold,
+    )
+    .await?;
+
+    let expect = reduce_block_metas(&block_metas, threshold)?;
+    let iter = locations.iter().zip(segment_infos.iter());
+    let (_, results) = FuseTable::merge_segments(iter)?;
+    assert_eq!(expect, results);
+    Ok(())
+}
+
 struct CtxDelegation {
     ctx: Arc<dyn TableContext>,
     catalog: Arc<FakedCatalog>,

src/query/service/tests/it/storages/fuse/operations/mutation/block_compact_mutator.rs

@@ -171,6 +171,7 @@ async fn test_safety() -> Result<()> {
             &segment_writer,
             &block_number_of_segments,
             &rows_per_blocks,
+            threshold,
         )
         .await?;
 

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

@@ -18,3 +18,4 @@ mod recluster_mutator;
 mod segments_compact_mutator;
 
 pub use deletion::do_deletion;
+pub use segments_compact_mutator::CompactSegmentTestFixture;

src/query/service/tests/it/storages/fuse/operations/mutation/segments_compact_mutator.rs

@@ -22,6 +22,7 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 use common_expression::types::number::NumberColumn;
 use common_expression::types::number::NumberScalar;
+use common_expression::BlockThresholds;
 use common_expression::Column;
 use common_expression::DataBlock;
 use common_expression::Scalar;
@@ -666,6 +667,7 @@ impl CompactSegmentTestFixture {
             &segment_writer,
             num_block_of_segments,
             &rows_per_block,
+            BlockThresholds::default(),
         )
         .await?;
         self.input_blocks = blocks;
@@ -682,6 +684,7 @@ impl CompactSegmentTestFixture {
         segment_writer: &SegmentWriter<'_>,
         block_num_of_segments: &[usize],
         rows_per_blocks: &[usize],
+        thresholds: BlockThresholds,
     ) -> Result<(Vec<Location>, Vec<BlockMeta>, Vec<SegmentInfo>)> {
         let mut locations = vec![];
         let mut collected_blocks = vec![];
@@ -693,7 +696,7 @@ impl CompactSegmentTestFixture {
         {
             let (schema, blocks) =
                 TestFixture::gen_sample_blocks_ex(*num_blocks, *rows_per_block, 1);
-            let mut stats_acc = StatisticsAccumulator::default();
+            let mut stats_acc = StatisticsAccumulator::new(thresholds);
             for block in blocks {
                 let block = block?;
                 let col_stats = gen_columns_statistics(&block, None, &schema)?;

src/query/storages/fuse/src/operations/commit.rs

@@ -466,9 +466,10 @@ impl FuseTable {
                 let stats = &segment_info.summary;
                 acc.row_count += stats.row_count;
                 acc.block_count += stats.block_count;
+                acc.perfect_block_count += stats.perfect_block_count;
                 acc.uncompressed_byte_size += stats.uncompressed_byte_size;
                 acc.compressed_byte_size += stats.compressed_byte_size;
-                acc.index_size = stats.index_size;
+                acc.index_size += stats.index_size;
                 acc.col_stats = if acc.col_stats.is_empty() {
                     stats.col_stats.clone()
                 } else {