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}