timely/dataflow/operators/
branch.rs

Help
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, StreamCore};
6use crate::{Container, Data};
7
8/// Extension trait for `Stream`.
9pub trait Branch<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 a reference to
12    /// the data and its time. If it returns `true`, the record will be sent
13    /// to the second returned stream, otherwise it will be sent to the first.
14    ///
15    /// If the result of the closure only depends on the time, not the data,
16    /// `branch_when` should be used instead.
17    ///
18    /// # Examples
19    /// ```
20    /// use timely::dataflow::operators::{ToStream, Branch, Inspect};
21    ///
22    /// timely::example(|scope| {
23    ///     let (odd, even) = (0..10)
24    ///         .to_stream(scope)
25    ///         .branch(|_time, x| *x % 2 == 0);
26    ///
27    ///     even.inspect(|x| println!("even numbers: {:?}", x));
28    ///     odd.inspect(|x| println!("odd numbers: {:?}", x));
29    /// });
30    /// ```
31    fn branch(
32        &self,
33        condition: impl Fn(&S::Timestamp, &D) -> bool + 'static,
34    ) -> (Stream<S, D>, Stream<S, D>);
35}
36
37impl<S: Scope, D: Data> Branch<S, D> for Stream<S, D> {
38    fn branch(
39        &self,
40        condition: impl Fn(&S::Timestamp, &D) -> bool + 'static,
41    ) -> (Stream<S, D>, Stream<S, D>) {
42        let mut builder = OperatorBuilder::new("Branch".to_owned(), self.scope());
43
44        let mut input = builder.new_input(self, Pipeline);
45        let (mut output1, stream1) = builder.new_output();
46        let (mut output2, stream2) = builder.new_output();
47
48        builder.build(move |_| {
49            move |_frontiers| {
50                let mut output1_handle = output1.activate();
51                let mut output2_handle = output2.activate();
52
53                input.for_each(|time, data| {
54                    let mut out1 = output1_handle.session(&time);
55                    let mut out2 = output2_handle.session(&time);
56                    for datum in data.drain(..) {
57                        if condition(time.time(), &datum) {
58                            out2.give(datum);
59                        } else {
60                            out1.give(datum);
61                        }
62                    }
63                });
64            }
65        });
66
67        (stream1, stream2)
68    }
69}
70
71/// Extension trait for `Stream`.
72pub trait BranchWhen<T>: Sized {
73    /// Takes one input stream and splits it into two output streams.
74    /// For each time, the supplied closure is called. If it returns `true`,
75    /// the records for that will be sent to the second returned stream, otherwise
76    /// they will be sent to the first.
77    ///
78    /// # Examples
79    /// ```
80    /// use timely::dataflow::operators::{ToStream, BranchWhen, Inspect, Delay};
81    ///
82    /// timely::example(|scope| {
83    ///     let (before_five, after_five) = (0..10)
84    ///         .to_stream(scope)
85    ///         .delay(|x,t| *x) // data 0..10 at time 0..10
86    ///         .branch_when(|time| time >= &5);
87    ///
88    ///     before_five.inspect(|x| println!("Times 0-4: {:?}", x));
89    ///     after_five.inspect(|x| println!("Times 5 and later: {:?}", x));
90    /// });
91    /// ```
92    fn branch_when(&self, condition: impl Fn(&T) -> bool + 'static) -> (Self, Self);
93}
94
95impl<S: Scope, C: Container + Data> BranchWhen<S::Timestamp> for StreamCore<S, C> {
96    fn branch_when(&self, condition: impl Fn(&S::Timestamp) -> bool + 'static) -> (Self, Self) {
97        let mut builder = OperatorBuilder::new("Branch".to_owned(), self.scope());
98
99        let mut input = builder.new_input(self, Pipeline);
100        let (mut output1, stream1) = builder.new_output();
101        let (mut output2, stream2) = builder.new_output();
102
103        builder.build(move |_| {
104
105            move |_frontiers| {
106                let mut output1_handle = output1.activate();
107                let mut output2_handle = output2.activate();
108
109                input.for_each(|time, data| {
110                    let mut out = if condition(time.time()) {
111                        output2_handle.session(&time)
112                    } else {
113                        output1_handle.session(&time)
114                    };
115                    out.give_container(data);
116                });
117            }
118        });
119
120        (stream1, stream2)
121    }
122}
timely/dataflow/operators/branch.rs

timely/dataflow/operators/
branch.rs
