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    /// Apply an operation from an upsert stream, retaining only the final operation per key.
187    /// Unlike retract input, an insertion can overwrite a row from an earlier batch, so a
188    /// subsequent deletion must retain its tombstone even if no earlier deletion was buffered.
189    pub fn apply_upsert_op_row(&mut self, op: Op, key: K, row: R) {
190        let slot = match op {
191            Op::Insert | Op::UpdateInsert => Slot {
192                old: None,
193                new: Some(row),
194            },
195            Op::Delete | Op::UpdateDelete => Slot {
196                old: Some(row),
197                new: None,
198            },
199        };
200        self.buffer.insert(key, slot);
201    }
202
203    /// Consume the buffer and produce a list of change records.
204    ///
205    /// No-op updates are filtered out.
206    pub fn into_records(self) -> impl Iterator<Item = Record<R>> {
207        self.buffer
208            .into_values()
209            .filter_map(|slot| match (slot.old, slot.new) {
210                (None, Some(new_row)) => Some(Record::Insert { new_row }),
211                (Some(old_row), None) => Some(Record::Delete { old_row }),
212                (Some(old_row), Some(new_row)) => {
213                    (old_row != new_row).then(|| Record::Update { old_row, new_row })
214                }
215                (None, None) => unreachable!("empty slot should have been removed"),
216            })
217    }
218}
219
220impl<K, R> Default for ChangeBuffer<K, R> {
221    fn default() -> Self {
222        Self::new()
223    }
224}
225
226impl<K, R> ChangeBuffer<K, R> {
227    /// Create a new `ChangeBuffer` that panics on inconsistency.
228    pub fn new() -> Self {
229        Self::with_capacity(0)
230    }
231
232    /// Create a new `ChangeBuffer` with the given capacity that panics on inconsistency.
233    pub fn with_capacity(capacity: usize) -> Self {
234        Self {
235            buffer: IndexMap::with_capacity(capacity),
236            ib: InconsistencyBehavior::Panic,
237        }
238    }
239
240    /// Set the inconsistency behavior.
241    pub fn with_inconsistency_behavior(mut self, ib: InconsistencyBehavior) -> Self {
242        self.ib = ib;
243        self
244    }
245
246    /// Get the number of keys that have pending changes in the buffer.
247    pub fn len(&self) -> usize {
248        self.buffer.len()
249    }
250
251    /// Check if the buffer is empty.
252    pub fn is_empty(&self) -> bool {
253        self.buffer.is_empty()
254    }
255}
256
257/// The kind of output for [`ChangeBuffer::into_chunk`] and [`ChangeBuffer::into_chunks`].
258/// Can be [`UPSERT`] or [`RETRACT`].
259pub type OutputKind = bool;
260pub mod output_kind {
261    use super::OutputKind;
262    /// For updates, only keep the new row with `Insert` operation.
263    ///
264    /// The output chunk can only be used in streams with `StreamKind::Upsert`. Refer to it for
265    /// more details.
266    pub const UPSERT: OutputKind = true;
267    /// For updates, keep both the old and new row with `UpdateDelete` and `UpdateInsert` operation.
268    pub const RETRACT: OutputKind = false;
269}
270use output_kind::*;
271
272impl<K, R> ChangeBuffer<K, R>
273where
274    K: private::Key,
275    R: private::Row + Row,
276{
277    /// Consume the buffer and produce a single compacted chunk.
278    pub fn into_chunk<const KIND: OutputKind>(
279        self,
280        data_types: Vec<DataType>,
281    ) -> Option<StreamChunk> {
282        let mut builder = StreamChunkBuilder::unlimited(data_types, Some(self.buffer.len()));
283        for record in self.into_records() {
284            let record = match KIND {
285                UPSERT => record.into_upsert(),
286                RETRACT => record,
287            };
288            let none = builder.append_record(record);
289            debug_assert!(none.is_none());
290        }
291        builder.take()
292    }
293
294    /// Consume the buffer and produce a single compacted chunk, with the given new key indices.
295    ///
296    /// The key must be a superset of the old key to ensure uniqueness. Otherwise, changes to the
297    /// same key might be written multiple times so that the correct ordering of the operations cannot
298    /// be guaranteed.
299    pub fn into_chunk_with_key(
300        self,
301        data_types: Vec<DataType>,
302        key_indices: &[usize],
303    ) -> Option<StreamChunk> {
304        let mut builder = StreamChunkBuilder::unlimited(data_types, Some(self.buffer.len()));
305        for record in self.into_records() {
306            macro_rules! append_record {
307                ($record:expr) => {
308                    let none = builder.append_record($record);
309                    debug_assert!(none.is_none());
310                };
311            }
312
313            if let Record::Update { old_row, new_row } = &record
314                && !Row::eq(&old_row.project(key_indices), new_row.project(key_indices))
315            {
316                append_record!(Record::Delete { old_row });
317                append_record!(Record::Insert { new_row });
318            } else {
319                append_record!(record);
320            }
321        }
322        builder.take()
323    }
324
325    /// Consume the buffer and produce a list of compacted chunks with the given size at most.
326    pub fn into_chunks<const KIND: OutputKind>(
327        self,
328        data_types: Vec<DataType>,
329        chunk_size: usize,
330    ) -> Vec<StreamChunk> {
331        let mut res = Vec::new();
332        let mut builder = StreamChunkBuilder::new(chunk_size, data_types);
333        for record in self.into_records() {
334            let record = match KIND {
335                UPSERT => record.into_upsert(),
336                RETRACT => record,
337            };
338            if let Some(chunk) = builder.append_record(record) {
339                res.push(chunk);
340            }
341        }
342        res.extend(builder.take());
343        res
344    }
345}
risingwave_stream/common/change_buffer.rs

risingwave_stream/common/
change_buffer.rs
