laminar_db/mv_store/
statement_snapshot.rs

1//! Statement planning that reads every scanned materialized view from one completed cycle, or
2//! from one committed checkpoint.
3
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7
8use arrow::array::RecordBatch;
9use datafusion::catalog::default_table_source::source_as_provider;
10use datafusion::dataframe::DataFrame;
11use datafusion::error::DataFusionError;
12use datafusion::execution::SessionState;
13use datafusion::logical_expr::{Expr, LogicalPlan};
14use datafusion::prelude::SessionContext;
15use datafusion_common::tree_node::TreeNodeRecursion;
16
17use super::append_log::LogReadError;
18use super::committed::{CommittedReadError, CommittedViews, PinnedCut, ViewShape};
19use super::{MvStore, ViewRead};
20use crate::error::DbError;
21use crate::table_provider::MvTableProvider;
22
23/// Results of every view a statement scans, captured under one store read guard. A cycle
24/// applies all views under one write guard, so the statement observes a single cycle.
25#[derive(Debug)]
26pub(crate) struct StatementViewSnapshot {
27    views: HashMap<String, RecordBatch>,
28}
29
30impl StatementViewSnapshot {
31    pub(crate) fn get(&self, view: &str) -> Option<&RecordBatch> {
32        self.views.get(view)
33    }
34}
35
36/// The committed cut a statement reads every view at, pinned until the statement is dropped.
37/// Each scan reads its view from the cut's artifacts with the scan's filters.
38pub(crate) struct CommittedStatement {
39    views: Arc<CommittedViews>,
40    cut: PinnedCut,
41    shapes: HashMap<String, ViewShape>,
42}
43
44impl std::fmt::Debug for CommittedStatement {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        f.debug_struct("CommittedStatement")
47            .field("epoch", &self.cut.epoch)
48            .field("views", &self.shapes.keys())
49            .finish_non_exhaustive()
50    }
51}
52
53impl CommittedStatement {
54    pub(crate) async fn read(
55        &self,
56        view: &str,
57        filters: &[Expr],
58    ) -> Result<Vec<RecordBatch>, DbError> {
59        let shape = self.shapes.get(view).ok_or_else(|| {
60            DbError::MaterializedView(format!("view '{view}' was not planned for this read"))
61        })?;
62        self.cut.read_view(&self.views, shape, filters).await
63    }
64}
65
66/// The committed checkpoint a read observed.
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68pub struct ReadCut {
69    /// Epoch of the cut; pass it back as [`super::CommittedRead::AtLeast`] to read no earlier.
70    pub epoch: u64,
71    /// Checkpoint that committed the cut.
72    pub checkpoint_id: u64,
73    /// When the cut was taken, in Unix milliseconds; the result is as old as now minus this.
74    pub taken_at_ms: u64,
75}
76
77/// Attempts at reading views whose captured checkpoint retention moves past before it is pinned.
78const VIEW_READ_ATTEMPTS: u32 = 4;
79const VIEW_READ_RETRY_DELAY: Duration = Duration::from_millis(20);
80/// How long a read may hold a committed checkpoint against retention.
81pub(crate) const VIEW_READ_DEADLINE: Duration = Duration::from_secs(30);
82
83/// Read `views` as of one completed cycle. Every view is captured under one store read guard;
84/// append views then pin the checkpoint they captured and load its segments. A read whose
85/// checkpoint retention moved past before it was pinned is captured again.
86pub(crate) async fn read_views(
87    store: &parking_lot::RwLock<MvStore>,
88    views: &[String],
89) -> Result<HashMap<String, RecordBatch>, DbError> {
90    let (views, ()) = read_captured(store, &views.join(", "), |guard| {
91        let mut captured = Vec::with_capacity(views.len());
92        for view in views {
93            if let Some(read) = guard.read(view)? {
94                captured.push((view.clone(), read));
95            }
96        }
97        Ok((captured, ()))
98    })
99    .await?;
100    Ok(views)
101}
102
103/// Read what `capture` captures under one store read guard, with what it takes beside the
104/// views under that guard, as [`read_views`] does. A capture whose checkpoint retention moved
105/// past is dropped and taken again.
106pub(crate) async fn read_captured<T>(
107    store: &parking_lot::RwLock<MvStore>,
108    described: &str,
109    mut capture: impl FnMut(&MvStore) -> Result<(Vec<(String, ViewRead)>, T), DbError>,
110) -> Result<(HashMap<String, RecordBatch>, T), DbError> {
111    let mut attempt = 1;
112    loop {
113        let (captured, taken) = capture(&store.read())?;
114        match materialize(captured, Instant::now() + VIEW_READ_DEADLINE).await {
115            Ok(views) => return Ok((views, taken)),
116            Err(LogReadError::Failed(error)) => return Err(error),
117            Err(LogReadError::Retired) if attempt < VIEW_READ_ATTEMPTS => {
118                drop(taken);
119                tokio::time::sleep(VIEW_READ_RETRY_DELAY * attempt).await;
120                attempt += 1;
121            }
122            Err(LogReadError::Retired) => {
123                return Err(DbError::MaterializedView(format!(
124                    "materialized views [{described}] were checkpointed {VIEW_READ_ATTEMPTS} \
125                     times while being read; retry the statement"
126                )))
127            }
128        }
129    }
130}
131
132async fn materialize(
133    captured: Vec<(String, ViewRead)>,
134    deadline: Instant,
135) -> Result<HashMap<String, RecordBatch>, LogReadError> {
136    let mut views = HashMap::with_capacity(captured.len());
137    for (view, read) in captured {
138        let batch = match read {
139            ViewRead::Ready(batch) => batch,
140            ViewRead::Log(log) => log.materialize(deadline).await?,
141        };
142        views.insert(view, batch);
143    }
144    Ok(views)
145}
146
147/// Plan `sql` so that every materialized view it scans reads the same completed cycle.
148/// DDL and session statements keep the context's own execution path.
149pub(crate) async fn plan_statement(
150    ctx: &SessionContext,
151    store: &parking_lot::RwLock<MvStore>,
152    sql: &str,
153) -> Result<DataFrame, DbError> {
154    let mut state = ctx.state();
155    let plan = state.create_logical_plan(sql).await?;
156    if matches!(plan, LogicalPlan::Ddl(_) | LogicalPlan::Statement(_)) {
157        return Ok(ctx.execute_logical_plan(plan).await?);
158    }
159    let views = scanned_tables(&plan)?.views;
160    if !views.is_empty() && store.read().holds_node_part() {
161        return Err(DbError::InvalidOperation(format!(
162            "a cluster node holds only its part of materialized views [{}]; read them at a \
163             committed cut",
164            views.join(", ")
165        )));
166    }
167    if !views.is_empty() {
168        let snapshot = StatementViewSnapshot {
169            views: read_views(store, &views).await?,
170        };
171        state
172            .config_mut()
173            .set_extension(std::sync::Arc::new(snapshot));
174    }
175    Ok(DataFrame::new(state, plan))
176}
177
178/// A query planned to read its materialized views at a committed cut not yet chosen.
179pub(crate) struct CommittedPlan {
180    state: SessionState,
181    plan: LogicalPlan,
182    views: Vec<String>,
183}
184
185/// Plan the query `sql` to read every materialized view it scans at one committed cut.
186pub(crate) async fn plan_committed_statement(
187    ctx: &SessionContext,
188    sql: &str,
189) -> Result<CommittedPlan, DbError> {
190    let state = ctx.state();
191    let plan = state.create_logical_plan(sql).await?;
192    if matches!(plan, LogicalPlan::Ddl(_) | LogicalPlan::Statement(_)) {
193        return Err(DbError::InvalidOperation(
194            "a committed read must be a query".into(),
195        ));
196    }
197    let scanned = scanned_tables(&plan)?;
198    if let Some(table) = scanned.live.into_iter().next() {
199        return Err(CommittedReadError::LiveTable { table }.into());
200    }
201    Ok(CommittedPlan {
202        state,
203        plan,
204        views: scanned.views,
205    })
206}
207
208impl CommittedPlan {
209    /// The materialized views the query scans, sorted and distinct.
210    pub(crate) fn views(&self) -> &[String] {
211        &self.views
212    }
213
214    /// Read every view at `cut`, which stays pinned until the returned frame and every stream
215    /// executed from it are dropped, or until its deadline.
216    pub(crate) fn at_cut(
217        self,
218        store: &parking_lot::RwLock<MvStore>,
219        views: Arc<CommittedViews>,
220        cut: PinnedCut,
221    ) -> (DataFrame, ReadCut) {
222        let Self {
223            mut state,
224            plan,
225            views: scanned,
226        } = self;
227        let shapes = {
228            let guard = store.read();
229            scanned
230                .iter()
231                .filter_map(|view| Some((view.clone(), guard.view_shape(view)?)))
232                .collect()
233        };
234        let read_cut = ReadCut {
235            epoch: cut.epoch,
236            checkpoint_id: cut.checkpoint_id,
237            taken_at_ms: cut.taken_at_ms,
238        };
239        state
240            .config_mut()
241            .set_extension(Arc::new(CommittedStatement { views, cut, shapes }));
242        (DataFrame::new(state, plan), read_cut)
243    }
244}
245
246/// Tables a plan scans: sorted, distinct materialized views, including subqueries and views
247/// that `DataFusion` inlines at planning, and every other table.
248#[derive(Debug, Default)]
249struct ScannedTables {
250    views: Vec<String>,
251    live: Vec<String>,
252}
253
254fn scanned_tables(plan: &LogicalPlan) -> Result<ScannedTables, DataFusionError> {
255    let mut scanned = ScannedTables::default();
256    collect_scanned_tables(plan, &mut scanned)?;
257    scanned.views.sort_unstable();
258    scanned.views.dedup();
259    Ok(scanned)
260}
261
262fn collect_scanned_tables(
263    plan: &LogicalPlan,
264    scanned: &mut ScannedTables,
265) -> Result<(), DataFusionError> {
266    plan.apply_with_subqueries(|node| {
267        let LogicalPlan::TableScan(scan) = node else {
268            return Ok(TreeNodeRecursion::Continue);
269        };
270        let Ok(provider) = source_as_provider(&scan.source) else {
271            scanned.live.push(scan.table_name.to_string());
272            return Ok(TreeNodeRecursion::Continue);
273        };
274        if let Some(view) = provider.as_any().downcast_ref::<MvTableProvider>() {
275            scanned.views.push(view.name().to_owned());
276        } else if let Some(inner) = provider.get_logical_plan() {
277            collect_scanned_tables(&inner, scanned)?;
278        } else {
279            scanned.live.push(scan.table_name.to_string());
280        }
281        Ok(TreeNodeRecursion::Continue)
282    })?;
283    Ok(())
284}
285
286#[cfg(test)]
287mod tests {
288    use std::sync::Arc;
289
290    use arrow::array::{Array, Int64Array};
291    use arrow::datatypes::{DataType, Field, Schema};
292    use datafusion::datasource::MemTable;
293
294    use super::*;
295    use crate::mv_store::{MvStorageMode, ViewKey};
296
297    fn schema() -> Arc<Schema> {
298        Arc::new(Schema::new(vec![
299            Field::new("k", DataType::Int64, false),
300            Field::new("n", DataType::Int64, false),
301        ]))
302    }
303
304    fn rows(k: &[i64], n: &[i64]) -> RecordBatch {
305        RecordBatch::try_new(
306            schema(),
307            vec![
308                Arc::new(Int64Array::from(k.to_vec())),
309                Arc::new(Int64Array::from(n.to_vec())),
310            ],
311        )
312        .unwrap()
313    }
314
315    /// Upserts of `(k, n)` rows into a keyed view.
316    fn upserts(k: &[i64], n: &[i64]) -> RecordBatch {
317        let plain = rows(k, n);
318        let mut fields = plain.schema().fields().to_vec();
319        fields.push(Arc::new(Field::new(
320            laminar_core::changelog::WEIGHT_COLUMN,
321            DataType::Int64,
322            false,
323        )));
324        let mut columns = plain.columns().to_vec();
325        columns.push(Arc::new(Int64Array::from(vec![1; k.len()])));
326        RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
327    }
328
329    /// Two keyed views registered on a fresh context, both at cycle 1.
330    fn two_views() -> (SessionContext, Arc<parking_lot::RwLock<MvStore>>) {
331        let store = Arc::new(parking_lot::RwLock::new(MvStore::new()));
332        let ctx = laminar_sql::create_session_context();
333        for view in ["a", "b"] {
334            store
335                .write()
336                .create_mv(
337                    view,
338                    schema(),
339                    MvStorageMode::changelog(ViewKey::Keyed(vec![0])),
340                )
341                .unwrap();
342            store.write().update(view, &upserts(&[1, 2], &[1, 1]));
343            let provider = MvTableProvider::new(view.into(), schema(), Arc::clone(&store));
344            ctx.register_table(view, Arc::new(provider)).unwrap();
345        }
346        (ctx, store)
347    }
348
349    /// Apply cycle 2 to both views under one write guard, as the coordinator does.
350    fn apply_second_cycle(store: &parking_lot::RwLock<MvStore>) {
351        let mut store = store.write();
352        store.update("a", &upserts(&[1, 2], &[2, 2]));
353        store.update("b", &upserts(&[1, 2], &[2, 2]));
354    }
355
356    async fn joined_counts(df: DataFrame) -> Vec<(i64, i64)> {
357        let mut out = Vec::new();
358        for batch in df.collect().await.unwrap() {
359            let a = batch
360                .column(0)
361                .as_any()
362                .downcast_ref::<Int64Array>()
363                .unwrap();
364            let b = batch
365                .column(1)
366                .as_any()
367                .downcast_ref::<Int64Array>()
368                .unwrap();
369            out.extend((0..batch.num_rows()).map(|row| (a.value(row), b.value(row))));
370        }
371        out.sort_unstable();
372        out
373    }
374
375    const JOIN: &str = "SELECT a.n, b.n FROM a JOIN b ON a.k = b.k";
376
377    #[tokio::test]
378    async fn a_statement_reads_every_view_from_the_cycle_it_was_planned_at() {
379        let (ctx, store) = two_views();
380        let df = plan_statement(&ctx, &store, JOIN).await.unwrap();
381        apply_second_cycle(&store);
382        assert_eq!(joined_counts(df).await, vec![(1, 1), (1, 1)]);
383    }
384
385    #[tokio::test]
386    async fn an_unsnapshotted_scan_reads_the_store_when_it_executes() {
387        let (ctx, store) = two_views();
388        let df = ctx.sql(JOIN).await.unwrap();
389        apply_second_cycle(&store);
390        assert_eq!(joined_counts(df).await, vec![(2, 2), (2, 2)]);
391    }
392
393    #[tokio::test]
394    async fn scanned_views_cover_subqueries_ctes_and_inlined_views() {
395        let (ctx, _store) = two_views();
396        let plain = MemTable::try_new(schema(), vec![vec![rows(&[1], &[1])]]).unwrap();
397        ctx.register_table("plain", Arc::new(plain)).unwrap();
398        ctx.sql("CREATE VIEW over_b AS SELECT k FROM b")
399            .await
400            .unwrap();
401        for (sql, expected) in [
402            ("SELECT * FROM plain", vec![]),
403            (
404                "SELECT * FROM a JOIN a AS again ON a.k = again.k",
405                vec!["a"],
406            ),
407            (
408                "SELECT k FROM plain WHERE k IN (SELECT k FROM b)",
409                vec!["b"],
410            ),
411            (
412                "WITH x AS (SELECT k FROM a) SELECT * FROM x JOIN b USING (k)",
413                vec!["a", "b"],
414            ),
415            (
416                "SELECT * FROM a WHERE k IN (SELECT k FROM over_b)",
417                vec!["a", "b"],
418            ),
419        ] {
420            let plan = ctx.state().create_logical_plan(sql).await.unwrap();
421            assert_eq!(scanned_tables(&plan).unwrap().views, expected, "{sql}");
422        }
423    }
424
425    #[tokio::test]
426    async fn ddl_keeps_the_context_execution_path() {
427        let (ctx, store) = two_views();
428        plan_statement(&ctx, &store, "CREATE VIEW over_a AS SELECT * FROM a")
429            .await
430            .unwrap();
431        let df = plan_statement(&ctx, &store, "SELECT COUNT(*) FROM over_a")
432            .await
433            .unwrap();
434        let batches = df.collect().await.unwrap();
435        let count = batches[0].column(0).as_any().downcast_ref::<Int64Array>();
436        assert_eq!(count.map(|c| c.value(0)), Some(2));
437    }
438}
laminar_db/mv_store/statement_snapshot.rs

laminar_db/mv_store/
statement_snapshot.rs
