support exclude

Signed-off-by: Little-Wallace <[email protected]>
risingwavelabs · May 29, 2024 · 9bda810 · 9bda810
1 parent 11aac57
commit 9bda810
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 16 deletions.
diff --git a/src/storage/src/hummock/iterator/backward_user.rs b/src/storage/src/hummock/iterator/backward_user.rs
@@ -212,48 +212,60 @@ impl<I: HummockIterator<Direction = Backward>> BackwardUserIterator<I> {
     pub async fn rewind(&mut self) -> HummockResult<()> {
         // Handle range scan
         match &self.key_range.1 {
-            Included(end_key) => {
+            Included(end_key) | Excluded(end_key) => {
                 let full_key = FullKey {
                     user_key: end_key.clone(),
                     epoch_with_gap: EpochWithGap::new_min_epoch(),
                 };
                 self.iterator.seek(full_key.to_ref()).await?;
             }
-            Excluded(_) => unimplemented!("excluded begin key is not supported"),
             Unbounded => self.iterator.rewind().await?,
         };
 
         // Handle multi-version
         self.reset();
         // Handle range scan when key < begin_key
-        self.next().await
+        self.next().await?;
+        if let Excluded(end_key) = &self.key_range.0 {
+            if self.is_valid() && self.last_key.user_key.as_ref() == end_key.as_ref() {
+                self.next().await?;
+            }
+        }
+        Ok(())
     }
 
     /// Resets the iterating position to the first position where the key >= provided key.
     pub async fn seek(&mut self, user_key: UserKey<&[u8]>) -> HummockResult<()> {
         // Handle range scan when key > end_key
-        let user_key = match &self.key_range.1 {
-            Included(end_key) => {
+        let seek_key = match &self.key_range.1 {
+            Included(end_key) | Excluded(end_key) => {
                 let end_key = end_key.as_ref();
                 if end_key < user_key {
                     end_key
                 } else {
                     user_key
                 }
             }
-            Excluded(_) => unimplemented!("excluded begin key is not supported"),
             Unbounded => user_key,
         };
         let full_key = FullKey {
-            user_key,
+            user_key: seek_key.clone(),
             epoch_with_gap: EpochWithGap::new_min_epoch(),
         };
         self.iterator.seek(full_key).await?;
 
         // Handle multi-version
         self.reset();
         // Handle range scan when key < begin_key
-        self.next().await
+        self.next().await?;
+        if let Excluded(end_key) = &self.key_range.0
+            && end_key.as_ref() >= user_key
+        {
+            if self.is_valid() && self.last_key.user_key.as_ref() == end_key.as_ref() {
+                self.next().await?;
+            }
+        }
+        Ok(())
     }
 
     /// Indicates whether the iterator can be used.

diff --git a/src/storage/src/hummock/iterator/forward_user.rs b/src/storage/src/hummock/iterator/forward_user.rs
@@ -124,20 +124,28 @@ impl<I: HummockIterator<Direction = Forward>> UserIterator<I> {
 
         // Handle range scan
         match &self.key_range.0 {
-            Included(begin_key) => {
+            Included(begin_key) | Excluded(begin_key) => {
                 let full_key = FullKey {
                     user_key: begin_key.clone(),
                     epoch_with_gap: EpochWithGap::new(self.read_epoch, MAX_SPILL_TIMES),
                 };
                 self.iterator.seek(full_key.to_ref()).await?;
             }
-            Excluded(_) => unimplemented!("excluded begin key is not supported"),
             Unbounded => {
                 self.iterator.rewind().await?;
             }
         };
 
-        self.try_advance_to_next_valid().await
+        self.try_advance_to_next_valid().await?;
+        match &self.key_range.0 {
+            Excluded(begin_key) => {
+                if self.key().user_key == begin_key.as_ref() {
+                    self.next().await?;
+                }
+            }
+            _ => {}
+        }
+        Ok(())
     }
 
     /// Resets the iterating position to the first position where the key >= provided key.
@@ -147,26 +155,32 @@ impl<I: HummockIterator<Direction = Forward>> UserIterator<I> {
         self.full_key_tracker = FullKeyTracker::new(FullKey::default());
 
         // Handle range scan when key < begin_key
-        let user_key = match &self.key_range.0 {
-            Included(begin_key) => {
+        let seek_key = match &self.key_range.0 {
+            Included(begin_key) | Excluded(begin_key) => {
                 let begin_key = begin_key.as_ref();
                 if begin_key > user_key {
                     begin_key
                 } else {
                     user_key
                 }
             }
-            Excluded(_) => unimplemented!("excluded begin key is not supported"),
             Unbounded => user_key,
         };
 
         let full_key = FullKey {
-            user_key,
+            user_key: seek_key.clone(),
             epoch_with_gap: EpochWithGap::new(self.read_epoch, MAX_SPILL_TIMES),
         };
         self.iterator.seek(full_key).await?;
 
-        self.try_advance_to_next_valid().await
+        self.try_advance_to_next_valid().await?;
+        if let Excluded(begin_key) = &self.key_range.0
+            && begin_key.as_ref() >= user_key
+            && self.key().user_key == begin_key.as_ref()
+        {
+            self.next().await?;
+        }
+        Ok(())
     }
 
     /// Indicates whether the iterator can be used.