risingwave_stream/executor/mview/
refresh_progress_table.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
15//! Refresh Progress Table
16//!
17//! This module implements a persistent table for tracking refresh operation progress.
18//! It stores the merge position of each `VirtualNode` during a refresh cycle. Within a cycle the
19//! merge resumes from it after every barrier; a refresh interrupted by recovery is currently
20//! abandoned by meta rather than resumed from it.
21
22use std::collections::{HashMap, HashSet};
23use std::sync::Arc;
24
25use risingwave_common::bitmap::Bitmap;
26use risingwave_common::hash::VirtualNode;
27use risingwave_common::row::{OwnedRow, Row};
28use risingwave_common::types::{DataType, ScalarImpl, ScalarRefImpl};
29use risingwave_common::util::epoch::EpochPair;
30use risingwave_storage::StateStore;
31
32use crate::common::table::state_table::StateTablePostCommit;
33use crate::executor::StreamExecutorResult;
34use crate::executor::prelude::StateTable;
35
36/// Schema for the refresh progress table (simplified, following backfill pattern):
37/// - `vnode` (i16): `VirtualNode` identifier
38/// - `current_pos` (variable): Current processing position (primary key of last processed row)
39/// - `is_completed` (bool): Whether this vnode has completed processing
40/// - `processed_rows` (i64): Number of rows processed so far in this vnode
41pub struct RefreshProgressTable<S: StateStore> {
42    /// The underlying state table for persistence
43    pub state_table: StateTable<S>,
44    /// In-memory cache of progress information for quick access
45    cache: HashMap<VirtualNode, RefreshProgressEntry>,
46    /// Primary key length for upstream table (needed for schema)
47    pk_len: usize,
48}
49
50/// Progress information for a single `VirtualNode`
51#[derive(Debug, Clone, PartialEq)]
52pub struct RefreshProgressEntry {
53    pub vnode: VirtualNode,
54    pub current_pos: Option<OwnedRow>,
55    pub is_completed: bool,
56    pub processed_rows: u64,
57}
58
59/// Refresh stages are now tracked by `MaterializeExecutor` in memory
60/// This enum is kept for compatibility but no longer stored in progress table
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62#[repr(i32)]
63pub enum ProgressRefreshStage {
64    Normal = 0,
65    Refreshing = 1,
66    Merging = 2,
67    Cleanup = 3,
68}
69
70impl From<i32> for ProgressRefreshStage {
71    fn from(value: i32) -> Self {
72        match value {
73            0 => ProgressRefreshStage::Normal,
74            1 => ProgressRefreshStage::Refreshing,
75            2 => ProgressRefreshStage::Merging,
76            3 => ProgressRefreshStage::Cleanup,
77            _ => unreachable!(),
78        }
79    }
80}
81
82impl<S: StateStore> RefreshProgressTable<S> {
83    /// Create a new `RefreshProgressTable`
84    pub fn new(state_table: StateTable<S>, pk_len: usize) -> Self {
85        Self {
86            state_table,
87            cache: HashMap::new(),
88            pk_len,
89        }
90    }
91
92    /// Initialize the progress table by loading existing entries from storage
93    pub async fn recover(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
94        self.state_table.init_epoch(epoch).await?;
95
96        // Load existing progress entries from storage into cache
97        let mut loaded_count = 0;
98
99        for vnode in self.state_table.vnodes().iter_ones() {
100            let row = self
101                .state_table
102                .get_row(OwnedRow::new(vec![
103                    VirtualNode::from_index(vnode).to_datum(),
104                ]))
105                .await?;
106            if row.is_some()
107                && let Some(entry) = self.parse_row_to_entry(&row, self.pk_len)
108            {
109                self.cache.insert(entry.vnode, entry);
110                loaded_count += 1;
111            }
112        }
113
114        tracing::debug!(
115            loaded_count,
116            "Loading existing progress entries during initialization"
117        );
118
119        tracing::info!(
120            loaded_entries = self.cache.len(),
121            "Initialized refresh progress table"
122        );
123
124        Ok(())
125    }
126
127    /// Set progress for a specific `VirtualNode`
128    pub fn set_progress(
129        &mut self,
130        vnode: VirtualNode,
131        current_pos: Option<OwnedRow>,
132        is_completed: bool,
133        processed_rows: u64,
134    ) -> StreamExecutorResult<()> {
135        let entry = RefreshProgressEntry {
136            vnode,
137            current_pos,
138            is_completed,
139            processed_rows,
140        };
141
142        // Update cache
143        self.cache.insert(vnode, entry.clone());
144
145        // Write to persistent storage
146        let row = self.entry_to_row(&entry, self.pk_len);
147        self.state_table.insert(&row);
148
149        Ok(())
150    }
151
152    /// Get progress for a specific `VirtualNode`
153    pub fn get_progress(&self, vnode: VirtualNode) -> Option<&RefreshProgressEntry> {
154        self.cache.get(&vnode)
155    }
156
157    /// Get all `VirtualNodes` that have completed processing
158    pub fn get_completed_vnodes(&self) -> HashSet<VirtualNode> {
159        self.cache
160            .iter()
161            .filter(|(_, entry)| entry.is_completed)
162            .map(|(&vnode, _)| vnode)
163            .collect()
164    }
165
166    /// Get all `VirtualNodes` in a specific stage (now tracked by executor memory)
167    /// This method is kept for compatibility but should not be used with simplified schema
168    pub fn get_vnodes_in_stage(&self, _stage: ProgressRefreshStage) -> Vec<VirtualNode> {
169        tracing::warn!(
170            "get_vnodes_in_stage called on simplified progress table - stage info no longer stored"
171        );
172        Vec::new()
173    }
174
175    /// Clear progress for a specific `VirtualNode`
176    pub fn clear_progress(&mut self, vnode: VirtualNode) -> StreamExecutorResult<()> {
177        if let Some(entry) = self.cache.remove(&vnode) {
178            let row = self.entry_to_row(&entry, self.pk_len);
179            self.state_table.delete(&row);
180        }
181
182        Ok(())
183    }
184
185    /// Clear all progress entries
186    pub fn clear_all_progress(&mut self) -> StreamExecutorResult<()> {
187        for vnode in self.cache.keys().copied().collect::<Vec<_>>() {
188            self.clear_progress(vnode)?;
189        }
190        Ok(())
191    }
192
193    /// Get the total number of processed rows across all `VirtualNodes`
194    pub fn get_total_processed_rows(&self) -> u64 {
195        self.cache.values().map(|entry| entry.processed_rows).sum()
196    }
197
198    /// Get progress statistics
199    pub fn get_progress_stats(&self) -> RefreshProgressStats {
200        let total_vnodes = self.cache.len();
201        let completed_vnodes = self.get_completed_vnodes().len();
202        let total_processed_rows = self.get_total_processed_rows();
203
204        RefreshProgressStats {
205            total_vnodes,
206            completed_vnodes,
207            total_processed_rows,
208        }
209    }
210
211    /// Commit changes to storage
212    pub async fn commit(
213        &mut self,
214        epoch: EpochPair,
215    ) -> StreamExecutorResult<StateTablePostCommit<'_, S>> {
216        self.state_table.commit(epoch).await
217    }
218
219    /// Convert `RefreshProgressEntry` to `OwnedRow` for storage
220    /// New schema: | vnode | `current_pos`... | `is_completed` | `processed_rows` |
221    fn entry_to_row(&self, entry: &RefreshProgressEntry, pk_len: usize) -> OwnedRow {
222        let mut row_data = vec![entry.vnode.to_datum()];
223
224        // Add current_pos fields (variable length based on primary key)
225        if let Some(ref pos) = entry.current_pos {
226            row_data.extend(pos.iter().map(|d| d.map(|s| s.into_scalar_impl())));
227        } else {
228            // Use None placeholders for empty position (similar to backfill finished state)
229            for _ in 0..pk_len {
230                row_data.push(None);
231            }
232        }
233
234        // Add metadata fields
235        row_data.push(Some(ScalarImpl::Bool(entry.is_completed)));
236        row_data.push(Some(ScalarImpl::Int64(entry.processed_rows as i64)));
237
238        OwnedRow::new(row_data)
239    }
240
241    /// Parse `OwnedRow` from storage to `RefreshProgressEntry`
242    /// New schema: | vnode | `current_pos`... | `is_completed` | `processed_rows` |
243    /// Note: The length depends on the primary key length of upstream table
244    fn parse_row_to_entry(&self, row: &impl Row, pk_len: usize) -> Option<RefreshProgressEntry> {
245        let datums = row.iter().collect::<Vec<_>>();
246        let expected_len = 1 + pk_len + 2; // vnode + pk_fields + is_completed + processed_rows
247
248        if datums.len() != expected_len {
249            tracing::warn!(
250                "Row length mismatch: got {}, expected {} (pk_len={}), row: {:?}",
251                datums.len(),
252                expected_len,
253                pk_len,
254                row,
255            );
256            return None;
257        }
258
259        // Parse vnode (first field)
260        let vnode = VirtualNode::from_datum_ref(datums[0]);
261
262        // Parse current_pos (pk_len fields after vnode)
263        let current_pos = if pk_len > 0 {
264            let pos_datums: Vec<_> = datums[1..1 + pk_len]
265                .iter()
266                .map(|d| d.map(|s| s.into_scalar_impl()))
267                .collect();
268            // Check if all position fields are None (indicating finished/empty state)
269            if pos_datums.iter().all(|d| d.is_none()) {
270                None
271            } else {
272                Some(OwnedRow::new(pos_datums))
273            }
274        } else {
275            None
276        };
277
278        // Parse is_completed (second to last field)
279        let is_completed = match datums[1 + pk_len]? {
280            ScalarRefImpl::Bool(val) => val,
281            _ => return None,
282        };
283
284        // Parse processed_rows (last field)
285        let processed_rows = match datums[1 + pk_len + 1]? {
286            ScalarRefImpl::Int64(val) => val as u64,
287            _ => return None,
288        };
289
290        Some(RefreshProgressEntry {
291            vnode,
292            current_pos,
293            is_completed,
294            processed_rows,
295        })
296    }
297
298    /// Get the expected schema for the progress table
299    /// Schema: | vnode | `current_pos`... | `is_completed` | `processed_rows` |
300    pub fn expected_schema(pk_data_types: &[DataType]) -> Vec<DataType> {
301        let mut schema = vec![VirtualNode::RW_TYPE]; // vnode
302        schema.extend(pk_data_types.iter().cloned()); // current_pos fields
303        schema.push(DataType::Boolean); // is_completed
304        schema.push(DataType::Int64); // processed_rows
305        schema
306    }
307
308    /// Get column names for the progress table schema
309    /// Schema: | vnode | `current_pos`... | `is_completed` | `processed_rows` |
310    pub fn column_names(pk_column_names: &[&str]) -> Vec<String> {
311        let mut names = vec!["vnode".to_owned()];
312        for pk_name in pk_column_names {
313            names.push(format!("pos_{}", pk_name));
314        }
315        names.push("is_completed".to_owned());
316        names.push("processed_rows".to_owned());
317        names
318    }
319
320    pub fn vnodes(&self) -> &Arc<Bitmap> {
321        self.state_table.vnodes()
322    }
323}
324
325/// Progress statistics for monitoring
326#[derive(Debug, Clone)]
327pub struct RefreshProgressStats {
328    pub total_vnodes: usize,
329    pub completed_vnodes: usize,
330    pub total_processed_rows: u64,
331}
332
333impl RefreshProgressStats {
334    /// Check if refresh is complete
335    pub fn is_complete(&self) -> bool {
336        self.total_vnodes > 0 && self.completed_vnodes == self.total_vnodes
337    }
338}
risingwave_stream/executor/mview/refresh_progress_table.rs

risingwave_stream/executor/mview/
refresh_progress_table.rs
