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