Open the Inspect trait for more stream kinds

timely/src/dataflow/operators/generic/operator.rs

@@ -594,7 +594,7 @@ where
 /// timely::example(|scope| {
 ///
 ///
-///     empty(scope)     // type required in this example
+///     empty::<_, Vec<_>>(scope)     // type required in this example
 ///         .inspect(|()| panic!("never called"));
 ///
 /// });

timely/src/dataflow/operators/inspect.rs

@@ -1,13 +1,15 @@
 //! Extension trait and implementation for observing and action on streamed data.
 
+use std::rc::Rc;
+use timely_container::columnation::{Columnation, TimelyStack};
 use crate::Container;
 use crate::Data;
 use crate::dataflow::channels::pact::Pipeline;
-use crate::dataflow::{Stream, Scope, StreamCore};
+use crate::dataflow::{Scope, StreamCore};
 use crate::dataflow::operators::generic::Operator;
 
 /// Methods to inspect records and batches of records on a stream.
-pub trait Inspect<G: Scope, D: Data>: InspectCore<G, Vec<D>> {
+pub trait Inspect<G: Scope, C: Container>: InspectCore<G, C> + Sized {
     /// Runs a supplied closure on each observed data element.
     ///
     /// # Examples
@@ -19,7 +21,7 @@ pub trait Inspect<G: Scope, D: Data>: InspectCore<G, Vec<D>> {
     ///            .inspect(|x| println!("seen: {:?}", x));
     /// });
     /// ```
-    fn inspect(&self, mut func: impl FnMut(&D)+'static) -> Stream<G, D> {
+    fn inspect(&self, mut func: impl FnMut(&C::Item)+'static) -> Self {
         self.inspect_batch(move |_, data| {
             for datum in data.iter() { func(datum); }
         })
@@ -36,7 +38,7 @@ pub trait Inspect<G: Scope, D: Data>: InspectCore<G, Vec<D>> {
     ///            .inspect_time(|t, x| println!("seen at: {:?}\t{:?}", t, x));
     /// });
     /// ```
-    fn inspect_time(&self, mut func: impl FnMut(&G::Timestamp, &D)+'static) -> Stream<G, D> {
+    fn inspect_time(&self, mut func: impl FnMut(&G::Timestamp, &C::Item)+'static) -> Self {
         self.inspect_batch(move |time, data| {
             for datum in data.iter() {
                 func(&time, &datum);
@@ -55,7 +57,7 @@ pub trait Inspect<G: Scope, D: Data>: InspectCore<G, Vec<D>> {
     ///            .inspect_batch(|t,xs| println!("seen at: {:?}\t{:?} records", t, xs.len()));
     /// });
     /// ```
-    fn inspect_batch(&self, mut func: impl FnMut(&G::Timestamp, &[D])+'static) -> Stream<G, D> {
+    fn inspect_batch(&self, mut func: impl FnMut(&G::Timestamp, &[C::Item])+'static) -> Self {
         self.inspect_core(move |event| {
             if let Ok((time, data)) = event {
                 func(time, data);
@@ -82,15 +84,29 @@ pub trait Inspect<G: Scope, D: Data>: InspectCore<G, Vec<D>> {
     ///             });
     /// });
     /// ```
-    fn inspect_core<F>(&self, func: F) -> Stream<G, D> where F: FnMut(Result<(&G::Timestamp, &[D]), &[G::Timestamp]>)+'static;
+    fn inspect_core<F>(&self, func: F) -> Self where F: FnMut(Result<(&G::Timestamp, &[C::Item]), &[G::Timestamp]>)+'static;
 }
 
-impl<G: Scope, D: Data> Inspect<G, D> for Stream<G, D> {
-    fn inspect_core<F>(&self, mut func: F) -> Stream<G, D> where F: FnMut(Result<(&G::Timestamp, &[D]), &[G::Timestamp]>) + 'static {
+impl<G: Scope, D: Data> Inspect<G, Vec<D>> for StreamCore<G, Vec<D>> {
+    fn inspect_core<F>(&self, mut func: F) -> Self where F: FnMut(Result<(&G::Timestamp, &[D]), &[G::Timestamp]>) + 'static {
         self.inspect_container(move |r| func(r.map(|(t, c)| (t, &c[..]))))
     }
 }
 
+impl<G: Scope, D: Data+Columnation> Inspect<G, TimelyStack<D>> for StreamCore<G, TimelyStack<D>> {
+    fn inspect_core<F>(&self, mut func: F) -> Self where F: FnMut(Result<(&G::Timestamp, &[D]), &[G::Timestamp]>) + 'static {
+        self.inspect_container(move |r| func(r.map(|(t, c)| (t, &c[..]))))
+    }
+}
+
+impl<G: Scope, C: Container> Inspect<G, Rc<C>> for StreamCore<G, Rc<C>>
+    where C: AsRef<[C::Item]>
+{
+    fn inspect_core<F>(&self, mut func: F) -> Self where F: FnMut(Result<(&G::Timestamp, &[C::Item]), &[G::Timestamp]>) + 'static {
+        self.inspect_container(move |r| func(r.map(|(t, c)| (t, c.as_ref().as_ref()))))
+    }
+}
+
 /// Inspect containers
 pub trait InspectCore<G: Scope, C: Container> {
     /// Runs a supplied closure on each observed container, and each frontier advancement.