Merge batcher announce records

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

Author: antiguru

src/logging.rs

@@ -52,6 +52,8 @@ impl From<BatchEvent> for DifferentialEvent { fn from(e: BatchEvent) -> Self { D
 pub struct BatcherEvent {
     /// Operator identifier.
     pub operator: usize,
+    /// Change in records.
+    pub records_diff: isize,
     /// Change in used size.
     pub size_diff: isize,
     /// Change in capacity.

src/trace/implementations/merge_batcher.rs

@@ -8,7 +8,7 @@ use timely::logging_core::Logger;
 use timely::progress::{frontier::Antichain, Timestamp};
 
 use crate::difference::Semigroup;
-use crate::logging::DifferentialEvent;
+use crate::logging::{BatcherEvent, DifferentialEvent};
 use crate::trace::{Batcher, Builder};
 
 /// Creates batches from unordered tuples.
@@ -28,11 +28,11 @@ where
     type Item = ((K,V),T,D);
     type Time = T;
 
-    fn new(_logger: Option<Logger<DifferentialEvent, WorkerIdentifier>>, _operator_id: usize) -> Self {
+    fn new(logger: Option<Logger<DifferentialEvent, WorkerIdentifier>>, operator_id: usize) -> Self {
         MergeBatcher {
-            sorter: MergeSorter::new(),
+            sorter: MergeSorter::new(logger, operator_id),
             frontier: Antichain::new(),
-            lower: Antichain::from_elem(<T as timely::progress::Timestamp>::minimum()),
+            lower: Antichain::from_elem(T::minimum()),
         }
     }
 
@@ -134,20 +134,23 @@ where
             self.sorter.push(&mut buffer);
         }
 
-        let seal = builder.done(self.lower.clone(), upper.clone(), Antichain::from_elem(<T as timely::progress::Timestamp>::minimum()));
+        let seal = builder.done(self.lower.clone(), upper.clone(), Antichain::from_elem(T::minimum()));
         self.lower = upper;
         seal
     }
 
-    // the frontier of elements remaining after the most recent call to `self.seal`.
+    /// The frontier of elements remaining after the most recent call to `self.seal`.
     fn frontier(&mut self) -> timely::progress::frontier::AntichainRef<T> {
         self.frontier.borrow()
     }
 }
 
 struct MergeSorter<D, T, R> {
-    queue: Vec<Vec<Vec<(D, T, R)>>>,    // each power-of-two length list of allocations.
+    /// each power-of-two length list of allocations. Do not push/pop directly but use the corresponding functions.
+    queue: Vec<Vec<Vec<(D, T, R)>>>,
     stash: Vec<Vec<(D, T, R)>>,
+    logger: Option<Logger<DifferentialEvent, WorkerIdentifier>>,
+    operator_id: usize,
 }
 
 impl<D: Ord, T: Ord, R: Semigroup> MergeSorter<D, T, R> {
@@ -166,21 +169,20 @@ impl<D: Ord, T: Ord, R: Semigroup> MergeSorter<D, T, R> {
     }
 
     #[inline]
-    pub fn new() -> Self { MergeSorter { queue: Vec::new(), stash: Vec::new() } }
+    fn new(logger: Option<Logger<DifferentialEvent, WorkerIdentifier>>, operator_id: usize) -> Self {
+        Self {
+            logger,
+            operator_id,
+            queue: Vec::new(),
+            stash: Vec::new(),
+        }
+    }
 
     #[inline]
     pub fn empty(&mut self) -> Vec<(D, T, R)> {
         self.stash.pop().unwrap_or_else(|| Vec::with_capacity(Self::buffer_size()))
     }
 
-    #[inline(never)]
-    pub fn _sort(&mut self, list: &mut Vec<Vec<(D, T, R)>>) {
-        for mut batch in list.drain(..) {
-            self.push(&mut batch);
-        }
-        self.finish_into(list);
-    }
-
     #[inline]
     pub fn push(&mut self, batch: &mut Vec<(D, T, R)>) {
         // TODO: Reason about possible unbounded stash growth. How to / should we return them?
@@ -194,12 +196,13 @@ impl<D: Ord, T: Ord, R: Semigroup> MergeSorter<D, T, R> {
 
         if !batch.is_empty() {
             crate::consolidation::consolidate_updates(&mut batch);
-            self.queue.push(vec![batch]);
+            self.account([batch.len()], 1);
+            self.queue_push(vec![batch]);
             while self.queue.len() > 1 && (self.queue[self.queue.len()-1].len() >= self.queue[self.queue.len()-2].len() / 2) {
-                let list1 = self.queue.pop().unwrap();
-                let list2 = self.queue.pop().unwrap();
+                let list1 = self.queue_pop().unwrap();
+                let list2 = self.queue_pop().unwrap();
                 let merged = self.merge_by(list1, list2);
-                self.queue.push(merged);
+                self.queue_push(merged);
             }
         }
     }
@@ -208,31 +211,32 @@ impl<D: Ord, T: Ord, R: Semigroup> MergeSorter<D, T, R> {
     // to break it down to be so.
     pub fn push_list(&mut self, list: Vec<Vec<(D, T, R)>>) {
         while self.queue.len() > 1 && self.queue[self.queue.len()-1].len() < list.len() {
-            let list1 = self.queue.pop().unwrap();
-            let list2 = self.queue.pop().unwrap();
+            let list1 = self.queue_pop().unwrap();
+            let list2 = self.queue_pop().unwrap();
             let merged = self.merge_by(list1, list2);
-            self.queue.push(merged);
+            self.queue_push(merged);
         }
-        self.queue.push(list);
+        self.queue_push(list);
     }
 
     #[inline(never)]
     pub fn finish_into(&mut self, target: &mut Vec<Vec<(D, T, R)>>) {
         while self.queue.len() > 1 {
-            let list1 = self.queue.pop().unwrap();
-            let list2 = self.queue.pop().unwrap();
+            let list1 = self.queue_pop().unwrap();
+            let list2 = self.queue_pop().unwrap();
             let merged = self.merge_by(list1, list2);
-            self.queue.push(merged);
+            self.queue_push(merged);
         }
 
-        if let Some(mut last) = self.queue.pop() {
+        if let Some(mut last) = self.queue_pop() {
             ::std::mem::swap(&mut last, target);
         }
     }
 
     // merges two sorted input lists into one sorted output list.
     #[inline(never)]
     fn merge_by(&mut self, list1: Vec<Vec<(D, T, R)>>, list2: Vec<Vec<(D, T, R)>>) -> Vec<Vec<(D, T, R)>> {
+        self.account(list1.iter().chain(list2.iter()).map(Vec::len), -1);
 
         use std::cmp::Ordering;
 
@@ -307,3 +311,46 @@ impl<D: Ord, T: Ord, R: Semigroup> MergeSorter<D, T, R> {
         output
     }
 }
+
+impl<D, T, R> MergeSorter<D, T, R> {
+    /// Pop a batch from `self.queue` and account size changes.
+    #[inline]
+    fn queue_pop(&mut self) -> Option<Vec<Vec<(D, T, R)>>> {
+        let batch = self.queue.pop();
+        self.account(batch.iter().flatten().map(Vec::len), -1);
+        batch
+    }
+
+    /// Push a batch to `self.queue` and account size changes.
+    #[inline]
+    fn queue_push(&mut self, batch: Vec<Vec<(D, T, R)>>) {
+        self.account(batch.iter().map(Vec::len), 1);
+        self.queue.push(batch);
+    }
+
+    /// Account size changes. Only performs work if a logger exists.
+    ///
+    /// Calculate the size based on the [`TimelyStack`]s passed along, with each attribute
+    /// multiplied by `diff`. Usually, one wants to pass 1 or -1 as the diff.
+    fn account<I: IntoIterator<Item=usize>>(&self, items: I, diff: isize) {
+        if let Some(logger) = &self.logger {
+            let mut records= 0isize;
+            for len in items {
+                records = records.saturating_add_unsigned(len);
+            }
+            logger.log(BatcherEvent {
+                operator: self.operator_id,
+                records_diff: records * diff,
+                size_diff: 0,
+                capacity_diff: 0,
+                allocations_diff: 0,
+            })
+        }
+    }
+}
+
+impl<D, T, R> Drop for MergeSorter<D, T, R> {
+    fn drop(&mut self) {
+        while self.queue_pop().is_some() { }
+    }
+}

src/trace/implementations/merge_batcher_col.rs

@@ -14,10 +14,10 @@ use crate::trace::{Batcher, Builder};
 /// Creates batches from unordered tuples.
 pub struct ColumnatedMergeBatcher<K, V, T, D>
 where
-    K: Columnation,
-    V: Columnation,
-    T: Columnation,
-    D: Columnation,
+    K: Columnation + 'static,
+    V: Columnation + 'static,
+    T: Columnation + 'static,
+    D: Columnation + 'static,
 {
     sorter: MergeSorterColumnation<(K, V), T, D>,
     lower: Antichain<T>,
@@ -132,12 +132,12 @@ where
         // Drain buffers (fast reclamation).
         self.sorter.clear_stash();
 
-        let seal = builder.done(self.lower.clone(), upper.clone(), Antichain::from_elem(<T as timely::progress::Timestamp>::minimum()));
+        let seal = builder.done(self.lower.clone(), upper.clone(), Antichain::from_elem(T::minimum()));
         self.lower = upper;
         seal
     }
 
-    // the frontier of elements remaining after the most recent call to `self.seal`.
+    /// The frontier of elements remaining after the most recent call to `self.seal`.
     fn frontier(&mut self) -> timely::progress::frontier::AntichainRef<T> {
         self.frontier.borrow()
     }
@@ -184,8 +184,9 @@ impl<T: Columnation> TimelyStackQueue<T> {
     }
 }
 
-struct MergeSorterColumnation<D: Columnation, T: Columnation, R: Columnation> {
-    queue: Vec<Vec<TimelyStack<(D, T, R)>>>,    // each power-of-two length list of allocations.
+struct MergeSorterColumnation<D: Columnation + 'static, T: Columnation + 'static, R: Columnation + 'static> {
+    /// each power-of-two length list of allocations. Do not push/pop directly but use the corresponding functions.
+    queue: Vec<Vec<TimelyStack<(D, T, R)>>>,
     stash: Vec<TimelyStack<(D, T, R)>>,
     pending: Vec<(D, T, R)>,
     logger: Option<Logger<DifferentialEvent, WorkerIdentifier>>,
@@ -223,29 +224,6 @@ impl<D: Ord+Columnation+'static, T: Ord+Columnation+'static, R: Semigroup+Column
         }
     }
 
-    /// Account size changes. Only performs work if a logger exists.
-    ///
-    /// Calculate the size based on the [`TimelyStack`]s passed along, with each attribute
-    /// multiplied by `diff`. Usually, one wants to pass 1 or -1 as the diff.
-    fn account<'a, I: IntoIterator<Item=&'a TimelyStack<(D, T, R)>>>(&self, items: I, diff: isize) {
-        if let Some(logger) = &self.logger {
-            let (mut siz, mut capacity, mut allocations) = (0isize, 0isize, 0isize);
-            for stack in items {
-                stack.heap_size(|s, c| {
-                    siz = siz.saturating_add_unsigned(s);
-                    capacity = capacity.saturating_add_unsigned(c);
-                    allocations += isize::from(c > 0);
-                });
-            }
-            logger.log(BatcherEvent {
-                operator: self.operator_id,
-                size_diff: siz * diff,
-                capacity_diff: capacity * diff,
-                allocations_diff: allocations * diff,
-            })
-        }
-    }
-
     fn empty(&mut self) -> TimelyStack<(D, T, R)> {
         self.stash.pop().unwrap_or_else(|| TimelyStack::with_capacity(Self::buffer_size()))
     }
@@ -290,14 +268,12 @@ impl<D: Ord+Columnation+'static, T: Ord+Columnation+'static, R: Semigroup+Column
             for tuple in self.pending.drain(..) {
                 stack.copy(&tuple);
             }
-            self.account([&stack], 1);
-            self.queue.push(vec![stack]);
+            self.queue_push(vec![stack]);
             while self.queue.len() > 1 && (self.queue[self.queue.len()-1].len() >= self.queue[self.queue.len()-2].len() / 2) {
-                let list1 = self.queue.pop().unwrap();
-                let list2 = self.queue.pop().unwrap();
+                let list1 = self.queue_pop().unwrap();
+                let list2 = self.queue_pop().unwrap();
                 let merged = self.merge_by(list1, list2);
-                self.account(&merged, 1);
-                self.queue.push(merged);
+                self.queue_push(merged);
             }
         }
     }
@@ -306,34 +282,31 @@ impl<D: Ord+Columnation+'static, T: Ord+Columnation+'static, R: Semigroup+Column
     // to break it down to be so.
     fn push_list(&mut self, list: Vec<TimelyStack<(D, T, R)>>) {
         while self.queue.len() > 1 && self.queue[self.queue.len()-1].len() < list.len() {
-            let list1 = self.queue.pop().unwrap();
-            let list2 = self.queue.pop().unwrap();
+            let list1 = self.queue_pop().unwrap();
+            let list2 = self.queue_pop().unwrap();
             let merged = self.merge_by(list1, list2);
-            self.account(&merged, 1);
-            self.queue.push(merged);
+            self.queue_push(merged);
         }
-        self.queue.push(list);
+        self.queue_push(list);
     }
 
     fn finish_into(&mut self, target: &mut Vec<TimelyStack<(D, T, R)>>) {
         crate::consolidation::consolidate_updates(&mut self.pending);
         self.flush_pending();
         while self.queue.len() > 1 {
-            let list1 = self.queue.pop().unwrap();
-            let list2 = self.queue.pop().unwrap();
+            let list1 = self.queue_pop().unwrap();
+            let list2 = self.queue_pop().unwrap();
             let merged = self.merge_by(list1, list2);
-            self.queue.push(merged);
+            self.queue_push(merged);
         }
 
-        if let Some(mut last) = self.queue.pop() {
+        if let Some(mut last) = self.queue_pop() {
             std::mem::swap(&mut last, target);
         }
     }
 
     // merges two sorted input lists into one sorted output list.
     fn merge_by(&mut self, list1: Vec<TimelyStack<(D, T, R)>>, list2: Vec<TimelyStack<(D, T, R)>>) -> Vec<TimelyStack<(D, T, R)>> {
-        self.account(list1.iter().chain(list2.iter()), -1);
-
         use std::cmp::Ordering;
 
         // TODO: `list1` and `list2` get dropped; would be better to reuse?
@@ -411,3 +384,52 @@ impl<D: Ord+Columnation+'static, T: Ord+Columnation+'static, R: Semigroup+Column
         output
     }
 }
+
+impl<D: Columnation + 'static, T: Columnation + 'static, R: Columnation + 'static> MergeSorterColumnation<D, T, R> {
+    /// Pop a batch from `self.queue` and account size changes.
+    #[inline]
+    fn queue_pop(&mut self) -> Option<Vec<TimelyStack<(D, T, R)>>> {
+        let batch = self.queue.pop();
+        self.account(batch.iter().flatten(), -1);
+        batch
+    }
+
+    /// Push a batch to `self.queue` and account size changes.
+    #[inline]
+    fn queue_push(&mut self, batch: Vec<TimelyStack<(D, T, R)>>) {
+        self.account(&batch, 1);
+        self.queue.push(batch);
+    }
+
+    /// Account size changes. Only performs work if a logger exists.
+    ///
+    /// Calculate the size based on the [`TimelyStack`]s passed along, with each attribute
+    /// multiplied by `diff`. Usually, one wants to pass 1 or -1 as the diff.
+    fn account<'a, I: IntoIterator<Item=&'a TimelyStack<(D, T, R)>>>(&self, items: I, diff: isize) {
+        if let Some(logger) = &self.logger {
+            let (mut records, mut siz, mut capacity, mut allocations) = (0isize, 0isize, 0isize, 0isize);
+            for stack in items {
+                records = records.saturating_add_unsigned(stack.len());
+                stack.heap_size(|s, c| {
+                    siz = siz.saturating_add_unsigned(s);
+                    capacity = capacity.saturating_add_unsigned(c);
+                    allocations += isize::from(c > 0);
+                });
+            }
+            logger.log(BatcherEvent {
+                operator: self.operator_id,
+                records_diff: records * diff,
+                size_diff: siz * diff,
+                capacity_diff: capacity * diff,
+                allocations_diff: allocations * diff,
+            })
+        }
+    }
+
+}
+
+impl<D: Columnation + 'static, T: Columnation + 'static, R: Columnation + 'static> Drop for MergeSorterColumnation<D, T, R> {
+    fn drop(&mut self) {
+        while self.queue_pop().is_some() { }
+    }
+}

src/trace/mod.rs

@@ -17,8 +17,8 @@ use timely::logging::WorkerIdentifier;
 use timely::logging_core::Logger;
 use timely::progress::{Antichain, frontier::AntichainRef};
 use timely::progress::Timestamp;
-use logging::DifferentialEvent;
 
+use crate::logging::DifferentialEvent;
 use crate::trace::cursor::MyTrait;
 
 // use ::difference::Semigroup;