mz_row_spine/
arc_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//! An `Arc`-backed batch newtype whose contents can be shared across timely runtimes.
11//!
12//! Differential's default spines reference-count their batches with `Rc`, which is worker-local.
13//! Sharing an arrangement with another runtime (a reader on a different worker thread) needs the
14//! batches behind an `Arc` so a batch whose contents are `Send + Sync` can be read from that other
15//! thread.
16//!
17//! The blanket `impl Trait for Arc<B>` that would express this lives outside this crate: both `Arc`
18//! and differential's `SpineBatch`/`Builder`/`Merger`/`Cursor` traits are foreign, so the orphan rule
19//! forbids it here. [`ArcBatch`] is a local newtype around `Arc<B>` that carries those impls
20//! instead. The impls delegate straight through to the inner batch, so `ArcBatch<B>` behaves
21//! exactly like `B` except that its handle is atomically reference counted.
22//!
23//! This mirrors differential's own `rc_blanket_impls` (for `Rc<B>`), swapping in `Arc`. Keeping it
24//! here as a newtype lets cross-thread arrangement sharing build against a released
25//! differential-dataflow, with no differential-side `Arc` batch impls required.
26
27use std::sync::Arc;
28
29use differential_dataflow::trace::implementations::merge_batcher::Sealer;
30use differential_dataflow::trace::implementations::spine_fueled::{Merger, SpineBatch};
31use differential_dataflow::trace::{Builder, Cursor, Navigable};
32use timely::progress::frontier::AntichainRef;
33
34/// An `Arc`-backed batch, shareable across threads when `B`'s contents are `Send + Sync`.
35///
36/// A transparent newtype around `Arc<B>`. Cloning shares the underlying batch, exactly like the
37/// `Rc`-backed default, but with atomic reference counting.
38pub struct ArcBatch<B>(pub Arc<B>);
39
40// Hand-written rather than derived: `#[derive(Clone)]` would bound `B: Clone`, but `Arc<B>` is
41// `Clone` for any `B` (it clones the handle, not the batch). The derived bound would make
42// `ArcBatch<B>: Clone` fail for a non-`Clone` batch such as `OrdValBatch`, which in turn breaks
43// `Spine<ArcBatch<B>>: TraceReader`.
44impl<B> Clone for ArcBatch<B> {
45    fn clone(&self) -> Self {
46        ArcBatch(Arc::clone(&self.0))
47    }
48}
49
50impl<B> ArcBatch<B> {
51    /// Wraps a batch in an `Arc`.
52    pub fn new(batch: B) -> Self {
53        ArcBatch(Arc::new(batch))
54    }
55}
56
57impl<B> std::ops::Deref for ArcBatch<B> {
58    type Target = B;
59    fn deref(&self) -> &B {
60        &self.0
61    }
62}
63
64impl<B: Navigable> Navigable for ArcBatch<B> {
65    type Cursor = ArcBatchCursor<B::Cursor>;
66    fn cursor(&self) -> Self::Cursor {
67        // Disambiguate to the inner batch's cursor, reached through the `Deref`, so the wrapper's
68        // `Cursor` is `B`'s rather than any impl that might exist on `Arc<B>` itself.
69        ArcBatchCursor::new(<B as Navigable>::cursor(&self.0))
70    }
71}
72
73impl<B: SpineBatch> SpineBatch for ArcBatch<B> {
74    type Time = B::Time;
75    type Merger = ArcMerger<B>;
76    fn len(&self) -> usize {
77        self.0.len()
78    }
79}
80
81/// Cursor over an [`ArcBatch`], delegating to the inner batch's cursor.
82pub struct ArcBatchCursor<C> {
83    cursor: C,
84}
85
86impl<C> ArcBatchCursor<C> {
87    fn new(cursor: C) -> Self {
88        ArcBatchCursor { cursor }
89    }
90}
91
92impl<C: Cursor> Cursor for ArcBatchCursor<C> {
93    type Storage = ArcBatch<C::Storage>;
94
95    type Key<'a> = C::Key<'a>;
96    type ValOwn = C::ValOwn;
97    type Val<'a> = C::Val<'a>;
98    type Time = C::Time;
99    type TimeGat<'a> = C::TimeGat<'a>;
100    type Diff = C::Diff;
101    type DiffGat<'a> = C::DiffGat<'a>;
102    type KeyContainer = C::KeyContainer;
103    type ValContainer = C::ValContainer;
104    type TimeContainer = C::TimeContainer;
105    type DiffContainer = C::DiffContainer;
106
107    #[inline]
108    fn key_valid(&self, storage: &Self::Storage) -> bool {
109        self.cursor.key_valid(&storage.0)
110    }
111    #[inline]
112    fn val_valid(&self, storage: &Self::Storage) -> bool {
113        self.cursor.val_valid(&storage.0)
114    }
115
116    #[inline]
117    fn key<'a>(&self, storage: &'a Self::Storage) -> Self::Key<'a> {
118        self.cursor.key(&storage.0)
119    }
120    #[inline]
121    fn val<'a>(&self, storage: &'a Self::Storage) -> Self::Val<'a> {
122        self.cursor.val(&storage.0)
123    }
124
125    #[inline]
126    fn get_key<'a>(&self, storage: &'a Self::Storage) -> Option<Self::Key<'a>> {
127        self.cursor.get_key(&storage.0)
128    }
129    #[inline]
130    fn get_val<'a>(&self, storage: &'a Self::Storage) -> Option<Self::Val<'a>> {
131        self.cursor.get_val(&storage.0)
132    }
133
134    #[inline]
135    fn map_times<L: FnMut(Self::TimeGat<'_>, Self::DiffGat<'_>)>(
136        &mut self,
137        storage: &Self::Storage,
138        logic: L,
139    ) {
140        self.cursor.map_times(&storage.0, logic)
141    }
142
143    #[inline]
144    fn step_key(&mut self, storage: &Self::Storage) {
145        self.cursor.step_key(&storage.0)
146    }
147    #[inline]
148    fn seek_key(&mut self, storage: &Self::Storage, key: Self::Key<'_>) {
149        self.cursor.seek_key(&storage.0, key)
150    }
151
152    #[inline]
153    fn step_val(&mut self, storage: &Self::Storage) {
154        self.cursor.step_val(&storage.0)
155    }
156    #[inline]
157    fn seek_val(&mut self, storage: &Self::Storage, val: Self::Val<'_>) {
158        self.cursor.seek_val(&storage.0, val)
159    }
160
161    #[inline]
162    fn rewind_keys(&mut self, storage: &Self::Storage) {
163        self.cursor.rewind_keys(&storage.0)
164    }
165    #[inline]
166    fn rewind_vals(&mut self, storage: &Self::Storage) {
167        self.cursor.rewind_vals(&storage.0)
168    }
169}
170
171/// Builds [`ArcBatch`]es, delegating to the inner batch's builder.
172pub struct ArcBuilder<B> {
173    builder: B,
174}
175
176impl<B: Default> Default for ArcBuilder<B> {
177    fn default() -> Self {
178        ArcBuilder {
179            builder: B::default(),
180        }
181    }
182}
183
184impl<B: Builder> Builder for ArcBuilder<B> {
185    type Input = B::Input;
186    type Time = B::Time;
187    type Output = ArcBatch<B::Output>;
188    fn push(&mut self, input: &mut Self::Input) {
189        self.builder.push(input)
190    }
191    fn done(self) -> Option<ArcBatch<B::Output>> {
192        self.builder.done().map(ArcBatch::new)
193    }
194}
195
196impl<C, B: Sealer<C>> Sealer<C> for ArcBuilder<B> {
197    type Output = ArcBatch<B::Output>;
198    fn seal(chain: &mut Vec<C>) -> Option<Self::Output> {
199        B::seal(chain).map(ArcBatch::new)
200    }
201}
202
203impl<B: Builder + mz_timely_util::columnar::chunk::ChainState>
204    mz_timely_util::columnar::chunk::ChainState for ArcBuilder<B>
205{
206    type State = B::State;
207
208    fn wants_bodies() -> bool {
209        B::wants_bodies()
210    }
211
212    fn observe(state: &mut Self::State, input: &Self::Input) {
213        B::observe(state, input)
214    }
215
216    fn observe_records(state: &mut Self::State, records: usize) {
217        B::observe_records(state, records)
218    }
219
220    fn from_state(state: Self::State) -> Self {
221        ArcBuilder {
222            builder: B::from_state(state),
223        }
224    }
225}
226
227/// Merges [`ArcBatch`]es, delegating to the inner batch's merger.
228pub struct ArcMerger<B: SpineBatch> {
229    merger: B::Merger,
230}
231
232impl<B: SpineBatch> Merger<ArcBatch<B>> for ArcMerger<B> {
233    fn new(
234        source1: &ArcBatch<B>,
235        source2: &ArcBatch<B>,
236        compaction_frontier: AntichainRef<B::Time>,
237    ) -> Self {
238        ArcMerger {
239            merger: B::Merger::new(&source1.0, &source2.0, compaction_frontier),
240        }
241    }
242    fn work(&mut self, source1: &ArcBatch<B>, source2: &ArcBatch<B>, fuel: &mut isize) {
243        self.merger.work(&source1.0, &source2.0, fuel)
244    }
245    fn done(self) -> Option<ArcBatch<B>> {
246        self.merger.done().map(ArcBatch::new)
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use differential_dataflow::batcher::{Batcher, Extract};
253    use differential_dataflow::trace::Navigable;
254    use differential_dataflow::trace::cursor::Cursor;
255    use timely::progress::Antichain;
256
257    use crate::ArcOrdValBatcher;
258
259    /// An `ArcBatch`'s cursor can be constructed and read from a thread other than the one that
260    /// built it, proving the newtype's batches are usable across a thread boundary. This is the
261    /// property that lets [`crate::ArcOrdValSpine`] (and the `RowRow`/`Err` spines built on
262    /// [`ArcBatch`]) back a cross-runtime shared trace; the default `Rc`-backed spines are
263    /// worker-local by design and do not have it.
264    ///
265    /// Mirrors differential-dataflow's own `tests/trace.rs` cross-thread batch read, over the local
266    /// [`ArcBatch`] newtype.
267    #[mz_ore::test]
268    fn arc_batch_reads_from_other_thread() {
269        fn assert_send_sync<T: Send + Sync>(_: &T) {}
270
271        let mut batcher = ArcOrdValBatcher::<u64, u64, usize, i64>::new(None, 0);
272        let mut updates: Vec<((u64, u64), usize, i64)> = vec![((1, 2), 0, 1), ((2, 3), 1, 1)];
273        batcher.insert(&mut updates);
274        let (batch, _frontier) = batcher.extract(Antichain::from_elem(2).borrow());
275        let batch = batch.expect("updates were pushed, so the batch is non-empty");
276
277        assert_send_sync(&batch);
278
279        let read = std::thread::spawn(move || {
280            let mut cursor = batch.cursor();
281            cursor.to_vec(&batch, |k| *k, |v| *v)
282        })
283        .join()
284        .expect("reader thread panicked");
285
286        assert_eq!(read, vec![((1, 2), vec![(0, 1)]), ((2, 3), vec![(1, 1)])]);
287    }
288}
mz_row_spine/arc_batch.rs

mz_row_spine/
arc_batch.rs
