timely/dataflow/operators/generic/
handles.rs

1//! Handles to an operator's input and output streams.
2//!
3//! These handles are used by the generic operator interfaces to allow user closures to interact as
4//! the operator would with its input and output streams.
5
6use std::rc::Rc;
7use std::cell::RefCell;
8use std::collections::VecDeque;
9
10use crate::progress::Timestamp;
11use crate::progress::ChangeBatch;
12use crate::progress::operate::PortConnectivity;
13use crate::dataflow::channels::pullers::Counter as PullCounter;
14use crate::dataflow::channels::Message;
15use crate::communication::Pull;
16use crate::{Container, ContainerBuilder, Accountable};
17use crate::container::{CapacityContainerBuilder, PushInto};
18
19use crate::dataflow::operators::InputCapability;
20use crate::dataflow::operators::capability::CapabilityTrait;
21
22/// Handle to an operator's input stream.
23pub struct InputHandleCore<T: Timestamp, C, P: Pull<Message<T, C>>> {
24    pull_counter: PullCounter<T, C, P>,
25    internal: Rc<RefCell<Vec<Rc<RefCell<ChangeBatch<T>>>>>>,
26    /// Timestamp summaries from this input to each output.
27    ///
28    /// Each timestamp received through this input may only produce output timestamps
29    /// greater or equal to the input timestamp subjected to at least one of these summaries.
30    summaries: Rc<RefCell<PortConnectivity<T::Summary>>>,
31    /// Staged capabilities and containers.
32    staging: VecDeque<(InputCapability<T>, C)>,
33    staged: Vec<C>,
34}
35
36impl<T: Timestamp, C: Accountable, P: Pull<Message<T, C>>> InputHandleCore<T, C, P> {
37    /// Reads the next input buffer (at some timestamp `t`) and a corresponding capability for `t`.
38    /// The timestamp `t` of the input buffer can be retrieved by invoking `.time()` on the capability.
39    /// Returns `None` when there's no more data available.
40    #[inline]
41    fn next(&mut self) -> Option<(InputCapability<T>, &mut C)> {
42        let internal = &self.internal;
43        let summaries = &self.summaries;
44        self.pull_counter.next_guarded().map(|(guard, bundle)| {
45            (InputCapability::new(Rc::clone(internal), Rc::clone(summaries), guard), &mut bundle.data)
46        })
47    }
48    /// Iterates through pairs of capability and container.
49    ///
50    /// The `for_each_time` method is equivalent, but groups containers by capability and is preferred,
51    /// in that it often leads to grouping work by capability, including the creation of output sessions.
52    pub fn for_each<F>(&mut self, mut logic: F) where F: FnMut(InputCapability<T>, &mut C) {
53        while let Some((cap, data)) = self.next() { logic(cap, data); }
54    }
55    /// Iterates through distinct capabilities and the lists of containers associated with each.
56    pub fn for_each_time<F>(&mut self, mut logic: F) where F: FnMut(InputCapability<T>, std::slice::IterMut::<C>), C: Default {
57        while let Some((cap, data)) = self.next() {
58            let data = std::mem::take(data);
59            self.staging.push_back((cap, data));
60        }
61        self.staging.make_contiguous().sort_by(|x,y| x.0.time().cmp(&y.0.time()));
62
63        while let Some((cap, data)) = self.staging.pop_front() {
64            self.staged.push(data);
65            let more = self.staging.iter().take_while(|(c,_)| c.time() == cap.time()).count();
66            self.staged.extend(self.staging.drain(..more).map(|(_,d)| d));
67            logic(cap, self.staged.iter_mut());
68            // Could return these back to the input ..
69            self.staged.clear();
70        }
71    }
72}
73
74/// Constructs an input handle.
75/// Declared separately so that it can be kept private when `InputHandle` is re-exported.
76pub fn new_input_handle<T: Timestamp, C: Accountable, P: Pull<Message<T, C>>>(
77    pull_counter: PullCounter<T, C, P>,
78    internal: Rc<RefCell<Vec<Rc<RefCell<ChangeBatch<T>>>>>>,
79    summaries: Rc<RefCell<PortConnectivity<T::Summary>>>,
80) -> InputHandleCore<T, C, P> {
81    InputHandleCore {
82        pull_counter,
83        internal,
84        summaries,
85        staging: Default::default(),
86        staged: Default::default(),
87    }
88}
89
90/// An owning pair of output pusher and container builder.
91pub struct OutputBuilder<T: Timestamp, CB: ContainerBuilder> {
92    output: crate::dataflow::channels::pushers::Output<T, CB::Container>,
93    builder: CB,
94}
95
96impl<T: Timestamp, CB: ContainerBuilder> OutputBuilder<T, CB> {
97    /// Constructs an output builder from an output and a default container builder.
98    pub fn from(output: crate::dataflow::channels::pushers::Output<T, CB::Container>) -> Self {
99        Self { output, builder: CB::default() }
100    }
101    /// An activated output buffer for building containers.
102    pub fn activate<'a>(&'a mut self) -> OutputBuilderSession<'a, T, CB> {
103        OutputBuilderSession {
104            session: self.output.activate(),
105            builder: &mut self.builder,
106        }
107    }
108}
109
110/// A wrapper around a live output session, with a container builder to buffer.
111pub struct OutputBuilderSession<'a, T: Timestamp, CB: ContainerBuilder> {
112    session: crate::dataflow::channels::pushers::OutputSession<'a, T, CB::Container>,
113    builder: &'a mut CB,
114}
115
116impl<'a, T: Timestamp, CB: ContainerBuilder> OutputBuilderSession<'a, T, CB> {
117    /// A container-building session associated with a capability.
118    ///
119    /// This method is the prefered way of sending records that must be accumulated into a container,
120    /// as it avoid the recurring overhead of capability validation.
121    pub fn session_with_builder<'b, CT: CapabilityTrait<T>>(&'b mut self, capability: &'b CT) -> Session<'a, 'b, T, CB, CT> where 'a: 'b {
122        debug_assert!(self.session.valid(capability));
123        Session {
124            buffer: self,
125            capability,
126        }
127    }
128}
129
130impl<'a, T: Timestamp, C: Container> OutputBuilderSession<'a, T, CapacityContainerBuilder<C>> {
131    /// A container-building session associated with a capability.
132    ///
133    /// This method is the prefered way of sending records that must be accumulated into a container,
134    /// as it avoid the recurring overhead of capability validation.
135    pub fn session<'b, CT: CapabilityTrait<T>>(&'b mut self, capability: &'b CT) -> Session<'a, 'b, T, CapacityContainerBuilder<C>, CT> where 'a: 'b {
136        debug_assert!(self.session.valid(capability));
137        Session {
138            buffer: self,
139            capability,
140        }
141    }
142}
143
144/// An active output building session, which accepts items and builds containers.
145pub struct Session<'a: 'b, 'b, T: Timestamp, CB: ContainerBuilder, CT: CapabilityTrait<T>> {
146    buffer: &'b mut OutputBuilderSession<'a, T, CB>,
147    capability: &'b CT,
148}
149
150impl<'a: 'b, 'b, T: Timestamp, CB: ContainerBuilder, CT: CapabilityTrait<T>> Session<'a, 'b, T, CB, CT> {
151
152    /// Provides access to the underlying container builder.
153    pub fn builder(&mut self) -> &mut CB { &mut self.buffer.builder }
154
155    /// Provides one record at the time specified by the `Session`.
156    #[inline] pub fn give<D>(&mut self, data: D) where CB: PushInto<D> {
157        self.buffer.builder.push_into(data);
158        self.extract_and_send();
159    }
160    /// Provides an iterator of records at the time specified by the `Session`.
161    #[inline] pub fn give_iterator<I>(&mut self, iter: I) where I: Iterator, CB: PushInto<I::Item> {
162        for item in iter { self.buffer.builder.push_into(item); }
163        self.extract_and_send();
164    }
165    /// Provide a container at the time specified by the [Session].
166    #[inline] pub fn give_container(&mut self, container: &mut CB::Container) {
167        self.buffer.session.give(&self.capability, container);
168    }
169    /// Provide multiple containers at the time specifid by the [Session].
170    #[inline] pub fn give_containers<'c>(&mut self, containers: impl Iterator<Item = &'c mut CB::Container>) {
171        for container in containers { self.buffer.session.give(&self.capability, container); }
172    }
173
174    /// Extracts built containers and sends them.
175    pub fn extract_and_send(&mut self) {
176        while let Some(container) = self.buffer.builder.extract() {
177            self.buffer.session.give(&self.capability, container);
178        }
179    }
180    /// Finalizes containers and sends them.
181    pub fn flush(&mut self) {
182        while let Some(container) = self.buffer.builder.finish() {
183            self.buffer.session.give(&self.capability, container);
184        }
185    }
186}
187
188impl<'a: 'b, 'b, T: Timestamp, CB: ContainerBuilder, CT: CapabilityTrait<T>> Drop for Session<'a, 'b, T, CB, CT> {
189    fn drop(&mut self) { self.flush() }
190}
timely/dataflow/operators/generic/handles.rs

timely/dataflow/operators/generic/
handles.rs
