Optimization: merge UDDSketch using iterators

The previous implementation would create a UDDSketch (with a backing
HashMap) for every possible merge, and then call `compact_buckets` on
that in order to ensure the number of compactions between the target and
the source were equal.

Profiling this, we found out that in a `rollup` call of a lot of data,
the `compact_buckets` was pretty much the main contributor to all the
CPU time.

However, if we merge a different sketch into this sketch, we don't need
to actually compact_buckets all the time, we can directly consume the
keys and counts, and apply some compact_key calls to it.

This prevents a lot of heap allocations, as compact_buckets does a fully
copy of the backing `HashMap`, and then rebuilding it.

For a particular workload, this reduced the execution time from 30 to 12
seconds.

Author: feikesteenbergen

crates/udd-sketch/src/lib.rs

@@ -46,6 +46,15 @@ impl std::cmp::PartialOrd for SketchHashKey {
     }
 }
 
+/// `UDDSketchMetadata` was created to avoid passing along many parameters
+/// to function calls.
+pub struct UDDSketchMetadata {
+    pub max_buckets: u32,
+    pub current_error: f64,
+    pub compactions: u32,
+    pub values: u64,
+    pub sum: f64,
+}
 impl SketchHashKey {
     /// This is the key corresponding to the current key after the SketchHashMap it refers to has gone through one compaction.
     /// Note that odd buckets get combined with the bucket after them (i.e. old buckets -3 and -2 become new bucket -1, {-1, 0} -> 0, {1, 2} -> 1)
@@ -211,22 +220,18 @@ impl UDDSketch {
 
     // This constructor is used to recreate a UddSketch from it's component data
     pub fn new_from_data(
-        max_buckets: u64,
-        current_error: f64,
-        compactions: u64,
-        values: u64,
-        sum: f64,
+        metadata: &UDDSketchMetadata,
         keys: impl Iterator<Item = SketchHashKey>,
         counts: impl Iterator<Item = u64>,
     ) -> Self {
         let mut sketch = UDDSketch {
             buckets: SketchHashMap::new(),
-            alpha: current_error,
-            gamma: gamma(current_error),
-            compactions: compactions as u32,
-            max_buckets,
-            num_values: values,
-            values_sum: sum,
+            alpha: metadata.current_error,
+            gamma: gamma(metadata.current_error),
+            compactions: metadata.compactions,
+            max_buckets: metadata.max_buckets as u64,
+            num_values: metadata.values,
+            values_sum: metadata.sum,
         };
         // TODO
         let keys: Vec<_> = keys.collect();
@@ -285,6 +290,51 @@ impl UDDSketch {
         self.values_sum += value;
     }
 
+    /// `merge_items` will merge these values into the current sketch
+    /// it requires less memory than `merge_sketch`, as that needs a fully serialized
+    /// `UDDSketch`, whereas this function relies on iterators to do its job.
+    pub fn merge_items(
+        &mut self,
+        other: &UDDSketchMetadata,
+        mut keys: impl Iterator<Item = SketchHashKey>,
+        mut counts: impl Iterator<Item = u64>,
+    ) {
+        let other_gamma = gamma(other.current_error);
+        // Require matching initial parameters
+        debug_assert!(
+            (self
+                .gamma
+                .powf(1.0 / f64::powi(2.0, self.compactions as i32))
+                - other_gamma.powf(1.0 / f64::powi(2.0, other.compactions as i32)))
+            .abs()
+                < 1e-9 // f64::EPSILON too small, see issue #396
+        );
+        debug_assert_eq!(self.max_buckets, other.max_buckets as u64);
+
+        if other.values == 0 {
+            return;
+        }
+
+        while self.compactions < other.compactions {
+            self.compact_buckets();
+        }
+
+        let extra_compactions = self.compactions - other.compactions;
+        while let (Some(mut key), Some(count)) = (keys.next(), counts.next()) {
+            for _ in 0..extra_compactions {
+                key = key.compact_key();
+            }
+            self.buckets.entry(key).count += count;
+        }
+
+        while self.buckets.len() > self.max_buckets as usize {
+            self.compact_buckets();
+        }
+
+        self.num_values += other.values;
+        self.values_sum += other.sum;
+    }
+
     pub fn merge_sketch(&mut self, other: &UDDSketch) {
         // Require matching initial parameters
         assert!(

extension/src/uddsketch.rs

@@ -2,7 +2,7 @@ use pgrx::*;
 
 use encodings::{delta, prefix_varint};
 
-use uddsketch::{SketchHashKey, UDDSketch as UddSketchInternal};
+use uddsketch::{SketchHashKey, UDDSketch as UddSketchInternal, UDDSketchMetadata};
 
 use crate::{
     accessors::{
@@ -148,11 +148,13 @@ impl From<&UddSketchInternal> for SerializedUddSketch {
 impl From<SerializedUddSketch> for UddSketchInternal {
     fn from(sketch: SerializedUddSketch) -> Self {
         UddSketchInternal::new_from_data(
-            sketch.max_buckets as u64,
-            sketch.alpha,
-            sketch.compactions as u64,
-            sketch.count,
-            sketch.sum,
+            &UDDSketchMetadata {
+                max_buckets: sketch.max_buckets,
+                current_error: sketch.alpha,
+                compactions: sketch.compactions,
+                values: sketch.count,
+                sum: sketch.sum,
+            },
             sketch.keys(),
             sketch.counts(),
         )
@@ -304,16 +306,18 @@ impl<'input> UddSketch<'input> {
         )
     }
 
+    fn metadata(&self) -> UDDSketchMetadata {
+        UDDSketchMetadata {
+            max_buckets: self.max_buckets,
+            current_error: self.alpha,
+            compactions: self.compactions as u32,
+            values: self.count,
+            sum: self.sum,
+        }
+    }
+
     fn to_uddsketch(&self) -> UddSketchInternal {
-        UddSketchInternal::new_from_data(
-            self.max_buckets as u64,
-            self.alpha,
-            self.compactions,
-            self.count,
-            self.sum,
-            self.keys(),
-            self.counts(),
-        )
+        UddSketchInternal::new_from_data(&self.metadata(), self.keys(), self.counts())
     }
 
     fn from_internal(state: &UddSketchInternal) -> Self {
@@ -520,15 +524,12 @@ pub fn uddsketch_compound_trans_inner(
 ) -> Option<Inner<UddSketchInternal>> {
     unsafe {
         in_aggregate_context(fcinfo, || {
-            let value = match value {
-                None => return state,
-                Some(value) => value.to_uddsketch(),
-            };
-            let mut state = match state {
-                None => return Some(value.into()),
-                Some(state) => state,
+            let Some(value) = value else { return state };
+            let Some(mut state) = state else {
+                return Some(value.to_uddsketch().into());
             };
-            state.merge_sketch(&value);
+
+            state.merge_items(&value.metadata(), value.keys(), value.counts());
             state.into()
         })
     }