Add keyed fold scan

src/operator/iteration/iterate.rs

@@ -297,7 +297,7 @@ impl<Out: ExchangeData, State: ExchangeData + Sync> Operator for Iterate<Out, St
     }
 }
 
-impl<Out: ExchangeData, State: ExchangeData + Sync> Display for Iterate<Out, State> {
+impl<Out: ExchangeData, State: ExchangeData> Display for Iterate<Out, State> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "Iterate<{}>", std::any::type_name::<Out>())
     }

src/operator/mod.rs

@@ -3,6 +3,7 @@
 //! The actual operator list can be found from the implemented methods of [`Stream`],
 //! [`KeyedStream`], [`crate::WindowedStream`]
 
+use std::collections::HashMap;
 use std::fmt::Display;
 use std::hash::Hash;
 use std::ops::{AddAssign, Div};
@@ -2685,6 +2686,92 @@ where
         self.0.split_block(End::new, NextStrategy::random())
     }
 
+    pub fn fold_scan<O, S, L, F>(
+        self,
+        init: S,
+        fold: L,
+        map: F,
+    ) -> KeyedStream<impl Operator<Out = (K, O)>>
+    where
+        Op::Out: ExchangeData,
+        I: Send,
+        K: ExchangeDataKey + Sync,
+        L: Fn(&K, &mut S, I) + Send + Clone + 'static,
+        F: Fn(&K, &S, I) -> O + Send + Clone + 'static,
+        S: ExchangeData + Sync,
+        O: ExchangeData,
+    {
+        #[derive(Serialize, Deserialize, Clone)]
+        enum TwoPass<I, O> {
+            First(I),
+            Second(I),
+            Output(O),
+        }
+
+        let (state, s) = self.map(|el| TwoPass::First(el.1)).unkey().iterate(
+            2,
+            HashMap::<K, S>::default(),
+            |s, state| {
+                s.to_keyed()
+                    .map(move |(k, el)| match el {
+                        TwoPass::First(el) => TwoPass::Second(el),
+                        TwoPass::Second(el) => {
+                            TwoPass::Output((map)(k, state.get().get(k).unwrap(), el))
+                        }
+                        TwoPass::Output(_) => unreachable!(),
+                    })
+                    .unkey()
+            },
+            move |local: &mut HashMap<K, S>, (k, el)| match el {
+                TwoPass::First(_) => {}
+                TwoPass::Second(el) => fold(
+                    &k,
+                    local.entry(k.clone()).or_insert_with(|| init.clone()),
+                    el,
+                ),
+                TwoPass::Output(_) => {}
+            },
+            move |global: &mut HashMap<K, S>, mut local| {
+                global.extend(local.drain());
+            },
+            |_| true,
+        );
+
+        state.for_each(std::mem::drop);
+        s.to_keyed().map(|(_, t)| match t {
+            TwoPass::First(_) | TwoPass::Second(_) => unreachable!(),
+            TwoPass::Output(o) => o,
+        })
+    }
+
+    pub fn reduce_scan<O, S, F1, F2, R>(
+        self,
+        first_map: F1,
+        reduce: R,
+        second_map: F2,
+    ) -> KeyedStream<impl Operator<Out = (K, O)>>
+    where
+        Op::Out: ExchangeData,
+        F1: Fn(&K, I) -> S + Send + Clone + 'static,
+        F2: Fn(&K, &S, I) -> O + Send + Clone + 'static,
+        R: Fn(&K, S, S) -> S + Send + Clone + 'static,
+        K: Sync,
+        S: ExchangeData + Sync,
+        O: ExchangeData,
+    {
+        self.fold_scan(
+            None,
+            move |k, acc: &mut Option<S>, x| {
+                let map = (first_map)(k, x);
+                *acc = Some(match acc.take() {
+                    Some(v) => (reduce)(k, v, map),
+                    None => map,
+                });
+            },
+            move |k, state, x| (second_map)(k, state.as_ref().unwrap(), x),
+        )
+    }
+
     /// Close the stream and send resulting items to a channel on a single host.
     ///
     /// If the stream is distributed among multiple replicas, parallelism will

src/operator/sink/avro.rs

@@ -134,30 +134,3 @@ where
             .finalize_block();
     }
 }
-
-// #[cfg(test)]
-// mod qtests {
-//     use std::AvroSinkions::HashSet;
-
-//     use crate::config::RuntimeConfig;
-//     use crate::environment::StreamContext;
-//     use crate::operator::source;
-
-//     #[test]
-//     fn AvroSink_vec() {
-//         let env = StreamContext::new(RuntimeConfig::local(4).unwrap());
-//         let source = source::IteratorSource::new(0..10u8);
-//         let res = env.stream(source).AvroSink::<Vec<_>>();
-//         env.execute_blocking();
-//         assert_eq!(res.get().unwrap(), (0..10).AvroSink::<Vec<_>>());
-//     }
-
-//     #[test]
-//     fn AvroSink_set() {
-//         let env = StreamContext::new(RuntimeConfig::local(4).unwrap());
-//         let source = source::IteratorSource::new(0..10u8);
-//         let res = env.stream(source).AvroSink::<HashSet<_>>();
-//         env.execute_blocking();
-//         assert_eq!(res.get().unwrap(), (0..10).AvroSink::<HashSet<_>>());
-//     }
-// }

tests/fold_scan.rs

@@ -56,3 +56,55 @@ fn reduce_scan() {
         }
     });
 }
+
+#[test]
+fn keyed_fold_scan() {
+    TestHelper::local_remote_env(|ctx| {
+        let res = ctx
+            .stream_par_iter(0..100i32)
+            .group_by(|e| e % 5)
+            .fold_scan(
+                0,
+                |_k, acc: &mut i32, x| {
+                    *acc += x;
+                },
+                |_k, acc, x| (x, *acc),
+            )
+            .unkey()
+            .map(|t| (t.0, t.1 .0, t.1 .1))
+            .collect_vec();
+
+        ctx.execute_blocking();
+        if let Some(mut res) = res.get() {
+            let mut expected = (0..100)
+                .map(|x| (x % 5, x, (0..100).filter(|m| m % 5 == x % 5).sum::<i32>()))
+                .collect::<Vec<_>>();
+            res.sort();
+            expected.sort();
+            assert_eq!(expected, res);
+        }
+    });
+}
+
+#[test]
+fn keyed_reduce_scan() {
+    TestHelper::local_remote_env(|ctx| {
+        let res = ctx
+            .stream_par_iter(0..100i32)
+            .group_by(|e| e % 5)
+            .reduce_scan(|_k, x| x, |_k, a, b| a + b, |_k, acc, x| (x, *acc))
+            .unkey()
+            .map(|t| (t.0, t.1 .0, t.1 .1))
+            .collect_vec();
+
+        ctx.execute_blocking();
+        if let Some(mut res) = res.get() {
+            let mut expected = (0..100)
+                .map(|x| (x % 5, x, (0..100).filter(|m| m % 5 == x % 5).sum::<i32>()))
+                .collect::<Vec<_>>();
+            res.sort();
+            expected.sort();
+            assert_eq!(expected, res);
+        }
+    });
+}