feat: add tuple sketch implementation (#152)

Co-authored-by: tison <wander4096@gmail.com>

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/error.rs

@@ -89,8 +89,7 @@ impl Error {
     }
 }
 
-// pub(crate) convenient constructors
-#[allow(dead_code)] // some constructors are only used for certain sketches
+#[allow(dead_code)] // some convenient constructors are only used for certain sketches
 impl Error {
     pub(crate) fn invalid_argument(msg: impl Into<String>) -> Self {
         Self::new(ErrorKind::InvalidArgument, msg)

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,11 @@ pub mod hll;
 pub mod tdigest;
 #[cfg(feature = "theta")]
 pub mod theta;
-#[cfg(feature = "theta")]
+#[cfg(any(feature = "theta", feature = "tuple"))]
+#[allow(dead_code)] // some utilities are only used for certain sketches
 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/theta/union.rs

@@ -60,17 +60,17 @@ impl ThetaUnion {
 
     /// Return this union as a compact sketch.
     pub fn to_sketch(&self, ordered: bool) -> CompactThetaSketch {
-        let result = self.raw.result(ordered);
+        let parts = self.raw.to_compact_parts(ordered);
         CompactThetaSketch::from_parts(
-            result
+            parts
                 .entries
                 .into_iter()
                 .map(|entry| entry.hash())
                 .collect(),
-            result.theta,
-            result.seed_hash,
-            result.ordered,
-            result.empty,
+            parts.theta,
+            parts.seed_hash,
+            parts.ordered,
+            parts.empty,
         )
     }