risingwave_stream/executor/over_window/
watermark_trigger.rs1use std::sync::Arc;
16
17use risingwave_common::bitmap::Bitmap;
18use risingwave_common::row::{self, OwnedRow, Row, RowExt};
19use risingwave_common::types::{DataType, ScalarImpl};
20use risingwave_common::util::epoch::EpochPair;
21use risingwave_common::util::memcmp_encoding::{self, MemcmpEncoded};
22use risingwave_common::util::sort_util::OrderType;
23use risingwave_common_estimate_size::EstimateSize;
24use risingwave_expr::window_function::WatermarkBound;
25use risingwave_storage::StateStore;
26
27use crate::common::table::state_table::{StateTable, StateTablePostCommit};
28use crate::executor::StreamExecutorResult;
29use crate::executor::eowc::BufferTableCache;
30
31#[derive(Clone, Debug, EstimateSize)]
33pub(super) struct TriggerEntry {
34 pub bound: WatermarkBound,
35 pub partition_key: OwnedRow,
36}
37
38pub(super) struct WatermarkTriggers<S: StateStore> {
41 table: StateTable<S>,
43 cache: BufferTableCache<MemcmpEncoded, TriggerEntry>,
45 order_key_data_type: DataType,
46 watermark: Option<(ScalarImpl, MemcmpEncoded)>,
48 last_popped: Option<OwnedRow>,
51}
52
53impl<S: StateStore> WatermarkTriggers<S> {
54 pub fn new(table: StateTable<S>, order_key_data_type: DataType) -> Self {
55 Self {
56 table,
57 cache: BufferTableCache::new(),
58 order_key_data_type,
59 watermark: None,
60 last_popped: None,
61 }
62 }
63
64 pub async fn init_epoch(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
65 self.table.init_epoch(epoch).await
66 }
67
68 pub fn watermark(&self) -> Option<&MemcmpEncoded> {
71 self.watermark.as_ref().map(|(_, encoded)| encoded)
72 }
73
74 pub fn advance_watermark(&mut self, watermark: ScalarImpl) -> StreamExecutorResult<bool> {
76 let encoded =
77 memcmp_encoding::encode_value(Some(watermark.clone()), OrderType::ascending())?;
78 if self.watermark().is_some_and(|w| w >= &encoded) {
79 return Ok(false);
80 }
81 self.watermark = Some((watermark, encoded));
82 self.last_popped = None;
83 Ok(true)
84 }
85
86 pub fn register(&mut self, entry: TriggerEntry) {
87 let row = self.table_row(&entry);
88 self.table.insert(&row);
89 self.cache.insert(cache_key(&self.table, &row), entry);
90 }
91
92 pub fn remove(&mut self, entry: &TriggerEntry) {
94 let row = self.table_row(entry);
95 self.table.delete(&row);
96 self.cache.delete(&cache_key(&self.table, &row));
97 }
98
99 pub async fn pop_due(&mut self, limit: usize) -> StreamExecutorResult<Vec<TriggerEntry>> {
102 let Some((_, watermark)) = self.watermark.clone() else {
103 return Ok(vec![]);
104 };
105 let mut due = vec![];
106 while due.len() < limit {
107 if !self.cache.is_synced() {
108 self.refill_cache().await?;
109 }
110 match self
111 .cache
112 .pop_first_if(|_, entry| entry.bound.is_satisfied_by(&watermark))
113 {
114 Some((_, entry)) => {
115 let row = self.table_row(&entry);
116 self.table.delete(&row);
117 self.last_popped = Some(row);
118 due.push(entry);
119 }
120 None if !self.cache.is_synced() => {}
122 None => break,
123 }
124 }
125 Ok(due)
126 }
127
128 pub fn clean(&mut self) {
130 if let Some((watermark, _)) = &self.watermark {
131 self.table.update_watermark(watermark.clone());
132 }
133 }
134
135 pub async fn try_flush(&mut self) -> StreamExecutorResult<()> {
136 self.table.try_flush().await
137 }
138
139 pub async fn commit(
140 &mut self,
141 epoch: EpochPair,
142 ) -> StreamExecutorResult<TriggersPostCommit<'_, S>> {
143 Ok(TriggersPostCommit {
144 post_commit: self.table.commit(epoch).await?,
145 cache: &mut self.cache,
146 })
147 }
148
149 async fn refill_cache(&mut self) -> StreamExecutorResult<()> {
150 let table = &self.table;
151 let complete = self
152 .cache
153 .refill(table, self.last_popped.clone(), |row| {
154 let entry = TriggerEntry {
155 bound: WatermarkBound {
156 value: memcmp_encoding::encode_value(
157 row.datum_at(0),
158 OrderType::ascending(),
159 )?,
160 exclusive: row
161 .datum_at(1)
162 .expect("trigger_exclusive must be non-NULL")
163 .into_bool(),
164 },
165 partition_key: OwnedRow::new(row.as_inner()[2..].to_vec()),
166 };
167 Ok((cache_key(table, &row), entry))
168 })
169 .await?;
170 if let Some(count) = complete {
171 self.cache.set_table_row_count(count);
173 }
174 Ok(())
175 }
176
177 fn table_row(&self, entry: &TriggerEntry) -> OwnedRow {
178 let value = memcmp_encoding::decode_value(
179 &self.order_key_data_type,
180 &entry.bound.value,
181 OrderType::ascending(),
182 )
183 .expect("trigger bound must be encoded from an order value");
184 row::once(value)
185 .chain(row::once(Some(ScalarImpl::Bool(entry.bound.exclusive))))
186 .chain(&entry.partition_key)
187 .into_owned_row()
188 }
189}
190
191fn cache_key<S: StateStore>(table: &StateTable<S>, row: &OwnedRow) -> MemcmpEncoded {
192 let mut key = vec![];
193 table
194 .pk_serde()
195 .serialize(row.project(table.pk_indices()), &mut key);
196 key.into()
197}
198
199pub(super) struct TriggersPostCommit<'a, S: StateStore> {
200 post_commit: StateTablePostCommit<'a, S>,
201 cache: &'a mut BufferTableCache<MemcmpEncoded, TriggerEntry>,
202}
203
204impl<S: StateStore> TriggersPostCommit<'_, S> {
205 pub async fn post_yield_barrier(
206 self,
207 new_vnodes: Option<Arc<Bitmap>>,
208 ) -> StreamExecutorResult<()> {
209 if self
210 .post_commit
211 .post_yield_barrier(new_vnodes)
212 .await?
213 .is_some()
214 {
215 self.cache.clear();
217 }
218 Ok(())
219 }
220}