Skip to main content

differential_dataflow/columnar/trace/
chunker.rs

1//! The container→batch chunker for the columnar trace.
2//!
3//! [`TrieChunker`] is the `ContainerBuilder` that turns a stream of
4//! [`RecordedUpdates`](crate::columnar::collection::RecordedUpdates) into
5//! [`ColChunk`](crate::columnar::trace::ColChunk) batches for the `Chunk`
6//! harness's batcher. It accumulates small inputs, consolidates, and ships
7//! chunks sized within 1-2x [`LINK_TARGET`](crate::columnar::LINK_TARGET).
8
9use super::ColChunk;
10
11use crate::columnar::updates::UpdatesTyped;
12use crate::columnar::collection::RecordedUpdates;
13
14/// A chunker that unwraps `RecordedUpdates` into bare `UpdatesTyped` for the merge batcher.
15///
16/// The intended behavior is to produce chunks whose size is within 1-2x `LINK_TARGET`.
17/// It ships large batches immediately, accumulates small batches, consolidates as they
18/// exceed 2xLINK_TARGET, and ships them unless they drop below 1xLINK_TARGET.
19///
20/// The flow is into (or around) `self.stage`, then consolidated blocks into `self.ready`,
21/// each of which is put in `self.stage`
22pub struct TrieChunker<U: crate::columnar::layout::ColumnarUpdate> {
23    /// Insufficiently large updates we haven't figured out how to ship yet.
24    blobs: Vec<(UpdatesTyped<U>, bool)>,
25    /// Sum of `len()` across `blobs`.
26    blob_records: usize,
27    /// Ready-to-emit chunks. Each is sorted and consolidated; size ≥ `LINK_TARGET`
28    /// (or smaller, only for the final chunk produced by `finish`).
29    ready: std::collections::VecDeque<UpdatesTyped<U>>,
30    /// Staging area for the next pull call: the ready trie wrapped as a `ColChunk`.
31    stage: Option<ColChunk<U>>,
32}
33
34impl<U: crate::columnar::layout::ColumnarUpdate> Default for TrieChunker<U> {
35    fn default() -> Self {
36        Self {
37            blobs: Default::default(),
38            blob_records: 0,
39            ready: Default::default(),
40            stage: None,
41        }
42    }
43}
44
45impl<U: crate::columnar::layout::ColumnarUpdate> TrieChunker<U> {
46    /// Consolidate and empty `self.blobs`, into `self.ready` if large enough or else return.
47    fn consolidate_blobs(&mut self) -> UpdatesTyped<U> {
48        // Single consolidated entry: pass through, no work.
49        if self.blobs.len() == 1 && self.blobs[0].1 {
50            let (result, _) = self.blobs.pop().unwrap();
51            self.blob_records = 0;
52            return result;
53        }
54
55        // TODO: Improve consolidation through column-oriented sorts.
56        let result = UpdatesTyped::<U>::form_unsorted(self.blobs.iter().flat_map(|(u, _)| u.iter()));
57        self.blobs.clear();
58        self.blob_records = 0;
59        result
60    }
61
62    /// Push a non-empty `UpdatesTyped` into blobs and update accounting.
63    fn absorb(&mut self, updates: UpdatesTyped<U>, consolidated: bool) {
64        self.blob_records += updates.len();
65        self.blobs.push((updates, consolidated));
66    }
67}
68
69impl<'a, U: crate::columnar::layout::ColumnarUpdate> timely::container::PushInto<&'a mut RecordedUpdates<U>> for TrieChunker<U> {
70    fn push_into(&mut self, container: &'a mut RecordedUpdates<U>) {
71        // Early return if an empty container (legit, for accountable progress tracking).
72        if container.updates.len() == 0 { return; }
73
74        // Our main goal is to only ship links that are 1-2 x LINK_TARGET, using blobs
75        // to accumulate updates until they are ready to go or we are asked to finish.
76        //
77        // Informally, we are aiming to move `container` into or around `self.blobs`.
78        // Into if small enough, as we can further consolidate, but if not we need to
79        // consolidate and then either ship (if large) or hold (if small) the results.
80
81        let updates = std::mem::take(&mut container.updates).into_typed();
82        let consolidated = container.consolidated;
83        let len = updates.len();
84
85        // The input may be ready to ship on its own.
86        // This is ideal, if we've used an accumulating container builder elsewhere.
87        if consolidated && len >= crate::columnar::LINK_TARGET { self.ready.push_back(updates); }
88        // Can move into blobs if the combined length is not too large.
89        else if self.blob_records + len < 2 * crate::columnar::LINK_TARGET { self.absorb(updates, consolidated); }
90        // Otherwise, we'll need to manage `self.blobs`.
91        else {
92            // Together `updates` and `self.blobs` exceed 2 * LINK_TARGET.
93            // At least one, perhaps both of them, are LINK_TARGET in size.
94            // We'll consolidate any that are, and ship or merge the results.
95            // We'll end up with at most LINK_TARGET in `self.blobs`, retiring
96            // a constant factor of the pending work we started with.
97
98            // Consolidate and move to ready if large; stash otherwise.
99            let input_residual = if len >= crate::columnar::LINK_TARGET {
100                let cons = if consolidated { updates } else { updates.consolidate() };
101                if cons.len() >= crate::columnar::LINK_TARGET { self.ready.push_back(cons); None }
102                else if cons.len() > 0 { Some((cons, true)) }
103                else { None }
104            }
105            else { Some((updates, consolidated)) };
106
107            // Consolidate and move to ready if large; stash otherwise.
108            let blobs_residual = if self.blob_records >= crate::columnar::LINK_TARGET {
109                let cons = self.consolidate_blobs();
110                if cons.len() >= crate::columnar::LINK_TARGET { self.ready.push_back(cons); None }
111                else if cons.len() > 0 { Some((cons, true)) }
112                else { None }
113            }
114            else { None };
115
116            // Return un-shipped
117            if let Some((r, c)) = input_residual { self.absorb(r, c); }
118            if let Some((r, c)) = blobs_residual { self.absorb(r, c); }
119        }
120    }
121}
122
123impl<U: crate::columnar::layout::ColumnarUpdate> timely::container::ContainerBuilder for TrieChunker<U> {
124    type Container = ColChunk<U>;
125    fn extract(&mut self) -> Option<&mut Self::Container> {
126        self.stage = self.ready.pop_front().map(ColChunk::from_trie);
127        self.stage.as_mut()
128    }
129    fn finish(&mut self) -> Option<&mut Self::Container> {
130        // Drain whatever's left in blobs as a single (possibly small) final chunk.
131        if !self.blobs.is_empty() {
132            let cons = self.consolidate_blobs();
133            if cons.len() > 0 { self.ready.push_back(cons); }
134        }
135        self.extract()
136    }
137}