risingwave_stream/executor/eowc/
sort_buffer.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::marker::PhantomData;
16
17use futures_async_stream::try_stream;
18use risingwave_common::array::StreamChunk;
19use risingwave_common::array::stream_record::Record;
20use risingwave_common::row::{OwnedRow, Row, RowExt};
21use risingwave_common::types::{
22    DefaultOrd, DefaultOrdered, ScalarImpl, ScalarRefImpl, ToOwnedDatum,
23};
24use risingwave_common::util::memcmp_encoding::MemcmpEncoded;
25use risingwave_storage::StateStore;
26
27use super::BufferTableCache;
28use crate::common::table::state_table::StateTable;
29use crate::executor::{StreamExecutorError, StreamExecutorResult};
30
31type CacheKey = (
32    DefaultOrdered<ScalarImpl>, // sort (watermark) column value
33    MemcmpEncoded,              // memcmp-encoded pk
34);
35
36fn row_to_cache_key<S: StateStore>(
37    sort_column_index: usize,
38    row: impl Row,
39    buffer_table: &StateTable<S>,
40) -> CacheKey {
41    let timestamp_val = row
42        .datum_at(sort_column_index)
43        .to_owned_datum()
44        .expect("watermark column is expected to be non-null");
45    let mut pk = vec![];
46    buffer_table
47        .pk_serde()
48        .serialize((&row).project(buffer_table.pk_indices()), &mut pk);
49    (timestamp_val.into(), pk.into())
50}
51
52/// [`SortBuffer`] is a common component that consume an unordered stream and produce an ordered
53/// stream by watermark. This component maintains a buffer table passed in, whose schema is same as
54/// [`SortBuffer`]'s input and output. Generally, the component acts as a buffer that output the
55/// data it received with a delay, commonly used to implement emit-on-window-close policy.
56pub struct SortBuffer<S: StateStore> {
57    /// The timestamp column to sort on.
58    sort_column_index: usize,
59
60    /// Cache of buffer table.
61    cache: BufferTableCache<CacheKey, OwnedRow>,
62
63    _phantom: PhantomData<S>,
64}
65
66impl<S: StateStore> SortBuffer<S> {
67    /// Create a new [`SortBuffer`].
68    pub fn new(sort_column_index: usize, buffer_table: &StateTable<S>) -> Self {
69        assert_eq!(
70            sort_column_index,
71            buffer_table.pk_indices()[0],
72            "the column to sort on must be the first pk column of the buffer table"
73        );
74
75        Self {
76            sort_column_index,
77            cache: BufferTableCache::new(),
78            _phantom: PhantomData,
79        }
80    }
81
82    /// Insert a new row into the buffer.
83    pub fn insert(&mut self, new_row: impl Row, buffer_table: &mut StateTable<S>) {
84        buffer_table.insert(&new_row);
85        let key = row_to_cache_key(self.sort_column_index, &new_row, buffer_table);
86        self.cache.insert(key, new_row.into_owned_row());
87    }
88
89    /// Delete a row from the buffer.
90    pub fn delete(&mut self, old_row: impl Row, buffer_table: &mut StateTable<S>) {
91        buffer_table.delete(&old_row);
92        let key = row_to_cache_key(self.sort_column_index, &old_row, buffer_table);
93        self.cache.delete(&key);
94    }
95
96    /// Update a row in the buffer.
97    pub fn update(
98        &mut self,
99        old_row: impl Row,
100        new_row: impl Row,
101        buffer_table: &mut StateTable<S>,
102    ) {
103        buffer_table.update(&old_row, &new_row);
104        let key = row_to_cache_key(self.sort_column_index, &old_row, buffer_table);
105        self.cache.delete(&key);
106        self.cache.insert(key, new_row.into_owned_row());
107    }
108
109    /// Apply a change to the buffer, insert/delete/update.
110    pub fn apply_change(&mut self, change: Record<impl Row>, buffer_table: &mut StateTable<S>) {
111        match change {
112            Record::Insert { new_row } => self.insert(new_row, buffer_table),
113            Record::Delete { old_row } => self.delete(old_row, buffer_table),
114            Record::Update { old_row, new_row } => self.update(old_row, new_row, buffer_table),
115        }
116    }
117
118    /// Apply a stream chunk to the buffer.
119    pub fn apply_chunk(&mut self, chunk: StreamChunk, buffer_table: &mut StateTable<S>) {
120        for record in chunk.records() {
121            self.apply_change(record, buffer_table);
122        }
123    }
124
125    /// Consume rows under `watermark` from the buffer.
126    #[try_stream(ok = OwnedRow, error = StreamExecutorError)]
127    pub async fn consume<'a>(
128        &'a mut self,
129        watermark: ScalarImpl,
130        buffer_table: &'a mut StateTable<S>,
131    ) {
132        let mut last_table_pk = None;
133        loop {
134            if !self.cache.is_synced() {
135                // Refill the cache, then consume from the cache, to ensure strong row ordering
136                // and prefetch for the next watermark.
137                self.refill_cache(last_table_pk.take(), buffer_table)
138                    .await?;
139            }
140
141            #[for_await]
142            for res in self.consume_from_cache(watermark.as_scalar_ref_impl()) {
143                let row = res?;
144                last_table_pk = Some((&row).project(buffer_table.pk_indices()).into_owned_row());
145                yield row;
146            }
147
148            if self.cache.is_synced() {
149                // The cache is still synced after consuming, meaning that there is no more rows
150                // under the watermark to yield.
151                break;
152            }
153        }
154
155        // TODO(rc): Need something like `table.range_delete()`. Here we call
156        // `update_watermark(watermark)` as an alternative to `range_delete((..watermark))`.
157        buffer_table.update_watermark(watermark);
158    }
159
160    #[try_stream(ok = OwnedRow, error = StreamExecutorError)]
161    async fn consume_from_cache<'a>(&'a mut self, watermark: ScalarRefImpl<'a>) {
162        while let Some((_, row)) = self
163            .cache
164            .pop_first_if(|key, _| key.0.as_scalar_ref_impl().default_cmp(&watermark).is_lt())
165        {
166            yield row;
167        }
168    }
169
170    /// Clear the cache and refill it with the current content of the buffer table.
171    pub async fn refill_cache(
172        &mut self,
173        last_table_pk: Option<OwnedRow>,
174        buffer_table: &StateTable<S>,
175    ) -> StreamExecutorResult<()> {
176        let sort_column_index = self.sort_column_index;
177        self.cache
178            .refill(buffer_table, last_table_pk, |row| {
179                let key = row_to_cache_key(sort_column_index, &row, buffer_table);
180                Ok((key, row))
181            })
182            .await?;
183        Ok(())
184    }
185}
risingwave_stream/executor/eowc/sort_buffer.rs

risingwave_stream/executor/eowc/
sort_buffer.rs
