risingwave_stream/executor/wrapper/
watermark_check.rs1use 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
26pub 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 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}