Skip to main content

mz_storage_operators/
stats.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Types and traits that connect up our mz-repr types with the stats that persist maintains.
11
12use mz_expr::{ResultSpec, SafeMfpPlan};
13use mz_persist_client::metrics::Metrics;
14use mz_persist_client::read::{Cursor, LazyPartStats, ReadHandle, Since};
15use mz_repr::{RelationDesc, Row, Timestamp};
16use mz_storage_types::StorageDiff;
17use mz_storage_types::controller::TxnsCodecRow;
18use mz_storage_types::errors::DataflowError;
19use mz_storage_types::sources::SourceData;
20use mz_storage_types::stats::RelationPartStats;
21use mz_txn_wal::txn_cache::TxnsCache;
22use timely::progress::Antichain;
23
24/// This is a streaming-consolidating cursor type specialized to `RelationDesc`.
25///
26/// Internally this maintains two separate cursors: one for errors and one for data.
27/// This is necessary so that errors are presented before data, which matches our usual
28/// lookup semantics. To avoid being ludicrously inefficient, this pushes down a filter
29/// on the stats. (In particular, in the common case of no errors, we don't do any extra
30/// fetching.)
31pub struct StatsCursor {
32    errors: Cursor<SourceData, (), Timestamp, StorageDiff>,
33    data: Cursor<SourceData, (), Timestamp, StorageDiff>,
34}
35
36impl StatsCursor {
37    pub async fn new(
38        handle: &mut ReadHandle<SourceData, (), Timestamp, StorageDiff>,
39        // If and only if we are using txn-wal to manage this shard, then
40        // this must be Some. This is because the upper might be advanced lazily
41        // and we have to go through txn-wal for reads.
42        txns_read: Option<&mut TxnsCache<Timestamp, TxnsCodecRow>>,
43        metrics: &Metrics,
44        mfp_plan: &SafeMfpPlan,
45        desc: &RelationDesc,
46        as_of: Antichain<Timestamp>,
47    ) -> Result<StatsCursor, Since<Timestamp>> {
48        let should_fetch = |name: &'static str, errors: bool| {
49            move |stats: Option<&LazyPartStats>| {
50                let Some(stats) = stats else { return true };
51                // Stats written by a newer version may not decode. The sound
52                // fallback is to fetch the part.
53                let Ok(stats) = stats.try_decode() else {
54                    return true;
55                };
56                let metrics = &metrics.pushdown.part_stats;
57                let relation_stats = RelationPartStats::new(name, metrics, desc, &stats);
58                if errors {
59                    relation_stats.err_count().map_or(true, |e| e > 0)
60                } else {
61                    relation_stats.may_match_mfp(ResultSpec::value_all(), mfp_plan)
62                }
63            }
64        };
65        let (errors, data) = match txns_read {
66            None => {
67                let errors = handle
68                    .snapshot_cursor(as_of.clone(), should_fetch("errors", true))
69                    .await?;
70                let data = handle
71                    .snapshot_cursor(as_of.clone(), should_fetch("data", false))
72                    .await?;
73                (errors, data)
74            }
75            Some(txns_read) => {
76                let as_of = as_of
77                    .as_option()
78                    .expect("reads as_of empty antichain block forever")
79                    .clone();
80                let _ = txns_read.update_gt(&as_of).await;
81                let data_snapshot = txns_read.data_snapshot(handle.shard_id(), as_of);
82                let errors: Cursor<SourceData, (), Timestamp, i64> = data_snapshot
83                    .snapshot_cursor(handle, should_fetch("errors", true))
84                    .await?;
85                let data = data_snapshot
86                    .snapshot_cursor(handle, should_fetch("data", false))
87                    .await?;
88                (errors, data)
89            }
90        };
91
92        Ok(StatsCursor { errors, data })
93    }
94
95    pub async fn next(
96        &mut self,
97    ) -> Option<impl Iterator<Item = (Result<Row, DataflowError>, Timestamp, StorageDiff)> + '_>
98    {
99        fn expect_decode(
100            raw: impl Iterator<Item = ((SourceData, ()), Timestamp, StorageDiff)>,
101            is_err: bool,
102        ) -> impl Iterator<Item = (Result<Row, DataflowError>, Timestamp, StorageDiff)> {
103            raw.map(|((k, v), t, d)| {
104                // NB: this matches the decode behaviour in sources
105                let SourceData(row) = k;
106                let () = v;
107                (row, t, d)
108            })
109            .filter(move |(r, _, _)| if is_err { r.is_err() } else { r.is_ok() })
110        }
111
112        if let Some(errors) = self.errors.next().await {
113            Some(expect_decode(errors, true))
114        } else if let Some(data) = self.data.next().await {
115            Some(expect_decode(data, false))
116        } else {
117            None
118        }
119    }
120}