feat: add tuple sketch implementation

Author: ariesdevil

datasketches/Cargo.toml

@@ -45,6 +45,7 @@ frequencies = []
 hll = []
 tdigest = []
 theta = []
+tuple = []
 
 [dev-dependencies]
 googletest = { workspace = true }

datasketches/src/codec/decode.rs

@@ -38,6 +38,16 @@ impl SketchSlice<'_> {
         self.slice.set_position(pos + n);
     }
 
+    /// Returns the not-yet-read portion of the underlying slice.
+    ///
+    /// Useful for handing the remaining bytes to a variable-length decoder that reports how many
+    /// bytes it consumed; pair it with [`advance`](Self::advance).
+    pub fn remaining(&self) -> &[u8] {
+        let buf = self.slice.get_ref();
+        let pos = (self.slice.position() as usize).min(buf.len());
+        &buf[pos..]
+    }
+
     /// Reads exactly `buf.len()` bytes from the slice into `buf`.
     pub fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
         self.slice.read_exact(buf)

datasketches/src/codec/family.rs

@@ -53,6 +53,15 @@ impl Family {
         max_pre_longs: 1,
     };
 
+    /// Tuple Sketch for cardinality estimation with per-key summaries.
+    #[cfg(feature = "tuple")]
+    pub const TUPLE: Family = Family {
+        id: 9,
+        name: "TUPLE",
+        min_pre_longs: 1,
+        max_pre_longs: 3,
+    };
+
     /// The Frequency family of sketches.
     #[cfg(feature = "frequencies")]
     pub const FREQUENCY: Family = Family {

datasketches/src/codec/mod.rs

@@ -29,7 +29,8 @@ pub use self::encode::SketchBytes;
     feature = "frequencies",
     feature = "hll",
     feature = "tdigest",
-    feature = "theta"
+    feature = "theta",
+    feature = "tuple"
 ))]
 #[allow(dead_code)] // some utilities are only used for certain sketches
 pub(crate) mod assert;
@@ -41,6 +42,7 @@ pub(crate) mod assert;
     feature = "frequencies",
     feature = "hll",
     feature = "tdigest",
-    feature = "theta"
+    feature = "theta",
+    feature = "tuple"
 ))]
 pub(crate) mod family;

datasketches/src/hash/mod.rs

@@ -20,15 +20,17 @@
     feature = "cpc",
     feature = "frequencies",
     feature = "hll",
-    feature = "theta"
+    feature = "theta",
+    feature = "tuple"
 ))]
 mod murmurhash;
 #[cfg(any(
     feature = "countmin",
     feature = "cpc",
     feature = "frequencies",
     feature = "hll",
-    feature = "theta"
+    feature = "theta",
+    feature = "tuple"
 ))]
 pub(crate) use self::murmurhash::MurmurHash3X64128;
 
@@ -56,7 +58,8 @@ pub(crate) use self::xxhash::XxHash64;
     feature = "cpc",
     feature = "frequencies",
     feature = "hll",
-    feature = "theta"
+    feature = "theta",
+    feature = "tuple"
 ))]
 pub(crate) const DEFAULT_UPDATE_SEED: u64 = 9001;
 
@@ -68,7 +71,12 @@ pub(crate) const DEFAULT_UPDATE_SEED: u64 = 9001;
 /// # Panics
 ///
 /// Panics if the computed seed hash is zero.
-#[cfg(any(feature = "countmin", feature = "cpc", feature = "theta"))]
+#[cfg(any(
+    feature = "countmin",
+    feature = "cpc",
+    feature = "theta",
+    feature = "tuple"
+))]
 pub(crate) fn compute_seed_hash(seed: u64) -> u16 {
     use std::hash::Hasher;
 
@@ -91,7 +99,8 @@ pub(crate) fn compute_seed_hash(seed: u64) -> u16 {
     feature = "cpc",
     feature = "frequencies",
     feature = "hll",
-    feature = "theta"
+    feature = "theta",
+    feature = "tuple"
 ))]
 fn read_u64_le(bytes: &[u8]) -> u64 {
     let mut buf = [0u8; 8];

datasketches/src/lib.rs

@@ -45,8 +45,10 @@ pub mod hll;
 pub mod tdigest;
 #[cfg(feature = "theta")]
 pub mod theta;
-#[cfg(feature = "theta")]
+#[cfg(any(feature = "theta", feature = "tuple"))]
 pub mod thetacommon;
+#[cfg(feature = "tuple")]
+pub mod tuple;
 
 // common modules
 pub mod codec;

datasketches/src/theta/serialization.rs

@@ -23,8 +23,3 @@ pub(super) const COMPRESSED_SERIAL_VERSION: u8 = 4;
 pub(super) const V2_PREAMBLE_EMPTY: u8 = 1;
 pub(super) const V2_PREAMBLE_PRECISE: u8 = 2;
 pub(super) const V2_PREAMBLE_ESTIMATE: u8 = 3;
-
-pub(super) const FLAGS_IS_READ_ONLY: u8 = 1 << 1;
-pub(super) const FLAGS_IS_EMPTY: u8 = 1 << 2;
-pub(super) const FLAGS_IS_COMPACT: u8 = 1 << 3;
-pub(super) const FLAGS_IS_ORDERED: u8 = 1 << 4;

datasketches/src/theta/sketch.rs

@@ -46,6 +46,10 @@ use crate::theta::serialization::V2_PREAMBLE_PRECISE;
 use crate::thetacommon::RawThetaSketchView;
 use crate::thetacommon::binomial_bounds;
 use crate::thetacommon::constants::DEFAULT_LG_K;
+use crate::thetacommon::constants::FLAGS_IS_COMPACT;
+use crate::thetacommon::constants::FLAGS_IS_EMPTY;
+use crate::thetacommon::constants::FLAGS_IS_ORDERED;
+use crate::thetacommon::constants::FLAGS_IS_READ_ONLY;
 use crate::thetacommon::constants::MAX_LG_K;
 use crate::thetacommon::constants::MAX_THETA;
 use crate::thetacommon::constants::MIN_LG_K;
@@ -220,25 +224,18 @@ impl ThetaSketch {
     /// assert_eq!(compact.num_retained(), 1);
     /// ```
     pub fn compact(&self, ordered: bool) -> CompactThetaSketch {
-        let mut entries: Vec<u64> = self.iter().map(|entry| entry.hash()).collect();
-
-        let empty = self.is_empty();
-        let theta = if empty {
-            // Match Java's correctThetaOnCompact() behavior for never-updated sketches
-            // initialized with p < 1.0.
-            MAX_THETA
-        } else {
-            self.table.theta()
-        };
-        let is_single = entries.len() == 1 && theta == MAX_THETA;
-        // Empty or Single-item sketches are always ordered (Java compatibility)
-        let ordered = ordered || empty || is_single;
-
-        if ordered && entries.len() > 1 {
-            entries.sort_unstable();
-        }
-
-        CompactThetaSketch::from_parts(entries, theta, self.table.seed_hash(), ordered, empty)
+        let parts = self.table.to_compact_parts(ordered);
+        CompactThetaSketch::from_parts(
+            parts
+                .entries
+                .into_iter()
+                .map(|entry| entry.hash())
+                .collect(),
+            parts.theta,
+            parts.seed_hash,
+            parts.ordered,
+            parts.empty,
+        )
     }
 
     /// Returns the approximate lower error bound given the specified number of Standard Deviations.
@@ -467,13 +464,13 @@ impl CompactThetaSketch {
         bytes.write_u16_be(0); // unused for compact
 
         let mut flags = 0u8;
-        flags |= serialization::FLAGS_IS_READ_ONLY;
-        flags |= serialization::FLAGS_IS_COMPACT;
+        flags |= FLAGS_IS_READ_ONLY;
+        flags |= FLAGS_IS_COMPACT;
         if self.is_empty() {
-            flags |= serialization::FLAGS_IS_EMPTY;
+            flags |= FLAGS_IS_EMPTY;
         }
         if self.is_ordered() {
-            flags |= serialization::FLAGS_IS_ORDERED;
+            flags |= FLAGS_IS_ORDERED;
         }
         bytes.write_u8(flags);
 
@@ -510,9 +507,9 @@ impl CompactThetaSketch {
         bytes.write_u8(num_entries_bytes);
 
         let mut flags = 0u8;
-        flags |= serialization::FLAGS_IS_READ_ONLY;
-        flags |= serialization::FLAGS_IS_COMPACT;
-        flags |= serialization::FLAGS_IS_ORDERED;
+        flags |= FLAGS_IS_READ_ONLY;
+        flags |= FLAGS_IS_COMPACT;
+        flags |= FLAGS_IS_ORDERED;
         bytes.write_u8(flags);
 
         bytes.write_u16_le(self.seed_hash);
@@ -748,7 +745,7 @@ impl CompactThetaSketch {
             .read_u16_le()
             .map_err(insufficient_data("seed_hash"))?;
 
-        let empty = (flags & serialization::FLAGS_IS_EMPTY) != 0;
+        let empty = (flags & FLAGS_IS_EMPTY) != 0;
         let mut theta = MAX_THETA;
         let num_entries;
         let mut entries = vec![];
@@ -776,7 +773,7 @@ impl CompactThetaSketch {
             }
             entries = Self::read_entries(&mut cursor, num_entries as usize, theta)?;
         }
-        let ordered = (flags & serialization::FLAGS_IS_ORDERED) != 0;
+        let ordered = (flags & FLAGS_IS_ORDERED) != 0;
         Ok(Self {
             entries,
             theta,
@@ -797,7 +794,7 @@ impl CompactThetaSketch {
         let seed_hash = cursor
             .read_u16_le()
             .map_err(insufficient_data("seed_hash"))?;
-        let empty = (flags & serialization::FLAGS_IS_EMPTY) != 0;
+        let empty = (flags & FLAGS_IS_EMPTY) != 0;
         if !empty {
             let expected_seed_hash = compute_seed_hash(expected_seed);
             if seed_hash != expected_seed_hash {
@@ -861,7 +858,7 @@ impl CompactThetaSketch {
             }
         }
 
-        let ordered = (flags & serialization::FLAGS_IS_ORDERED) != 0;
+        let ordered = (flags & FLAGS_IS_ORDERED) != 0;
 
         Ok(Self {
             entries,

datasketches/src/thetacommon/constants.rs

@@ -34,3 +34,13 @@ pub const HASH_TABLE_REBUILD_THRESHOLD: f64 = 15.0 / 16.0;
 
 pub const STRIDE_HASH_BITS: u8 = 7;
 pub const STRIDE_MASK: u64 = (1 << STRIDE_HASH_BITS) - 1;
+
+// Flag bits of the flags byte in the Theta-family wire format, shared by Theta and Tuple sketches.
+/// Flags byte bit: the sketch is read-only.
+pub const FLAGS_IS_READ_ONLY: u8 = 1 << 1;
+/// Flags byte bit: the sketch is logically empty.
+pub const FLAGS_IS_EMPTY: u8 = 1 << 2;
+/// Flags byte bit: the sketch is in compact form.
+pub const FLAGS_IS_COMPACT: u8 = 1 << 3;
+/// Flags byte bit: retained entries are ordered by ascending hash.
+pub const FLAGS_IS_ORDERED: u8 = 1 << 4;

datasketches/src/thetacommon/hash_table.rs

@@ -20,6 +20,7 @@ use std::hash::Hash;
 use crate::common::ResizeFactor;
 use crate::hash::MurmurHash3X64128;
 use crate::hash::compute_seed_hash;
+use crate::thetacommon::RawCompactParts;
 use crate::thetacommon::RawHashTableEntry;
 use crate::thetacommon::constants::HASH_TABLE_REBUILD_THRESHOLD;
 use crate::thetacommon::constants::HASH_TABLE_RESIZE_THRESHOLD;
@@ -168,6 +169,7 @@ where
     }
 
     /// Returns a reference to the entry stored for `hash`, or `None` if the hash is not retained.
+    #[allow(dead_code)] // used by the set operations; tuple set operations arrive in follow-ups
     pub fn get_entry(&self, hash: u64) -> Option<&E> {
         if hash == 0 {
             return None;
@@ -234,6 +236,33 @@ where
         self.entries.iter().filter_map(Option::as_ref)
     }
 
+    /// Returns the retained entries and theta as raw compact-sketch parts.
+    ///
+    /// An empty table reports `MAX_THETA` rather than its current theta, matching Java's
+    /// `correctThetaOnCompact()` behavior for never-updated sketches initialized with p < 1.0.
+    /// Empty and single-entry exact-mode results are always marked ordered (Java/C++
+    /// compatibility).
+    pub fn to_compact_parts(&self, ordered: bool) -> RawCompactParts<E>
+    where
+        E: Clone,
+    {
+        let mut entries: Vec<E> = self.iter_entries().cloned().collect();
+        let empty = self.is_empty();
+        let theta = if empty { MAX_THETA } else { self.theta() };
+        let is_single = entries.len() == 1 && theta == MAX_THETA;
+        let ordered = ordered || empty || is_single;
+        if ordered && entries.len() > 1 {
+            entries.sort_unstable_by_key(RawHashTableEntry::hash);
+        }
+        RawCompactParts {
+            entries,
+            theta,
+            seed_hash: self.seed_hash(),
+            ordered,
+            empty,
+        }
+    }
+
     /// Return number of all entries.
     #[cfg(test)]
     pub fn num_entries(&self) -> usize {
@@ -257,16 +286,19 @@ where
     }
 
     /// Set empty flag.
+    #[allow(dead_code)] // used by the set operations; tuple set operations arrive in follow-ups
     pub fn set_empty(&mut self, is_empty: bool) {
         self.is_empty = is_empty;
     }
 
     /// Get the hash seed used by this table.
+    #[allow(dead_code)] // used by the set operations; tuple set operations arrive in follow-ups
     pub fn hash_seed(&self) -> u64 {
         self.hash_seed
     }
 
     /// Sets theta value.
+    #[allow(dead_code)] // used by the set operations and tests; tuple set operations arrive in follow-ups
     pub fn set_theta(&mut self, theta: u64) {
         assert!(
             (1..=MAX_THETA).contains(&theta),
@@ -276,6 +308,7 @@ where
     }
 
     /// Returns minimal lg_size where rebuild-capacity can hold `count`.
+    #[allow(dead_code)] // used by the set operations; tuple set operations arrive in follow-ups
     pub fn lg_size_from_count_for_rebuild(count: usize, load_factor: f64) -> u8 {
         let log2 = |n: usize| {
             if n == 0 { 0_u8 } else { n.ilog2() as u8 }