Skip to main content

timely/dataflow/operators/vec/
count.rs

1//! Counts the number of records at each time.
2use std::collections::HashMap;
3
4use crate::dataflow::channels::pact::Pipeline;
5use crate::dataflow::{StreamVec, Scope};
6use crate::dataflow::operators::generic::operator::Operator;
7
8/// Accumulates records within a timestamp.
9pub trait Accumulate<G: Scope, D: 'static> : Sized {
10    /// Accumulates records within a timestamp.
11    ///
12    /// # Examples
13    ///
14    /// ```
15    /// use timely::dataflow::operators::{ToStream, Capture};
16    /// use timely::dataflow::operators::vec::count::Accumulate;
17    /// use timely::dataflow::operators::capture::Extract;
18    ///
19    /// let captured = timely::example(|scope| {
20    ///     (0..10).to_stream(scope)
21    ///            .accumulate(0, |sum, data| { for &x in data.iter() { *sum += x; } })
22    ///            .capture()
23    /// });
24    ///
25    /// let extracted = captured.extract();
26    /// assert_eq!(extracted, vec![(0, vec![45])]);
27    /// ```
28    fn accumulate<A: Clone+'static>(self, default: A, logic: impl Fn(&mut A, &mut Vec<D>)+'static) -> StreamVec<G, A>;
29    /// Counts the number of records observed at each time.
30    ///
31    /// # Examples
32    ///
33    /// ```
34    /// use timely::dataflow::operators::{ToStream, Capture};
35    /// use timely::dataflow::operators::vec::count::Accumulate;
36    /// use timely::dataflow::operators::capture::Extract;
37    ///
38    /// let captured = timely::example(|scope| {
39    ///     (0..10).to_stream(scope)
40    ///            .count()
41    ///            .capture()
42    /// });
43    ///
44    /// let extracted = captured.extract();
45    /// assert_eq!(extracted, vec![(0, vec![10])]);
46    /// ```
47    fn count(self) -> StreamVec<G, usize> {
48        self.accumulate(0, |sum, data| *sum += data.len())
49    }
50}
51
52impl<G: Scope<Timestamp: ::std::hash::Hash>, D: 'static> Accumulate<G, D> for StreamVec<G, D> {
53    fn accumulate<A: Clone+'static>(self, default: A, logic: impl Fn(&mut A, &mut Vec<D>)+'static) -> StreamVec<G, A> {
54
55        let mut accums = HashMap::new();
56        self.unary_notify(Pipeline, "Accumulate", vec![], move |input, output, notificator| {
57            input.for_each_time(|time, data| {
58                for data in data {
59                    logic(accums.entry(time.time().clone()).or_insert_with(|| default.clone()), data);
60                }
61                notificator.notify_at(time.retain(output.output_index()));
62            });
63
64            notificator.for_each(|time,_,_| {
65                if let Some(accum) = accums.remove(&time) {
66                    output.session(&time).give(accum);
67                }
68            });
69        })
70    }
71}