From baada44cf28f4699fe71b577422f1bb23623750f Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Sun, 19 Jul 2026 04:05:08 -0400 Subject: [PATCH 1/4] feat(kll): ASAPv1 wire payload for KLL and KLLDynamic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add self-describing ASAPv1 serialization for the two KLL quantile sketches, mirroring the HLL/CMS pattern: compact `KLL` (kind_id `0x06 0x00`) and `KLLDynamic` (`0x06 0x01`) share one payload `[levels, items, coin]`, with `serialize_to_bytes` / `deserialize_from_bytes` living in per-sketch `wire` submodules. KLL never hashes (it orders raw values with `total_cmp`), so its metadata carries no hash-spec group — only the structural params `{metadata_version, k, m, item_type}` (decision Q-KLL). `item_type` (`f64`/`i64`) is a metadata param, not a separate kind_id (mirrors CMS `counter_type`). Retained samples use the top-most-level-first layout that matches sketchlib-go's `KLLState`, so the compact KLL's leftward-grown L0 is reversed to input order on the wire. Decode fails closed: kind_id + metadata + item_type checks, level- layout consistency, coin `remaining_bits <= 64`, and a `2 <= m <= k <= MAX_CACHEABLE_K` bound so crafted `k`/`m` cannot drive `compute_max_capacity` into a huge allocation. The nested serde `Serialize`/`Deserialize` on `KLL` (used by `HydraCounter`) is kept. Add golden byte-vectors `kll_f64_k200` / `kll_i64_k200` (1..=50, seed 42, no compaction — deterministic) and cross-language tests; document the payload in docs/asapv1_wire_format.md §3.3. KLLDynamic golden is deferred (it has no seeded constructor); the portable `KllSketch` double-wrap and the Go mirror are follow-ups. Co-Authored-By: Claude Opus 4.8 (1M context) --- asapv1_golden/README.md | 11 + asapv1_golden/kll_f64_k200.hex | 1 + asapv1_golden/kll_i64_k200.hex | 1 + docs/asapv1_wire_format.md | 49 +- src/message_pack_format/native/kll.rs | 4 +- src/message_pack_format/native/kll_dynamic.rs | 3 +- src/sketches/kll.rs | 52 +- src/sketches/kll/wire.rs | 476 ++++++++++++++++++ src/sketches/kll_dynamic.rs | 23 +- src/sketches/kll_dynamic/wire.rs | 151 ++++++ tests/asapv1_golden.rs | 50 +- 11 files changed, 767 insertions(+), 54 deletions(-) create mode 100644 asapv1_golden/kll_f64_k200.hex create mode 100644 asapv1_golden/kll_i64_k200.hex create mode 100644 src/sketches/kll/wire.rs create mode 100644 src/sketches/kll_dynamic/wire.rs diff --git a/asapv1_golden/README.md b/asapv1_golden/README.md index cc39625..e6084b2 100644 --- a/asapv1_golden/README.md +++ b/asapv1_golden/README.md @@ -29,11 +29,22 @@ golden tests the **wire encoding**, isolated from the hash functions. | `hll_hip_p12` | HLL HIP, P12 | `01 03` | same registers + `hip_kxq0=1.5, hip_kxq1=2.5, hip_est=3.0` | | `cms_i64_regular_2x3` | Count-Min i64, RegularPath | `02 00` | 2×3 row-major `[[0,1,127],[128,300,65536]]` | | `cms_f64_fast_2x3` | Count-Min f64, FastPath | `02 00` | 2×3 row-major `[[0.0,1.5,2.25],[3.75,4.125,5.0625]]` | +| `kll_f64_k200` | KLL f64, k=200 | `06 00` | integers `1..=50`, compaction seed 42 | +| `kll_i64_k200` | KLL i64, k=200 | `06 00` | integers `1..=50`, compaction seed 42 | The i64 fixture deliberately spans the msgpack integer width boundaries (positive fixint / uint8 / uint16 / uint32) to lock the "non-negative integer → uint family, minimal width" rule (`docs/asapv1_wire_format.md` §5). +The KLL fixtures are a special case of "state is fixed, not hashed": KLL never +hashes — it orders raw numeric values — so inserting `1..=50` places exactly +those retained samples. `k=200` keeps the input below the level-0 capacity, so no +compaction fires (`num_levels = 1`, one level `[1..50]`) and the state is fully +deterministic. The fixed compaction seed (42) pins the carried coin state, which +must match `sketchlib-go`'s coin for the same input. Only the compact KLL +(`06 00`) has a golden; the dynamic variant (`06 01`) shares the payload shape but +lacks a seeded constructor, so its cross-language golden is deferred. + ## Tests that consume these - Rust: `tests/asapv1_golden.rs` — builds each fixture from known state, diff --git a/asapv1_golden/kll_f64_k200.hex b/asapv1_golden/kll_f64_k200.hex new file mode 100644 index 0000000..354d3ce --- /dev/null +++ b/asapv1_golden/kll_f64_k200.hex @@ -0,0 +1 @@ +4153415076310102060000000028000001cd84b06d657461646174615f76657273696f6e01a16bccc8a16d08a96974656d5f74797065a366363493920032dc0032cb3ff0000000000000cb4000000000000000cb4008000000000000cb4010000000000000cb4014000000000000cb4018000000000000cb401c000000000000cb4020000000000000cb4022000000000000cb4024000000000000cb4026000000000000cb4028000000000000cb402a000000000000cb402c000000000000cb402e000000000000cb4030000000000000cb4031000000000000cb4032000000000000cb4033000000000000cb4034000000000000cb4035000000000000cb4036000000000000cb4037000000000000cb4038000000000000cb4039000000000000cb403a000000000000cb403b000000000000cb403c000000000000cb403d000000000000cb403e000000000000cb403f000000000000cb4040000000000000cb4040800000000000cb4041000000000000cb4041800000000000cb4042000000000000cb4042800000000000cb4043000000000000cb4043800000000000cb4044000000000000cb4044800000000000cb4045000000000000cb4045800000000000cb4046000000000000cb4046800000000000cb4047000000000000cb4047800000000000cb4048000000000000cb4048800000000000cb4049000000000000932a0000 diff --git a/asapv1_golden/kll_i64_k200.hex b/asapv1_golden/kll_i64_k200.hex new file mode 100644 index 0000000..ac77317 --- /dev/null +++ b/asapv1_golden/kll_i64_k200.hex @@ -0,0 +1 @@ +41534150763101020600000000280000003d84b06d657461646174615f76657273696f6e01a16bccc8a16d08a96974656d5f74797065a369363493920032dc00320102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132932a0000 diff --git a/docs/asapv1_wire_format.md b/docs/asapv1_wire_format.md index 880a3bc..701d661 100644 --- a/docs/asapv1_wire_format.md +++ b/docs/asapv1_wire_format.md @@ -12,7 +12,7 @@ If the doc feels long, these sections carry the key points: - [**Section 1, Envelope**](#section-1-envelope): the Layout table. - [**Section 2, Metadata**](#section-2-metadata): the fields table, the hash-spec table, the structural-params table. -- [**Section 3, Payload**](#section-3-payload): the HLL payload and the Count-Min payload. +- [**Section 3, Payload**](#section-3-payload): the HLL payload, the Count-Min payload, and the KLL payload. ## Terms @@ -26,7 +26,7 @@ If the doc feels long, these sections carry the key points: ## Status -- **Implementing (Rust).** HLL and Count-Min serialize through the shared `message_pack_format::envelope` module per this spec. +- **Implementing (Rust).** HLL, Count-Min, and KLL (compact + dynamic) serialize through the shared `message_pack_format::envelope` module per this spec. - **Self-describing.** The hash-spec metadata is derived from the hasher's `HashProfile` (read live, never hardcoded), so the bytes truthfully describe how a sketch was hashed; custom hash profiles are supported (Section 2). - **Byte-level encoding** is pinned in Section 4; the resolved decisions are summarized at the end. - **`sketchlib-go`** is aligned separately (see Cross-language contract). @@ -73,6 +73,7 @@ erDiagram PAYLOAD ||--o| HLL_PAYLOAD : kind-0x0101-0x0102 PAYLOAD ||--o| HLL_HIP_PAYLOAD : kind-0x0103 PAYLOAD ||--o| COUNTMIN_PAYLOAD : kind-0x0200 + PAYLOAD ||--o| KLL_PAYLOAD : kind-0x0600-0x0601 ENVELOPE { bytes magic u8 version @@ -88,8 +89,8 @@ erDiagram string seed_derivation string input_encoding array seed_list - u32 seed_index "hash-spec group, per-sketch" - mixed structural_params "structural group, per-sketch: precision(HLL) or rows+cols+counter_type+mode(CMS)" + u32 seed_index "hash-spec group, per-sketch; absent for non-hashing sketches (KLL)" + mixed structural_params "structural group, per-sketch: precision(HLL); rows+cols+counter_type+mode(CMS); k+m+item_type(KLL)" } PAYLOAD { msgpack_array raw_state @@ -106,6 +107,11 @@ erDiagram COUNTMIN_PAYLOAD { array counts } + KLL_PAYLOAD { + array levels + array items + array coin + } ``` --- @@ -166,7 +172,7 @@ Today `kind_id` is `[family, variant]` and names the sketch's **algorithm** (its ### kind_id registry (single source of truth, mirrored verbatim in `sketchlib-go`) The **family** bytes match `sketchlib-go`'s `wire/asapmsgpack/magic_ids.go` verbatim; `0x0a`+ are new allocations for sketches in [`apis.md`](./apis.md) that Go has not assigned yet. -Only the HLL variants and Count-Min have designed payloads today; every other row reserves its family byte with variant `0x00` and payload **TBD**. +The HLL variants, Count-Min, and both KLL variants have designed payloads today; every other row reserves its family byte with variant `0x00` and payload **TBD**. Variant sub-ids are not invented ahead of a payload design; a family that later needs several algorithms allocates its variants when it is designed (as HLL did). This registry is the master list of algorithms still to design payloads for. @@ -180,8 +186,8 @@ This registry is the master list of algorithms still to design payloads for. | `0x03 0x00` | Count-Min-with-heap (CMSHeap) | - | TBD | assigned in Go / payload not designed | | `0x04 0x00` | Count Sketch | - | TBD | assigned in Go / payload not designed | | `0x05 0x00` | DDSketch | - | TBD | assigned in Go / payload not designed | -| `0x06 0x00` | KLL | - | TBD | assigned in Go / payload not designed | -| `0x06 0x01` | KLL dynamic | - | TBD | assigned in Go / payload not designed | +| `0x06 0x00` | KLL | Compact | Section 3.3 | implemented | +| `0x06 0x01` | KLL dynamic | Dynamic | Section 3.3 | implemented | | `0x07 0x00` | Hydra-KLL | - | TBD | assigned in Go / payload not designed | | `0x08 0x00` | SetAggregator | - | TBD | assigned in Go / payload not designed | | `0x09 0x00` | DeltaResult | - | TBD | assigned in Go / payload not designed | @@ -247,10 +253,12 @@ The metadata map is **two groups** of fields, written on the wire in this order: | Group | Role | Fields | | ----- | ---- | ------ | | **Hash spec** | how keys were hashed (check mergeability + re-hash a query key) | `metadata_version`, `hash_profile_id`, `hash_algorithm`, `seed_derivation`, `input_encoding`, `seed_list`, + the seed index(es) it uses | -| **Structural params** | parameters that shape the payload | `precision` (HLL); `rows`, `cols`, `counter_type`, `mode` (Count-Min) | +| **Structural params** | parameters that shape the payload | `precision` (HLL); `rows`, `cols`, `counter_type`, `mode` (Count-Min); `k`, `m`, `item_type` (KLL) | The two tables below are the field-by-field detail of each group. +> **Non-hashing sketches omit the hash-spec group entirely** (Q-KLL). The hash spec answers "how were keys hashed"; a sketch that does not hash its inputs has no truthful answer. KLL is comparison-based — it orders raw numeric values with `total_cmp` and never invokes a hasher — so its metadata carries **only** structural params (`metadata_version`, `k`, `m`, `item_type`). This keeps the bytes honest (no meaningless `seed_list`) and is the reason KLL's metadata schema is not built from a `HashProfile` the way HLL's / Count-Min's are. + **Hash spec** | Key | Type | Required | Meaning | @@ -275,6 +283,9 @@ The two tables below are the field-by-field detail of each group. | `cols` | u32 | Count-Min | matrix width (number of columns) | | `counter_type` | string | Count-Min | `"i64"` or `"f64"`; element type of `counts` | | `mode` | string | Count-Min | `"fast"` or `"regular"`; key-to-column derivation | +| `k` | u32 | KLL | compactor capacity (accuracy parameter) | +| `m` | u32 | KLL | minimum level capacity | +| `item_type` | string | KLL | `"f64"` or `"i64"`; element type of `items` | Count-Min's matrix dimensions are **configuration** (they shape the payload, like HLL's `precision`), so per the config-to-metadata rule they live here. Count-Min's canonical structural-param order is `... matrix_seed_index, rows, cols, counter_type, mode`; this is the wire contract and Go must mirror it verbatim. @@ -361,7 +372,23 @@ Wire counter types are `i64` and `f64` only (`i32` widens to `i64`; `i128` and e A counter type other than i64/f64, or non-`Vector2D` storage, must be converted first; see Section 5, "Converting an exotic in-memory sketch". Both modes, `FastPath` and `RegularPath`, serialize directly (you'd only "convert" a mode to *change* it, which needs re-inserting the data). -### 3.3 onward: payloads not yet designed +### 3.3: KLL payload (`0x06 0x00` compact / `0x06 0x01` dynamic) + +Both KLL variants (the compact fixed-buffer `KLL` and the growable `KLLDynamic`) share **one payload shape**; they differ only by `kind_id` (like HLL's Classic vs Ertl-MLE), because their in-memory buffers differ but the serialized quantile state is the same. The accuracy params `k` / `m` and the `item_type` live in the metadata (§2), so the payload is just the retained state: + +| Pos | Field | Type | Notes | +| ----- | ------- | ------ | ------- | +| 0 | `levels` | array | level boundary indices, length `num_levels + 1`; `levels[0] == 0`, `levels[last] == len(items)` | +| 1 | `items` | array | retained samples in level order; element type = `item_type` | +| 2 | `coin` | array | compaction RNG state `[state:u64, bit_cache:u64, remaining_bits:u32]` | + +`num_levels` is `levels.len() - 1` (derived, so not stored). The `coin` is the randomized-compaction RNG; it is carried so a decoded sketch can keep compacting deterministically (a query-only consumer may ignore it). It is a nested 3-element array mirroring `sketchlib-go`'s `CoinState`. + +**Item order (cross-language contract).** `levels` / `items` use the **top-most-level-first** layout, byte-for-byte matching `sketchlib-go`'s `KLLState`: index `i` in `levels` maps to compactor level `num_levels - 1 - i`, and level 0's run is in **input order**. The compact `KLL` grows its buffer leftward and stores level 0 reverse-input, so its encoder reverses level 0 back to input order (and its decoder reverses it in); `KLLDynamic` already stores this layout natively. Within a level, order past the first compaction is not guaranteed byte-identical across the two Rust variants (or across languages), but the retained set and quantiles agree — see the caveat on `KLL::wire_items`. + +> Note: this is a distinct wire format from the proto `KLLState` (`proto/kll/kll.proto`), which is the delta-transmission / cross-language *protobuf* path and carries extra machinery (the value-offset fixed-point item encoding). The two share the top-most-first item layout but are separate formats; the ASAPv1 KLL payload above is self-contained and carries no fixed-point form. + +### 3.4 onward: payloads not yet designed The remaining `kind_id`s reserve a family byte with payload TBD (Section 1 registry has their "assigned in Go" status). Likely shape when designed: @@ -370,8 +397,7 @@ The remaining `kind_id`s reserve a family byte with payload TBD (Section 1 regis | `0x03 0x00` | Count-Min-with-heap (CMSHeap) | similar to current CMS | | `0x04 0x00` | Count Sketch | similar to current CMS | | `0x05 0x00` | DDSketch | straightforward bucket | -| `0x06 0x00` | KLL | next on the list; no CDF serialization | -| `0x07 0x00` | Hydra-KLL | same challenge as KLL | +| `0x07 0x00` | Hydra-KLL | wraps KLL payloads (§3.3); nest one per counter | | `0x08 0x00` | SetAggregator | aggregation envelope, distinct from a stand-alone sketch (Section 1 mapping notes) | | `0x09 0x00` | DeltaResult | delta-result envelope, distinct from a stand-alone sketch (Section 1 mapping notes) | | `0x0a 0x00` | Count-Sketch-with-heap (CSHeap) | similar to current CMS | @@ -483,6 +509,7 @@ Keep it through the transition, retire `portable` once goldens are in place. - **Q-SEEDS**: `seed_list` is **inlined** in v1 so the bytes self-describe the hash. Resolving seeds from `hash_profile_id` alone is a v2 space optimization. Each sketch still carries only the seed *index* it uses. - **Q-PROFILE**: the hash-spec metadata is **derived from the hasher's `HashProfile`** (`hll_metadata::` / `cms_metadata::`) and never hardcoded, so it is always truthful to the hasher. Custom hash profiles are **supported and self-describing**. Merge compatibility is hash-spec equality, so a custom-profile sketch is not mergeable with a standard one. - **Q-CMS**: Count-Min is one `kind_id` (`0x02 0x00`); counter type and mode live in the metadata, so the id stays single. +- **Q-KLL**: KLL metadata carries **no hash-spec group** — KLL is comparison-based and never hashes, so those fields have no truthful value. Its metadata is structural-only (`metadata_version`, `k`, `m`, `item_type`) and is *not* `HashProfile`-derived. The two KLL variants (compact `0x06 0x00`, dynamic `0x06 0x01`) share one payload `[levels, items, coin]` and differ only by `kind_id`. `item_type` (`"f64"`/`"i64"`) is a metadata param, not a separate `kind_id` (mirrors Q-CMS's `counter_type`). Retained samples use the top-most-level-first layout that matches `sketchlib-go`'s `KLLState`. - **Q-CMS-DIMS**: Count-Min `rows`/`cols` are **metadata** and the payload omits them. They are configuration that shapes the payload (like HLL's `precision`), so per the config-to-metadata rule they belong in the descriptor. The payload is then just `[counts]`. Canonical structural-param order: `... matrix_seed_index, rows, cols, counter_type, mode`. - **Q-VER**: no payload version field. A new incompatible encoding gets a **new `kind_id`**; retired ids are reserved forever and never recycled. - **Encoding**: metadata + payload are both msgpack; payload is a positional array. Byte-level rules in Section 4. diff --git a/src/message_pack_format/native/kll.rs b/src/message_pack_format/native/kll.rs index 7ea9b85..e56487b 100644 --- a/src/message_pack_format/native/kll.rs +++ b/src/message_pack_format/native/kll.rs @@ -4,11 +4,11 @@ use serde::{Deserialize, Serialize}; use crate::common::numerical::NumericalValue; use crate::message_pack_format::{Error, MessagePackCodec}; -use crate::sketches::kll::KLL; +use crate::sketches::kll::{KLL, KllWireItem}; impl MessagePackCodec for KLL where - T: NumericalValue + Serialize + for<'de> Deserialize<'de>, + T: NumericalValue + KllWireItem + Serialize + for<'de> Deserialize<'de>, { fn to_msgpack(&self) -> Result, Error> { Ok(self.serialize_to_bytes()?) diff --git a/src/message_pack_format/native/kll_dynamic.rs b/src/message_pack_format/native/kll_dynamic.rs index 1873c34..7818a49 100644 --- a/src/message_pack_format/native/kll_dynamic.rs +++ b/src/message_pack_format/native/kll_dynamic.rs @@ -4,11 +4,12 @@ use serde::{Deserialize, Serialize}; use crate::common::numerical::NumericalValue; use crate::message_pack_format::{Error, MessagePackCodec}; +use crate::sketches::kll::KllWireItem; use crate::sketches::kll_dynamic::KLLDynamic; impl MessagePackCodec for KLLDynamic where - T: NumericalValue + Serialize + for<'de> Deserialize<'de>, + T: NumericalValue + KllWireItem + Serialize + for<'de> Deserialize<'de>, { fn to_msgpack(&self) -> Result, Error> { Ok(self.serialize_to_bytes()?) diff --git a/src/sketches/kll.rs b/src/sketches/kll.rs index cef5f51..64637f5 100644 --- a/src/sketches/kll.rs +++ b/src/sketches/kll.rs @@ -11,14 +11,18 @@ //! - use rand::{Rng, rng}; -use rmp_serde::decode::Error as RmpDecodeError; -use rmp_serde::encode::Error as RmpEncodeError; use serde::{Deserialize, Serialize}; use crate::common::input::data_input_to_f64; use crate::common::numerical::NumericalValue; use crate::{DataInput, Vector1D}; +mod wire; +pub(crate) use wire::{ + KLL_KIND_DYNAMIC, KllCoinWire, KllPayload, KllWireItem, kll_metadata, split_and_validate_meta, + validate_kll_payload, +}; + const MAX_LEVELS: usize = 61; const CAPACITY_CACHE_LEN: usize = 20; @@ -79,6 +83,27 @@ impl Coin { self.remaining_bits -= 1; bit } + + /// The coin's raw state in `sketchlib-go::CoinState` shape: + /// `(state, bit_cache, remaining_bits)`. Used by the ASAPv1 wire payload + /// (both KLL variants) to carry the compaction RNG so a decoded sketch can + /// continue compacting deterministically. + #[inline] + pub(crate) fn to_wire(&self) -> (u64, u64, u32) { + (self.state, self.bit_cache, self.remaining_bits as u32) + } + + /// Rebuilds a coin from its raw wire state. `remaining_bits` is validated by + /// the caller (it must be `<= 64`) so decode fails closed on crafted bytes + /// rather than silently truncating into the `u8` field. + #[inline] + pub(crate) fn from_wire(state: u64, bit_cache: u64, remaining_bits: u8) -> Self { + Self { + state, + bit_cache, + remaining_bits, + } + } } /// A single (value, cumulative-quantile) pair in a CDF table. @@ -786,22 +811,13 @@ impl KLL { } // -- Serialization ------------------------------------------------------- - - /// Serializes the sketch to a MessagePack byte vector. - pub fn serialize_to_bytes(&self) -> Result, RmpEncodeError> - where - T: Serialize, - { - rmp_serde::to_vec(self) - } - - /// Deserializes a KLL sketch from a MessagePack byte slice. - pub fn deserialize_from_bytes(bytes: &[u8]) -> Result - where - T: for<'de> Deserialize<'de>, - { - rmp_serde::from_slice(bytes) - } + // + // The ASAPv1 wire methods (`serialize_to_bytes` / `deserialize_from_bytes`) + // live in the `wire` submodule, which authors the envelope + metadata + + // payload. The `serde::{Serialize, Deserialize}` impls below stay: they are + // the *nested* codec used when a `KLL` is embedded in a larger serde value + // (e.g. `HydraCounter::KLL`), which is a different concern from the + // standalone ASAPv1 envelope. fn ensure_levels_sorted(&mut self) { if self.num_levels <= 1 { diff --git a/src/sketches/kll/wire.rs b/src/sketches/kll/wire.rs new file mode 100644 index 0000000..854367f --- /dev/null +++ b/src/sketches/kll/wire.rs @@ -0,0 +1,476 @@ +//! ASAPv1 wire serialization for the KLL quantile sketches. +//! +//! Child submodule of [`crate::sketches::kll`]: it holds the shared KLL wire +//! DTOs (metadata + payload + coin), the kind_id constants for both KLL +//! variants, the [`KllWireItem`] marker trait, and the `serialize_to_bytes` / +//! `deserialize_from_bytes` impls for the compact [`KLL`]. Being a descendant +//! module, it reads the sketch's private fields (`self.items`, `self.levels`, +//! `self.co`, …) and constructs the struct directly without widening any field +//! visibility. The sibling `kll_dynamic::wire` reuses these DTOs for the dynamic +//! variant. See `docs/asapv1_wire_format.md` §3. +//! +//! ## Why KLL metadata has no hash spec +//! +//! Unlike HLL and Count-Min, KLL never hashes its inputs — it orders raw numeric +//! values with `total_cmp`. So the hash-spec metadata group (profile id, +//! algorithm, inlined `seed_list`, …) that HLL/CMS carry does not apply, and KLL +//! metadata is **structural params only** (`k`, `m`, `item_type`). See the wire +//! doc's KLL section for the recorded decision. +//! +//! ## Payload item order (cross-language contract) +//! +//! The payload's `levels` / `items` use the **top-most-level-first** layout, +//! byte-for-byte matching `sketchlib-go`'s `KLLState` (index `i` in `levels` +//! maps to compactor level `num_levels - 1 - i`; level 0's run is in input +//! order). For the compact KLL this is exactly what [`KLL::wire_levels`] / +//! [`KLL::wire_items`] emit (they reverse the leftward-grown L0 buffer back to +//! input order); decode inverts that mapping. + +use rmp_serde::{decode::Error as RmpDecodeError, encode::Error as RmpEncodeError, from_slice}; +use serde::{Deserialize, Serialize}; + +use crate::common::numerical::NumericalValue; +use crate::message_pack_format::envelope; + +use super::{CAPACITY_CACHE_LEN, Coin, KLL, MAX_CACHEABLE_K, MAX_LEVELS, compute_max_capacity}; + +const KLL_KIND_FAMILY: u8 = 0x06; +/// kind_id for the compact, fixed-buffer [`KLL`] (`0x06 0x00`). +pub(crate) const KLL_KIND_COMPACT: &[u8] = &[KLL_KIND_FAMILY, 0x00]; +/// kind_id for the dynamic-buffer `KLLDynamic` (`0x06 0x01`). +pub(crate) const KLL_KIND_DYNAMIC: &[u8] = &[KLL_KIND_FAMILY, 0x01]; + +/// Names the wire item type carried in the KLL metadata (`item_type`). +/// Implemented only for the two wire-eligible retained-sample types; an exotic +/// `KLL` (any other `NumericalValue`) is not wire-serializable and must be +/// converted to one of these first. +pub trait KllWireItem: Copy { + /// Metadata `item_type` string — `"f64"` or `"i64"`. + const ITEM_TYPE: &'static str; +} +impl KllWireItem for f64 { + const ITEM_TYPE: &'static str = "f64"; +} +impl KllWireItem for i64 { + const ITEM_TYPE: &'static str = "i64"; +} + +/// KLL descriptor metadata (ASAPv1 §2), a msgpack **map** (`to_vec_named`) with +/// keys in this declaration order — the canonical order the wire spec fixes (Go +/// must mirror it). KLL does not hash, so there is **no hash-spec group**: the +/// fields are structural params only. `deny_unknown_fields` makes decode fail +/// closed on any unexpected key. +#[derive(Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct KllMetadata { + pub(crate) metadata_version: u8, + pub(crate) k: u32, + pub(crate) m: u32, + pub(crate) item_type: String, +} + +/// Builds the KLL descriptor metadata for a wire-eligible item type `T`. +pub(crate) fn kll_metadata(k: u32, m: u32) -> KllMetadata { + KllMetadata { + metadata_version: 1, + k, + m, + item_type: T::ITEM_TYPE.to_string(), + } +} + +/// The compaction coin's raw RNG state, carried in the payload so a decoded +/// sketch continues compacting deterministically. Serialized positionally (as +/// part of [`KllPayload`], `to_vec`) so it lands as a nested 3-element array +/// `[state, bit_cache, remaining_bits]`, mirroring `sketchlib-go`'s `CoinState`. +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct KllCoinWire { + pub(crate) state: u64, + pub(crate) bit_cache: u64, + pub(crate) remaining_bits: u32, +} + +/// KLL payload (ASAPv1 §3), a msgpack **array** (`to_vec`, positional): +/// `[levels, items, coin]`. `num_levels` is `levels.len() - 1` (derived, so not +/// stored); `k` / `m` / `item_type` live in the metadata. Element type of +/// `items` is fixed by the metadata `item_type`. +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct KllPayload { + pub(crate) levels: Vec, + pub(crate) items: Vec, + pub(crate) coin: KllCoinWire, +} + +/// Validates the shared parts of a decoded KLL payload's level layout against +/// the retained items, and the coin's `remaining_bits` bound. Returns +/// `(num_levels, Coin)` on success. Shared by both KLL variants' decoders so the +/// fail-closed rules stay identical. `k` / `m` are echoed structural params (the +/// sketch is sized from them), so they are not cross-checked here. +pub(crate) fn validate_kll_payload( + levels: &[u32], + items: &[T], + coin: &KllCoinWire, +) -> Result { + if levels.len() < 2 { + return Err(RmpDecodeError::Uncategorized(format!( + "KLL payload: levels too short (len {}, need >= 2)", + levels.len() + ))); + } + let num_levels = levels.len() - 1; + if num_levels > MAX_LEVELS { + return Err(RmpDecodeError::Uncategorized(format!( + "KLL payload: num_levels {num_levels} exceeds MAX_LEVELS {MAX_LEVELS}" + ))); + } + if levels[0] != 0 { + return Err(RmpDecodeError::Uncategorized(format!( + "KLL payload: levels[0] must be 0, got {}", + levels[0] + ))); + } + if levels.windows(2).any(|w| w[0] > w[1]) { + return Err(RmpDecodeError::Uncategorized( + "KLL payload: levels must be non-decreasing".to_string(), + )); + } + if *levels.last().unwrap() as usize != items.len() { + return Err(RmpDecodeError::Uncategorized(format!( + "KLL payload: levels[last] {} != items.len() {}", + levels.last().unwrap(), + items.len() + ))); + } + if coin.remaining_bits > u64::BITS { + return Err(RmpDecodeError::Uncategorized(format!( + "KLL payload: coin remaining_bits {} exceeds 64", + coin.remaining_bits + ))); + } + Ok(num_levels) +} + +/// Splits the envelope, checks the `kind_id`, and validates the metadata against +/// the target item type `T` (fail closed on any mismatch). `k` / `m` are +/// echoed back into the expected metadata (they are structural and sizing the +/// sketch from them is the whole point), so only `metadata_version` and +/// `item_type` are effectively pinned. Returns the decoded metadata plus the +/// raw payload bytes. Shared by both KLL variants. +pub(crate) fn split_and_validate_meta<'a, T: KllWireItem>( + bytes: &'a [u8], + expected_kind_id: &[u8], +) -> Result<(KllMetadata, &'a [u8]), RmpDecodeError> { + let (kind_id, metadata, payload) = + envelope::split(bytes).map_err(RmpDecodeError::Uncategorized)?; + if kind_id != expected_kind_id { + return Err(RmpDecodeError::Uncategorized(format!( + "KLL kind_id mismatch: stored {kind_id:?}, expected {expected_kind_id:?}" + ))); + } + let meta: KllMetadata = from_slice(metadata)?; + if meta != kll_metadata::(meta.k, meta.m) { + return Err(RmpDecodeError::Uncategorized( + "ASAPv1 KLL envelope: metadata mismatch".to_string(), + )); + } + // Fail closed on out-of-range `k` / `m`. `k` / `m` are echoed structural + // params (the sketch is sized from them), so the equality check above cannot + // pin them — without this bound a crafted `k`/`m` near `u32::MAX` would drive + // `compute_max_capacity` (compact decode) into a multi-terabyte allocation + // and abort the process. A legitimately serialized sketch always has + // `2 <= m <= k <= MAX_CACHEABLE_K` because the constructor clamps to that + // range (`init_internal`), so this never rejects real bytes. + if meta.m < 2 || meta.m > meta.k || meta.k > MAX_CACHEABLE_K as u32 { + return Err(RmpDecodeError::Uncategorized(format!( + "ASAPv1 KLL envelope: k={}, m={} outside valid range (2 <= m <= k <= {MAX_CACHEABLE_K})", + meta.k, meta.m + ))); + } + Ok((meta, payload)) +} + +// Wire serialization for the compact KLL. `wire` is a descendant of the sketch +// module, so these impls read the private fields and construct the struct +// directly. +impl KLL +where + T: NumericalValue + KllWireItem + Serialize + for<'de> Deserialize<'de>, +{ + /// Serializes the sketch into an ASAPv1 MessagePack envelope + /// (kind_id `0x06 0x00`). The retained samples use the top-most-level-first + /// layout that matches `sketchlib-go`'s `KLLState`. + pub fn serialize_to_bytes(&self) -> Result, RmpEncodeError> { + let metadata = rmp_serde::to_vec_named(&kll_metadata::(self.k as u32, self.m as u32))?; + let (state, bit_cache, remaining_bits) = self.co.to_wire(); + let payload = rmp_serde::to_vec(&KllPayload { + levels: self.wire_levels(), + items: self.wire_items(), + coin: KllCoinWire { + state, + bit_cache, + remaining_bits, + }, + })?; + Ok(envelope::encode(KLL_KIND_COMPACT, &metadata, &payload)) + } + + /// Deserializes a compact KLL from an ASAPv1 MessagePack envelope. Bytes + /// whose metadata does not match this item type are rejected (fail closed), + /// as are inconsistent level layouts. + pub fn deserialize_from_bytes(bytes: &[u8]) -> Result { + let (meta, payload_bytes) = split_and_validate_meta::(bytes, KLL_KIND_COMPACT)?; + let payload: KllPayload = from_slice(payload_bytes)?; + let num_levels = validate_kll_payload(&payload.levels, &payload.items, &payload.coin)?; + Self::from_wire_top_first(meta.k as usize, meta.m as usize, num_levels, payload) + } + + /// Rebuilds the compact leftward-grown buffer from the top-most-first wire + /// payload. Inverse of [`KLL::wire_levels`] / [`KLL::wire_items`]: L0's + /// input-order run is reversed back into the buffer (which stores L0 + /// reverse-input), and higher levels are copied as-is, so a round-trip + /// re-serializes to byte-identical output. + fn from_wire_top_first( + k: usize, + m: usize, + num_levels: usize, + payload: KllPayload, + ) -> Result { + let KllPayload { + levels, + items, + coin, + } = payload; + + let max_cap = compute_max_capacity(k, m); + let total = items.len(); + if total > max_cap { + return Err(RmpDecodeError::Uncategorized(format!( + "KLL payload: {total} items exceed max_capacity {max_cap} for k={k}, m={m}" + ))); + } + let offset = max_cap - total; + + let mut buf = vec![T::default(); max_cap].into_boxed_slice(); + let mut internal_levels = vec![0usize; MAX_LEVELS + 1].into_boxed_slice(); + + // Place levels bottom-first into the buffer starting at `offset` (free + // space stays at the front). Wire slot `top_i = num_levels - 1 - h` + // holds compactor level `h`. + let mut cursor = offset; + for h in 0..num_levels { + let top_i = num_levels - 1 - h; + let s = levels[top_i] as usize; + let e = levels[top_i + 1] as usize; + internal_levels[h] = cursor; + if h == 0 { + // Wire L0 is input order; the compact buffer stores L0 + // reverse-input, so reverse it back on the way in. + for (j, &v) in items[s..e].iter().rev().enumerate() { + buf[cursor + j] = v; + } + } else { + buf[cursor..cursor + (e - s)].copy_from_slice(&items[s..e]); + } + cursor += e - s; + } + // Sentinel boundary for the top level; slots beyond stay 0, matching a + // live sketch (which only writes `levels[num_levels]`), so a decoded + // sketch's level array is identical to the source's. + internal_levels[num_levels] = cursor; + debug_assert_eq!(cursor, max_cap); + + let mut sketch = KLL { + items: buf, + levels: internal_levels, + k, + m, + num_levels, + max_capacity: max_cap, + co: Coin::from_wire(coin.state, coin.bit_cache, coin.remaining_bits as u8), + // The wire form does not carry the explicit-seed flag; a + // round-tripped sketch keeps its coin state but re-randomizes on a + // later clear(). Callers needing clear()-determinism rebuild via + // init_with_seed. + seed: None, + capacity_cache: [0; CAPACITY_CACHE_LEN], + top_height: 0, + level0_capacity: 0, + merge_buf: Vec::with_capacity(k), + }; + sketch.rebuild_capacity_cache(); + sketch.ensure_levels_sorted(); + Ok(sketch) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sketches::kll::KLL; + + fn build_kll(k: i32, seed: u64, n: u64) -> KLL { + let mut sketch = KLL::::init_kll_with_seed(k, seed); + for v in 1..=n { + sketch.update(&(v as f64)); + } + sketch + } + + #[test] + fn kll_envelope_structure_and_round_trip() { + // Enough data to force several compaction levels. + let sketch = build_kll(200, 42, 200_000); + let bytes = sketch.serialize_to_bytes().expect("serialize"); + + assert!(bytes.starts_with(envelope::MAGIC)); + assert_eq!(bytes[6], envelope::VERSION); + assert_eq!(bytes[7], 2, "kind_id_len"); + assert_eq!(&bytes[8..10], KLL_KIND_COMPACT); + + let decoded = KLL::::deserialize_from_bytes(&bytes).expect("decode"); + // Byte-stable round trip. + assert_eq!( + decoded.serialize_to_bytes().expect("re-serialize"), + bytes, + "KLL serialized bytes differed after round trip" + ); + // Quantiles preserved exactly (same retained state). + for &q in &[0.0, 0.01, 0.25, 0.5, 0.75, 0.99, 1.0] { + assert_eq!( + decoded.quantile(q), + sketch.quantile(q), + "quantile mismatch at q={q} after round trip" + ); + } + } + + #[test] + fn kll_empty_round_trip() { + let sketch = KLL::::init_kll_with_seed(200, 7); + let bytes = sketch.serialize_to_bytes().expect("serialize"); + let decoded = KLL::::deserialize_from_bytes(&bytes).expect("decode"); + assert_eq!(decoded.count(), 0); + assert_eq!(decoded.serialize_to_bytes().expect("re-serialize"), bytes); + } + + #[test] + fn kll_i64_round_trip() { + let mut sketch = KLL::::init_kll_with_seed(200, 5); + for v in 1..=50_000i64 { + sketch.update(&v); + } + let bytes = sketch.serialize_to_bytes().expect("serialize"); + assert_eq!(&bytes[8..10], KLL_KIND_COMPACT); + let decoded = KLL::::deserialize_from_bytes(&bytes).expect("decode"); + assert_eq!(decoded.serialize_to_bytes().expect("re-serialize"), bytes); + assert_eq!(decoded.count(), sketch.count()); + } + + #[test] + fn kll_item_type_cross_rejection() { + // f64 bytes must not decode into an i64 KLL (metadata item_type mismatch). + let sketch = build_kll(200, 1, 1000); + let bytes = sketch.serialize_to_bytes().expect("serialize"); + assert!( + KLL::::deserialize_from_bytes(&bytes).is_err(), + "f64 KLL bytes must be rejected by an i64 decoder" + ); + } + + #[test] + fn kll_metadata_rejects_unknown_keys() { + #[derive(Serialize)] + struct WithExtra { + metadata_version: u8, + k: u32, + m: u32, + item_type: String, + bogus_field: u8, + } + let extra = WithExtra { + metadata_version: 1, + k: 200, + m: 8, + item_type: "f64".to_string(), + bogus_field: 7, + }; + let bytes = rmp_serde::to_vec_named(&extra).expect("encode"); + assert!( + rmp_serde::from_slice::(&bytes).is_err(), + "an unexpected metadata key must be rejected" + ); + } + + #[test] + fn kll_rejects_inconsistent_levels() { + // Valid envelope + metadata, but levels[last] != items.len(). + let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8)).unwrap(); + let payload = rmp_serde::to_vec(&KllPayload:: { + levels: vec![0, 3], + items: vec![1.0, 2.0], // len 2, but levels claim 3 + coin: KllCoinWire { + state: 1, + bit_cache: 0, + remaining_bits: 0, + }, + }) + .unwrap(); + let bytes = envelope::encode(KLL_KIND_COMPACT, &metadata, &payload); + assert!( + KLL::::deserialize_from_bytes(&bytes).is_err(), + "inconsistent level layout must be rejected, not panic" + ); + } + + #[test] + fn kll_rejects_out_of_range_k_m() { + // A crafted envelope with an enormous `k`/`m` must fail closed, not drive + // compute_max_capacity into a giant allocation. Empty payload so nothing + // downstream depends on the dimensions. + let empty_payload = || { + rmp_serde::to_vec(&KllPayload:: { + levels: vec![0, 0], + items: Vec::new(), + coin: KllCoinWire { + state: 1, + bit_cache: 0, + remaining_bits: 0, + }, + }) + .unwrap() + }; + for (k, m) in [ + (u32::MAX, u32::MAX), + (MAX_CACHEABLE_K as u32 + 1, 8), + (200, 1), + ] { + let metadata = rmp_serde::to_vec_named(&kll_metadata::(k, m)).unwrap(); + let bytes = envelope::encode(KLL_KIND_COMPACT, &metadata, &empty_payload()); + assert!( + KLL::::deserialize_from_bytes(&bytes).is_err(), + "k={k}, m={m} must be rejected, not allocated" + ); + } + } + + #[test] + fn kll_dynamic_kind_id_rejected_by_compact() { + // A compact decoder must reject the dynamic kind_id. + let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8)).unwrap(); + let payload = rmp_serde::to_vec(&KllPayload:: { + levels: vec![0, 0], + items: Vec::::new(), + coin: KllCoinWire { + state: 1, + bit_cache: 0, + remaining_bits: 0, + }, + }) + .unwrap(); + let bytes = envelope::encode(KLL_KIND_DYNAMIC, &metadata, &payload); + assert!( + KLL::::deserialize_from_bytes(&bytes).is_err(), + "dynamic kind_id must be rejected by the compact decoder" + ); + } +} diff --git a/src/sketches/kll_dynamic.rs b/src/sketches/kll_dynamic.rs index 6459a95..6011aaa 100644 --- a/src/sketches/kll_dynamic.rs +++ b/src/sketches/kll_dynamic.rs @@ -12,8 +12,6 @@ //! //! - -use rmp_serde::decode::Error as RmpDecodeError; -use rmp_serde::encode::Error as RmpEncodeError; use serde::{Deserialize, Serialize}; use crate::common::input::data_input_to_f64; @@ -22,6 +20,8 @@ use crate::{DataInput, Vector1D}; use super::kll::{Coin, merge_sorted_runs, randomly_halve_up}; +mod wire; + const CAPACITY_CACHE_LEN: usize = 20; const MAX_CACHEABLE_K: usize = 26_602; const CAPACITY_DECAY: f64 = 2.0 / 3.0; @@ -455,25 +455,6 @@ impl KLLDynamic { fn buffer_size(&self) -> usize { self.items.len() } - - /// Serialize the sketch into MessagePack bytes. - pub fn serialize_to_bytes(&self) -> Result, RmpEncodeError> - where - T: Serialize, - { - rmp_serde::to_vec(self) - } - - /// Deserialize a sketch from MessagePack bytes. - pub fn deserialize_from_bytes(bytes: &[u8]) -> Result - where - T: for<'de> Deserialize<'de>, - { - rmp_serde::from_slice(bytes).map(|mut sketch: KLLDynamic| { - sketch.rebuild_capacity_cache(); - sketch - }) - } } impl KLLDynamic { diff --git a/src/sketches/kll_dynamic/wire.rs b/src/sketches/kll_dynamic/wire.rs new file mode 100644 index 0000000..28f5740 --- /dev/null +++ b/src/sketches/kll_dynamic/wire.rs @@ -0,0 +1,151 @@ +//! ASAPv1 wire serialization for the dynamic KLL quantile sketch. +//! +//! Child submodule of [`crate::sketches::kll_dynamic`]: it holds the +//! `serialize_to_bytes` / `deserialize_from_bytes` impls for [`KLLDynamic`] +//! (kind_id `0x06 0x01`). The wire DTOs (metadata, payload, coin) and the +//! fail-closed validation are **shared** with the compact KLL and live in +//! [`crate::sketches::kll::wire`]; only the kind_id differs. See +//! `docs/asapv1_wire_format.md` §3. +//! +//! Unlike the compact KLL, `KLLDynamic` already stores its retained samples in +//! the wire's top-most-level-first order with level 0 in input order (it appends +//! to the end of a growable buffer), so encode/decode are a direct copy — no +//! buffer reversal is needed. + +use rmp_serde::{decode::Error as RmpDecodeError, encode::Error as RmpEncodeError, from_slice}; + +use crate::Vector1D; +use crate::message_pack_format::envelope; +use crate::sketches::kll::{ + Coin, KLL_KIND_DYNAMIC, KllCoinWire, KllPayload, KllWireItem, kll_metadata, + split_and_validate_meta, validate_kll_payload, +}; + +use super::{CAPACITY_CACHE_LEN, KLLDynamic}; + +// `wire` is a descendant of the sketch module, so these impls read the private +// fields and construct the struct directly. +impl KLLDynamic +where + T: crate::common::numerical::NumericalValue + + KllWireItem + + serde::Serialize + + for<'de> serde::Deserialize<'de>, +{ + /// Serializes the sketch into an ASAPv1 MessagePack envelope + /// (kind_id `0x06 0x01`). `KLLDynamic`'s buffer is already in the wire's + /// top-most-level-first / input-order-L0 layout, so `levels` and `items` are + /// copied straight out. + pub fn serialize_to_bytes(&self) -> Result, RmpEncodeError> { + let metadata = rmp_serde::to_vec_named(&kll_metadata::(self.k as u32, self.m as u32))?; + let (state, bit_cache, remaining_bits) = self.co.to_wire(); + let payload = rmp_serde::to_vec(&KllPayload { + levels: self.levels.as_slice().iter().map(|&l| l as u32).collect(), + items: self.items.as_slice().to_vec(), + coin: KllCoinWire { + state, + bit_cache, + remaining_bits, + }, + })?; + Ok(envelope::encode(KLL_KIND_DYNAMIC, &metadata, &payload)) + } + + /// Deserializes a `KLLDynamic` from an ASAPv1 MessagePack envelope. Bytes + /// whose metadata does not match this item type are rejected (fail closed), + /// as are inconsistent level layouts. + pub fn deserialize_from_bytes(bytes: &[u8]) -> Result { + let (meta, payload_bytes) = split_and_validate_meta::(bytes, KLL_KIND_DYNAMIC)?; + let payload: KllPayload = from_slice(payload_bytes)?; + let num_levels = validate_kll_payload(&payload.levels, &payload.items, &payload.coin)?; + + let KllPayload { + levels, + items, + coin, + } = payload; + let mut sketch = KLLDynamic { + items: Vector1D::from_vec(items), + levels: Vector1D::from_vec(levels.into_iter().map(|l| l as usize).collect()), + k: meta.k as usize, + m: meta.m as usize, + num_levels, + co: Coin::from_wire(coin.state, coin.bit_cache, coin.remaining_bits as u8), + capacity_cache: [0; CAPACITY_CACHE_LEN], + top_height: 0, + level0_capacity: 0, + }; + sketch.rebuild_capacity_cache(); + Ok(sketch) + } +} + +#[cfg(test)] +mod tests { + use crate::message_pack_format::envelope; + use crate::sketches::kll::KLL_KIND_DYNAMIC; + use crate::sketches::kll_dynamic::KLLDynamic; + + fn build_dynamic(k: i32, n: u64) -> KLLDynamic { + let mut sketch = KLLDynamic::::init_kll(k); + for v in 1..=n { + sketch.update(&(v as f64)); + } + sketch + } + + #[test] + fn kll_dynamic_envelope_structure_and_round_trip() { + let sketch = build_dynamic(200, 200_000); + let bytes = sketch.serialize_to_bytes().expect("serialize"); + + assert!(bytes.starts_with(envelope::MAGIC)); + assert_eq!(bytes[6], envelope::VERSION); + assert_eq!(bytes[7], 2, "kind_id_len"); + assert_eq!(&bytes[8..10], KLL_KIND_DYNAMIC); + + let decoded = KLLDynamic::::deserialize_from_bytes(&bytes).expect("decode"); + assert_eq!( + decoded.serialize_to_bytes().expect("re-serialize"), + bytes, + "KLLDynamic serialized bytes differed after round trip" + ); + for &q in &[0.0, 0.01, 0.25, 0.5, 0.75, 0.99, 1.0] { + assert_eq!( + decoded.quantile(q), + sketch.quantile(q), + "quantile mismatch at q={q} after round trip" + ); + } + } + + #[test] + fn kll_dynamic_empty_round_trip() { + let sketch = KLLDynamic::::init_kll(200); + let bytes = sketch.serialize_to_bytes().expect("serialize"); + let decoded = KLLDynamic::::deserialize_from_bytes(&bytes).expect("decode"); + assert_eq!(decoded.serialize_to_bytes().expect("re-serialize"), bytes); + } + + #[test] + fn kll_dynamic_i64_round_trip() { + let mut sketch = KLLDynamic::::init_kll(200); + for v in 1..=50_000i64 { + sketch.update(&v); + } + let bytes = sketch.serialize_to_bytes().expect("serialize"); + assert_eq!(&bytes[8..10], KLL_KIND_DYNAMIC); + let decoded = KLLDynamic::::deserialize_from_bytes(&bytes).expect("decode"); + assert_eq!(decoded.serialize_to_bytes().expect("re-serialize"), bytes); + } + + #[test] + fn kll_dynamic_item_type_cross_rejection() { + let sketch = build_dynamic(200, 1000); + let bytes = sketch.serialize_to_bytes().expect("serialize"); + assert!( + KLLDynamic::::deserialize_from_bytes(&bytes).is_err(), + "f64 KLLDynamic bytes must be rejected by an i64 decoder" + ); + } +} diff --git a/tests/asapv1_golden.rs b/tests/asapv1_golden.rs index a3fe105..b73a703 100644 --- a/tests/asapv1_golden.rs +++ b/tests/asapv1_golden.rs @@ -11,7 +11,7 @@ use asap_sketchlib::{ Classic, CountMin, ErtlMLE, FastPath, HllSketch, HllVariant, HyperLogLogHIPP12, HyperLogLogP12, - MessagePackCodec, RegularPath, Vector2D, + KLL, MessagePackCodec, RegularPath, Vector2D, }; fn decode_hex(s: &str) -> Vec { @@ -28,6 +28,8 @@ const GOLDEN_ERTL: &str = include_str!("../asapv1_golden/hll_ertl_mle_p12.hex"); const GOLDEN_HIP: &str = include_str!("../asapv1_golden/hll_hip_p12.hex"); const GOLDEN_CMS_I64: &str = include_str!("../asapv1_golden/cms_i64_regular_2x3.hex"); const GOLDEN_CMS_F64: &str = include_str!("../asapv1_golden/cms_f64_fast_2x3.hex"); +const GOLDEN_KLL_F64: &str = include_str!("../asapv1_golden/kll_f64_k200.hex"); +const GOLDEN_KLL_I64: &str = include_str!("../asapv1_golden/kll_i64_k200.hex"); /// The known P12 register pattern shared by all three HLL fixtures. fn p12_registers() -> Vec { @@ -151,3 +153,49 @@ fn cms_f64_fast_2x3_matches_golden() { assert_eq!(decoded.as_storage().as_slice(), flat.as_slice()); assert_eq!(decoded.serialize_to_bytes().expect("re-serialize"), want); } + +// --------------------------------------------------------------------------- +// KLL: build known state (k=200, seed 42, integers 1..=50 — below the level-0 +// capacity, so no compaction fires and the retained set is deterministic) -> +// serialize == golden, and golden round-trips. Matches the deterministic +// scenario the proto parity test uses (sketchlib-go's KLLSketch over the same +// input), so the coin state (42) lines up cross-language. +// --------------------------------------------------------------------------- + +/// A k=200 KLL over `1..=50` with a fixed compaction seed: fully deterministic. +fn kll_1to50 + asap_sketchlib::NumericalValue>(seed: u64) -> KLL { + let mut sketch = KLL::::init_kll_with_seed(200, seed); + for v in 1..=50u8 { + sketch.update(&T::from(v)); + } + sketch +} + +#[test] +fn kll_f64_k200_matches_golden() { + let want = decode_hex(GOLDEN_KLL_F64); + + let sketch = kll_1to50::(42); + let got = sketch.serialize_to_bytes().expect("serialize"); + assert_eq!(got, want, "KLL f64 bytes diverge from golden"); + + // Golden round-trips: decode, and re-encode is byte-identical. + let decoded = KLL::::deserialize_from_bytes(&want).expect("decode"); + assert_eq!(decoded.serialize_to_bytes().expect("re-serialize"), want); + assert_eq!(decoded.quantile(0.0), 1.0); + assert_eq!(decoded.quantile(1.0), 50.0); +} + +#[test] +fn kll_i64_k200_matches_golden() { + let want = decode_hex(GOLDEN_KLL_I64); + + let sketch = kll_1to50::(42); + let got = sketch.serialize_to_bytes().expect("serialize"); + assert_eq!(got, want, "KLL i64 bytes diverge from golden"); + + let decoded = KLL::::deserialize_from_bytes(&want).expect("decode"); + assert_eq!(decoded.serialize_to_bytes().expect("re-serialize"), want); + assert_eq!(decoded.quantile(0.0), 1.0); + assert_eq!(decoded.quantile(1.0), 50.0); +} From 464fa49838ff22dbe39468a31a73e340ea88f252 Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Sun, 19 Jul 2026 04:37:49 -0400 Subject: [PATCH 2/4] harden(kll): fail closed on crafted dimensions and level layouts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second-review follow-up. Two decode-side robustness gaps: 1. The `k`/`m` bound added to the ASAPv1 decoder was missing on the retained nested serde `Deserialize` for `KLL` (the `HydraCounter` path, reachable via `Hydra::deserialize_from_bytes` and direct `rmp_serde::from_slice::>`). A crafted `k` near `usize::MAX` overflowed `compute_max_capacity` / drove a huge allocation. Apply the same `2 <= m <= k <= MAX_CACHEABLE_K` bound plus level-layout validation (so the buffer math can't underflow or index OOB). 2. `validate_kll_payload` (both ASAPv1 variants) accepted a structurally-valid but adversarial level distribution — e.g. many items parked at a high compactor level — that decoded fine but then overflowed the weighted `count()`/`rank()`/`cdf()` (`size * 2^h`) at query time. Add a shared `checked_weighted_count` guard so decode rejects any layout whose weighted count would overflow; a live sketch's weighted count equals its ingested item count, so real bytes are never rejected. Adds regression tests for both (nested-serde crafted `k`, ASAPv1 weighted-count overflow). Full suite + hydra green; clippy clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/sketches/kll.rs | 75 ++++++++++++++++++++++++++++++++++++++++ src/sketches/kll/wire.rs | 50 ++++++++++++++++++++++++++- 2 files changed, 124 insertions(+), 1 deletion(-) diff --git a/src/sketches/kll.rs b/src/sketches/kll.rs index 64637f5..858a5f6 100644 --- a/src/sketches/kll.rs +++ b/src/sketches/kll.rs @@ -126,6 +126,25 @@ fn compute_max_capacity(k: usize, m: usize) -> usize { total } +/// Checked total weighted item count for compactor-level sizes given +/// **bottom-first** (`sizes[h]` is the item count at compactor level `h`, whose +/// weight is `2^h`). Returns `None` on `usize` overflow. +/// +/// A live sketch's weighted count equals the number of ingested items, so it +/// always fits. Decoders use this to reject a crafted-but-structurally-valid +/// level layout (e.g. many items parked at a high level) that would otherwise +/// overflow `count()` / `rank()` / `cdf()` at query time — a fail-closed guard +/// so decode never yields a sketch that panics on a later query. +pub(crate) fn checked_weighted_count(sizes_bottom_first: &[usize]) -> Option { + let mut total = 0usize; + for (h, &size) in sizes_bottom_first.iter().enumerate() { + // `h < MAX_LEVELS (61) < 64`, so the shift amount is always valid. + let weight = 1usize.checked_shl(h as u32)?; + total = total.checked_add(size.checked_mul(weight)?)?; + } + Some(total) +} + /// Halves a sorted run, placing survivors in the **upper** (right) half of /// `items[begin..begin+pop]` so they are contiguous with the level above. /// Traverses backwards to avoid overwriting unread source elements. @@ -877,9 +896,45 @@ impl KLL { impl<'de, T: NumericalValue + Deserialize<'de>> Deserialize<'de> for KLL { fn deserialize>(deserializer: D) -> Result { + use serde::de::Error as _; let wire = KLLWire::::deserialize(deserializer)?; + + // Fail closed on crafted bytes. This nested serde path is reachable with + // untrusted input via `HydraCounter::KLL` (`Hydra::deserialize_from_bytes`) + // and direct `rmp_serde::from_slice::>`, so it must guard the same + // way the ASAPv1 decoder does: bound `k`/`m` to the range the constructor + // clamps to (else `compute_max_capacity` / `Vec::with_capacity` blow up), + // and validate the level layout (else the buffer math below underflows or + // indexes out of bounds, and `count()` overflows on a later query). + if wire.m < 2 || wire.m > wire.k || wire.k > MAX_CACHEABLE_K { + return Err(D::Error::custom(format!( + "KLL: k={}, m={} outside valid range (2 <= m <= k <= {MAX_CACHEABLE_K})", + wire.k, wire.m + ))); + } + if wire.num_levels == 0 + || wire.num_levels > MAX_LEVELS + || wire.levels.len() != wire.num_levels + 1 + || wire.levels.first() != Some(&0) + || wire.levels.windows(2).any(|w| w[0] > w[1]) + || wire.levels.last() != Some(&wire.items.len()) + { + return Err(D::Error::custom("KLL: inconsistent level layout")); + } + let sizes: Vec = wire.levels.windows(2).map(|w| w[1] - w[0]).collect(); + if checked_weighted_count(&sizes).is_none() { + return Err(D::Error::custom( + "KLL: level layout overflows weighted count", + )); + } + let max_cap = compute_max_capacity(wire.k, wire.m); let used_len = wire.items.len(); + if used_len > max_cap { + return Err(D::Error::custom(format!( + "KLL: {used_len} items exceed max_capacity {max_cap}" + ))); + } let offset = max_cap - used_len; let mut items = vec![T::default(); max_cap].into_boxed_slice(); @@ -1014,6 +1069,26 @@ mod tests { use super::*; use crate::test_utils::{sample_uniform_f64, sample_zipf_f64}; + // Crafted nested-serde bytes (the `HydraCounter::KLL` path, reachable with + // untrusted input) with an out-of-range `k` must fail closed, not overflow + // `compute_max_capacity` / drive a huge allocation. + #[test] + fn kllwire_serde_rejects_crafted_dimensions() { + let wire = KLLWire:: { + items: vec![], + levels: vec![0, 0], + k: usize::MAX, + m: 8, + num_levels: 1, + co: Coin::from_seed(1), + }; + let bytes = rmp_serde::to_vec(&wire).expect("encode crafted wire"); + assert!( + rmp_serde::from_slice::>(&bytes).is_err(), + "an out-of-range k must be rejected by the nested serde decoder" + ); + } + // Direct reconstruction from portable wire state must be BIT-EXACT: // identical quantiles to the source sketch (unlike a lossy // replay-through-`update()` reconstruction). diff --git a/src/sketches/kll/wire.rs b/src/sketches/kll/wire.rs index 854367f..acd6c4e 100644 --- a/src/sketches/kll/wire.rs +++ b/src/sketches/kll/wire.rs @@ -32,7 +32,10 @@ use serde::{Deserialize, Serialize}; use crate::common::numerical::NumericalValue; use crate::message_pack_format::envelope; -use super::{CAPACITY_CACHE_LEN, Coin, KLL, MAX_CACHEABLE_K, MAX_LEVELS, compute_max_capacity}; +use super::{ + CAPACITY_CACHE_LEN, Coin, KLL, MAX_CACHEABLE_K, MAX_LEVELS, checked_weighted_count, + compute_max_capacity, +}; const KLL_KIND_FAMILY: u8 = 0x06; /// kind_id for the compact, fixed-buffer [`KLL`] (`0x06 0x00`). @@ -147,6 +150,22 @@ pub(crate) fn validate_kll_payload( coin.remaining_bits ))); } + // Reject a structurally-valid but adversarial level distribution (e.g. many + // items parked at a high compactor level) that would overflow the weighted + // `count()` / `rank()` / `cdf()` at query time. `levels` is top-most-first, + // so map slot `i` to compactor level `num_levels - 1 - i` to get sizes + // bottom-first for the weight check. + let sizes: Vec = (0..num_levels) + .map(|h| { + let i = num_levels - 1 - h; + (levels[i + 1] - levels[i]) as usize + }) + .collect(); + if checked_weighted_count(&sizes).is_none() { + return Err(RmpDecodeError::Uncategorized( + "KLL payload: level layout overflows weighted count".to_string(), + )); + } Ok(num_levels) } @@ -453,6 +472,35 @@ mod tests { } } + #[test] + fn kll_rejects_weighted_count_overflow() { + // Structurally valid (monotonic levels, levels[last]==items.len()) but + // adversarial: 16 items parked at compactor level 60 makes 16 * 2^60 + // overflow usize in count(). Decode must reject, not hand back a sketch + // that panics on the first query. + let num_levels = 61usize; + // Top-most-first cumulative levels: the top level (slot 0) holds all 16 + // items, every lower level is empty. + let mut levels = vec![16u32; num_levels + 1]; + levels[0] = 0; + let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8)).unwrap(); + let payload = rmp_serde::to_vec(&KllPayload:: { + levels, + items: vec![1.0; 16], + coin: KllCoinWire { + state: 1, + bit_cache: 0, + remaining_bits: 0, + }, + }) + .unwrap(); + let bytes = envelope::encode(KLL_KIND_COMPACT, &metadata, &payload); + assert!( + KLL::::deserialize_from_bytes(&bytes).is_err(), + "a level layout that overflows the weighted count must be rejected" + ); + } + #[test] fn kll_dynamic_kind_id_rejected_by_compact() { // A compact decoder must reject the dynamic kind_id. From 8be5eb1db7df5d54fb2e44083409ea9fbb1ecd52 Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Mon, 20 Jul 2026 02:01:14 -0400 Subject: [PATCH 3/4] feat(kll): record reproducible compaction seed in KLL metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an optional `seed` key to the KLL ASAPv1 metadata so a decoded sketch keeps clear()-reproducibility instead of falling back to wall-clock. It's construction config, so it lives in the metadata (config→metadata rule), and it's the first optional key in v1: present only when the sketch carries a seed (Some), omitted otherwise (skip_serializing_if + default). The payload's `coin` already carries the RNG's current position (enough to resume compaction); `seed` is what a later clear() re-seeds from. Scope: compact KLL populates and restores it. KLLDynamic has no seed concept and always omits the key — the two variants are deliberately not forced to be symmetric here. A consumer that doesn't use the key (including Go) must still preserve it verbatim on re-encode. Regenerate the kll_f64/kll_i64 goldens (metadata gains seed=42) and add tests: seed present-when-seeded / omitted-when-unseeded, and that it survives a round trip so a decoded sketch's clear() stays deterministic. Document as decision Q-KLL-SEED. Relates to determinism gap tracked in #77 (tumbling KLL uses unseeded init) — that fix is separate (framework-side). Co-Authored-By: Claude Opus 4.8 (1M context) --- asapv1_golden/README.md | 4 +- asapv1_golden/kll_f64_k200.hex | 2 +- asapv1_golden/kll_i64_k200.hex | 2 +- docs/asapv1_wire_format.md | 10 ++-- src/sketches/kll/wire.rs | 94 ++++++++++++++++++++++++++------ src/sketches/kll_dynamic/wire.rs | 5 +- 6 files changed, 92 insertions(+), 25 deletions(-) diff --git a/asapv1_golden/README.md b/asapv1_golden/README.md index e6084b2..fe009f2 100644 --- a/asapv1_golden/README.md +++ b/asapv1_golden/README.md @@ -29,8 +29,8 @@ golden tests the **wire encoding**, isolated from the hash functions. | `hll_hip_p12` | HLL HIP, P12 | `01 03` | same registers + `hip_kxq0=1.5, hip_kxq1=2.5, hip_est=3.0` | | `cms_i64_regular_2x3` | Count-Min i64, RegularPath | `02 00` | 2×3 row-major `[[0,1,127],[128,300,65536]]` | | `cms_f64_fast_2x3` | Count-Min f64, FastPath | `02 00` | 2×3 row-major `[[0.0,1.5,2.25],[3.75,4.125,5.0625]]` | -| `kll_f64_k200` | KLL f64, k=200 | `06 00` | integers `1..=50`, compaction seed 42 | -| `kll_i64_k200` | KLL i64, k=200 | `06 00` | integers `1..=50`, compaction seed 42 | +| `kll_f64_k200` | KLL f64, k=200 | `06 00` | integers `1..=50`, compaction seed 42 (recorded in metadata as `seed`) | +| `kll_i64_k200` | KLL i64, k=200 | `06 00` | integers `1..=50`, compaction seed 42 (recorded in metadata as `seed`) | The i64 fixture deliberately spans the msgpack integer width boundaries (positive fixint / uint8 / uint16 / uint32) to lock the "non-negative integer → diff --git a/asapv1_golden/kll_f64_k200.hex b/asapv1_golden/kll_f64_k200.hex index 354d3ce..1df1f66 100644 --- a/asapv1_golden/kll_f64_k200.hex +++ b/asapv1_golden/kll_f64_k200.hex @@ -1 +1 @@ -4153415076310102060000000028000001cd84b06d657461646174615f76657273696f6e01a16bccc8a16d08a96974656d5f74797065a366363493920032dc0032cb3ff0000000000000cb4000000000000000cb4008000000000000cb4010000000000000cb4014000000000000cb4018000000000000cb401c000000000000cb4020000000000000cb4022000000000000cb4024000000000000cb4026000000000000cb4028000000000000cb402a000000000000cb402c000000000000cb402e000000000000cb4030000000000000cb4031000000000000cb4032000000000000cb4033000000000000cb4034000000000000cb4035000000000000cb4036000000000000cb4037000000000000cb4038000000000000cb4039000000000000cb403a000000000000cb403b000000000000cb403c000000000000cb403d000000000000cb403e000000000000cb403f000000000000cb4040000000000000cb4040800000000000cb4041000000000000cb4041800000000000cb4042000000000000cb4042800000000000cb4043000000000000cb4043800000000000cb4044000000000000cb4044800000000000cb4045000000000000cb4045800000000000cb4046000000000000cb4046800000000000cb4047000000000000cb4047800000000000cb4048000000000000cb4048800000000000cb4049000000000000932a0000 +415341507631010206000000002e000001cd85b06d657461646174615f76657273696f6e01a16bccc8a16d08a96974656d5f74797065a3663634a4736565642a93920032dc0032cb3ff0000000000000cb4000000000000000cb4008000000000000cb4010000000000000cb4014000000000000cb4018000000000000cb401c000000000000cb4020000000000000cb4022000000000000cb4024000000000000cb4026000000000000cb4028000000000000cb402a000000000000cb402c000000000000cb402e000000000000cb4030000000000000cb4031000000000000cb4032000000000000cb4033000000000000cb4034000000000000cb4035000000000000cb4036000000000000cb4037000000000000cb4038000000000000cb4039000000000000cb403a000000000000cb403b000000000000cb403c000000000000cb403d000000000000cb403e000000000000cb403f000000000000cb4040000000000000cb4040800000000000cb4041000000000000cb4041800000000000cb4042000000000000cb4042800000000000cb4043000000000000cb4043800000000000cb4044000000000000cb4044800000000000cb4045000000000000cb4045800000000000cb4046000000000000cb4046800000000000cb4047000000000000cb4047800000000000cb4048000000000000cb4048800000000000cb4049000000000000932a0000 diff --git a/asapv1_golden/kll_i64_k200.hex b/asapv1_golden/kll_i64_k200.hex index ac77317..9ab8d06 100644 --- a/asapv1_golden/kll_i64_k200.hex +++ b/asapv1_golden/kll_i64_k200.hex @@ -1 +1 @@ -41534150763101020600000000280000003d84b06d657461646174615f76657273696f6e01a16bccc8a16d08a96974656d5f74797065a369363493920032dc00320102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132932a0000 +415341507631010206000000002e0000003d85b06d657461646174615f76657273696f6e01a16bccc8a16d08a96974656d5f74797065a3693634a4736565642a93920032dc00320102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132932a0000 diff --git a/docs/asapv1_wire_format.md b/docs/asapv1_wire_format.md index 701d661..4599262 100644 --- a/docs/asapv1_wire_format.md +++ b/docs/asapv1_wire_format.md @@ -253,11 +253,11 @@ The metadata map is **two groups** of fields, written on the wire in this order: | Group | Role | Fields | | ----- | ---- | ------ | | **Hash spec** | how keys were hashed (check mergeability + re-hash a query key) | `metadata_version`, `hash_profile_id`, `hash_algorithm`, `seed_derivation`, `input_encoding`, `seed_list`, + the seed index(es) it uses | -| **Structural params** | parameters that shape the payload | `precision` (HLL); `rows`, `cols`, `counter_type`, `mode` (Count-Min); `k`, `m`, `item_type` (KLL) | +| **Structural params** | parameters that shape the payload | `precision` (HLL); `rows`, `cols`, `counter_type`, `mode` (Count-Min); `k`, `m`, `item_type`, optional `seed` (KLL) | The two tables below are the field-by-field detail of each group. -> **Non-hashing sketches omit the hash-spec group entirely** (Q-KLL). The hash spec answers "how were keys hashed"; a sketch that does not hash its inputs has no truthful answer. KLL is comparison-based — it orders raw numeric values with `total_cmp` and never invokes a hasher — so its metadata carries **only** structural params (`metadata_version`, `k`, `m`, `item_type`). This keeps the bytes honest (no meaningless `seed_list`) and is the reason KLL's metadata schema is not built from a `HashProfile` the way HLL's / Count-Min's are. +> **Non-hashing sketches omit the hash-spec group entirely** (Q-KLL). The hash spec answers "how were keys hashed"; a sketch that does not hash its inputs has no truthful answer. KLL is comparison-based — it orders raw numeric values with `total_cmp` and never invokes a hasher — so its metadata carries **only** structural params (`metadata_version`, `k`, `m`, `item_type`, and an optional `seed`). This keeps the bytes honest (no meaningless `seed_list`) and is the reason KLL's metadata schema is not built from a `HashProfile` the way HLL's / Count-Min's are. Note `seed` is unrelated to the hash `seed_list`: it is the KLL compaction RNG's reproducible seed, and it is the **only optional key** in v1 (present only when the sketch was built with one; the key is omitted otherwise, and `KLLDynamic` never emits it). A consumer that does not use it (including Go) MUST still preserve it verbatim on re-encode so the bytes round-trip identically. **Hash spec** @@ -286,6 +286,7 @@ The two tables below are the field-by-field detail of each group. | `k` | u32 | KLL | compactor capacity (accuracy parameter) | | `m` | u32 | KLL | minimum level capacity | | `item_type` | string | KLL | `"f64"` or `"i64"`; element type of `items` | +| `seed` | u64 | KLL | **optional**; the reproducible compaction seed, present only when the sketch carries one, else the key is omitted. Compact KLL only (`KLLDynamic` never emits it). The one optional key in v1. | Count-Min's matrix dimensions are **configuration** (they shape the payload, like HLL's `precision`), so per the config-to-metadata rule they live here. Count-Min's canonical structural-param order is `... matrix_seed_index, rows, cols, counter_type, mode`; this is the wire contract and Go must mirror it verbatim. @@ -374,7 +375,7 @@ Both modes, `FastPath` and `RegularPath`, serialize directly (you'd only "conver ### 3.3: KLL payload (`0x06 0x00` compact / `0x06 0x01` dynamic) -Both KLL variants (the compact fixed-buffer `KLL` and the growable `KLLDynamic`) share **one payload shape**; they differ only by `kind_id` (like HLL's Classic vs Ertl-MLE), because their in-memory buffers differ but the serialized quantile state is the same. The accuracy params `k` / `m` and the `item_type` live in the metadata (§2), so the payload is just the retained state: +Both KLL variants (the compact fixed-buffer `KLL` and the growable `KLLDynamic`) share **one payload shape**; they differ only by `kind_id` (like HLL's Classic vs Ertl-MLE), because their in-memory buffers differ but the serialized quantile state is the same. The accuracy params `k` / `m`, the `item_type`, and the optional `seed` live in the metadata (§2), so the payload is just the retained state: | Pos | Field | Type | Notes | | ----- | ------- | ------ | ------- | @@ -509,7 +510,8 @@ Keep it through the transition, retire `portable` once goldens are in place. - **Q-SEEDS**: `seed_list` is **inlined** in v1 so the bytes self-describe the hash. Resolving seeds from `hash_profile_id` alone is a v2 space optimization. Each sketch still carries only the seed *index* it uses. - **Q-PROFILE**: the hash-spec metadata is **derived from the hasher's `HashProfile`** (`hll_metadata::` / `cms_metadata::`) and never hardcoded, so it is always truthful to the hasher. Custom hash profiles are **supported and self-describing**. Merge compatibility is hash-spec equality, so a custom-profile sketch is not mergeable with a standard one. - **Q-CMS**: Count-Min is one `kind_id` (`0x02 0x00`); counter type and mode live in the metadata, so the id stays single. -- **Q-KLL**: KLL metadata carries **no hash-spec group** — KLL is comparison-based and never hashes, so those fields have no truthful value. Its metadata is structural-only (`metadata_version`, `k`, `m`, `item_type`) and is *not* `HashProfile`-derived. The two KLL variants (compact `0x06 0x00`, dynamic `0x06 0x01`) share one payload `[levels, items, coin]` and differ only by `kind_id`. `item_type` (`"f64"`/`"i64"`) is a metadata param, not a separate `kind_id` (mirrors Q-CMS's `counter_type`). Retained samples use the top-most-level-first layout that matches `sketchlib-go`'s `KLLState`. +- **Q-KLL**: KLL metadata carries **no hash-spec group** — KLL is comparison-based and never hashes, so those fields have no truthful value. Its metadata is structural-only (`metadata_version`, `k`, `m`, `item_type`, optional `seed`) and is *not* `HashProfile`-derived. The two KLL variants (compact `0x06 0x00`, dynamic `0x06 0x01`) share one payload `[levels, items, coin]` and differ only by `kind_id`. `item_type` (`"f64"`/`"i64"`) is a metadata param, not a separate `kind_id` (mirrors Q-CMS's `counter_type`). Retained samples use the top-most-level-first layout that matches `sketchlib-go`'s `KLLState`. +- **Q-KLL-SEED**: KLL records its reproducible compaction `seed` as an **optional** metadata key. It is construction config (so metadata, not payload, per the config→metadata rule), and it is the first optional key in v1: present only when the sketch carries a seed, omitted otherwise. Rationale: the payload's `coin` already carries the RNG's *current* position (enough to resume compaction), but `seed` is what a later `clear()` re-seeds from — so serializing it lets a decoded sketch keep `clear()`-reproducibility instead of falling back to wall-clock. Cost of omitting it is bounded (only a decoded-then-`clear()`ed sketch loses cross-run byte reproducibility — never correctness), but it is cheap to carry and future-proofs the checkpoint/restore path. `KLLDynamic` has no seed concept and never emits the key; the two variants are deliberately **not** forced to be symmetric here. Go carries and preserves the key without interpreting it. - **Q-CMS-DIMS**: Count-Min `rows`/`cols` are **metadata** and the payload omits them. They are configuration that shapes the payload (like HLL's `precision`), so per the config-to-metadata rule they belong in the descriptor. The payload is then just `[counts]`. Canonical structural-param order: `... matrix_seed_index, rows, cols, counter_type, mode`. - **Q-VER**: no payload version field. A new incompatible encoding gets a **new `kind_id`**; retired ids are reserved forever and never recycled. - **Encoding**: metadata + payload are both msgpack; payload is a positional array. Byte-level rules in Section 4. diff --git a/src/sketches/kll/wire.rs b/src/sketches/kll/wire.rs index acd6c4e..b2d8924 100644 --- a/src/sketches/kll/wire.rs +++ b/src/sketches/kll/wire.rs @@ -14,8 +14,8 @@ //! Unlike HLL and Count-Min, KLL never hashes its inputs — it orders raw numeric //! values with `total_cmp`. So the hash-spec metadata group (profile id, //! algorithm, inlined `seed_list`, …) that HLL/CMS carry does not apply, and KLL -//! metadata is **structural params only** (`k`, `m`, `item_type`). See the wire -//! doc's KLL section for the recorded decision. +//! metadata is **structural params only** (`k`, `m`, `item_type`, and an optional +//! `seed`). See the wire doc's KLL section for the recorded decision. //! //! ## Payload item order (cross-language contract) //! @@ -63,6 +63,13 @@ impl KllWireItem for i64 { /// must mirror it). KLL does not hash, so there is **no hash-spec group**: the /// fields are structural params only. `deny_unknown_fields` makes decode fail /// closed on any unexpected key. +/// +/// `seed` is the **only optional** field: it is the compaction RNG's reproducible +/// seed (see [`crate::sketches::kll::KLL::init_with_seed`]). It is present only +/// when the sketch carries one (`Some`), and the key is **omitted** otherwise +/// (`skip_serializing_if` + `default`), so an unseeded sketch's bytes are +/// unchanged. It is echoed on decode, never pinned. The compact `KLL` populates +/// it; `KLLDynamic` has no seed and always omits it. #[derive(Debug, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub(crate) struct KllMetadata { @@ -70,15 +77,20 @@ pub(crate) struct KllMetadata { pub(crate) k: u32, pub(crate) m: u32, pub(crate) item_type: String, + #[serde(skip_serializing_if = "Option::is_none", default)] + pub(crate) seed: Option, } -/// Builds the KLL descriptor metadata for a wire-eligible item type `T`. -pub(crate) fn kll_metadata(k: u32, m: u32) -> KllMetadata { +/// Builds the KLL descriptor metadata for a wire-eligible item type `T`. `seed` +/// is the sketch's reproducible compaction seed when it has one (`Some`), +/// otherwise `None` (the key is then omitted from the wire). +pub(crate) fn kll_metadata(k: u32, m: u32, seed: Option) -> KllMetadata { KllMetadata { metadata_version: 1, k, m, item_type: T::ITEM_TYPE.to_string(), + seed, } } @@ -187,7 +199,7 @@ pub(crate) fn split_and_validate_meta<'a, T: KllWireItem>( ))); } let meta: KllMetadata = from_slice(metadata)?; - if meta != kll_metadata::(meta.k, meta.m) { + if meta != kll_metadata::(meta.k, meta.m, meta.seed) { return Err(RmpDecodeError::Uncategorized( "ASAPv1 KLL envelope: metadata mismatch".to_string(), )); @@ -219,7 +231,8 @@ where /// (kind_id `0x06 0x00`). The retained samples use the top-most-level-first /// layout that matches `sketchlib-go`'s `KLLState`. pub fn serialize_to_bytes(&self) -> Result, RmpEncodeError> { - let metadata = rmp_serde::to_vec_named(&kll_metadata::(self.k as u32, self.m as u32))?; + let metadata = + rmp_serde::to_vec_named(&kll_metadata::(self.k as u32, self.m as u32, self.seed))?; let (state, bit_cache, remaining_bits) = self.co.to_wire(); let payload = rmp_serde::to_vec(&KllPayload { levels: self.wire_levels(), @@ -240,7 +253,13 @@ where let (meta, payload_bytes) = split_and_validate_meta::(bytes, KLL_KIND_COMPACT)?; let payload: KllPayload = from_slice(payload_bytes)?; let num_levels = validate_kll_payload(&payload.levels, &payload.items, &payload.coin)?; - Self::from_wire_top_first(meta.k as usize, meta.m as usize, num_levels, payload) + Self::from_wire_top_first( + meta.k as usize, + meta.m as usize, + meta.seed, + num_levels, + payload, + ) } /// Rebuilds the compact leftward-grown buffer from the top-most-first wire @@ -251,6 +270,7 @@ where fn from_wire_top_first( k: usize, m: usize, + seed: Option, num_levels: usize, payload: KllPayload, ) -> Result { @@ -306,11 +326,10 @@ where num_levels, max_capacity: max_cap, co: Coin::from_wire(coin.state, coin.bit_cache, coin.remaining_bits as u8), - // The wire form does not carry the explicit-seed flag; a - // round-tripped sketch keeps its coin state but re-randomizes on a - // later clear(). Callers needing clear()-determinism rebuild via - // init_with_seed. - seed: None, + // The reproducible compaction seed is restored from the metadata (it + // is `None` for a sketch that never carried one), so a decoded sketch + // keeps clear()-determinism when the producer had it. + seed, capacity_cache: [0; CAPACITY_CACHE_LEN], top_height: 0, level0_capacity: 0, @@ -423,7 +442,7 @@ mod tests { #[test] fn kll_rejects_inconsistent_levels() { // Valid envelope + metadata, but levels[last] != items.len(). - let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8)).unwrap(); + let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8, None)).unwrap(); let payload = rmp_serde::to_vec(&KllPayload:: { levels: vec![0, 3], items: vec![1.0, 2.0], // len 2, but levels claim 3 @@ -463,7 +482,7 @@ mod tests { (MAX_CACHEABLE_K as u32 + 1, 8), (200, 1), ] { - let metadata = rmp_serde::to_vec_named(&kll_metadata::(k, m)).unwrap(); + let metadata = rmp_serde::to_vec_named(&kll_metadata::(k, m, None)).unwrap(); let bytes = envelope::encode(KLL_KIND_COMPACT, &metadata, &empty_payload()); assert!( KLL::::deserialize_from_bytes(&bytes).is_err(), @@ -472,6 +491,49 @@ mod tests { } } + #[test] + fn kll_seed_present_when_seeded_omitted_when_unseeded() { + // A seeded sketch records its seed in the metadata; an unseeded one omits + // the key entirely (so its bytes are unaffected by the optional field). + let seeded = KLL::::init_kll_with_seed(200, 42); + let bytes = seeded.serialize_to_bytes().expect("serialize"); + let (_k, meta_bytes, _p) = envelope::split(&bytes).expect("split"); + let meta: KllMetadata = rmp_serde::from_slice(meta_bytes).expect("meta"); + assert_eq!(meta.seed, Some(42), "seeded KLL must record its seed"); + + let mut unseeded = KLL::::init_kll(200); // Coin::new(); seed = None + unseeded.update(&1.0); + let bytes = unseeded.serialize_to_bytes().expect("serialize"); + let (_k, meta_bytes, _p) = envelope::split(&bytes).expect("split"); + let meta: KllMetadata = rmp_serde::from_slice(meta_bytes).expect("meta"); + assert_eq!(meta.seed, None, "unseeded KLL must omit the seed key"); + } + + #[test] + fn kll_seed_survives_round_trip_so_clear_stays_deterministic() { + // The point of carrying seed: a decoded seeded sketch must re-seed clear() + // from the original seed (not wall-clock). If seed were dropped on decode, + // `a` below would re-randomize on clear() and diverge from `b`. + let mut src = KLL::::init_kll_with_seed(200, 42); + for v in 1..=5000u64 { + src.update(&(v as f64)); + } + let mut a = KLL::::deserialize_from_bytes(&src.serialize_to_bytes().unwrap()).unwrap(); + let mut b = KLL::::init_kll_with_seed(200, 42); + + a.clear(); + b.clear(); + for v in 1..=3000u64 { + a.update(&(v as f64)); + b.update(&(v as f64)); + } + assert_eq!( + a.serialize_to_bytes().unwrap(), + b.serialize_to_bytes().unwrap(), + "decoded sketch lost its seed: clear() diverged from a fresh seeded sketch" + ); + } + #[test] fn kll_rejects_weighted_count_overflow() { // Structurally valid (monotonic levels, levels[last]==items.len()) but @@ -483,7 +545,7 @@ mod tests { // items, every lower level is empty. let mut levels = vec![16u32; num_levels + 1]; levels[0] = 0; - let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8)).unwrap(); + let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8, None)).unwrap(); let payload = rmp_serde::to_vec(&KllPayload:: { levels, items: vec![1.0; 16], @@ -504,7 +566,7 @@ mod tests { #[test] fn kll_dynamic_kind_id_rejected_by_compact() { // A compact decoder must reject the dynamic kind_id. - let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8)).unwrap(); + let metadata = rmp_serde::to_vec_named(&kll_metadata::(200, 8, None)).unwrap(); let payload = rmp_serde::to_vec(&KllPayload:: { levels: vec![0, 0], items: Vec::::new(), diff --git a/src/sketches/kll_dynamic/wire.rs b/src/sketches/kll_dynamic/wire.rs index 28f5740..ec79740 100644 --- a/src/sketches/kll_dynamic/wire.rs +++ b/src/sketches/kll_dynamic/wire.rs @@ -37,7 +37,10 @@ where /// top-most-level-first / input-order-L0 layout, so `levels` and `items` are /// copied straight out. pub fn serialize_to_bytes(&self) -> Result, RmpEncodeError> { - let metadata = rmp_serde::to_vec_named(&kll_metadata::(self.k as u32, self.m as u32))?; + // `KLLDynamic` has no reproducible-seed concept, so `seed` is always + // `None` and the metadata key is omitted. + let metadata = + rmp_serde::to_vec_named(&kll_metadata::(self.k as u32, self.m as u32, None))?; let (state, bit_cache, remaining_bits) = self.co.to_wire(); let payload = rmp_serde::to_vec(&KllPayload { levels: self.levels.as_slice().iter().map(|&l| l as u32).collect(), From c29dd109a3d5535c27c03a41dc1be8b1a6ee16c2 Mon Sep 17 00:00:00 2001 From: GordonYuanyc Date: Mon, 20 Jul 2026 02:05:13 -0400 Subject: [PATCH 4/4] test(kll): assert unseeded compact KLL decodes and round-trips byte-stable Closes a coverage gap from the third review: existing positive round-trip tests all used init_kll_with_seed, so the seed-absent metadata path was exercised on encode but never fully decoded and re-serialized. Add an unseeded decode + byte-stable re-serialize assertion. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/sketches/kll/wire.rs | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/sketches/kll/wire.rs b/src/sketches/kll/wire.rs index b2d8924..d8b2dea 100644 --- a/src/sketches/kll/wire.rs +++ b/src/sketches/kll/wire.rs @@ -509,6 +509,25 @@ mod tests { assert_eq!(meta.seed, None, "unseeded KLL must omit the seed key"); } + #[test] + fn kll_unseeded_round_trip_byte_stable() { + // An unseeded compact KLL (seed = None, so the metadata key is omitted) + // must decode end-to-end and re-serialize byte-identically — exercises the + // optional-key-absent path through decode, not just encode. + let mut s = KLL::::init_kll(200); // Coin::new(); seed = None + for v in 1..=2000u64 { + s.update(&(v as f64)); + } + let bytes = s.serialize_to_bytes().expect("serialize"); + let decoded = KLL::::deserialize_from_bytes(&bytes).expect("decode"); + assert_eq!( + decoded.serialize_to_bytes().expect("re-serialize"), + bytes, + "unseeded KLL round trip not byte-stable" + ); + assert_eq!(decoded.count(), s.count()); + } + #[test] fn kll_seed_survives_round_trip_so_clear_stays_deterministic() { // The point of carrying seed: a decoded seeded sketch must re-seed clear()