refactor(storage): sstable iter compare fullkey struct instead of enc…

…oded key to avoid memory allocation (risingwavelabs#8607)
Little-Wallace · Mar 17, 2023 · f4a2f8d · f4a2f8d
1 parent ab701c7
commit f4a2f8d
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 37 deletions.
diff --git a/src/storage/src/hummock/compactor/iterator.rs b/src/storage/src/hummock/compactor/iterator.rs
@@ -388,14 +388,12 @@ impl HummockIterator for ConcatSstableIterator {
     /// Resets the iterator and seeks to the first position where the stored key >= `key`.
     fn seek<'a>(&'a mut self, key: FullKey<&'a [u8]>) -> Self::SeekFuture<'a> {
         async move {
-            let encoded_key = key.encode();
-            let key_slice = encoded_key.as_slice();
-            let seek_key: &[u8] = if self.key_range.left.is_empty() {
-                key_slice
+            let seek_key = if self.key_range.left.is_empty() {
+                key
             } else {
-                match KeyComparator::compare_encoded_full_key(key_slice, &self.key_range.left) {
-                    Ordering::Less | Ordering::Equal => &self.key_range.left,
-                    Ordering::Greater => key_slice,
+                match key.cmp(&FullKey::decode(&self.key_range.left)) {
+                    Ordering::Less | Ordering::Equal => FullKey::decode(&self.key_range.left),
+                    Ordering::Greater => key,
                 }
             };
             let table_idx = self.tables.partition_point(|table| {
@@ -406,7 +404,7 @@ impl HummockIterator for ConcatSstableIterator {
                 // Note that we need to use `<` instead of `<=` to ensure that all keys in an SST
                 // (including its max. key) produce the same search result.
                 let max_sst_key = &table.key_range.as_ref().unwrap().right;
-                KeyComparator::compare_encoded_full_key(max_sst_key, seek_key) == Ordering::Less
+                FullKey::decode(max_sst_key).cmp(&seek_key) == Ordering::Less
             });
 
             self.seek_idx(table_idx, Some(key)).await
@@ -424,7 +422,6 @@ mod tests {
 
     use risingwave_hummock_sdk::key::{next_full_key, prev_full_key, FullKey};
     use risingwave_hummock_sdk::key_range::KeyRange;
-    use risingwave_hummock_sdk::KeyComparator;
 
     use crate::hummock::compactor::ConcatSstableIterator;
     use crate::hummock::iterator::test_utils::mock_sstable_store;
@@ -597,11 +594,9 @@ mod tests {
         let mut iter =
             ConcatSstableIterator::new(table_infos.clone(), kr.clone(), sstable_store.clone());
         // Use block_2_smallest_key as seek key and result in invalid iterator.
-        let seek_key = block_2_smallest_key.clone();
-        assert!(KeyComparator::compare_encoded_full_key(&seek_key, &kr.right) == Ordering::Greater);
-        iter.seek_idx(0, Some(FullKey::decode(&seek_key)))
-            .await
-            .unwrap();
+        let seek_key = FullKey::decode(&block_2_smallest_key);
+        assert!(seek_key.cmp(&FullKey::decode(&kr.right)) == Ordering::Greater);
+        iter.seek_idx(0, Some(seek_key)).await.unwrap();
         assert!(!iter.is_valid());
         // Use a small enough seek key and result in the second KV of block 1.
         let seek_key = test_key_of(0).encode();

diff --git a/src/storage/src/hummock/iterator/concat_inner.rs b/src/storage/src/hummock/iterator/concat_inner.rs
@@ -19,7 +19,6 @@ use std::sync::Arc;
 use itertools::Itertools;
 use risingwave_common::must_match;
 use risingwave_hummock_sdk::key::FullKey;
-use risingwave_hummock_sdk::KeyComparator;
 use risingwave_pb::hummock::SstableInfo;
 
 use crate::hummock::iterator::{DirectionEnum, HummockIterator, HummockIteratorDirection};
@@ -195,22 +194,17 @@ impl<TI: SstableIteratorType> HummockIterator for ConcatIteratorInner<TI> {
 
     fn seek<'a>(&'a mut self, key: FullKey<&'a [u8]>) -> Self::SeekFuture<'a> {
         async move {
-            let encoded_key = key.encode();
             let table_idx = self
                 .tables
                 .partition_point(|table| match Self::Direction::direction() {
                     DirectionEnum::Forward => {
-                        let ord = KeyComparator::compare_encoded_full_key(
-                            table.smallest_key(),
-                            &encoded_key[..],
-                        );
+                        let ord = FullKey::decode(table.smallest_key()).cmp(&key);
+
                         ord == Less || ord == Equal
                     }
                     DirectionEnum::Backward => {
-                        let ord = KeyComparator::compare_encoded_full_key(
-                            table.largest_key(),
-                            &encoded_key[..],
-                        );
+                        let ord = FullKey::decode(table.largest_key()).cmp(&key);
+
                         ord == Greater || ord == Equal
                     }
                 })

diff --git a/src/storage/src/hummock/sstable/backward_sstable_iterator.rs b/src/storage/src/hummock/sstable/backward_sstable_iterator.rs
@@ -17,7 +17,6 @@ use std::future::Future;
 use std::sync::Arc;
 
 use risingwave_hummock_sdk::key::FullKey;
-use risingwave_hummock_sdk::KeyComparator;
 
 use crate::hummock::iterator::{Backward, HummockIterator};
 use crate::hummock::sstable::SstableIteratorReadOptions;
@@ -132,8 +131,6 @@ impl HummockIterator for BackwardSstableIterator {
 
     fn seek<'a>(&'a mut self, key: FullKey<&'a [u8]>) -> Self::SeekFuture<'a> {
         async move {
-            let encoded_key = key.encode();
-            let encoded_key_slice = encoded_key.as_slice();
             let block_idx = self
                 .sst
                 .value()
@@ -143,10 +140,7 @@ impl HummockIterator for BackwardSstableIterator {
                     // Compare by version comparator
                     // Note: we are comparing against the `smallest_key` of the `block`, thus the
                     // partition point should be `prev(<=)` instead of `<`.
-                    let ord = KeyComparator::compare_encoded_full_key(
-                        block_meta.smallest_key.as_slice(),
-                        encoded_key_slice,
-                    );
+                    let ord = FullKey::decode(&block_meta.smallest_key).cmp(&key);
                     ord == Less || ord == Equal
                 })
                 .saturating_sub(1); // considering the boundary of 0

diff --git a/src/storage/src/hummock/sstable/forward_sstable_iterator.rs b/src/storage/src/hummock/sstable/forward_sstable_iterator.rs
@@ -19,7 +19,6 @@ use std::ops::Bound::*;
 use std::sync::Arc;
 
 use risingwave_hummock_sdk::key::FullKey;
-use risingwave_hummock_sdk::KeyComparator;
 
 use super::super::{HummockResult, HummockValue};
 use super::Sstable;
@@ -301,7 +300,6 @@ impl HummockIterator for SstableIterator {
 
     fn seek<'a>(&'a mut self, key: FullKey<&'a [u8]>) -> Self::SeekFuture<'a> {
         async move {
-            let encoded_key = key.encode();
             let block_idx = self
                 .sst
                 .value()
@@ -311,10 +309,7 @@ impl HummockIterator for SstableIterator {
                     // compare by version comparator
                     // Note: we are comparing against the `smallest_key` of the `block`, thus the
                     // partition point should be `prev(<=)` instead of `<`.
-                    let ord = KeyComparator::compare_encoded_full_key(
-                        block_meta.smallest_key.as_slice(),
-                        encoded_key.as_slice(),
-                    );
+                    let ord = FullKey::decode(&block_meta.smallest_key).cmp(&key);
                     ord == Less || ord == Equal
                 })
                 .saturating_sub(1); // considering the boundary of 0