feat: Expose Sketch Memory Footprint

datasketches/Cargo.toml

@@ -35,7 +35,7 @@ all-features = true
 rustdoc-args = ["--cfg", "docsrs"]
 
 [features]
-default = []
+default = ["hll", "cpc", "bloom", "countmin", "frequencies", "tdigest", "theta"]
 
 # Each sketch has its own feature, so that users can opt in to only the sketches they need.
 bloom = []

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> {

datasketches/src/theta/hash_table.rs

@@ -381,6 +381,11 @@ impl ThetaHashTable {
     fn get_stride(key: u64, lg_size: u8) -> usize {
         (2 * ((key >> (lg_size)) & STRIDE_MASK) + 1) as usize
     }
+
+    /// Returns the estimated size of the heap allocations in bytes
+    pub fn estimated_size(&self) -> usize {
+        self.entries.capacity() * std::mem::size_of::<u64>()
+    }
 }
 
 /// Compute initial lg_size for hash table based on target lg_size, minimum lg_size, and resize

datasketches/src/theta/sketch.rs

@@ -313,6 +313,11 @@ impl ThetaSketch {
         )
         .expect("theta should always be valid")
     }
+
+    /// Returns the estimated size of the sketch in bytes
+    pub fn estimated_size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.table.estimated_size()
+    }
 }
 
 impl ThetaSketchView for ThetaSketch {
@@ -888,6 +893,11 @@ impl CompactThetaSketch {
             empty,
         })
     }
+
+    /// Returns the estimated size of the sketch in bytes
+    pub fn estimated_size(&self) -> usize {
+        std::mem::size_of::<Self>() + self.entries.capacity() * std::mem::size_of::<u64>()
+    }
 }
 
 impl ThetaSketchView for CompactThetaSketch {