risingwave_stream/executor/over_window/
range_cache.rs

1// Copyright 2025 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::BTreeMap;
16use std::ops::{Bound, RangeInclusive};
17
18use risingwave_common::config::streaming::OverWindowCachePolicy as CachePolicy;
19use risingwave_common::row::OwnedRow;
20use risingwave_common::types::Sentinelled;
21use risingwave_common_estimate_size::EstimateSize;
22use risingwave_common_estimate_size::collections::EstimatedBTreeMap;
23use risingwave_expr::window_function::StateKey;
24use static_assertions::const_assert;
25
26pub(super) type CacheKey = Sentinelled<StateKey>;
27
28/// Range cache for one over window partition.
29/// The cache entries can be:
30///
31/// - `(Normal)*`
32/// - `Smallest, (Normal)*, Largest`
33/// - `(Normal)+, Largest`
34/// - `Smallest, (Normal)+`
35///
36/// This means it's impossible to only have one sentinel in the cache without any normal entry,
37/// and, each of the two types of sentinel can only appear once. Also, since sentinels are either
38/// smallest or largest, they always appear at the beginning or the end of the cache.
39#[derive(Clone, Debug, Default)]
40pub(super) struct PartitionCache {
41    inner: EstimatedBTreeMap<CacheKey, OwnedRow>,
42}
43
44impl PartitionCache {
45    /// Create a new empty partition cache without sentinel values.
46    pub fn new_without_sentinels() -> Self {
47        Self {
48            inner: EstimatedBTreeMap::new(),
49        }
50    }
51
52    /// Create a new empty partition cache with sentinel values.
53    pub fn new() -> Self {
54        let mut cache = Self {
55            inner: EstimatedBTreeMap::new(),
56        };
57        cache.insert(CacheKey::Smallest, OwnedRow::empty());
58        cache.insert(CacheKey::Largest, OwnedRow::empty());
59        cache
60    }
61
62    /// Get access to the inner `BTreeMap` for cursor operations.
63    pub fn inner(&self) -> &BTreeMap<CacheKey, OwnedRow> {
64        self.inner.inner()
65    }
66
67    /// Insert a key-value pair into the cache.
68    pub fn insert(&mut self, key: CacheKey, value: OwnedRow) -> Option<OwnedRow> {
69        self.inner.insert(key, value)
70    }
71
72    /// Remove a key from the cache.
73    pub fn remove(&mut self, key: &CacheKey) -> Option<OwnedRow> {
74        self.inner.remove(key)
75    }
76
77    /// Get the number of entries in the cache.
78    pub fn len(&self) -> usize {
79        self.inner.len()
80    }
81
82    /// Check if the cache is empty.
83    pub fn is_empty(&self) -> bool {
84        self.inner.is_empty()
85    }
86
87    /// Get the first key-value pair in the cache.
88    pub fn first_key_value(&self) -> Option<(&CacheKey, &OwnedRow)> {
89        self.inner.first_key_value()
90    }
91
92    /// Get the last key-value pair in the cache.
93    pub fn last_key_value(&self) -> Option<(&CacheKey, &OwnedRow)> {
94        self.inner.last_key_value()
95    }
96
97    /// Retain entries in the given range, removing others.
98    /// Returns `(left_removed, right_removed)` where sentinels are filtered out.
99    /// Sentinels are preserved in the cache.
100    fn retain_range(
101        &mut self,
102        range: RangeInclusive<&CacheKey>,
103    ) -> (BTreeMap<CacheKey, OwnedRow>, BTreeMap<CacheKey, OwnedRow>) {
104        // Check if we had sentinels before the operation
105        let had_smallest = self.inner.inner().contains_key(&CacheKey::Smallest);
106        let had_largest = self.inner.inner().contains_key(&CacheKey::Largest);
107
108        let (left_removed, right_removed) = self.inner.retain_range(range);
109
110        // Restore sentinels if they were present before
111        if had_smallest {
112            self.inner.insert(CacheKey::Smallest, OwnedRow::empty());
113        }
114        if had_largest {
115            self.inner.insert(CacheKey::Largest, OwnedRow::empty());
116        }
117
118        // Filter out sentinels from the returned maps
119        let left_removed = left_removed
120            .into_iter()
121            .filter(|(k, _)| k.is_normal())
122            .collect();
123        let right_removed = right_removed
124            .into_iter()
125            .filter(|(k, _)| k.is_normal())
126            .collect();
127
128        (left_removed, right_removed)
129    }
130
131    /// Get the number of cached `Sentinel::Normal` entries.
132    pub fn normal_len(&self) -> usize {
133        let len = self.inner().len();
134        if len <= 1 {
135            debug_assert!(
136                self.inner()
137                    .first_key_value()
138                    .map(|(k, _)| k.is_normal())
139                    .unwrap_or(true)
140            );
141            return len;
142        }
143        // len >= 2
144        let cache_inner = self.inner();
145        let sentinels = [
146            // sentinels only appear at the beginning and/or the end
147            cache_inner.first_key_value().unwrap().0.is_sentinel(),
148            cache_inner.last_key_value().unwrap().0.is_sentinel(),
149        ];
150        len - sentinels.into_iter().filter(|x| *x).count()
151    }
152
153    /// Get the first normal key in the cache, if any.
154    pub fn first_normal_key(&self) -> Option<&StateKey> {
155        self.inner()
156            .iter()
157            .find(|(k, _)| k.is_normal())
158            .map(|(k, _)| k.as_normal_expect())
159    }
160
161    /// Get the last normal key in the cache, if any.
162    pub fn last_normal_key(&self) -> Option<&StateKey> {
163        self.inner()
164            .iter()
165            .rev()
166            .find(|(k, _)| k.is_normal())
167            .map(|(k, _)| k.as_normal_expect())
168    }
169
170    /// Whether the leftmost entry is a sentinel.
171    pub fn left_is_sentinel(&self) -> bool {
172        self.first_key_value()
173            .map(|(k, _)| k.is_sentinel())
174            .unwrap_or(false)
175    }
176
177    /// Whether the rightmost entry is a sentinel.
178    pub fn right_is_sentinel(&self) -> bool {
179        self.last_key_value()
180            .map(|(k, _)| k.is_sentinel())
181            .unwrap_or(false)
182    }
183
184    /// Shrink the partition cache based on the given policy and recently accessed range.
185    pub fn shrink(
186        &mut self,
187        deduped_part_key: &OwnedRow,
188        cache_policy: CachePolicy,
189        recently_accessed_range: RangeInclusive<StateKey>,
190    ) {
191        const MAGIC_CACHE_SIZE: usize = 1024;
192        const MAGIC_JITTER_PREVENTION: usize = MAGIC_CACHE_SIZE / 8;
193
194        // The cache can have zero entry (not even sentinels), e.g., after all rows of the
195        // fully-cached partition are deleted. The `Recent` branch below can't handle this.
196        if self.is_empty() {
197            return;
198        }
199
200        tracing::trace!(
201            partition=?deduped_part_key,
202            cache_policy=?cache_policy,
203            recently_accessed_range=?recently_accessed_range,
204            "find the range to retain in the range cache"
205        );
206
207        let (start, end) = match cache_policy {
208            CachePolicy::Full => {
209                // evict nothing if the policy is to cache full partition
210                return;
211            }
212            CachePolicy::Recent => {
213                let (sk_start, sk_end) = recently_accessed_range.into_inner();
214                let (ck_start, ck_end) = (CacheKey::from(sk_start), CacheKey::from(sk_end));
215
216                // find the cursor just before `ck_start`
217                let mut cursor = self.inner().upper_bound(Bound::Excluded(&ck_start));
218                for _ in 0..MAGIC_JITTER_PREVENTION {
219                    if cursor.prev().is_none() {
220                        // already at the beginning
221                        break;
222                    }
223                }
224                let start = cursor
225                    .peek_prev()
226                    .map(|(k, _)| k)
227                    .unwrap_or_else(|| self.first_key_value().unwrap().0)
228                    .clone();
229
230                // find the cursor just after `ck_end`
231                let mut cursor = self.inner().lower_bound(Bound::Excluded(&ck_end));
232                for _ in 0..MAGIC_JITTER_PREVENTION {
233                    if cursor.next().is_none() {
234                        // already at the end
235                        break;
236                    }
237                }
238                let end = cursor
239                    .peek_next()
240                    .map(|(k, _)| k)
241                    .unwrap_or_else(|| self.last_key_value().unwrap().0)
242                    .clone();
243
244                (start, end)
245            }
246            CachePolicy::RecentFirstN => {
247                if self.len() <= MAGIC_CACHE_SIZE {
248                    // no need to evict if cache len <= N
249                    return;
250                } else {
251                    let (sk_start, _sk_end) = recently_accessed_range.into_inner();
252                    let ck_start = CacheKey::from(sk_start);
253
254                    let mut capacity_remain = MAGIC_CACHE_SIZE; // precision is not important here, code simplicity is the first
255                    const_assert!(MAGIC_JITTER_PREVENTION < MAGIC_CACHE_SIZE);
256
257                    // find the cursor just before `ck_start`
258                    let cursor_just_before_ck_start =
259                        self.inner().upper_bound(Bound::Excluded(&ck_start));
260
261                    let mut cursor = cursor_just_before_ck_start.clone();
262                    // go back for at most `MAGIC_JITTER_PREVENTION` entries
263                    for _ in 0..MAGIC_JITTER_PREVENTION {
264                        if cursor.prev().is_none() {
265                            // already at the beginning
266                            break;
267                        }
268                        capacity_remain -= 1;
269                    }
270                    let start = cursor
271                        .peek_prev()
272                        .map(|(k, _)| k)
273                        .unwrap_or_else(|| self.first_key_value().unwrap().0)
274                        .clone();
275
276                    let mut cursor = cursor_just_before_ck_start;
277                    // go forward for at most `capacity_remain` entries
278                    for _ in 0..capacity_remain {
279                        if cursor.next().is_none() {
280                            // already at the end
281                            break;
282                        }
283                    }
284                    let end = cursor
285                        .peek_next()
286                        .map(|(k, _)| k)
287                        .unwrap_or_else(|| self.last_key_value().unwrap().0)
288                        .clone();
289
290                    (start, end)
291                }
292            }
293            CachePolicy::RecentLastN => {
294                if self.len() <= MAGIC_CACHE_SIZE {
295                    // no need to evict if cache len <= N
296                    return;
297                } else {
298                    let (_sk_start, sk_end) = recently_accessed_range.into_inner();
299                    let ck_end = CacheKey::from(sk_end);
300
301                    let mut capacity_remain = MAGIC_CACHE_SIZE; // precision is not important here, code simplicity is the first
302                    const_assert!(MAGIC_JITTER_PREVENTION < MAGIC_CACHE_SIZE);
303
304                    // find the cursor just after `ck_end`
305                    let cursor_just_after_ck_end =
306                        self.inner().lower_bound(Bound::Excluded(&ck_end));
307
308                    let mut cursor = cursor_just_after_ck_end.clone();
309                    // go forward for at most `MAGIC_JITTER_PREVENTION` entries
310                    for _ in 0..MAGIC_JITTER_PREVENTION {
311                        if cursor.next().is_none() {
312                            // already at the end
313                            break;
314                        }
315                        capacity_remain -= 1;
316                    }
317                    let end = cursor
318                        .peek_next()
319                        .map(|(k, _)| k)
320                        .unwrap_or_else(|| self.last_key_value().unwrap().0)
321                        .clone();
322
323                    let mut cursor = cursor_just_after_ck_end;
324                    // go back for at most `capacity_remain` entries
325                    for _ in 0..capacity_remain {
326                        if cursor.prev().is_none() {
327                            // already at the beginning
328                            break;
329                        }
330                    }
331                    let start = cursor
332                        .peek_prev()
333                        .map(|(k, _)| k)
334                        .unwrap_or_else(|| self.first_key_value().unwrap().0)
335                        .clone();
336
337                    (start, end)
338                }
339            }
340        };
341
342        tracing::trace!(
343            partition=?deduped_part_key,
344            retain_range=?(&start..=&end),
345            "retain range in the range cache"
346        );
347
348        let (left_removed, right_removed) = self.retain_range(&start..=&end);
349        if self.is_empty() {
350            if !left_removed.is_empty() || !right_removed.is_empty() {
351                self.insert(CacheKey::Smallest, OwnedRow::empty());
352                self.insert(CacheKey::Largest, OwnedRow::empty());
353            }
354        } else {
355            if !left_removed.is_empty() {
356                self.insert(CacheKey::Smallest, OwnedRow::empty());
357            }
358            if !right_removed.is_empty() {
359                self.insert(CacheKey::Largest, OwnedRow::empty());
360            }
361        }
362    }
363}
364
365impl EstimateSize for PartitionCache {
366    fn estimated_heap_size(&self) -> usize {
367        self.inner.estimated_heap_size()
368    }
369}
370
371#[cfg(test)]
372mod tests {
373    use risingwave_common::row::OwnedRow;
374    use risingwave_common::types::{DefaultOrdered, ScalarImpl};
375    use risingwave_common::util::memcmp_encoding::encode_value;
376    use risingwave_common::util::sort_util::OrderType;
377    use risingwave_expr::window_function::StateKey;
378
379    use super::*;
380
381    fn create_test_state_key(value: i32) -> StateKey {
382        let row = OwnedRow::new(vec![Some(ScalarImpl::Int32(value))]);
383        StateKey {
384            order_key: encode_value(Some(ScalarImpl::Int32(value)), OrderType::ascending())
385                .unwrap(),
386            pk: DefaultOrdered::new(row),
387        }
388    }
389
390    fn create_test_cache_key(value: i32) -> CacheKey {
391        CacheKey::from(create_test_state_key(value))
392    }
393
394    fn create_test_row(value: i32) -> OwnedRow {
395        OwnedRow::new(vec![Some(ScalarImpl::Int32(value))])
396    }
397
398    #[test]
399    fn test_partition_cache_new() {
400        let cache = PartitionCache::new_without_sentinels();
401        assert!(cache.is_empty());
402        assert_eq!(cache.len(), 0);
403    }
404
405    #[test]
406    fn test_partition_cache_new_with_sentinels() {
407        let cache = PartitionCache::new();
408        assert!(!cache.is_empty());
409        assert_eq!(cache.len(), 2);
410
411        // Should have smallest and largest sentinels
412        let first = cache.first_key_value().unwrap();
413        let last = cache.last_key_value().unwrap();
414
415        assert_eq!(*first.0, CacheKey::Smallest);
416        assert_eq!(*last.0, CacheKey::Largest);
417    }
418
419    #[test]
420    fn test_partition_cache_insert_and_remove() {
421        let mut cache = PartitionCache::new_without_sentinels();
422        let key = create_test_cache_key(1);
423        let value = create_test_row(100);
424
425        // Insert
426        assert!(cache.insert(key.clone(), value.clone()).is_none());
427        assert_eq!(cache.len(), 1);
428        assert!(!cache.is_empty());
429
430        // Remove
431        let removed = cache.remove(&key);
432        assert!(removed.is_some());
433        assert_eq!(removed.unwrap(), value);
434        assert!(cache.is_empty());
435        assert_eq!(cache.len(), 0);
436    }
437
438    #[test]
439    fn test_partition_cache_first_last_key_value() {
440        let mut cache = PartitionCache::new_without_sentinels();
441
442        // Empty cache
443        assert!(cache.first_key_value().is_none());
444        assert!(cache.last_key_value().is_none());
445
446        // Add some entries
447        cache.insert(create_test_cache_key(2), create_test_row(200));
448        cache.insert(create_test_cache_key(1), create_test_row(100));
449        cache.insert(create_test_cache_key(3), create_test_row(300));
450
451        let first = cache.first_key_value().unwrap();
452        let last = cache.last_key_value().unwrap();
453
454        // BTreeMap should order by key
455        assert_eq!(*first.0, create_test_cache_key(1));
456        assert_eq!(*first.1, create_test_row(100));
457
458        assert_eq!(*last.0, create_test_cache_key(3));
459        assert_eq!(*last.1, create_test_row(300));
460    }
461
462    #[test]
463    fn test_partition_cache_retain_range() {
464        let mut cache = PartitionCache::new();
465
466        // Add some entries
467        for i in 1..=5 {
468            cache.insert(create_test_cache_key(i), create_test_row(i * 100));
469        }
470
471        assert_eq!(cache.len(), 7); // 5 normal entries + 2 sentinels
472
473        // Retain range [2, 4]
474        let start = create_test_cache_key(2);
475        let end = create_test_cache_key(4);
476        let (left_removed, right_removed) = cache.retain_range(&start..=&end);
477
478        // Should have removed key 1 on the left and key 5 on the right
479        assert_eq!(left_removed.len(), 1);
480        assert_eq!(right_removed.len(), 1);
481        assert!(left_removed.contains_key(&create_test_cache_key(1)));
482        assert!(right_removed.contains_key(&create_test_cache_key(5)));
483
484        // Cache should now contain keys 2, 3, 4 plus sentinels
485        assert_eq!(cache.len(), 5);
486        for i in 2..=4 {
487            let key = create_test_cache_key(i);
488            assert!(cache.inner.iter().any(|(k, _)| *k == key));
489        }
490    }
491
492    #[test]
493    fn test_partition_cache_shrink_full_policy() {
494        let mut cache = PartitionCache::new();
495
496        // Add many entries
497        for i in 1..=10 {
498            cache.insert(create_test_cache_key(i), create_test_row(i * 100));
499        }
500
501        let initial_len = cache.len();
502        let deduped_part_key = OwnedRow::empty();
503        let recently_accessed_range = create_test_state_key(3)..=create_test_state_key(7);
504
505        // Full policy should not shrink anything
506        cache.shrink(
507            &deduped_part_key,
508            CachePolicy::Full,
509            recently_accessed_range,
510        );
511
512        assert_eq!(cache.len(), initial_len);
513    }
514
515    #[test]
516    fn test_partition_cache_shrink_recent_policy() {
517        let mut cache = PartitionCache::new();
518
519        // Add entries
520        for i in 1..=10 {
521            cache.insert(create_test_cache_key(i), create_test_row(i * 100));
522        }
523
524        let deduped_part_key = OwnedRow::empty();
525        let recently_accessed_range = create_test_state_key(4)..=create_test_state_key(6);
526
527        // Recent policy should keep entries around the accessed range
528        cache.shrink(
529            &deduped_part_key,
530            CachePolicy::Recent,
531            recently_accessed_range,
532        );
533
534        // Cache should still contain the accessed range and some nearby entries
535        let remaining_keys: Vec<_> = cache
536            .inner
537            .iter()
538            .filter_map(|(k, _)| match k {
539                CacheKey::Normal(state_key) => Some(state_key),
540                _ => None,
541            })
542            .collect();
543
544        // Should contain at least the accessed range
545        for i in 4..=6 {
546            let target_key = create_test_state_key(i);
547            assert!(
548                remaining_keys
549                    .iter()
550                    .any(|k| k.order_key == target_key.order_key)
551            );
552        }
553    }
554
555    #[test]
556    fn test_partition_cache_shrink_with_small_cache() {
557        let mut cache = PartitionCache::new();
558
559        // Add only a few entries (less than MAGIC_CACHE_SIZE)
560        for i in 1..=5 {
561            cache.insert(create_test_cache_key(i), create_test_row(i * 100));
562        }
563
564        let initial_len = cache.len();
565        let deduped_part_key = OwnedRow::empty();
566        let recently_accessed_range = create_test_state_key(2)..=create_test_state_key(4);
567
568        // RecentFirstN and RecentLastN should not shrink small caches
569        cache.shrink(
570            &deduped_part_key,
571            CachePolicy::RecentFirstN,
572            recently_accessed_range.clone(),
573        );
574        assert_eq!(cache.len(), initial_len);
575
576        cache.shrink(
577            &deduped_part_key,
578            CachePolicy::RecentLastN,
579            recently_accessed_range,
580        );
581        assert_eq!(cache.len(), initial_len);
582    }
583
584    #[test]
585    fn test_partition_cache_estimate_size() {
586        let cache = PartitionCache::new_without_sentinels();
587        let size_without_sentinels = cache.estimated_heap_size();
588
589        let mut cache = PartitionCache::new();
590        let size_with_sentinels = cache.estimated_heap_size();
591
592        // Size should increase when adding entries
593        assert!(size_with_sentinels >= size_without_sentinels);
594
595        cache.insert(create_test_cache_key(1), create_test_row(100));
596        let size_with_entry = cache.estimated_heap_size();
597
598        assert!(size_with_entry > size_with_sentinels);
599    }
600}
risingwave_stream/executor/over_window/range_cache.rs

risingwave_stream/executor/over_window/
range_cache.rs
