refactor: move segment generation into TableMutationAggregator for MERGE INTO (#17800)

feat: fix merge into get too many segments

Author: zhyass

benchmark/clickbench/merge_into/queries.sql

@@ -1,3 +1,3 @@
-set enable_experimental_merge_into = 1;merge into target_table as t1 using (select * from source_table as t2)  on t1.l_partkey = t2.l_partkey and t1.l_orderkey = t2.l_orderkey and t1.l_suppkey = t2.l_suppkey and t1.l_linenumber = t2.l_linenumber when matched and t1.l_partkey >= 200000 then update * when matched then delete when not matched then insert *;
-set enable_experimental_merge_into = 1;merge into target_table as t1 using (select * from source_table as t2)  on t1.l_partkey = t2.l_partkey and t1.l_orderkey = t2.l_orderkey and t1.l_suppkey = t2.l_suppkey and t1.l_linenumber = t2.l_linenumber when matched and t1.l_partkey >= 200000 then update * when matched then delete when not matched then insert *;
-set enable_experimental_merge_into = 1;merge into target_table as t1 using (select * from source_table as t2)  on t1.l_partkey = t2.l_partkey and t1.l_orderkey = t2.l_orderkey and t1.l_suppkey = t2.l_suppkey and t1.l_linenumber = t2.l_linenumber when matched and t1.l_partkey >= 200000 then update * when matched then delete when not matched then insert *;
+merge into target_table as t1 using (select * from source_table as t2)  on t1.l_partkey = t2.l_partkey and t1.l_orderkey = t2.l_orderkey and t1.l_suppkey = t2.l_suppkey and t1.l_linenumber = t2.l_linenumber when matched and t1.l_partkey >= 200000 then update * when matched then delete when not matched then insert *;
+merge into target_table as t1 using (select * from source_table as t2)  on t1.l_partkey = t2.l_partkey and t1.l_orderkey = t2.l_orderkey and t1.l_suppkey = t2.l_suppkey and t1.l_linenumber = t2.l_linenumber when matched and t1.l_partkey >= 200000 then update * when matched then delete when not matched then insert *;
+merge into target_table as t1 using (select * from source_table as t2)  on t1.l_partkey = t2.l_partkey and t1.l_orderkey = t2.l_orderkey and t1.l_suppkey = t2.l_suppkey and t1.l_linenumber = t2.l_linenumber when matched and t1.l_partkey >= 200000 then update * when matched then delete when not matched then insert *;

src/query/service/src/pipelines/builders/builder_mutation.rs

@@ -34,7 +34,6 @@ use databend_common_pipeline_transforms::processors::TransformPipelineHelper;
 use databend_common_sql::binder::MutationStrategy;
 use databend_common_sql::executor::physical_plans::Mutation;
 use databend_common_sql::executor::physical_plans::MutationKind;
-use databend_common_storages_fuse::operations::new_serialize_segment_pipe_item;
 use databend_common_storages_fuse::operations::TransformSerializeBlock;
 use databend_common_storages_fuse::operations::UnMatchedExprs;
 use databend_common_storages_fuse::FuseTable;
@@ -61,14 +60,6 @@ impl PipelineBuilder {
         let io_request_semaphore =
             Arc::new(Semaphore::new(self.settings.get_max_threads()? as usize));
 
-        let serialize_segment_transform = new_serialize_segment_pipe_item(
-            InputPort::create(),
-            OutputPort::create(),
-            table,
-            block_thresholds,
-            merge_into.table_meta_timestamps,
-        )?;
-
         // For row_id port, create rowid_aggregate_mutator
         // For matched data port and unmatched port, do serialize
         let serialize_len = match merge_into.strategy {
@@ -131,51 +122,6 @@ impl PipelineBuilder {
             pipe_items,
         ));
 
-        // The complete pipeline:
-        // aggregate_mutator port               aggregate_mutator port
-        // serialize_block port0
-        // serialize_block port1     ======>    serialize_block port
-        // .......
-        let mut ranges = Vec::with_capacity(self.main_pipeline.output_len());
-        // row id port
-        let row_id_offset = if merge_into.need_match {
-            ranges.push(vec![0]);
-            1
-        } else {
-            0
-        };
-
-        // Resize data ports
-        debug_assert!(serialize_len > 0);
-        let mut vec = Vec::with_capacity(self.main_pipeline.output_len());
-        for idx in 0..serialize_len {
-            vec.push(idx + row_id_offset);
-        }
-        ranges.push(vec);
-        self.main_pipeline.resize_partial_one(ranges)?;
-
-        let pipe_items = {
-            let mut vec = Vec::with_capacity(2);
-            if merge_into.need_match {
-                // row_id port
-                vec.push(create_dummy_item());
-            }
-            // data port
-            vec.push(serialize_segment_transform);
-            vec
-        };
-
-        // The complete pipeline:
-        // output_port0: MutationLogs(row_id)
-        // output_port1: MutationLogs(data)
-        // 1. FullOperation and MatchedOnly: same as above
-        // 2. InsertOnly: no output_port0
-        let output_len = pipe_items.iter().map(|item| item.outputs_port.len()).sum();
-        self.main_pipeline.add_pipe(Pipe::create(
-            self.main_pipeline.output_len(),
-            output_len,
-            pipe_items,
-        ));
         Ok(())
     }
 

src/query/service/src/pipelines/builders/builder_replace_into.rs

@@ -49,7 +49,6 @@ use databend_common_sql::BindContext;
 use databend_common_sql::Metadata;
 use databend_common_sql::MetadataRef;
 use databend_common_sql::NameResolutionContext;
-use databend_common_storages_fuse::operations::new_serialize_segment_pipe_item;
 use databend_common_storages_fuse::operations::BroadcastProcessor;
 use databend_common_storages_fuse::operations::ReplaceIntoProcessor;
 use databend_common_storages_fuse::operations::TransformSerializeBlock;
@@ -134,13 +133,6 @@ impl PipelineBuilder {
         let mut block_builder = serialize_block_transform.get_block_builder();
         block_builder.source_schema = table.schema_with_stream();
 
-        let serialize_segment_transform = new_serialize_segment_pipe_item(
-            InputPort::create(),
-            OutputPort::create(),
-            table,
-            *block_thresholds,
-            replace.table_meta_timestamps,
-        )?;
         if !*need_insert {
             if segment_partition_num == 0 {
                 return Ok(());
@@ -204,15 +196,9 @@ impl PipelineBuilder {
         };
 
         // 4. connect with MergeIntoOperationAggregators
-        if segment_partition_num == 0 {
-            let dummy_item = create_dummy_item();
-            self.main_pipeline.add_pipe(Pipe::create(2, 2, vec![
-                serialize_segment_transform,
-                dummy_item,
-            ]));
-        } else {
+        if segment_partition_num != 0 {
             //      ┌──────────────────┐               ┌────────────────┐
-            // ────►│  SerializeBlock  ├──────────────►│SerializeSegment│
+            // ────►│  SerializeBlock  ├──────────────►│ DummyTransform │
             //      └──────────────────┘               └────────────────┘
             //
             //      ┌───────────────────┐              ┌──────────────────────┐
@@ -230,7 +216,7 @@ impl PipelineBuilder {
             let item_size = segment_partition_num + 1;
             let mut pipe_items = Vec::with_capacity(item_size);
             // setup the dummy transform
-            pipe_items.push(serialize_segment_transform);
+            pipe_items.push(create_dummy_item());
 
             let max_threads = self.settings.get_max_threads()?;
             let io_request_semaphore = Arc::new(Semaphore::new(max_threads as usize));

src/query/settings/src/settings_default.rs

@@ -680,7 +680,7 @@ impl DefaultSettings {
                 }),
                 ("enable_experimental_merge_into", DefaultSettingValue {
                     value: UserSettingValue::UInt64(1),
-                    desc: "Enables the experimental feature for 'MERGE INTO'.",
+                    desc: "Deprecated setting",
                     mode: SettingMode::Both,
                     scope: SettingScope::Both,
                     range: Some(SettingRange::Numeric(0..=1)),

src/query/settings/src/settings_getter_setter.rs

@@ -554,10 +554,6 @@ impl Settings {
         Ok(self.try_get_u64("enable_distributed_copy_into")? != 0)
     }
 
-    pub fn get_enable_experimental_merge_into(&self) -> Result<bool> {
-        Ok(self.try_get_u64("enable_experimental_merge_into")? != 0)
-    }
-
     pub fn get_enable_distributed_merge_into(&self) -> Result<bool> {
         Ok(self.try_get_u64("enable_distributed_merge_into")? != 0)
     }

src/query/storages/fuse/src/operations/common/meta/mutation_log.rs

@@ -39,7 +39,7 @@ pub enum MutationLogEntry {
         format_version: FormatVersion,
         summary: Statistics,
     },
-    ReclusterAppendBlock {
+    AppendBlock {
         block_meta: Arc<BlockMeta>,
     },
     DeletedBlock {

src/query/storages/fuse/src/operations/common/processors/mod.rs

@@ -26,6 +26,5 @@ pub use transform_block_writer::TransformBlockWriter;
 pub use transform_merge_commit_meta::TransformMergeCommitMeta;
 pub use transform_mutation_aggregator::TableMutationAggregator;
 pub use transform_serialize_block::TransformSerializeBlock;
-pub use transform_serialize_segment::new_serialize_segment_pipe_item;
 pub use transform_serialize_segment::new_serialize_segment_processor;
 pub use transform_serialize_segment::TransformSerializeSegment;

src/query/storages/fuse/src/operations/common/processors/transform_mutation_aggregator.rs

@@ -26,7 +26,6 @@ use databend_common_expression::BlockMetaInfoPtr;
 use databend_common_expression::BlockThresholds;
 use databend_common_expression::DataBlock;
 use databend_common_expression::TableSchemaRef;
-use databend_common_metrics::storage::metrics_inc_recluster_write_block_nums;
 use databend_common_pipeline_transforms::processors::AsyncAccumulatingTransform;
 use databend_common_sql::executor::physical_plans::MutationKind;
 use databend_storages_common_table_meta::meta::BlockMeta;
@@ -69,8 +68,7 @@ pub struct TableMutationAggregator {
 
     default_cluster_key_id: Option<u32>,
     base_segments: Vec<Location>,
-    // Used for recluster.
-    recluster_merged_blocks: Vec<Arc<BlockMeta>>,
+    merged_blocks: Vec<Arc<BlockMeta>>,
     set_hilbert_level: bool,
 
     mutations: HashMap<SegmentIndex, BlockMutations>,
@@ -103,6 +101,7 @@ impl AsyncAccumulatingTransform for TableMutationAggregator {
 
     #[async_backtrace::framed]
     async fn on_finish(&mut self, _output: bool) -> Result<Option<DataBlock>> {
+        self.generate_append_segments().await?;
         let mut new_segment_locs = Vec::new();
         new_segment_locs.extend(self.appended_segments.clone());
 
@@ -114,7 +113,35 @@ impl AsyncAccumulatingTransform for TableMutationAggregator {
                 },
                 self.schema.clone(),
             )),
-            MutationKind::Recluster => self.apply_recluster(&mut new_segment_locs).await?,
+            MutationKind::Recluster => {
+                let mut new_segments = std::mem::take(&mut self.appended_segments);
+                let new_segments_len = new_segments.len();
+                let removed_segments_len = self.removed_segment_indexes.len();
+                let replaced_segments_len = new_segments_len.min(removed_segments_len);
+                let mut appended_segments = Vec::new();
+                let mut replaced_segments = HashMap::with_capacity(replaced_segments_len);
+                if new_segments_len > removed_segments_len {
+                    // The remain new segments will be appended.
+                    let appended = new_segments.split_off(removed_segments_len);
+                    for location in appended.into_iter().rev() {
+                        appended_segments.push(location);
+                    }
+                }
+
+                for (i, location) in new_segments.into_iter().enumerate() {
+                    // The old segments will be replaced with the news.
+                    replaced_segments.insert(self.removed_segment_indexes[i], location);
+                }
+
+                ConflictResolveContext::ModifiedSegmentExistsInLatest(SnapshotChanges {
+                    appended_segments,
+                    removed_segment_indexes: self.removed_segment_indexes[replaced_segments_len..]
+                        .to_vec(),
+                    replaced_segments,
+                    removed_statistics: self.removed_statistics.clone(),
+                    merged_statistics: std::mem::take(&mut self.appended_statistics),
+                })
+            }
             _ => self.apply_mutation(&mut new_segment_locs).await?,
         };
 
@@ -131,7 +158,7 @@ impl TableMutationAggregator {
         table: &FuseTable,
         ctx: Arc<dyn TableContext>,
         base_segments: Vec<Location>,
-        recluster_merged_blocks: Vec<Arc<BlockMeta>>,
+        merged_blocks: Vec<Arc<BlockMeta>>,
         removed_segment_indexes: Vec<usize>,
         removed_statistics: Statistics,
         kind: MutationKind,
@@ -158,7 +185,7 @@ impl TableMutationAggregator {
             mutations: HashMap::new(),
             appended_segments: vec![],
             base_segments,
-            recluster_merged_blocks,
+            merged_blocks,
             appended_statistics: Statistics::default(),
             removed_segment_indexes,
             removed_statistics,
@@ -197,9 +224,8 @@ impl TableMutationAggregator {
                     }
                 }
             }
-            MutationLogEntry::ReclusterAppendBlock { block_meta } => {
-                metrics_inc_recluster_write_block_nums();
-                self.recluster_merged_blocks.push(block_meta);
+            MutationLogEntry::AppendBlock { block_meta } => {
+                self.merged_blocks.push(block_meta);
             }
             MutationLogEntry::DeletedBlock { index } => {
                 self.mutations
@@ -215,7 +241,6 @@ impl TableMutationAggregator {
                     self.default_cluster_key_id,
                 );
             }
-            MutationLogEntry::DoNothing => (),
             MutationLogEntry::AppendSegment {
                 segment_location,
                 format_version,
@@ -251,25 +276,26 @@ impl TableMutationAggregator {
                     self.default_cluster_key_id,
                 );
             }
+            MutationLogEntry::DoNothing => (),
         }
     }
 
-    async fn apply_recluster(
-        &mut self,
-        new_segment_locs: &mut Vec<Location>,
-    ) -> Result<ConflictResolveContext> {
-        // safe to unwrap.
-        let default_cluster_key_id = self.default_cluster_key_id.unwrap();
-        // sort ascending.
-        self.recluster_merged_blocks.sort_by(|a, b| {
-            sort_by_cluster_stats(&a.cluster_stats, &b.cluster_stats, default_cluster_key_id)
-        });
+    async fn generate_append_segments(&mut self) -> Result<()> {
+        if self.merged_blocks.is_empty() {
+            return Ok(());
+        }
+
+        if let Some(id) = self.default_cluster_key_id {
+            // sort ascending.
+            self.merged_blocks
+                .sort_by(|a, b| sort_by_cluster_stats(&a.cluster_stats, &b.cluster_stats, id));
+        }
 
         let mut tasks = Vec::new();
-        let merged_blocks = std::mem::take(&mut self.recluster_merged_blocks);
+        let merged_blocks = std::mem::take(&mut self.merged_blocks);
         let segments_num = (merged_blocks.len() / self.thresholds.block_per_segment).max(1);
         let chunk_size = merged_blocks.len().div_ceil(segments_num);
-        let default_cluster_key = Some(default_cluster_key_id);
+        let default_cluster_key = self.default_cluster_key_id;
         let thresholds = self.thresholds;
         let set_hilbert_level = self.set_hilbert_level;
         let kind = self.kind;
@@ -297,8 +323,7 @@ impl TableMutationAggregator {
         }
 
         let threads_nums = self.ctx.get_settings().get_max_threads()? as usize;
-
-        let mut new_segments = execute_futures_in_parallel(
+        let new_segments = execute_futures_in_parallel(
             tasks,
             threads_nums,
             threads_nums * 2,
@@ -308,41 +333,16 @@ impl TableMutationAggregator {
         .into_iter()
         .collect::<Result<Vec<_>>>()?;
 
-        let new_segments_len = new_segments.len();
-        let removed_segments_len = self.removed_segment_indexes.len();
-        let replaced_segments_len = new_segments_len.min(removed_segments_len);
-        let mut merged_statistics = Statistics::default();
-        let mut appended_segments = Vec::new();
-        let mut replaced_segments = HashMap::with_capacity(replaced_segments_len);
-        if new_segments_len > removed_segments_len {
-            // The remain new segments will be appended.
-            let appended = new_segments.split_off(removed_segments_len);
-            for (location, stats) in appended.into_iter().rev() {
-                let segment_loc = (location, SegmentInfo::VERSION);
-                new_segment_locs.push(segment_loc.clone());
-                appended_segments.push(segment_loc);
-                merge_statistics_mut(&mut merged_statistics, &stats, self.default_cluster_key_id);
-            }
-        }
-
-        for (i, (location, stats)) in new_segments.into_iter().enumerate() {
-            // The old segments will be replaced with the news.
-            let segment_loc = (location, SegmentInfo::VERSION);
-            new_segment_locs.push(segment_loc.clone());
-            replaced_segments.insert(self.removed_segment_indexes[i], segment_loc);
-            merge_statistics_mut(&mut merged_statistics, &stats, self.default_cluster_key_id);
+        for (location, stats) in new_segments {
+            merge_statistics_mut(
+                &mut self.appended_statistics,
+                &stats,
+                self.default_cluster_key_id,
+            );
+            self.appended_segments
+                .push((location, SegmentInfo::VERSION));
         }
-
-        let conflict_resolve_context =
-            ConflictResolveContext::ModifiedSegmentExistsInLatest(SnapshotChanges {
-                appended_segments,
-                removed_segment_indexes: self.removed_segment_indexes[replaced_segments_len..]
-                    .to_vec(),
-                replaced_segments,
-                removed_statistics: self.removed_statistics.clone(),
-                merged_statistics,
-            });
-        Ok(conflict_resolve_context)
+        Ok(())
     }
 
     async fn apply_mutation(

src/query/storages/fuse/src/operations/common/processors/transform_serialize_block.rs

@@ -24,6 +24,7 @@ use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::ComputedExpr;
 use databend_common_expression::DataBlock;
 use databend_common_expression::TableSchema;
+use databend_common_metrics::storage::metrics_inc_recluster_write_block_nums;
 use databend_common_pipeline_core::processors::Event;
 use databend_common_pipeline_core::processors::InputPort;
 use databend_common_pipeline_core::processors::OutputPort;
@@ -359,12 +360,15 @@ impl Processor for TransformSerializeBlock {
                         }
                     }
 
-                    if matches!(self.kind, MutationKind::Recluster) {
-                        Self::mutation_logs(MutationLogEntry::ReclusterAppendBlock {
+                    if matches!(self.kind, MutationKind::Insert) {
+                        DataBlock::empty_with_meta(Box::new(block_meta))
+                    } else {
+                        if matches!(self.kind, MutationKind::Recluster) {
+                            metrics_inc_recluster_write_block_nums();
+                        }
+                        Self::mutation_logs(MutationLogEntry::AppendBlock {
                             block_meta: Arc::new(block_meta),
                         })
-                    } else {
-                        DataBlock::empty_with_meta(Box::new(block_meta))
                     }
                 };
                 self.output_data = Some(mutation_log_data_block);