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;
33pub(crate) const ENC_SPARSE: u8 = 1;
34pub(crate) const HEADER_LEN: usize = 8;
35
36pub(crate) fn write_header(out: &mut Vec<u8>, encoding: u8, p: u32) {
37    out.extend_from_slice(&MAGIC);
38    out.push(FORMAT_VERSION);
39    out.push(encoding);
40    out.push(p as u8);
41    out.push(0);
42}
43
44/// Validates magic, version and precision, returning `(encoding, p)`.
45pub(crate) fn read_header(bytes: &[u8]) -> Result<(u8, u32), HllError> {
46    if bytes.len() < HEADER_LEN {
47        return Err(HllError::Truncated {
48            expected: HEADER_LEN,
49            actual: bytes.len(),
50        });
51    }
52    if bytes[..4] != MAGIC {
53        return Err(HllError::BadMagic);
54    }
55    if bytes[4] != FORMAT_VERSION {
56        return Err(HllError::UnsupportedVersion(bytes[4]));
57    }
58    let p = u32::from(bytes[6]);
59    if !(MIN_PRECISION..=MAX_PRECISION).contains(&p) {
60        return Err(HllError::InvalidPrecision(p));
61    }
62    Ok((bytes[5], p))
63}
64
65pub(crate) fn read_u32(bytes: &[u8], at: usize) -> u32 {
66    u32::from_be_bytes([bytes[at], bytes[at + 1], bytes[at + 2], bytes[at + 3]])
67}
68
69impl HyperLogLog {
70    /// Serialise to the canonical dense form: an 8-byte header then the raw
71    /// register array. Length is always `8 + 2^p`, so a reader can size the
72    /// allocation from the header alone.
73    pub fn to_bytes(&self) -> Vec<u8> {
74        let mut out = Vec::with_capacity(HEADER_LEN + self.registers.len());
75        write_header(&mut out, ENC_DENSE, self.precision());
76        out.extend_from_slice(&self.registers);
77        out
78    }
79
80    /// Parse a dense buffer. A sparse buffer is rejected with
81    /// `UnsupportedEncoding` rather than silently densified - use
82    /// `SparseHyperLogLog::from_bytes`, which reads both.
83    pub fn from_bytes(bytes: &[u8]) -> Result<Self, HllError> {
84        let (encoding, p) = read_header(bytes)?;
85        if encoding != ENC_DENSE {
86            return Err(HllError::UnsupportedEncoding(encoding));
87        }
88        let m = 1usize << p;
89        let expected = HEADER_LEN + m;
90        if bytes.len() < expected {
91            return Err(HllError::Truncated {
92                expected,
93                actual: bytes.len(),
94            });
95        }
96        let mut hll = HyperLogLog::new(p);
97        hll.registers
98            .copy_from_slice(&bytes[HEADER_LEN..HEADER_LEN + m]);
99        Ok(hll)
100    }
101}
102
103#[cfg(test)]
104#[path = "codec_tests.rs"]
105mod tests;
subms_hyperloglog/codec.rs

subms_hyperloglog/
codec.rs
