risingwave_stream/executor/eowc/
buffer_table_cache.rs

1// Copyright 2026 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeSet;
16use std::ops::Bound;
17
18use anyhow::Context;
19use bytes::Bytes;
20use futures::StreamExt;
21use futures_async_stream::for_await;
22use risingwave_common::hash::VnodeBitmapExt;
23use risingwave_common::row::{self, OwnedRow};
24use risingwave_common_estimate_size::EstimateSize;
25use risingwave_storage::StateStore;
26use risingwave_storage::store::PrefetchOptions;
27use risingwave_storage::table::KeyedRow;
28use risingwave_storage::table::merge_sort::merge_sort;
29
30use crate::common::state_cache::{StateCache, StateCacheFiller, TopNStateCache};
31use crate::common::table::state_table::StateTable;
32use crate::executor::StreamExecutorResult;
33
34// TODO(rc): need to make this configurable?
35const CACHE_CAPACITY: usize = 2048;
36
37/// Caches the smallest rows of a buffer table, whose rows are consumed in the pk order by
38/// watermark, so that consuming doesn't scan the table every time. The cache key must be ordered
39/// as the pk of the table, and the writes to the table must be applied to the cache too.
40pub struct BufferTableCache<K: Ord + EstimateSize, V: EstimateSize> {
41    cache: TopNStateCache<K, V>,
42}
43
44impl<K: Ord + EstimateSize + Clone, V: EstimateSize> BufferTableCache<K, V> {
45    pub fn new() -> Self {
46        Self {
47            cache: TopNStateCache::new(CACHE_CAPACITY),
48        }
49    }
50
51    pub fn is_synced(&self) -> bool {
52        self.cache.is_synced()
53    }
54
55    pub fn insert(&mut self, key: K, value: V) {
56        self.cache.insert(key, value);
57    }
58
59    pub fn delete(&mut self, key: &K) -> Option<V> {
60        self.cache.delete(key)
61    }
62
63    /// Pop the first cached entry if `pred` holds for it. `None` is also returned if the cache is
64    /// not synced, e.g., after it runs out of entries, when it should be refilled after the last
65    /// popped entry.
66    pub fn pop_first_if(&mut self, pred: impl FnOnce(&K, &V) -> bool) -> Option<(K, V)> {
67        if !self.cache.is_synced() {
68            return None;
69        }
70        let (key, value) = self.cache.first_key_value()?;
71        if !pred(key, value) {
72            return None;
73        }
74        let key = key.clone();
75        let value = self.cache.delete(&key).expect("the first key must exist");
76        Some((key, value))
77    }
78
79    /// The table row count is known when a refill reads all rows, after which the cache stays
80    /// synced as long as the count is maintained by the writes.
81    pub fn set_table_row_count(&mut self, table_row_count: usize) {
82        self.cache.set_table_row_count(table_row_count);
83    }
84
85    pub fn clear(&mut self) {
86        self.cache.clear();
87    }
88
89    /// Clear the cache and fill it with the smallest rows of `table` after `start_after`, converted
90    /// by `entry`. Returns the number of rows read if they are all the rows after `start_after`.
91    pub async fn refill<S: StateStore>(
92        &mut self,
93        table: &StateTable<S>,
94        start_after: Option<OwnedRow>,
95        mut entry: impl FnMut(OwnedRow) -> StreamExecutorResult<(K, V)>,
96    ) -> StreamExecutorResult<Option<usize>> {
97        // Start over, as the row count of the table is known only if the refill reads all rows.
98        self.cache = TopNStateCache::new(CACHE_CAPACITY);
99        let mut filler = self.cache.begin_syncing();
100
101        let pk_range = (
102            start_after.map(Bound::Excluded).unwrap_or(Bound::Unbounded),
103            Bound::<row::Empty>::Unbounded,
104        );
105
106        let streams: Vec<_> =
107            futures::future::try_join_all(table.vnodes().iter_vnodes().map(|vnode| {
108                table.iter_keyed_row_with_vnode(
109                    vnode,
110                    &pk_range,
111                    PrefetchOptions::new(filler.capacity().is_none(), false),
112                )
113            }))
114            .await?
115            .into_iter()
116            .map(Box::pin)
117            .collect();
118
119        let capacity = filler.capacity().unwrap_or(usize::MAX);
120        let mut count = 0;
121        #[for_await]
122        for kv in merge_sort(streams).take(capacity) {
123            let (key, value) = entry(key_value_to_full_row(kv?, table)?)?;
124            filler.insert_unchecked(key, value);
125            count += 1;
126        }
127
128        filler.finish();
129        Ok((count < capacity).then_some(count))
130    }
131}
132
133impl<K: Ord + EstimateSize + Clone, V: EstimateSize> Default for BufferTableCache<K, V> {
134    fn default() -> Self {
135        Self::new()
136    }
137}
138
139/// Merge the key part and value part of a row into a full row. This is needed for state table with
140/// non-None value indices.
141fn key_value_to_full_row<S: StateStore>(
142    keyed_row: KeyedRow<Bytes>,
143    table: &StateTable<S>,
144) -> StreamExecutorResult<OwnedRow> {
145    let Some(val_indices) = table.value_indices() else {
146        return Ok(keyed_row.into_owned_row());
147    };
148    let pk_indices = table.pk_indices();
149    let indices: BTreeSet<_> = val_indices
150        .iter()
151        .chain(pk_indices.iter())
152        .copied()
153        .collect();
154    let len = indices.iter().max().unwrap() + 1;
155    assert!(indices.iter().copied().eq(0..len));
156
157    let mut row = vec![None; len];
158    let key = table
159        .pk_serde()
160        .deserialize(keyed_row.key())
161        .context("failed to deserialize pk")?;
162    for (i, v) in key.into_iter().enumerate() {
163        row[pk_indices[i]] = v;
164    }
165    for (i, v) in keyed_row.into_owned_row().into_iter().enumerate() {
166        row[val_indices[i]] = v;
167    }
168    Ok(OwnedRow::new(row))
169}
risingwave_stream/executor/eowc/buffer_table_cache.rs

risingwave_stream/executor/eowc/
buffer_table_cache.rs
