feat(storage): add may_exist interface in state table (#7489)

This PR adds the `may_exist` interface in both local state store and state table, which will be used by the streaming executor to check the existence of a given key range in storage efficiently. A typical use case is for hash join to check existence of a join key before refilling its operator cache on write (#7393 (comment))

Note that there can be false positive when using this interface but the false positive rate can be improved by providing a `prefix_hint` in ReadOptions to check SST bloom filters. In the current implementation, state table will always provide the `prefix_hint` in its `may_exist` implementation.

LocalStateStore:
```rust
/// Check existence of a given `key_range`.
/// It is better to provide `prefix_hint` in `read_options`, which will be used
/// for checking bloom filter if hummock is used. If `prefix_hint` is not provided,
/// the false positive rate can be significantly higher because bloom filter cannot
/// be used.
///
/// Returns:
/// - false: `key_range` is guaranteed to be absent in storage.
/// - true: `key_range` may or may not exist in storage.
fn may_exist(
&self,
key_range: (Bound<Vec<u8>>, Bound<Vec<u8>>),
read_options: ReadOptions,
) -> Self::MayExistFuture<'_>;
```

StateTable:
```rust
/// Returns:
/// false: the provided pk prefix is absent in state store.
/// true: the provided pk prefix may or may not be present in state store.
pub async fn may_exist(&self, pk_prefix: impl Row) -> StreamExecutorResult<bool> {
```

Approved-By: Li0k
Approved-By: wenym1
Approved-By: KeXiangWang
Approved-By: wcy-fdu
Approved-By: Little-Wallace

Author: hzxa21

grafana/risingwave-dashboard.dashboard.py

@@ -1696,6 +1696,23 @@ def section_hummock(panels):
                 ),
             ],
         ),
+        panels.timeseries_latency(
+            "Read Duration - MayExist",
+            "",
+            [
+                *quantile(
+                    lambda quantile, legend: panels.target(
+                        f"histogram_quantile({quantile}, sum(rate({metric('state_store_may_exist_duration_bucket')}[$__rate_interval])) by (le, job, instance, table_id))",
+                        f"p{legend}" + " - {{table_id}} @ {{job}} @ {{instance}}",
+                    ),
+                    [50, 90, 99, "max"],
+                ),
+                panels.target(
+                    f"sum by(le, job, instance, table_id)(rate({metric('state_store_may_exist_duration_sum')}[$__rate_interval])) / sum by(le, job, instance, table_id) (rate({metric('state_store_may_exist_duration_count')}[$__rate_interval]))",
+                    "avg - {{table_id}} {{job}} @ {{instance}}",
+                ),
+            ],
+        ),
         panels.timeseries_ops(
             "Read Bloom Filter",
             "",

grafana/risingwave-dashboard.json

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::cmp::Ordering;
 use std::fmt::Debug;
 use std::future::Future;
 use std::marker::PhantomData;
@@ -23,7 +24,7 @@ use std::sync::{Arc, LazyLock};
 use bytes::Bytes;
 use itertools::Itertools;
 use risingwave_common::catalog::TableId;
-use risingwave_hummock_sdk::key::{FullKey, TableKey, UserKey};
+use risingwave_hummock_sdk::key::{FullKey, TableKey, TableKeyRange, UserKey};
 
 use crate::hummock::iterator::{
     Backward, DeleteRangeIterator, DirectionEnum, Forward, HummockIterator,
@@ -194,6 +195,33 @@ impl SharedBufferBatch {
                 .le(table_key.as_ref())
     }
 
+    pub fn range_exists(&self, table_key_range: &TableKeyRange) -> bool {
+        self.inner
+            .binary_search_by(|m| {
+                let key = &m.0;
+                let too_left = match &table_key_range.0 {
+                    std::ops::Bound::Included(range_start) => range_start.as_ref() > key.as_ref(),
+                    std::ops::Bound::Excluded(range_start) => range_start.as_ref() >= key.as_ref(),
+                    std::ops::Bound::Unbounded => false,
+                };
+                if too_left {
+                    return Ordering::Less;
+                }
+
+                let too_right = match &table_key_range.1 {
+                    std::ops::Bound::Included(range_end) => range_end.as_ref() < key.as_ref(),
+                    std::ops::Bound::Excluded(range_end) => range_end.as_ref() <= key.as_ref(),
+                    std::ops::Bound::Unbounded => false,
+                };
+                if too_right {
+                    return Ordering::Greater;
+                }
+
+                Ordering::Equal
+            })
+            .is_ok()
+    }
+
     pub fn into_directed_iter<D: HummockIteratorDirection>(self) -> SharedBufferBatchIterator<D> {
         SharedBufferBatchIterator::<D>::new(self.inner, self.table_id, self.epoch)
     }
@@ -498,7 +526,10 @@ impl DeleteRangeIterator for SharedBufferDeleteRangeIterator {
 
 #[cfg(test)]
 mod tests {
+    use std::ops::Bound::{Excluded, Included};
+
     use itertools::Itertools;
+    use risingwave_hummock_sdk::key::map_table_key_range;
 
     use super::*;
     use crate::hummock::iterator::test_utils::{
@@ -764,4 +795,48 @@ mod tests {
             .await
             .unwrap();
     }
+
+    #[tokio::test]
+    async fn test_shared_buffer_batch_range_existx() {
+        let epoch = 1;
+        let shared_buffer_items = vec![
+            (Vec::from("a_1"), HummockValue::put(Bytes::from("value1"))),
+            (Vec::from("a_3"), HummockValue::put(Bytes::from("value2"))),
+            (Vec::from("a_5"), HummockValue::put(Bytes::from("value3"))),
+            (Vec::from("b_2"), HummockValue::put(Bytes::from("value3"))),
+        ];
+        let shared_buffer_batch = SharedBufferBatch::for_test(
+            transform_shared_buffer(shared_buffer_items),
+            epoch,
+            Default::default(),
+        );
+
+        let range = (Included(Vec::from("a")), Excluded(Vec::from("b")));
+        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("a_")), Excluded(Vec::from("b_")));
+        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("a_1")), Included(Vec::from("a_1")));
+        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("a_1")), Included(Vec::from("a_2")));
+        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("a_0x")), Included(Vec::from("a_2x")));
+        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("a_")), Excluded(Vec::from("c_")));
+        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("b_0x")), Included(Vec::from("b_2x")));
+        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("b_2")), Excluded(Vec::from("c_1x")));
+        assert!(shared_buffer_batch.range_exists(&map_table_key_range(range)));
+
+        let range = (Included(Vec::from("a_0")), Excluded(Vec::from("a_1")));
+        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("a__0")), Excluded(Vec::from("a__5")));
+        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("b_1")), Excluded(Vec::from("b_2")));
+        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("b_3")), Excluded(Vec::from("c_1")));
+        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
+        let range = (Included(Vec::from("b__x")), Excluded(Vec::from("c__x")));
+        assert!(!shared_buffer_batch.range_exists(&map_table_key_range(range)));
+    }
 }

src/storage/src/hummock/shared_buffer/shared_buffer_batch.rs

@@ -14,7 +14,6 @@
 
 use std::cmp::Ordering;
 use std::future::Future;
-use std::ops::Bound::{Excluded, Included};
 use std::ops::{Bound, RangeBounds};
 use std::sync::Arc;
 use std::time::Duration;
@@ -30,7 +29,7 @@ use risingwave_hummock_sdk::key::{
     bound_table_key_range, map_table_key_range, user_key, FullKey, TableKey, TableKeyRange, UserKey,
 };
 use risingwave_hummock_sdk::key_range::KeyRangeCommon;
-use risingwave_hummock_sdk::{can_concat, HummockReadEpoch};
+use risingwave_hummock_sdk::HummockReadEpoch;
 use risingwave_pb::hummock::LevelType;
 use tokio::sync::oneshot;
 use tracing::log::warn;
@@ -51,7 +50,7 @@ use crate::hummock::iterator::{
 use crate::hummock::local_version::ReadVersion;
 use crate::hummock::shared_buffer::build_ordered_merge_iter;
 use crate::hummock::sstable::SstableIteratorReadOptions;
-use crate::hummock::utils::{prune_ssts, search_sst_idx};
+use crate::hummock::utils::{prune_nonoverlapping_ssts, prune_overlapping_ssts};
 use crate::hummock::{
     DeleteRangeAggregator, ForwardIter, HummockEpoch, HummockError, HummockIteratorType,
     HummockResult, Sstable,
@@ -62,8 +61,8 @@ use crate::monitor::{
 use crate::storage_value::StorageValue;
 use crate::store::*;
 use crate::{
-    define_state_store_associated_type, define_state_store_read_associated_type,
-    define_state_store_write_associated_type,
+    define_local_state_store_associated_type, define_state_store_associated_type,
+    define_state_store_read_associated_type, define_state_store_write_associated_type,
 };
 
 impl HummockStorageV1 {
@@ -139,8 +138,12 @@ impl HummockStorageV1 {
             }
             match level.level_type() {
                 LevelType::Overlapping | LevelType::Unspecified => {
-                    let sstable_infos =
-                        prune_ssts(level.table_infos.iter(), table_id, &(table_key..=table_key));
+                    let single_table_key_range = table_key..=table_key;
+                    let sstable_infos = prune_overlapping_ssts(
+                        &level.table_infos,
+                        table_id,
+                        &single_table_key_range,
+                    );
                     for sstable_info in sstable_infos {
                         stats_guard.local_stats.sub_iter_count += 1;
                         if let Some(v) = get_from_sstable_info(
@@ -300,24 +303,12 @@ impl HummockStorageV1 {
                 continue;
             }
             if level.level_type == LevelType::Nonoverlapping as i32 {
-                debug_assert!(can_concat(&level.table_infos));
-                let start_table_idx = match encoded_user_key_range.start_bound() {
-                    Included(key) | Excluded(key) => search_sst_idx(&level.table_infos, key),
-                    _ => 0,
-                };
-                let end_table_idx = match encoded_user_key_range.end_bound() {
-                    Included(key) | Excluded(key) => search_sst_idx(&level.table_infos, key),
-                    _ => level.table_infos.len().saturating_sub(1),
-                };
-                assert!(
-                    start_table_idx < level.table_infos.len()
-                        && end_table_idx < level.table_infos.len()
-                );
-                let matched_table_infos = &level.table_infos[start_table_idx..=end_table_idx];
+                let matched_table_infos =
+                    prune_nonoverlapping_ssts(&level.table_infos, &encoded_user_key_range);
 
                 let pruned_sstables = match T::Direction::direction() {
-                    DirectionEnum::Backward => matched_table_infos.iter().rev().collect_vec(),
-                    DirectionEnum::Forward => matched_table_infos.iter().collect_vec(),
+                    DirectionEnum::Backward => matched_table_infos.rev().collect_vec(),
+                    DirectionEnum::Forward => matched_table_infos.collect_vec(),
                 };
 
                 let mut sstables = vec![];
@@ -346,7 +337,8 @@ impl HummockStorageV1 {
                     iter_read_options.clone(),
                 )));
             } else {
-                let table_infos = prune_ssts(level.table_infos.iter(), table_id, &table_key_range);
+                let table_infos =
+                    prune_overlapping_ssts(&level.table_infos, table_id, &table_key_range);
                 for table_info in table_infos.into_iter().rev() {
                     let sstable = self
                         .sstable_store
@@ -463,7 +455,17 @@ impl StateStoreWrite for HummockStorageV1 {
     }
 }
 
-impl LocalStateStore for HummockStorageV1 {}
+impl LocalStateStore for HummockStorageV1 {
+    define_local_state_store_associated_type!();
+
+    fn may_exist(
+        &self,
+        _key_range: (Bound<Vec<u8>>, Bound<Vec<u8>>),
+        _read_options: ReadOptions,
+    ) -> Self::MayExistFuture<'_> {
+        async move { Ok(true) }
+    }
+}
 
 impl StateStore for HummockStorageV1 {
     type Local = Self;

src/storage/src/hummock/state_store_v1.rs

@@ -20,6 +20,7 @@ use minitrace::future::FutureExt;
 use parking_lot::RwLock;
 use risingwave_common::catalog::TableId;
 use risingwave_hummock_sdk::key::{map_table_key_range, TableKey, TableKeyRange};
+use risingwave_hummock_sdk::HummockEpoch;
 use tokio::sync::mpsc;
 use tracing::warn;
 
@@ -39,8 +40,8 @@ use crate::monitor::{HummockStateStoreMetrics, IterLocalMetricsGuard, StoreLocal
 use crate::storage_value::StorageValue;
 use crate::store::*;
 use crate::{
-    define_state_store_read_associated_type, define_state_store_write_associated_type,
-    StateStoreIter,
+    define_local_state_store_associated_type, define_state_store_read_associated_type,
+    define_state_store_write_associated_type, StateStoreIter,
 };
 
 pub struct LocalHummockStorage {
@@ -107,6 +108,25 @@ impl LocalHummockStorage {
             .iter(table_key_range, epoch, read_options, read_snapshot)
             .await
     }
+
+    pub async fn may_exist_inner(
+        &self,
+        key_range: (Bound<Vec<u8>>, Bound<Vec<u8>>),
+        read_options: ReadOptions,
+    ) -> StorageResult<bool> {
+        let table_key_range = map_table_key_range(key_range);
+
+        let read_snapshot = read_filter_for_local(
+            HummockEpoch::MAX, // Use MAX epoch to make sure we read from latest
+            read_options.table_id,
+            &table_key_range,
+            self.read_version.clone(),
+        )?;
+
+        self.hummock_version_reader
+            .may_exist(table_key_range, read_options, read_snapshot)
+            .await
+    }
 }
 
 impl StateStoreRead for LocalHummockStorage {
@@ -195,7 +215,17 @@ impl StateStoreWrite for LocalHummockStorage {
     }
 }
 
-impl LocalStateStore for LocalHummockStorage {}
+impl LocalStateStore for LocalHummockStorage {
+    define_local_state_store_associated_type!();
+
+    fn may_exist(
+        &self,
+        key_range: (Bound<Vec<u8>>, Bound<Vec<u8>>),
+        read_options: ReadOptions,
+    ) -> Self::MayExistFuture<'_> {
+        self.may_exist_inner(key_range, read_options)
+    }
+}
 
 impl LocalHummockStorage {
     pub fn new(