mz_compute/render/
sinks.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//! Logic related to the creation of dataflow sinks.
11
12use std::any::Any;
13use std::collections::{BTreeMap, BTreeSet};
14use std::rc::Rc;
15
16use differential_dataflow::VecCollection;
17use mz_compute_types::sinks::{ComputeSinkConnection, ComputeSinkDesc};
18use mz_expr::EvalError;
19use mz_ore::soft_assert_or_log;
20use mz_ore::str::StrExt;
21use mz_ore::vec::PartialOrdVecExt;
22use mz_repr::{Diff, GlobalId, Row};
23use mz_storage_types::controller::CollectionMetadata;
24use mz_timely_util::operator::CollectionExt;
25use mz_timely_util::probe::Handle;
26use timely::container::CapacityContainerBuilder;
27use timely::dataflow::Scope;
28use timely::progress::Antichain;
29
30use crate::compute_state::SinkToken;
31use crate::logging::compute::LogDataflowErrors;
32use crate::render::columnar::columnar_to_vec;
33use crate::render::context::{Context, distinct_errs_collection};
34use crate::render::errors::DataflowErrorSer;
35use crate::render::{RenderTimestamp, StartSignal};
36
37impl<'g, T: RenderTimestamp> Context<'g, T> {
38    /// Export the sink described by `sink` from the rendering context.
39    pub(crate) fn export_sink(
40        &self,
41        compute_state: &mut crate::compute_state::ComputeState,
42        tokens: &BTreeMap<GlobalId, Rc<dyn Any>>,
43        dependency_ids: BTreeSet<GlobalId>,
44        sink_id: GlobalId,
45        sink: &ComputeSinkDesc<CollectionMetadata>,
46        start_signal: StartSignal,
47        output_probe: &Handle<mz_repr::Timestamp>,
48        outer_scope: Scope<'g, mz_repr::Timestamp>,
49    ) {
50        soft_assert_or_log!(
51            sink.non_null_assertions.is_strictly_sorted(),
52            "non-null assertions not sorted"
53        );
54
55        // put together tokens that belong to the export
56        let mut needed_tokens = Vec::new();
57        for dep_id in dependency_ids {
58            if let Some(token) = tokens.get(&dep_id) {
59                needed_tokens.push(Rc::clone(token))
60            }
61        }
62
63        // Which arrangement of the `from` collection to consume, and the MFP that reconstructs
64        // full rows from it, were decided during LIR lowering and recorded in
65        // `sink.from_arrangement`. An unarranged form takes precedence when the bundle provides
66        // one. Lowering promises one exactly when `from_arrangement` is `None`, but the renderer
67        // can produce one where lowering promised only an arrangement: a snapshot-excluded
68        // subscribe imports an index as a filtered collection instead (see `import_index`).
69        let bundle = self
70            .lookup_id(mz_expr::Id::Global(sink.from))
71            .expect("Sink source collection not loaded");
72        let (ok_collection, mut err_collection) = match (&bundle.collection, &sink.from_arrangement)
73        {
74            (Some((oks, errs)), _) => (columnar_to_vec(oks.clone()), errs.clone()),
75            (None, Some((key, mfp_plan))) => {
76                // The sink serializes rows, so decode to `Vec` here. This is the
77                // sanctioned sink leaf, the same seam as the raw-collection arm
78                // above.
79                let (oks, errs) = bundle.as_collection_core(
80                    mfp_plan.clone(),
81                    Some((key.clone(), None)),
82                    self.until.clone(),
83                );
84                (columnar_to_vec(oks), errs)
85            }
86            (None, None) => panic!(
87                "sink source {} has neither a collection nor a planned arrangement",
88                sink.from
89            ),
90        };
91
92        // Attach logging of dataflow errors.
93        if let Some(logger) = compute_state.compute_logger.clone() {
94            err_collection = err_collection.log_dataflow_errors(logger, sink_id);
95        }
96
97        let mut ok_collection = ok_collection.leave(outer_scope);
98        let mut err_collection = err_collection.leave(outer_scope);
99
100        // Ensure that the frontier does not advance past the expiration time, if set. Otherwise,
101        // we might write down incorrect data.
102        if let Some(&expiration) = self.dataflow_expiration.as_option() {
103            ok_collection = ok_collection
104                .expire_collection_at(&format!("{}_export_sink_oks", self.debug_name), expiration);
105            err_collection = err_collection
106                .expire_collection_at(&format!("{}_export_sink_errs", self.debug_name), expiration);
107        }
108
109        let non_null_assertions = sink.non_null_assertions.clone();
110        let from_desc = sink.from_desc.clone();
111        if !non_null_assertions.is_empty() {
112            let name = format!("NullAssertions({sink_id:?})");
113            type CB<C> = CapacityContainerBuilder<C>;
114            let (oks, null_errs) =
115                ok_collection.map_fallible::<CB<_>, CB<_>, _, _, _>(&name, move |row| {
116                    let mut idx = 0;
117                    let mut iter = row.iter();
118                    for &i in &non_null_assertions {
119                        let skip = i - idx;
120                        let datum = iter.nth(skip).unwrap();
121                        idx += skip + 1;
122                        if datum.is_null() {
123                            return Err(DataflowErrorSer::from(EvalError::MustNotBeNull(
124                                format!("column {}", from_desc.get_name(i).quoted()).into(),
125                            )));
126                        }
127                    }
128                    Ok(row)
129                });
130            ok_collection = oks;
131            err_collection = err_collection.concat(null_errs);
132        }
133
134        // Normalize what the sink persists. A materialized view's error multiplicity is durable
135        // state another dataflow reads back verbatim, so leaving it at this dataflow's fan-out both
136        // multiplies across the object graph and makes the written value depend on plan shape.
137        // Placed after the null assertions so their errors, whose multiplicity follows the ok row's,
138        // are covered too. Only the persist-backed sink is re-importable: a subscribe or a one-shot
139        // copy is read once by a client, so neither pays for the arrangement.
140        if matches!(sink.connection, ComputeSinkConnection::MaterializedView(_)) {
141            err_collection = distinct_errs_collection(err_collection);
142        }
143
144        let region_name = match sink.connection {
145            ComputeSinkConnection::Subscribe(_) => format!("SubscribeSink({:?})", sink_id),
146            ComputeSinkConnection::MaterializedView(_) => {
147                format!("MaterializedViewSink({:?})", sink_id)
148            }
149            ComputeSinkConnection::CopyToS3Oneshot(_) => {
150                format!("CopyToS3OneshotSink({:?})", sink_id)
151            }
152            ComputeSinkConnection::MetricSink(_) => format!("MetricSink({:?})", sink_id),
153        };
154        outer_scope.clone().region_named(&region_name, |inner| {
155            let sink_render = get_sink_render_for(&sink.connection);
156
157            let sink_token = sink_render.render_sink(
158                compute_state,
159                sink,
160                sink_id,
161                self.as_of_frontier.clone(),
162                start_signal,
163                ok_collection.enter_region(inner),
164                err_collection.enter_region(inner),
165                output_probe,
166            );
167
168            if let Some(sink_token) = sink_token {
169                needed_tokens.push(sink_token);
170            }
171
172            let collection = compute_state.expect_collection_mut(sink_id);
173            collection.sink_token = Some(SinkToken::new(Box::new(needed_tokens)));
174        });
175    }
176}
177
178/// A type that can be rendered as a dataflow sink.
179pub(crate) trait SinkRender<'scope> {
180    fn render_sink(
181        &self,
182        compute_state: &mut crate::compute_state::ComputeState,
183        sink: &ComputeSinkDesc<CollectionMetadata>,
184        sink_id: GlobalId,
185        as_of: Antichain<mz_repr::Timestamp>,
186        start_signal: StartSignal,
187        sinked_collection: VecCollection<'scope, mz_repr::Timestamp, Row, Diff>,
188        err_collection: VecCollection<'scope, mz_repr::Timestamp, DataflowErrorSer, Diff>,
189        output_probe: &Handle<mz_repr::Timestamp>,
190    ) -> Option<Rc<dyn Any>>;
191}
192
193fn get_sink_render_for<'scope>(
194    connection: &ComputeSinkConnection<CollectionMetadata>,
195) -> Box<dyn SinkRender<'scope>> {
196    match connection {
197        ComputeSinkConnection::Subscribe(connection) => Box::new(connection.clone()),
198        ComputeSinkConnection::MaterializedView(connection) => Box::new(connection.clone()),
199        ComputeSinkConnection::CopyToS3Oneshot(connection) => Box::new(connection.clone()),
200        ComputeSinkConnection::MetricSink(connection) => Box::new(connection.clone()),
201    }
202}
mz_compute/render/sinks.rs

mz_compute/render/
sinks.rs
