join_traces uses timely's container builder (#478)

join_traces can produce arbitrary containers.

---------

Signed-off-by: Moritz Hoffmann <[email protected]>

Author: antiguru

examples/spines.rs

@@ -55,11 +55,11 @@ fn main() {
                     let data =
                     data.map(|x| (x.clone().into_bytes(), x.into_bytes()))
                         .arrange::<PreferredSpine<[u8],[u8],_,_>>()
-                        .reduce_abelian::<_, _, _, PreferredSpine<[u8],(),_,_>>("distinct", |v| v.clone(), |_,_,output| output.push(((), 1)));
+                        .reduce_abelian::<_, _, _, PreferredSpine<[u8],(),_,_>>("distinct", |_| (), |_,_,output| output.push(((), 1)));
                     let keys =
                     keys.map(|x| (x.clone().into_bytes(), 7))
                         .arrange::<PreferredSpine<[u8],u8,_,_>>()
-                        .reduce_abelian::<_, _, _, PreferredSpine<[u8],(),_,_>>("distinct", |v| v.clone(), |_,_,output| output.push(((), 1)));
+                        .reduce_abelian::<_, _, _, PreferredSpine<[u8],(),_,_>>("distinct", |_| (), |_,_,output| output.push(((), 1)));
 
                     keys.join_core(&data, |k,_v1,_v2| {
                         println!("{:?}", k.text);

server/Cargo.toml

@@ -10,7 +10,7 @@ path="../"
 [dependencies]
 rand="0.3.13"
 libloading="*"
-timely = { git = "https://github.com/TimelyDataflow/timely-dataflow" }
+timely = { workspace = true }
 
 #[workspace]
 #members = [

src/capture.rs

@@ -399,7 +399,7 @@ pub mod source {
                     for message in source.by_ref() {
                         match message {
                             Message::Updates(mut updates) => {
-                                updates_session.give_vec(&mut updates);
+                                updates_session.give_container(&mut updates);
                             }
                             Message::Progress(progress) => {
                                 // We must send a copy of each progress message to all workers,

src/consolidation.rs

@@ -10,6 +10,9 @@
 //! you need specific behavior, it may be best to defensively copy, paste, and maintain the
 //! specific behavior you require.
 
+use std::collections::VecDeque;
+use timely::container::{ContainerBuilder, PushContainer, PushInto};
+use crate::Data;
 use crate::difference::Semigroup;
 
 /// Sorts and consolidates `vec`.
@@ -145,6 +148,84 @@ pub fn consolidate_updates_slice<D: Ord, T: Ord, R: Semigroup>(slice: &mut [(D,
     offset
 }
 
+/// A container builder that consolidates data in-places into fixed-sized containers. Does not
+/// maintain FIFO ordering.
+#[derive(Default)]
+pub struct ConsolidatingContainerBuilder<C>{
+    current: C,
+    empty: Vec<C>,
+    outbound: VecDeque<C>,
+}
+
+impl<D,T,R> ConsolidatingContainerBuilder<Vec<(D, T, R)>>
+where
+    D: Data,
+    T: Data,
+    R: Semigroup,
+{
+    /// Flush `self.current` up to the biggest `multiple` of elements. Pass 1 to flush all elements.
+    // TODO: Can we replace `multiple` by a bool?
+    fn consolidate_and_flush_through(&mut self, multiple: usize) {
+        let preferred_capacity = <Vec<(D,T,R)>>::preferred_capacity();
+        consolidate_updates(&mut self.current);
+        let mut drain = self.current.drain(..(self.current.len()/multiple)*multiple).peekable();
+        while drain.peek().is_some() {
+            let mut container = self.empty.pop().unwrap_or_else(|| Vec::with_capacity(preferred_capacity));
+            container.extend((&mut drain).take(preferred_capacity));
+            self.outbound.push_back(container);
+        }
+    }
+}
+
+impl<D,T,R> ContainerBuilder for ConsolidatingContainerBuilder<Vec<(D, T, R)>>
+where
+    D: Data,
+    T: Data,
+    R: Semigroup,
+{
+    type Container = Vec<(D,T,R)>;
+
+    /// Push an element.
+    ///
+    /// Precondition: `current` is not allocated or has space for at least one element.
+    #[inline]
+    fn push<P: PushInto<Self::Container>>(&mut self, item: P) {
+        let preferred_capacity = <Vec<(D,T,R)>>::preferred_capacity();
+        if self.current.capacity() < preferred_capacity * 2 {
+            self.current.reserve(preferred_capacity * 2 - self.current.capacity());
+        }
+        item.push_into(&mut self.current);
+        if self.current.len() == self.current.capacity() {
+            // Flush complete containers.
+            self.consolidate_and_flush_through(preferred_capacity);
+        }
+    }
+
+    fn push_container(&mut self, container: &mut Self::Container) {
+        for item in container.drain(..) {
+            self.push(item);
+        }
+    }
+
+    fn extract(&mut self) -> Option<&mut Vec<(D,T,R)>> {
+        if let Some(container) = self.outbound.pop_front() {
+            self.empty.push(container);
+            self.empty.last_mut()
+        } else {
+            None
+        }
+    }
+
+    fn finish(&mut self) -> Option<&mut Vec<(D,T,R)>> {
+        // Flush all
+        self.consolidate_and_flush_through(1);
+        // Remove all but two elements from the stash of empty to avoid memory leaks. We retain
+        // two to match `current` capacity.
+        self.empty.truncate(2);
+        self.extract()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -211,4 +292,46 @@ mod tests {
             assert_eq!(input, output);
         }
     }
+
+    #[test]
+    fn test_consolidating_container_builder() {
+        let mut ccb = <ConsolidatingContainerBuilder<Vec<(usize, usize, usize)>>>::default();
+        for _ in 0..1024 {
+            ccb.push((0, 0, 0));
+        }
+        assert_eq!(ccb.extract(), None);
+        assert_eq!(ccb.finish(), None);
+
+        for i in 0..1024 {
+            ccb.push((i, 0, 1));
+        }
+
+        let mut collected = Vec::default();
+        while let Some(container) = ccb.finish() {
+            collected.append(container);
+        }
+        // The output happens to be sorted, but it's not guaranteed.
+        collected.sort();
+        for i in 0..1024 {
+            assert_eq!((i, 0, 1), collected[i]);
+        }
+
+        ccb = Default::default();
+        ccb.push_container(&mut Vec::default());
+        assert_eq!(ccb.extract(), None);
+        assert_eq!(ccb.finish(), None);
+
+        ccb.push_container(&mut Vec::from_iter((0..1024).map(|i| (i, 0, 1))));
+        ccb.push_container(&mut Vec::from_iter((0..1024).map(|i| (i, 0, 1))));
+        collected.clear();
+        while let Some(container) = ccb.finish() {
+            collected.append(container);
+        }
+        // The output happens to be sorted, but it's not guaranteed.
+        consolidate_updates(&mut collected);
+        for i in 0..1024 {
+            assert_eq!((i, 0, 2), collected[i]);
+        }
+
+    }
 }

src/dynamic/mod.rs

@@ -60,7 +60,7 @@ where
                     vec.truncate(level - 1);
                     time.inner = PointStamp::new(vec);
                 }
-                output.session(&new_cap).give_vec(&mut vector);
+                output.session(&new_cap).give_container(&mut vector);
             });
         });
 

src/operators/arrange/arrangement.rs

@@ -255,7 +255,7 @@ where
         self.join_core_internal_unsafe(other, result)
     }
     /// A direct implementation of the `JoinCore::join_core_internal_unsafe` method.
-    pub fn join_core_internal_unsafe<T2,I,L,D,ROut> (&self, other: &Arranged<G,T2>, result: L) -> Collection<G,D,ROut>
+    pub fn join_core_internal_unsafe<T2,I,L,D,ROut> (&self, other: &Arranged<G,T2>, mut result: L) -> Collection<G,D,ROut>
     where
         T2: for<'a> TraceReader<Key<'a>=T1::Key<'a>, Time=T1::Time>+Clone+'static,
         D: Data,
@@ -264,7 +264,15 @@ where
         L: FnMut(T1::Key<'_>, T1::Val<'_>,T2::Val<'_>,&G::Timestamp,&T1::Diff,&T2::Diff)->I+'static,
     {
         use crate::operators::join::join_traces;
-        join_traces(self, other, result)
+        join_traces::<_, _, _, _, crate::consolidation::ConsolidatingContainerBuilder<_>>(
+            self,
+            other,
+            move |k, v1, v2, t, d1, d2, c| {
+                for datum in result(k, v1, v2, t, d1, d2) {
+                    c.give(datum);
+                }
+            }
+        )
             .as_collection()
     }
 }

src/operators/consolidate.rs

@@ -99,7 +99,7 @@ where
                     input.for_each(|time, data| {
                         data.swap(&mut vector);
                         crate::consolidation::consolidate_updates(&mut vector);
-                        output.session(&time).give_vec(&mut vector);
+                        output.session(&time).give_container(&mut vector);
                     })
                 }
             })