feat(storage): support table watermark filter for in-mem state store

src/storage/hummock_sdk/src/table_watermark.rs

@@ -331,6 +331,10 @@ impl VnodeWatermark {
     pub fn vnode_bitmap(&self) -> &Bitmap {
         &self.vnode_bitmap
     }
+
+    pub fn watermark(&self) -> &Bytes {
+        &self.watermark
+    }
 }
 
 #[derive(Clone, Debug, PartialEq)]

src/storage/src/mem_table.rs

@@ -16,17 +16,20 @@ use std::cmp::Ordering;
 use std::collections::btree_map::Entry;
 use std::collections::BTreeMap;
 use std::future::Future;
-use std::ops::Bound::{Included, Unbounded};
+use std::ops::Bound::{Excluded, Included, Unbounded};
 use std::ops::{Bound, RangeBounds};
 
 use bytes::Bytes;
 use futures::{pin_mut, StreamExt};
 use futures_async_stream::try_stream;
+use itertools::Itertools;
 use risingwave_common::catalog::{TableId, TableOption};
 use risingwave_common::estimate_size::{EstimateSize, KvSize};
-use risingwave_hummock_sdk::key::{FullKey, TableKey, TableKeyRange};
+use risingwave_common::hash::VnodeBitmapExt;
+use risingwave_hummock_sdk::key::{prefixed_range_with_vnode, FullKey, TableKey, TableKeyRange};
+use risingwave_hummock_sdk::table_watermark::WatermarkDirection;
 use thiserror::Error;
-use tracing::warn;
+use tracing::error;
 
 use crate::error::{StorageError, StorageResult};
 use crate::hummock::iterator::{FromRustIterator, RustIteratorBuilder};
@@ -571,16 +574,14 @@ impl<S: StateStoreWrite + StateStoreRead> LocalStateStore for MemtableLocalState
                 }
             }
         }
-        self.inner
-            .ingest_batch(
-                kv_pairs,
-                delete_ranges,
-                WriteOptions {
-                    epoch: self.epoch(),
-                    table_id: self.table_id,
-                },
-            )
-            .await
+        self.inner.ingest_batch(
+            kv_pairs,
+            delete_ranges,
+            WriteOptions {
+                epoch: self.epoch(),
+                table_id: self.table_id,
+            },
+        )
     }
 
     fn epoch(&self) -> u64 {
@@ -613,8 +614,38 @@ impl<S: StateStoreWrite + StateStoreRead> LocalStateStore for MemtableLocalState
             next_epoch,
             prev_epoch
         );
-        if opts.table_watermarks.is_some() {
-            warn!("table watermark only supported in hummock state store");
+        if let Some((direction, watermarks)) = opts.table_watermarks {
+            let delete_ranges = watermarks
+                .iter()
+                .flat_map(|vnode_watermark| {
+                    let inner_range = match direction {
+                        WatermarkDirection::Ascending => {
+                            (Unbounded, Excluded(vnode_watermark.watermark().clone()))
+                        }
+                        WatermarkDirection::Descending => {
+                            (Excluded(vnode_watermark.watermark().clone()), Unbounded)
+                        }
+                    };
+                    vnode_watermark
+                        .vnode_bitmap()
+                        .iter_vnodes()
+                        .map(move |vnode| {
+                            let (start, end) =
+                                prefixed_range_with_vnode(inner_range.clone(), vnode);
+                            (start.map(|key| key.0.clone()), end.map(|key| key.0.clone()))
+                        })
+                })
+                .collect_vec();
+            if let Err(e) = self.inner.ingest_batch(
+                Vec::new(),
+                delete_ranges,
+                WriteOptions {
+                    epoch: self.epoch(),
+                    table_id: self.table_id,
+                },
+            ) {
+                error!(err = ?e, "failed to write delete ranges of table watermark");
+            }
         }
     }
 

src/storage/src/memory.rs

@@ -575,8 +575,7 @@ impl<R: RangeKv> StateStoreRead for RangeKvStateStore<R> {
 }
 
 impl<R: RangeKv> StateStoreWrite for RangeKvStateStore<R> {
-    #[allow(clippy::unused_async)]
-    async fn ingest_batch(
+    fn ingest_batch(
         &self,
         mut kv_pairs: Vec<(TableKey<Bytes>, StorageValue)>,
         delete_ranges: Vec<(Bound<Bytes>, Bound<Bytes>)>,
@@ -747,7 +746,6 @@ mod tests {
                     table_id: Default::default(),
                 },
             )
-            .await
             .unwrap();
         state_store
             .ingest_batch(
@@ -767,7 +765,6 @@ mod tests {
                     table_id: Default::default(),
                 },
             )
-            .await
             .unwrap();
         assert_eq!(
             state_store

src/storage/src/panic_store.rs

@@ -55,8 +55,7 @@ impl StateStoreRead for PanicStateStore {
 }
 
 impl StateStoreWrite for PanicStateStore {
-    #[allow(clippy::unused_async)]
-    async fn ingest_batch(
+    fn ingest_batch(
         &self,
         _kv_pairs: Vec<(TableKey<Bytes>, StorageValue)>,
         _delete_ranges: Vec<(Bound<Bytes>, Bound<Bytes>)>,

src/storage/src/store.rs

@@ -155,7 +155,7 @@ pub trait StateStoreWrite: StaticSendSync {
         kv_pairs: Vec<(TableKey<Bytes>, StorageValue)>,
         delete_ranges: Vec<(Bound<Bytes>, Bound<Bytes>)>,
         write_options: WriteOptions,
-    ) -> impl Future<Output = StorageResult<usize>> + Send + '_;
+    ) -> StorageResult<usize>;
 }
 
 #[derive(Default, Debug)]

src/storage/src/store_impl.rs

@@ -317,24 +317,19 @@ pub mod verify {
     }
 
     impl<A: StateStoreWrite, E: StateStoreWrite> StateStoreWrite for VerifyStateStore<A, E> {
-        async fn ingest_batch(
+        fn ingest_batch(
             &self,
             kv_pairs: Vec<(TableKey<Bytes>, StorageValue)>,
             delete_ranges: Vec<(Bound<Bytes>, Bound<Bytes>)>,
             write_options: WriteOptions,
         ) -> StorageResult<usize> {
-            let actual = self
-                .actual
-                .ingest_batch(
-                    kv_pairs.clone(),
-                    delete_ranges.clone(),
-                    write_options.clone(),
-                )
-                .await;
+            let actual = self.actual.ingest_batch(
+                kv_pairs.clone(),
+                delete_ranges.clone(),
+                write_options.clone(),
+            );
             if let Some(expected) = &self.expected {
-                let expected = expected
-                    .ingest_batch(kv_pairs, delete_ranges, write_options)
-                    .await;
+                let expected = expected.ingest_batch(kv_pairs, delete_ranges, write_options);
                 assert_eq!(actual.is_err(), expected.is_err());
             }
             actual