Tighten up constraints on UDDSketch config

Buckets should be within 1 and 44 million (practical PostgreSQL limit
for column size), which is better caught by a NonZeroU32.

Compactions are limited to at most 65. Using a u8 for that allows some
further information in the struct without requiring more memory.

Author: feikesteenbergen

crates/udd-sketch/src/lib.rs

@@ -7,9 +7,10 @@ use std::collections::HashMap;
 
 #[cfg(test)]
 use ordered_float::OrderedFloat;
+use serde::{Deserialize, Serialize};
 #[cfg(test)]
 use std::collections::HashSet;
-use serde::{Deserialize, Serialize};
+use std::num::NonZeroU32;
 
 #[cfg(test)]
 extern crate quickcheck;
@@ -57,9 +58,13 @@ impl PartialOrd for SketchHashKey {
 /// `UDDSketchMetadata` was created to avoid passing along many parameters
 /// to function calls.
 pub struct UDDSketchMetadata {
-    pub max_buckets: u32,
+    // `max_buckets` is limited to  `PostgreSQL`'s limit: At most 1GB per column.
+    // As we need 24 bytes to represent any bucket, there are at most something like
+    // 44 million values possible. It can however never be actually zero.
+    pub max_buckets: NonZeroU32,
     pub current_error: f64,
-    pub compactions: u32,
+    // Compactions are limited to at most 65
+    pub compactions: u8,
     pub values: u64,
     pub sum: f64,
 }
@@ -259,14 +264,14 @@ pub struct UDDSketch {
     buckets: SketchHashMap,
     alpha: f64,
     gamma: f64,
-    compactions: u32, // should always be smaller than 64
-    max_buckets: u64,
+    compactions: u8,         // should always be smaller than 64
+    max_buckets: NonZeroU32, // PostgreSQL limit is 1GB per column, which is roughly 44 million buckets
     num_values: u64,
     values_sum: f64,
 }
 
 impl UDDSketch {
-    pub fn new(max_buckets: u64, initial_error: f64) -> Self {
+    pub fn new(max_buckets: NonZeroU32, initial_error: f64) -> Self {
         assert!((1e-12..1.0).contains(&initial_error));
         UDDSketch {
             buckets: SketchHashMap::new(),
@@ -290,7 +295,7 @@ impl UDDSketch {
             alpha: metadata.current_error,
             gamma: gamma(metadata.current_error),
             compactions: metadata.compactions,
-            max_buckets: metadata.max_buckets as u64,
+            max_buckets: metadata.max_buckets,
             num_values: metadata.values,
             values_sum: metadata.sum,
         };
@@ -343,7 +348,7 @@ impl UDDSketch {
     pub fn add_value(&mut self, value: f64) {
         self.buckets.increment(self.key(value));
 
-        while self.buckets.len() > self.max_buckets as usize {
+        while self.buckets.len() > self.max_buckets.get() as usize {
             self.compact_buckets();
         }
 
@@ -370,7 +375,7 @@ impl UDDSketch {
             .abs()
                 < 1e-9 // f64::EPSILON too small, see issue #396
         );
-        debug_assert_eq!(self.max_buckets, other.max_buckets as u64);
+        debug_assert_eq!(self.max_buckets, other.max_buckets);
 
         if other.values == 0 {
             return;
@@ -389,7 +394,7 @@ impl UDDSketch {
             self.buckets.entry_upsert(key, count);
         }
 
-        while self.buckets.len() > self.max_buckets as usize {
+        while self.buckets.len() > self.max_buckets.get() as usize {
             self.compact_buckets();
         }
 
@@ -433,19 +438,19 @@ impl UDDSketch {
             self.buckets.entry_upsert(key, value);
         }
 
-        while self.buckets.len() > self.max_buckets as usize {
+        while self.buckets.len() > self.max_buckets.get() as usize {
             self.compact_buckets();
         }
 
         self.num_values += other.num_values;
         self.values_sum += other.values_sum;
     }
 
-    pub fn max_allowed_buckets(&self) -> u64 {
-        self.max_buckets
+    pub fn max_allowed_buckets(&self) -> u32 {
+        self.max_buckets.get()
     }
 
-    pub fn times_compacted(&self) -> u32 {
+    pub fn times_compacted(&self) -> u8 {
         self.compactions
     }
 
@@ -591,7 +596,7 @@ mod tests {
 
     #[test]
     fn build_and_add_values() {
-        let mut sketch = UDDSketch::new(20, 0.1);
+        let mut sketch = UDDSketch::new(NonZeroU32::new(20).unwrap(), 0.1);
         sketch.add_value(1.0);
         sketch.add_value(3.0);
         sketch.add_value(0.5);
@@ -603,7 +608,7 @@ mod tests {
 
     #[test]
     fn exceed_buckets() {
-        let mut sketch = UDDSketch::new(20, 0.1);
+        let mut sketch = UDDSketch::new(NonZeroU32::new(20).unwrap(), 0.1);
         sketch.add_value(1.1); // Bucket #1
         sketch.add_value(400.0); // Bucket #30
         let a2 = 0.2 / 1.01;
@@ -641,7 +646,7 @@ mod tests {
         }
 
         let metadata = UDDSketchMetadata {
-            max_buckets: other.max_buckets as u32,
+            max_buckets: other.max_buckets,
             current_error: other.alpha,
             compactions: other.compactions,
             values: other.num_values,
@@ -662,7 +667,7 @@ mod tests {
         let a4 = 2.0 * a3 / (1.0 + f64::powi(a3, 2)); // alpha for 3 compactions
         let a5 = 2.0 * a4 / (1.0 + f64::powi(a4, 2)); // alpha for 4 compactions
 
-        let mut sketch1 = UDDSketch::new(20, 0.1);
+        let mut sketch1 = UDDSketch::new(NonZeroU32::new(20).unwrap(), 0.1);
         sketch1.add_value(1.1); // Bucket #1
         sketch1.add_value(1.5); // Bucket #3
         sketch1.add_value(1.6); // Bucket #3
@@ -672,7 +677,7 @@ mod tests {
         assert_eq!(sketch1.count(), 5);
         assert_eq!(sketch1.max_error(), a1);
 
-        let mut sketch2 = UDDSketch::new(20, 0.1);
+        let mut sketch2 = UDDSketch::new(NonZeroU32::new(20).unwrap(), 0.1);
         sketch2.add_value(5.1); // Bucket #9
         sketch2.add_value(7.5); // Bucket #11
         sketch2.add_value(10.6); // Bucket #12
@@ -685,7 +690,7 @@ mod tests {
         assert_eq!(sketch1.count(), 10);
         assert_eq!(sketch1.max_error(), a1);
 
-        let mut sketch3 = UDDSketch::new(20, 0.1);
+        let mut sketch3 = UDDSketch::new(NonZeroU32::new(20).unwrap(), 0.1);
         sketch3.add_value(0.8); // Bucket #-1
         sketch3.add_value(3.7); // Bucket #7
         sketch3.add_value(15.2); // Bucket #14
@@ -698,7 +703,7 @@ mod tests {
         assert_eq!(sketch1.count(), 15);
         assert_eq!(sketch1.max_error(), a1);
 
-        let mut sketch4 = UDDSketch::new(20, 0.1);
+        let mut sketch4 = UDDSketch::new(NonZeroU32::new(20).unwrap(), 0.1);
         sketch4.add_value(400.0); // Bucket #30
         sketch4.add_value(0.004); // Bucket #-27
         sketch4.add_value(0.0); // Zero Bucket
@@ -716,7 +721,7 @@ mod tests {
         assert_eq!(sketch1.count(), 24);
         assert_eq!(sketch1.max_error(), a2);
 
-        let mut sketch5 = UDDSketch::new(20, 0.1);
+        let mut sketch5 = UDDSketch::new(NonZeroU32::new(20).unwrap(), 0.1);
         for i in 100..220 {
             sketch5.add_value(1.23_f64.powi(i));
         }
@@ -731,7 +736,7 @@ mod tests {
 
     #[test]
     fn test_quantile_and_value_estimates() {
-        let mut sketch = UDDSketch::new(50, 0.1);
+        let mut sketch = UDDSketch::new(NonZeroU32::new(50).unwrap(), 0.1);
         for v in 1..=10000 {
             sketch.add_value(v as f64 / 100.0);
         }
@@ -771,7 +776,7 @@ mod tests {
 
     #[test]
     fn test_extreme_quantile_at_value() {
-        let mut sketch = UDDSketch::new(50, 0.1);
+        let mut sketch = UDDSketch::new(NonZeroU32::new(50).unwrap(), 0.1);
         for v in 1..=10000 {
             sketch.add_value(v as f64 / 100.0);
         }
@@ -786,7 +791,7 @@ mod tests {
 
     #[test]
     fn random_stress() {
-        let mut sketch = UDDSketch::new(1000, 0.01);
+        let mut sketch = UDDSketch::new(NonZeroU32::new(1000).unwrap(), 0.01);
         let seed = rand::thread_rng().gen();
         let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
         let mut bounds = Vec::new();
@@ -917,7 +922,7 @@ mod tests {
         if master.len() < 100 {
             return TestResult::discard();
         }
-        let mut sketch = UDDSketch::new(100, 0.000001);
+        let mut sketch = UDDSketch::new(NonZeroU32::new(100).unwrap(), 0.000001);
         for value in &master {
             sketch.add_value(*value);
         }

extension/src/uddsketch.rs

@@ -1,9 +1,13 @@
 use pgrx::*;
+use std::num::NonZeroU32;
 
 use encodings::{delta, prefix_varint};
 
 use uddsketch::{SketchHashKey, UDDSketch as UddSketchInternal, UDDSketchMetadata};
 
+const ERROR_UDDSKETCH_ZERO_BUCKETS: &str = "uddsketch size most be 1 or more";
+const ERROR_COMPACTIONS_LIMITED: &str = "compactions are limited to at most 65";
+
 use crate::{
     accessors::{
         AccessorApproxPercentile, AccessorApproxPercentileRank, AccessorError, AccessorMean,
@@ -25,12 +29,20 @@ pub fn uddsketch_trans(
     value: Option<f64>,
     fcinfo: pg_sys::FunctionCallInfo,
 ) -> Option<Internal> {
-    uddsketch_trans_inner(unsafe { state.to_inner() }, size, max_error, value, fcinfo).internal()
+    uddsketch_trans_inner(
+        unsafe { state.to_inner() },
+        NonZeroU32::new(u32::try_from(size).expect(ERROR_UDDSKETCH_ZERO_BUCKETS))
+            .expect(ERROR_UDDSKETCH_ZERO_BUCKETS),
+        max_error,
+        value,
+        fcinfo,
+    )
+    .internal()
 }
 
 pub fn uddsketch_trans_inner(
     state: Option<Inner<UddSketchInternal>>,
-    size: i32,
+    size: NonZeroU32,
     max_error: f64,
     value: Option<f64>,
     fcinfo: pg_sys::FunctionCallInfo,
@@ -42,7 +54,7 @@ pub fn uddsketch_trans_inner(
                 Some(value) => value,
             };
             let mut state = match state {
-                None => UddSketchInternal::new(size as u64, max_error).into(),
+                None => UddSketchInternal::new(size, max_error).into(),
                 Some(state) => state,
             };
             state.add_value(value);
@@ -51,7 +63,7 @@ pub fn uddsketch_trans_inner(
     }
 }
 
-const PERCENTILE_AGG_DEFAULT_SIZE: u32 = 200;
+const PERCENTILE_AGG_DEFAULT_SIZE: NonZeroU32 = NonZeroU32::new(200).unwrap();
 const PERCENTILE_AGG_DEFAULT_ERROR: f64 = 0.001;
 
 // transition function for the simpler percentile_agg aggregate, which doesn't
@@ -135,9 +147,9 @@ impl From<&UddSketchInternal> for SerializedUddSketch {
         let buckets = compress_buckets(sketch.bucket_iter());
         SerializedUddSketch {
             alpha: sketch.max_error(),
-            max_buckets: sketch.max_allowed_buckets() as u32,
+            max_buckets: sketch.max_allowed_buckets(),
             num_buckets: sketch.current_buckets_count() as u32,
-            compactions: sketch.times_compacted(),
+            compactions: sketch.times_compacted() as u32,
             count: sketch.count(),
             sum: sketch.sum(),
             buckets,
@@ -149,9 +161,10 @@ impl From<SerializedUddSketch> for UddSketchInternal {
     fn from(sketch: SerializedUddSketch) -> Self {
         UddSketchInternal::new_from_data(
             &UDDSketchMetadata {
-                max_buckets: sketch.max_buckets,
+                max_buckets: NonZeroU32::new(sketch.max_buckets)
+                    .expect(ERROR_UDDSKETCH_ZERO_BUCKETS),
                 current_error: sketch.alpha,
-                compactions: sketch.compactions,
+                compactions: u8::try_from(sketch.compactions).expect(ERROR_COMPACTIONS_LIMITED),
                 values: sketch.count,
                 sum: sketch.sum,
             },
@@ -308,9 +321,9 @@ impl<'input> UddSketch<'input> {
 
     fn metadata(&self) -> UDDSketchMetadata {
         UDDSketchMetadata {
-            max_buckets: self.max_buckets,
+            max_buckets: NonZeroU32::new(self.max_buckets).expect(ERROR_UDDSKETCH_ZERO_BUCKETS),
             current_error: self.alpha,
-            compactions: self.compactions as u32,
+            compactions: u8::try_from(self.compactions).expect(ERROR_COMPACTIONS_LIMITED),
             values: self.count,
             sum: self.sum,
         }
@@ -1262,11 +1275,41 @@ mod tests {
     fn uddsketch_byte_io_test() {
         unsafe {
             use std::ptr;
-            let state = uddsketch_trans_inner(None, 100, 0.005, Some(14.0), ptr::null_mut());
-            let state = uddsketch_trans_inner(state, 100, 0.005, Some(18.0), ptr::null_mut());
-            let state = uddsketch_trans_inner(state, 100, 0.005, Some(22.7), ptr::null_mut());
-            let state = uddsketch_trans_inner(state, 100, 0.005, Some(39.42), ptr::null_mut());
-            let state = uddsketch_trans_inner(state, 100, 0.005, Some(-43.0), ptr::null_mut());
+            let state = uddsketch_trans_inner(
+                None,
+                NonZeroU32::new(100).unwrap(),
+                0.005,
+                Some(14.0),
+                ptr::null_mut(),
+            );
+            let state = uddsketch_trans_inner(
+                state,
+                NonZeroU32::new(100).unwrap(),
+                0.005,
+                Some(18.0),
+                ptr::null_mut(),
+            );
+            let state = uddsketch_trans_inner(
+                state,
+                NonZeroU32::new(100).unwrap(),
+                0.005,
+                Some(22.7),
+                ptr::null_mut(),
+            );
+            let state = uddsketch_trans_inner(
+                state,
+                NonZeroU32::new(100).unwrap(),
+                0.005,
+                Some(39.42),
+                ptr::null_mut(),
+            );
+            let state = uddsketch_trans_inner(
+                state,
+                NonZeroU32::new(100).unwrap(),
+                0.005,
+                Some(-43.0),
+                ptr::null_mut(),
+            );
 
             let control = state.unwrap();
             let buffer = uddsketch_serialize(Inner::from(control.clone()).internal().unwrap());