timely/dataflow/operators/
count.rs

1//! Counts the number of records at each time.
2use std::collections::HashMap;
3
4use crate::communication::message::RefOrMut;
5
6use crate::Data;
7use crate::dataflow::channels::pact::Pipeline;
8use crate::dataflow::{Stream, Scope};
9use crate::dataflow::operators::generic::operator::Operator;
10
11/// Accumulates records within a timestamp.
12pub trait Accumulate<G: Scope, D: Data> {
13    /// Accumulates records within a timestamp.
14    ///
15    /// # Examples
16    ///
17    /// ```
18    /// use timely::dataflow::operators::{ToStream, Accumulate, Capture};
19    /// use timely::dataflow::operators::capture::Extract;
20    ///
21    /// let captured = timely::example(|scope| {
22    ///     (0..10).to_stream(scope)
23    ///            .accumulate(0, |sum, data| { for &x in data.iter() { *sum += x; } })
24    ///            .capture()
25    /// });
26    ///
27    /// let extracted = captured.extract();
28    /// assert_eq!(extracted, vec![(0, vec![45])]);
29    /// ```
30    fn accumulate<A: Data>(&self, default: A, logic: impl Fn(&mut A, RefOrMut<Vec<D>>)+'static) -> Stream<G, A>;
31    /// Counts the number of records observed at each time.
32    ///
33    /// # Examples
34    ///
35    /// ```
36    /// use timely::dataflow::operators::{ToStream, Accumulate, Capture};
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) -> Stream<G, usize> {
49        self.accumulate(0, |sum, data| *sum += data.len())
50    }
51}
52
53impl<G: Scope, D: Data> Accumulate<G, D> for Stream<G, D> {
54    fn accumulate<A: Data>(&self, default: A, logic: impl Fn(&mut A, RefOrMut<Vec<D>>)+'static) -> Stream<G, A> {
55
56        let mut accums = HashMap::new();
57        self.unary_notify(Pipeline, "Accumulate", vec![], move |input, output, notificator| {
58            input.for_each(|time, data| {
59                logic(&mut accums.entry(time.time().clone()).or_insert_with(|| default.clone()), data);
60                notificator.notify_at(time.retain());
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}
timely/dataflow/operators/count.rs

timely/dataflow/operators/
count.rs
