feat: Expose Sketch Memory Footprint (#136)

Author: tabac

datasketches/src/bloom/sketch.rs

@@ -561,6 +561,11 @@ impl BloomFilter {
             self.num_bits_set += 1;
         }
     }
+
+    /// Returns the estimated size of the filter in bytes
+    pub fn estimated_size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.bit_array.len() * std::mem::size_of::<u64>()
+    }
 }
 
 #[cfg(test)]

datasketches/src/cpc/pair_table.rs

@@ -246,4 +246,9 @@ impl PairTable {
             }
         }
     }
+
+    /// Returns the estimated size of the heap allocations in bytes
+    pub fn estimated_size(&self) -> usize {
+        self.slots.capacity() * std::mem::size_of::<u32>()
+    }
 }

datasketches/src/cpc/sketch.rs

@@ -450,6 +450,18 @@ impl CpcSketch {
 
         matrix
     }
+
+    /// Returns the estimated size of the sketch in bytes
+    pub fn estimated_size(&self) -> usize {
+        let heap_size = self.sliding_window.capacity()
+            + self
+                .surprising_value_table
+                .as_ref()
+                .map(|t| t.estimated_size())
+                .unwrap_or(0);
+
+        std::mem::size_of::<Self>() + heap_size
+    }
 }
 
 impl CpcSketch {

datasketches/src/hll/array4.rs

@@ -425,6 +425,16 @@ impl Array4 {
 
         bytes.into_bytes()
     }
+
+    /// Returns the estimated size of the heap allocations in bytes
+    pub fn estimated_size(&self) -> usize {
+        self.bytes.len()
+            + self
+                .aux_map
+                .as_ref()
+                .map(|a| a.estimated_size())
+                .unwrap_or(0)
+    }
 }
 
 #[cfg(test)]

datasketches/src/hll/array6.rs

@@ -272,6 +272,11 @@ impl Array6 {
 
         bytes.into_bytes()
     }
+
+    /// Returns the estimated size of the heap allocations in bytes
+    pub fn estimated_size(&self) -> usize {
+        self.bytes.len()
+    }
 }
 
 /// Calculate number of bytes needed for k slots with 6 bits each

datasketches/src/hll/array8.rs

@@ -344,6 +344,11 @@ impl Array8 {
 
         bytes.into_bytes()
     }
+
+    /// Returns the estimated size of the heap allocations in bytes
+    pub fn estimated_size(&self) -> usize {
+        self.bytes.len()
+    }
 }
 
 #[cfg(test)]

datasketches/src/hll/aux_map.rs

@@ -226,6 +226,11 @@ impl AuxMap {
             }
         })
     }
+
+    /// Returns the estimated size of the heap allocations in bytes
+    pub fn estimated_size(&self) -> usize {
+        self.entries.len() * std::mem::size_of::<Coupon>()
+    }
 }
 
 /// Iterator over AuxMap entries

datasketches/src/hll/container.rs

@@ -135,4 +135,9 @@ impl Container {
     pub fn iter(&self) -> impl Iterator<Item = Coupon> + '_ {
         self.coupons.iter().filter(|&&c| !c.is_empty()).copied()
     }
+
+    /// Returns the estimated size of the heap allocations in bytes
+    pub fn estimated_size(&self) -> usize {
+        self.coupons.len() * std::mem::size_of::<Coupon>()
+    }
 }

datasketches/src/hll/sketch.rs

@@ -423,6 +423,19 @@ impl HllSketch {
             Mode::Array8(arr) => arr.serialize(self.lg_config_k),
         }
     }
+
+    /// Returns the estimated size of the sketch in bytes
+    pub fn estimated_size(&self) -> usize {
+        let heap_size = match &self.mode {
+            Mode::List { list, .. } => list.container().estimated_size(),
+            Mode::Set { set, .. } => set.container().estimated_size(),
+            Mode::Array4(arr) => arr.estimated_size(),
+            Mode::Array6(arr) => arr.estimated_size(),
+            Mode::Array8(arr) => arr.estimated_size(),
+        };
+
+        std::mem::size_of::<Self>() + heap_size
+    }
 }
 
 fn promote_container_to_set(container: &Container, hll_type: HllType) -> Mode {

datasketches/src/tdigest/sketch.rs

@@ -793,6 +793,13 @@ impl TDigestMut {
         self.reverse_merge = !self.reverse_merge;
         self.buffer.clear();
     }
+
+    /// Returns the estimated size of the sketch in bytes
+    pub fn estimated_size(&self) -> usize {
+        std::mem::size_of::<Self>()
+            + self.centroids.capacity() * std::mem::size_of::<Centroid>()
+            + self.buffer.capacity() * std::mem::size_of::<f64>()
+    }
 }
 
 /// Immutable (frozen) T-Digest sketch for estimating quantiles and ranks.
@@ -1001,6 +1008,11 @@ impl TDigest {
             vec![],
         )
     }
+
+    /// Returns the estimated size of the sketch in bytes
+    pub fn estimated_size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.centroids.capacity() * std::mem::size_of::<Centroid>()
+    }
 }
 
 struct TDigestView<'a> {