risingwave_stream/common/
change_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::sync::LazyLock;
16
17use indexmap::IndexMap;
18use indexmap::map::Entry;
19use risingwave_common::array::stream_record::Record;
20use risingwave_common::array::{Op, StreamChunk, StreamChunkBuilder};
21use risingwave_common::log::LogSuppressor;
22use risingwave_common::row::{Row, RowExt as _};
23use risingwave_common::types::DataType;
24
25use crate::consistency::consistency_panic;
26
27/// Behavior when inconsistency is detected when aggregating changes to [`ChangeBuffer`].
28#[derive(Clone, Copy, Debug, PartialEq, Eq)]
29pub enum InconsistencyBehavior {
30    Panic,
31    Warn,
32    Tolerate,
33}
34
35impl InconsistencyBehavior {
36    /// Report an inconsistency.
37    #[track_caller]
38    pub fn report(self, msg: &str) {
39        match self {
40            InconsistencyBehavior::Panic => consistency_panic!("{}", msg),
41            InconsistencyBehavior::Warn => {
42                static LOG_SUPPRESSOR: LazyLock<LogSuppressor> =
43                    LazyLock::new(LogSuppressor::default);
44
45                if let Ok(suppressed_count) = LOG_SUPPRESSOR.check() {
46                    tracing::warn!(suppressed_count, "{}", msg);
47                }
48            }
49            InconsistencyBehavior::Tolerate => {}
50        }
51    }
52}
53
54mod private {
55    pub trait Key: Eq + std::hash::Hash {}
56    impl<K> Key for K where K: Eq + std::hash::Hash {}
57
58    pub trait Row: Eq {}
59    impl<R> Row for R where R: Eq {}
60}
61
62/// The accumulated change of a key, with each side in an [`Option`] so that state
63/// transitions can take one side out in place. `(None, None)` entries are removed eagerly.
64#[derive(Debug)]
65struct Slot<R> {
66    old: Option<R>,
67    new: Option<R>,
68}
69
70/// A buffer that accumulates changes and produce compacted changes.
71#[derive(Debug)]
72pub struct ChangeBuffer<K, R> {
73    // We use an `IndexMap` to preserve the original order of the changes as much as possible.
74    buffer: IndexMap<K, Slot<R>>,
75    ib: InconsistencyBehavior,
76}
77
78impl<K, R> ChangeBuffer<K, R>
79where
80    K: private::Key,
81    R: private::Row,
82{
83    /// Apply an insertion of a row with the given key.
84    pub fn insert(&mut self, key: K, new_row: R) {
85        let entry = self.buffer.entry(key);
86        match entry {
87            Entry::Vacant(e) => {
88                e.insert(Slot {
89                    old: None,
90                    new: Some(new_row),
91                });
92            }
93            Entry::Occupied(mut e) => {
94                let slot = e.get_mut();
95                if slot.new.is_some() {
96                    self.ib.report("inconsistent changes: double-inserting");
97                }
98                slot.new = Some(new_row);
99            }
100        }
101    }
102
103    /// Apply a deletion of a row with the given key.
104    pub fn delete(&mut self, key: K, old_row: R) {
105        let entry = self.buffer.entry(key);
106        match entry {
107            Entry::Vacant(e) => {
108                e.insert(Slot {
109                    old: Some(old_row),
110                    new: None,
111                });
112            }
113            Entry::Occupied(mut e) => {
114                let slot = e.get_mut();
115                if slot.new.take().is_some() {
116                    if slot.old.is_none() {
117                        // The previous `Insert` is fully cancelled by this deletion.
118                        // FIXME: though preserving the order well,
119                        // this is not performant compared to `swap_remove`
120                        e.shift_remove();
121                    }
122                } else {
123                    self.ib.report("inconsistent changes: double-deleting");
124                    slot.old = Some(old_row);
125                }
126            }
127        }
128    }
129
130    /// Apply an update of a row with the given key.
131    pub fn update(&mut self, key: K, old_row: R, new_row: R) {
132        let entry = self.buffer.entry(key);
133        match entry {
134            Entry::Vacant(e) => {
135                e.insert(Slot {
136                    old: Some(old_row),
137                    new: Some(new_row),
138                });
139            }
140            Entry::Occupied(mut e) => {
141                let slot = e.get_mut();
142                if slot.new.is_some() {
143                    slot.new = Some(new_row);
144                } else {
145                    self.ib.report("inconsistent changes: update after delete");
146                    slot.old = Some(old_row);
147                    slot.new = Some(new_row);
148                }
149            }
150        }
151    }
152
153    /// Apply a change record, with the key extracted by the given function.
154    ///
155    /// For `Record::Update`, inconsistency is reported if the old key and the new key are different.
156    /// Further behavior is determined by the `InconsistencyBehavior`.
157    pub fn apply_record(&mut self, record: Record<R>, key_fn: impl Fn(&R) -> K) {
158        match record {
159            Record::Insert { new_row } => self.insert(key_fn(&new_row), new_row),
160            Record::Delete { old_row } => self.delete(key_fn(&old_row), old_row),
161            Record::Update { old_row, new_row } => {
162                let old_key = key_fn(&old_row);
163                let new_key = key_fn(&new_row);
164
165                // As long as `ib` is not `Panic`, we still gracefully handle the mismatched key.
166                if old_key != new_key {
167                    self.ib
168                        .report("inconsistent changes: mismatched key in update");
169                    self.delete(old_key, old_row);
170                    self.insert(new_key, new_row);
171                } else {
172                    self.update(old_key, old_row, new_row);
173                }
174            }
175        }
176    }
177
178    /// Apply an `Op` of a row with the given key.
179    pub fn apply_op_row(&mut self, op: Op, key: K, row: R) {
180        match op {
181            Op::Insert | Op::UpdateInsert => self.insert(key, row),
182            Op::Delete | Op::UpdateDelete => self.delete(key, row),
183        }
184    }
185
186    /// Consume the buffer and produce a list of change records.
187    ///
188    /// No-op updates are filtered out.
189    pub fn into_records(self) -> impl Iterator<Item = Record<R>> {
190        self.buffer
191            .into_values()
192            .filter_map(|slot| match (slot.old, slot.new) {
193                (None, Some(new_row)) => Some(Record::Insert { new_row }),
194                (Some(old_row), None) => Some(Record::Delete { old_row }),
195                (Some(old_row), Some(new_row)) => {
196                    (old_row != new_row).then(|| Record::Update { old_row, new_row })
197                }
198                (None, None) => unreachable!("empty slot should have been removed"),
199            })
200    }
201}
202
203impl<K, R> Default for ChangeBuffer<K, R> {
204    fn default() -> Self {
205        Self::new()
206    }
207}
208
209impl<K, R> ChangeBuffer<K, R> {
210    /// Create a new `ChangeBuffer` that panics on inconsistency.
211    pub fn new() -> Self {
212        Self::with_capacity(0)
213    }
214
215    /// Create a new `ChangeBuffer` with the given capacity that panics on inconsistency.
216    pub fn with_capacity(capacity: usize) -> Self {
217        Self {
218            buffer: IndexMap::with_capacity(capacity),
219            ib: InconsistencyBehavior::Panic,
220        }
221    }
222
223    /// Set the inconsistency behavior.
224    pub fn with_inconsistency_behavior(mut self, ib: InconsistencyBehavior) -> Self {
225        self.ib = ib;
226        self
227    }
228
229    /// Get the number of keys that have pending changes in the buffer.
230    pub fn len(&self) -> usize {
231        self.buffer.len()
232    }
233
234    /// Check if the buffer is empty.
235    pub fn is_empty(&self) -> bool {
236        self.buffer.is_empty()
237    }
238}
239
240/// The kind of output for [`ChangeBuffer::into_chunk`] and [`ChangeBuffer::into_chunks`].
241/// Can be [`UPSERT`] or [`RETRACT`].
242pub type OutputKind = bool;
243pub mod output_kind {
244    use super::OutputKind;
245    /// For updates, only keep the new row with `Insert` operation.
246    ///
247    /// The output chunk can only be used in streams with `StreamKind::Upsert`. Refer to it for
248    /// more details.
249    pub const UPSERT: OutputKind = true;
250    /// For updates, keep both the old and new row with `UpdateDelete` and `UpdateInsert` operation.
251    pub const RETRACT: OutputKind = false;
252}
253use output_kind::*;
254
255impl<K, R> ChangeBuffer<K, R>
256where
257    K: private::Key,
258    R: private::Row + Row,
259{
260    /// Consume the buffer and produce a single compacted chunk.
261    pub fn into_chunk<const KIND: OutputKind>(
262        self,
263        data_types: Vec<DataType>,
264    ) -> Option<StreamChunk> {
265        let mut builder = StreamChunkBuilder::unlimited(data_types, Some(self.buffer.len()));
266        for record in self.into_records() {
267            let record = match KIND {
268                UPSERT => record.into_upsert(),
269                RETRACT => record,
270            };
271            let none = builder.append_record(record);
272            debug_assert!(none.is_none());
273        }
274        builder.take()
275    }
276
277    /// Consume the buffer and produce a single compacted chunk, with the given new key indices.
278    ///
279    /// The key must be a superset of the old key to ensure uniqueness. Otherwise, changes to the
280    /// same key might be written multiple times so that the correct ordering of the operations cannot
281    /// be guaranteed.
282    pub fn into_chunk_with_key(
283        self,
284        data_types: Vec<DataType>,
285        key_indices: &[usize],
286    ) -> Option<StreamChunk> {
287        let mut builder = StreamChunkBuilder::unlimited(data_types, Some(self.buffer.len()));
288        for record in self.into_records() {
289            macro_rules! append_record {
290                ($record:expr) => {
291                    let none = builder.append_record($record);
292                    debug_assert!(none.is_none());
293                };
294            }
295
296            if let Record::Update { old_row, new_row } = &record
297                && !Row::eq(&old_row.project(key_indices), new_row.project(key_indices))
298            {
299                append_record!(Record::Delete { old_row });
300                append_record!(Record::Insert { new_row });
301            } else {
302                append_record!(record);
303            }
304        }
305        builder.take()
306    }
307
308    /// Consume the buffer and produce a list of compacted chunks with the given size at most.
309    pub fn into_chunks<const KIND: OutputKind>(
310        self,
311        data_types: Vec<DataType>,
312        chunk_size: usize,
313    ) -> Vec<StreamChunk> {
314        let mut res = Vec::new();
315        let mut builder = StreamChunkBuilder::new(chunk_size, data_types);
316        for record in self.into_records() {
317            let record = match KIND {
318                UPSERT => record.into_upsert(),
319                RETRACT => record,
320            };
321            if let Some(chunk) = builder.append_record(record) {
322                res.push(chunk);
323            }
324        }
325        res.extend(builder.take());
326        res
327    }
328}
risingwave_stream/common/change_buffer.rs

risingwave_stream/common/
change_buffer.rs
