timely/dataflow/operators/aggregation/state_machine.rs
1//! General purpose state transition operator.
2use std::hash::Hash;
3use std::collections::HashMap;
4
5use crate::{Data, ExchangeData};
6use crate::dataflow::{Stream, Scope};
7use crate::dataflow::operators::generic::operator::Operator;
8use crate::dataflow::channels::pact::Exchange;
9
10/// Generic state-transition machinery: each key has a state, and receives a sequence of events.
11/// Events are applied in time-order, but no other promises are made. Each state transition can
12/// produce output, which is sent.
13///
14/// `state_machine` will buffer inputs if earlier inputs may still arrive. it will directly apply
15/// updates for the current time reflected in the notificator, though. In the case of partially
16/// ordered times, the only guarantee is that updates are not applied out of order, not that there
17/// is some total order on times respecting the total order (updates may be interleaved).
18
19/// Provides the `state_machine` method.
20pub trait StateMachine<S: Scope, K: ExchangeData+Hash+Eq, V: ExchangeData> {
21 /// Tracks a state for each presented key, using user-supplied state transition logic.
22 ///
23 /// The transition logic `fold` may mutate the state, and produce both output records and
24 /// a `bool` indicating that it is appropriate to deregister the state, cleaning up once
25 /// the state is no longer helpful.
26 ///
27 /// # Examples
28 /// ```
29 /// use timely::dataflow::operators::{ToStream, Map, Inspect};
30 /// use timely::dataflow::operators::aggregation::StateMachine;
31 ///
32 /// timely::example(|scope| {
33 ///
34 /// // these results happen to be right, but aren't guaranteed.
35 /// // the system is at liberty to re-order within a timestamp.
36 /// let result = vec![(0,0), (0,2), (0,6), (0,12), (0,20),
37 /// (1,1), (1,4), (1,9), (1,16), (1,25)];
38 ///
39 /// (0..10).to_stream(scope)
40 /// .map(|x| (x % 2, x))
41 /// .state_machine(
42 /// |_key, val, agg| { *agg += val; (false, Some((*_key, *agg))) },
43 /// |key| *key as u64
44 /// )
45 /// .inspect(move |x| assert!(result.contains(x)));
46 /// });
47 /// ```
48 fn state_machine<
49 R: Data, // output type
50 D: Default+'static, // per-key state (data)
51 I: IntoIterator<Item=R>, // type of output iterator
52 F: Fn(&K, V, &mut D)->(bool, I)+'static, // state update logic
53 H: Fn(&K)->u64+'static, // "hash" function for keys
54 >(&self, fold: F, hash: H) -> Stream<S, R> where S::Timestamp : Hash+Eq ;
55}
56
57impl<S: Scope, K: ExchangeData+Hash+Eq, V: ExchangeData> StateMachine<S, K, V> for Stream<S, (K, V)> {
58 fn state_machine<
59 R: Data, // output type
60 D: Default+'static, // per-key state (data)
61 I: IntoIterator<Item=R>, // type of output iterator
62 F: Fn(&K, V, &mut D)->(bool, I)+'static, // state update logic
63 H: Fn(&K)->u64+'static, // "hash" function for keys
64 >(&self, fold: F, hash: H) -> Stream<S, R> where S::Timestamp : Hash+Eq {
65
66 let mut pending: HashMap<_, Vec<(K, V)>> = HashMap::new(); // times -> (keys -> state)
67 let mut states = HashMap::new(); // keys -> state
68
69 let mut vector = Vec::new();
70
71 self.unary_notify(Exchange::new(move |&(ref k, _)| hash(k)), "StateMachine", vec![], move |input, output, notificator| {
72
73 // go through each time with data, process each (key, val) pair.
74 notificator.for_each(|time,_,_| {
75 if let Some(pend) = pending.remove(time.time()) {
76 let mut session = output.session(&time);
77 for (key, val) in pend {
78 let (remove, output) = {
79 let state = states.entry(key.clone()).or_insert_with(Default::default);
80 fold(&key, val, state)
81 };
82 if remove { states.remove(&key); }
83 session.give_iterator(output.into_iter());
84 }
85 }
86 });
87
88 // stash each input and request a notification when ready
89 input.for_each(|time, data| {
90
91 data.swap(&mut vector);
92
93 // stash if not time yet
94 if notificator.frontier(0).less_than(time.time()) {
95 pending.entry(time.time().clone()).or_insert_with(Vec::new).extend(vector.drain(..));
96 notificator.notify_at(time.retain());
97 }
98 else {
99 // else we can process immediately
100 let mut session = output.session(&time);
101 for (key, val) in vector.drain(..) {
102 let (remove, output) = {
103 let state = states.entry(key.clone()).or_insert_with(Default::default);
104 fold(&key, val, state)
105 };
106 if remove { states.remove(&key); }
107 session.give_iterator(output.into_iter());
108 }
109 }
110 });
111 })
112 }
113}