timely/dataflow/operators/
count.rs

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

timely/dataflow/operators/
count.rs
