subms_hyperloglog/
codec.rs

1//! Canonical wire format. A sketch is worth ~1% of the bytes the raw ids
2//! would cost, which only pays off if it can leave the process: checkpoint to
3//! disk, ship a per-shard partial to a collector, cache a per-key sketch in
4//! Redis. That needs a format both ports agree on byte for byte, and both
5//! ports emit exactly the bytes below.
6//!
7//! ```text
8//! 0..4   magic  "SHLL"
9//! 4      format version (1)
10//! 5      encoding: 0 dense, 1 sparse
11//! 6      precision p
12//! 7      reserved, zero
13//! dense  8..8+m           m register bytes, one per register
14//! sparse 8..12            u32 BE promotion threshold
15//!        12..16           u32 BE entry count n
16//!        16..16+5n        n * (u32 BE register index, u8 rho)
17//! ```
18//!
19//! Multi-byte fields are big-endian, so a hex dump reads left to right and a
20//! Java `DataOutputStream` needs no byte-order argument.
21//!
22//! This is the recipe's own format. It is not Redis's `PFADD` string and it is
23//! not a DataSketches `HllSketch` image; neither will read these bytes.
24
25use crate::{HllError, HyperLogLog, MAX_PRECISION, MIN_PRECISION};
26
27/// Leading bytes of every buffer this codec writes.
28pub const MAGIC: [u8; 4] = *b"SHLL";
29/// Format version. Bumped only on a breaking layout change.
30pub const FORMAT_VERSION: u8 = 1;
31
32pub(crate) const ENC_DENSE: u8 = 0;
33#[cfg(feature = "sparse")]
34pub(crate) const ENC_SPARSE: u8 = 1;
35pub(crate) const HEADER_LEN: usize = 8;
36
37pub(crate) fn write_header(out: &mut Vec<u8>, encoding: u8, p: u32) {
38    out.extend_from_slice(&MAGIC);
39    out.push(FORMAT_VERSION);
40    out.push(encoding);
41    out.push(p as u8);
42    out.push(0);
43}
44
45/// Validates magic, version and precision, returning `(encoding, p)`.
46pub(crate) fn read_header(bytes: &[u8]) -> Result<(u8, u32), HllError> {
47    if bytes.len() < HEADER_LEN {
48        return Err(HllError::Truncated {
49            expected: HEADER_LEN,
50            actual: bytes.len(),
51        });
52    }
53    if bytes[..4] != MAGIC {
54        return Err(HllError::BadMagic);
55    }
56    if bytes[4] != FORMAT_VERSION {
57        return Err(HllError::UnsupportedVersion(bytes[4]));
58    }
59    let p = u32::from(bytes[6]);
60    if !(MIN_PRECISION..=MAX_PRECISION).contains(&p) {
61        return Err(HllError::InvalidPrecision(p));
62    }
63    Ok((bytes[5], p))
64}
65
66#[cfg(feature = "sparse")]
67pub(crate) fn read_u32(bytes: &[u8], at: usize) -> u32 {
68    u32::from_be_bytes([bytes[at], bytes[at + 1], bytes[at + 2], bytes[at + 3]])
69}
70
71impl HyperLogLog {
72    /// Serialise to the canonical dense form: an 8-byte header then the raw
73    /// register array. Length is always `8 + 2^p`, so a reader can size the
74    /// allocation from the header alone.
75    pub fn to_bytes(&self) -> Vec<u8> {
76        let mut out = Vec::with_capacity(HEADER_LEN + self.registers.len());
77        write_header(&mut out, ENC_DENSE, self.precision());
78        out.extend_from_slice(&self.registers);
79        out
80    }
81
82    /// Parse a dense buffer. A sparse buffer is rejected with
83    /// `UnsupportedEncoding` rather than silently densified - use
84    /// `SparseHyperLogLog::from_bytes`, which reads both.
85    pub fn from_bytes(bytes: &[u8]) -> Result<Self, HllError> {
86        let (encoding, p) = read_header(bytes)?;
87        if encoding != ENC_DENSE {
88            return Err(HllError::UnsupportedEncoding(encoding));
89        }
90        let m = 1usize << p;
91        let expected = HEADER_LEN + m;
92        if bytes.len() < expected {
93            return Err(HllError::Truncated {
94                expected,
95                actual: bytes.len(),
96            });
97        }
98        let mut hll = HyperLogLog::new(p);
99        hll.registers
100            .copy_from_slice(&bytes[HEADER_LEN..HEADER_LEN + m]);
101        Ok(hll)
102    }
103}
104
105#[cfg(test)]
106#[path = "codec_tests.rs"]
107mod tests;
subms_hyperloglog/codec.rs

subms_hyperloglog/
codec.rs
