differential_dataflow/columnar/collection/
exchange.rs

1//! Exchange / parallelization contract for `RecordedUpdates`.
2//!
3//! `ValPact` is the PACT used when shuffling columnar updates across workers;
4//! `ValDistributor` is the per-worker partitioner it constructs.
5
6use std::rc::Rc;
7
8use columnar::{Index, Len};
9use timely::logging::TimelyLogger;
10use timely::dataflow::channels::pushers::{Exchange, exchange::Distributor};
11use timely::dataflow::channels::Message;
12use timely::dataflow::channels::pact::{LogPuller, LogPusher, ParallelizationContract};
13use timely::progress::Timestamp;
14use timely::worker::Worker;
15
16use crate::columnar::layout::ColumnarUpdate as Update;
17use crate::columnar::updates::UpdatesTyped;
18use super::RecordedUpdates;
19
20/// Distributor that routes `RecordedUpdates` records to workers by hashing keys.
21pub struct ValDistributor<U: Update, H> {
22    marker: std::marker::PhantomData<U>,
23    hashfunc: H,
24    pre_lens: Vec<usize>,
25    worker: usize,
26}
27
28impl<U: Update, H: for<'a> FnMut(columnar::Ref<'a, U::Key>)->u64> Distributor<RecordedUpdates<U>> for ValDistributor<U, H> {
29    // TODO: For unsorted UpdatesTyped (stride-1 outer keys), each key is its own outer group,
30    // so the per-group pre_lens snapshot and seal check costs O(keys × workers). Should
31    // either batch keys by destination first, or detect stride-1 outer bounds and use a
32    // simpler single-pass partitioning that seals once at the end.
33    fn partition<T: Clone, P: timely::communication::Push<Message<T, RecordedUpdates<U>>>>(&mut self, container: &mut RecordedUpdates<U>, stamp: &timely::progress::Stamp<T>, pushers: &mut [P]) {
34        use crate::columnar::updates::child_range;
35
36        let view = container.updates.view();
37        let keys_b = view.keys;
38        let mut outputs: Vec<UpdatesTyped<U>> = (0..pushers.len()).map(|_| UpdatesTyped::default()).collect();
39
40        // Each outer key group becomes a separate run in the destination.
41        for outer in 0..Len::len(&keys_b) {
42            self.pre_lens.clear();
43            self.pre_lens.extend(outputs.iter().map(|o| o.keys.values.len()));
44            if pushers.len().is_power_of_two() {
45                let mask = (pushers.len() - 1) as u64;
46                for k in child_range(keys_b.bounds, outer) {
47                    let key = keys_b.values.get(k);
48                    let h = (self.hashfunc)(key);
49                    let idx = (h & mask) as usize;
50                    outputs[idx].extend_from_keys(view, k..k+1);
51                }
52            }
53            else {
54                let pushers_len = pushers.len() as u64;
55                for k in child_range(keys_b.bounds, outer) {
56                    let key = keys_b.values.get(k);
57                    let h = (self.hashfunc)(key);
58                    let idx = (h % pushers_len) as usize;
59                    outputs[idx].extend_from_keys(view, k..k+1);
60                }
61            }
62            for (output, &pre) in outputs.iter_mut().zip(self.pre_lens.iter()) {
63                if output.keys.values.len() > pre {
64                    output.keys.bounds.push(output.keys.values.len() as u64);
65                }
66            }
67        }
68
69        // Distribute the input's record count across non-empty outputs.
70        let total_records = container.records;
71        let non_empty: usize = outputs.iter().filter(|o| !o.keys.values.is_empty()).count();
72        // N.B. An input whose updates all consolidated away needs to emit a message
73        // carrying a `records` count for timely's bookkeeping.
74        if non_empty == 0 && total_records > 0 {
75            let mut recorded = RecordedUpdates::<U> {
76                updates: Default::default(),
77                records: total_records,
78                consolidated: true,
79            };
80            // Push the empty update to the worker that produced the original
81            // values so the send stays local. Not needed for correctness, but
82            // a reasonable choice.
83            Message::push_at(&mut recorded, stamp.clone(), &mut pushers[self.worker % pushers.len()]);
84            return;
85        }
86
87        let mut first_records = total_records.saturating_sub(non_empty.saturating_sub(1));
88        for (pusher, output) in pushers.iter_mut().zip(outputs) {
89            if !output.keys.values.is_empty() {
90                let recorded = RecordedUpdates { updates: output.into(), records: first_records, consolidated: container.consolidated };
91                first_records = 1;
92                let mut recorded = recorded;
93                Message::push_at(&mut recorded, stamp.clone(), pusher);
94            }
95        }
96    }
97    fn flush<T: Clone, P: timely::communication::Push<Message<T, RecordedUpdates<U>>>>(&mut self, _stamp: &timely::progress::Stamp<T>, _pushers: &mut [P]) { }
98    fn relax(&mut self) { }
99}
100
101/// PACT for shuffling `RecordedUpdates` containers by hashing keys.
102pub struct ValPact<H> {
103    /// Hash function applied to each key reference.
104    pub hashfunc: H,
105}
106
107impl<T, U, H> ParallelizationContract<T, RecordedUpdates<U>> for ValPact<H>
108where
109    T: Timestamp,
110    U: Update,
111    H: for<'a> FnMut(columnar::Ref<'a, U::Key>)->u64 + 'static,
112{
113    type Pusher = Exchange<
114        T,
115        LogPusher<Box<dyn timely::communication::Push<Message<T, RecordedUpdates<U>>>>>,
116        ValDistributor<U, H>
117    >;
118    type Puller = LogPuller<Box<dyn timely::communication::Pull<Message<T, RecordedUpdates<U>>>>>;
119
120    fn connect(self, worker: &Worker, identifier: usize, address: Rc<[usize]>, logging: Option<TimelyLogger>) -> (Self::Pusher, Self::Puller) {
121        let (senders, receiver) = worker.allocate::<Message<T, RecordedUpdates<U>>>(identifier, address);
122        let senders = senders.into_iter().enumerate().map(|(i,x)| LogPusher::new(x, worker.index(), i, identifier, logging.clone())).collect::<Vec<_>>();
123        let distributor = ValDistributor {
124            marker: std::marker::PhantomData,
125            hashfunc: self.hashfunc,
126            pre_lens: Vec::new(),
127            worker: worker.index(),
128        };
129        (Exchange::new(senders, distributor), LogPuller::new(receiver, worker.index(), identifier, logging.clone()))
130    }
131}
differential_dataflow/columnar/collection/exchange.rs

differential_dataflow/columnar/collection/
exchange.rs
