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}