differential_dataflow/batcher/merge/
mod.rs

1//! A `Batcher` implementation based on merge sort.
2//!
3//! The `MergeBatcher` requires a "merger" that implements the [`Merger`] trait, which provides
4//! hooks for manipulating sorted "chains" of chunks: merging chunks, and through [`Extractor`]
5//! also splitting them apart based on time.
6//!
7//! Raw input containers are fed to the batcher via [`Batcher::insert`], which chunks them with
8//! its `Chu` before merging: forming sorted, consolidated chunks is the first stage of the
9//! batcher's own work rather than something a caller arranges.
10//!
11//! The geometric ladder of chains is a [`MergeLadder`], which holds no opinion about how
12//! updates leave it, and can serve other batchers.
13
14pub mod chunker;
15
16use timely::container::{ContainerBuilder, PushInto};
17use timely::progress::frontier::AntichainRef;
18use timely::progress::{frontier::Antichain, Timestamp};
19
20use crate::logging::{BatcherEvent, Logger};
21use crate::batcher::{Batcher, Extract};
22
23/// Creates batches from chunks of sorted, consolidated tuples.
24///
25/// Chunking input is `Chu`'s business, merging chunks is `M`'s, and sealing the extracted chain
26/// into a batch is `S`'s; the batcher's own work is the carve-by-frontier over its
27/// [`MergeLadder`].
28pub struct MergeBatcher<Chu, M: Extractor, S> {
29    /// Melds input containers into sorted, consolidated chunks.
30    chunker: Chu,
31    /// Sorted, consolidated chains.
32    ladder: MergeLadder<M>,
33    /// The lower-bound frontier of the data, after the last call to extract.
34    frontier: Antichain<M::Time>,
35    /// Seals each extracted chain into a batch.
36    sealer: std::marker::PhantomData<S>,
37}
38
39impl<C, Chu, M, S> Batcher<C> for MergeBatcher<Chu, M, S>
40where
41    M: Extractor<Time: Timestamp>,
42    Chu: ContainerBuilder<Container = M::Chunk> + for<'a> PushInto<&'a mut C>,
43    S: Sealer<M::Chunk>,
44{
45    fn insert(&mut self, container: &mut C) {
46        self.chunker.push_into(container);
47        while let Some(chunk) = self.chunker.extract().map(std::mem::take) {
48            self.ladder.insert_chain(vec![chunk]);
49        }
50    }
51}
52
53impl<Chu, M, S> Extract for MergeBatcher<Chu, M, S>
54where
55    M: Extractor<Time: Timestamp>,
56    Chu: ContainerBuilder<Container = M::Chunk>,
57    S: Sealer<M::Chunk>,
58{
59    type Time = M::Time;
60    type Output = S::Output;
61
62    // Extraction means finding those updates with times not greater or equal to any time in
63    // `upper`. All updates must have time greater or equal to the previously used `upper`, by
64    // assumption that after extracting from a batcher we receive no more updates with times not
65    // greater or equal to `upper`.
66    fn extract<'a>(&'a mut self, upper: AntichainRef<'_, M::Time>) -> (Option<S::Output>, AntichainRef<'a, M::Time>) {
67        // Flush whatever the chunker is still accumulating: a partial final chunk would
68        // otherwise never reach the merge ladder.
69        while let Some(chunk) = self.chunker.finish().map(std::mem::take) {
70            self.ladder.insert_chain(vec![chunk]);
71        }
72
73        // Merge all remaining chains into a single chain.
74        let merged = self.ladder.merge_all();
75
76        // Extract readied data.
77        let mut kept = Vec::new();
78        let mut readied = Vec::new();
79        self.frontier.clear();
80
81        let ladder = &mut self.ladder;
82        ladder.merger.extract(merged, upper, &mut self.frontier, &mut readied, &mut kept, &mut ladder.stash);
83
84        if !kept.is_empty() {
85            ladder.chain_push(kept);
86        }
87
88        ladder.stash.clear();
89
90        (S::seal(&mut readied), self.frontier.borrow())
91    }
92}
93
94impl<Chu: Default, M: Extractor, S> MergeBatcher<Chu, M, S> {
95    /// Allocates a new empty batcher.
96    ///
97    /// The logger and operator identifier are used to report the batcher's memory footprint,
98    /// attributed to the operator that owns it.
99    pub fn new(logger: Option<Logger>, operator_id: usize) -> Self {
100        Self {
101            chunker: Chu::default(),
102            ladder: MergeLadder::new(logger, operator_id),
103            frontier: Antichain::new(),
104            sealer: std::marker::PhantomData,
105        }
106    }
107}
108
109/// A geometric ladder of sorted, consolidated chains, merged as they arrive.
110///
111/// Chains are kept in order of decreasing update count, each at least twice the next, so that
112/// there are logarithmically many and each update is merged logarithmically often.
113pub struct MergeLadder<M: Merger> {
114    /// Sorted, consolidated chains, each paired with its cached summed update count.
115    ///
116    /// The cached count is the chain's *merge weight*: the geometric ladder weighs
117    /// chains by updates, not chunk counts, since regrading decouples the two. A
118    /// chain is immutable until merged, so the weight is computed once at push.
119    ///
120    /// Do not push/pop directly but use the corresponding functions ([`Self::chain_push`]/[`Self::chain_pop`]).
121    chains: Vec<(usize, Vec<M::Chunk>)>,
122    /// Stash of empty chunks, recycled through the merging process.
123    stash: Vec<M::Chunk>,
124    /// Merges consolidated chunks.
125    merger: M,
126    /// Logger for size accounting.
127    logger: Option<Logger>,
128    /// Timely operator ID.
129    operator_id: usize,
130}
131
132impl<M: Merger> MergeLadder<M> {
133    /// Allocates a new empty ladder.
134    ///
135    /// The logger and operator identifier are used to report the ladder's memory footprint,
136    /// attributed to the operator that owns it.
137    pub fn new(logger: Option<Logger>, operator_id: usize) -> Self {
138        Self {
139            chains: Vec::new(),
140            stash: Vec::new(),
141            merger: M::default(),
142            logger,
143            operator_id,
144        }
145    }
146
147    /// Insert a chain and maintain chain properties: Chains are geometrically sized
148    /// (by summed updates) and ordered by decreasing update weight.
149    pub fn insert_chain(&mut self, chain: Vec<M::Chunk>) {
150        if !chain.is_empty() {
151            self.chain_push(chain);
152            while self.chains.len() > 1 && (self.chains[self.chains.len() - 1].0 >= self.chains[self.chains.len() - 2].0 / 2) {
153                let list1 = self.chain_pop().unwrap();
154                let list2 = self.chain_pop().unwrap();
155                let merged = self.merge_by(list1, list2);
156                self.chain_push(merged);
157            }
158        }
159    }
160
161    /// Merge all chains into a single chain, and remove it from the ladder.
162    pub fn merge_all(&mut self) -> Vec<M::Chunk> {
163        while self.chains.len() > 1 {
164            let list1 = self.chain_pop().unwrap();
165            let list2 = self.chain_pop().unwrap();
166            let merged = self.merge_by(list1, list2);
167            self.chain_push(merged);
168        }
169        self.chain_pop().unwrap_or_default()
170    }
171
172    /// The chains, largest first.
173    pub fn chains(&self) -> impl Iterator<Item = &[M::Chunk]> {
174        self.chains.iter().map(|(_, chain)| &chain[..])
175    }
176
177    /// Lets `logic` remove updates from each chain, largest first, for batchers that release
178    /// updates from the fronts of chains rather than by merging them all.
179    ///
180    /// A chain keeps the weight it was pushed with, so the ladder merges as though no updates had
181    /// left it, and stays logarithmic in the updates inserted. Chains left empty are removed.
182    pub fn retain_chains(&mut self, mut logic: impl FnMut(&mut Vec<M::Chunk>)) {
183        let mut chains = std::mem::take(&mut self.chains);
184        for (_weight, chain) in chains.iter_mut() {
185            self.account(chain.iter().map(Self::record), -1);
186            logic(chain);
187            self.account(chain.iter().map(Self::record), 1);
188        }
189        chains.retain(|(_weight, chain)| !chain.is_empty());
190        self.chains = chains;
191    }
192
193    // merges two sorted input lists into one sorted output list.
194    fn merge_by(&mut self, list1: Vec<M::Chunk>, list2: Vec<M::Chunk>) -> Vec<M::Chunk> {
195        // TODO: `list1` and `list2` get dropped; would be better to reuse?
196        let mut output = Vec::with_capacity(list1.len() + list2.len());
197        self.merger.merge(list1, list2, &mut output, &mut self.stash);
198
199        output
200    }
201
202    /// Pop a chain and account size changes.
203    #[inline]
204    fn chain_pop(&mut self) -> Option<Vec<M::Chunk>> {
205        let (_weight, chain) = self.chains.pop()?;
206        self.account(chain.iter().map(Self::record), -1);
207        Some(chain)
208    }
209
210    /// Push a chain and account size changes.
211    ///
212    /// Caches the chain's summed update count alongside it for the ladder.
213    #[inline]
214    fn chain_push(&mut self, chain: Vec<M::Chunk>) {
215        let weight = chain.iter().map(M::len).sum();
216        self.account(chain.iter().map(Self::record), 1);
217        self.chains.push((weight, chain));
218    }
219
220    /// The `(records, size, capacity, allocations)` logger tuple for one chunk,
221    /// assembled from the two focused `Merger` methods.
222    #[inline]
223    fn record(chunk: &M::Chunk) -> (usize, usize, usize, usize) {
224        let (size, capacity, allocations) = M::allocation(chunk);
225        (M::len(chunk), size, capacity, allocations)
226    }
227
228    /// Account size changes. Only performs work if a logger exists.
229    ///
230    /// Calculate the size based on the iterator passed along, with each attribute
231    /// multiplied by `diff`. Usually, one wants to pass 1 or -1 as the diff.
232    #[inline]
233    fn account<I: IntoIterator<Item = (usize, usize, usize, usize)>>(&self, items: I, diff: isize) {
234        if let Some(logger) = &self.logger {
235            let (mut records, mut size, mut capacity, mut allocations) = (0isize, 0isize, 0isize, 0isize);
236            for (records_, size_, capacity_, allocations_) in items {
237                records = records.saturating_add_unsigned(records_);
238                size = size.saturating_add_unsigned(size_);
239                capacity = capacity.saturating_add_unsigned(capacity_);
240                allocations = allocations.saturating_add_unsigned(allocations_);
241            }
242            logger.log(BatcherEvent {
243                operator: self.operator_id,
244                records_diff: records * diff,
245                size_diff: size * diff,
246                capacity_diff: capacity * diff,
247                allocations_diff: allocations * diff,
248            })
249        }
250    }
251}
252
253impl<M: Merger> Drop for MergeLadder<M> {
254    fn drop(&mut self) {
255        // Cleanup chain to retract accounting information.
256        while self.chain_pop().is_some() {}
257    }
258}
259
260/// A trait to describe interesting moments in a merge ladder.
261pub trait Merger: Default {
262    /// The internal representation of chunks of data.
263    type Chunk: Default;
264    /// Merge chains into an output chain.
265    fn merge(&mut self, list1: Vec<Self::Chunk>, list2: Vec<Self::Chunk>, output: &mut Vec<Self::Chunk>, stash: &mut Vec<Self::Chunk>);
266
267    /// The number of updates in a chunk.
268    ///
269    /// Drives the geometric ladder (chains are weighed by summed updates, not chunk
270    /// counts, since regrading decouples the two) and the `records` field of the
271    /// size logger.
272    fn len(chunk: &Self::Chunk) -> usize;
273
274    /// Backing-allocation figures for a chunk: `(size, capacity, allocations)`, for
275    /// the size logger's memory telemetry.
276    ///
277    /// Defaults to zero — most chunk types do not track this. Override to report
278    /// real figures (e.g. Materialize's memory accounting).
279    fn allocation(_chunk: &Self::Chunk) -> (usize, usize, usize) { (0, 0, 0) }
280}
281
282/// A [`Merger`] that can also split a chain apart by time, as the merge batcher extracts.
283pub trait Extractor: Merger {
284    /// The type of time in frontiers to extract updates.
285    type Time;
286    /// Extract ready updates based on the `upper` frontier.
287    fn extract(
288        &mut self,
289        merged: Vec<Self::Chunk>,
290        upper: AntichainRef<Self::Time>,
291        frontier: &mut Antichain<Self::Time>,
292        readied: &mut Vec<Self::Chunk>,
293        kept: &mut Vec<Self::Chunk>,
294        stash: &mut Vec<Self::Chunk>,
295    );
296}
297
298/// Forms a batch from a whole chain of updates at once.
299///
300/// Named rather than a bare `fn(&mut Vec<C>) -> Option<B>` so that implementors can name the
301/// batch they produce. There is no receiver: the chain goes in and the batch comes out, leaving
302/// nowhere for an update to be retained.
303pub trait Sealer<C> {
304    /// Output batch type.
305    type Output;
306
307    /// Builds a batch from a chain of updates.
308    ///
309    /// This method relies on the chain only containing updates greater or equal to the lower frontier,
310    /// and not greater or equal to the upper frontier, of the interval the caller means to describe.
311    /// Chains must also be sorted and consolidated.
312    ///
313    /// Having the whole chain in hand, an implementor can size itself before it fills.
314    fn seal(chain: &mut Vec<C>) -> Option<Self::Output>;
315}
316
317/// A `Merger` implementation for vector update containers.
318pub mod vec {
319
320    use std::marker::PhantomData;
321    use timely::container::SizableContainer;
322    use timely::progress::frontier::{Antichain, AntichainRef};
323    use timely::PartialOrder;
324    use crate::batcher::merge::{Extractor, Merger};
325
326    /// A `Merger` implementation for `Vec<(D, T, R)>` that drains owned inputs.
327    pub struct VecMerger<D, T, R> {
328        _marker: PhantomData<(D, T, R)>,
329    }
330
331    impl<D, T, R> Default for VecMerger<D, T, R> {
332        fn default() -> Self { Self { _marker: PhantomData } }
333    }
334
335    impl<D, T, R> VecMerger<D, T, R> {
336        /// The target capacity for output buffers, as a power of two.
337        ///
338        /// This amount is used to size vectors, where vectors not exactly this capacity are dropped.
339        /// If this is mis-set, there is the potential for more memory churn than anticipated.
340        fn target_capacity() -> usize {
341            timely::container::buffer::default_capacity::<(D, T, R)>().next_power_of_two()
342        }
343        /// Acquire a buffer with the target capacity.
344        fn empty(&self, stash: &mut Vec<Vec<(D, T, R)>>) -> Vec<(D, T, R)> {
345            let target = Self::target_capacity();
346            let mut container = stash.pop().unwrap_or_default();
347            container.clear();
348            // Reuse if at target; otherwise allocate fresh.
349            if container.capacity() != target {
350                container = Vec::with_capacity(target);
351            }
352            container
353        }
354        /// Refill `queue` from `iter` if empty. Recycles drained queues into `stash`.
355        fn refill(queue: &mut std::collections::VecDeque<(D, T, R)>, iter: &mut impl Iterator<Item = Vec<(D, T, R)>>, stash: &mut Vec<Vec<(D, T, R)>>) {
356            if queue.is_empty() {
357                let target = Self::target_capacity();
358                if stash.len() < 2 {
359                    let mut recycled = Vec::from(std::mem::take(queue));
360                    recycled.clear();
361                    if recycled.capacity() == target {
362                        stash.push(recycled);
363                    }
364                }
365                if let Some(chunk) = iter.next() {
366                    *queue = std::collections::VecDeque::from(chunk);
367                }
368            }
369        }
370    }
371
372    impl<D, T, R> Merger for VecMerger<D, T, R>
373    where
374        D: Ord + Clone + 'static,
375        T: Ord + Clone + PartialOrder + 'static,
376        R: crate::difference::Semigroup + 'static,
377    {
378        type Chunk = Vec<(D, T, R)>;
379
380        fn merge(
381            &mut self,
382            list1: Vec<Vec<(D, T, R)>>,
383            list2: Vec<Vec<(D, T, R)>>,
384            output: &mut Vec<Vec<(D, T, R)>>,
385            stash: &mut Vec<Vec<(D, T, R)>>,
386        ) {
387            use std::cmp::Ordering;
388            use std::collections::VecDeque;
389
390            let mut iter1 = list1.into_iter();
391            let mut iter2 = list2.into_iter();
392            let mut q1 = VecDeque::<(D,T,R)>::from(iter1.next().unwrap_or_default());
393            let mut q2 = VecDeque::<(D,T,R)>::from(iter2.next().unwrap_or_default());
394
395            let mut result = self.empty(stash);
396
397            // Merge while both queues are non-empty.
398            while let (Some((d1, t1, _)), Some((d2, t2, _))) = (q1.front(), q2.front()) {
399                match (d1, t1).cmp(&(d2, t2)) {
400                    Ordering::Less => {
401                        result.push(q1.pop_front().unwrap());
402                    }
403                    Ordering::Greater => {
404                        result.push(q2.pop_front().unwrap());
405                    }
406                    Ordering::Equal => {
407                        let (d, t, mut r1) = q1.pop_front().unwrap();
408                        let (_, _, r2) = q2.pop_front().unwrap();
409                        r1.plus_equals(&r2);
410                        if !r1.is_zero() {
411                            result.push((d, t, r1));
412                        }
413                    }
414                }
415
416                if result.at_capacity() {
417                    output.push(std::mem::take(&mut result));
418                    result = self.empty(stash);
419                }
420
421                // Refill emptied queues from their chains.
422                if q1.is_empty() { Self::refill(&mut q1, &mut iter1, stash); }
423                if q2.is_empty() { Self::refill(&mut q2, &mut iter2, stash); }
424            }
425
426            // Push partial result and remaining data from both sides.
427            if !result.is_empty() { output.push(result); }
428            for q in [q1, q2] {
429                if !q.is_empty() { output.push(Vec::from(q)); }
430            }
431            output.extend(iter1);
432            output.extend(iter2);
433        }
434
435        fn len(chunk: &Vec<(D, T, R)>) -> usize { chunk.len() }
436    }
437
438    impl<D, T, R> Extractor for VecMerger<D, T, R>
439    where
440        D: Ord + Clone + 'static,
441        T: Ord + Clone + PartialOrder + 'static,
442        R: crate::difference::Semigroup + 'static,
443    {
444        type Time = T;
445
446        fn extract(
447            &mut self,
448            merged: Vec<Vec<(D, T, R)>>,
449            upper: AntichainRef<T>,
450            frontier: &mut Antichain<T>,
451            ship: &mut Vec<Vec<(D, T, R)>>,
452            kept: &mut Vec<Vec<(D, T, R)>>,
453            stash: &mut Vec<Vec<(D, T, R)>>,
454        ) {
455            let mut keep = self.empty(stash);
456            let mut ready = self.empty(stash);
457
458            for mut chunk in merged {
459                // Go update-by-update to swap out full containers.
460                for (data, time, diff) in chunk.drain(..) {
461                    if upper.less_equal(&time) {
462                        frontier.insert_with(&time, |time| time.clone());
463                        keep.push((data, time, diff));
464                    } else {
465                        ready.push((data, time, diff));
466                    }
467                    if keep.at_capacity() {
468                        kept.push(std::mem::take(&mut keep));
469                        keep = self.empty(stash);
470                    }
471                    if ready.at_capacity() {
472                        ship.push(std::mem::take(&mut ready));
473                        ready = self.empty(stash);
474                    }
475                }
476                // Recycle the now-empty chunk if it has the right capacity.
477                if chunk.capacity() == Self::target_capacity() {
478                    stash.push(chunk);
479                }
480            }
481            if !keep.is_empty() { kept.push(keep); }
482            if !ready.is_empty() { ship.push(ready); }
483        }
484    }
485}
486
487#[cfg(test)]
488mod test {
489    use timely::progress::frontier::Antichain;
490    use crate::batcher::Extract;
491    use super::{MergeBatcher, Sealer};
492    use super::vec::VecMerger;
493    use crate::batcher::merge::chunker::ContainerChunker;
494
495    type In = Vec<((u64, ()), u64, i64)>;
496    type Bt = MergeBatcher<ContainerChunker<In>, VecMerger<(u64, ()), u64, i64>, TestSealer>;
497
498    struct TestSealer;
499
500    impl Sealer<In> for TestSealer {
501        type Output = Vec<In>;
502
503        fn seal(chain: &mut Vec<In>) -> Option<Self::Output> {
504            (!chain.is_empty()).then(|| std::mem::take(chain))
505        }
506    }
507
508    /// The sealed frontier must reflect the POST-CONSOLIDATION set of distinct kept times:
509    /// two chains carry cancelling updates at a kept time (`t=5`), plus a survivor at a later
510    /// kept time (`t=7`). After `extract(upper=[3])` the frontier must be `{7}` — `(100, 5)` nets
511    /// to zero and needs no capability. (A per-chain extract that folds the frontier before
512    /// consolidating would wrongly report `{5}`.)
513    #[test]
514    fn frontier_is_post_consolidation() {
515        let mut b = Bt::new(None, 0);
516        b.ladder.chain_push(vec![vec![((100u64, ()), 5u64, 1i64), ((200u64, ()), 7u64, 1i64)]]);
517        b.ladder.chain_push(vec![vec![((100u64, ()), 5u64, -1i64)]]);
518        let (_, retained) = Extract::extract(&mut b, Antichain::from_elem(3).borrow());
519        let got: Vec<u64> = retained.iter().cloned().collect();
520        assert_eq!(got, vec![7u64],
521            "frontier held a capability at t=5, which consolidates to zero (got {got:?})");
522    }
523
524    /// Sanity: with no cross-chain cancellation, the frontier is the minimal kept time.
525    #[test]
526    fn frontier_survivor_minimum() {
527        let mut b = Bt::new(None, 0);
528        b.ladder.chain_push(vec![vec![((100u64, ()), 5u64, 1i64)]]);
529        b.ladder.chain_push(vec![vec![((200u64, ()), 7u64, 1i64)]]);
530        let (_, retained) = Extract::extract(&mut b, Antichain::from_elem(3).borrow());
531        let got: Vec<u64> = retained.iter().cloned().collect();
532        assert_eq!(got, vec![5u64]);
533    }
534}
differential_dataflow/batcher/merge/mod.rs

differential_dataflow/batcher/merge/
mod.rs
