Add finalize step to the Aggregator trait.

A recent extension of the `Aggregator` trait requires aggregated values
to form a semigroup, which is not the case for instance for the AVG
aggregate.  The standard solution is to use an intermediate
representation of the aggregate that forms a semigroup and only compute
the actual value of the aggregate after all partitions have been combined.

To support such aggregates we extend the `Aggregator` trait with the
`Accumulator` associated type returned by `aggregate`,
which must implement `trait Semigroup`.  The
`Aggregator::Output` associated type represents the final output of the
aggregator computed by applying the `finalize` method to the final value
of the accumulator.

Author: ryzhyk

benches/fraud.rs

@@ -195,8 +195,8 @@ struct Args {
 type Weight = i32;
 
 type EnrichedTransactions = OrdIndexedZSet<(F64, i64), (Transaction, Demographics), Weight>;
-type AverageSpendingPerWeek = OrdPartitionedIndexedZSet<F64, i64, Option<Avg<F64, i32>>, Weight>;
-type AverageSpendingPerMonth = OrdPartitionedIndexedZSet<F64, i64, Option<Avg<F64, i32>>, Weight>;
+type AverageSpendingPerWeek = OrdPartitionedIndexedZSet<F64, i64, Option<F64>, Weight>;
+type AverageSpendingPerMonth = OrdPartitionedIndexedZSet<F64, i64, Option<F64>, Weight>;
 type TransactionFrequency = OrdPartitionedIndexedZSet<F64, i64, Option<i32>, Weight>;
 
 struct FraudBenchmark {
@@ -233,16 +233,12 @@ impl FraudBenchmark {
             let avg_spend_pw: Stream<_, AverageSpendingPerWeek> = amounts
                 .partitioned_rolling_aggregate_linear(
                     |amt| Avg::new(*amt, 1),
+                    |avg| avg.compute_avg().unwrap(),
                     RelRange::new(RelOffset::Before(DAY_IN_SECONDS * 7), RelOffset::Before(1)),
                 );
 
-            // TODO: this should be returned directly by `partitioned_rolling_aggregate`
-            let avg_spend_pw_indexed = avg_spend_pw.map_index(|(cc_num, (ts, avg_amt))| {
-                (
-                    (*cc_num, *ts),
-                    avg_amt.as_ref().and_then(|avg| avg.compute_avg()),
-                )
-            });
+            let avg_spend_pw_indexed =
+                avg_spend_pw.map_index(|(cc_num, (ts, avg_amt))| ((*cc_num, *ts), *avg_amt));
 
             // AVG(amt) OVER(
             //     PARTITION BY  CAST(cc_num AS NUMERIC)
@@ -252,15 +248,12 @@ impl FraudBenchmark {
             let avg_spend_pm: Stream<_, AverageSpendingPerMonth> = amounts
                 .partitioned_rolling_aggregate_linear(
                     |amt| Avg::new(*amt, 1),
+                    |avg| avg.compute_avg().unwrap(),
                     RelRange::new(RelOffset::Before(DAY_IN_SECONDS * 30), RelOffset::Before(1)),
                 );
 
-            let avg_spend_pm_indexed = avg_spend_pm.map_index(|(cc_num, (ts, avg_amt))| {
-                (
-                    (*cc_num, *ts),
-                    avg_amt.as_ref().and_then(|avg| avg.compute_avg()),
-                )
-            });
+            let avg_spend_pm_indexed =
+                avg_spend_pm.map_index(|(cc_num, (ts, avg_amt))| ((*cc_num, *ts), *avg_amt));
 
             // COUNT(*) OVER(
             //     PARTITION BY  CAST(cc_num AS NUMERIC)
@@ -270,6 +263,7 @@ impl FraudBenchmark {
             let trans_freq_24: Stream<_, TransactionFrequency> = amounts
                 .partitioned_rolling_aggregate_linear(
                     |_amt| 1,
+                    |cnt| cnt,
                     RelRange::new(RelOffset::Before(DAY_IN_SECONDS), RelOffset::Before(1)),
                 );
 

src/operator/aggregate/fold.rs

@@ -63,16 +63,17 @@ impl<V, T, R, A, S, O, SF, OF> Aggregator<V, T, R> for Fold<A, S, SF, OF>
 where
     T: Timestamp,
     R: MonoidValue,
-    A: Clone + 'static,
+    A: DBData,
     SF: Fn(&mut A, &V, R) + Clone + 'static,
     OF: Fn(A) -> O + Clone + 'static,
-    S: Semigroup<O> + Clone + 'static,
+    S: Semigroup<A> + Clone + 'static,
     O: DBData,
 {
+    type Accumulator = A;
     type Output = O;
     type Semigroup = S;
 
-    fn aggregate<'s, C>(&self, cursor: &mut C) -> Option<Self::Output>
+    fn aggregate<'s, C>(&self, cursor: &mut C) -> Option<Self::Accumulator>
     where
         C: Cursor<'s, V, (), T, R>,
     {
@@ -91,8 +92,10 @@ where
             cursor.step_key();
         }
 
-        // Aggregator must return None iff the input cursor is empty (all keys have
-        // weight 0).
-        non_empty.then(|| (self.output)(acc))
+        non_empty.then_some(acc)
+    }
+
+    fn finalize(&self, acc: Self::Accumulator) -> Self::Output {
+        (self.output)(acc)
     }
 }

src/operator/aggregate/max.rs

@@ -28,11 +28,12 @@ where
     T: Timestamp,
     R: MonoidValue,
 {
+    type Accumulator = V;
     type Output = V;
     type Semigroup = MaxSemigroup<V>;
 
     // TODO: this can be more efficient with reverse iterator.
-    fn aggregate<'s, C>(&self, cursor: &mut C) -> Option<Self::Output>
+    fn aggregate<'s, C>(&self, cursor: &mut C) -> Option<Self::Accumulator>
     where
         C: Cursor<'s, V, (), T, R>,
     {
@@ -52,4 +53,8 @@ where
 
         result
     }
+
+    fn finalize(&self, accumulator: Self::Accumulator) -> Self::Output {
+        accumulator
+    }
 }

src/operator/aggregate/min.rs

@@ -30,6 +30,7 @@ where
     T: Timestamp,
     R: MonoidValue,
 {
+    type Accumulator = V;
     type Output = V;
     type Semigroup = MinSemigroup<V>;
 
@@ -53,4 +54,8 @@ where
 
         None
     }
+
+    fn finalize(&self, accumulator: Self::Accumulator) -> Self::Output {
+        accumulator
+    }
 }

src/operator/aggregate/mod.rs

@@ -39,13 +39,42 @@ pub use min::Min;
 /// A trait for aggregator objects.  An aggregator summarizes the contents
 /// of a Z-set into a single value.
 ///
+/// This trait supports two aggregation methods that can be combined in
+/// various ways to compute the aggregate efficiently.  First, the
+/// [`aggregate`](`Self::aggregate`) method takes a cursor pointing to
+/// the first key of a Z-set and scans the cursor to compute the aggregate
+/// over all values in the cursor.  Second, the
+/// [`Semigroup`](`Self::Semigroup`) associated type allows aggregating
+/// partitioned Z-sets by combining aggregates computed over individual
+/// partitions (e.g., computed by different worker threads).
+///
+/// This design requires aggregate values to form a semigroup with a `+`
+/// operation such that `Agg(x + y) = Agg(x) + Agg(y)`, i.e., aggregating
+/// a union of two Z-sets must produce the same result as the sum of
+/// individual aggregates.  Not all aggregates have this property.  E.g.,
+/// the average value of a Z-set cannot be computed by combining the
+/// averages of its subsets.  We can get around this problem by representing
+/// average as `(sum, count)` tuple with point-wise `+` operator
+/// `(sum1, count1) + (sum2, count2) = (sum1+sum2, count1+count2)`.
+/// The final result is converted to an actual average value by dividing
+/// the first element of the tuple by the second.  This is a general technique
+/// that works for all aggregators (although it may not always be optimal).
+///
+/// To support such aggregates, this trait distinguishes between the
+/// `Accumulator` type returned by [`aggregate`](`Self::aggregate`),
+/// which must implement
+/// [`trait Semigroup`](`crate::algebra::Semigroup`), and the final
+/// [`Output`](`Self::Output`) of the aggregator.  The latter is
+/// computed by applying the [`finalize`](`Self::finalize`) method
+/// to the final value of the accumulator.
+///
 /// This is a low-level trait that is mostly used to build libraries of
 /// aggregators.  Users will typicaly work with ready-made implementations
 /// like [`Min`] and [`Fold`].
 // TODO: Owned aggregation using `Consumer`
 pub trait Aggregator<K, T, R>: Clone + 'static {
-    /// Aggregate type output by this aggregator.
-    type Output: DBData;
+    /// Accumulator type returned by
+    type Accumulator: DBData;
 
     /// Semigroup structure over aggregate values.
     ///
@@ -61,10 +90,11 @@ pub trait Aggregator<K, T, R>: Clone + 'static {
     // per-worker aggregates computes over arbitrary subsets of values,
     // which additionally requires commutativity.  Do we want to introduce
     // the `CommutativeSemigroup` trait?
-    type Semigroup: Semigroup<Self::Output>;
+    type Semigroup: Semigroup<Self::Accumulator>;
+
+    /// Aggregate type produced by this aggregator.
+    type Output: DBData;
 
-    /// Compute an aggregate over a Z-set.
-    ///
     /// Takes a cursor pointing to the first key of a Z-set and outputs
     /// an aggregate of the Z-set.
     ///
@@ -75,9 +105,20 @@ pub trait Aggregator<K, T, R>: Clone + 'static {
     ///
     /// * The method must return `None` if the total weight of each key is zero.
     ///   It must return `Some` otherwise.
-    fn aggregate<'s, C>(&self, cursor: &mut C) -> Option<Self::Output>
+    fn aggregate<'s, C>(&self, cursor: &mut C) -> Option<Self::Accumulator>
     where
         C: Cursor<'s, K, (), T, R>;
+
+    /// Compute the final value of the aggregate.
+    fn finalize(&self, accumulator: Self::Accumulator) -> Self::Output;
+
+    /// Applies `aggregate` to `cursor` followed by `finalize`.
+    fn aggregate_and_finalize<'s, C>(&self, cursor: &mut C) -> Option<Self::Output>
+    where
+        C: Cursor<'s, K, (), T, R>,
+    {
+        self.aggregate(cursor).map(|x| self.finalize(x))
+    }
 }
 
 /// Aggregator used internally by [`Stream::aggregate_linear`].  Computes
@@ -90,6 +131,7 @@ where
     T: Timestamp,
     R: DBWeight,
 {
+    type Accumulator = R;
     type Output = R;
     type Semigroup = DefaultSemigroup<R>;
 
@@ -107,6 +149,10 @@ where
             Some(weight)
         }
     }
+
+    fn finalize(&self, accumulator: Self::Accumulator) -> Self::Output {
+        accumulator
+    }
 }
 
 impl<P, Z> Stream<Circuit<P>, Z>
@@ -329,7 +375,7 @@ where
         while cursor.key_valid() {
             if let Some(agg) = self
                 .aggregator
-                .aggregate(&mut CursorGroup::new(&mut cursor, ()))
+                .aggregate_and_finalize(&mut CursorGroup::new(&mut cursor, ()))
             {
                 builder.push((O::item_from(cursor.key().clone(), agg), O::R::one()));
             }
@@ -461,7 +507,7 @@ where
             // Z-set associated with `input_cursor.key()` at time `self.time`.
             if let Some(aggregate) = self
                 .aggregator
-                .aggregate(&mut CursorGroup::new(input_cursor, self.time.clone()))
+                .aggregate_and_finalize(&mut CursorGroup::new(input_cursor, self.time.clone()))
             {
                 output.push((key.clone(), Some(aggregate)));
             } else {

src/operator/time_series/radix_tree/partitioned_tree_aggregate.rs

@@ -115,15 +115,15 @@ where
     pub fn partitioned_tree_aggregate<TS, V, Agg>(
         &self,
         aggregator: Agg,
-    ) -> OrdPartitionedRadixTreeStream<Z::Key, TS, Agg::Output, isize>
+    ) -> OrdPartitionedRadixTreeStream<Z::Key, TS, Agg::Accumulator, isize>
     where
         Z: PartitionedIndexedZSet<TS, V> + SizeOf,
         TS: DBData + PrimInt,
         V: DBData,
         Agg: Aggregator<V, (), Z::R>,
-        Agg::Output: Default,
+        Agg::Accumulator: Default,
     {
-        self.partitioned_tree_aggregate_generic::<TS, V, Agg, OrdPartitionedRadixTree<Z::Key, TS, Agg::Output, isize>>(
+        self.partitioned_tree_aggregate_generic::<TS, V, Agg, OrdPartitionedRadixTree<Z::Key, TS, Agg::Accumulator, isize>>(
             aggregator,
         )
     }
@@ -139,8 +139,8 @@ where
         TS: DBData + PrimInt,
         V: DBData,
         Agg: Aggregator<V, (), Z::R>,
-        Agg::Output: Default,
-        O: PartitionedRadixTreeBatch<TS, Agg::Output, Key = Z::Key>,
+        Agg::Accumulator: Default,
+        O: PartitionedRadixTreeBatch<TS, Agg::Accumulator, Key = Z::Key>,
         O::R: ZRingValue,
     {
         self.circuit()
@@ -345,10 +345,10 @@ where
     TS: DBData + PrimInt,
     V: DBData,
     IT: PartitionedBatchReader<TS, V, Key = Z::Key, R = Z::R> + Clone,
-    OT: PartitionedRadixTreeReader<TS, Agg::Output, Key = Z::Key, R = O::R> + Clone,
+    OT: PartitionedRadixTreeReader<TS, Agg::Accumulator, Key = Z::Key, R = O::R> + Clone,
     Agg: Aggregator<V, (), Z::R>,
-    Agg::Output: Default,
-    O: PartitionedRadixTreeBatch<TS, Agg::Output, Key = Z::Key>,
+    Agg::Accumulator: Default,
+    O: PartitionedRadixTreeBatch<TS, Agg::Accumulator, Key = Z::Key>,
     O::R: ZRingValue,
 {
     fn eval<'a>(

src/operator/time_series/radix_tree/tree_aggregate.rs

@@ -55,14 +55,16 @@ where
     pub fn tree_aggregate<Agg>(
         &self,
         aggregator: Agg,
-    ) -> Stream<Circuit<P>, OrdRadixTree<Z::Key, Agg::Output, isize>>
+    ) -> Stream<Circuit<P>, OrdRadixTree<Z::Key, Agg::Accumulator, isize>>
     where
         Z: IndexedZSet + SizeOf + NumEntries + Send,
         Z::Key: PrimInt,
         Agg: Aggregator<Z::Val, (), Z::R>,
-        Agg::Output: Default,
+        Agg::Accumulator: Default,
     {
-        self.tree_aggregate_generic::<Agg, OrdRadixTree<Z::Key, Agg::Output, isize>>(aggregator)
+        self.tree_aggregate_generic::<Agg, OrdRadixTree<Z::Key, Agg::Accumulator, isize>>(
+            aggregator,
+        )
     }
 
     /// Like [`Self::tree_aggregate`], but can return any batch type.
@@ -71,8 +73,8 @@ where
         Z: IndexedZSet + SizeOf + NumEntries + Send,
         Z::Key: PrimInt,
         Agg: Aggregator<Z::Val, (), Z::R>,
-        Agg::Output: Default,
-        O: RadixTreeBatch<Z::Key, Agg::Output>,
+        Agg::Accumulator: Default,
+        O: RadixTreeBatch<Z::Key, Agg::Accumulator>,
         O::R: ZRingValue,
     {
         self.circuit()
@@ -186,10 +188,10 @@ where
     Z: IndexedZSet,
     Z::Key: PrimInt,
     IT: BatchReader<Key = Z::Key, Val = Z::Val, Time = (), R = Z::R> + Clone,
-    OT: RadixTreeReader<Z::Key, Agg::Output, R = O::R> + Clone,
+    OT: RadixTreeReader<Z::Key, Agg::Accumulator, R = O::R> + Clone,
     Agg: Aggregator<Z::Val, (), Z::R>,
-    Agg::Output: Default,
-    O: RadixTreeBatch<Z::Key, Agg::Output>,
+    Agg::Accumulator: Default,
+    O: RadixTreeBatch<Z::Key, Agg::Accumulator>,
     O::R: ZRingValue,
 {
     fn eval<'a>(

src/operator/time_series/radix_tree/updater.rs

@@ -369,19 +369,19 @@ pub(super) fn radix_tree_update<'a, 'b, TS, V, R, Agg, UC, IC, TC, OR>(
     mut input: IC,
     tree: TC,
     aggregator: &Agg,
-    output_updates: &'a mut Vec<TreeNodeUpdate<TS, Agg::Output>>,
+    output_updates: &'a mut Vec<TreeNodeUpdate<TS, Agg::Accumulator>>,
 ) where
     TS: PrimInt + Debug,
     R: MonoidValue,
     Agg: Aggregator<V, (), R>,
-    Agg::Output: Clone + Default + Eq + Debug,
+    Agg::Accumulator: Clone + Default + Eq + Debug,
     UC: Cursor<'b, TS, V, (), R>,
     IC: Cursor<'b, TS, V, (), R>,
-    TC: RadixTreeCursor<'b, TS, Agg::Output, OR>,
+    TC: RadixTreeCursor<'b, TS, Agg::Accumulator, OR>,
     OR: MonoidValue,
 {
     let mut tree_updater =
-        <TreeUpdater<'a, TS, Agg::Output, OR, Agg::Semigroup, TC>>::new(tree, output_updates);
+        <TreeUpdater<'a, TS, Agg::Accumulator, OR, Agg::Semigroup, TC>>::new(tree, output_updates);
 
     while input_delta.key_valid() {
         //println!("affected key {:x?}", input_delta.key());