Container-invariant BranchWhen operator
As the title says, this changes the BranchWhen operator to work on streams
of arbitrary containers.
use crate::dataflow::channels::pact::Pipeline;
use crate::dataflow::operators::generic::builder_rc::OperatorBuilder;
-use crate::dataflow::{Scope, Stream};
-use crate::Data;
+use crate::dataflow::{Scope, Stream, StreamCore};
+use crate::{Container, Data};
/// Extension trait for `Stream`.
pub trait Branch<S: Scope, D: Data> {
@@ -71,7 +71,7 @@ impl<S: Scope, D: Data> Branch<S, D> for Stream<S, D> {
}
/// Extension trait for `Stream`.
-pub trait BranchWhen<S: Scope, D: Data> {
+pub trait BranchWhen<T>: Sized {
/// Takes one input stream and splits it into two output streams.
/// For each time, the supplied closure is called. If it returns true,
/// the records for that will be sent to the second returned stream, otherwise
@@ -91,17 +91,11 @@ pub trait BranchWhen<S: Scope, D: Data> {
/// after_five.inspect(|x| println!("Times 5 and later: {:?}", x));
/// });
/// ```
- fn branch_when(
- &self,
- condition: impl Fn(&S::Timestamp) -> bool + 'static,
- ) -> (Stream<S, D>, Stream<S, D>);
+ fn branch_when(&self, condition: impl Fn(&T) -> bool + 'static) -> (Self, Self);
}
-impl<S: Scope, D: Data> BranchWhen<S, D> for Stream<S, D> {
- fn branch_when(
- &self,
- condition: impl Fn(&S::Timestamp) -> bool + 'static,
- ) -> (Stream<S, D>, Stream<S, D>) {
+impl<S: Scope, C: Container> BranchWhen<S::Timestamp> for StreamCore<S, C> {
+ fn branch_when(&self, condition: impl Fn(&S::Timestamp) -> bool + 'static) -> (Self, Self) {
let mut builder = OperatorBuilder::new("Branch".to_owned(), self.scope());
let mut input = builder.new_input(self, Pipeline);
@@ -110,19 +104,19 @@ impl<S: Scope, D: Data> BranchWhen<S, D> for Stream<S, D> {
builder.build(move |_| {
- let mut vector = Vec::new();
+ let mut container = Default::default();
move |_frontiers| {
let mut output1_handle = output1.activate();
let mut output2_handle = output2.activate();
input.for_each(|time, data| {
- data.swap(&mut vector);
+ data.swap(&mut container);
let mut out = if condition(&time.time()) {
output2_handle.session(&time)
} else {
output1_handle.session(&time)
};
- out.give_vec(&mut vector);
+ out.give_container(&mut container);
});
}
});
