mz_storage/upsert_continual_feedback_v2/
feedback.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//! The persist-feedback side of the upsert-v2 operator: the feedback
11//! arrangement's types, and the `UpsertValue` to `Row` encoding it stores.
12
13use differential_dataflow::hashable::Hashable;
14use differential_dataflow::trace::chunk::ChunkBatch;
15use differential_dataflow::{AsCollection, VecCollection};
16use mz_repr::{Datum, Diff, Row};
17// Only the fuzzing-gated `datum_seq_to_upsert_value` takes a `DatumSeq`.
18#[cfg(feature = "fuzzing")]
19use mz_row_spine::DatumSeq;
20use mz_storage_types::errors::{DataflowError, UpsertError};
21use mz_timely_util::columnar::Column;
22use mz_timely_util::columnar::builder::ColumnBuilder;
23use mz_timely_util::columnar::chunk::{ColumnChunk, StorageSpill};
24use timely::dataflow::Stream;
25use timely::dataflow::channels::pact::Pipeline;
26use timely::dataflow::operators::Exchange as _;
27use timely::dataflow::operators::generic::Operator;
28use timely::order::TotalOrder;
29use timely::progress::Timestamp;
30use timely::progress::timestamp::Refines;
31
32use crate::statistics::SourceStatistics;
33use crate::upsert::{UpsertKey, UpsertValue};
34
35/// One persist-feedback update: a key, its current value row, the time, and
36/// an additive count.
37pub(super) type FeedbackUpdate<T> = ((UpsertKey, Row), T, Diff);
38
39/// One feedback-arrangement chunk: a sorted, consolidated run of updates,
40/// resident or spilled to the buffer pool. Both batcher chains and sealed
41/// spine batches are sequences of these, so the arrangement's state pages out
42/// of RSS under the pool's budget, and the drain reads it back through the
43/// bulk [`UnloadChunk`](mz_timely_util::columnar::unload::UnloadChunk)
44/// surface: copy-out probes, no cursor borrows.
45pub(super) type FeedbackChunk<T> = ColumnChunk<(UpsertKey, Row), T, Diff, StorageSpill>;
46
47/// The feedback arrangement's trace: a funded spine of `Rc`-shared chunk
48/// batches. See [`mz_timely_util::funded_spine`].
49pub(super) type FeedbackSpine<T> =
50    mz_timely_util::funded_spine::Spine<std::rc::Rc<ChunkBatch<FeedbackChunk<T>>>>;
51
52// The persist-feedback arrangement stores `(UpsertKey, Row)` pairs in columnar
53// chunk batches ([`FeedbackSpine`]). `UpsertValue` is
54// `Result<Row, Box<UpsertError>>`, so we fold both arms into a single `Row`
55// with a leading tag column so they share the value column.
56
57/// Encode an [`UpsertValue`] as a `Row` with a leading tag column so both `Ok`
58/// and `Err` payloads round-trip through `Row` byte storage.
59///
60/// Used on the render path. `pub` only so [`crate::fuzz_exports`] can re-export
61/// it under the `fuzzing` feature for the storage fuzz crate. The enclosing
62/// module is crate-private, so it is not otherwise reachable. Not a stable
63/// public API.
64#[doc(hidden)]
65pub fn upsert_value_to_row(value: &UpsertValue) -> Row {
66    let mut row = Row::default();
67    let mut packer = row.packer();
68    match value {
69        Ok(ok) => {
70            packer.push(Datum::UInt8(0));
71            packer.extend(ok.iter());
72        }
73        Err(err) => {
74            packer.push(Datum::UInt8(1));
75            let bytes =
76                bincode::serialize(err.as_ref()).expect("UpsertError is serializable via bincode");
77            packer.push(Datum::Bytes(&bytes));
78        }
79    }
80    row
81}
82
83/// Heap-size estimate for an emitted [`UpsertValue`], used to drive
84/// `give_fueled` yielding on the output edge.
85pub(super) fn upsert_value_byte_len(value: &UpsertValue) -> usize {
86    match value {
87        Ok(row) => row.byte_len(),
88        Err(err) => std::mem::size_of_val(err.as_ref()),
89    }
90}
91
92/// Decode an [`UpsertValue`] produced by [`upsert_value_to_row`] back from the
93/// `DatumSeq` view returned by a `ValRowSpine` cursor.
94///
95/// Exists only for the storage fuzz crate, so it is gated behind the `fuzzing`
96/// feature. Not a stable public API.
97#[cfg(feature = "fuzzing")]
98#[doc(hidden)]
99pub fn datum_seq_to_upsert_value(seq: DatumSeq<'_>) -> UpsertValue {
100    decode_upsert_value(seq)
101}
102
103/// Decode an [`UpsertValue`] produced by [`upsert_value_to_row`] from any
104/// datum iterator, whether over a columnar `Row` reference or an owned `Row`.
105pub(super) fn decode_upsert_value<'a>(mut iter: impl Iterator<Item = Datum<'a>>) -> UpsertValue {
106    let tag = match iter.next() {
107        Some(Datum::UInt8(tag)) => tag,
108        other => panic!("upsert value missing UInt8 tag, got {:?}", other),
109    };
110    match tag {
111        0 => {
112            let mut row = Row::default();
113            row.packer().extend(iter);
114            Ok(row)
115        }
116        1 => {
117            let bytes = match iter.next() {
118                Some(Datum::Bytes(b)) => b,
119                other => panic!("upsert error tag missing Bytes payload, got {:?}", other),
120            };
121            let err: UpsertError =
122                bincode::deserialize(bytes).expect("UpsertError bincode round-trip");
123            Err(Box::new(err))
124        }
125        tag => panic!("unknown upsert value tag {tag}"),
126    }
127}
128
129/// Key the persist feedback by [`UpsertKey`], record source statistics, and
130/// encode `(UpsertKey, UpsertValue)` as `(UpsertKey, Row)` `Column` chunks,
131/// the input both flavors' feedback arrangements consume. Built with
132/// `Pipeline` downstream of an `UpsertKey::hashed` exchange, so the
133/// arrangement keeps that locality.
134pub(super) fn encode_feedback<'scope, T>(
135    persist_ok: VecCollection<'scope, T, Row, Diff>,
136    persist_err: VecCollection<'scope, T, DataflowError, Diff>,
137    key_indices: Vec<usize>,
138    source_statistics: SourceStatistics,
139) -> Stream<'scope, T, Column<((UpsertKey, Row), T, Diff)>>
140where
141    T: Timestamp + TotalOrder + Sync,
142    T: Refines<mz_repr::Timestamp> + differential_dataflow::lattice::Lattice,
143    T: columnation::Columnation,
144    T: columnar::Columnar + Default,
145    for<'a> columnar::Ref<'a, T>: Copy + Ord,
146{
147    let persist_keyed = crate::upsert::key_persist_feedback(persist_ok, persist_err, key_indices);
148    let persist_keyed = persist_keyed
149        .inner
150        // The arrangement already implicitly exchanges by key, so this is redundant, but we want to
151        // do it earlier so that we can inspect the stream properly for source statistics.
152        .exchange(move |((key, _), _, _)| UpsertKey::hashed(key))
153        .as_collection()
154        .inspect(move |((_, row), _, diff)| {
155            source_statistics.update_records_indexed_by(diff.into_inner());
156            source_statistics.update_bytes_indexed_by(
157                row.as_ref().map_or(0, |r| r.byte_len().try_into().unwrap()) * diff.into_inner(),
158            );
159        });
160    persist_keyed
161        .inner
162        .unary::<ColumnBuilder<((UpsertKey, Row), T, Diff)>, _, _, _>(
163            Pipeline,
164            "Persist feedback encode",
165            |_, _| {
166                move |input, output| {
167                    input.for_each(|time, data| {
168                        let mut session = output.session_with_builder(&time);
169                        for ((key, value), ts, diff) in data.drain(..) {
170                            let row = upsert_value_to_row(&value);
171                            session.give(((&key, &row), &ts, &diff));
172                        }
173                    });
174                }
175            },
176        )
177}
mz_storage/upsert_continual_feedback_v2/feedback.rs

mz_storage/upsert_continual_feedback_v2/
feedback.rs
