Add size_hint to windows and collect_to operator

Cargo.toml

@@ -78,6 +78,7 @@ indexmap = "2.2.6"
 tracing = { version = "0.1.40", features = ["log"] }
 quick_cache = "0.5.1"
 dashmap = "5.5.3"
+arrow = "51.0.0"
 
 [dev-dependencies]
 # for the tests

src/operator/window/aggr/collect.rs

@@ -1,58 +1,63 @@
 /// TODO
-
-
 use super::super::*;
 use crate::operator::{Data, DataKey, Operator};
 use crate::stream::{KeyedStream, WindowedStream};
 
-#[derive(Clone)]
-struct CollectVec<I, O, F>
+struct CollectTo<I, O>
 where
-    F: Fn(Vec<I>) -> O,
+    O: Extend<I> + Default,
 {
-    vec: Vec<I>,
-    f: F,
-    _o: PhantomData<O>,
+    collection: O,
+    _i: PhantomData<I>,
+}
+
+impl<I, O> Clone for CollectTo<I, O>
+where
+    O: Extend<I> + Default,
+{
+    fn clone(&self) -> Self {
+        Self {
+            collection: Default::default(),
+            _i: PhantomData,
+        }
+    }
 }
 
-impl<I, O, F> WindowAccumulator for CollectVec<I, O, F>
+impl<I, O> WindowAccumulator for CollectTo<I, O>
 where
-    F: Fn(Vec<I>) -> O + Send + Clone + 'static,
+    O:,
     I: Clone + Send + 'static,
-    O: Clone + Send + 'static,
+    O: Extend<I> + Default + Clone+ Send + 'static,
 {
     type In = I;
 
     type Out = O;
 
     #[inline]
     fn process(&mut self, el: Self::In) {
-        self.vec.push(el);
+        self.collection.extend(std::iter::once(el));
     }
 
     #[inline]
     fn output(self) -> Self::Out {
-        (self.f)(self.vec)
+        self.collection
     }
 }
 
-impl<Key, Out, WindowDescr, OperatorChain> WindowedStream<Key, Out, OperatorChain, Out, WindowDescr>
+impl<Key, Out, WindowDescr, OperatorChain> WindowedStream<OperatorChain, Out, WindowDescr>
 where
-    WindowDescr: WindowDescription,
-    OperatorChain: Operator<KeyValue<Key, Out>> + 'static,
+    WindowDescr: WindowDescription<Out>,
+    OperatorChain: Operator<Out = (Key, Out)> + 'static,
     Key: DataKey,
     Out: Data + Ord,
 {
-    /// Prefer other aggregators if possible as they don't save all elements
-    pub fn map<NewOut: Data, F: Fn(Vec<Out>) -> NewOut + Send + Clone + 'static>(
+    pub fn collect_to<C: Extend<Out> + Default + Clone + Send + 'static>(
         self,
-        f: F,
-    ) -> KeyedStream<Key, NewOut, impl Operator<KeyValue<Key, NewOut>>> {
-        let acc = CollectVec::<Out, NewOut, _> {
-            vec: Default::default(),
-            f,
-            _o: PhantomData,
+    ) -> KeyedStream<impl Operator<Out = (Key, C)>> {
+        let acc = CollectTo::<Out, C> {
+            collection: Default::default(),
+            _i: PhantomData,
         };
-        self.add_window_operator("WindowMap", acc)
+        self.add_window_operator("WindowCollect", acc)
     }
 }

src/operator/window/aggr/collect_vec.rs

@@ -21,6 +21,7 @@ where
     type In = I;
 
     type Out = O;
+
 
     #[inline]
     fn process(&mut self, el: Self::In) {
@@ -31,6 +32,11 @@ where
     fn output(self) -> Self::Out {
         (self.f)(self.vec)
     }
+
+    #[inline]
+    fn size_hint(&mut self, size: usize) {
+        self.vec.reserve(size);
+    }
 }
 
 impl<Key, Out, WindowDescr, OperatorChain> WindowedStream<OperatorChain, Out, WindowDescr>
@@ -59,6 +65,6 @@ where
             f: |v| v,
             _o: PhantomData,
         };
-        self.add_window_operator("WindowMap", acc)
+        self.add_window_operator("WindowCollectVec", acc)
     }
 }

src/operator/window/aggr/mod.rs

@@ -2,6 +2,7 @@ mod fold;
 pub(super) use fold::{Fold, FoldFirst};
 
 mod collect_vec;
+mod collect;
 mod count;
 mod join;
 mod max;

src/operator/window/descr/count.rs

@@ -142,8 +142,10 @@ impl<T: Data> WindowDescription<T> for CountWindow {
 
     #[inline]
     fn build<A: WindowAccumulator<In = T>>(&self, accumulator: A) -> Self::Manager<A> {
+        let mut init = accumulator;
+        init.size_hint(self.size);
         CountWindowManager {
-            init: accumulator,
+            init,
             size: self.size,
             slide: self.slide,
             exact: self.exact,

src/operator/window/descr/event_time.rs

@@ -3,7 +3,7 @@ use std::collections::VecDeque;
 use super::super::*;
 use crate::operator::{Data, StreamElement, Timestamp};
 
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 pub struct EventTimeWindowManager<A>
 where
     A: WindowAccumulator,
@@ -14,6 +14,22 @@ where
     last_watermark: Option<Timestamp>,
     ws: VecDeque<Slot<A>>,
 }
+
+impl<A> Clone for EventTimeWindowManager<A>
+where
+    A: WindowAccumulator,
+{
+    fn clone(&self) -> Self {
+        Self {
+            init: self.init.clone(),
+            size: self.size.clone(),
+            slide: self.slide.clone(),
+            last_watermark: self.last_watermark.clone(),
+            ws: self.ws.clone(),
+        }
+    }
+}
+
 impl<A: WindowAccumulator> EventTimeWindowManager<A> {
     fn alloc_windows(&mut self, ts: Timestamp) {
         assert!(self.last_watermark.map(|w| ts >= w).unwrap_or(true));

src/operator/window/mod.rs

@@ -34,9 +34,11 @@ pub trait WindowDescription<T> {
 ///
 /// Convention: output will always be called after at least one element has been processed
 pub trait WindowAccumulator: Clone + Send + 'static {
-    type In: Data;
-    type Out: Data;
+    type In: Clone + Send + 'static;
+    type Out: Clone + Send + 'static;
 
+    #[allow(unused)]
+    fn size_hint(&mut self, size: usize) {}
     /// Process a single input element updating the state of the accumulator
     fn process(&mut self, el: Self::In);
     /// Finalize the accumulator and produce a result