Skip to main content

risingwave_hummock_sdk/compaction_group/
hummock_version_ext.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::borrow::Borrow;
16use std::cmp::Ordering;
17use std::collections::hash_map::Entry;
18use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
19use std::iter::once;
20use std::sync::{Arc, LazyLock};
21
22use bytes::Bytes;
23use itertools::Itertools;
24use risingwave_common::catalog::TableId;
25use risingwave_common::hash::VnodeBitmapExt;
26use risingwave_common::log::LogSuppressor;
27use risingwave_pb::hummock::{
28    CompactionConfig, CompatibilityVersion, PbLevelType, StateTableInfo, StateTableInfoDelta,
29};
30use tracing::warn;
31
32use super::group_split::split_sst_with_table_ids;
33use super::{StateTableId, group_split};
34use crate::change_log::{EpochNewChangeLogCommon, TableChangeLogCommon};
35use crate::compact_task::is_compaction_task_expired;
36use crate::compaction_group::StaticCompactionGroupId;
37use crate::key_range::KeyRangeCommon;
38use crate::level::{Level, LevelCommon, Levels, OverlappingLevel};
39use crate::sstable_info::SstableInfo;
40use crate::table_watermark::{ReadTableWatermark, TableWatermarks};
41use crate::vector_index::apply_vector_index_delta;
42use crate::version::{
43    GroupDelta, GroupDeltaCommon, HummockVersion, HummockVersionCommon, HummockVersionDeltaCommon,
44    IntraLevelDelta, IntraLevelDeltaCommon, ObjectIdReader, SstableIdReader,
45};
46use crate::{
47    CompactionGroupId, HummockObjectId, HummockSstableId, HummockSstableObjectId, can_concat,
48};
49
50#[derive(Debug, Clone, Default)]
51pub struct SstDeltaInfo {
52    pub insert_sst_level: u32,
53    pub insert_sst_infos: Vec<SstableInfo>,
54    pub delete_sst_object_ids: Vec<HummockSstableObjectId>,
55}
56
57pub type BranchedSstInfo = HashMap<CompactionGroupId, Vec<HummockSstableId>>;
58
59impl HummockVersionCommon<SstableInfo> {
60    pub fn get_compaction_group_levels(&self, compaction_group_id: CompactionGroupId) -> &Levels {
61        self.levels
62            .get(&compaction_group_id)
63            .unwrap_or_else(|| panic!("compaction group {} does not exist", compaction_group_id))
64    }
65
66    pub fn get_compaction_group_levels_mut(
67        &mut self,
68        compaction_group_id: CompactionGroupId,
69    ) -> &mut Levels {
70        self.levels
71            .get_mut(&compaction_group_id)
72            .unwrap_or_else(|| panic!("compaction group {} does not exist", compaction_group_id))
73    }
74
75    // Only scan SST infos from levels in the specified compaction group.
76    pub fn get_sst_ids_by_group_id(
77        &self,
78        compaction_group_id: CompactionGroupId,
79    ) -> impl Iterator<Item = HummockSstableId> + '_ {
80        self.levels
81            .get(&compaction_group_id)
82            .into_iter()
83            .flat_map(|level| level.l0.sub_levels.iter().rev().chain(level.levels.iter()))
84            .flat_map(|level| level.table_infos.iter())
85            .map(|s| s.sst_id)
86    }
87
88    /// Prune stale table ids that no longer exist in `state_table_info` from SST metadata.
89    ///
90    /// This is used to normalize recovered versions from old clusters where dropped table ids
91    /// may still exist in persisted SST metadata.
92    pub fn prune_stale_table_ids_from_ssts(&mut self) -> usize {
93        let live_table_ids: HashSet<_> = self.state_table_info.info().keys().copied().collect();
94        // Older checkpoints may rely on deprecated `member_table_ids` before `state_table_info`
95        // is backfilled.
96        if live_table_ids.is_empty()
97            && self.levels.values().any(|levels| {
98                #[expect(deprecated)]
99                {
100                    !levels.member_table_ids.is_empty()
101                }
102            })
103        {
104            return 0;
105        }
106
107        let mut pruned_table_ids = HashSet::new();
108
109        for levels in self.levels.values_mut() {
110            let stale_table_ids = levels
111                .l0
112                .sub_levels
113                .iter()
114                .chain(levels.levels.iter())
115                .flat_map(|level| level.table_infos.iter())
116                .flat_map(|sst| sst.table_ids.iter().copied())
117                .filter(|table_id| !live_table_ids.contains(table_id))
118                .collect::<HashSet<_>>();
119
120            if stale_table_ids.is_empty() {
121                continue;
122            }
123
124            pruned_table_ids.extend(stale_table_ids.iter().copied());
125            levels.prune_table_ids_from_ssts(&stale_table_ids);
126        }
127
128        pruned_table_ids.len()
129    }
130
131    pub fn level_iter<F: FnMut(&Level) -> bool>(
132        &self,
133        compaction_group_id: CompactionGroupId,
134        mut f: F,
135    ) {
136        if let Some(levels) = self.levels.get(&compaction_group_id) {
137            for sub_level in &levels.l0.sub_levels {
138                if !f(sub_level) {
139                    return;
140                }
141            }
142            for level in &levels.levels {
143                if !f(level) {
144                    return;
145                }
146            }
147        }
148    }
149
150    pub fn num_levels(&self, compaction_group_id: CompactionGroupId) -> usize {
151        // l0 is currently separated from all levels
152        self.levels
153            .get(&compaction_group_id)
154            .map(|group| group.levels.len() + 1)
155            .unwrap_or(0)
156    }
157
158    pub fn safe_epoch_table_watermarks(
159        &self,
160        existing_table_ids: &[TableId],
161    ) -> BTreeMap<TableId, TableWatermarks> {
162        safe_epoch_table_watermarks_impl(&self.table_watermarks, existing_table_ids)
163    }
164}
165
166pub fn safe_epoch_table_watermarks_impl(
167    table_watermarks: &HashMap<TableId, Arc<TableWatermarks>>,
168    existing_table_ids: &[TableId],
169) -> BTreeMap<TableId, TableWatermarks> {
170    fn extract_single_table_watermark(
171        table_watermarks: &TableWatermarks,
172    ) -> Option<TableWatermarks> {
173        if let Some((first_epoch, first_epoch_watermark)) = table_watermarks.watermarks.first() {
174            Some(TableWatermarks {
175                watermarks: vec![(*first_epoch, first_epoch_watermark.clone())],
176                direction: table_watermarks.direction,
177                watermark_type: table_watermarks.watermark_type,
178            })
179        } else {
180            None
181        }
182    }
183    table_watermarks
184        .iter()
185        .filter_map(|(table_id, table_watermarks)| {
186            if !existing_table_ids.contains(table_id) {
187                None
188            } else {
189                extract_single_table_watermark(table_watermarks)
190                    .map(|table_watermarks| (*table_id, table_watermarks))
191            }
192        })
193        .collect()
194}
195
196pub fn safe_epoch_read_table_watermarks_impl(
197    safe_epoch_watermarks: BTreeMap<TableId, TableWatermarks>,
198) -> BTreeMap<TableId, ReadTableWatermark> {
199    safe_epoch_watermarks
200        .into_iter()
201        .map(|(table_id, watermarks)| {
202            assert_eq!(watermarks.watermarks.len(), 1);
203            let vnode_watermarks = &watermarks.watermarks.first().expect("should exist").1;
204            let mut vnode_watermark_map = BTreeMap::new();
205            for vnode_watermark in vnode_watermarks.iter() {
206                let watermark = Bytes::copy_from_slice(vnode_watermark.watermark());
207                for vnode in vnode_watermark.vnode_bitmap().iter_vnodes() {
208                    assert!(
209                        vnode_watermark_map
210                            .insert(vnode, watermark.clone())
211                            .is_none(),
212                        "duplicate table watermark on vnode {}",
213                        vnode.to_index()
214                    );
215                }
216            }
217            (
218                table_id,
219                ReadTableWatermark {
220                    direction: watermarks.direction,
221                    vnode_watermarks: vnode_watermark_map,
222                },
223            )
224        })
225        .collect()
226}
227
228impl HummockVersionCommon<SstableInfo> {
229    pub fn count_new_ssts_in_group_split(
230        &self,
231        parent_group_id: CompactionGroupId,
232        split_key: Bytes,
233    ) -> u64 {
234        self.levels
235            .get(&parent_group_id)
236            .map_or(0, |parent_levels| {
237                let l0 = &parent_levels.l0;
238                let mut split_count = 0;
239                for sub_level in &l0.sub_levels {
240                    assert!(!sub_level.table_infos.is_empty());
241
242                    if sub_level.level_type == PbLevelType::Overlapping {
243                        // TODO: use table_id / vnode / key_range filter
244                        split_count += sub_level
245                            .table_infos
246                            .iter()
247                            .map(|sst| {
248                                if let group_split::SstSplitType::Both =
249                                    group_split::need_to_split(sst, split_key.clone())
250                                {
251                                    2
252                                } else {
253                                    0
254                                }
255                            })
256                            .sum::<u64>();
257                        continue;
258                    }
259
260                    let pos = group_split::get_split_pos(&sub_level.table_infos, split_key.clone());
261                    let sst = sub_level.table_infos.get(pos).unwrap();
262
263                    if let group_split::SstSplitType::Both =
264                        group_split::need_to_split(sst, split_key.clone())
265                    {
266                        split_count += 2;
267                    }
268                }
269
270                for level in &parent_levels.levels {
271                    if level.table_infos.is_empty() {
272                        continue;
273                    }
274                    let pos = group_split::get_split_pos(&level.table_infos, split_key.clone());
275                    let sst = level.table_infos.get(pos).unwrap();
276                    if let group_split::SstSplitType::Both =
277                        group_split::need_to_split(sst, split_key.clone())
278                    {
279                        split_count += 2;
280                    }
281                }
282
283                split_count
284            })
285    }
286
287    pub fn init_with_parent_group(
288        &mut self,
289        parent_group_id: CompactionGroupId,
290        group_id: CompactionGroupId,
291        member_table_ids: BTreeSet<StateTableId>,
292        new_sst_start_id: HummockSstableId,
293    ) {
294        let mut new_sst_id = new_sst_start_id;
295        if parent_group_id == StaticCompactionGroupId::NewCompactionGroup {
296            if new_sst_start_id != 0 {
297                if cfg!(debug_assertions) {
298                    panic!(
299                        "non-zero sst start id {} for NewCompactionGroup",
300                        new_sst_start_id
301                    );
302                } else {
303                    warn!(
304                        %new_sst_start_id,
305                        "non-zero sst start id for NewCompactionGroup"
306                    );
307                }
308            }
309            return;
310        } else if !self.levels.contains_key(&parent_group_id) {
311            unreachable!(
312                "non-existing parent group id {} to init from",
313                parent_group_id
314            );
315        }
316        let [parent_levels, cur_levels] = self
317            .levels
318            .get_disjoint_mut([&parent_group_id, &group_id])
319            .map(|res| res.unwrap());
320        // After certain compaction group operation, e.g. split, any ongoing compaction tasks created prior to that should be rejected due to expiration.
321        // By incrementing the compaction_group_version_id of the compaction group, and comparing it with the one recorded in compaction task, expired compaction tasks can be identified.
322        parent_levels.compaction_group_version_id += 1;
323        cur_levels.compaction_group_version_id += 1;
324        let l0 = &mut parent_levels.l0;
325        {
326            for sub_level in &mut l0.sub_levels {
327                let target_l0 = &mut cur_levels.l0;
328                // Remove SST from sub level may result in empty sub level. It will be purged
329                // whenever another compaction task is finished.
330                let insert_table_infos =
331                    split_sst_info_for_level(&member_table_ids, sub_level, &mut new_sst_id);
332                sub_level.normalize();
333                if insert_table_infos.is_empty() {
334                    continue;
335                }
336                match group_split::get_sub_level_insert_hint(&target_l0.sub_levels, sub_level) {
337                    Ok(idx) => {
338                        add_ssts_to_sub_level(target_l0, idx, insert_table_infos);
339                    }
340                    Err(idx) => {
341                        insert_new_sub_level(
342                            target_l0,
343                            sub_level.sub_level_id,
344                            sub_level.level_type,
345                            insert_table_infos,
346                            Some(idx),
347                        );
348                    }
349                }
350            }
351            l0.normalize();
352        }
353        for (idx, level) in parent_levels.levels.iter_mut().enumerate() {
354            let insert_table_infos =
355                split_sst_info_for_level(&member_table_ids, level, &mut new_sst_id);
356            cur_levels.levels[idx].total_file_size += insert_table_infos
357                .iter()
358                .map(|sst| sst.sst_size)
359                .sum::<u64>();
360            cur_levels.levels[idx].uncompressed_file_size += insert_table_infos
361                .iter()
362                .map(|sst| sst.uncompressed_file_size)
363                .sum::<u64>();
364            cur_levels.levels[idx]
365                .table_infos
366                .extend(insert_table_infos);
367            cur_levels.levels[idx]
368                .table_infos
369                .sort_by(|sst1, sst2| sst1.key_range.cmp(&sst2.key_range));
370            assert!(can_concat(&cur_levels.levels[idx].table_infos));
371            level.normalize();
372        }
373
374        assert!(
375            parent_levels
376                .l0
377                .sub_levels
378                .iter()
379                .all(|level| !level.table_infos.is_empty())
380        );
381        assert!(
382            cur_levels
383                .l0
384                .sub_levels
385                .iter()
386                .all(|level| !level.table_infos.is_empty())
387        );
388    }
389
390    pub fn build_sst_delta_infos(
391        &self,
392        version_delta: &HummockVersionDeltaCommon<SstableInfo>,
393    ) -> Vec<SstDeltaInfo> {
394        let mut infos = vec![];
395
396        // Skip trivial move delta for refiller
397        // The trivial move task only changes the position of the sst in the lsm, it does not modify the object information corresponding to the sst, and does not need to re-execute the refill.
398        if version_delta.trivial_move {
399            return infos;
400        }
401
402        for (group_id, group_deltas) in &version_delta.group_deltas {
403            let mut info = SstDeltaInfo::default();
404
405            let mut removed_l0_ssts: BTreeSet<HummockSstableId> = BTreeSet::new();
406            let mut removed_ssts: BTreeMap<u32, BTreeSet<HummockSstableId>> = BTreeMap::new();
407
408            // Group rewrites cannot be represented by an SST refill summary.
409            if !group_deltas.group_deltas.iter().all(|delta| {
410                matches!(
411                    delta,
412                    GroupDelta::IntraLevel(_)
413                        | GroupDelta::NewL0SubLevel(_)
414                        | GroupDelta::PruneTableIdsFromSsts(_)
415                )
416            }) {
417                continue;
418            }
419
420            // TODO(MrCroxx): At most one insert delta is allowed here. It's okay for now with the
421            // current `hummock::manager::gen_version_delta` implementation. Better refactor the
422            // struct to reduce conventions.
423            for group_delta in &group_deltas.group_deltas {
424                match group_delta {
425                    GroupDeltaCommon::NewL0SubLevel(inserted_table_infos) => {
426                        if !inserted_table_infos.is_empty() {
427                            info.insert_sst_level = 0;
428                            info.insert_sst_infos
429                                .extend(inserted_table_infos.iter().cloned());
430                        }
431                    }
432                    GroupDeltaCommon::IntraLevel(intra_level) => {
433                        if !intra_level.inserted_table_infos.is_empty() {
434                            info.insert_sst_level = intra_level.level_idx;
435                            info.insert_sst_infos
436                                .extend(intra_level.inserted_table_infos.iter().cloned());
437                        }
438                        if !intra_level.removed_table_ids.is_empty() {
439                            for id in &intra_level.removed_table_ids {
440                                if intra_level.level_idx == 0 {
441                                    removed_l0_ssts.insert(*id);
442                                } else {
443                                    removed_ssts
444                                        .entry(intra_level.level_idx)
445                                        .or_default()
446                                        .insert(*id);
447                                }
448                            }
449                        }
450                    }
451                    GroupDeltaCommon::PruneTableIdsFromSsts(table_ids) => {
452                        // Follow delta order, as version application does: a commit can append
453                        // SSTs and then truncate a refresh staging table in the same group.
454                        info.insert_sst_infos.retain_mut(|sst| {
455                            if sst.table_ids.iter().any(|id| table_ids.contains(id)) {
456                                let mut inner = sst.get_inner();
457                                inner.table_ids.retain(|id| !table_ids.contains(id));
458                                sst.set_inner(inner);
459                            }
460                            !sst.table_ids.is_empty()
461                        });
462                    }
463                    GroupDeltaCommon::GroupConstruct(_)
464                    | GroupDeltaCommon::GroupDestroy(_)
465                    | GroupDeltaCommon::GroupMerge(_) => {}
466                }
467            }
468
469            let group = self.levels.get(group_id).unwrap();
470            if !removed_l0_ssts.is_empty() {
471                for l0_sub_level in &group.level0().sub_levels {
472                    for sst_info in &l0_sub_level.table_infos {
473                        if removed_l0_ssts.remove(&sst_info.sst_id) {
474                            info.delete_sst_object_ids.push(sst_info.object_id);
475                        }
476                    }
477                }
478            }
479            for level in &group.levels {
480                if let Some(mut removed_level_ssts) = removed_ssts.remove(&level.level_idx) {
481                    for sst_info in &level.table_infos {
482                        if removed_level_ssts.remove(&sst_info.sst_id) {
483                            info.delete_sst_object_ids.push(sst_info.object_id);
484                        }
485                    }
486                    if !removed_level_ssts.is_empty() {
487                        tracing::error!(
488                            "removed_level_ssts is not empty: {:?}",
489                            removed_level_ssts,
490                        );
491                    }
492                    debug_assert!(removed_level_ssts.is_empty());
493                }
494            }
495
496            if !removed_l0_ssts.is_empty() || !removed_ssts.is_empty() {
497                tracing::error!(
498                    "not empty removed_l0_ssts: {:?}, removed_ssts: {:?}",
499                    removed_l0_ssts,
500                    removed_ssts
501                );
502            }
503            debug_assert!(removed_l0_ssts.is_empty());
504            debug_assert!(removed_ssts.is_empty());
505
506            if !info.insert_sst_infos.is_empty() || !info.delete_sst_object_ids.is_empty() {
507                infos.push(info);
508            }
509        }
510
511        infos
512    }
513
514    pub fn apply_version_delta(
515        &mut self,
516        version_delta: &HummockVersionDeltaCommon<SstableInfo>,
517    ) -> HashMap<TableId, Option<StateTableInfo>> {
518        assert_eq!(self.id, version_delta.prev_id);
519
520        let (changed_table_info, mut is_commit_epoch) = self.state_table_info.apply_delta(
521            &version_delta.state_table_info_delta,
522            &version_delta.removed_table_ids,
523        );
524
525        #[expect(deprecated)]
526        {
527            if !is_commit_epoch && self.max_committed_epoch < version_delta.max_committed_epoch {
528                is_commit_epoch = true;
529                tracing::trace!(
530                    "max committed epoch bumped but no table committed epoch is changed"
531                );
532            }
533        }
534
535        // apply to `levels`, which is different compaction groups
536        for (compaction_group_id, group_deltas) in &version_delta.group_deltas {
537            let mut is_applied_l0_compact = false;
538            for group_delta in &group_deltas.group_deltas {
539                match group_delta {
540                    GroupDeltaCommon::GroupConstruct(group_construct) => {
541                        let mut new_levels = build_initial_compaction_group_levels(
542                            *compaction_group_id,
543                            group_construct.get_group_config().unwrap(),
544                        );
545                        let parent_group_id = group_construct.parent_group_id;
546                        new_levels.parent_group_id = parent_group_id;
547                        #[expect(deprecated)]
548                        // for backward-compatibility of previous hummock version delta
549                        new_levels
550                            .member_table_ids
551                            .clone_from(&group_construct.table_ids);
552                        self.levels.insert(*compaction_group_id, new_levels);
553                        if group_construct.version() >= CompatibilityVersion::SplitGroupByTableId {
554                            let split_key = group_construct.split_key.clone().map(Bytes::from);
555                            self.init_with_parent_group_v2(
556                                parent_group_id,
557                                *compaction_group_id,
558                                group_construct.new_sst_start_id,
559                                split_key,
560                            );
561                        } else {
562                            // for backward-compatibility of previous hummock version delta
563                            let member_table_ids = if group_construct.version()
564                                >= CompatibilityVersion::NoMemberTableIds
565                            {
566                                self.state_table_info
567                                    .compaction_group_member_table_ids(*compaction_group_id)
568                                    .iter()
569                                    .copied()
570                                    .collect()
571                            } else {
572                                #[expect(deprecated)]
573                                BTreeSet::from_iter(
574                                    group_construct.table_ids.iter().copied().map(Into::into),
575                                )
576                            };
577                            self.init_with_parent_group(
578                                parent_group_id,
579                                *compaction_group_id,
580                                member_table_ids,
581                                group_construct.new_sst_start_id,
582                            );
583                        }
584                    }
585                    GroupDeltaCommon::GroupMerge(group_merge) => {
586                        tracing::info!(
587                            "group_merge left {:?} right {:?}",
588                            group_merge.left_group_id,
589                            group_merge.right_group_id
590                        );
591                        self.merge_compaction_group(
592                            group_merge.left_group_id,
593                            group_merge.right_group_id,
594                        )
595                    }
596                    GroupDeltaCommon::IntraLevel(level_delta) => {
597                        let levels =
598                            self.levels.get_mut(compaction_group_id).unwrap_or_else(|| {
599                                panic!("compaction group {} does not exist", compaction_group_id)
600                            });
601                        if is_commit_epoch {
602                            assert!(
603                                level_delta.removed_table_ids.is_empty(),
604                                "no sst should be deleted when committing an epoch"
605                            );
606
607                            let IntraLevelDelta {
608                                level_idx,
609                                l0_sub_level_id,
610                                inserted_table_infos,
611                                ..
612                            } = level_delta;
613                            {
614                                assert_eq!(
615                                    *level_idx, 0,
616                                    "we should only add to L0 when we commit an epoch."
617                                );
618                                if !inserted_table_infos.is_empty() {
619                                    insert_new_sub_level(
620                                        &mut levels.l0,
621                                        *l0_sub_level_id,
622                                        PbLevelType::Overlapping,
623                                        inserted_table_infos.clone(),
624                                        None,
625                                    );
626                                }
627                            }
628                        } else {
629                            // The delta is caused by compaction.
630                            levels.apply_compact_ssts(
631                                level_delta,
632                                self.state_table_info
633                                    .compaction_group_member_table_ids(*compaction_group_id),
634                            );
635                            if level_delta.level_idx == 0 {
636                                is_applied_l0_compact = true;
637                            }
638                        }
639                    }
640                    GroupDeltaCommon::NewL0SubLevel(inserted_table_infos) => {
641                        let levels =
642                            self.levels.get_mut(compaction_group_id).unwrap_or_else(|| {
643                                panic!("compaction group {} does not exist", compaction_group_id)
644                            });
645                        assert!(is_commit_epoch);
646
647                        if !inserted_table_infos.is_empty() {
648                            let next_l0_sub_level_id = levels
649                                .l0
650                                .sub_levels
651                                .last()
652                                .map(|level| level.sub_level_id + 1)
653                                .unwrap_or(1);
654
655                            insert_new_sub_level(
656                                &mut levels.l0,
657                                next_l0_sub_level_id,
658                                PbLevelType::Overlapping,
659                                inserted_table_infos.clone(),
660                                None,
661                            );
662                        }
663                    }
664                    GroupDeltaCommon::GroupDestroy(_) => {
665                        self.levels.remove(compaction_group_id);
666                    }
667
668                    GroupDeltaCommon::PruneTableIdsFromSsts(table_ids) => {
669                        self.levels
670                            .get_mut(compaction_group_id)
671                            .unwrap_or_else(|| {
672                                panic!("compaction group {} does not exist", compaction_group_id)
673                            })
674                            .prune_table_ids_from_ssts(table_ids);
675                    }
676                }
677            }
678
679            if is_applied_l0_compact && let Some(levels) = self.levels.get_mut(compaction_group_id)
680            {
681                levels.l0.normalize();
682            }
683        }
684        self.id = version_delta.id;
685        #[expect(deprecated)]
686        {
687            self.max_committed_epoch = version_delta.max_committed_epoch;
688        }
689
690        // apply to table watermark
691
692        // Store the table watermarks that needs to be updated. None means to remove the table watermark of the table id
693        let mut modified_table_watermarks: HashMap<TableId, Option<TableWatermarks>> =
694            HashMap::new();
695
696        // apply to table watermark
697        for (table_id, table_watermarks) in &version_delta.new_table_watermarks {
698            if let Some(current_table_watermarks) = self.table_watermarks.get(table_id) {
699                if version_delta.removed_table_ids.contains(table_id) {
700                    modified_table_watermarks.insert(*table_id, None);
701                } else {
702                    let mut current_table_watermarks = (**current_table_watermarks).clone();
703                    current_table_watermarks.apply_new_table_watermarks(table_watermarks);
704                    modified_table_watermarks.insert(*table_id, Some(current_table_watermarks));
705                }
706            } else {
707                modified_table_watermarks.insert(*table_id, Some(table_watermarks.clone()));
708            }
709        }
710        for table_id in &version_delta.removed_table_ids {
711            modified_table_watermarks.insert(*table_id, None);
712        }
713        for (table_id, table_watermarks) in &self.table_watermarks {
714            let safe_epoch = if let Some(state_table_info) =
715                self.state_table_info.info().get(table_id)
716                && let Some((oldest_epoch, _)) = table_watermarks.watermarks.first()
717                && state_table_info.committed_epoch > *oldest_epoch
718            {
719                // safe epoch has progressed, need further clear.
720                state_table_info.committed_epoch
721            } else {
722                // safe epoch not progressed or the table has been removed. No need to truncate
723                continue;
724            };
725            let table_watermarks = modified_table_watermarks
726                .entry(*table_id)
727                .or_insert_with(|| Some((**table_watermarks).clone()));
728            if let Some(table_watermarks) = table_watermarks {
729                table_watermarks.clear_stale_epoch_watermark(safe_epoch);
730            }
731        }
732        // apply the staging table watermark to hummock version
733        for (table_id, table_watermarks) in modified_table_watermarks {
734            if let Some(table_watermarks) = table_watermarks {
735                self.table_watermarks
736                    .insert(table_id, Arc::new(table_watermarks));
737            } else {
738                self.table_watermarks.remove(&table_id);
739            }
740        }
741        // apply to vector index
742        apply_vector_index_delta(
743            &mut self.vector_indexes,
744            &version_delta.vector_index_delta,
745            &version_delta.removed_table_ids,
746        );
747
748        changed_table_info
749    }
750
751    pub fn apply_change_log_delta<T: Clone>(
752        table_change_log: &mut HashMap<TableId, TableChangeLogCommon<T>>,
753        change_log_delta: &HashMap<TableId, EpochNewChangeLogCommon<T>>,
754    ) {
755        for (table_id, new_change_log) in change_log_delta {
756            match table_change_log.entry(*table_id) {
757                Entry::Occupied(entry) => {
758                    let change_log = entry.into_mut();
759                    change_log.add_change_log(new_change_log.clone());
760                }
761                Entry::Vacant(entry) => {
762                    entry.insert(TableChangeLogCommon::new(once(new_change_log.clone())));
763                }
764            };
765        }
766    }
767
768    /// Returns the log deltas required to truncate the entire change log for a table.
769    pub fn collect_gc_change_log_delta<'a, T: Clone>(
770        current_change_log_table_ids: impl Iterator<Item = &'a TableId>,
771        change_log_delta: &HashMap<TableId, EpochNewChangeLogCommon<T>>,
772        removed_table_ids: &HashSet<TableId>,
773        state_table_info_delta: &HashMap<TableId, StateTableInfoDelta>,
774        changed_table_info: &HashMap<TableId, Option<StateTableInfo>>,
775    ) -> HashSet<TableId> {
776        let mut gc_change_log_delta = HashSet::new();
777        // If a table has no new change log entry (even an empty one), it means we have stopped maintained
778        // the change log for the table, and then we will remove the table change log.
779        // The table change log will also be removed when the table id is removed.
780        for table_id in current_change_log_table_ids {
781            if removed_table_ids.contains(table_id) {
782                gc_change_log_delta.insert(*table_id);
783                continue;
784            }
785            if let Some(table_info_delta) = state_table_info_delta.get(table_id)
786                && let Some(Some(prev_table_info)) = changed_table_info.get(table_id)
787                && table_info_delta.committed_epoch > prev_table_info.committed_epoch
788            {
789                // the table exists previously, and its committed epoch has progressed.
790            } else {
791                // otherwise, the table change log should be kept anyway
792                continue;
793            }
794            let contains = change_log_delta.contains_key(table_id);
795            if !contains {
796                gc_change_log_delta.insert(*table_id);
797                static LOG_SUPPRESSOR: LazyLock<LogSuppressor> =
798                    LazyLock::new(|| LogSuppressor::per_second(1));
799                if let Ok(suppressed_count) = LOG_SUPPRESSOR.check() {
800                    warn!(
801                        suppressed_count,
802                        %table_id,
803                        "table change log dropped due to no further change log at newly committed epoch"
804                    );
805                }
806            }
807        }
808        gc_change_log_delta
809    }
810
811    pub fn build_branched_sst_info(&self) -> BTreeMap<HummockSstableObjectId, BranchedSstInfo> {
812        let mut ret: BTreeMap<_, _> = BTreeMap::new();
813        for (compaction_group_id, group) in &self.levels {
814            for level in group.l0.sub_levels.iter().chain(&group.levels) {
815                for table_info in &level.table_infos {
816                    if table_info.sst_id.as_raw_id() == table_info.object_id.as_raw_id() {
817                        continue;
818                    }
819                    let object_id = table_info.object_id;
820                    let entry: &mut BranchedSstInfo = ret.entry(object_id).or_default();
821                    entry
822                        .entry(*compaction_group_id)
823                        .or_default()
824                        .push(table_info.sst_id)
825                }
826            }
827        }
828        ret
829    }
830
831    pub fn merge_compaction_group(
832        &mut self,
833        left_group_id: CompactionGroupId,
834        right_group_id: CompactionGroupId,
835    ) {
836        // Double check
837        let left_group_id_table_ids = self
838            .state_table_info
839            .compaction_group_member_table_ids(left_group_id)
840            .iter();
841        let right_group_id_table_ids = self
842            .state_table_info
843            .compaction_group_member_table_ids(right_group_id)
844            .iter();
845
846        assert!(
847            left_group_id_table_ids
848                .chain(right_group_id_table_ids)
849                .is_sorted()
850        );
851
852        let total_cg = self.levels.keys().cloned().collect::<Vec<_>>();
853        let right_levels = self.levels.remove(&right_group_id).unwrap_or_else(|| {
854            panic!(
855                "compaction group should exist right {} all {:?}",
856                right_group_id, total_cg
857            )
858        });
859
860        let left_levels = self.levels.get_mut(&left_group_id).unwrap_or_else(|| {
861            panic!(
862                "compaction group should exist left {} all {:?}",
863                left_group_id, total_cg
864            )
865        });
866
867        group_split::merge_levels(left_levels, right_levels);
868    }
869
870    pub fn init_with_parent_group_v2(
871        &mut self,
872        parent_group_id: CompactionGroupId,
873        group_id: CompactionGroupId,
874        new_sst_start_id: HummockSstableId,
875        split_key: Option<Bytes>,
876    ) {
877        let mut new_sst_id = new_sst_start_id;
878        if parent_group_id == StaticCompactionGroupId::NewCompactionGroup {
879            if new_sst_start_id != 0 {
880                if cfg!(debug_assertions) {
881                    panic!(
882                        "non-zero sst start id {} for NewCompactionGroup",
883                        new_sst_start_id
884                    );
885                } else {
886                    warn!(
887                        %new_sst_start_id,
888                        "non-zero sst start id for NewCompactionGroup"
889                    );
890                }
891            }
892            return;
893        } else if !self.levels.contains_key(&parent_group_id) {
894            unreachable!(
895                "non-existing parent group id {} to init from (V2)",
896                parent_group_id
897            );
898        }
899
900        let [parent_levels, cur_levels] = self
901            .levels
902            .get_disjoint_mut([&parent_group_id, &group_id])
903            .map(|res| res.unwrap());
904        // After certain compaction group operation, e.g. split, any ongoing compaction tasks created prior to that should be rejected due to expiration.
905        // By incrementing the compaction_group_version_id of the compaction group, and comparing it with the one recorded in compaction task, expired compaction tasks can be identified.
906        parent_levels.compaction_group_version_id += 1;
907        cur_levels.compaction_group_version_id += 1;
908
909        let l0 = &mut parent_levels.l0;
910        {
911            for sub_level in &mut l0.sub_levels {
912                let target_l0 = &mut cur_levels.l0;
913                // Remove SST from sub level may result in empty sub level. It will be purged
914                // whenever another compaction task is finished.
915                let insert_table_infos = if let Some(split_key) = &split_key {
916                    group_split::split_sst_info_for_level_v2(
917                        sub_level,
918                        &mut new_sst_id,
919                        split_key.clone(),
920                    )
921                } else {
922                    vec![]
923                };
924
925                if insert_table_infos.is_empty() {
926                    continue;
927                }
928
929                sub_level.normalize();
930                match group_split::get_sub_level_insert_hint(&target_l0.sub_levels, sub_level) {
931                    Ok(idx) => {
932                        add_ssts_to_sub_level(target_l0, idx, insert_table_infos);
933                    }
934                    Err(idx) => {
935                        insert_new_sub_level(
936                            target_l0,
937                            sub_level.sub_level_id,
938                            sub_level.level_type,
939                            insert_table_infos,
940                            Some(idx),
941                        );
942                    }
943                }
944            }
945            l0.normalize();
946        }
947
948        for (idx, level) in parent_levels.levels.iter_mut().enumerate() {
949            let insert_table_infos = if let Some(split_key) = &split_key {
950                group_split::split_sst_info_for_level_v2(level, &mut new_sst_id, split_key.clone())
951            } else {
952                vec![]
953            };
954
955            if insert_table_infos.is_empty() {
956                continue;
957            }
958
959            cur_levels.levels[idx].total_file_size += insert_table_infos
960                .iter()
961                .map(|sst| sst.sst_size)
962                .sum::<u64>();
963            cur_levels.levels[idx].uncompressed_file_size += insert_table_infos
964                .iter()
965                .map(|sst| sst.uncompressed_file_size)
966                .sum::<u64>();
967            cur_levels.levels[idx]
968                .table_infos
969                .extend(insert_table_infos);
970            cur_levels.levels[idx]
971                .table_infos
972                .sort_by(|sst1, sst2| sst1.key_range.cmp(&sst2.key_range));
973            assert!(can_concat(&cur_levels.levels[idx].table_infos));
974            level.normalize();
975        }
976
977        assert!(
978            parent_levels
979                .l0
980                .sub_levels
981                .iter()
982                .all(|level| !level.table_infos.is_empty())
983        );
984        assert!(
985            cur_levels
986                .l0
987                .sub_levels
988                .iter()
989                .all(|level| !level.table_infos.is_empty())
990        );
991    }
992}
993
994impl<T> HummockVersionCommon<T>
995where
996    T: SstableIdReader + ObjectIdReader,
997{
998    pub fn get_object_ids(&self) -> impl Iterator<Item = HummockObjectId> + '_ {
999        // DO NOT REMOVE THIS LINE
1000        // This is to ensure that when adding new variant to `HummockObjectId`,
1001        // the compiler will warn us if we forget to handle it here.
1002        match HummockObjectId::Sstable(0.into()) {
1003            HummockObjectId::Sstable(_) => {}
1004            HummockObjectId::VectorFile(_) => {}
1005            HummockObjectId::HnswGraphFile(_) => {}
1006        };
1007        self.get_sst_infos()
1008            .map(|s| HummockObjectId::Sstable(s.object_id()))
1009            .chain(
1010                self.vector_indexes
1011                    .values()
1012                    .flat_map(|index| index.get_objects().map(|(object_id, _)| object_id)),
1013            )
1014    }
1015
1016    pub fn get_sst_ids(&self) -> HashSet<HummockSstableId> {
1017        self.get_sst_infos().map(|s| s.sst_id()).collect()
1018    }
1019
1020    pub fn get_sst_infos(&self) -> impl Iterator<Item = &T> {
1021        self.get_combined_levels()
1022            .flat_map(|level| level.table_infos.iter())
1023    }
1024}
1025
1026impl Levels {
1027    pub(crate) fn apply_compact_ssts(
1028        &mut self,
1029        level_delta: &IntraLevelDeltaCommon<SstableInfo>,
1030        member_table_ids: &BTreeSet<TableId>,
1031    ) {
1032        let IntraLevelDeltaCommon {
1033            level_idx,
1034            l0_sub_level_id,
1035            inserted_table_infos: insert_table_infos,
1036            vnode_partition_count,
1037            removed_table_ids: delete_sst_ids_set,
1038            compaction_group_version_id,
1039        } = level_delta;
1040        let new_vnode_partition_count = *vnode_partition_count;
1041
1042        if is_compaction_task_expired(
1043            self.compaction_group_version_id,
1044            *compaction_group_version_id,
1045        ) {
1046            warn!(
1047                current_compaction_group_version_id = self.compaction_group_version_id,
1048                delta_compaction_group_version_id = compaction_group_version_id,
1049                level_idx,
1050                l0_sub_level_id,
1051                insert_table_infos = ?insert_table_infos
1052                    .iter()
1053                    .map(|sst| (sst.sst_id, sst.object_id))
1054                    .collect_vec(),
1055                ?delete_sst_ids_set,
1056                "This VersionDelta may be committed by an expired compact task. Please check it."
1057            );
1058            return;
1059        }
1060        if !delete_sst_ids_set.is_empty() {
1061            if *level_idx == 0 {
1062                for level in &mut self.l0.sub_levels {
1063                    level.delete_ssts(delete_sst_ids_set);
1064                }
1065            } else {
1066                let idx = *level_idx as usize - 1;
1067                self.levels[idx].delete_ssts(delete_sst_ids_set);
1068            }
1069        }
1070
1071        if !insert_table_infos.is_empty() {
1072            let insert_sst_level_id = *level_idx;
1073            let insert_sub_level_id = *l0_sub_level_id;
1074            if insert_sst_level_id == 0 {
1075                let l0 = &mut self.l0;
1076                let index = l0
1077                    .sub_levels
1078                    .partition_point(|level| level.sub_level_id < insert_sub_level_id);
1079                assert!(
1080                    index < l0.sub_levels.len()
1081                        && l0.sub_levels[index].sub_level_id == insert_sub_level_id,
1082                    "should find the level to insert into when applying compaction generated delta. sub level idx: {},  removed sst ids: {:?}, sub levels: {:?},",
1083                    insert_sub_level_id,
1084                    delete_sst_ids_set,
1085                    l0.sub_levels
1086                        .iter()
1087                        .map(|level| level.sub_level_id)
1088                        .collect_vec()
1089                );
1090                if l0.sub_levels[index].table_infos.is_empty()
1091                    && member_table_ids.len() == 1
1092                    && insert_table_infos.iter().all(|sst| {
1093                        sst.table_ids.len() == 1
1094                            && sst.table_ids[0]
1095                                == *member_table_ids.iter().next().expect("non-empty")
1096                    })
1097                {
1098                    // Only change vnode_partition_count for group which has only one state-table.
1099                    // Only change vnode_partition_count for level which update all sst files in this compact task.
1100                    l0.sub_levels[index].vnode_partition_count = new_vnode_partition_count;
1101                }
1102                level_insert_ssts(&mut l0.sub_levels[index], insert_table_infos);
1103            } else {
1104                let idx = insert_sst_level_id as usize - 1;
1105                if self.levels[idx].table_infos.is_empty()
1106                    && insert_table_infos
1107                        .iter()
1108                        .all(|sst| sst.table_ids.len() == 1)
1109                {
1110                    self.levels[idx].vnode_partition_count = new_vnode_partition_count;
1111                } else if self.levels[idx].vnode_partition_count != 0
1112                    && new_vnode_partition_count == 0
1113                    && member_table_ids.len() > 1
1114                {
1115                    self.levels[idx].vnode_partition_count = 0;
1116                }
1117                level_insert_ssts(&mut self.levels[idx], insert_table_infos);
1118            }
1119        }
1120    }
1121
1122    /// Prune specified table ids from all SST metadata, remove emptied SSTs and sub-levels,
1123    /// then bump `compaction_group_version_id`.
1124    pub(crate) fn prune_table_ids_from_ssts(&mut self, table_ids: &HashSet<TableId>) {
1125        for level in self.l0.sub_levels.iter_mut().chain(self.levels.iter_mut()) {
1126            level.prune_table_ids_from_ssts(table_ids);
1127        }
1128        self.l0.normalize();
1129        self.compaction_group_version_id += 1;
1130    }
1131}
1132
1133impl<T> HummockVersionCommon<T> {
1134    pub fn get_combined_levels(&self) -> impl Iterator<Item = &'_ LevelCommon<T>> + '_ {
1135        self.levels
1136            .values()
1137            .flat_map(|level| level.l0.sub_levels.iter().rev().chain(level.levels.iter()))
1138    }
1139}
1140
1141pub fn build_initial_compaction_group_levels(
1142    group_id: impl Into<CompactionGroupId>,
1143    compaction_config: &CompactionConfig,
1144) -> Levels {
1145    let mut levels = vec![];
1146    for l in 0..compaction_config.get_max_level() {
1147        levels.push(Level {
1148            level_idx: (l + 1) as u32,
1149            level_type: PbLevelType::Nonoverlapping,
1150            table_infos: vec![],
1151            total_file_size: 0,
1152            sub_level_id: 0,
1153            uncompressed_file_size: 0,
1154            vnode_partition_count: 0,
1155        });
1156    }
1157    #[expect(deprecated)] // for backward-compatibility of previous hummock version delta
1158    Levels {
1159        levels,
1160        l0: OverlappingLevel {
1161            sub_levels: vec![],
1162            total_file_size: 0,
1163            uncompressed_file_size: 0,
1164        },
1165        group_id: group_id.into(),
1166        parent_group_id: 0.into(),
1167        member_table_ids: vec![],
1168        compaction_group_version_id: 0,
1169    }
1170}
1171
1172fn split_sst_info_for_level(
1173    member_table_ids: &BTreeSet<TableId>,
1174    level: &mut Level,
1175    new_sst_id: &mut HummockSstableId,
1176) -> Vec<SstableInfo> {
1177    // Remove SST from sub level may result in empty sub level. It will be purged
1178    // whenever another compaction task is finished.
1179    let mut insert_table_infos = vec![];
1180    for sst_info in &mut level.table_infos {
1181        let contains_member_table = sst_info
1182            .table_ids
1183            .iter()
1184            .any(|table_id| member_table_ids.contains(table_id));
1185        let sst_size = sst_info.sst_size;
1186        if sst_size / 2 == 0 {
1187            tracing::warn!(
1188                id = %sst_info.sst_id,
1189                object_id = %sst_info.object_id,
1190                sst_size = sst_info.sst_size,
1191                file_size = sst_info.file_size,
1192                "Sstable sst_size is under expected",
1193            );
1194        };
1195        if contains_member_table {
1196            let (modified_sst, branch_sst) = split_sst_with_table_ids(
1197                sst_info,
1198                new_sst_id,
1199                sst_size / 2,
1200                sst_size / 2,
1201                member_table_ids.iter().cloned().collect_vec(),
1202            );
1203            *sst_info = modified_sst;
1204            insert_table_infos.push(branch_sst);
1205        }
1206    }
1207    insert_table_infos
1208}
1209
1210/// Gets all compaction group ids.
1211pub fn get_compaction_group_ids(
1212    version: &HummockVersion,
1213) -> impl Iterator<Item = CompactionGroupId> + '_ {
1214    version.levels.keys().cloned()
1215}
1216
1217pub fn get_table_compaction_group_id_mapping(
1218    version: &HummockVersion,
1219) -> HashMap<StateTableId, CompactionGroupId> {
1220    version
1221        .state_table_info
1222        .info()
1223        .iter()
1224        .map(|(table_id, info)| (*table_id, info.compaction_group_id))
1225        .collect()
1226}
1227
1228/// Gets all SSTs in `group_id`
1229pub fn get_compaction_group_ssts(
1230    version: &HummockVersion,
1231    group_id: CompactionGroupId,
1232) -> impl Iterator<Item = (HummockSstableObjectId, HummockSstableId)> + '_ {
1233    let group_levels = version.get_compaction_group_levels(group_id);
1234    group_levels
1235        .l0
1236        .sub_levels
1237        .iter()
1238        .rev()
1239        .chain(group_levels.levels.iter())
1240        .flat_map(|level| {
1241            level
1242                .table_infos
1243                .iter()
1244                .map(|table_info| (table_info.object_id, table_info.sst_id))
1245        })
1246}
1247
1248pub fn new_sub_level(
1249    sub_level_id: u64,
1250    level_type: PbLevelType,
1251    table_infos: Vec<SstableInfo>,
1252) -> Level {
1253    if level_type == PbLevelType::Nonoverlapping {
1254        debug_assert!(
1255            can_concat(&table_infos),
1256            "sst of non-overlapping level is not concat-able: {:?}",
1257            table_infos
1258        );
1259    }
1260    let total_file_size = table_infos.iter().map(|table| table.sst_size).sum();
1261    let uncompressed_file_size = table_infos
1262        .iter()
1263        .map(|table| table.uncompressed_file_size)
1264        .sum();
1265    Level {
1266        level_idx: 0,
1267        level_type,
1268        table_infos,
1269        total_file_size,
1270        sub_level_id,
1271        uncompressed_file_size,
1272        vnode_partition_count: 0,
1273    }
1274}
1275
1276pub fn add_ssts_to_sub_level(
1277    l0: &mut OverlappingLevel,
1278    sub_level_idx: usize,
1279    insert_table_infos: Vec<SstableInfo>,
1280) {
1281    insert_table_infos.iter().for_each(|sst| {
1282        l0.sub_levels[sub_level_idx].total_file_size += sst.sst_size;
1283        l0.sub_levels[sub_level_idx].uncompressed_file_size += sst.uncompressed_file_size;
1284        l0.total_file_size += sst.sst_size;
1285        l0.uncompressed_file_size += sst.uncompressed_file_size;
1286    });
1287    l0.sub_levels[sub_level_idx]
1288        .table_infos
1289        .extend(insert_table_infos);
1290    if l0.sub_levels[sub_level_idx].level_type == PbLevelType::Nonoverlapping {
1291        l0.sub_levels[sub_level_idx]
1292            .table_infos
1293            .sort_by(|sst1, sst2| sst1.key_range.cmp(&sst2.key_range));
1294        assert!(
1295            can_concat(&l0.sub_levels[sub_level_idx].table_infos),
1296            "sstable ids: {:?}",
1297            l0.sub_levels[sub_level_idx]
1298                .table_infos
1299                .iter()
1300                .map(|sst| sst.sst_id)
1301                .collect_vec()
1302        );
1303    }
1304}
1305
1306/// `None` value of `sub_level_insert_hint` means append.
1307pub fn insert_new_sub_level(
1308    l0: &mut OverlappingLevel,
1309    insert_sub_level_id: u64,
1310    level_type: PbLevelType,
1311    insert_table_infos: Vec<SstableInfo>,
1312    sub_level_insert_hint: Option<usize>,
1313) {
1314    if insert_sub_level_id == u64::MAX {
1315        return;
1316    }
1317    let insert_pos = if let Some(insert_pos) = sub_level_insert_hint {
1318        insert_pos
1319    } else {
1320        if let Some(newest_level) = l0.sub_levels.last() {
1321            assert!(
1322                newest_level.sub_level_id < insert_sub_level_id,
1323                "inserted new level is not the newest: prev newest: {}, insert: {}. L0: {:?}",
1324                newest_level.sub_level_id,
1325                insert_sub_level_id,
1326                l0,
1327            );
1328        }
1329        l0.sub_levels.len()
1330    };
1331    #[cfg(debug_assertions)]
1332    {
1333        if insert_pos > 0
1334            && let Some(smaller_level) = l0.sub_levels.get(insert_pos - 1)
1335        {
1336            debug_assert!(smaller_level.sub_level_id < insert_sub_level_id);
1337        }
1338        if let Some(larger_level) = l0.sub_levels.get(insert_pos) {
1339            debug_assert!(larger_level.sub_level_id > insert_sub_level_id);
1340        }
1341    }
1342    // All files will be committed in one new Overlapping sub-level and become
1343    // Nonoverlapping  after at least one compaction.
1344    let level = new_sub_level(insert_sub_level_id, level_type, insert_table_infos);
1345    l0.total_file_size += level.total_file_size;
1346    l0.uncompressed_file_size += level.uncompressed_file_size;
1347    l0.sub_levels.insert(insert_pos, level);
1348}
1349
1350impl Level {
1351    fn recompute_size(&mut self) {
1352        self.total_file_size = self
1353            .table_infos
1354            .iter()
1355            .map(|table| table.sst_size)
1356            .sum::<u64>();
1357        self.uncompressed_file_size = self
1358            .table_infos
1359            .iter()
1360            .map(|table| table.uncompressed_file_size)
1361            .sum::<u64>();
1362    }
1363
1364    /// Remove SSTs with empty `table_ids`, then recompute sizes.
1365    fn normalize(&mut self) {
1366        self.table_infos
1367            .retain(|sst_info| !sst_info.table_ids.is_empty());
1368        self.recompute_size();
1369    }
1370
1371    /// Return `true` if any SST was actually removed.
1372    fn delete_ssts(&mut self, ids: &HashSet<HummockSstableId>) -> bool {
1373        let original_len = self.table_infos.len();
1374        self.table_infos
1375            .retain(|table| !ids.contains(&table.sst_id));
1376        self.recompute_size();
1377        original_len != self.table_infos.len()
1378    }
1379
1380    /// Prune specified `table_ids` from each SST's metadata,
1381    /// remove SSTs that become empty, then recompute sizes.
1382    fn prune_table_ids_from_ssts(&mut self, table_ids: &HashSet<TableId>) {
1383        for sstable_info in &mut self.table_infos {
1384            if !sstable_info
1385                .table_ids
1386                .iter()
1387                .any(|table_id| table_ids.contains(table_id))
1388            {
1389                continue;
1390            }
1391
1392            let mut inner = sstable_info.get_inner();
1393            inner.table_ids.retain(|id| !table_ids.contains(id));
1394            sstable_info.set_inner(inner);
1395        }
1396        self.normalize();
1397    }
1398}
1399
1400impl OverlappingLevel {
1401    /// Remove empty sub-levels, then recompute aggregated sizes.
1402    fn normalize(&mut self) {
1403        self.sub_levels
1404            .retain(|level| !level.table_infos.is_empty());
1405        self.total_file_size = self
1406            .sub_levels
1407            .iter()
1408            .map(|level| level.total_file_size)
1409            .sum::<u64>();
1410        self.uncompressed_file_size = self
1411            .sub_levels
1412            .iter()
1413            .map(|level| level.uncompressed_file_size)
1414            .sum::<u64>();
1415    }
1416}
1417
1418fn level_insert_ssts(operand: &mut Level, insert_table_infos: &Vec<SstableInfo>) {
1419    fn display_sstable_infos(ssts: &[impl Borrow<SstableInfo>]) -> String {
1420        format!(
1421            "sstable ids: {:?}",
1422            ssts.iter().map(|s| s.borrow().sst_id).collect_vec()
1423        )
1424    }
1425    operand.total_file_size += insert_table_infos
1426        .iter()
1427        .map(|sst| sst.sst_size)
1428        .sum::<u64>();
1429    operand.uncompressed_file_size += insert_table_infos
1430        .iter()
1431        .map(|sst| sst.uncompressed_file_size)
1432        .sum::<u64>();
1433    if operand.level_type == PbLevelType::Overlapping {
1434        operand.level_type = PbLevelType::Nonoverlapping;
1435        operand
1436            .table_infos
1437            .extend(insert_table_infos.iter().cloned());
1438        operand
1439            .table_infos
1440            .sort_by(|sst1, sst2| sst1.key_range.cmp(&sst2.key_range));
1441        assert!(
1442            can_concat(&operand.table_infos),
1443            "{}",
1444            display_sstable_infos(&operand.table_infos)
1445        );
1446    } else if !insert_table_infos.is_empty() {
1447        let sorted_insert: Vec<_> = insert_table_infos
1448            .iter()
1449            .sorted_by(|sst1, sst2| sst1.key_range.cmp(&sst2.key_range))
1450            .cloned()
1451            .collect();
1452        let first = &sorted_insert[0];
1453        let last = &sorted_insert[sorted_insert.len() - 1];
1454        let pos = operand
1455            .table_infos
1456            .partition_point(|b| b.key_range.cmp(&first.key_range) == Ordering::Less);
1457        if pos >= operand.table_infos.len()
1458            || last.key_range.cmp(&operand.table_infos[pos].key_range) == Ordering::Less
1459        {
1460            operand.table_infos.splice(pos..pos, sorted_insert);
1461            // Validate the inserted SST batch along with the two SSTs that precede and follow it.
1462            let start = pos.saturating_sub(1);
1463            let end = (start + insert_table_infos.len() + 2).min(operand.table_infos.len());
1464            let validate_range = &operand.table_infos[start..end];
1465            assert!(
1466                can_concat(validate_range),
1467                "{}",
1468                display_sstable_infos(validate_range),
1469            );
1470        } else {
1471            // If this branch is reached, it indicates some unexpected behavior in compaction.
1472            // Here we issue a warning and fall back to insert one by one.
1473            warn!(insert = ?insert_table_infos, level = ?operand.table_infos, "unexpected overlap");
1474            for i in insert_table_infos {
1475                let pos = operand
1476                    .table_infos
1477                    .partition_point(|b| b.key_range.cmp(&i.key_range) == Ordering::Less);
1478                operand.table_infos.insert(pos, i.clone());
1479            }
1480            assert!(
1481                can_concat(&operand.table_infos),
1482                "{}",
1483                display_sstable_infos(&operand.table_infos)
1484            );
1485        }
1486    }
1487}
1488
1489pub fn version_object_size_map(version: &HummockVersion) -> HashMap<HummockObjectId, u64> {
1490    // DO NOT REMOVE THIS LINE
1491    // This is to ensure that when adding new variant to `HummockObjectId`,
1492    // the compiler will warn us if we forget to handle it here.
1493    match HummockObjectId::Sstable(0.into()) {
1494        HummockObjectId::Sstable(_) => {}
1495        HummockObjectId::VectorFile(_) => {}
1496        HummockObjectId::HnswGraphFile(_) => {}
1497    };
1498    version
1499        .levels
1500        .values()
1501        .flat_map(|cg| {
1502            cg.level0()
1503                .sub_levels
1504                .iter()
1505                .chain(cg.levels.iter())
1506                .flat_map(|level| level.table_infos.iter().map(|t| (t.object_id, t.file_size)))
1507        })
1508        .map(|(object_id, size)| (HummockObjectId::Sstable(object_id), size))
1509        .chain(
1510            version
1511                .vector_indexes
1512                .values()
1513                .flat_map(|index| index.get_objects()),
1514        )
1515        .collect()
1516}
1517
1518/// Verify the validity of a `HummockVersion` and return a list of violations if any.
1519/// Currently this method is only used by risectl validate-version.
1520pub fn validate_version(version: &HummockVersion) -> Vec<String> {
1521    let mut res = Vec::new();
1522    // Ensure each table maps to only one compaction group
1523    for (group_id, levels) in &version.levels {
1524        // Ensure compaction group id matches
1525        if levels.group_id != *group_id {
1526            res.push(format!(
1527                "GROUP {}: inconsistent group id {} in Levels",
1528                group_id, levels.group_id
1529            ));
1530        }
1531
1532        let validate_level = |group: CompactionGroupId,
1533                              expected_level_idx: u32,
1534                              level: &Level,
1535                              res: &mut Vec<String>| {
1536            let mut level_identifier = format!("GROUP {} LEVEL {}", group, level.level_idx);
1537            if level.level_idx == 0 {
1538                level_identifier.push_str(format!("SUBLEVEL {}", level.sub_level_id).as_str());
1539                // Ensure sub-level is not empty
1540                if level.table_infos.is_empty() {
1541                    res.push(format!("{}: empty level", level_identifier));
1542                }
1543            } else if level.level_type != PbLevelType::Nonoverlapping {
1544                // Ensure non-L0 level is non-overlapping level
1545                res.push(format!(
1546                    "{}: level type {:?} is not non-overlapping",
1547                    level_identifier, level.level_type
1548                ));
1549            }
1550
1551            // Ensure level idx matches
1552            if level.level_idx != expected_level_idx {
1553                res.push(format!(
1554                    "{}: mismatched level idx {}",
1555                    level_identifier, expected_level_idx
1556                ));
1557            }
1558
1559            let mut prev_table_info: Option<&SstableInfo> = None;
1560            for table_info in &level.table_infos {
1561                // Ensure table_ids are sorted and unique
1562                if !table_info.table_ids.is_sorted_by(|a, b| a < b) {
1563                    res.push(format!(
1564                        "{} SST {}: table_ids not sorted",
1565                        level_identifier, table_info.object_id
1566                    ));
1567                }
1568
1569                // Ensure SSTs in non-overlapping level have non-overlapping key range
1570                if level.level_type == PbLevelType::Nonoverlapping {
1571                    if let Some(prev) = prev_table_info.take()
1572                        && prev
1573                            .key_range
1574                            .compare_right_with(&table_info.key_range.left)
1575                            != Ordering::Less
1576                    {
1577                        res.push(format!(
1578                            "{} SST {}: key range should not overlap. prev={:?}, cur={:?}",
1579                            level_identifier, table_info.object_id, prev, table_info
1580                        ));
1581                    }
1582                    let _ = prev_table_info.insert(table_info);
1583                }
1584            }
1585        };
1586
1587        let l0 = &levels.l0;
1588        let mut prev_sub_level_id = None;
1589        for sub_level in &l0.sub_levels {
1590            // Ensure sub_level_id is strictly increasing.
1591            if let Some(prev_sub_level_id) = prev_sub_level_id
1592                && sub_level.sub_level_id <= prev_sub_level_id
1593            {
1594                res.push(format!(
1595                    "GROUP {} LEVEL 0: sub_level_id {} <= prev_sub_level {}",
1596                    group_id, sub_level.sub_level_id, prev_sub_level_id
1597                ));
1598            }
1599            prev_sub_level_id = Some(sub_level.sub_level_id);
1600
1601            validate_level(*group_id, 0, sub_level, &mut res);
1602        }
1603
1604        for idx in 1..=levels.levels.len() {
1605            validate_level(*group_id, idx as u32, levels.get_level(idx), &mut res);
1606        }
1607    }
1608    res
1609}
1610
1611#[cfg(test)]
1612mod tests {
1613    use std::collections::{HashMap, HashSet};
1614    use std::sync::Arc;
1615
1616    use bytes::Bytes;
1617    use risingwave_common::bitmap::Bitmap;
1618    use risingwave_common::catalog::TableId;
1619    use risingwave_common::hash::VirtualNode;
1620    use risingwave_common::util::epoch::test_epoch;
1621    use risingwave_pb::hummock::{
1622        CompactionConfig, GroupConstruct, GroupDestroy, LevelType, StateTableInfo,
1623    };
1624
1625    use super::group_split;
1626    use crate::HummockVersionId;
1627    use crate::compaction_group::group_split::*;
1628    use crate::compaction_group::hummock_version_ext::build_initial_compaction_group_levels;
1629    use crate::key::{FullKey, gen_key_from_str};
1630    use crate::key_range::KeyRange;
1631    use crate::level::{Level, Levels, OverlappingLevel};
1632    use crate::sstable_info::{SstableInfo, SstableInfoInner};
1633    use crate::table_watermark::{
1634        TableWatermarks, VnodeWatermark, WatermarkDirection, WatermarkSerdeType,
1635    };
1636    use crate::version::{
1637        GroupDelta, GroupDeltas, HummockVersion, HummockVersionDelta, HummockVersionStateTableInfo,
1638        IntraLevelDelta,
1639    };
1640
1641    fn gen_sstable_info(sst_id: u64, table_ids: Vec<u32>, epoch: u64) -> SstableInfo {
1642        gen_sstable_info_impl(sst_id, table_ids, epoch).into()
1643    }
1644
1645    fn gen_sstable_info_impl(sst_id: u64, table_ids: Vec<u32>, epoch: u64) -> SstableInfoInner {
1646        let table_key_l = gen_key_from_str(VirtualNode::ZERO, "1");
1647        let table_key_r = gen_key_from_str(VirtualNode::MAX_FOR_TEST, "1");
1648        let full_key_l = FullKey::for_test(
1649            TableId::new(*table_ids.first().unwrap()),
1650            table_key_l,
1651            epoch,
1652        )
1653        .encode();
1654        let full_key_r =
1655            FullKey::for_test(TableId::new(*table_ids.last().unwrap()), table_key_r, epoch)
1656                .encode();
1657
1658        SstableInfoInner {
1659            sst_id: sst_id.into(),
1660            key_range: KeyRange {
1661                left: full_key_l.into(),
1662                right: full_key_r.into(),
1663                right_exclusive: false,
1664            },
1665            table_ids: table_ids.into_iter().map(Into::into).collect(),
1666            object_id: sst_id.into(),
1667            min_epoch: 20,
1668            max_epoch: 20,
1669            file_size: 100,
1670            sst_size: 100,
1671            ..Default::default()
1672        }
1673    }
1674
1675    #[test]
1676    fn test_get_sst_object_ids() {
1677        let mut version = HummockVersion {
1678            id: HummockVersionId::new(0),
1679            levels: HashMap::from_iter([(
1680                0.into(),
1681                Levels {
1682                    levels: vec![],
1683                    l0: OverlappingLevel {
1684                        sub_levels: vec![],
1685                        total_file_size: 0,
1686                        uncompressed_file_size: 0,
1687                    },
1688                    ..Default::default()
1689                },
1690            )]),
1691            ..Default::default()
1692        };
1693        assert_eq!(version.get_object_ids().count(), 0);
1694
1695        // Add to sub level
1696        version
1697            .levels
1698            .get_mut(&0)
1699            .unwrap()
1700            .l0
1701            .sub_levels
1702            .push(Level {
1703                table_infos: vec![
1704                    SstableInfoInner {
1705                        object_id: 11.into(),
1706                        sst_id: 11.into(),
1707                        ..Default::default()
1708                    }
1709                    .into(),
1710                ],
1711                ..Default::default()
1712            });
1713        assert_eq!(version.get_object_ids().count(), 1);
1714
1715        // Add to non sub level
1716        version.levels.get_mut(&0).unwrap().levels.push(Level {
1717            table_infos: vec![
1718                SstableInfoInner {
1719                    object_id: 22.into(),
1720                    sst_id: 22.into(),
1721                    ..Default::default()
1722                }
1723                .into(),
1724            ],
1725            ..Default::default()
1726        });
1727        assert_eq!(version.get_object_ids().count(), 2);
1728    }
1729
1730    #[test]
1731    fn test_apply_version_delta() {
1732        let mut version = HummockVersion {
1733            id: HummockVersionId::new(0),
1734            levels: HashMap::from_iter([
1735                (
1736                    0.into(),
1737                    build_initial_compaction_group_levels(
1738                        0,
1739                        &CompactionConfig {
1740                            max_level: 6,
1741                            ..Default::default()
1742                        },
1743                    ),
1744                ),
1745                (
1746                    1.into(),
1747                    build_initial_compaction_group_levels(
1748                        1,
1749                        &CompactionConfig {
1750                            max_level: 6,
1751                            ..Default::default()
1752                        },
1753                    ),
1754                ),
1755            ]),
1756            ..Default::default()
1757        };
1758        let version_delta = HummockVersionDelta {
1759            id: HummockVersionId::new(1),
1760            group_deltas: HashMap::from_iter([
1761                (
1762                    2.into(),
1763                    GroupDeltas {
1764                        group_deltas: vec![GroupDelta::GroupConstruct(Box::new(GroupConstruct {
1765                            group_config: Some(CompactionConfig {
1766                                max_level: 6,
1767                                ..Default::default()
1768                            }),
1769                            ..Default::default()
1770                        }))],
1771                    },
1772                ),
1773                (
1774                    0.into(),
1775                    GroupDeltas {
1776                        group_deltas: vec![GroupDelta::GroupDestroy(GroupDestroy {})],
1777                    },
1778                ),
1779                (
1780                    1.into(),
1781                    GroupDeltas {
1782                        group_deltas: vec![GroupDelta::IntraLevel(IntraLevelDelta::new(
1783                            1,
1784                            0,
1785                            HashSet::new(),
1786                            vec![
1787                                SstableInfoInner {
1788                                    object_id: 1.into(),
1789                                    sst_id: 1.into(),
1790                                    ..Default::default()
1791                                }
1792                                .into(),
1793                            ],
1794                            0,
1795                            version
1796                                .levels
1797                                .get(&1)
1798                                .as_ref()
1799                                .unwrap()
1800                                .compaction_group_version_id,
1801                        ))],
1802                    },
1803                ),
1804            ]),
1805            ..Default::default()
1806        };
1807        let version_delta = version_delta;
1808
1809        version.apply_version_delta(&version_delta);
1810        let mut cg1 = build_initial_compaction_group_levels(
1811            1,
1812            &CompactionConfig {
1813                max_level: 6,
1814                ..Default::default()
1815            },
1816        );
1817        cg1.levels[0] = Level {
1818            level_idx: 1,
1819            level_type: LevelType::Nonoverlapping,
1820            table_infos: vec![
1821                SstableInfoInner {
1822                    object_id: 1.into(),
1823                    sst_id: 1.into(),
1824                    ..Default::default()
1825                }
1826                .into(),
1827            ],
1828            ..Default::default()
1829        };
1830        assert_eq!(
1831            version,
1832            HummockVersion {
1833                id: HummockVersionId::new(1),
1834                levels: HashMap::from_iter([
1835                    (
1836                        2.into(),
1837                        build_initial_compaction_group_levels(
1838                            2,
1839                            &CompactionConfig {
1840                                max_level: 6,
1841                                ..Default::default()
1842                            },
1843                        ),
1844                    ),
1845                    (1.into(), cg1),
1846                ]),
1847                ..Default::default()
1848            }
1849        );
1850    }
1851
1852    fn gen_sst_info(object_id: u64, table_ids: Vec<u32>, left: Bytes, right: Bytes) -> SstableInfo {
1853        gen_sst_info_impl(object_id, table_ids, left, right).into()
1854    }
1855
1856    fn gen_sst_info_impl(
1857        object_id: u64,
1858        table_ids: Vec<u32>,
1859        left: Bytes,
1860        right: Bytes,
1861    ) -> SstableInfoInner {
1862        SstableInfoInner {
1863            object_id: object_id.into(),
1864            sst_id: object_id.into(),
1865            key_range: KeyRange {
1866                left,
1867                right,
1868                right_exclusive: false,
1869            },
1870            table_ids: table_ids.into_iter().map(Into::into).collect(),
1871            file_size: 100,
1872            sst_size: 100,
1873            uncompressed_file_size: 100,
1874            ..Default::default()
1875        }
1876    }
1877
1878    #[test]
1879    fn test_merge_levels() {
1880        let mut left_levels = build_initial_compaction_group_levels(
1881            1,
1882            &CompactionConfig {
1883                max_level: 6,
1884                ..Default::default()
1885            },
1886        );
1887
1888        let mut right_levels = build_initial_compaction_group_levels(
1889            2,
1890            &CompactionConfig {
1891                max_level: 6,
1892                ..Default::default()
1893            },
1894        );
1895
1896        left_levels.levels[0] = Level {
1897            level_idx: 1,
1898            level_type: LevelType::Nonoverlapping,
1899            table_infos: vec![
1900                gen_sst_info(
1901                    1,
1902                    vec![3],
1903                    FullKey::for_test(
1904                        TableId::new(3),
1905                        gen_key_from_str(VirtualNode::from_index(1), "1"),
1906                        0,
1907                    )
1908                    .encode()
1909                    .into(),
1910                    FullKey::for_test(
1911                        TableId::new(3),
1912                        gen_key_from_str(VirtualNode::from_index(200), "1"),
1913                        0,
1914                    )
1915                    .encode()
1916                    .into(),
1917                ),
1918                gen_sst_info(
1919                    10,
1920                    vec![3, 4],
1921                    FullKey::for_test(
1922                        TableId::new(3),
1923                        gen_key_from_str(VirtualNode::from_index(201), "1"),
1924                        0,
1925                    )
1926                    .encode()
1927                    .into(),
1928                    FullKey::for_test(
1929                        TableId::new(4),
1930                        gen_key_from_str(VirtualNode::from_index(10), "1"),
1931                        0,
1932                    )
1933                    .encode()
1934                    .into(),
1935                ),
1936                gen_sst_info(
1937                    11,
1938                    vec![4],
1939                    FullKey::for_test(
1940                        TableId::new(4),
1941                        gen_key_from_str(VirtualNode::from_index(11), "1"),
1942                        0,
1943                    )
1944                    .encode()
1945                    .into(),
1946                    FullKey::for_test(
1947                        TableId::new(4),
1948                        gen_key_from_str(VirtualNode::from_index(200), "1"),
1949                        0,
1950                    )
1951                    .encode()
1952                    .into(),
1953                ),
1954            ],
1955            total_file_size: 300,
1956            ..Default::default()
1957        };
1958
1959        left_levels.l0.sub_levels.push(Level {
1960            level_idx: 0,
1961            table_infos: vec![gen_sst_info(
1962                3,
1963                vec![3],
1964                FullKey::for_test(
1965                    TableId::new(3),
1966                    gen_key_from_str(VirtualNode::from_index(1), "1"),
1967                    0,
1968                )
1969                .encode()
1970                .into(),
1971                FullKey::for_test(
1972                    TableId::new(3),
1973                    gen_key_from_str(VirtualNode::from_index(200), "1"),
1974                    0,
1975                )
1976                .encode()
1977                .into(),
1978            )],
1979            sub_level_id: 101,
1980            level_type: LevelType::Overlapping,
1981            total_file_size: 100,
1982            ..Default::default()
1983        });
1984
1985        left_levels.l0.sub_levels.push(Level {
1986            level_idx: 0,
1987            table_infos: vec![gen_sst_info(
1988                3,
1989                vec![3],
1990                FullKey::for_test(
1991                    TableId::new(3),
1992                    gen_key_from_str(VirtualNode::from_index(1), "1"),
1993                    0,
1994                )
1995                .encode()
1996                .into(),
1997                FullKey::for_test(
1998                    TableId::new(3),
1999                    gen_key_from_str(VirtualNode::from_index(200), "1"),
2000                    0,
2001                )
2002                .encode()
2003                .into(),
2004            )],
2005            sub_level_id: 103,
2006            level_type: LevelType::Overlapping,
2007            total_file_size: 100,
2008            ..Default::default()
2009        });
2010
2011        left_levels.l0.sub_levels.push(Level {
2012            level_idx: 0,
2013            table_infos: vec![gen_sst_info(
2014                3,
2015                vec![3],
2016                FullKey::for_test(
2017                    TableId::new(3),
2018                    gen_key_from_str(VirtualNode::from_index(1), "1"),
2019                    0,
2020                )
2021                .encode()
2022                .into(),
2023                FullKey::for_test(
2024                    TableId::new(3),
2025                    gen_key_from_str(VirtualNode::from_index(200), "1"),
2026                    0,
2027                )
2028                .encode()
2029                .into(),
2030            )],
2031            sub_level_id: 105,
2032            level_type: LevelType::Nonoverlapping,
2033            total_file_size: 100,
2034            ..Default::default()
2035        });
2036
2037        right_levels.levels[0] = Level {
2038            level_idx: 1,
2039            level_type: LevelType::Nonoverlapping,
2040            table_infos: vec![
2041                gen_sst_info(
2042                    1,
2043                    vec![5],
2044                    FullKey::for_test(
2045                        TableId::new(5),
2046                        gen_key_from_str(VirtualNode::from_index(1), "1"),
2047                        0,
2048                    )
2049                    .encode()
2050                    .into(),
2051                    FullKey::for_test(
2052                        TableId::new(5),
2053                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2054                        0,
2055                    )
2056                    .encode()
2057                    .into(),
2058                ),
2059                gen_sst_info(
2060                    10,
2061                    vec![5, 6],
2062                    FullKey::for_test(
2063                        TableId::new(5),
2064                        gen_key_from_str(VirtualNode::from_index(201), "1"),
2065                        0,
2066                    )
2067                    .encode()
2068                    .into(),
2069                    FullKey::for_test(
2070                        TableId::new(6),
2071                        gen_key_from_str(VirtualNode::from_index(10), "1"),
2072                        0,
2073                    )
2074                    .encode()
2075                    .into(),
2076                ),
2077                gen_sst_info(
2078                    11,
2079                    vec![6],
2080                    FullKey::for_test(
2081                        TableId::new(6),
2082                        gen_key_from_str(VirtualNode::from_index(11), "1"),
2083                        0,
2084                    )
2085                    .encode()
2086                    .into(),
2087                    FullKey::for_test(
2088                        TableId::new(6),
2089                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2090                        0,
2091                    )
2092                    .encode()
2093                    .into(),
2094                ),
2095            ],
2096            total_file_size: 300,
2097            ..Default::default()
2098        };
2099
2100        right_levels.l0.sub_levels.push(Level {
2101            level_idx: 0,
2102            table_infos: vec![gen_sst_info(
2103                3,
2104                vec![5],
2105                FullKey::for_test(
2106                    TableId::new(5),
2107                    gen_key_from_str(VirtualNode::from_index(1), "1"),
2108                    0,
2109                )
2110                .encode()
2111                .into(),
2112                FullKey::for_test(
2113                    TableId::new(5),
2114                    gen_key_from_str(VirtualNode::from_index(200), "1"),
2115                    0,
2116                )
2117                .encode()
2118                .into(),
2119            )],
2120            sub_level_id: 101,
2121            level_type: LevelType::Overlapping,
2122            total_file_size: 100,
2123            ..Default::default()
2124        });
2125
2126        right_levels.l0.sub_levels.push(Level {
2127            level_idx: 0,
2128            table_infos: vec![gen_sst_info(
2129                5,
2130                vec![5],
2131                FullKey::for_test(
2132                    TableId::new(5),
2133                    gen_key_from_str(VirtualNode::from_index(1), "1"),
2134                    0,
2135                )
2136                .encode()
2137                .into(),
2138                FullKey::for_test(
2139                    TableId::new(5),
2140                    gen_key_from_str(VirtualNode::from_index(200), "1"),
2141                    0,
2142                )
2143                .encode()
2144                .into(),
2145            )],
2146            sub_level_id: 102,
2147            level_type: LevelType::Overlapping,
2148            total_file_size: 100,
2149            ..Default::default()
2150        });
2151
2152        right_levels.l0.sub_levels.push(Level {
2153            level_idx: 0,
2154            table_infos: vec![gen_sst_info(
2155                3,
2156                vec![5],
2157                FullKey::for_test(
2158                    TableId::new(5),
2159                    gen_key_from_str(VirtualNode::from_index(1), "1"),
2160                    0,
2161                )
2162                .encode()
2163                .into(),
2164                FullKey::for_test(
2165                    TableId::new(5),
2166                    gen_key_from_str(VirtualNode::from_index(200), "1"),
2167                    0,
2168                )
2169                .encode()
2170                .into(),
2171            )],
2172            sub_level_id: 103,
2173            level_type: LevelType::Nonoverlapping,
2174            total_file_size: 100,
2175            ..Default::default()
2176        });
2177
2178        {
2179            // test empty
2180            let mut left_levels = Levels::default();
2181            let right_levels = Levels::default();
2182
2183            group_split::merge_levels(&mut left_levels, right_levels);
2184        }
2185
2186        {
2187            // test empty left
2188            let mut left_levels = build_initial_compaction_group_levels(
2189                1,
2190                &CompactionConfig {
2191                    max_level: 6,
2192                    ..Default::default()
2193                },
2194            );
2195            let right_levels = right_levels.clone();
2196
2197            group_split::merge_levels(&mut left_levels, right_levels);
2198
2199            assert!(left_levels.l0.sub_levels.len() == 3);
2200            assert!(left_levels.l0.sub_levels[0].sub_level_id == 101);
2201            assert_eq!(100, left_levels.l0.sub_levels[0].total_file_size);
2202            assert!(left_levels.l0.sub_levels[1].sub_level_id == 102);
2203            assert_eq!(100, left_levels.l0.sub_levels[1].total_file_size);
2204            assert!(left_levels.l0.sub_levels[2].sub_level_id == 103);
2205            assert_eq!(100, left_levels.l0.sub_levels[2].total_file_size);
2206
2207            assert!(left_levels.levels[0].level_idx == 1);
2208            assert_eq!(300, left_levels.levels[0].total_file_size);
2209        }
2210
2211        {
2212            // test empty right
2213            let mut left_levels = left_levels.clone();
2214            let right_levels = build_initial_compaction_group_levels(
2215                2,
2216                &CompactionConfig {
2217                    max_level: 6,
2218                    ..Default::default()
2219                },
2220            );
2221
2222            group_split::merge_levels(&mut left_levels, right_levels);
2223
2224            assert!(left_levels.l0.sub_levels.len() == 3);
2225            assert!(left_levels.l0.sub_levels[0].sub_level_id == 101);
2226            assert_eq!(100, left_levels.l0.sub_levels[0].total_file_size);
2227            assert!(left_levels.l0.sub_levels[1].sub_level_id == 103);
2228            assert_eq!(100, left_levels.l0.sub_levels[1].total_file_size);
2229            assert!(left_levels.l0.sub_levels[2].sub_level_id == 105);
2230            assert_eq!(100, left_levels.l0.sub_levels[2].total_file_size);
2231
2232            assert!(left_levels.levels[0].level_idx == 1);
2233            assert_eq!(300, left_levels.levels[0].total_file_size);
2234        }
2235
2236        {
2237            let mut left_levels = left_levels.clone();
2238            let right_levels = right_levels.clone();
2239
2240            group_split::merge_levels(&mut left_levels, right_levels);
2241
2242            assert!(left_levels.l0.sub_levels.len() == 6);
2243            assert!(left_levels.l0.sub_levels[0].sub_level_id == 101);
2244            assert_eq!(100, left_levels.l0.sub_levels[0].total_file_size);
2245            assert!(left_levels.l0.sub_levels[1].sub_level_id == 103);
2246            assert_eq!(100, left_levels.l0.sub_levels[1].total_file_size);
2247            assert!(left_levels.l0.sub_levels[2].sub_level_id == 105);
2248            assert_eq!(100, left_levels.l0.sub_levels[2].total_file_size);
2249            assert!(left_levels.l0.sub_levels[3].sub_level_id == 106);
2250            assert_eq!(100, left_levels.l0.sub_levels[3].total_file_size);
2251            assert!(left_levels.l0.sub_levels[4].sub_level_id == 107);
2252            assert_eq!(100, left_levels.l0.sub_levels[4].total_file_size);
2253            assert!(left_levels.l0.sub_levels[5].sub_level_id == 108);
2254            assert_eq!(100, left_levels.l0.sub_levels[5].total_file_size);
2255
2256            assert!(left_levels.levels[0].level_idx == 1);
2257            assert_eq!(600, left_levels.levels[0].total_file_size);
2258        }
2259    }
2260
2261    #[test]
2262    fn test_get_split_pos() {
2263        let epoch = test_epoch(1);
2264        let s1 = gen_sstable_info(1, vec![1, 2], epoch);
2265        let s2 = gen_sstable_info(2, vec![3, 4, 5], epoch);
2266        let s3 = gen_sstable_info(3, vec![6, 7], epoch);
2267
2268        let ssts = vec![s1, s2, s3];
2269        let split_key = group_split::build_split_key(4.into(), VirtualNode::ZERO);
2270
2271        let pos = group_split::get_split_pos(&ssts, split_key.clone());
2272        assert_eq!(1, pos);
2273
2274        let pos = group_split::get_split_pos(&vec![], split_key);
2275        assert_eq!(0, pos);
2276    }
2277
2278    #[test]
2279    fn test_split_sst() {
2280        let epoch = test_epoch(1);
2281        let sst = gen_sstable_info(1, vec![1, 2, 3, 5], epoch);
2282
2283        {
2284            let split_key = group_split::build_split_key(3.into(), VirtualNode::ZERO);
2285            let origin_sst = sst.clone();
2286            let sst_size = origin_sst.sst_size;
2287            let split_type = group_split::need_to_split(&origin_sst, split_key.clone());
2288            assert_eq!(SstSplitType::Both, split_type);
2289
2290            let mut new_sst_id = 10.into();
2291            let (origin_sst, branched_sst) = group_split::split_sst(
2292                origin_sst,
2293                &mut new_sst_id,
2294                split_key,
2295                sst_size / 2,
2296                sst_size / 2,
2297            );
2298
2299            let origin_sst = origin_sst.unwrap();
2300            let branched_sst = branched_sst.unwrap();
2301
2302            assert!(origin_sst.key_range.right_exclusive);
2303            assert!(
2304                origin_sst
2305                    .key_range
2306                    .right
2307                    .cmp(&branched_sst.key_range.left)
2308                    .is_le()
2309            );
2310            assert!(origin_sst.table_ids.is_sorted());
2311            assert!(branched_sst.table_ids.is_sorted());
2312            assert!(origin_sst.table_ids.last().unwrap() < branched_sst.table_ids.first().unwrap());
2313            assert!(branched_sst.sst_size < origin_sst.file_size);
2314            assert_eq!(10, branched_sst.sst_id);
2315            assert_eq!(11, origin_sst.sst_id);
2316            assert_eq!(3, branched_sst.table_ids.first().unwrap().as_raw_id()); // split table_id to right
2317        }
2318
2319        {
2320            // test un-exist table_id
2321            let split_key = group_split::build_split_key(4.into(), VirtualNode::ZERO);
2322            let origin_sst = sst.clone();
2323            let sst_size = origin_sst.sst_size;
2324            let split_type = group_split::need_to_split(&origin_sst, split_key.clone());
2325            assert_eq!(SstSplitType::Both, split_type);
2326
2327            let mut new_sst_id = 10.into();
2328            let (origin_sst, branched_sst) = group_split::split_sst(
2329                origin_sst,
2330                &mut new_sst_id,
2331                split_key,
2332                sst_size / 2,
2333                sst_size / 2,
2334            );
2335
2336            let origin_sst = origin_sst.unwrap();
2337            let branched_sst = branched_sst.unwrap();
2338
2339            assert!(origin_sst.key_range.right_exclusive);
2340            assert!(origin_sst.key_range.right.le(&branched_sst.key_range.left));
2341            assert!(origin_sst.table_ids.is_sorted());
2342            assert!(branched_sst.table_ids.is_sorted());
2343            assert!(origin_sst.table_ids.last().unwrap() < branched_sst.table_ids.first().unwrap());
2344            assert!(branched_sst.sst_size < origin_sst.file_size);
2345            assert_eq!(10, branched_sst.sst_id);
2346            assert_eq!(11, origin_sst.sst_id);
2347            assert_eq!(5, branched_sst.table_ids.first().unwrap().as_raw_id()); // split table_id to right
2348        }
2349
2350        {
2351            let split_key = group_split::build_split_key(6.into(), VirtualNode::ZERO);
2352            let split_type = group_split::need_to_split(&sst, split_key);
2353            assert_eq!(SstSplitType::Left, split_type);
2354        }
2355
2356        {
2357            let split_key = group_split::build_split_key(4.into(), VirtualNode::ZERO);
2358            let origin_sst = sst.clone();
2359            let split_type = group_split::need_to_split(&origin_sst, split_key);
2360            assert_eq!(SstSplitType::Both, split_type);
2361
2362            let split_key = group_split::build_split_key(1.into(), VirtualNode::ZERO);
2363            let origin_sst = sst;
2364            let split_type = group_split::need_to_split(&origin_sst, split_key);
2365            assert_eq!(SstSplitType::Right, split_type);
2366        }
2367
2368        {
2369            // test key_range left = right
2370            let mut sst = gen_sstable_info_impl(1, vec![1], epoch);
2371            sst.key_range.right = sst.key_range.left.clone();
2372            let sst: SstableInfo = sst.into();
2373            let split_key = group_split::build_split_key(1.into(), VirtualNode::ZERO);
2374            let origin_sst = sst;
2375            let sst_size = origin_sst.sst_size;
2376
2377            let mut new_sst_id = 10.into();
2378            let (origin_sst, branched_sst) = group_split::split_sst(
2379                origin_sst,
2380                &mut new_sst_id,
2381                split_key,
2382                sst_size / 2,
2383                sst_size / 2,
2384            );
2385
2386            assert!(origin_sst.is_none());
2387            assert!(branched_sst.is_some());
2388        }
2389    }
2390
2391    #[test]
2392    fn test_split_sst_info_for_level() {
2393        let mut version = HummockVersion {
2394            id: HummockVersionId::new(0),
2395            levels: HashMap::from_iter([(
2396                1.into(),
2397                build_initial_compaction_group_levels(
2398                    1,
2399                    &CompactionConfig {
2400                        max_level: 6,
2401                        ..Default::default()
2402                    },
2403                ),
2404            )]),
2405            ..Default::default()
2406        };
2407
2408        let cg1 = version.levels.get_mut(&1).unwrap();
2409
2410        cg1.levels[0] = Level {
2411            level_idx: 1,
2412            level_type: LevelType::Nonoverlapping,
2413            table_infos: vec![
2414                gen_sst_info(
2415                    1,
2416                    vec![3],
2417                    FullKey::for_test(
2418                        TableId::new(3),
2419                        gen_key_from_str(VirtualNode::from_index(1), "1"),
2420                        0,
2421                    )
2422                    .encode()
2423                    .into(),
2424                    FullKey::for_test(
2425                        TableId::new(3),
2426                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2427                        0,
2428                    )
2429                    .encode()
2430                    .into(),
2431                ),
2432                gen_sst_info(
2433                    10,
2434                    vec![3, 4],
2435                    FullKey::for_test(
2436                        TableId::new(3),
2437                        gen_key_from_str(VirtualNode::from_index(201), "1"),
2438                        0,
2439                    )
2440                    .encode()
2441                    .into(),
2442                    FullKey::for_test(
2443                        TableId::new(4),
2444                        gen_key_from_str(VirtualNode::from_index(10), "1"),
2445                        0,
2446                    )
2447                    .encode()
2448                    .into(),
2449                ),
2450                gen_sst_info(
2451                    11,
2452                    vec![4],
2453                    FullKey::for_test(
2454                        TableId::new(4),
2455                        gen_key_from_str(VirtualNode::from_index(11), "1"),
2456                        0,
2457                    )
2458                    .encode()
2459                    .into(),
2460                    FullKey::for_test(
2461                        TableId::new(4),
2462                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2463                        0,
2464                    )
2465                    .encode()
2466                    .into(),
2467                ),
2468            ],
2469            total_file_size: 300,
2470            ..Default::default()
2471        };
2472
2473        cg1.l0.sub_levels.push(Level {
2474            level_idx: 0,
2475            table_infos: vec![
2476                gen_sst_info(
2477                    2,
2478                    vec![2],
2479                    FullKey::for_test(
2480                        TableId::new(0),
2481                        gen_key_from_str(VirtualNode::from_index(1), "1"),
2482                        0,
2483                    )
2484                    .encode()
2485                    .into(),
2486                    FullKey::for_test(
2487                        TableId::new(2),
2488                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2489                        0,
2490                    )
2491                    .encode()
2492                    .into(),
2493                ),
2494                gen_sst_info(
2495                    22,
2496                    vec![2],
2497                    FullKey::for_test(
2498                        TableId::new(0),
2499                        gen_key_from_str(VirtualNode::from_index(1), "1"),
2500                        0,
2501                    )
2502                    .encode()
2503                    .into(),
2504                    FullKey::for_test(
2505                        TableId::new(2),
2506                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2507                        0,
2508                    )
2509                    .encode()
2510                    .into(),
2511                ),
2512                gen_sst_info(
2513                    23,
2514                    vec![2],
2515                    FullKey::for_test(
2516                        TableId::new(0),
2517                        gen_key_from_str(VirtualNode::from_index(1), "1"),
2518                        0,
2519                    )
2520                    .encode()
2521                    .into(),
2522                    FullKey::for_test(
2523                        TableId::new(2),
2524                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2525                        0,
2526                    )
2527                    .encode()
2528                    .into(),
2529                ),
2530                gen_sst_info(
2531                    24,
2532                    vec![2],
2533                    FullKey::for_test(
2534                        TableId::new(2),
2535                        gen_key_from_str(VirtualNode::from_index(1), "1"),
2536                        0,
2537                    )
2538                    .encode()
2539                    .into(),
2540                    FullKey::for_test(
2541                        TableId::new(2),
2542                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2543                        0,
2544                    )
2545                    .encode()
2546                    .into(),
2547                ),
2548                gen_sst_info(
2549                    25,
2550                    vec![2],
2551                    FullKey::for_test(
2552                        TableId::new(0),
2553                        gen_key_from_str(VirtualNode::from_index(1), "1"),
2554                        0,
2555                    )
2556                    .encode()
2557                    .into(),
2558                    FullKey::for_test(
2559                        TableId::new(0),
2560                        gen_key_from_str(VirtualNode::from_index(200), "1"),
2561                        0,
2562                    )
2563                    .encode()
2564                    .into(),
2565                ),
2566            ],
2567            sub_level_id: 101,
2568            level_type: LevelType::Overlapping,
2569            total_file_size: 300,
2570            ..Default::default()
2571        });
2572
2573        {
2574            // split Overlapping level
2575            let split_key = group_split::build_split_key(1.into(), VirtualNode::ZERO);
2576
2577            let mut new_sst_id = 100.into();
2578            let x = group_split::split_sst_info_for_level_v2(
2579                &mut cg1.l0.sub_levels[0],
2580                &mut new_sst_id,
2581                split_key,
2582            );
2583            // assert_eq!(3, x.len());
2584            // assert_eq!(100, x[0].sst_id);
2585            // assert_eq!(100, x[0].sst_size);
2586            // assert_eq!(101, x[1].sst_id);
2587            // assert_eq!(100, x[1].sst_size);
2588            // assert_eq!(102, x[2].sst_id);
2589            // assert_eq!(100, x[2].sst_size);
2590
2591            let mut right_l0 = OverlappingLevel {
2592                sub_levels: vec![],
2593                total_file_size: 0,
2594                uncompressed_file_size: 0,
2595            };
2596
2597            right_l0.sub_levels.push(Level {
2598                level_idx: 0,
2599                table_infos: x,
2600                sub_level_id: 101,
2601                total_file_size: 100,
2602                level_type: LevelType::Overlapping,
2603                ..Default::default()
2604            });
2605
2606            let right_levels = Levels {
2607                levels: vec![],
2608                l0: right_l0,
2609                ..Default::default()
2610            };
2611
2612            merge_levels(cg1, right_levels);
2613        }
2614
2615        {
2616            // test split empty level
2617            let mut new_sst_id = 100.into();
2618            let split_key = group_split::build_split_key(1.into(), VirtualNode::ZERO);
2619            let x = group_split::split_sst_info_for_level_v2(
2620                &mut cg1.levels[2],
2621                &mut new_sst_id,
2622                split_key,
2623            );
2624
2625            assert!(x.is_empty());
2626        }
2627
2628        {
2629            // test split to right Nonoverlapping level
2630            let mut cg1 = cg1.clone();
2631            let split_key = group_split::build_split_key(1.into(), VirtualNode::ZERO);
2632
2633            let mut new_sst_id = 100.into();
2634            let x = group_split::split_sst_info_for_level_v2(
2635                &mut cg1.levels[0],
2636                &mut new_sst_id,
2637                split_key,
2638            );
2639
2640            assert_eq!(3, x.len());
2641            assert_eq!(1, x[0].sst_id);
2642            assert_eq!(100, x[0].sst_size);
2643            assert_eq!(10, x[1].sst_id);
2644            assert_eq!(100, x[1].sst_size);
2645            assert_eq!(11, x[2].sst_id);
2646            assert_eq!(100, x[2].sst_size);
2647
2648            assert_eq!(0, cg1.levels[0].table_infos.len());
2649        }
2650
2651        {
2652            // test split to left Nonoverlapping level
2653            let mut cg1 = cg1.clone();
2654            let split_key = group_split::build_split_key(5.into(), VirtualNode::ZERO);
2655
2656            let mut new_sst_id = 100.into();
2657            let x = group_split::split_sst_info_for_level_v2(
2658                &mut cg1.levels[0],
2659                &mut new_sst_id,
2660                split_key,
2661            );
2662
2663            assert_eq!(0, x.len());
2664            assert_eq!(3, cg1.levels[0].table_infos.len());
2665        }
2666
2667        // {
2668        //     // test split to both Nonoverlapping level
2669        //     let mut cg1 = cg1.clone();
2670        //     let split_key = build_split_key(3, VirtualNode::MAX);
2671
2672        //     let mut new_sst_id = 100;
2673        //     let x = group_split::split_sst_info_for_level_v2(
2674        //         &mut cg1.levels[0],
2675        //         &mut new_sst_id,
2676        //         split_key,
2677        //     );
2678
2679        //     assert_eq!(2, x.len());
2680        //     assert_eq!(100, x[0].sst_id);
2681        //     assert_eq!(100 / 2, x[0].sst_size);
2682        //     assert_eq!(11, x[1].sst_id);
2683        //     assert_eq!(100, x[1].sst_size);
2684        //     assert_eq!(vec![3, 4], x[0].table_ids);
2685
2686        //     assert_eq!(2, cg1.levels[0].table_infos.len());
2687        //     assert_eq!(101, cg1.levels[0].table_infos[1].sst_id);
2688        //     assert_eq!(100 / 2, cg1.levels[0].table_infos[1].sst_size);
2689        //     assert_eq!(vec![3], cg1.levels[0].table_infos[1].table_ids);
2690        // }
2691
2692        {
2693            // test split to both Nonoverlapping level
2694            let mut cg1 = cg1.clone();
2695            let split_key = group_split::build_split_key(4.into(), VirtualNode::ZERO);
2696
2697            let mut new_sst_id = 100.into();
2698            let x = group_split::split_sst_info_for_level_v2(
2699                &mut cg1.levels[0],
2700                &mut new_sst_id,
2701                split_key,
2702            );
2703
2704            assert_eq!(2, x.len());
2705            assert_eq!(100, x[0].sst_id);
2706            assert_eq!(100 / 2, x[0].sst_size);
2707            assert_eq!(11, x[1].sst_id);
2708            assert_eq!(100, x[1].sst_size);
2709            assert_eq!(vec![TableId::new(4)], x[1].table_ids);
2710
2711            assert_eq!(2, cg1.levels[0].table_infos.len());
2712            assert_eq!(101, cg1.levels[0].table_infos[1].sst_id);
2713            assert_eq!(100 / 2, cg1.levels[0].table_infos[1].sst_size);
2714            assert_eq!(
2715                vec![TableId::new(3)],
2716                cg1.levels[0].table_infos[1].table_ids
2717            );
2718        }
2719    }
2720
2721    #[test]
2722    fn test_validate_version_l0_sublevel_order() {
2723        let group_id = 1.into();
2724        let mut version = HummockVersion {
2725            levels: HashMap::from([(
2726                group_id,
2727                Levels {
2728                    group_id,
2729                    ..Default::default()
2730                },
2731            )]),
2732            ..Default::default()
2733        };
2734        for sub_level_id in [1, 2] {
2735            super::insert_new_sub_level(
2736                &mut version.levels.get_mut(&group_id).unwrap().l0,
2737                sub_level_id,
2738                LevelType::Overlapping,
2739                vec![gen_sstable_info(sub_level_id, vec![1], test_epoch(1))],
2740                None,
2741            );
2742        }
2743        let violations = super::validate_version(&version);
2744        assert!(violations.is_empty(), "{violations:?}");
2745
2746        version
2747            .levels
2748            .get_mut(&group_id)
2749            .unwrap()
2750            .l0
2751            .sub_levels
2752            .swap(0, 1);
2753        assert_eq!(
2754            super::validate_version(&version),
2755            vec!["GROUP 1 LEVEL 0: sub_level_id 1 <= prev_sub_level 2"],
2756        );
2757
2758        version.levels.get_mut(&group_id).unwrap().l0.sub_levels[0].sub_level_id = 1;
2759        assert_eq!(
2760            super::validate_version(&version),
2761            vec!["GROUP 1 LEVEL 0: sub_level_id 1 <= prev_sub_level 1"],
2762        );
2763    }
2764
2765    fn make_sst(sst_id: u64, table_ids: Vec<u32>, sst_size: u64) -> SstableInfo {
2766        SstableInfoInner {
2767            sst_id: sst_id.into(),
2768            object_id: sst_id.into(),
2769            table_ids: table_ids.into_iter().map(TableId::new).collect(),
2770            file_size: sst_size,
2771            sst_size,
2772            uncompressed_file_size: sst_size * 2,
2773            ..Default::default()
2774        }
2775        .into()
2776    }
2777
2778    #[test]
2779    fn test_level_normalize() {
2780        // Mixed: some SSTs empty, some not.
2781        let mut level = Level {
2782            level_idx: 1,
2783            level_type: LevelType::Nonoverlapping,
2784            table_infos: vec![
2785                make_sst(1, vec![1, 2], 100),
2786                make_sst(2, vec![], 200), // empty → removed
2787                make_sst(3, vec![3], 300),
2788                make_sst(4, vec![], 400), // empty → removed
2789            ],
2790            total_file_size: 9999, // intentionally wrong
2791            uncompressed_file_size: 9999,
2792            ..Default::default()
2793        };
2794
2795        level.normalize();
2796
2797        assert_eq!(level.table_infos.len(), 2);
2798        assert_eq!(1, level.table_infos[0].sst_id);
2799        assert_eq!(3, level.table_infos[1].sst_id);
2800        assert_eq!(level.total_file_size, 400);
2801        assert_eq!(level.uncompressed_file_size, 800);
2802
2803        // No empty SSTs: just recompute sizes.
2804        level.total_file_size = 0;
2805        level.uncompressed_file_size = 0;
2806        level.normalize();
2807        assert_eq!(level.table_infos.len(), 2);
2808        assert_eq!(level.total_file_size, 400);
2809        assert_eq!(level.uncompressed_file_size, 800);
2810
2811        // All empty → level becomes empty.
2812        level.table_infos = vec![make_sst(10, vec![], 100), make_sst(11, vec![], 200)];
2813        level.normalize();
2814        assert!(level.table_infos.is_empty());
2815        assert_eq!(level.total_file_size, 0);
2816        assert_eq!(level.uncompressed_file_size, 0);
2817    }
2818
2819    #[test]
2820    fn test_level_delete_ssts() {
2821        let mut level = Level {
2822            level_idx: 1,
2823            level_type: LevelType::Nonoverlapping,
2824            table_infos: vec![
2825                make_sst(1, vec![1], 100),
2826                make_sst(2, vec![2], 200),
2827                make_sst(3, vec![3], 300),
2828            ],
2829            total_file_size: 600,
2830            uncompressed_file_size: 1200,
2831            ..Default::default()
2832        };
2833
2834        let delete_ids: HashSet<crate::HummockSstableId> = HashSet::from([2.into()]);
2835        let changed = level.delete_ssts(&delete_ids);
2836
2837        assert!(changed);
2838        assert_eq!(level.table_infos.len(), 2);
2839        assert_eq!(1, level.table_infos[0].sst_id);
2840        assert_eq!(3, level.table_infos[1].sst_id);
2841        assert_eq!(level.total_file_size, 400);
2842        assert_eq!(level.uncompressed_file_size, 800);
2843
2844        // Delete non-existent id → no change, returns false.
2845        let delete_ids: HashSet<crate::HummockSstableId> = HashSet::from([999.into()]);
2846        let changed = level.delete_ssts(&delete_ids);
2847        assert!(!changed);
2848        assert_eq!(level.table_infos.len(), 2);
2849    }
2850
2851    #[test]
2852    fn test_level_prune_table_ids_from_ssts() {
2853        let mut level = Level {
2854            level_idx: 1,
2855            level_type: LevelType::Nonoverlapping,
2856            table_infos: vec![
2857                make_sst(1, vec![1, 2], 100),
2858                make_sst(2, vec![2, 3], 200),
2859                make_sst(3, vec![2], 300),
2860            ],
2861            total_file_size: 600,
2862            uncompressed_file_size: 1200,
2863            ..Default::default()
2864        };
2865
2866        let pruned_table_ids = HashSet::from([TableId::new(2)]);
2867        level.prune_table_ids_from_ssts(&pruned_table_ids);
2868
2869        // SST 3 should be removed (was only table_id=2)
2870        assert_eq!(level.table_infos.len(), 2);
2871        assert_eq!(level.table_infos[0].table_ids, vec![TableId::new(1)]);
2872        assert_eq!(level.table_infos[1].table_ids, vec![TableId::new(3)]);
2873        assert_eq!(level.total_file_size, 100 + 200);
2874        assert_eq!(level.uncompressed_file_size, 200 + 400);
2875
2876        // Prune remaining tables, so all SSTs are removed.
2877        let pruned_table_ids = HashSet::from([TableId::new(1), TableId::new(3)]);
2878        level.prune_table_ids_from_ssts(&pruned_table_ids);
2879        assert!(level.table_infos.is_empty());
2880        assert_eq!(level.total_file_size, 0);
2881        assert_eq!(level.uncompressed_file_size, 0);
2882    }
2883
2884    #[test]
2885    fn test_overlapping_level_normalize() {
2886        let mut l0 = OverlappingLevel {
2887            sub_levels: vec![
2888                Level {
2889                    level_idx: 0,
2890                    table_infos: vec![make_sst(1, vec![1], 100)],
2891                    total_file_size: 100,
2892                    uncompressed_file_size: 200,
2893                    sub_level_id: 1,
2894                    ..Default::default()
2895                },
2896                Level {
2897                    level_idx: 0,
2898                    table_infos: vec![], // empty → should be removed
2899                    total_file_size: 0,
2900                    uncompressed_file_size: 0,
2901                    sub_level_id: 2,
2902                    ..Default::default()
2903                },
2904                Level {
2905                    level_idx: 0,
2906                    table_infos: vec![make_sst(3, vec![3], 300)],
2907                    total_file_size: 300,
2908                    uncompressed_file_size: 600,
2909                    sub_level_id: 3,
2910                    ..Default::default()
2911                },
2912            ],
2913            total_file_size: 9999, // intentionally wrong
2914            uncompressed_file_size: 9999,
2915        };
2916
2917        l0.normalize();
2918
2919        assert_eq!(l0.sub_levels.len(), 2);
2920        assert_eq!(l0.sub_levels[0].sub_level_id, 1);
2921        assert_eq!(l0.sub_levels[1].sub_level_id, 3);
2922        assert_eq!(l0.total_file_size, 100 + 300);
2923        assert_eq!(l0.uncompressed_file_size, 200 + 600);
2924
2925        // All sub-levels empty → normalize clears everything.
2926        l0.sub_levels = vec![Level {
2927            level_idx: 0,
2928            table_infos: vec![],
2929            ..Default::default()
2930        }];
2931        l0.normalize();
2932        assert!(l0.sub_levels.is_empty());
2933        assert_eq!(l0.total_file_size, 0);
2934        assert_eq!(l0.uncompressed_file_size, 0);
2935    }
2936
2937    #[test]
2938    fn test_levels_prune_table_ids_from_ssts() {
2939        #[expect(deprecated)]
2940        let mut levels = Levels {
2941            l0: OverlappingLevel {
2942                sub_levels: vec![
2943                    Level {
2944                        level_idx: 0,
2945                        table_infos: vec![
2946                            make_sst(1, vec![10], 100), // table 10
2947                            make_sst(2, vec![20], 200), // table 20
2948                        ],
2949                        total_file_size: 300,
2950                        uncompressed_file_size: 600,
2951                        sub_level_id: 1,
2952                        ..Default::default()
2953                    },
2954                    Level {
2955                        level_idx: 0,
2956                        table_infos: vec![
2957                            make_sst(3, vec![10], 150), // table 10
2958                        ],
2959                        total_file_size: 150,
2960                        uncompressed_file_size: 300,
2961                        sub_level_id: 2,
2962                        ..Default::default()
2963                    },
2964                ],
2965                total_file_size: 450,
2966                uncompressed_file_size: 900,
2967            },
2968            levels: vec![Level {
2969                level_idx: 1,
2970                level_type: LevelType::Nonoverlapping,
2971                table_infos: vec![
2972                    make_sst(4, vec![10, 20], 400), // shared SST
2973                    make_sst(5, vec![10], 500),     // table 10 only
2974                ],
2975                total_file_size: 900,
2976                uncompressed_file_size: 1800,
2977                ..Default::default()
2978            }],
2979            group_id: 1.into(),
2980            parent_group_id: 0.into(),
2981            member_table_ids: vec![],
2982            compaction_group_version_id: 0,
2983        };
2984
2985        // Prune table 10 from SST metadata.
2986        levels.prune_table_ids_from_ssts(&HashSet::from([TableId::new(10)]));
2987
2988        assert_eq!(levels.l0.sub_levels.len(), 1);
2989        assert_eq!(levels.l0.sub_levels[0].sub_level_id, 1);
2990        assert_eq!(levels.l0.sub_levels[0].table_infos.len(), 1);
2991        assert_eq!(2, levels.l0.sub_levels[0].table_infos[0].sst_id);
2992        assert_eq!(levels.l0.sub_levels[0].total_file_size, 200);
2993        assert_eq!(levels.l0.sub_levels[0].uncompressed_file_size, 400);
2994
2995        assert_eq!(levels.l0.total_file_size, 200);
2996        assert_eq!(levels.l0.uncompressed_file_size, 400);
2997
2998        assert_eq!(levels.levels[0].table_infos.len(), 1);
2999        assert_eq!(4, levels.levels[0].table_infos[0].sst_id);
3000        assert_eq!(
3001            levels.levels[0].table_infos[0].table_ids,
3002            vec![TableId::new(20)]
3003        );
3004        assert_eq!(levels.levels[0].total_file_size, 400);
3005        assert_eq!(levels.levels[0].uncompressed_file_size, 800);
3006
3007        assert_eq!(levels.compaction_group_version_id, 1);
3008    }
3009
3010    #[test]
3011    fn test_refill_summary_keeps_live_inserts_when_committing_a_table_truncation() {
3012        let group = 1.into();
3013        let mut version = HummockVersion::default();
3014        version.levels.insert(
3015            group,
3016            Levels {
3017                group_id: group,
3018                ..Default::default()
3019            },
3020        );
3021        version.state_table_info = HummockVersionStateTableInfo::from_protobuf_owned(
3022            [100, 200]
3023                .map(TableId::new)
3024                .map(|table_id| {
3025                    (
3026                        table_id,
3027                        StateTableInfo {
3028                            compaction_group_id: group,
3029                            ..Default::default()
3030                        },
3031                    )
3032                })
3033                .into(),
3034        );
3035        let mut delta = version.version_delta_after();
3036        for table_id in [100, 200].map(TableId::new) {
3037            delta.state_table_info_delta.insert(
3038                table_id,
3039                risingwave_pb::hummock::StateTableInfoDelta {
3040                    committed_epoch: test_epoch(1),
3041                    compaction_group_id: group,
3042                },
3043            );
3044        }
3045        // Meta appends committed L0 SSTs before pruning the refresh staging table.
3046        delta.group_deltas.insert(
3047            group,
3048            GroupDeltas {
3049                group_deltas: vec![
3050                    GroupDelta::NewL0SubLevel(vec![
3051                        make_sst(1, vec![200], 10),
3052                        make_sst(2, vec![100], 20),
3053                        make_sst(3, vec![100, 200], 30),
3054                    ]),
3055                    GroupDelta::PruneTableIdsFromSsts([TableId::new(100)].into()),
3056                ],
3057            },
3058        );
3059        let infos = version.build_sst_delta_infos(&delta);
3060        version.apply_version_delta(&delta);
3061        let live = &version.levels[&group].l0.sub_levels[0].table_infos;
3062        assert_eq!(
3063            live.iter()
3064                .map(|sst| sst.sst_id.as_raw_id())
3065                .collect::<Vec<_>>(),
3066            vec![1, 3]
3067        );
3068        assert!(
3069            live.iter()
3070                .all(|sst| sst.table_ids == vec![TableId::new(200)])
3071        );
3072        assert_eq!(infos.len(), 1);
3073        assert_eq!(infos[0].insert_sst_level, 0);
3074        assert_eq!(&infos[0].insert_sst_infos, live);
3075        assert!(infos[0].delete_sst_object_ids.is_empty());
3076    }
3077
3078    #[test]
3079    fn test_apply_version_delta_prune_table_ids_from_ssts() {
3080        let mut version = HummockVersion {
3081            id: HummockVersionId::new(0),
3082            levels: HashMap::from_iter([(1.into(), {
3083                #[expect(deprecated)]
3084                let levels = Levels {
3085                    l0: OverlappingLevel {
3086                        sub_levels: vec![
3087                            Level {
3088                                level_idx: 0,
3089                                level_type: LevelType::Overlapping,
3090                                table_infos: vec![
3091                                    make_sst(1, vec![100], 50), // only table 100
3092                                    make_sst(2, vec![200], 60), // only table 200
3093                                ],
3094                                total_file_size: 110,
3095                                uncompressed_file_size: 220,
3096                                sub_level_id: 1,
3097                                ..Default::default()
3098                            },
3099                            Level {
3100                                level_idx: 0,
3101                                level_type: LevelType::Overlapping,
3102                                table_infos: vec![
3103                                    make_sst(3, vec![100], 70), // only table 100
3104                                ],
3105                                total_file_size: 70,
3106                                uncompressed_file_size: 140,
3107                                sub_level_id: 2,
3108                                ..Default::default()
3109                            },
3110                        ],
3111                        total_file_size: 180,
3112                        uncompressed_file_size: 360,
3113                    },
3114                    levels: vec![Level {
3115                        level_idx: 1,
3116                        level_type: LevelType::Nonoverlapping,
3117                        table_infos: vec![
3118                            make_sst(4, vec![100, 200], 80), // shared
3119                            make_sst(5, vec![100], 90),      // only table 100
3120                        ],
3121                        total_file_size: 170,
3122                        uncompressed_file_size: 340,
3123                        ..Default::default()
3124                    }],
3125                    group_id: 1.into(),
3126                    parent_group_id: 0.into(),
3127                    member_table_ids: vec![],
3128                    compaction_group_version_id: 0,
3129                };
3130                levels
3131            })]),
3132            ..Default::default()
3133        };
3134
3135        let version_delta = HummockVersionDelta {
3136            id: HummockVersionId::new(1),
3137            group_deltas: HashMap::from_iter([(
3138                1.into(),
3139                GroupDeltas {
3140                    group_deltas: vec![GroupDelta::PruneTableIdsFromSsts(HashSet::from([
3141                        TableId::new(100),
3142                    ]))],
3143                },
3144            )]),
3145            ..Default::default()
3146        };
3147
3148        assert!(
3149            version.build_sst_delta_infos(&version_delta).is_empty(),
3150            "prune-only deltas must not create SST refill work"
3151        );
3152        version.apply_version_delta(&version_delta);
3153
3154        let cg = version.get_compaction_group_levels(1.into());
3155
3156        assert_eq!(
3157            cg.l0.sub_levels.len(),
3158            1,
3159            "empty sub-level should be removed"
3160        );
3161        assert_eq!(cg.l0.sub_levels[0].sub_level_id, 1);
3162        assert_eq!(cg.l0.sub_levels[0].table_infos.len(), 1);
3163        assert_eq!(2, cg.l0.sub_levels[0].table_infos[0].sst_id);
3164        assert_eq!(
3165            cg.l0.sub_levels[0].table_infos[0].table_ids,
3166            vec![TableId::new(200)]
3167        );
3168
3169        assert_eq!(cg.l0.total_file_size, 60);
3170        assert_eq!(cg.l0.uncompressed_file_size, 120);
3171
3172        assert_eq!(cg.levels[0].table_infos.len(), 1);
3173        assert_eq!(4, cg.levels[0].table_infos[0].sst_id);
3174        assert_eq!(
3175            cg.levels[0].table_infos[0].table_ids,
3176            vec![TableId::new(200)]
3177        );
3178        assert_eq!(cg.levels[0].total_file_size, 80);
3179        assert_eq!(cg.levels[0].uncompressed_file_size, 160);
3180
3181        assert_eq!(cg.compaction_group_version_id, 1);
3182    }
3183
3184    #[test]
3185    fn test_apply_version_delta_removes_dropped_table_watermark() {
3186        let table_id = TableId::new(100);
3187        let mut version = HummockVersion {
3188            id: HummockVersionId::new(0),
3189            table_watermarks: HashMap::from([(
3190                table_id,
3191                Arc::new(TableWatermarks::single_epoch(
3192                    test_epoch(1),
3193                    vec![VnodeWatermark::new(
3194                        Arc::new(Bitmap::ones(VirtualNode::COUNT_FOR_TEST)),
3195                        Bytes::from_static(b"watermark"),
3196                    )],
3197                    WatermarkDirection::Ascending,
3198                    WatermarkSerdeType::PkPrefix,
3199                )),
3200            )]),
3201            ..Default::default()
3202        };
3203
3204        let version_delta = HummockVersionDelta {
3205            id: HummockVersionId::new(1),
3206            prev_id: HummockVersionId::new(0),
3207            removed_table_ids: HashSet::from([table_id]),
3208            ..Default::default()
3209        };
3210
3211        version.apply_version_delta(&version_delta);
3212
3213        assert!(!version.table_watermarks.contains_key(&table_id));
3214    }
3215
3216    #[test]
3217    fn test_prune_stale_table_ids_from_ssts() {
3218        let live_table_id = TableId::new(100);
3219        let stale_table_id = TableId::new(200);
3220        let mut version = HummockVersion {
3221            id: HummockVersionId::new(0),
3222            levels: HashMap::from_iter([(1.into(), {
3223                #[expect(deprecated)]
3224                let levels = Levels {
3225                    l0: OverlappingLevel {
3226                        sub_levels: vec![Level {
3227                            level_idx: 0,
3228                            level_type: LevelType::Overlapping,
3229                            table_infos: vec![make_sst(
3230                                1,
3231                                vec![live_table_id.as_raw_id(), stale_table_id.as_raw_id()],
3232                                50,
3233                            )],
3234                            total_file_size: 50,
3235                            uncompressed_file_size: 100,
3236                            sub_level_id: 1,
3237                            ..Default::default()
3238                        }],
3239                        total_file_size: 50,
3240                        uncompressed_file_size: 100,
3241                    },
3242                    levels: vec![Level {
3243                        level_idx: 1,
3244                        level_type: LevelType::Nonoverlapping,
3245                        table_infos: vec![make_sst(2, vec![stale_table_id.as_raw_id()], 60)],
3246                        total_file_size: 60,
3247                        uncompressed_file_size: 120,
3248                        ..Default::default()
3249                    }],
3250                    group_id: 1.into(),
3251                    parent_group_id: 0.into(),
3252                    member_table_ids: vec![],
3253                    compaction_group_version_id: 0,
3254                };
3255                levels
3256            })]),
3257            state_table_info: HummockVersionStateTableInfo::from_protobuf_owned(
3258                HashMap::from_iter([(
3259                    live_table_id,
3260                    StateTableInfo {
3261                        committed_epoch: 1,
3262                        compaction_group_id: 1.into(),
3263                    },
3264                )]),
3265            ),
3266            ..Default::default()
3267        };
3268
3269        assert_eq!(version.prune_stale_table_ids_from_ssts(), 1);
3270
3271        let cg = version.get_compaction_group_levels(1.into());
3272        assert_eq!(cg.l0.sub_levels.len(), 1);
3273        assert_eq!(cg.l0.sub_levels[0].table_infos.len(), 1);
3274        assert_eq!(cg.l0.sub_levels[0].table_infos[0].sst_id, 1);
3275        assert_eq!(
3276            cg.l0.sub_levels[0].table_infos[0].table_ids,
3277            vec![live_table_id]
3278        );
3279        assert!(cg.levels[0].table_infos.is_empty());
3280        assert_eq!(cg.compaction_group_version_id, 1);
3281    }
3282
3283    #[test]
3284    fn test_prune_stale_table_ids_from_ssts_skips_legacy_member_table_ids() {
3285        let mut version = HummockVersion {
3286            id: HummockVersionId::new(0),
3287            levels: HashMap::from_iter([(1.into(), {
3288                #[expect(deprecated)]
3289                let levels = Levels {
3290                    l0: OverlappingLevel {
3291                        sub_levels: vec![Level {
3292                            level_idx: 0,
3293                            level_type: LevelType::Overlapping,
3294                            table_infos: vec![make_sst(1, vec![100, 200], 50)],
3295                            total_file_size: 50,
3296                            uncompressed_file_size: 100,
3297                            sub_level_id: 1,
3298                            ..Default::default()
3299                        }],
3300                        total_file_size: 50,
3301                        uncompressed_file_size: 100,
3302                    },
3303                    group_id: 1.into(),
3304                    parent_group_id: 0.into(),
3305                    member_table_ids: vec![100, 200],
3306                    compaction_group_version_id: 0,
3307                    ..Default::default()
3308                };
3309                levels
3310            })]),
3311            ..Default::default()
3312        };
3313
3314        assert_eq!(version.prune_stale_table_ids_from_ssts(), 0);
3315
3316        let cg = version.get_compaction_group_levels(1.into());
3317        assert_eq!(
3318            cg.l0.sub_levels[0].table_infos[0].table_ids,
3319            vec![TableId::new(100), TableId::new(200)]
3320        );
3321        assert_eq!(cg.compaction_group_version_id, 0);
3322    }
3323
3324    #[test]
3325    fn test_apply_version_delta_compact_l0() {
3326        let mut version = HummockVersion {
3327            id: HummockVersionId::new(0),
3328            levels: HashMap::from_iter([(1.into(), {
3329                #[expect(deprecated)]
3330                let levels = Levels {
3331                    l0: OverlappingLevel {
3332                        sub_levels: vec![
3333                            Level {
3334                                level_idx: 0,
3335                                level_type: LevelType::Nonoverlapping,
3336                                table_infos: vec![
3337                                    make_sst(1, vec![1], 100),
3338                                    make_sst(2, vec![2], 200),
3339                                ],
3340                                total_file_size: 300,
3341                                uncompressed_file_size: 600,
3342                                sub_level_id: 1,
3343                                ..Default::default()
3344                            },
3345                            Level {
3346                                level_idx: 0,
3347                                level_type: LevelType::Nonoverlapping,
3348                                table_infos: vec![make_sst(3, vec![3], 300)],
3349                                total_file_size: 300,
3350                                uncompressed_file_size: 600,
3351                                sub_level_id: 2,
3352                                ..Default::default()
3353                            },
3354                        ],
3355                        total_file_size: 600,
3356                        uncompressed_file_size: 1200,
3357                    },
3358                    levels: vec![Level {
3359                        level_idx: 1,
3360                        level_type: LevelType::Nonoverlapping,
3361                        table_infos: vec![],
3362                        total_file_size: 0,
3363                        uncompressed_file_size: 0,
3364                        ..Default::default()
3365                    }],
3366                    group_id: 1.into(),
3367                    parent_group_id: 0.into(),
3368                    member_table_ids: vec![],
3369                    compaction_group_version_id: 0,
3370                };
3371                levels
3372            })]),
3373            ..Default::default()
3374        };
3375
3376        let version_delta = HummockVersionDelta {
3377            id: HummockVersionId::new(1),
3378            group_deltas: HashMap::from_iter([(
3379                1.into(),
3380                GroupDeltas {
3381                    group_deltas: vec![
3382                        GroupDelta::IntraLevel(IntraLevelDelta::new(
3383                            0, // L0
3384                            0,
3385                            HashSet::from([1.into(), 2.into(), 3.into()]),
3386                            vec![],
3387                            0,
3388                            0,
3389                        )),
3390                        GroupDelta::IntraLevel(IntraLevelDelta::new(
3391                            1, // L1
3392                            0,
3393                            HashSet::new(),
3394                            vec![make_sst(10, vec![1, 2, 3], 500)],
3395                            0,
3396                            0,
3397                        )),
3398                    ],
3399                },
3400            )]),
3401            ..Default::default()
3402        };
3403
3404        version.apply_version_delta(&version_delta);
3405
3406        let cg = version.get_compaction_group_levels(1.into());
3407
3408        assert!(cg.l0.sub_levels.is_empty());
3409        assert_eq!(cg.l0.total_file_size, 0);
3410        assert_eq!(cg.l0.uncompressed_file_size, 0);
3411
3412        assert_eq!(cg.levels[0].table_infos.len(), 1);
3413        assert_eq!(10, cg.levels[0].table_infos[0].sst_id);
3414        assert_eq!(cg.levels[0].total_file_size, 500);
3415    }
3416}