Skip to main content

risingwave_stream/executor/over_window/
watermark_trigger.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::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/// A partition to check once the watermark of the order key satisfies the bound.
32#[derive(Clone, Debug, EstimateSize)]
33pub(super) struct TriggerEntry {
34    pub bound: WatermarkBound,
35    pub partition_key: OwnedRow,
36}
37
38/// The triggers of the partitions whose next rows to output may be closed by watermark, indexed by
39/// bound, so that a watermark only visits the partitions whose triggers it satisfies.
40pub(super) struct WatermarkTriggers<S: StateStore> {
41    /// See `StreamEowcOverWindow::infer_watermark_trigger_table` for the schema.
42    table: StateTable<S>,
43    /// The smallest entries of `table`, keyed by the memcmp-encoded pk.
44    cache: BufferTableCache<MemcmpEncoded, TriggerEntry>,
45    order_key_data_type: DataType,
46    /// The latest watermark of the order key, with its memcmp encoding.
47    watermark: Option<(ScalarImpl, MemcmpEncoded)>,
48    /// The table row of the last entry popped for the current watermark, after which the cache is
49    /// refilled to skip the entries deleted before.
50    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    /// The latest watermark of the order key, memcmp-encoded, which is unknown after recovery until
69    /// the next watermark arrives.
70    pub fn watermark(&self) -> Option<&MemcmpEncoded> {
71        self.watermark.as_ref().map(|(_, encoded)| encoded)
72    }
73
74    /// Returns whether the watermark advances.
75    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    /// Remove a registered entry that hasn't fired.
93    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    /// Pop at most `limit` entries satisfied by the watermark, in the order of their bounds, and
100    /// delete them from the table.
101    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                // The cache runs out, so refill it.
121                None if !self.cache.is_synced() => {}
122                None => break,
123            }
124        }
125        Ok(due)
126    }
127
128    /// Clean the entries below the watermark from the table, after all satisfied ones are popped.
129    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            // The entries before `last_popped` are all deleted, so the cache holds all entries.
172            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            // The cached entries may be of the vnodes no longer owned.
216            self.cache.clear();
217        }
218        Ok(())
219    }
220}