Skip to main content

risingwave_stream/executor/wrapper/
watermark_check.rs

1// Copyright 2026 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::sync::Arc;
17
18use futures::StreamExt;
19use risingwave_common::array::Op;
20use risingwave_common::row::{Row, RowExt};
21use risingwave_common::types::{DefaultOrd, ScalarImpl};
22use risingwave_pb::stream_plan::stream_node::PbStreamKind;
23
24use crate::executor::{ExecutorInfo, Message, MessageStream};
25
26/// Streams wrapped by `watermark_check` will check that no row falls below a watermark emitted
27/// earlier on the same column. A delete of an upsert stream is not checked, as only its stream key
28/// identifies the row it deletes.
29///
30/// It is a combinator rather than a generator, which would deepen the stack of every executor in
31/// debug builds.
32pub fn watermark_check(info: Arc<ExecutorInfo>, input: impl MessageStream) -> impl MessageStream {
33    let mut watermarks = BTreeMap::<usize, ScalarImpl>::new();
34    input.inspect(move |message| match message {
35        Ok(Message::Watermark(watermark)) => {
36            let emitted = watermarks
37                .entry(watermark.col_idx)
38                .or_insert_with(|| watermark.val.clone());
39            if watermark.val.default_cmp(emitted).is_gt() {
40                *emitted = watermark.val.clone();
41            }
42        }
43        Ok(Message::Chunk(chunk)) if !watermarks.is_empty() => {
44            for (op, row) in chunk.rows() {
45                if info.stream_kind == PbStreamKind::Upsert
46                    && matches!(op, Op::Delete | Op::UpdateDelete)
47                {
48                    continue;
49                }
50                for (&col_idx, watermark) in &watermarks {
51                    if let Some(value) = row.datum_at(col_idx) {
52                        debug_assert!(
53                            value.default_cmp(&watermark.as_scalar_ref_impl()).is_ge(),
54                            "row below a watermark emitted earlier
55row: {}
56op: {:?}
57column: {}
58watermark: {:?}
59executor: {}",
60                            row.display(),
61                            op,
62                            col_idx,
63                            watermark,
64                            info.identity
65                        );
66                    }
67                }
68            }
69        }
70        _ => {}
71    })
72}
73
74#[cfg(test)]
75mod tests {
76    use futures::pin_mut;
77    use risingwave_common::array::StreamChunk;
78    use risingwave_common::array::stream_chunk::StreamChunkTestExt;
79    use risingwave_common::catalog::{Field, Schema};
80    use risingwave_common::types::DataType;
81
82    use super::*;
83    use crate::executor::test_utils::MockSource;
84
85    /// Runs `chunk` after `Watermark(10)` on the only column through the check.
86    async fn check_after_watermark(stream_kind: PbStreamKind, chunk: &str) {
87        let (mut tx, source) = MockSource::channel();
88        let source =
89            source.into_executor(Schema::new(vec![Field::unnamed(DataType::Int64)]), vec![]);
90        let mut info = source.info().clone();
91        info.stream_kind = stream_kind;
92        tx.push_int64_watermark(0, 10);
93        tx.push_chunk(StreamChunk::from_pretty(chunk));
94
95        let checked = watermark_check(info.into(), source.execute());
96        pin_mut!(checked);
97        checked.next().await.unwrap().unwrap();
98        checked.next().await.unwrap().unwrap();
99    }
100
101    #[should_panic(expected = "row below a watermark emitted earlier")]
102    #[tokio::test]
103    async fn test_row_below_watermark() {
104        check_after_watermark(PbStreamKind::Retract, " I\n + 10\n - 9").await;
105    }
106
107    #[tokio::test]
108    async fn test_rows_not_below_watermark() {
109        check_after_watermark(PbStreamKind::Retract, " I\n + 10\n + 11\n + .").await;
110        check_after_watermark(PbStreamKind::Upsert, " I\n - 9").await;
111    }
112}