mz_compute/sink/
shared_batch.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//! Writing MV sink batches through the process-global [`SharedBatches`].
11//!
12//! The workers running in one process coalesce their updates for a given batch interval into a
13//! single batch instead of each writing its own small batch. All workers building an interval
14//! share `desc.shared_id`, because the `mint` operator broadcasts one description.
15//!
16//! One task per process encodes and uploads a shared batch, so routing large intervals through it
17//! serializes work the workers otherwise do in parallel. Each worker therefore stages its updates
18//! and hands them to the shared batch only if they stay below a size threshold. Past the
19//! threshold, it writes its own batch.
20
21use std::sync::Arc;
22
23use mz_dyncfg::ConfigValHandle;
24use mz_persist_client::batch::{BatchBuilder, ProtoBatch};
25use mz_persist_client::write::WriteHandle;
26use mz_persist_client::{PersistClient, Schemas};
27use mz_persist_types::ShardId;
28use mz_persist_types::codec_impls::UnitSchema;
29use mz_persist_types::part::PartBuilder;
30use mz_repr::{RelationDesc, Timestamp};
31use mz_storage_operators::persist::{SharedBatchBuilder, SharedBatches};
32use mz_storage_types::StorageDiff;
33use mz_storage_types::sources::SourceData;
34
35use crate::sink::materialized_view::{BatchDescription, PersistApi};
36
37/// A sink's handle on the process-global [`SharedBatches`] for its output shard.
38pub(super) struct SharedBatchWriter {
39    persist_client: PersistClient,
40    shard_id: ShardId,
41    shard_name: String,
42    schemas: Schemas<SourceData, ()>,
43    shared_batches: SharedBatches,
44    max_bytes: ConfigValHandle<usize>,
45}
46
47impl SharedBatchWriter {
48    /// Creates a writer for the shard `writer` writes to, which shares a worker's updates for an
49    /// interval while they stay below `max_bytes`. Zero disables sharing.
50    pub(super) async fn new(
51        persist_api: &PersistApi,
52        writer: &WriteHandle<SourceData, (), Timestamp, StorageDiff>,
53        max_bytes: ConfigValHandle<usize>,
54    ) -> Self {
55        Self {
56            persist_client: persist_api.open_client().await,
57            shard_id: persist_api.collection.data_shard,
58            shard_name: persist_api.shard_name.clone(),
59            // The handle resolves the registered id for this sink's schema.
60            schemas: Schemas {
61                id: writer.schema_id(),
62                key: Arc::new(persist_api.collection.relation_desc.clone()),
63                val: Arc::new(UnitSchema),
64            },
65            shared_batches: persist_api.persist_batches.clone(),
66            max_bytes,
67        }
68    }
69
70    /// Starts this worker's write for `desc`.
71    pub(super) fn start(&self, desc: &BatchDescription) -> SharedBatchWrite<'_> {
72        // Read per interval, so a threshold change applies without re-rendering the sink. Workers may
73        // disagree on an interval: each either contributes to the shared batch or writes its own,
74        // and the append operator appends all of them together. With a threshold of zero, a
75        // writer joins no shared batch at all, so it starts no shared builder task.
76        let max_bytes = self.max_bytes.get();
77        let shared = (max_bytes > 0).then(|| {
78            self.shared_batches.builder(
79                desc.shared_id,
80                self.persist_client.clone(),
81                self.shard_id,
82                self.shard_name.clone(),
83                self.schemas.clone(),
84                desc.lower.clone(),
85                desc.upper.clone(),
86            )
87        });
88        SharedBatchWrite {
89            writer: self,
90            desc: desc.clone(),
91            max_bytes,
92            staged: PartBuilder::new(&*self.schemas.key, &*self.schemas.val),
93            own: None,
94            shared,
95        }
96    }
97
98    fn own_builder(
99        &self,
100        desc: &BatchDescription,
101    ) -> BatchBuilder<SourceData, (), Timestamp, StorageDiff> {
102        self.persist_client.batch_builder(
103            self.shard_id,
104            &self.shard_name,
105            self.schemas.clone(),
106            desc.lower.clone(),
107            None,
108        )
109    }
110}
111
112/// One worker's write for one interval, which contributes to the shared batch until the worker's
113/// updates reach the size threshold and writes the worker's own batch after. With sharing disabled,
114/// all updates go to the worker's own batch.
115///
116/// The caller must [`finish`](Self::finish) it, also when it pushed nothing (see
117/// [`SharedBatchBuilder::finish`]).
118pub(super) struct SharedBatchWrite<'a> {
119    writer: &'a SharedBatchWriter,
120    desc: BatchDescription,
121    max_bytes: usize,
122    /// Updates not yet handed to either batch. Empty once `own` exists, and always empty without
123    /// `shared`.
124    staged: PartBuilder<SourceData, RelationDesc, (), UnitSchema>,
125    /// Created on the first update that does not go to the shared batch, so a worker without
126    /// updates writes no batch.
127    own: Option<BatchBuilder<SourceData, (), Timestamp, StorageDiff>>,
128    shared: Option<SharedBatchBuilder>,
129}
130
131impl SharedBatchWrite<'_> {
132    /// Adds an update.
133    pub(super) async fn push(&mut self, (k, v): &(SourceData, ()), t: Timestamp, d: StorageDiff) {
134        if self.own.is_none() {
135            if self.shared.is_some() {
136                self.staged.push(k, v, t, d);
137                if self.staged.goodbytes() < self.max_bytes {
138                    return;
139                }
140                let writer = self.writer;
141                let staged = self
142                    .staged
143                    .finish_and_replace(&*writer.schemas.key, &*writer.schemas.val);
144                let mut own = writer.own_builder(&self.desc);
145                own.add_part(staged).await.expect("valid usage");
146                self.own = Some(own);
147                return;
148            }
149            self.own = Some(self.writer.own_builder(&self.desc));
150        }
151        let own = self.own.as_mut().expect("created above");
152        own.add(k, v, &t, &d).await.expect("valid usage");
153    }
154
155    /// Hands over the staged updates and returns the batches this worker delivers: its own batch
156    /// if it has one, and the shared batch if it is the last worker to finish.
157    pub(super) async fn finish(self) -> Vec<ProtoBatch> {
158        let own = match self.own {
159            Some(own) => Some(
160                own.finish(self.desc.upper.clone())
161                    .await
162                    .expect("valid usage"),
163            ),
164            None => None,
165        };
166        let shared = match self.shared {
167            Some(shared) => {
168                shared.push(self.staged.finish()).await;
169                shared.finish().await
170            }
171            None => None,
172        };
173        own.into_iter()
174            .chain(shared)
175            .map(|batch| batch.into_transmittable_batch())
176            .collect()
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use mz_dyncfg::ConfigValHandle;
183    use mz_persist_client::Diagnostics;
184    use mz_repr::{Datum, Row, SqlScalarType};
185    use timely::progress::Antichain;
186
187    use super::*;
188
189    fn desc() -> RelationDesc {
190        RelationDesc::builder()
191            .with_column("a", SqlScalarType::Int64.nullable(false))
192            .finish()
193    }
194
195    fn writer(
196        client: &PersistClient,
197        shared_batches: &SharedBatches,
198        shard_id: ShardId,
199        max_bytes: usize,
200    ) -> SharedBatchWriter {
201        SharedBatchWriter {
202            persist_client: client.clone(),
203            shard_id,
204            shard_name: "test".to_string(),
205            schemas: Schemas {
206                id: None,
207                key: Arc::new(desc()),
208                val: Arc::new(UnitSchema),
209            },
210            shared_batches: shared_batches.clone(),
211            max_bytes: ConfigValHandle::disconnected(max_bytes),
212        }
213    }
214
215    fn update(i: i64) -> (SourceData, ()) {
216        (SourceData(Ok(Row::pack_slice(&[Datum::Int64(i)]))), ())
217    }
218
219    async fn update_counts(
220        client: &PersistClient,
221        shard_id: ShardId,
222        batches: Vec<ProtoBatch>,
223    ) -> Vec<usize> {
224        let write = client
225            .open_writer::<SourceData, (), Timestamp, StorageDiff>(
226                shard_id,
227                Arc::new(desc()),
228                Arc::new(UnitSchema),
229                Diagnostics::for_tests(),
230            )
231            .await
232            .expect("valid usage");
233        batches
234            .into_iter()
235            .map(|proto| {
236                write
237                    .batch_from_transmittable_batch(proto)
238                    .into_hollow_batch()
239                    .len
240            })
241            .collect()
242    }
243
244    #[mz_ore::test(tokio::test)]
245    async fn test_shared_batch_write_crosses_threshold() {
246        let client = PersistClient::new_for_tests().await;
247        let shared_batches = SharedBatches::new();
248        let shard_id = ShardId::new();
249        let small = writer(&client, &shared_batches, shard_id, usize::MAX);
250        let large = writer(&client, &shared_batches, shard_id, 1);
251        let desc = BatchDescription::new(
252            Antichain::from_elem(Timestamp::from(0)),
253            Antichain::from_elem(Timestamp::from(1)),
254            0,
255        );
256
257        let mut small_write = small.start(&desc);
258        let mut large_write = large.start(&desc);
259        small_write.push(&update(0), Timestamp::from(0), 1).await;
260        // The first update moves the staged part to the worker's own batch via `add_part`, and
261        // the second goes there via `add`.
262        large_write.push(&update(1), Timestamp::from(0), 1).await;
263        large_write.push(&update(2), Timestamp::from(0), 1).await;
264
265        let small_batches = small_write.finish().await;
266        assert!(
267            small_batches.is_empty(),
268            "a worker below the threshold that finishes first delivers nothing"
269        );
270        let large_batches = large_write.finish().await;
271        assert_eq!(
272            update_counts(&client, shard_id, large_batches).await,
273            vec![2, 1],
274            "the last worker delivers its own batch and the shared batch"
275        );
276    }
277
278    #[mz_ore::test(tokio::test)]
279    async fn test_shared_batch_write_zero_threshold() {
280        let client = PersistClient::new_for_tests().await;
281        let shared_batches = SharedBatches::new();
282        let shard_id = ShardId::new();
283        let desc = BatchDescription::new(
284            Antichain::from_elem(Timestamp::from(0)),
285            Antichain::from_elem(Timestamp::from(1)),
286            0,
287        );
288        let writer = writer(&client, &shared_batches, shard_id, 0);
289        let mut write = writer.start(&desc);
290        assert!(
291            write.shared.is_none(),
292            "a zero threshold joins no shared batch"
293        );
294        write.push(&update(0), Timestamp::from(0), 1).await;
295        assert_eq!(
296            update_counts(&client, shard_id, write.finish().await).await,
297            vec![1],
298            "a zero threshold writes one own batch"
299        );
300    }
301}
mz_compute/sink/shared_batch.rs

mz_compute/sink/
shared_batch.rs
