Skip to main content

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}