timely/dataflow/operators/ok_err.rs
1//! Operators that separate one stream into two streams based on some condition
2
3use crate::dataflow::channels::pact::Pipeline;
4use crate::dataflow::operators::generic::builder_rc::OperatorBuilder;
5use crate::dataflow::{Scope, Stream};
6use crate::Data;
7
8/// Extension trait for `Stream`.
9pub trait OkErr<S: Scope, D: Data> {
10 /// Takes one input stream and splits it into two output streams.
11 /// For each record, the supplied closure is called with the data.
12 /// If it returns `Ok(x)`, then `x` will be sent
13 /// to the first returned stream; otherwise, if it returns `Err(e)`,
14 /// then `e` will be sent to the second.
15 ///
16 /// If the result of the closure only depends on the time, not the data,
17 /// `branch_when` should be used instead.
18 ///
19 /// # Examples
20 /// ```
21 /// use timely::dataflow::operators::{ToStream, OkErr, Inspect};
22 ///
23 /// timely::example(|scope| {
24 /// let (odd, even) = (0..10)
25 /// .to_stream(scope)
26 /// .ok_err(|x| if x % 2 == 0 { Ok(x) } else { Err(x) });
27 ///
28 /// even.inspect(|x| println!("even numbers: {:?}", x));
29 /// odd.inspect(|x| println!("odd numbers: {:?}", x));
30 /// });
31 /// ```
32 fn ok_err<D1, D2, L>(
33 &self,
34 logic: L,
35 ) -> (Stream<S, D1>, Stream<S, D2>)
36
37 where
38 D1: Data,
39 D2: Data,
40 L: FnMut(D) -> Result<D1,D2>+'static
41 ;
42}
43
44impl<S: Scope, D: Data> OkErr<S, D> for Stream<S, D> {
45 fn ok_err<D1, D2, L>(
46 &self,
47 mut logic: L,
48 ) -> (Stream<S, D1>, Stream<S, D2>)
49
50 where
51 D1: Data,
52 D2: Data,
53 L: FnMut(D) -> Result<D1,D2>+'static
54 {
55 let mut builder = OperatorBuilder::new("OkErr".to_owned(), self.scope());
56
57 let mut input = builder.new_input(self, Pipeline);
58 let (mut output1, stream1) = builder.new_output();
59 let (mut output2, stream2) = builder.new_output();
60
61 builder.build(move |_| {
62 let mut vector = Vec::new();
63 move |_frontiers| {
64 let mut output1_handle = output1.activate();
65 let mut output2_handle = output2.activate();
66
67 input.for_each(|time, data| {
68 data.swap(&mut vector);
69 let mut out1 = output1_handle.session(&time);
70 let mut out2 = output2_handle.session(&time);
71 for datum in vector.drain(..) {
72 match logic(datum) {
73 Ok(datum) => out1.give(datum),
74 Err(datum) => out2.give(datum),
75 }
76 }
77 });
78 }
79 });
80
81 (stream1, stream2)
82 }
83}