risingwave_expr_impl/window_function/
buffer.rs

1// Copyright 2024 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::VecDeque;
16use std::ops::Range;
17
18use educe::Educe;
19use risingwave_common::array::Op;
20use risingwave_common::types::Sentinelled;
21use risingwave_common::util::memcmp_encoding::{self, MemcmpEncoded};
22use risingwave_expr::window_function::{
23    FrameExclusion, RangeFrameBounds, Readiness, RowsFrameBounds, SessionFrameBounds, StateKey,
24    WatermarkBound,
25};
26
27use super::range_utils::{range_diff, range_except};
28
29/// A common sliding window buffer.
30pub(super) struct WindowBuffer<W: WindowImpl> {
31    window_impl: W,
32    frame_exclusion: FrameExclusion,
33    buffer: VecDeque<Entry<W::Key, W::Value>>,
34    curr_idx: usize,
35    left_idx: usize,       // inclusive, note this can be > `curr_idx`
36    right_excl_idx: usize, // exclusive, note this can be <= `curr_idx` and even <= `left_idx`
37    curr_delta: Option<Vec<(Op, W::Value)>>,
38}
39
40/// A key-value pair in the buffer.
41struct Entry<K: Ord, V> {
42    key: K,
43    value: V,
44}
45
46/// Note: A window frame can be pure preceding, pure following, or acrossing the _current row_.
47pub(super) struct CurrWindow<'a, K> {
48    pub key: Option<&'a K>,
49
50    // XXX(rc): Maybe will be used in the future, let's keep it for now.
51    #[cfg_attr(not(test), expect(dead_code))]
52    /// The preceding half of the current window is saturated.
53    pub preceding_saturated: bool,
54    /// The following half of the current window is saturated.
55    pub following_saturated: bool,
56}
57
58impl<W: WindowImpl> WindowBuffer<W> {
59    pub fn new(window_impl: W, frame_exclusion: FrameExclusion, enable_delta: bool) -> Self {
60        if enable_delta {
61            // TODO(rc): currently only support `FrameExclusion::NoOthers` for delta
62            assert!(frame_exclusion.is_no_others());
63        }
64
65        Self {
66            window_impl,
67            frame_exclusion,
68            buffer: Default::default(),
69            curr_idx: 0,
70            left_idx: 0,
71            right_excl_idx: 0,
72            curr_delta: if enable_delta {
73                Some(Default::default())
74            } else {
75                None
76            },
77        }
78    }
79
80    /// Get the smallest key that is still kept in the buffer.
81    /// Returns `None` if there's nothing yet.
82    pub fn smallest_key(&self) -> Option<&W::Key> {
83        self.buffer.front().map(|Entry { key, .. }| key)
84    }
85
86    /// Get the key part of the current row.
87    pub fn curr_key(&self) -> Option<&W::Key> {
88        self.buffer.get(self.curr_idx).map(|Entry { key, .. }| key)
89    }
90
91    fn buffer_ref(&self) -> BufferRef<'_, W::Key, W::Value> {
92        BufferRef {
93            buffer: &self.buffer,
94            curr_idx: self.curr_idx,
95            left_idx: self.left_idx,
96            right_excl_idx: self.right_excl_idx,
97        }
98    }
99
100    /// Get the current window info.
101    pub fn curr_window(&self) -> CurrWindow<'_, W::Key> {
102        let window = self.buffer_ref();
103        CurrWindow {
104            key: self.curr_key(),
105            preceding_saturated: self.window_impl.preceding_saturated(window),
106            following_saturated: self.window_impl.following_saturated(window),
107        }
108    }
109
110    /// Get the readiness of the current window.
111    pub fn readiness(&self) -> Readiness {
112        let window = self.buffer_ref();
113        if self.curr_key().is_none() {
114            Readiness::WaitForInput
115        } else if self.window_impl.closed_by_input(window) {
116            Readiness::Ready
117        } else if let Some(bound) = self.window_impl.watermark_bound(window) {
118            Readiness::OnWatermark(bound)
119        } else {
120            Readiness::WaitForInput
121        }
122    }
123
124    pub fn is_empty(&self) -> bool {
125        self.buffer.is_empty()
126    }
127
128    fn curr_window_outer(&self) -> Range<usize> {
129        self.left_idx..std::cmp::max(self.right_excl_idx, self.left_idx)
130    }
131
132    fn curr_window_exclusion(&self) -> Range<usize> {
133        // TODO(rc): should intersect with `curr_window_outer` to be more accurate
134        match self.frame_exclusion {
135            FrameExclusion::CurrentRow => self.curr_idx..self.curr_idx + 1,
136            FrameExclusion::NoOthers => self.curr_idx..self.curr_idx,
137        }
138    }
139
140    fn curr_window_ranges(&self) -> (Range<usize>, Range<usize>) {
141        let selection = self.curr_window_outer();
142        let exclusion = self.curr_window_exclusion();
143        range_except(selection, exclusion)
144    }
145
146    /// Iterate over values in the current window.
147    pub fn curr_window_values(&self) -> impl DoubleEndedIterator<Item = &W::Value> {
148        assert!(self.left_idx <= self.buffer.len()); // `left_idx` can be the same as `buffer.len()` when the buffer is empty
149        assert!(self.right_excl_idx <= self.buffer.len());
150
151        let (left, right) = self.curr_window_ranges();
152        self.buffer
153            .range(left)
154            .chain(self.buffer.range(right))
155            .map(|Entry { value, .. }| value)
156    }
157
158    /// Consume the delta of values comparing the current window to the previous window.
159    /// The delta is not guaranteed to be sorted, especially when frame exclusion is not `NoOthers`.
160    pub fn consume_curr_window_values_delta(
161        &mut self,
162    ) -> impl Iterator<Item = (Op, W::Value)> + '_ {
163        self.curr_delta
164            .as_mut()
165            .expect("delta mode should be enabled")
166            .drain(..)
167    }
168
169    /// Append a key-value pair to the buffer.
170    pub fn append(&mut self, key: W::Key, value: W::Value) {
171        let old_outer = self.curr_window_outer();
172
173        self.buffer.push_back(Entry { key, value });
174        self.recalculate_left_right(RecalculateHint::Append);
175
176        if self.curr_delta.is_some() {
177            self.maintain_delta(old_outer, self.curr_window_outer());
178        }
179    }
180
181    /// Slide the current window forward.
182    /// Returns the keys that are removed from the buffer.
183    pub fn slide(&mut self) -> impl Iterator<Item = (W::Key, W::Value)> + '_ {
184        let old_outer = self.curr_window_outer();
185
186        self.curr_idx += 1;
187        self.recalculate_left_right(RecalculateHint::Slide);
188
189        if self.curr_delta.is_some() {
190            self.maintain_delta(old_outer, self.curr_window_outer());
191        }
192
193        let min_needed_idx = [self.left_idx, self.curr_idx, self.right_excl_idx]
194            .iter()
195            .min()
196            .copied()
197            .unwrap();
198        self.curr_idx -= min_needed_idx;
199        self.left_idx -= min_needed_idx;
200        self.right_excl_idx -= min_needed_idx;
201
202        self.window_impl.shift_indices(min_needed_idx);
203
204        self.buffer
205            .drain(0..min_needed_idx)
206            .map(|Entry { key, value }| (key, value))
207    }
208
209    fn maintain_delta(&mut self, old_outer: Range<usize>, new_outer: Range<usize>) {
210        debug_assert!(self.frame_exclusion.is_no_others());
211
212        let (outer_removed, outer_added) = range_diff(old_outer, new_outer);
213        let delta = self.curr_delta.as_mut().unwrap();
214        for idx in outer_removed.iter().cloned().flatten() {
215            delta.push((Op::Delete, self.buffer[idx].value.clone()));
216        }
217        for idx in outer_added.iter().cloned().flatten() {
218            delta.push((Op::Insert, self.buffer[idx].value.clone()));
219        }
220    }
221
222    fn recalculate_left_right(&mut self, hint: RecalculateHint) {
223        let window = BufferRefMut {
224            buffer: &self.buffer,
225            curr_idx: &mut self.curr_idx,
226            left_idx: &mut self.left_idx,
227            right_excl_idx: &mut self.right_excl_idx,
228        };
229        self.window_impl.recalculate_left_right(window, hint);
230    }
231}
232
233/// Wraps a reference to the buffer and some indices, to be used by [`WindowImpl`]s.
234#[derive(Educe)]
235#[educe(Clone, Copy)]
236pub(super) struct BufferRef<'a, K: Ord, V: Clone> {
237    buffer: &'a VecDeque<Entry<K, V>>,
238    curr_idx: usize,
239    left_idx: usize,
240    right_excl_idx: usize,
241}
242
243/// Wraps a reference to the buffer and some mutable indices, to be used by [`WindowImpl`]s.
244pub(super) struct BufferRefMut<'a, K: Ord, V: Clone> {
245    buffer: &'a VecDeque<Entry<K, V>>,
246    curr_idx: &'a mut usize,
247    left_idx: &'a mut usize,
248    right_excl_idx: &'a mut usize,
249}
250
251#[derive(Clone, Copy, PartialEq, Eq)]
252pub(super) enum RecalculateHint {
253    Append,
254    Slide,
255}
256
257/// A trait for sliding window implementations. This trait is used by [`WindowBuffer`] to
258/// determine the status of current window and how to slide the window.
259pub(super) trait WindowImpl {
260    type Key: Ord;
261    type Value: Clone;
262
263    /// Whether the preceding half of the current window is saturated.
264    /// By "saturated" we mean that every row that is possible to be in the preceding half of the
265    /// current window is already in the buffer.
266    fn preceding_saturated(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool;
267
268    /// Whether the following half of the current window is saturated.
269    fn following_saturated(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool;
270
271    /// Whether the current row exists and its window is closed by the rows in the buffer, given
272    /// that no future row is less than the last one in the order of the order key.
273    fn closed_by_input(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool {
274        self.following_saturated(window)
275    }
276
277    /// The watermark of the order key that closes the current window, called only when the
278    /// current row exists and its window is not closed by input. `None` means only future rows
279    /// can close it, which is also the case for a non-ascending order, since a watermark only
280    /// bounds future order values from below.
281    fn watermark_bound(
282        &self,
283        _window: BufferRef<'_, Self::Key, Self::Value>,
284    ) -> Option<WatermarkBound> {
285        None
286    }
287
288    /// Recalculate the left and right indices of the current window, according to the latest
289    /// `curr_idx` and the hint.
290    fn recalculate_left_right(
291        &mut self,
292        window: BufferRefMut<'_, Self::Key, Self::Value>,
293        hint: RecalculateHint,
294    );
295
296    /// Shift the indices stored by the [`WindowImpl`] by `n` positions. This should be called
297    /// after evicting rows from the buffer.
298    fn shift_indices(&mut self, n: usize);
299}
300
301/// The sliding window implementation for `ROWS` frames.
302pub(super) struct RowsWindow<K: Ord, V: Clone> {
303    frame_bounds: RowsFrameBounds,
304    _phantom: std::marker::PhantomData<K>,
305    _phantom2: std::marker::PhantomData<V>,
306}
307
308impl<K: Ord, V: Clone> RowsWindow<K, V> {
309    pub fn new(frame_bounds: RowsFrameBounds) -> Self {
310        Self {
311            frame_bounds,
312            _phantom: std::marker::PhantomData,
313            _phantom2: std::marker::PhantomData,
314        }
315    }
316}
317
318impl<K: Ord, V: Clone> WindowImpl for RowsWindow<K, V> {
319    type Key = K;
320    type Value = V;
321
322    fn preceding_saturated(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool {
323        window.curr_idx < window.buffer.len() && {
324            let start_off = self.frame_bounds.start.to_offset();
325            if let Some(start_off) = start_off {
326                if start_off >= 0 {
327                    true // pure following frame, always preceding-saturated
328                } else {
329                    // FIXME(rc): Clippy rule `clippy::nonminimal_bool` is misreporting that
330                    // the following can be simplified.
331
332                    {
333                        assert!(window.curr_idx >= window.left_idx);
334                    }
335                    window.curr_idx - window.left_idx >= start_off.unsigned_abs()
336                }
337            } else {
338                false // unbounded frame start, never preceding-saturated
339            }
340        }
341    }
342
343    fn following_saturated(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool {
344        window.curr_idx < window.buffer.len() && {
345            let end_off = self.frame_bounds.end.to_offset();
346            if let Some(end_off) = end_off {
347                if end_off <= 0 {
348                    true // pure preceding frame, always following-saturated
349                } else {
350                    // FIXME(rc): Ditto.
351
352                    {
353                        assert!(window.right_excl_idx > 0);
354                        assert!(window.right_excl_idx > window.curr_idx);
355                        assert!(window.right_excl_idx <= window.buffer.len());
356                    }
357                    window.right_excl_idx - 1 - window.curr_idx >= end_off as usize
358                }
359            } else {
360                false // unbounded frame end, never following-saturated
361            }
362        }
363    }
364
365    fn recalculate_left_right(
366        &mut self,
367        window: BufferRefMut<'_, Self::Key, Self::Value>,
368        _hint: RecalculateHint,
369    ) {
370        if window.buffer.is_empty() {
371            *window.left_idx = 0;
372            *window.right_excl_idx = 0;
373        }
374
375        let start_off = self.frame_bounds.start.to_offset();
376        let end_off = self.frame_bounds.end.to_offset();
377        if let Some(start_off) = start_off {
378            let logical_left_idx = *window.curr_idx as isize + start_off;
379            if logical_left_idx >= 0 {
380                *window.left_idx = std::cmp::min(logical_left_idx as usize, window.buffer.len());
381            } else {
382                *window.left_idx = 0;
383            }
384        } else {
385            // unbounded start
386            *window.left_idx = 0;
387        }
388        if let Some(end_off) = end_off {
389            let logical_right_excl_idx = *window.curr_idx as isize + end_off + 1;
390            if logical_right_excl_idx >= 0 {
391                *window.right_excl_idx =
392                    std::cmp::min(logical_right_excl_idx as usize, window.buffer.len());
393            } else {
394                *window.right_excl_idx = 0;
395            }
396        } else {
397            // unbounded end
398            *window.right_excl_idx = window.buffer.len();
399        }
400    }
401
402    fn shift_indices(&mut self, _n: usize) {}
403}
404
405/// The sliding window implementation for `RANGE` frames.
406pub(super) struct RangeWindow<V: Clone> {
407    frame_bounds: RangeFrameBounds,
408    /// The memcmp-encoded frame end of the current row, `None` if unbounded. It's maintained by
409    /// `recalculate_left_right`, so it's valid while the current row exists.
410    curr_frame_end: Option<MemcmpEncoded>,
411    _phantom: std::marker::PhantomData<V>,
412}
413
414impl<V: Clone> RangeWindow<V> {
415    pub fn new(frame_bounds: RangeFrameBounds) -> Self {
416        Self {
417            frame_bounds,
418            curr_frame_end: None,
419            _phantom: std::marker::PhantomData,
420        }
421    }
422}
423
424impl<V: Clone> WindowImpl for RangeWindow<V> {
425    type Key = StateKey;
426    type Value = V;
427
428    fn preceding_saturated(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool {
429        window.curr_idx < window.buffer.len() && {
430            // XXX(rc): It seems that preceding saturation is not important, may remove later.
431            true
432        }
433    }
434
435    fn following_saturated(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool {
436        window.curr_idx < window.buffer.len()
437            && {
438                // Left OK? (note that `left_idx` can be greater than `right_idx`)
439                // The following line checks whether the left value is the last one in the buffer.
440                // Here we adopt a conservative approach, which means we assume the next future value
441                // is likely to be the same as the last value in the current window, in which case
442                // we can't say the current window is saturated.
443                window.left_idx < window.buffer.len() /* non-zero */ - 1
444            }
445            && {
446                // Right OK? Ditto.
447                window.right_excl_idx < window.buffer.len()
448            }
449    }
450
451    fn closed_by_input(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool {
452        // A row after the frame end is enough, no matter whether there are rows in the frame.
453        window.curr_idx < window.buffer.len() && window.right_excl_idx < window.buffer.len()
454    }
455
456    fn watermark_bound(
457        &self,
458        _window: BufferRef<'_, Self::Key, Self::Value>,
459    ) -> Option<WatermarkBound> {
460        if !self.frame_bounds.order_type.is_ascending() {
461            return None;
462        }
463        // Future rows equal to the frame end still fall into the frame.
464        self.curr_frame_end.clone().map(|value| WatermarkBound {
465            value,
466            exclusive: true,
467        })
468    }
469
470    fn recalculate_left_right(
471        &mut self,
472        window: BufferRefMut<'_, Self::Key, Self::Value>,
473        _hint: RecalculateHint,
474    ) {
475        if window.buffer.is_empty() {
476            *window.left_idx = 0;
477            *window.right_excl_idx = 0;
478        }
479
480        let Some(entry) = window.buffer.get(*window.curr_idx) else {
481            // If the current index has been moved to a future position, we can't touch anything
482            // because the next coming key may equal to the previous one which means the left and
483            // right indices will be the same.
484            return;
485        };
486        let curr_key = &entry.key;
487
488        let curr_order_value = memcmp_encoding::decode_value(
489            &self.frame_bounds.order_data_type,
490            &curr_key.order_key,
491            self.frame_bounds.order_type,
492        )
493        .expect("no reason to fail here because we just encoded it in memory");
494
495        match self.frame_bounds.frame_start_of(&curr_order_value) {
496            Sentinelled::Smallest => {
497                // unbounded frame start
498                assert_eq!(
499                    *window.left_idx, 0,
500                    "for unbounded start, left index should always be 0"
501                );
502            }
503            Sentinelled::Normal(value) => {
504                // bounded, find the start position
505                let value_enc = memcmp_encoding::encode_value(value, self.frame_bounds.order_type)
506                    .expect("no reason to fail here");
507                *window.left_idx = window
508                    .buffer
509                    .partition_point(|elem| elem.key.order_key < value_enc);
510            }
511            Sentinelled::Largest => unreachable!("frame start never be UNBOUNDED FOLLOWING"),
512        }
513
514        match self.frame_bounds.frame_end_of(curr_order_value) {
515            Sentinelled::Largest => {
516                // unbounded frame end
517                *window.right_excl_idx = window.buffer.len();
518                self.curr_frame_end = None;
519            }
520            Sentinelled::Normal(value) => {
521                // bounded, find the end position
522                let value_enc = memcmp_encoding::encode_value(value, self.frame_bounds.order_type)
523                    .expect("no reason to fail here");
524                *window.right_excl_idx = window
525                    .buffer
526                    .partition_point(|elem| elem.key.order_key <= value_enc);
527                self.curr_frame_end = Some(value_enc);
528            }
529            Sentinelled::Smallest => unreachable!("frame end never be UNBOUNDED PRECEDING"),
530        }
531    }
532
533    fn shift_indices(&mut self, _n: usize) {}
534}
535
536pub(super) struct SessionWindow<V: Clone> {
537    frame_bounds: SessionFrameBounds,
538    /// The latest session is the rightmost session in the buffer, which is updated during appending.
539    latest_session: Option<LatestSession>,
540    /// The sizes of recognized but not consumed sessions in the buffer. It's updated during appending.
541    /// The first element, if any, should be the size of the "current session window". When sliding,
542    /// the front should be popped.
543    recognized_session_sizes: VecDeque<usize>,
544    _phantom: std::marker::PhantomData<V>,
545}
546
547#[derive(Debug)]
548struct LatestSession {
549    /// The starting index of the latest session.
550    start_idx: usize,
551
552    /// Minimal next start means the minimal order value that can start a new session.
553    /// If a row has an order value less than this, it should be in the current session.
554    minimal_next_start: MemcmpEncoded,
555}
556
557impl<V: Clone> SessionWindow<V> {
558    pub fn new(frame_bounds: SessionFrameBounds) -> Self {
559        Self {
560            frame_bounds,
561            latest_session: None,
562            recognized_session_sizes: Default::default(),
563            _phantom: std::marker::PhantomData,
564        }
565    }
566}
567
568impl<V: Clone> WindowImpl for SessionWindow<V> {
569    type Key = StateKey;
570    type Value = V;
571
572    fn preceding_saturated(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool {
573        window.curr_idx < window.buffer.len() && {
574            // XXX(rc): It seems that preceding saturation is not important, may remove later.
575            true
576        }
577    }
578
579    fn following_saturated(&self, window: BufferRef<'_, Self::Key, Self::Value>) -> bool {
580        window.curr_idx < window.buffer.len() && {
581            // For session window, `left_idx` is always smaller than `right_excl_idx`.
582            assert!(window.left_idx <= window.curr_idx);
583            assert!(window.curr_idx < window.right_excl_idx);
584
585            // The following expression checks whether the current window is the latest session.
586            // If it is, we don't think it's saturated because the next row may be still in the
587            // same session. Otherwise, we can safely say it's saturated.
588            self.latest_session
589                .as_ref()
590                .is_some_and(|LatestSession { start_idx, .. }| window.left_idx < *start_idx)
591        }
592    }
593
594    fn watermark_bound(
595        &self,
596        _window: BufferRef<'_, Self::Key, Self::Value>,
597    ) -> Option<WatermarkBound> {
598        if !self.frame_bounds.order_type.is_ascending() {
599            return None;
600        }
601        // The current window is the latest session, which is closed once any future row would
602        // start a new session.
603        self.latest_session.as_ref().map(
604            |LatestSession {
605                 minimal_next_start, ..
606             }| WatermarkBound {
607                value: minimal_next_start.clone(),
608                exclusive: false,
609            },
610        )
611    }
612
613    fn recalculate_left_right(
614        &mut self,
615        window: BufferRefMut<'_, Self::Key, Self::Value>,
616        hint: RecalculateHint,
617    ) {
618        // Terms:
619        // - Session: A continuous range of rows among any two of which the difference of order values
620        //   is less than the session gap. This is a concept on the whole stream. Sessions are recognized
621        //   during appending.
622        // - Current window: The range of rows that are represented by the indices in `window`. It is a
623        //   status of the `WindowBuffer`. If the current window happens to be the last session in the
624        //   buffer, it will be updated during appending. Otherwise it will only be updated during sliding.
625
626        match hint {
627            RecalculateHint::Append => {
628                assert!(!window.buffer.is_empty()); // because we just appended a row
629                let appended_idx = window.buffer.len() - 1;
630                let appended_key = &window.buffer[appended_idx].key;
631
632                let minimal_next_start_of_appended = self.frame_bounds.minimal_next_start_of(
633                    memcmp_encoding::decode_value(
634                        &self.frame_bounds.order_data_type,
635                        &appended_key.order_key,
636                        self.frame_bounds.order_type,
637                    )
638                    .expect("no reason to fail here because we just encoded it in memory"),
639                );
640                let minimal_next_start_enc_of_appended = memcmp_encoding::encode_value(
641                    minimal_next_start_of_appended,
642                    self.frame_bounds.order_type,
643                )
644                .expect("no reason to fail here");
645
646                if let Some(&mut LatestSession {
647                    ref start_idx,
648                    ref mut minimal_next_start,
649                }) = self.latest_session.as_mut()
650                {
651                    if &appended_key.order_key >= minimal_next_start {
652                        // the appended row starts a new session
653                        self.recognized_session_sizes
654                            .push_back(appended_idx - start_idx);
655                        self.latest_session = Some(LatestSession {
656                            start_idx: appended_idx,
657                            minimal_next_start: minimal_next_start_enc_of_appended,
658                        });
659                        // no need to update the current window because it's now corresponding
660                        // to some previous session
661                    } else {
662                        // the appended row belongs to the latest session
663                        *minimal_next_start = minimal_next_start_enc_of_appended;
664
665                        if *start_idx == *window.left_idx {
666                            // the current window is the latest session, we should extend it
667                            *window.right_excl_idx = appended_idx + 1;
668                        }
669                    }
670                } else {
671                    // no session yet, the current window should be empty
672                    let left_idx = *window.left_idx;
673                    let curr_idx = *window.curr_idx;
674                    let old_right_excl_idx = *window.right_excl_idx;
675                    assert_eq!(left_idx, curr_idx);
676                    assert_eq!(left_idx, old_right_excl_idx);
677                    assert_eq!(old_right_excl_idx, window.buffer.len() - 1);
678
679                    // now we put the first row into the current window
680                    *window.right_excl_idx = window.buffer.len();
681
682                    // and start to recognize the latest session
683                    self.latest_session = Some(LatestSession {
684                        start_idx: left_idx,
685                        minimal_next_start: minimal_next_start_enc_of_appended,
686                    });
687                }
688            }
689            RecalculateHint::Slide => {
690                let old_left_idx = *window.left_idx;
691                let new_curr_idx = *window.curr_idx;
692                let old_right_excl_idx = *window.right_excl_idx;
693
694                if new_curr_idx < old_right_excl_idx {
695                    // the current row is still in the current session window, no need to slide
696                } else {
697                    let old_session_size = self.recognized_session_sizes.pop_front();
698                    let next_session_size = self.recognized_session_sizes.front().copied();
699
700                    if let Some(old_session_size) = old_session_size {
701                        assert_eq!(old_session_size, old_right_excl_idx - old_left_idx);
702
703                        // slide the window to the next session
704                        if let Some(next_session_size) = next_session_size {
705                            // the next session is fully recognized, so we know the ending index
706                            *window.left_idx = old_right_excl_idx;
707                            *window.right_excl_idx = old_right_excl_idx + next_session_size;
708                        } else {
709                            // the next session is still in recognition, so we end the window at the end of buffer
710                            *window.left_idx = old_right_excl_idx;
711                            *window.right_excl_idx = window.buffer.len();
712                        }
713                    } else {
714                        // no recognized session yet, meaning the current window is the last session in the buffer
715                        assert_eq!(old_right_excl_idx, window.buffer.len());
716                        *window.left_idx = old_right_excl_idx;
717                        *window.right_excl_idx = old_right_excl_idx;
718                        self.latest_session = None;
719                    }
720                }
721            }
722        }
723    }
724
725    fn shift_indices(&mut self, n: usize) {
726        if let Some(LatestSession { start_idx, .. }) = self.latest_session.as_mut() {
727            *start_idx -= n;
728        }
729    }
730}
731
732#[cfg(test)]
733mod tests {
734    use itertools::Itertools;
735    use risingwave_common::row::OwnedRow;
736    use risingwave_common::types::{DataType, ScalarImpl};
737    use risingwave_common::util::sort_util::OrderType;
738    use risingwave_expr::window_function::FrameBound::{
739        CurrentRow, Following, Preceding, UnboundedFollowing, UnboundedPreceding,
740    };
741    use risingwave_expr::window_function::{FrameBound, RangeFrameOffset, SessionFrameGap};
742
743    use super::*;
744
745    #[test]
746    fn test_rows_frame_unbounded_preceding_to_current_row() {
747        let mut buffer = WindowBuffer::<RowsWindow<_, _>>::new(
748            RowsWindow::new(RowsFrameBounds {
749                start: UnboundedPreceding,
750                end: CurrentRow,
751            }),
752            FrameExclusion::NoOthers,
753            true,
754        );
755
756        let window = buffer.curr_window();
757        assert!(window.key.is_none());
758        assert!(!window.preceding_saturated);
759        assert!(!window.following_saturated);
760        buffer.append(1, "hello");
761        let window = buffer.curr_window();
762        assert_eq!(window.key, Some(&1));
763        assert!(!window.preceding_saturated); // unbounded preceding is never saturated
764        assert!(window.following_saturated);
765        buffer.append(2, "world");
766        let window = buffer.curr_window();
767        assert_eq!(window.key, Some(&1));
768        assert!(!window.preceding_saturated);
769        assert!(window.following_saturated);
770        assert_eq!(
771            buffer.curr_window_values().cloned().collect_vec(),
772            vec!["hello"]
773        );
774        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
775        assert!(removed_keys.is_empty()); // unbouded preceding, nothing can ever be removed
776        let window = buffer.curr_window();
777        assert_eq!(window.key, Some(&2));
778        assert!(!window.preceding_saturated);
779        assert!(window.following_saturated);
780        assert_eq!(buffer.smallest_key(), Some(&1));
781    }
782
783    #[test]
784    fn test_rows_frame_preceding_to_current_row() {
785        let mut buffer = WindowBuffer::<RowsWindow<_, _>>::new(
786            RowsWindow::new(RowsFrameBounds {
787                start: Preceding(1),
788                end: CurrentRow,
789            }),
790            FrameExclusion::NoOthers,
791            true,
792        );
793
794        let window = buffer.curr_window();
795        assert!(window.key.is_none());
796        assert!(!window.preceding_saturated);
797        assert!(!window.following_saturated);
798        buffer.append(1, "hello");
799        let window = buffer.curr_window();
800        assert_eq!(window.key, Some(&1));
801        assert!(!window.preceding_saturated);
802        assert!(window.following_saturated);
803        assert_eq!(
804            buffer.curr_window_values().cloned().collect_vec(),
805            vec!["hello"]
806        );
807        buffer.append(2, "world");
808        let window = buffer.curr_window();
809        assert_eq!(window.key, Some(&1));
810        assert!(!window.preceding_saturated);
811        assert!(window.following_saturated);
812        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
813        assert!(removed_keys.is_empty());
814        let window = buffer.curr_window();
815        assert_eq!(window.key, Some(&2));
816        assert!(window.preceding_saturated);
817        assert!(window.following_saturated);
818        assert_eq!(buffer.smallest_key(), Some(&1));
819        buffer.append(3, "!");
820        let window = buffer.curr_window();
821        assert_eq!(window.key, Some(&2));
822        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
823        assert_eq!(removed_keys, vec![1]);
824    }
825
826    #[test]
827    fn test_rows_frame_preceding_to_preceding() {
828        let mut buffer = WindowBuffer::<RowsWindow<_, _>>::new(
829            RowsWindow::new(RowsFrameBounds {
830                start: Preceding(2),
831                end: Preceding(1),
832            }),
833            FrameExclusion::NoOthers,
834            true,
835        );
836
837        buffer.append(1, "RisingWave");
838        let window = buffer.curr_window();
839        assert_eq!(window.key, Some(&1));
840        assert!(!window.preceding_saturated);
841        assert!(window.following_saturated);
842        assert!(buffer.curr_window_values().collect_vec().is_empty());
843        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
844        assert!(removed_keys.is_empty());
845        assert_eq!(buffer.smallest_key(), Some(&1));
846        buffer.append(2, "is the best");
847        let window = buffer.curr_window();
848        assert_eq!(window.key, Some(&2));
849        assert!(!window.preceding_saturated);
850        assert!(window.following_saturated);
851        assert_eq!(
852            buffer.curr_window_values().cloned().collect_vec(),
853            vec!["RisingWave"]
854        );
855        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
856        assert!(removed_keys.is_empty());
857        assert_eq!(buffer.smallest_key(), Some(&1));
858        buffer.append(3, "streaming platform");
859        let window = buffer.curr_window();
860        assert_eq!(window.key, Some(&3));
861        assert!(window.preceding_saturated);
862        assert!(window.following_saturated);
863        assert_eq!(
864            buffer.curr_window_values().cloned().collect_vec(),
865            vec!["RisingWave", "is the best"]
866        );
867        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
868        assert_eq!(removed_keys, vec![1]);
869        assert_eq!(buffer.smallest_key(), Some(&2));
870    }
871
872    #[test]
873    fn test_rows_frame_current_row_to_unbounded_following() {
874        let mut buffer = WindowBuffer::<RowsWindow<_, _>>::new(
875            RowsWindow::new(RowsFrameBounds {
876                start: CurrentRow,
877                end: UnboundedFollowing,
878            }),
879            FrameExclusion::NoOthers,
880            true,
881        );
882
883        buffer.append(1, "RisingWave");
884        let window = buffer.curr_window();
885        assert_eq!(window.key, Some(&1));
886        assert!(window.preceding_saturated);
887        assert!(!window.following_saturated);
888        assert_eq!(
889            buffer.curr_window_values().cloned().collect_vec(),
890            vec!["RisingWave"]
891        );
892        buffer.append(2, "is the best");
893        let window = buffer.curr_window();
894        assert_eq!(window.key, Some(&1));
895        assert!(window.preceding_saturated);
896        assert!(!window.following_saturated);
897        assert_eq!(
898            buffer.curr_window_values().cloned().collect_vec(),
899            vec!["RisingWave", "is the best"]
900        );
901        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
902        assert_eq!(removed_keys, vec![1]);
903        assert_eq!(buffer.smallest_key(), Some(&2));
904        let window = buffer.curr_window();
905        assert_eq!(window.key, Some(&2));
906        assert!(window.preceding_saturated);
907        assert!(!window.following_saturated);
908        assert_eq!(
909            buffer.curr_window_values().cloned().collect_vec(),
910            vec!["is the best"]
911        );
912    }
913
914    #[test]
915    fn test_rows_frame_current_row_to_following() {
916        let mut buffer = WindowBuffer::<RowsWindow<_, _>>::new(
917            RowsWindow::new(RowsFrameBounds {
918                start: CurrentRow,
919                end: Following(1),
920            }),
921            FrameExclusion::NoOthers,
922            true,
923        );
924
925        buffer.append(1, "RisingWave");
926        let window = buffer.curr_window();
927        assert_eq!(window.key, Some(&1));
928        assert!(window.preceding_saturated);
929        assert!(!window.following_saturated);
930        assert_eq!(
931            buffer.curr_window_values().cloned().collect_vec(),
932            vec!["RisingWave"]
933        );
934        buffer.append(2, "is the best");
935        let window = buffer.curr_window();
936        assert_eq!(window.key, Some(&1));
937        assert!(window.preceding_saturated);
938        assert!(window.following_saturated);
939        assert_eq!(
940            buffer.curr_window_values().cloned().collect_vec(),
941            vec!["RisingWave", "is the best"]
942        );
943        buffer.append(3, "streaming platform");
944        let window = buffer.curr_window();
945        assert_eq!(window.key, Some(&1));
946        assert!(window.preceding_saturated);
947        assert!(window.following_saturated);
948        assert_eq!(
949            buffer.curr_window_values().cloned().collect_vec(),
950            vec!["RisingWave", "is the best"]
951        );
952        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
953        assert_eq!(removed_keys, vec![1]);
954        let window = buffer.curr_window();
955        assert_eq!(window.key, Some(&2));
956        assert!(window.preceding_saturated);
957        assert!(window.following_saturated);
958        assert_eq!(
959            buffer.curr_window_values().cloned().collect_vec(),
960            vec!["is the best", "streaming platform"]
961        );
962    }
963
964    #[test]
965    fn test_rows_frame_following_to_following() {
966        let mut buffer = WindowBuffer::<RowsWindow<_, _>>::new(
967            RowsWindow::new(RowsFrameBounds {
968                start: Following(1),
969                end: Following(2),
970            }),
971            FrameExclusion::NoOthers,
972            true,
973        );
974
975        buffer.append(1, "RisingWave");
976        let window = buffer.curr_window();
977        assert_eq!(window.key, Some(&1));
978        assert!(window.preceding_saturated);
979        assert!(!window.following_saturated);
980        assert!(buffer.curr_window_values().collect_vec().is_empty());
981        buffer.append(2, "is the best");
982        let window = buffer.curr_window();
983        assert_eq!(window.key, Some(&1));
984        assert!(window.preceding_saturated);
985        assert!(!window.following_saturated);
986        assert_eq!(
987            buffer.curr_window_values().cloned().collect_vec(),
988            vec!["is the best"]
989        );
990        buffer.append(3, "streaming platform");
991        let window = buffer.curr_window();
992        assert_eq!(window.key, Some(&1));
993        assert!(window.preceding_saturated);
994        assert!(window.following_saturated);
995        assert_eq!(
996            buffer.curr_window_values().cloned().collect_vec(),
997            vec!["is the best", "streaming platform"]
998        );
999        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
1000        assert_eq!(removed_keys, vec![1]);
1001        let window = buffer.curr_window();
1002        assert_eq!(window.key, Some(&2));
1003        assert!(window.preceding_saturated);
1004        assert!(!window.following_saturated);
1005        assert_eq!(
1006            buffer.curr_window_values().cloned().collect_vec(),
1007            vec!["streaming platform"]
1008        );
1009    }
1010
1011    #[test]
1012    fn test_rows_frame_exclude_current_row() {
1013        let mut buffer = WindowBuffer::<RowsWindow<_, _>>::new(
1014            RowsWindow::new(RowsFrameBounds {
1015                start: UnboundedPreceding,
1016                end: CurrentRow,
1017            }),
1018            FrameExclusion::CurrentRow,
1019            false,
1020        );
1021
1022        buffer.append(1, "hello");
1023        assert!(
1024            buffer
1025                .curr_window_values()
1026                .cloned()
1027                .collect_vec()
1028                .is_empty()
1029        );
1030        buffer.append(2, "world");
1031        let _ = buffer.slide();
1032        assert_eq!(
1033            buffer.curr_window_values().cloned().collect_vec(),
1034            vec!["hello"]
1035        );
1036    }
1037
1038    #[test]
1039    fn test_session_frame() {
1040        let order_data_type = DataType::Int64;
1041        let order_type = OrderType::ascending();
1042        let gap_data_type = DataType::Int64;
1043
1044        let mut buffer = WindowBuffer::<SessionWindow<_>>::new(
1045            SessionWindow::new(SessionFrameBounds {
1046                order_data_type: order_data_type.clone(),
1047                order_type,
1048                gap_data_type: gap_data_type.clone(),
1049                gap: SessionFrameGap::new_for_test(
1050                    ScalarImpl::Int64(5),
1051                    &order_data_type,
1052                    &gap_data_type,
1053                ),
1054            }),
1055            FrameExclusion::NoOthers,
1056            true,
1057        );
1058
1059        let key = |key: i64| -> StateKey {
1060            StateKey {
1061                order_key: memcmp_encoding::encode_value(Some(ScalarImpl::from(key)), order_type)
1062                    .unwrap(),
1063                pk: OwnedRow::empty().into(),
1064            }
1065        };
1066
1067        assert!(buffer.curr_key().is_none());
1068
1069        buffer.append(key(1), "hello");
1070        buffer.append(key(3), "session");
1071        let window = buffer.curr_window();
1072        assert_eq!(window.key, Some(&key(1)));
1073        assert!(window.preceding_saturated);
1074        assert!(!window.following_saturated);
1075        assert_eq!(
1076            buffer.curr_window_values().cloned().collect_vec(),
1077            vec!["hello", "session"]
1078        );
1079
1080        buffer.append(key(8), "window"); // start a new session
1081        let window = buffer.curr_window();
1082        assert!(window.following_saturated);
1083        assert_eq!(
1084            buffer.curr_window_values().cloned().collect_vec(),
1085            vec!["hello", "session"]
1086        );
1087
1088        buffer.append(key(15), "and");
1089        buffer.append(key(16), "world");
1090        assert_eq!(
1091            buffer.curr_window_values().cloned().collect_vec(),
1092            vec!["hello", "session"]
1093        );
1094
1095        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
1096        assert!(removed_keys.is_empty());
1097        let window = buffer.curr_window();
1098        assert_eq!(window.key, Some(&key(3)));
1099        assert!(window.preceding_saturated);
1100        assert!(window.following_saturated);
1101        assert_eq!(
1102            buffer.curr_window_values().cloned().collect_vec(),
1103            vec!["hello", "session"]
1104        );
1105
1106        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
1107        assert_eq!(removed_keys, vec![key(1), key(3)]);
1108        assert_eq!(buffer.smallest_key(), Some(&key(8)));
1109        let window = buffer.curr_window();
1110        assert_eq!(window.key, Some(&key(8)));
1111        assert!(window.preceding_saturated);
1112        assert!(window.following_saturated);
1113        assert_eq!(
1114            buffer.curr_window_values().cloned().collect_vec(),
1115            vec!["window"]
1116        );
1117
1118        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
1119        assert_eq!(removed_keys, vec![key(8)]);
1120        assert_eq!(buffer.smallest_key(), Some(&key(15)));
1121        let window = buffer.curr_window();
1122        assert_eq!(window.key, Some(&key(15)));
1123        assert!(window.preceding_saturated);
1124        assert!(!window.following_saturated);
1125        assert_eq!(
1126            buffer.curr_window_values().cloned().collect_vec(),
1127            vec!["and", "world"]
1128        );
1129
1130        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
1131        assert!(removed_keys.is_empty());
1132        assert_eq!(buffer.curr_key(), Some(&key(16)));
1133        assert_eq!(
1134            buffer.curr_window_values().cloned().collect_vec(),
1135            vec!["and", "world"]
1136        );
1137
1138        let removed_keys = buffer.slide().map(|(k, _)| k).collect_vec();
1139        assert_eq!(removed_keys, vec![key(15), key(16)]);
1140        assert!(buffer.curr_key().is_none());
1141        assert!(
1142            buffer
1143                .curr_window_values()
1144                .cloned()
1145                .collect_vec()
1146                .is_empty()
1147        );
1148    }
1149
1150    fn key_of(order: ScalarImpl) -> StateKey {
1151        StateKey {
1152            order_key: memcmp_encoding::encode_value(Some(order), OrderType::ascending()).unwrap(),
1153            pk: OwnedRow::empty().into(),
1154        }
1155    }
1156
1157    fn key(order: i64) -> StateKey {
1158        key_of(order.into())
1159    }
1160
1161    fn range_buffer_of(
1162        data_type: DataType,
1163        start: FrameBound<ScalarImpl>,
1164        end: FrameBound<ScalarImpl>,
1165    ) -> WindowBuffer<RangeWindow<&'static str>> {
1166        let offset = |bound: FrameBound<ScalarImpl>| {
1167            bound.map(|offset| RangeFrameOffset::new_for_test(offset, &data_type, &data_type))
1168        };
1169        WindowBuffer::new(
1170            RangeWindow::new(RangeFrameBounds {
1171                order_data_type: data_type.clone(),
1172                order_type: OrderType::ascending(),
1173                offset_data_type: data_type.clone(),
1174                start: offset(start),
1175                end: offset(end),
1176            }),
1177            FrameExclusion::NoOthers,
1178            true,
1179        )
1180    }
1181
1182    fn range_buffer(
1183        start: FrameBound<i64>,
1184        end: FrameBound<i64>,
1185    ) -> WindowBuffer<RangeWindow<&'static str>> {
1186        range_buffer_of(
1187            DataType::Int64,
1188            start.map(ScalarImpl::from),
1189            end.map(ScalarImpl::from),
1190        )
1191    }
1192
1193    #[test]
1194    fn test_range_frame_readiness() {
1195        // The frame ending before the current row is closed once the row arrives, while the
1196        // conservative `following_saturated` waits for the next row.
1197        let mut buffer = range_buffer(Preceding(10), Preceding(5));
1198        buffer.append(key(102), "a");
1199        assert_eq!(buffer.readiness(), Readiness::Ready);
1200        assert!(!buffer.curr_window().following_saturated);
1201
1202        let mut buffer = range_buffer(CurrentRow, UnboundedFollowing);
1203        buffer.append(key(1), "a");
1204        buffer.append(key(2), "b");
1205        assert_eq!(buffer.readiness(), Readiness::WaitForInput);
1206
1207        // The bound follows the current row as the window slides.
1208        let bound = |order: i64| {
1209            Readiness::OnWatermark(WatermarkBound {
1210                value: key(order).order_key,
1211                exclusive: true,
1212            })
1213        };
1214        let mut buffer = range_buffer(CurrentRow, Following(2));
1215        buffer.append(key(1), "a");
1216        buffer.append(key(2), "b");
1217        assert_eq!(buffer.readiness(), bound(3));
1218        buffer.append(key(4), "c");
1219        assert_eq!(buffer.readiness(), Readiness::Ready);
1220        let _ = buffer.slide();
1221        assert_eq!(buffer.readiness(), bound(4));
1222
1223        // A float frame end may round back to the current row, `2^54 - 1.0 == 2^54` in `float8`,
1224        // so the row stays in its own frame until the watermark passes it.
1225        let x = ScalarImpl::from(2f64.powi(54));
1226        let one = ScalarImpl::from(1.0f64);
1227        let mut buffer = range_buffer_of(DataType::Float64, Preceding(one.clone()), Preceding(one));
1228        buffer.append(key_of(x.clone()), "a");
1229        assert_eq!(
1230            buffer.readiness(),
1231            Readiness::OnWatermark(WatermarkBound {
1232                value: key_of(x).order_key,
1233                exclusive: true,
1234            })
1235        );
1236    }
1237}
risingwave_expr_impl/window_function/buffer.rs

risingwave_expr_impl/window_function/
buffer.rs
