add ut

Signed-off-by: Little-Wallace <[email protected]>

src/meta/src/hummock/compaction/level_selector.rs

@@ -615,6 +615,7 @@ pub mod tests {
             uncompressed_file_size: sst.uncompressed_file_size,
             sub_level_id: sst.get_sst_id(),
             table_infos: vec![sst],
+            vnode_partition_count: 0,
         });
     }
 
@@ -645,6 +646,7 @@ pub mod tests {
             sub_level_id,
             table_infos,
             uncompressed_file_size,
+            vnode_partition_count: 0,
         });
     }
 
@@ -729,6 +731,7 @@ pub mod tests {
             total_file_size,
             sub_level_id: 0,
             uncompressed_file_size,
+            vnode_partition_count: 0,
         }
     }
 
@@ -751,6 +754,7 @@ pub mod tests {
                     uncompressed_file_size: table.uncompressed_file_size,
                     sub_level_id: idx as u64,
                     table_infos: vec![table],
+                    vnode_partition_count: 0,
                 })
                 .collect_vec(),
             total_file_size,
@@ -775,6 +779,7 @@ pub mod tests {
                         .sum::<u64>(),
                     sub_level_id: idx as u64,
                     table_infos: table,
+                    vnode_partition_count: 0,
                 })
                 .collect_vec(),
             total_file_size: 0,
@@ -809,6 +814,7 @@ pub mod tests {
                         .iter()
                         .map(|sst| sst.uncompressed_file_size)
                         .sum::<u64>(),
+                    vnode_partition_count: 0,
                 })
                 .collect_vec(),
             total_file_size: 0,

src/meta/src/hummock/compaction/mod.rs

@@ -38,6 +38,7 @@ pub use crate::hummock::compaction::level_selector::{
     ManualCompactionSelector, SpaceReclaimCompactionSelector, TtlCompactionSelector,
 };
 use crate::hummock::compaction::overlap_strategy::{OverlapStrategy, RangeOverlapStrategy};
+pub use crate::hummock::compaction::picker::{partition_level, SubLevelPartition};
 use crate::hummock::compaction::picker::{CompactionInput, LocalPickerStatistic};
 use crate::hummock::level_handler::LevelHandler;
 use crate::hummock::model::CompactionGroup;

src/meta/src/hummock/compaction/picker/intral_sub_level_picker.rs

@@ -1,3 +1,18 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::ops::Bound;
 use std::sync::Arc;
 
 use risingwave_common::hash::VirtualNode;
@@ -119,7 +134,7 @@ impl CompactionPicker for IntraSubLevelPicker {
                 input_levels,
                 target_level: 0,
                 target_sub_level_id: level.sub_level_id,
-                vnode_partition_count: vnode_partition_count,
+                vnode_partition_count,
             });
         }
 
@@ -224,16 +239,23 @@ pub fn partition_level(
     level: &Level,
     partitions: &mut Vec<SubLevelPartition>,
 ) -> bool {
+    assert_eq!(partition_vnode_count, partitions.len());
     let mut left_idx = 0;
     let mut can_partition = true;
     let partition_size = VirtualNode::COUNT / partition_vnode_count;
-    for partition_id in 0..partition_vnode_count {
+    for (partition_id, partition) in partitions.iter_mut().enumerate() {
         let smallest_vnode = partition_id * partition_size;
         let largest_vnode = (partition_id + 1) * partition_size;
         let smallest_table_key =
             UserKey::prefix_of_vnode(table_id, VirtualNode::from_index(smallest_vnode));
-        let largest_table_key =
-            UserKey::prefix_of_vnode(table_id, VirtualNode::from_index(largest_vnode));
+        let largest_table_key = if largest_vnode >= VirtualNode::COUNT {
+            Bound::Unbounded
+        } else {
+            Bound::Excluded(UserKey::prefix_of_vnode(
+                table_id,
+                VirtualNode::from_index(largest_vnode),
+            ))
+        };
         while left_idx < level.table_infos.len() {
             let key_range = level.table_infos[left_idx].key_range.as_ref().unwrap();
             let ret = key_range.compare_right_with_user_key(smallest_table_key.as_ref());
@@ -243,7 +265,7 @@ pub fn partition_level(
             left_idx += 1;
         }
         if left_idx >= level.table_infos.len() {
-            partitions[partition_id].sub_levels.push(PartitionInfo {
+            partition.sub_levels.push(PartitionInfo {
                 sub_level_id: level.sub_level_id,
                 left_idx: 0,
                 right_idx: 0,
@@ -264,29 +286,45 @@ pub fn partition_level(
         let mut right_idx = left_idx;
         while right_idx < level.table_infos.len() {
             let key_range = level.table_infos[right_idx].key_range.as_ref().unwrap();
-            let ret = key_range.compare_right_with_user_key(largest_table_key.as_ref());
+            let ret = match &largest_table_key {
+                Bound::Excluded(key) => key_range.compare_right_with_user_key(key.as_ref()),
+                Bound::Unbounded => {
+                    let right_key = FullKey::decode(&key_range.right);
+                    assert!(right_key.user_key.table_id.table_id == table_id);
+                    // We would assign partition_vnode_count to a level only when we compact all
+                    // sstable of it, so there will never be another stale table in this sstable
+                    // file.
+                    std::cmp::Ordering::Less
+                }
+                _ => unreachable!(),
+            };
+
             if ret != std::cmp::Ordering::Less {
                 break;
             }
             total_file_size += level.table_infos[right_idx].file_size;
             right_idx += 1;
         }
+
         if right_idx < level.table_infos.len()
-            && FullKey::decode(
-                &level.table_infos[right_idx]
-                    .key_range
-                    .as_ref()
-                    .unwrap()
-                    .left,
-            )
-            .user_key
-            .lt(&largest_table_key.as_ref())
+            && match &largest_table_key {
+                Bound::Excluded(key) => FullKey::decode(
+                    &level.table_infos[right_idx]
+                        .key_range
+                        .as_ref()
+                        .unwrap()
+                        .left,
+                )
+                .user_key
+                .lt(&key.as_ref()),
+                _ => unreachable!(),
+            }
         {
             can_partition = false;
             break;
         }
         left_idx = right_idx;
-        partitions[partition_id].sub_levels.push(PartitionInfo {
+        partition.sub_levels.push(PartitionInfo {
             sub_level_id: level.sub_level_id,
             left_idx,
             right_idx,

src/meta/src/hummock/compaction/picker/manual_compaction_picker.rs

@@ -382,9 +382,7 @@ pub mod tests {
                     generate_table(1, 1, 101, 200, 1),
                     generate_table(2, 1, 222, 300, 1),
                 ],
-                total_file_size: 0,
-                sub_level_id: 0,
-                uncompressed_file_size: 0,
+                ..Default::default()
             },
             Level {
                 level_idx: 2,
@@ -396,9 +394,7 @@ pub mod tests {
                     generate_table(7, 1, 501, 800, 1),
                     generate_table(8, 2, 301, 400, 1),
                 ],
-                total_file_size: 0,
-                sub_level_id: 0,
-                uncompressed_file_size: 0,
+                ..Default::default()
             },
         ];
         let mut levels = Levels {
@@ -560,9 +556,7 @@ pub mod tests {
                     generate_table(3, 1, 0, 100, 1),
                     generate_table(4, 2, 2000, 3000, 1),
                 ],
-                total_file_size: 0,
-                sub_level_id: 0,
-                uncompressed_file_size: 0,
+                ..Default::default()
             },
             Level {
                 level_idx: 2,
@@ -571,9 +565,7 @@ pub mod tests {
                     generate_table(1, 1, 0, 100, 1),
                     generate_table(2, 2, 2000, 3000, 1),
                 ],
-                total_file_size: 0,
-                sub_level_id: 0,
-                uncompressed_file_size: 0,
+                ..Default::default()
             },
         ];
         // Set internal_table_ids.
@@ -615,9 +607,7 @@ pub mod tests {
                 generate_table(3, 2, 200, 300, 1),
                 generate_table(4, 2, 300, 400, 1),
             ],
-            total_file_size: 0,
-            sub_level_id: 0,
-            uncompressed_file_size: 0,
+            ..Default::default()
         }];
         let levels = Levels {
             levels,
@@ -635,10 +625,7 @@ pub mod tests {
         let levels = vec![Level {
             level_idx: 1,
             level_type: LevelType::Nonoverlapping as i32,
-            table_infos: vec![],
-            total_file_size: 0,
-            sub_level_id: 0,
-            uncompressed_file_size: 0,
+            ..Default::default()
         }];
         let levels = Levels {
             levels,
@@ -1175,9 +1162,7 @@ pub mod tests {
                     generate_table(3, 1, 101, 200, 1),
                     generate_table(4, 1, 222, 300, 1),
                 ],
-                total_file_size: 0,
-                sub_level_id: 0,
-                uncompressed_file_size: 0,
+                ..Default::default()
             },
         ];
         assert_eq!(levels.len(), 4);
@@ -1285,9 +1270,7 @@ pub mod tests {
                     generate_table(6, 1, 444, 500, 1),
                     generate_table(7, 1, 555, 600, 1),
                 ],
-                total_file_size: 0,
-                sub_level_id: 0,
-                uncompressed_file_size: 0,
+                ..Default::default()
             },
         ];
         assert_eq!(levels.len(), 4);