Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(storage): sstable iter compare fullkey struct instead of encoded key to avoid memory allocation #8607

Merged
merged 3 commits into from
Mar 17, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 9 additions & 14 deletions src/storage/src/hummock/compactor/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -341,14 +341,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| {
Expand All @@ -359,7 +357,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
Expand All @@ -377,7 +375,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;
Expand Down Expand Up @@ -547,11 +544,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();
Expand Down
14 changes: 4 additions & 10 deletions src/storage/src/hummock/iterator/concat_inner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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
}
})
Expand Down
8 changes: 1 addition & 7 deletions src/storage/src/hummock/sstable/backward_sstable_iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand Down
7 changes: 1 addition & 6 deletions src/storage/src/hummock/sstable/forward_sstable_iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand Down