From f15ec54dd78614dc272499008567d4cab4926eeb Mon Sep 17 00:00:00 2001 From: Patrick McGleenon Date: Sat, 30 May 2026 18:56:49 +0100 Subject: [PATCH 1/2] feat(req): port REQ sketch from reqsketch-rs --- Cargo.lock | 187 +++- Cargo.toml | 5 +- datasketches/Cargo.toml | 11 + datasketches/src/codec/family.rs | 9 + datasketches/src/codec/mod.rs | 2 + datasketches/src/lib.rs | 2 + datasketches/src/req/compactor.rs | 604 +++++++++++++ datasketches/src/req/iter.rs | 86 ++ datasketches/src/req/mod.rs | 66 ++ datasketches/src/req/serialization.rs | 47 + datasketches/src/req/sketch.rs | 833 ++++++++++++++++++ datasketches/src/req/sorted_view.rs | 372 ++++++++ datasketches/src/req/union.rs | 95 ++ datasketches/src/req/value.rs | 191 ++++ datasketches/tests/req_test/accuracy.rs | 54 ++ datasketches/tests/req_test/bounds.rs | 171 ++++ datasketches/tests/req_test/core.rs | 243 +++++ datasketches/tests/req_test/main.rs | 54 ++ datasketches/tests/req_test/merge.rs | 141 +++ datasketches/tests/req_test/property.rs | 113 +++ datasketches/tests/req_test/query.rs | 194 ++++ datasketches/tests/req_test/serialization.rs | 374 ++++++++ .../tests/req_test/sorted_view_api.rs | 165 ++++ datasketches/tests/req_test/structure.rs | 77 ++ datasketches/tests/req_test/union.rs | 116 +++ 25 files changed, 4210 insertions(+), 2 deletions(-) create mode 100644 datasketches/src/req/compactor.rs create mode 100644 datasketches/src/req/iter.rs create mode 100644 datasketches/src/req/mod.rs create mode 100644 datasketches/src/req/serialization.rs create mode 100644 datasketches/src/req/sketch.rs create mode 100644 datasketches/src/req/sorted_view.rs create mode 100644 datasketches/src/req/union.rs create mode 100644 datasketches/src/req/value.rs create mode 100644 datasketches/tests/req_test/accuracy.rs create mode 100644 datasketches/tests/req_test/bounds.rs create mode 100644 datasketches/tests/req_test/core.rs create mode 100644 datasketches/tests/req_test/main.rs create mode 100644 datasketches/tests/req_test/merge.rs create mode 100644 datasketches/tests/req_test/property.rs create mode 100644 datasketches/tests/req_test/query.rs create mode 100644 datasketches/tests/req_test/serialization.rs create mode 100644 datasketches/tests/req_test/sorted_view_api.rs create mode 100644 datasketches/tests/req_test/structure.rs create mode 100644 datasketches/tests/req_test/union.rs diff --git a/Cargo.lock b/Cargo.lock index d18361b5..1eb3735e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -67,6 +67,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + [[package]] name = "autocfg" version = "1.5.1" @@ -79,6 +88,21 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + [[package]] name = "bitflags" version = "2.13.1" @@ -206,8 +230,11 @@ dependencies = [ name = "datasketches" version = "0.4.0" dependencies = [ + "approx", "googletest", "insta", + "proptest", + "rand", ] [[package]] @@ -258,6 +285,12 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "getrandom" version = "0.2.17" @@ -269,6 +302,18 @@ dependencies = [ "wasi", ] +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + [[package]] name = "getrandom" version = "0.4.3" @@ -277,7 +322,7 @@ checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", "libc", - "r-efi", + "r-efi 6.0.0", ] [[package]] @@ -410,6 +455,15 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -419,6 +473,31 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "proptest" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" +dependencies = [ + "bit-set", + "bit-vec", + "bitflags", + "num-traits", + "rand", + "rand_chacha", + "rand_xorshift", + "regex-syntax", + "rusty-fork", + "tempfile", + "unarray", +] + +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + [[package]] name = "quote" version = "1.0.47" @@ -428,12 +507,56 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_xorshift" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" +dependencies = [ + "rand_core", +] + [[package]] name = "regex" version = "1.13.1" @@ -531,6 +654,18 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" +[[package]] +name = "rusty-fork" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2" +dependencies = [ + "fnv", + "quick-error", + "tempfile", + "wait-timeout", +] + [[package]] name = "semver" version = "1.0.28" @@ -690,6 +825,12 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "unarray" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + [[package]] name = "unicode-ident" version = "1.0.24" @@ -743,12 +884,30 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "wait-timeout" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11" +dependencies = [ + "libc", +] + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasip2" +version = "1.0.1+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "webpki-roots" version = "1.0.9" @@ -877,6 +1036,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "wit-bindgen" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" + [[package]] name = "x" version = "0.0.0" @@ -889,6 +1054,26 @@ dependencies = [ "which", ] +[[package]] +name = "zerocopy" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "zeroize" version = "1.9.0" diff --git a/Cargo.toml b/Cargo.toml index 04f944e6..863e677c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,10 +32,13 @@ rust-version = "1.86.0" datasketches = { path = "datasketches" } # Crates.io dependencies +approx = { version = "0.5" } clap = { version = "4.6.5", features = ["derive"] } -insta = { version = "1.48.0" } googletest = { version = "0.14.3" } +insta = { version = "1.48.0" } cargo_metadata = { version = "0.23.1" } +proptest = { version = "1" } +rand = { version = "0.9.2" } which = { version = "8.0.5" } flate2 = { version = "1.1.9", default-features = false, features = ["zlib-rs"] } tar = { version = "0.4.46", default-features = false } diff --git a/datasketches/Cargo.toml b/datasketches/Cargo.toml index 169ee7e1..33610a65 100644 --- a/datasketches/Cargo.toml +++ b/datasketches/Cargo.toml @@ -43,6 +43,7 @@ countmin = [] cpc = [] frequencies = [] hll = [] +req = ["dep:rand"] tdigest = [] theta = [] tuple = [] @@ -72,6 +73,11 @@ name = "hll_test" path = "tests/hll_test/main.rs" required-features = ["hll"] +[[test]] +name = "req_test" +path = "tests/req_test/main.rs" +required-features = ["req"] + [[test]] name = "tdigest_test" path = "tests/tdigest_test/main.rs" @@ -87,9 +93,14 @@ name = "tuple_test" path = "tests/tuple_test/main.rs" required-features = ["tuple"] +[dependencies] +rand = { workspace = true, optional = true } + [dev-dependencies] +approx = { workspace = true } googletest = { workspace = true } insta = { workspace = true } +proptest = { workspace = true } [lints] workspace = true diff --git a/datasketches/src/codec/family.rs b/datasketches/src/codec/family.rs index c6ab908c..2e3f4193 100644 --- a/datasketches/src/codec/family.rs +++ b/datasketches/src/codec/family.rs @@ -80,6 +80,15 @@ impl Family { max_pre_longs: 5, }; + /// Relative Error Quantiles (REQ) sketch. + #[cfg(feature = "req")] + pub const REQ: Family = Family { + id: 17, + name: "REQ", + min_pre_longs: 2, + max_pre_longs: 4, + }; + /// CountMin Sketch #[cfg(feature = "countmin")] pub const COUNTMIN: Family = Family { diff --git a/datasketches/src/codec/mod.rs b/datasketches/src/codec/mod.rs index 30d10753..bd466882 100644 --- a/datasketches/src/codec/mod.rs +++ b/datasketches/src/codec/mod.rs @@ -28,6 +28,7 @@ pub use self::encode::SketchBytes; feature = "cpc", feature = "frequencies", feature = "hll", + feature = "req", feature = "tdigest", feature = "theta", feature = "tuple", @@ -41,6 +42,7 @@ pub(crate) mod assert; feature = "cpc", feature = "frequencies", feature = "hll", + feature = "req", feature = "tdigest", feature = "theta", feature = "tuple", diff --git a/datasketches/src/lib.rs b/datasketches/src/lib.rs index e169c4c7..066b2376 100644 --- a/datasketches/src/lib.rs +++ b/datasketches/src/lib.rs @@ -41,6 +41,8 @@ pub mod cpc; pub mod frequencies; #[cfg(feature = "hll")] pub mod hll; +#[cfg(feature = "req")] +pub mod req; #[cfg(feature = "tdigest")] pub mod tdigest; #[cfg(any(feature = "theta", feature = "tuple"))] diff --git a/datasketches/src/req/compactor.rs b/datasketches/src/req/compactor.rs new file mode 100644 index 00000000..55e68264 --- /dev/null +++ b/datasketches/src/req/compactor.rs @@ -0,0 +1,604 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Compactor implementation for REQ sketch levels. +//! +//! Each level in the REQ sketch uses a compactor to maintain a bounded set of items +//! with deterministic compaction when capacity is exceeded. + +use super::RankAccuracy; +use super::value::ReqValue; +use crate::error::Error; + +fn nearest_even(value: f32) -> u32 { + ((value / 2.0).round() as u32) << 1 +} + +/// A compactor maintains items at a specific level of the REQ sketch. +/// +/// When the compactor reaches its nominal capacity, it performs compaction +/// by keeping approximately half the items and promoting the rest to the next level. +#[derive(Debug, Clone)] +pub(super) struct Compactor { + /// Current items in the compactor + items: Vec, + /// Whether items are currently sorted + is_sorted: bool, + /// State for deterministic compaction + state: u64, + /// Reusable scratch buffer for compaction operations + scratch_buffer: Vec, + + /// Actual section size (rounded to integer) + section_size: u32, + /// Number of sections in this compactor + num_sections: u8, + /// The level of this compactor (0 = base level) + lg_weight: u8, + + /// Whether this compactor is configured for high rank accuracy + rank_accuracy: RankAccuracy, + /// Raw section size (may be fractional) + section_size_raw: f32, + /// Random bit for compaction + coin: bool, +} + +impl Compactor +where + T: Clone + ReqValue, +{ + /// Creates a new compactor for the given level. + /// + /// # Arguments + /// * `lg_weight` - The level (log weight) of this compactor + /// * `k` - The k parameter from the parent sketch + /// * `rank_accuracy` - Rank accuracy configuration + pub(super) fn new(lg_weight: u8, k: u16, rank_accuracy: RankAccuracy) -> Self { + let section_size_raw = k as f32; + let section_size = nearest_even(section_size_raw); + let num_sections = 3u8; + + let nominal: usize = (2 * section_size * num_sections as u32) as usize; + + Self { + items: Vec::with_capacity(nominal), + is_sorted: true, + state: 0, + scratch_buffer: Vec::with_capacity(nominal / 2 + 8), + + section_size, + num_sections, + lg_weight, + + rank_accuracy, + section_size_raw, + coin: false, + } + } + + /// Returns the number of items currently in this compactor. + pub(super) fn num_items(&self) -> u32 { + self.items.len() as u32 + } + + /// Returns the nominal capacity of this compactor. + pub(super) fn nominal_capacity(&self) -> u32 { + 2 * self.section_size * self.num_sections as u32 + } + + /// Returns whether the items are currently sorted. + pub(super) fn is_sorted(&self) -> bool { + self.is_sorted + } + + /// Appends an item to this compactor. + #[inline(always)] + pub(super) fn append(&mut self, item: T) { + self.items.push(item); + if self.items.len() > 1 { + self.is_sorted = false; + } + } + + /// Merges items from another compactor into this one. + pub(super) fn merge(&mut self, other: &Self) { + self.state |= other.state; + self.items.extend_from_slice(&other.items); + if !other.items.is_empty() { + self.is_sorted = false; + } + // OR-ing the schedule counters can advance state past several doubling + // thresholds at once. Loop until no more doublings are needed (C++: + // req_compactor_impl.hpp:250 — `while (ensure_enough_sections()) {}`). + while self.ensure_enough_sections() {} + } + + /// Counts the items at-or-below (`inclusive`) or strictly below `item`. + /// + /// Uses binary search when this compactor is sorted, and a linear scan + /// otherwise. This lets [`ReqSketch::rank`](super::ReqSketch::rank) sum + /// per-level weights directly without first building a sorted view. + pub(super) fn count_below(&self, item: &T, inclusive: bool) -> usize { + if self.is_sorted { + if inclusive { + self.items.partition_point(|x| x.total_cmp(item).is_le()) + } else { + self.items.partition_point(|x| x.total_cmp(item).is_lt()) + } + } else { + self.items + .iter() + .filter(|x| { + let ord = x.total_cmp(item); + if inclusive { ord.is_le() } else { ord.is_lt() } + }) + .count() + } + } + + /// Merges pre-sorted items into this compactor. + /// Merges sorted items into this compactor using scratch buffer to avoid allocation. + /// Both this compactor's items and the input must be sorted. + #[inline(always)] + pub(super) fn merge_sorted(&mut self, items: &[T]) { + if items.is_empty() { + return; + } + + if self.items.is_empty() { + self.items.extend_from_slice(items); + self.is_sorted = true; + return; + } + + // Ensure sorted on both inputs by contract + let total = self.items.len() + items.len(); + self.scratch_buffer.clear(); + if self.scratch_buffer.capacity() < total { + self.scratch_buffer + .reserve(total - self.scratch_buffer.capacity()); + } + + let (mut i, mut j) = (0usize, 0usize); + let (a, b) = (&self.items, items); + + // Two-pointer merge into scratch buffer + while i < a.len() && j < b.len() { + if a[i].total_cmp(&b[j]).is_le() { + self.scratch_buffer.push(a[i].clone()); + i += 1; + } else { + self.scratch_buffer.push(b[j].clone()); + j += 1; + } + } + + // Add remaining elements + if i < a.len() { + self.scratch_buffer.extend_from_slice(&a[i..]); + } + if j < b.len() { + self.scratch_buffer.extend_from_slice(&b[j..]); + } + + // Swap scratch buffer with items (zero-copy) + self.items.clear(); + std::mem::swap(&mut self.items, &mut self.scratch_buffer); + self.is_sorted = true; + } + + /// Sorts the items in this compactor if not already sorted. + #[inline(always)] + pub(super) fn sort(&mut self) { + if !self.is_sorted { + // Use unstable sort for better performance (stable not needed for REQ sketch) + self.items.sort_unstable_by(|a, b| a.total_cmp(b)); + self.is_sorted = true; + } + } + + /// Compacts into the provided output buffer without allocating. + /// Writes promoted items into `out` and removes the compacted range in-place via `copy_within + + /// truncate`. + #[inline(always)] + pub(super) fn compact_into(&mut self, _rank_accuracy: RankAccuracy, out: &mut Vec) { + if self.items.is_empty() { + out.clear(); + return; + } + + // Sort entire buffer (C++ sorts full buffer before compaction) + self.sort(); + + // Calculate sections to compact based on state + let secs_to_compact = + ((!self.state).trailing_zeros() + 1).min(self.num_sections as u32) as u8; + let compaction_range = self.compute_compaction_range(secs_to_compact); + + // Must have at least 2 items to compact + if compaction_range.1 <= compaction_range.0 || (compaction_range.1 - compaction_range.0) < 2 + { + out.clear(); + return; + } + + if (self.state & 1) == 1 { + self.coin = !self.coin; // flip coin for odd states + } else { + self.coin = rand::random::(); // random coin flip for even states + } + let odds = self.coin; + + // Build promoted items directly into output buffer (no alloc) + out.clear(); + let (start, end) = compaction_range; + let mut i = start + if odds { 1 } else { 0 }; + while i < end { + out.push(self.items[i].clone()); // TODO: use Copy fast-path for numeric types + i += 2; + } + + // Remove the compacted range in-place by rotating elements left + let removed = end - start; + if end < self.items.len() { + // Use rotate_left to move tail elements to fill the gap + self.items[start..].rotate_left(removed); + } + self.items.truncate(self.items.len() - removed); + + // Update state, then ensure enough sections (C++ order) + self.state += 1; + self.ensure_enough_sections(); + } + + /// Returns an iterator over the items in this compactor. + pub(super) fn iter(&self) -> impl Iterator { + self.items.iter() + } + + /// Returns a slice of items for zero-allocation iteration. + pub(super) fn items_slice(&self) -> &[T] { + &self.items + } + + /// Returns the weight (2^lg_weight) for items in this compactor. + pub(super) fn weight(&self) -> u64 { + 1u64 << self.lg_weight + } + + // Private helper methods + + fn ensure_enough_sections(&mut self) -> bool { + let ssr = self.section_size_raw / (2.0_f32).sqrt(); + let ne = nearest_even(ssr); + + const MIN_K: u32 = 4; // matches datasketches-cpp + + if self.num_sections <= 64 && self.state >= (1u64 << (self.num_sections - 1)) && ne >= MIN_K + { + self.section_size_raw = ssr; + self.section_size = ne; + self.num_sections <<= 1; // Double the sections + true + } else { + false + } + } + + #[inline(always)] + fn compute_compaction_range(&self, secs_to_compact: u8) -> (usize, usize) { + let nom_capacity = self.nominal_capacity() as usize; + let mut non_compact = nom_capacity / 2 + + (self.num_sections - secs_to_compact) as usize * self.section_size as usize; + + // if (((num_items_ - non_compact) & 1) == 1) ++non_compact; + if self.items.len() >= non_compact && ((self.items.len() - non_compact) & 1) == 1 { + non_compact += 1; + } + + let (low, high) = match self.rank_accuracy { + RankAccuracy::HighRank => { + // HRA: Protect high ranks by compacting LOW sections (low values) + // This means we compact from [0, num_items - non_compact] (bottom end) + let high = if self.items.len() >= non_compact { + self.items.len() - non_compact + } else { + 0 + }; + (0, high) + } + RankAccuracy::LowRank => { + // LRA: Protect low ranks by compacting HIGH sections (high values) + // This means we compact from [non_compact, num_items] (top end) + let low = non_compact.min(self.items.len()); + (low, self.items.len()) + } + }; + + // Empty window safety: ensure we have at least 2 items to compact + if high <= low || (high - low) < 2 { + return (0, 0); // Signal no compaction needed + } + + (low, high) + } + + /// Serialize this compactor (preamble + items) into the byte buffer. + pub(super) fn serialize_into(&self, bytes: &mut crate::codec::SketchBytes) { + bytes.write_u64_le(self.state); + bytes.write_f32_le(self.section_size_raw); + bytes.write_u8(self.lg_weight); + bytes.write_u8(self.num_sections); + bytes.write_u16_le(0); // padding + bytes.write_u32_le(self.num_items()); + for item in self.iter() { + item.serialize_value(bytes); + } + } + + /// Deserialize a compactor (preamble + items) from the byte cursor. + pub(super) fn deserialize( + cursor: &mut crate::codec::SketchSlice<'_>, + rank_accuracy: super::RankAccuracy, + is_level_zero_sorted: bool, + ) -> Result { + use crate::codec::assert::insufficient_data; + let state = cursor + .read_u64_le() + .map_err(insufficient_data("compactor.state"))?; + let section_size_raw = cursor + .read_f32_le() + .map_err(insufficient_data("compactor.section_size_raw"))?; + let lg_weight = cursor + .read_u8() + .map_err(insufficient_data("compactor.lg_weight"))?; + let num_sections = cursor + .read_u8() + .map_err(insufficient_data("compactor.num_sections"))?; + let _padding = cursor + .read_u16_le() + .map_err(insufficient_data("compactor.padding"))?; + let num_items = cursor + .read_u32_le() + .map_err(insufficient_data("compactor.num_items"))?; + + // Validate the wire-controlled fields before they feed capacity/weight + // arithmetic. A legitimate compactor always satisfies these bounds + // (`section_size` derives from k ≤ MAX_K and only shrinks; `lg_weight` is the + // level index), so rejecting anything else keeps `nominal_capacity` and + // `weight` from overflowing on crafted input. + if !(0.0..=super::MAX_K as f32).contains(§ion_size_raw) { + return Err(Error::invalid_argument(format!( + "REQ compactor section_size {section_size_raw} out of range" + ))); + } + // `weight()` computes `1u64 << lg_weight`, which overflows once lg_weight ≥ 64. + if lg_weight >= 64 { + return Err(Error::invalid_argument(format!( + "REQ compactor lg_weight {lg_weight} exceeds maximum" + ))); + } + + // Don't trust `num_items` for the allocation: a malformed length could request + // a multi-gigabyte reservation before the per-item reads below fail. The buffer + // holds at most `remaining` more items (each item is ≥ 1 byte), so cap the + // pre-allocation there; `push` still grows the Vec as the validated data needs. + let capacity = (num_items as usize).min(cursor.remaining().len()); + let mut items = Vec::with_capacity(capacity); + for _ in 0..num_items { + items.push(T::deserialize_value(cursor)?); + } + + Ok(Compactor::from_serialized_state( + lg_weight, + section_size_raw, + num_sections, + state, + items, + is_level_zero_sorted, + rank_accuracy, + )) + } + + /// Build a level-0 compactor from raw items (used by the `RAW_ITEMS` deserialize path). + /// + /// The wire format omits the compactor preamble for tiny sketches (n ≤ 4); this + /// helper synthesises a fresh compactor and seeds it with the deserialized items. + /// `is_sorted` is taken verbatim from the wire — both C++ and Java produce valid + /// sketches but disagree on this flag for n=1 (C++ sets true, Java sets false), + /// so faithful round-trip requires preserving whatever the input said. + pub(super) fn raw_items_compactor( + k: u16, + rank_accuracy: super::RankAccuracy, + items: Vec, + is_sorted: bool, + ) -> Self { + let mut c = Self::new(0, k, rank_accuracy); + for item in items { + c.append(item); + } + // append() may have flipped is_sorted off; restore the wire flag verbatim. + c.is_sorted = is_sorted; + c + } + + /// Reconstruct a Compactor from deserialized state. + /// + /// Used by [`Compactor::deserialize`]. Transient state (random coin, scratch + /// buffer) is reset; the deterministic `state` counter and the persistent + /// configuration (`lg_weight`, `section_size_raw`, `num_sections`) are preserved + /// from the wire data. + pub(super) fn from_serialized_state( + lg_weight: u8, + section_size_raw: f32, + num_sections: u8, + state: u64, + items: Vec, + is_sorted: bool, + rank_accuracy: super::RankAccuracy, + ) -> Self { + Self { + items, + is_sorted, + state, + scratch_buffer: Vec::new(), + section_size: nearest_even(section_size_raw), + num_sections, + lg_weight, + rank_accuracy, + section_size_raw, + coin: false, + } + } +} + +#[cfg(test)] +impl Compactor +where + T: Clone + ReqValue, +{ + /// Returns the level (log weight) of this compactor. Test-only accessor. + pub(super) fn lg_weight(&self) -> u8 { + self.lg_weight + } + + /// Returns the current state for deterministic compaction. Test-only accessor. + pub(super) fn state(&self) -> u64 { + self.state + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_compactor() { + let compactor: Compactor = Compactor::new(0, 12, RankAccuracy::HighRank); + assert_eq!(compactor.lg_weight(), 0); + assert_eq!(compactor.num_items(), 0); + assert!(compactor.is_sorted()); + assert_eq!(compactor.weight(), 1); + } + + #[test] + fn test_append_and_sort() { + let mut compactor = Compactor::new(0, 12, RankAccuracy::HighRank); + + compactor.append(5); + assert_eq!(compactor.num_items(), 1); + assert!(compactor.is_sorted()); // Single item is sorted + + compactor.append(3); + assert_eq!(compactor.num_items(), 2); + assert!(!compactor.is_sorted()); // Multiple items, not sorted + + compactor.sort(); + assert!(compactor.is_sorted()); + + let items: Vec<&i32> = compactor.iter().collect(); + assert_eq!(items, vec![&3, &5]); + } + + #[test] + fn test_nearest_even() { + assert_eq!(nearest_even(0.0), 0); // 0/2=0, round(0)=0, 0<<1=0 + assert_eq!(nearest_even(1.0), 2); // 1/2=0.5, round(0.5)=1, 1<<1=2 + assert_eq!(nearest_even(2.0), 2); // 2/2=1, round(1)=1, 1<<1=2 + assert_eq!(nearest_even(3.0), 4); // 3/2=1.5, round(1.5)=2, 2<<1=4 + assert_eq!(nearest_even(4.0), 4); // 4/2=2, round(2)=2, 2<<1=4 + assert_eq!(nearest_even(4.6), 4); // 4.6/2=2.3, round(2.3)=2, 2<<1=4 + assert_eq!(nearest_even(5.6), 6); // 5.6/2=2.8, round(2.8)=3, 3<<1=6 + assert_eq!(nearest_even(13.0), 14); // 13/2=6.5, round(6.5)=7, 7<<1=14 + } + + #[test] + fn test_merge_sorted() { + let mut compactor = Compactor::new(0, 12, RankAccuracy::HighRank); + + compactor.append(1); + compactor.append(3); + compactor.append(5); + compactor.sort(); + + let other_items = vec![2, 4, 6]; + compactor.merge_sorted(&other_items); + + assert!(compactor.is_sorted()); + let items: Vec<&i32> = compactor.iter().collect(); + assert_eq!(items, vec![&1, &2, &3, &4, &5, &6]); + } + + #[test] + fn compactor_serialization_round_trip() { + use crate::codec::SketchBytes; + use crate::codec::SketchSlice; + + let mut c: Compactor = Compactor::new(0, 12, RankAccuracy::HighRank); + for i in 0..30 { + c.append(i as f32); + } + c.sort(); + + let mut bytes = SketchBytes::with_capacity(256); + c.serialize_into(&mut bytes); + let raw = bytes.into_bytes(); + + let mut cursor = SketchSlice::new(&raw); + let c2 = Compactor::::deserialize(&mut cursor, RankAccuracy::HighRank, true).unwrap(); + + assert_eq!(c.num_items(), c2.num_items()); + assert_eq!(c.lg_weight(), c2.lg_weight()); + assert_eq!(c.state(), c2.state()); + let xs: Vec = c.iter().copied().collect(); + let ys: Vec = c2.iter().copied().collect(); + assert_eq!(xs, ys); + } + + #[test] + fn merge_loops_ensure_enough_sections_for_high_state() { + // Regression test for the bug where Compactor::merge called + // ensure_enough_sections() once instead of looping. Without the loop, + // num_sections doubles at most once per merge — but OR-ing a high state + // can advance past several doubling thresholds at once and require + // multiple doublings (matching the C++ reference at + // req_compactor_impl.hpp:250 — `while (ensure_enough_sections()) {}`). + // + // Setup: a fresh compactor (state=0, num_sections=3) merged with another + // whose state is 0xFFFF. After merge, state |= 0xFFFF = 0xFFFF. + // ensure_enough_sections doublings (k=12, section_size_raw=12): + // - state=0xFFFF >= (1<<2)=4 ✓ → num_sections=6, ssr≈8.49 + // - state=0xFFFF >= (1<<5)=32 ✓ → num_sections=12, ssr≈6.00 + // - state=0xFFFF >= (1<<11)=2048 ✓ → num_sections=24, ssr≈4.24 + // - state=0xFFFF >= (1<<23)=8388608 ✗ → stop + // Expected: num_sections == 24 with the fix; == 6 with only one call. + let mut a: Compactor = Compactor::new(0, 12, RankAccuracy::HighRank); + let mut b: Compactor = Compactor::new(0, 12, RankAccuracy::HighRank); + b.state = 0xFFFF; + + assert_eq!(a.num_sections, 3, "default num_sections sanity"); + + a.merge(&b); + + assert!( + a.num_sections >= 12, + "merge must loop ensure_enough_sections; got num_sections={} (single-call would yield 6)", + a.num_sections + ); + } +} diff --git a/datasketches/src/req/iter.rs b/datasketches/src/req/iter.rs new file mode 100644 index 00000000..a3c98630 --- /dev/null +++ b/datasketches/src/req/iter.rs @@ -0,0 +1,86 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Iterator implementations for REQ sketch inspection. + +use super::compactor::Compactor; +use super::value::ReqValue; + +/// Iterator over (item, weight) pairs in a REQ sketch. +/// +/// Provides access to all items in the sketch along with their weights, +/// which depend on the level of the compactor they're stored in. +/// +/// Zero-allocation implementation that works directly with slices. +pub struct ReqSketchIterator<'a, T: ReqValue> { + compactors: &'a [Compactor], + current_level: usize, + current_level_iter: Option>, + current_weight: u64, +} + +impl<'a, T: ReqValue> ReqSketchIterator<'a, T> { + /// Creates a new iterator over the compactors. + pub(super) fn new(compactors: &'a [Compactor]) -> Self { + let mut iter = Self { + compactors, + current_level: 0, + current_level_iter: None, + current_weight: 0, + }; + iter.advance_to_next_level(); + iter + } + + fn advance_to_next_level(&mut self) { + while self.current_level < self.compactors.len() { + let compactor = &self.compactors[self.current_level]; + // Access items slice directly without allocation + let items_slice = compactor.items_slice(); + + if !items_slice.is_empty() { + self.current_level_iter = Some(items_slice.iter()); + self.current_weight = compactor.weight(); + return; + } + + self.current_level += 1; + } + + self.current_level_iter = None; + } +} + +impl Iterator for ReqSketchIterator<'_, T> { + type Item = (T, u64); + + fn next(&mut self) -> Option { + loop { + if let Some(ref mut level_iter) = self.current_level_iter { + if let Some(item) = level_iter.next() { + return Some((item.clone(), self.current_weight)); + } + } + + // Current level exhausted, move to next + self.current_level += 1; + self.advance_to_next_level(); + + self.current_level_iter.as_ref()?; + } + } +} diff --git a/datasketches/src/req/mod.rs b/datasketches/src/req/mod.rs new file mode 100644 index 00000000..c17b441d --- /dev/null +++ b/datasketches/src/req/mod.rs @@ -0,0 +1,66 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Relative Error Quantiles (REQ) sketch. +//! +//! Provides approximate quantile estimation with relative error guarantees, especially +//! useful for streaming scenarios needing bounded memory. Based on the paper +//! [Relative Error Streaming Quantiles](https://arxiv.org/abs/2004.01668) by Cormode, +//! Karnin, Liberty, Thaler and Veselý, and on the Apache DataSketches C++ reference +//! implementation. + +mod compactor; +mod iter; +mod serialization; +mod sketch; +mod sorted_view; +mod union; +mod value; + +pub use self::iter::ReqSketchIterator; +pub use self::sketch::ReqSketch; +pub use self::sketch::ReqSketchBuilder; +pub use self::sorted_view::SortedView; +pub use self::union::ReqUnion; +pub use self::value::ReqValue; + +/// Default value of `k` if not specified. Roughly 1% relative error at 95% confidence. +pub const DEFAULT_K: u16 = 12; +/// Minimum allowed value of `k`. +pub const MIN_K: u16 = 4; +/// Maximum allowed value of `k`. +pub const MAX_K: u16 = 1024; + +/// Selects which tail of the rank domain the sketch optimizes for. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RankAccuracy { + /// Optimize for accuracy at high ranks (near 1.0). + #[default] + HighRank, + /// Optimize for accuracy at low ranks (near 0.0). + LowRank, +} + +/// Whether queries include the weight of the search item. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum SearchCriteria { + /// Include the weight of the search item in the result. + #[default] + Inclusive, + /// Exclude the weight of the search item from the result. + Exclusive, +} diff --git a/datasketches/src/req/serialization.rs b/datasketches/src/req/serialization.rs new file mode 100644 index 00000000..19061d84 --- /dev/null +++ b/datasketches/src/req/serialization.rs @@ -0,0 +1,47 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! REQ sketch wire format — constants and helpers shared by sketch + compactor serdes. + +use crate::codec::assert::ensure_preamble_longs_in; +use crate::codec::assert::ensure_serial_version_is; +use crate::error::Error; + +pub(super) const SERIAL_VERSION: u8 = 1; +pub(super) const PREAMBLE_INTS_EXACT: u8 = 2; +pub(super) const PREAMBLE_INTS_ESTIMATION: u8 = 4; +pub(super) const RAW_ITEMS_THRESHOLD: u64 = 4; + +/// Flag bits — match the C++ enum order: RESERVED1, RESERVED2, IS_EMPTY, IS_HIGH_RANK, RAW_ITEMS, +/// IS_LEVEL_ZERO_SORTED. +pub(super) const FLAG_IS_EMPTY: u8 = 1 << 2; +pub(super) const FLAG_IS_HIGH_RANK: u8 = 1 << 3; +pub(super) const FLAG_RAW_ITEMS: u8 = 1 << 4; +pub(super) const FLAG_IS_LEVEL_ZERO_SORTED: u8 = 1 << 5; + +pub(super) fn check_serial_version(actual: u8) -> Result<(), Error> { + ensure_serial_version_is(SERIAL_VERSION, actual) +} + +pub(super) fn check_preamble_ints(actual: u8, num_levels: u8) -> Result<(), Error> { + let expected = if num_levels > 1 { + PREAMBLE_INTS_ESTIMATION + } else { + PREAMBLE_INTS_EXACT + }; + ensure_preamble_longs_in(&[expected], actual) +} diff --git a/datasketches/src/req/sketch.rs b/datasketches/src/req/sketch.rs new file mode 100644 index 00000000..3b4b0b87 --- /dev/null +++ b/datasketches/src/req/sketch.rs @@ -0,0 +1,833 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! REQ sketch — generic over `T: ReqValue`. + +use std::fmt; + +use super::DEFAULT_K; +use super::MAX_K; +use super::MIN_K; +use super::RankAccuracy; +use super::SearchCriteria; +use super::compactor::Compactor; +use super::iter::ReqSketchIterator; +use super::sorted_view::SortedView; +use super::value::ReqValue; +use crate::error::Error; + +/// A Relative Error Quantiles sketch for approximate quantile estimation. +/// +/// See the [module-level documentation](super) for background. +#[derive(Debug, Clone)] +pub struct ReqSketch { + pub(super) k: u16, + pub(super) rank_accuracy: RankAccuracy, + pub(super) n: u64, + pub(super) max_nom_size: u32, + pub(super) num_retained: u32, + pub(super) compactors: Vec>, + pub(super) promotion_buf: Vec, + pub(super) min_item: Option, + pub(super) max_item: Option, +} + +impl ReqSketch { + /// Creates a new sketch with default parameters (`k = 12`, `RankAccuracy::HighRank`). + pub fn new() -> Self { + let mut s = Self { + k: DEFAULT_K, + rank_accuracy: RankAccuracy::HighRank, + n: 0, + max_nom_size: 0, + num_retained: 0, + compactors: Vec::new(), + promotion_buf: Vec::with_capacity(DEFAULT_K as usize), + min_item: None, + max_item: None, + }; + // C++ parity: an empty sketch has a level-0 compactor present from the start. + // This makes is_raw_items() and flags_byte() byte-compatible with the C++/Java + // wire format for the empty case. + s.grow(); + s + } + + /// Creates a new sketch with the given `k` and rank accuracy. + /// + /// # Errors + /// + /// Returns an error if `k` is odd or outside `[MIN_K, MAX_K]`. + pub fn try_new(k: u16, rank_accuracy: RankAccuracy) -> Result { + if !(MIN_K..=MAX_K).contains(&k) { + return Err(Error::invalid_argument(format!( + "k must be in [{}, {}], got {k}", + MIN_K, MAX_K + ))); + } + if k % 2 != 0 { + return Err(Error::invalid_argument(format!("k must be even, got {k}"))); + } + let mut s = Self { + k, + rank_accuracy, + n: 0, + max_nom_size: 0, + num_retained: 0, + compactors: Vec::new(), + promotion_buf: Vec::with_capacity(k as usize), + min_item: None, + max_item: None, + }; + s.grow(); + Ok(s) + } + + /// Returns a builder for chained configuration. + pub fn builder() -> ReqSketchBuilder { + ReqSketchBuilder::new() + } + + /// Returns the configured `k` parameter. + pub fn k(&self) -> u16 { + self.k + } + + /// Returns the configured rank accuracy. + pub fn rank_accuracy(&self) -> RankAccuracy { + self.rank_accuracy + } + + /// Returns the total number of items observed (matches C++ `get_n`). + pub fn n(&self) -> u64 { + self.n + } + + /// Returns true if the sketch has observed no items. + pub fn is_empty(&self) -> bool { + self.n == 0 + } + + /// Returns true if compaction has occurred. + pub fn is_estimation_mode(&self) -> bool { + self.compactors.len() > 1 + } + + /// Returns the number of items currently stored across all compactors. + pub fn num_retained(&self) -> u32 { + self.num_retained + } + + /// Returns the smallest item ever observed, or `None` if empty. + pub fn min_item(&self) -> Option<&T> { + self.min_item.as_ref() + } + + /// Returns the largest item ever observed, or `None` if empty. + pub fn max_item(&self) -> Option<&T> { + self.max_item.as_ref() + } + + /// Updates the sketch with a new item. + /// + /// NaN inputs are silently ignored for floating-point types, matching the behavior + /// of the Java reference implementation (`checkNaNUpdate`). This is intentional and + /// documented in the cross-language differences doc. + pub fn update(&mut self, item: T) { + if item.is_nan() { + return; + } + match &mut self.min_item { + None => self.min_item = Some(item.clone()), + Some(cur) if item.total_cmp(cur).is_lt() => *cur = item.clone(), + _ => {} + } + match &mut self.max_item { + None => self.max_item = Some(item.clone()), + Some(cur) if item.total_cmp(cur).is_gt() => *cur = item.clone(), + _ => {} + } + + self.compactors[0].append(item); + self.n += 1; + self.num_retained += 1; + + if self.num_retained == self.max_nom_size { + self.compress(); + } + } + + /// Resets the sketch to the empty state. + pub fn reset(&mut self) { + self.n = 0; + self.num_retained = 0; + self.max_nom_size = 0; + self.min_item = None; + self.max_item = None; + self.compactors.clear(); + self.grow(); + } + + /// Returns an iterator over `(item, weight)` pairs. + pub fn iter(&self) -> ReqSketchIterator<'_, T> { + ReqSketchIterator::new(&self.compactors) + } + + /// Returns the approximate rank of `item` in `[0.0, 1.0]`. + /// + /// Computed directly from the retained items in a single `O(retained)` pass, + /// without building a sorted view. The result is identical to + /// [`SortedView::rank`] on [`Self::sorted_view`]. + /// + /// # Errors + /// Returns an error if the sketch is empty or `item` is NaN. + pub fn rank(&self, item: &T, criteria: SearchCriteria) -> Result { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + if item.is_nan() { + return Err(Error::invalid_argument("query item is NaN")); + } + let inclusive = matches!(criteria, SearchCriteria::Inclusive); + let weight: u64 = self + .compactors + .iter() + .map(|c| c.count_below(item, inclusive) as u64 * c.weight()) + .sum(); + Ok(weight as f64 / self.n as f64) + } + + /// Returns the approximate quantile at the given normalized rank. + /// + /// Builds a transient [`SortedView`] internally. For repeated quantile + /// queries, take one snapshot with [`Self::sorted_view`] and query it. + pub fn quantile(&self, rank: f64, criteria: SearchCriteria) -> Result { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + if !(0.0..=1.0).contains(&rank) { + return Err(Error::invalid_argument(format!( + "rank {rank} must be in [0, 1]" + ))); + } + self.sorted_view().quantile(rank, criteria) + } + + /// Returns approximate quantiles for the given normalized ranks. + /// + /// The sorted view is built once and shared across all ranks. + pub fn quantiles(&self, ranks: &[f64], criteria: SearchCriteria) -> Result, Error> { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + // Reject invalid ranks before paying for the view build. + for &r in ranks { + if !(0.0..=1.0).contains(&r) { + return Err(Error::invalid_argument(format!( + "rank {r} must be in [0, 1]" + ))); + } + } + let view = self.sorted_view(); + ranks.iter().map(|&r| view.quantile(r, criteria)).collect() + } + + /// Returns the Probability Mass Function over the given split points. + pub fn pmf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + self.sorted_view().pmf(split_points, criteria) + } + + /// Returns the Cumulative Distribution Function over the given split points. + pub fn cdf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + self.sorted_view().cdf(split_points, criteria) + } + + /// Returns an owned, sorted snapshot of the sketch's current state. + /// + /// An empty sketch yields an empty view; queries on it return an error. The + /// view is independent of the sketch — it can be queried (and sent to other + /// threads) while the sketch keeps receiving updates, and it keeps answering + /// from the state it was taken at. + /// + /// Building the view costs `O(retained · log retained)`; each query on it is + /// then `O(log retained)`. Prefer taking one view for repeated queries over + /// calling [`Self::quantile`]/[`Self::pmf`]/[`Self::cdf`], which each build a + /// transient view. + pub fn sorted_view(&self) -> SortedView { + let mut weighted_items = Vec::with_capacity(self.num_retained as usize); + for compactor in &self.compactors { + let weight = compactor.weight(); + for item in compactor.iter() { + weighted_items.push((item.clone(), weight)); + } + } + SortedView::new(weighted_items) + } + + /// Merges another sketch into this one. + /// + /// # Errors + /// + /// Returns an error if the two sketches have different `rank_accuracy`. + pub fn merge(&mut self, other: &Self) -> Result<(), Error> { + if self.rank_accuracy != other.rank_accuracy { + return Err(Error::invalid_argument( + "sketches must have the same rank_accuracy", + )); + } + + if other.is_empty() { + return Ok(()); + } + + self.n += other.n; + + if let Some(m) = &other.min_item { + match &self.min_item { + None => self.min_item = Some(m.clone()), + Some(cur) if m.total_cmp(cur).is_lt() => self.min_item = Some(m.clone()), + _ => {} + } + } + if let Some(m) = &other.max_item { + match &self.max_item { + None => self.max_item = Some(m.clone()), + Some(cur) if m.total_cmp(cur).is_gt() => self.max_item = Some(m.clone()), + _ => {} + } + } + + while self.compactors.len() < other.compactors.len() { + self.grow(); + } + + for (i, other_c) in other.compactors.iter().enumerate() { + self.compactors[i].merge(other_c); + } + + self.update_max_nom_size(); + self.update_num_retained(); + + if self.num_retained >= self.max_nom_size { + self.compress(); + } + + Ok(()) + } + + /// Returns the lower bound for the rank of a quantile at `num_std_dev` confidence. + pub fn rank_lower_bound(&self, rank: f64, num_std_dev: u8) -> f64 { + self.compute_rank_lower_bound( + self.k, + self.compactors.len() as u8, + rank, + num_std_dev, + self.n, + matches!(self.rank_accuracy, RankAccuracy::HighRank), + ) + } + + /// Returns the upper bound for the rank of a quantile at `num_std_dev` confidence. + pub fn rank_upper_bound(&self, rank: f64, num_std_dev: u8) -> f64 { + self.compute_rank_upper_bound( + self.k, + self.compactors.len() as u8, + rank, + num_std_dev, + self.n, + matches!(self.rank_accuracy, RankAccuracy::HighRank), + ) + } + + const FIXED_RSE_FACTOR: f64 = 0.084; + const INIT_NUM_SECTIONS: u8 = 3; + + fn relative_rse_factor() -> f64 { + (0.0512 / Self::INIT_NUM_SECTIONS as f64).sqrt() + } + + fn compute_rank_lower_bound( + &self, + k: u16, + num_levels: u8, + rank: f64, + num_std_dev: u8, + n: u64, + hra: bool, + ) -> f64 { + if self.is_exact_rank_threshold(k, num_levels, rank, n, hra) { + return rank; + } + let relative = Self::relative_rse_factor() / k as f64 * if hra { 1.0 - rank } else { rank }; + let fixed = Self::FIXED_RSE_FACTOR / k as f64; + let lb_rel = rank - num_std_dev as f64 * relative; + let lb_fix = rank - num_std_dev as f64 * fixed; + lb_rel.max(lb_fix).max(0.0) + } + + fn compute_rank_upper_bound( + &self, + k: u16, + num_levels: u8, + rank: f64, + num_std_dev: u8, + n: u64, + hra: bool, + ) -> f64 { + if self.is_exact_rank_threshold(k, num_levels, rank, n, hra) { + return rank; + } + let relative = Self::relative_rse_factor() / k as f64 * if hra { 1.0 - rank } else { rank }; + let fixed = Self::FIXED_RSE_FACTOR / k as f64; + let ub_rel = rank + num_std_dev as f64 * relative; + let ub_fix = rank + num_std_dev as f64 * fixed; + ub_rel.min(ub_fix).min(1.0) + } + + fn is_exact_rank_threshold( + &self, + k: u16, + num_levels: u8, + rank: f64, + n: u64, + hra: bool, + ) -> bool { + let base_cap = k as u64 * Self::INIT_NUM_SECTIONS as u64; + if num_levels == 1 || n <= base_cap { + return true; + } + let exact_rank_thresh = base_cap as f64 / n as f64; + if hra { + rank >= 1.0 - exact_rank_thresh + } else { + rank <= exact_rank_thresh + } + } + + /// Returns per-level info: `(level_index, num_items, capacity, weight)`. + /// Internal/test API; subject to change. + #[doc(hidden)] + pub fn level_info(&self) -> Vec<(usize, u32, u32, u64)> { + self.compactors + .iter() + .enumerate() + .map(|(i, c)| (i, c.num_items(), c.nominal_capacity(), c.weight())) + .collect() + } + + /// Total nominal capacity across all levels. Internal/test API. + #[doc(hidden)] + pub fn total_nominal_capacity(&self) -> u32 { + self.compactors.iter().map(|c| c.nominal_capacity()).sum() + } + + /// Total retained items across all levels. Internal/test API. + #[doc(hidden)] + pub fn total_retained_items(&self) -> u32 { + self.compactors.iter().map(|c| c.num_items()).sum() + } + + /// Sum of `level_items × level_weight` across compactors. Internal/test API. + #[doc(hidden)] + pub fn computed_total_weight(&self) -> u64 { + self.compactors + .iter() + .map(|c| c.num_items() as u64 * c.weight()) + .sum() + } + + pub(super) fn flags_byte(&self) -> u8 { + use super::serialization::FLAG_IS_EMPTY; + use super::serialization::FLAG_IS_HIGH_RANK; + use super::serialization::FLAG_IS_LEVEL_ZERO_SORTED; + use super::serialization::FLAG_RAW_ITEMS; + let mut flags = 0u8; + if self.is_empty() { + flags |= FLAG_IS_EMPTY; + } + if matches!(self.rank_accuracy, RankAccuracy::HighRank) { + flags |= FLAG_IS_HIGH_RANK; + } + if self.is_raw_items() { + flags |= FLAG_RAW_ITEMS; + } + if self.compactors[0].is_sorted() { + flags |= FLAG_IS_LEVEL_ZERO_SORTED; + } + flags + } + + pub(super) fn is_raw_items(&self) -> bool { + use super::serialization::RAW_ITEMS_THRESHOLD; + self.n <= RAW_ITEMS_THRESHOLD && self.compactors.len() == 1 + } + + /// Number of bytes required to serialize the sketch. + pub fn serialized_size_bytes(&self) -> usize { + // Fixed sketch preamble: 8 bytes (preamble_ints, serial_version, family, + // flags, k(2), num_levels, num_raw_items). + let mut size = 8usize; + if self.is_empty() { + return size; + } + if self.is_estimation_mode() { + size += 8; // n + size += T::serialize_size(self.min_item.as_ref().unwrap()); + size += T::serialize_size(self.max_item.as_ref().unwrap()); + } + if self.is_raw_items() { + for item in self.compactors[0].iter() { + size += T::serialize_size(item); + } + } else { + for c in &self.compactors { + // 20-byte compactor preamble + items + size += 20; + for item in c.iter() { + size += T::serialize_size(item); + } + } + } + size + } + + /// Serialize the sketch into a `Vec` matching the C++/Java REQ wire format. + pub fn serialize(&self) -> Vec { + use super::serialization::PREAMBLE_INTS_ESTIMATION; + use super::serialization::PREAMBLE_INTS_EXACT; + use super::serialization::SERIAL_VERSION; + use crate::codec::SketchBytes; + use crate::codec::family::Family; + + let mut out = SketchBytes::with_capacity(self.serialized_size_bytes()); + let preamble_ints = if self.is_estimation_mode() { + PREAMBLE_INTS_ESTIMATION + } else { + PREAMBLE_INTS_EXACT + }; + out.write_u8(preamble_ints); + out.write_u8(SERIAL_VERSION); + out.write_u8(Family::REQ.id); + out.write_u8(self.flags_byte()); + out.write_u16_le(self.k); + let num_levels = if self.is_empty() { + 0 + } else { + self.compactors.len() as u8 + }; + out.write_u8(num_levels); + let num_raw_items = if self.is_raw_items() { self.n as u8 } else { 0 }; + out.write_u8(num_raw_items); + + if self.is_empty() { + return out.into_bytes(); + } + + if self.is_estimation_mode() { + out.write_u64_le(self.n); + self.min_item.as_ref().unwrap().serialize_value(&mut out); + self.max_item.as_ref().unwrap().serialize_value(&mut out); + } + + if self.is_raw_items() { + for item in self.compactors[0].iter() { + item.serialize_value(&mut out); + } + } else { + for c in &self.compactors { + c.serialize_into(&mut out); + } + } + + out.into_bytes() + } + + /// Deserialize a sketch from bytes produced by [`Self::serialize`] or by the + /// C++/Java reference implementations. + pub fn deserialize(bytes: &[u8]) -> Result { + use super::compactor::Compactor; + use super::serialization::FLAG_IS_EMPTY; + use super::serialization::FLAG_IS_HIGH_RANK; + use super::serialization::FLAG_IS_LEVEL_ZERO_SORTED; + use super::serialization::FLAG_RAW_ITEMS; + use super::serialization::RAW_ITEMS_THRESHOLD; + use super::serialization::check_preamble_ints; + use super::serialization::check_serial_version; + use crate::codec::SketchSlice; + use crate::codec::assert::insufficient_data; + use crate::codec::family::Family; + + let mut cursor = SketchSlice::new(bytes); + let preamble_ints = cursor + .read_u8() + .map_err(insufficient_data("preamble_ints"))?; + let serial_version = cursor + .read_u8() + .map_err(insufficient_data("serial_version"))?; + let family_id = cursor.read_u8().map_err(insufficient_data("family_id"))?; + let flags = cursor.read_u8().map_err(insufficient_data("flags"))?; + let k = cursor.read_u16_le().map_err(insufficient_data("k"))?; + let num_levels = cursor.read_u8().map_err(insufficient_data("num_levels"))?; + let num_raw_items = cursor + .read_u8() + .map_err(insufficient_data("num_raw_items"))?; + + check_preamble_ints(preamble_ints, num_levels)?; + check_serial_version(serial_version)?; + Family::REQ.validate_id(family_id)?; + + let is_empty = flags & FLAG_IS_EMPTY != 0; + let hra = flags & FLAG_IS_HIGH_RANK != 0; + let raw_items = flags & FLAG_RAW_ITEMS != 0; + let is_level_zero_sorted = flags & FLAG_IS_LEVEL_ZERO_SORTED != 0; + + let rank_accuracy = if hra { + RankAccuracy::HighRank + } else { + RankAccuracy::LowRank + }; + if !(MIN_K..=MAX_K).contains(&k) || k % 2 != 0 { + return Err(Error::invalid_argument(format!( + "k {k} is not a valid REQ k value" + ))); + } + + if is_empty { + if num_levels != 0 { + return Err(Error::invalid_argument(format!( + "empty REQ sketch must have 0 levels, got {num_levels}" + ))); + } + if num_raw_items != 0 { + return Err(Error::invalid_argument(format!( + "empty REQ sketch must have 0 raw items, got {num_raw_items}" + ))); + } + return ReqSketch::try_new(k, rank_accuracy); + } + + if num_levels == 0 { + return Err(Error::invalid_argument( + "non-empty REQ sketch must have at least one level", + )); + } + + if raw_items { + if num_levels != 1 { + return Err(Error::invalid_argument(format!( + "raw-items REQ sketch must have exactly 1 level, got {num_levels}" + ))); + } + if num_raw_items == 0 || num_raw_items as u64 > RAW_ITEMS_THRESHOLD { + return Err(Error::invalid_argument(format!( + "raw-items REQ sketch must contain 1..={RAW_ITEMS_THRESHOLD} items, got {num_raw_items}" + ))); + } + } else if num_raw_items != 0 { + return Err(Error::invalid_argument(format!( + "non-raw REQ sketch must have 0 raw items, got {num_raw_items}" + ))); + } + + let mut min_item: Option = None; + let mut max_item: Option = None; + let mut n: u64 = 1; + + if num_levels > 1 { + n = cursor.read_u64_le().map_err(insufficient_data("n"))?; + min_item = Some(T::deserialize_value(&mut cursor)?); + max_item = Some(T::deserialize_value(&mut cursor)?); + } + + let mut compactors: Vec> = Vec::with_capacity(num_levels as usize); + + if raw_items { + // Single compactor at level 0; items follow directly. + let mut items = Vec::with_capacity(num_raw_items as usize); + for _ in 0..num_raw_items { + items.push(T::deserialize_value(&mut cursor)?); + } + let c = + Compactor::::raw_items_compactor(k, rank_accuracy, items, is_level_zero_sorted); + compactors.push(c); + } else { + for i in 0..num_levels { + let level_sorted = if i == 0 { is_level_zero_sorted } else { true }; + let c = Compactor::::deserialize(&mut cursor, rank_accuracy, level_sorted)?; + compactors.push(c); + } + } + + if num_levels == 1 { + // Recover n / min / max from level 0 (these aren't in the preamble for exact mode). + let level0 = &compactors[0]; + n = level0.num_items() as u64; + let mut iter = level0.iter(); + if let Some(first) = iter.next() { + let mut mn = first.clone(); + let mut mx = first.clone(); + for x in iter { + if x.total_cmp(&mn).is_lt() { + mn = x.clone(); + } + if x.total_cmp(&mx).is_gt() { + mx = x.clone(); + } + } + min_item = Some(mn); + max_item = Some(mx); + } + } + + if n == 0 || min_item.is_none() || max_item.is_none() { + return Err(Error::invalid_argument( + "non-empty REQ sketch contains no items", + )); + } + + let mut sketch = ReqSketch::try_new(k, rank_accuracy)?; + sketch.n = n; + sketch.min_item = min_item; + sketch.max_item = max_item; + sketch.compactors = compactors; + sketch.update_max_nom_size(); + sketch.update_num_retained(); + Ok(sketch) + } + + // --- Internal --- + + pub(super) fn grow(&mut self) { + let level = self.compactors.len() as u8; + let compactor = Compactor::new(level, self.k, self.rank_accuracy); + self.compactors.push(compactor); + self.update_max_nom_size(); + } + + pub(super) fn compress(&mut self) { + for h in 0..self.compactors.len() { + if self.compactors[h].num_items() >= self.compactors[h].nominal_capacity() { + if h == 0 { + self.compactors[0].sort(); + } + if h + 1 >= self.compactors.len() { + self.grow(); + } + self.promotion_buf.clear(); + self.compactors[h].compact_into(self.rank_accuracy, &mut self.promotion_buf); + if !self.promotion_buf.is_empty() { + self.compactors[h + 1].sort(); + self.compactors[h + 1].merge_sorted(&self.promotion_buf); + } + self.update_max_nom_size(); + self.update_num_retained(); + } + } + } + + pub(super) fn update_max_nom_size(&mut self) { + self.max_nom_size = self.compactors.iter().map(|c| c.nominal_capacity()).sum(); + } + + pub(super) fn update_num_retained(&mut self) { + self.num_retained = self.compactors.iter().map(|c| c.num_items()).sum(); + } +} + +impl Default for ReqSketch { + fn default() -> Self { + Self::new() + } +} + +/// Builder for [`ReqSketch`]. +#[derive(Debug, Clone)] +pub struct ReqSketchBuilder { + k: u16, + rank_accuracy: RankAccuracy, + _marker: std::marker::PhantomData, +} + +impl Default for ReqSketchBuilder { + fn default() -> Self { + Self::new() + } +} + +impl ReqSketchBuilder { + /// Creates a new builder with default parameters. + pub fn new() -> Self { + Self { + k: DEFAULT_K, + rank_accuracy: RankAccuracy::HighRank, + _marker: std::marker::PhantomData, + } + } + + /// Sets the `k` parameter. + /// + /// # Errors + /// + /// Returns an error if `k` is odd or outside `[MIN_K, MAX_K]`. + pub fn k(mut self, k: u16) -> Result { + if !(MIN_K..=MAX_K).contains(&k) { + return Err(Error::invalid_argument(format!( + "k must be in [{}, {}], got {k}", + MIN_K, MAX_K + ))); + } + if k % 2 != 0 { + return Err(Error::invalid_argument(format!("k must be even, got {k}"))); + } + self.k = k; + Ok(self) + } + + /// Sets the rank accuracy. + pub fn rank_accuracy(mut self, ra: RankAccuracy) -> Self { + self.rank_accuracy = ra; + self + } + + /// Builds the sketch. + pub fn build(self) -> Result, Error> { + ReqSketch::try_new(self.k, self.rank_accuracy) + } +} + +impl fmt::Display for ReqSketch { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "REQ Sketch Summary:")?; + writeln!(f, " k : {}", self.k)?; + writeln!(f, " rank accuracy : {:?}", self.rank_accuracy)?; + writeln!(f, " n : {}", self.n)?; + writeln!(f, " num retained : {}", self.num_retained)?; + writeln!(f, " num levels : {}", self.compactors.len())?; + writeln!(f, " estimation mode : {}", self.is_estimation_mode())?; + if let (Some(min), Some(max)) = (&self.min_item, &self.max_item) { + writeln!(f, " min item : {min}")?; + writeln!(f, " max item : {max}")?; + } + Ok(()) + } +} diff --git a/datasketches/src/req/sorted_view.rs b/datasketches/src/req/sorted_view.rs new file mode 100644 index 00000000..337f0c68 --- /dev/null +++ b/datasketches/src/req/sorted_view.rs @@ -0,0 +1,372 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Sorted view implementation for efficient quantile queries. + +use super::SearchCriteria; +use super::value::ReqValue; +use crate::error::Error; + +/// An owned, sorted snapshot of a [`ReqSketch`](super::ReqSketch)'s items with +/// their cumulative weights. +/// +/// Obtain one with [`ReqSketch::sorted_view`](super::ReqSketch::sorted_view). +/// The view is independent of the sketch: it can be queried (and sent to other +/// threads) while the sketch keeps receiving updates, and it keeps answering +/// from the state it was taken at. Building it costs `O(retained · log retained)`; +/// each subsequent query is `O(log retained)`, so it is the right tool for +/// repeated quantile/rank queries. +#[derive(Debug, Clone)] +pub struct SortedView { + /// Items in sorted order + items: Vec, + /// Cumulative weights for each item + cumulative_weights: Vec, + /// Total weight of all items + total_weight: u64, +} + +impl SortedView +where + T: ReqValue, +{ + /// Creates a new sorted view from weighted items. + /// + /// # Arguments + /// * `weighted_items` - Vector of (item, weight) pairs + /// + /// The items will be sorted and cumulative weights computed. + pub(super) fn new(mut weighted_items: Vec<(T, u64)>) -> Self { + if weighted_items.is_empty() { + return Self { + items: Vec::new(), + cumulative_weights: Vec::new(), + total_weight: 0, + }; + } + + // Sort by item value - use unstable sort for better performance + weighted_items.sort_unstable_by(|a, b| a.0.total_cmp(&b.0)); + + let mut items: Vec = Vec::with_capacity(weighted_items.len()); + let mut cumulative_weights = Vec::with_capacity(weighted_items.len()); + let mut cumulative_weight = 0u64; + + for (item, weight) in weighted_items { + if let Some(last) = items.last() { + if matches!(last.total_cmp(&item), std::cmp::Ordering::Equal) { + cumulative_weight += weight; + let last_idx = cumulative_weights.len() - 1; + cumulative_weights[last_idx] = cumulative_weight; + continue; + } + } + cumulative_weight += weight; + items.push(item); + cumulative_weights.push(cumulative_weight); + } + + Self { + items, + cumulative_weights, + total_weight: cumulative_weight, + } + } + + /// Returns true if the sorted view is empty. + pub fn is_empty(&self) -> bool { + self.items.is_empty() + } + + /// Returns the number of distinct items in the sorted view. + pub fn len(&self) -> usize { + self.items.len() + } + + /// Returns the total weight (stream length captured) of all items. + pub fn total_weight(&self) -> u64 { + self.total_weight + } + + /// Returns the approximate normalized rank of the given item in `[0.0, 1.0]`. + /// + /// # Arguments + /// * `item` - The item to find the rank for + /// * `criteria` - Whether to include the item's weight in the rank + /// + /// # Errors + /// Returns an error if the view is empty or `item` is NaN. + pub fn rank(&self, item: &T, criteria: SearchCriteria) -> Result { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + if item.is_nan() { + return Err(Error::invalid_argument("query item is NaN")); + } + + match criteria { + SearchCriteria::Inclusive => { + // Find the last position where items[i] <= item + // partition_point finds first index where predicate is false + let pos = self.items.partition_point(|x| x.total_cmp(item).is_le()); + if pos == 0 { + Ok(0.0) + } else { + Ok(self.cumulative_weights[pos - 1] as f64 / self.total_weight as f64) + } + } + SearchCriteria::Exclusive => { + // Find the last position where items[i] < item + let pos = self.items.partition_point(|x| x.total_cmp(item).is_lt()); + if pos == 0 { + Ok(0.0) + } else { + Ok(self.cumulative_weights[pos - 1] as f64 / self.total_weight as f64) + } + } + } + } + + /// Returns the approximate quantile for the given normalized rank. + /// + /// # Arguments + /// * `rank` - A normalized rank in [0.0, 1.0] + /// * `criteria` - Search criteria for quantile selection + /// + /// # Returns + /// The item at approximately the given rank + pub fn quantile(&self, rank: f64, criteria: SearchCriteria) -> Result { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + + if !(0.0..=1.0).contains(&rank) { + return Err(Error::invalid_argument(format!( + "rank {rank} must be in [0, 1]" + ))); + } + + // Handle edge cases + if rank == 0.0 { + match criteria { + SearchCriteria::Inclusive => return Ok(self.items[0].clone()), + SearchCriteria::Exclusive => return Ok(self.items[0].clone()), + } + } + if rank == 1.0 { + return Ok(self.items[self.items.len() - 1].clone()); + } + + // Convert rank to target cumulative weight + // uint64_t weight = static_cast(inclusive ? std::ceil(rank * total_weight_) : + // rank * total_weight_); + let target_weight = match criteria { + SearchCriteria::Inclusive => (rank * self.total_weight as f64).ceil() as u64, + SearchCriteria::Exclusive => (rank * self.total_weight as f64) as u64, + }; + + let index = match criteria { + SearchCriteria::Inclusive => { + // Equivalent to C++ lower_bound: first index where cumulative_weight >= target + self.cumulative_weights + .partition_point(|&w| w < target_weight) + } + SearchCriteria::Exclusive => { + // Equivalent to C++ upper_bound: first index where cumulative_weight > target + self.cumulative_weights + .partition_point(|&w| w <= target_weight) + } + }; + + if index >= self.items.len() { + return Ok(self.items[self.items.len() - 1].clone()); + } + + Ok(self.items[index].clone()) + } + + /// Returns the Probability Mass Function (PMF) for the given split points. + /// + /// # Arguments + /// * `split_points` - Array of split points that divide the domain + /// * `criteria` - Search criteria for boundary handling + /// + /// # Returns + /// Array of probabilities for each interval defined by the split points + pub fn pmf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + + self.validate_split_points(split_points)?; + + let mut result = Vec::with_capacity(split_points.len() + 1); + let mut prev_rank = 0.0; + + for split_point in split_points { + let rank = self.rank(split_point, criteria)?; + result.push(rank - prev_rank); + prev_rank = rank; + } + + // Add the final interval + result.push(1.0 - prev_rank); + + Ok(result) + } + + /// Returns the Cumulative Distribution Function (CDF) for the given split points. + /// + /// # Arguments + /// * `split_points` - Array of split points that divide the domain + /// * `criteria` - Search criteria for boundary handling + /// + /// # Returns + /// Array of cumulative probabilities at each split point + pub fn cdf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { + if self.is_empty() { + return Err(Error::invalid_argument("sketch is empty")); + } + + self.validate_split_points(split_points)?; + + let mut result = Vec::with_capacity(split_points.len() + 1); + let mut cumulative = 0.0; + + let pmf = self.pmf(split_points, criteria)?; + for mass in pmf { + cumulative += mass; + result.push(cumulative); + } + + Ok(result) + } + + // Private helper methods + + fn validate_split_points(&self, split_points: &[T]) -> Result<(), Error> { + // Check that split points are monotonically increasing + for i in 1..split_points.len() { + if split_points[i - 1].total_cmp(&split_points[i]).is_ge() { + return Err(Error::invalid_argument( + "Split points must be unique and monotonically increasing".to_string(), + )); + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn create_test_view() -> SortedView { + let weighted_items = vec![(1, 1), (3, 1), (5, 1), (7, 1), (9, 1)]; + SortedView::new(weighted_items) + } + + #[test] + fn test_sorted_view_creation() { + let view = create_test_view(); + assert_eq!(view.len(), 5); + assert_eq!(view.total_weight(), 5); + assert!(!view.is_empty()); + } + + #[test] + fn test_rank_queries() -> Result<(), Error> { + let view = create_test_view(); + + // Test exact matches + assert!((view.rank(&1, SearchCriteria::Inclusive)? - 0.2).abs() < 1e-10); + assert!((view.rank(&1, SearchCriteria::Exclusive)? - 0.0).abs() < 1e-10); + + // Test values between items + assert!((view.rank(&2, SearchCriteria::Inclusive)? - 0.2).abs() < 1e-10); + assert!((view.rank(&6, SearchCriteria::Inclusive)? - 0.6).abs() < 1e-10); + + // Test edge cases + assert!((view.rank(&0, SearchCriteria::Inclusive)? - 0.0).abs() < 1e-10); + assert!((view.rank(&10, SearchCriteria::Inclusive)? - 1.0).abs() < 1e-10); + Ok(()) + } + + #[test] + fn test_quantile_queries() -> Result<(), Error> { + let view = create_test_view(); + + // Test edge cases + assert_eq!(view.quantile(0.0, SearchCriteria::Inclusive)?, 1); + assert_eq!(view.quantile(1.0, SearchCriteria::Inclusive)?, 9); + + // Test middle values + let median = view.quantile(0.5, SearchCriteria::Inclusive)?; + assert!((3..=7).contains(&median)); // Should be around the middle (values are 1,3,5,7,9) + + // Test various ranks + let q25 = view.quantile(0.25, SearchCriteria::Inclusive)?; + let q75 = view.quantile(0.75, SearchCriteria::Inclusive)?; + assert!(q25 <= median); + assert!(median <= q75); + Ok(()) + } + + #[test] + fn test_pmf() -> Result<(), Error> { + let view = create_test_view(); + let split_points = vec![3, 7]; + + let pmf = view.pmf(&split_points, SearchCriteria::Inclusive)?; + assert_eq!(pmf.len(), 3); // 2 split points create 3 intervals + + // Sum should be approximately 1.0 + let sum: f64 = pmf.iter().sum(); + assert!((sum - 1.0).abs() < 1e-10); + Ok(()) + } + + #[test] + fn test_cdf() -> Result<(), Error> { + let view = create_test_view(); + let split_points = vec![3, 7]; + + let cdf = view.cdf(&split_points, SearchCriteria::Inclusive)?; + assert_eq!(cdf.len(), 3); + + // CDF should be monotonically increasing + for i in 1..cdf.len() { + assert!(cdf[i] >= cdf[i - 1]); + } + + // Last value should be 1.0 + assert!((cdf[cdf.len() - 1] - 1.0).abs() < 1e-10); + Ok(()) + } + + #[test] + fn test_empty_view() { + let view: SortedView = SortedView::new(vec![]); + assert!(view.is_empty()); + assert_eq!(view.len(), 0); + assert_eq!(view.total_weight(), 0); + + // Operations on empty view should return errors + assert!(view.rank(&5, SearchCriteria::Inclusive).is_err()); + assert!(view.quantile(0.5, SearchCriteria::Inclusive).is_err()); + } +} diff --git a/datasketches/src/req/union.rs b/datasketches/src/req/union.rs new file mode 100644 index 00000000..493189da --- /dev/null +++ b/datasketches/src/req/union.rs @@ -0,0 +1,95 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! REQ union — combines REQ sketches into a single result. + +use super::RankAccuracy; +use super::sketch::ReqSketch; +use super::value::ReqValue; +use crate::error::Error; + +/// Combines multiple [`ReqSketch`]es into a single result. +/// +/// Internally wraps a `ReqSketch` configured for union semantics. The C++ +/// equivalent is `req_union`. +#[derive(Debug, Clone)] +pub struct ReqUnion { + inner: ReqSketch, +} + +impl ReqUnion { + /// Creates a new union with default `k = 12` and `RankAccuracy::HighRank`. + pub fn new() -> Self { + Self { + inner: ReqSketch::new(), + } + } + + /// Creates a new union with the given `k` and rank accuracy. + /// + /// # Errors + /// + /// Returns an error if `k` is invalid (see [`ReqSketch::try_new`]). + pub fn try_new(k: u16, rank_accuracy: RankAccuracy) -> Result { + Ok(Self { + inner: ReqSketch::try_new(k, rank_accuracy)?, + }) + } + + /// Returns the configured `k` parameter. + pub fn k(&self) -> u16 { + self.inner.k() + } + + /// Returns the configured rank accuracy. + pub fn rank_accuracy(&self) -> RankAccuracy { + self.inner.rank_accuracy() + } + + /// Returns true if the union has not yet absorbed any sketch. + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Merges a sketch into the union. + /// + /// # Errors + /// + /// Returns an error if the sketch's `rank_accuracy` differs from the union's. + pub fn merge(&mut self, sketch: &ReqSketch) -> Result<(), Error> { + self.inner.merge(sketch) + } + + /// Extracts the merged result as a [`ReqSketch`]. + /// + /// Equivalent to C++ `req_union::get_result`. Renamed per the workspace's + /// CPC PR #81 precedent. + pub fn to_sketch(&self) -> ReqSketch { + self.inner.clone() + } + + /// Resets the union to empty. + pub fn reset(&mut self) { + self.inner.reset(); + } +} + +impl Default for ReqUnion { + fn default() -> Self { + Self::new() + } +} diff --git a/datasketches/src/req/value.rs b/datasketches/src/req/value.rs new file mode 100644 index 00000000..a1db3fbe --- /dev/null +++ b/datasketches/src/req/value.rs @@ -0,0 +1,191 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Trait for types storable in a [`ReqSketch`](super::ReqSketch). + +use std::cmp::Ordering; + +use crate::codec::SketchBytes; +use crate::codec::SketchSlice; +use crate::error::Error; + +/// Trait for types that can be stored in a [`ReqSketch`](super::ReqSketch). +/// +/// Provides total ordering (so floating-point types with NaN are well-defined under +/// sketch operations) and binary serialization compatible with the Apache DataSketches +/// REQ wire format used by the C++ and Java reference implementations. +pub trait ReqValue: Sized + Clone + PartialOrd { + /// Total ordering used for sketch operations (sort, compaction, rank, quantile). + /// + /// For integer types this is equivalent to [`Ord::cmp`]. For floating-point types + /// this delegates to [`f32::total_cmp`] / [`f64::total_cmp`] so NaN comparisons are + /// deterministic. + fn total_cmp(&self, other: &Self) -> Ordering; + + /// Returns true if this value is the floating-point NaN sentinel. + /// + /// Default: false (integer types are never NaN). Float impls override + /// to delegate to [`f32::is_nan`] / [`f64::is_nan`]. + #[inline(always)] + fn is_nan(&self) -> bool { + false + } + + /// Number of bytes this value will occupy when serialized. + fn serialize_size(item: &Self) -> usize; + + /// Serialize this value into the byte buffer. + fn serialize_value(&self, bytes: &mut SketchBytes); + + /// Deserialize a value from the byte cursor. + fn deserialize_value(cursor: &mut SketchSlice<'_>) -> Result; +} + +macro_rules! impl_req_value_primitive { + // Form with explicit is_nan body (for float types). + ($t:ty, $read:ident, $write:ident, $cmp:expr, nan: $nan:expr) => { + impl ReqValue for $t { + #[inline(always)] + fn total_cmp(&self, other: &Self) -> Ordering { + $cmp(self, other) + } + + fn serialize_size(_item: &Self) -> usize { + std::mem::size_of::<$t>() + } + + fn serialize_value(&self, bytes: &mut SketchBytes) { + bytes.$write(*self); + } + + fn deserialize_value(cursor: &mut SketchSlice<'_>) -> Result { + cursor.$read().map_err(|_| { + Error::insufficient_data(concat!( + "failed to read ", + stringify!($t), + " from REQ sketch" + )) + }) + } + + #[inline(always)] + fn is_nan(&self) -> bool { + $nan(self) + } + } + }; + // Form without is_nan (for integer types — default returns false). + ($t:ty, $read:ident, $write:ident, $cmp:expr) => { + impl ReqValue for $t { + #[inline(always)] + fn total_cmp(&self, other: &Self) -> Ordering { + $cmp(self, other) + } + + fn serialize_size(_item: &Self) -> usize { + std::mem::size_of::<$t>() + } + + fn serialize_value(&self, bytes: &mut SketchBytes) { + bytes.$write(*self); + } + + fn deserialize_value(cursor: &mut SketchSlice<'_>) -> Result { + cursor.$read().map_err(|_| { + Error::insufficient_data(concat!( + "failed to read ", + stringify!($t), + " from REQ sketch" + )) + }) + } + } + }; +} + +impl_req_value_primitive!(i32, read_i32_le, write_i32_le, Ord::cmp); +impl_req_value_primitive!(i64, read_i64_le, write_i64_le, Ord::cmp); +impl_req_value_primitive!(u32, read_u32_le, write_u32_le, Ord::cmp); +impl_req_value_primitive!(u64, read_u64_le, write_u64_le, Ord::cmp); +impl_req_value_primitive!(f32, read_f32_le, write_f32_le, + |a: &f32, b: &f32| if let Some(o) = a.partial_cmp(b) { o } else { f32::total_cmp(a, b) }, + nan: |x: &f32| f32::is_nan(*x)); +impl_req_value_primitive!(f64, read_f64_le, write_f64_le, + |a: &f64, b: &f64| if let Some(o) = a.partial_cmp(b) { o } else { f64::total_cmp(a, b) }, + nan: |x: &f64| f64::is_nan(*x)); + +#[cfg(test)] +mod tests { + use super::*; + + fn round_trip(v: T) { + let mut bytes = SketchBytes::with_capacity(T::serialize_size(&v)); + v.serialize_value(&mut bytes); + let raw = bytes.into_bytes(); + assert_eq!(raw.len(), T::serialize_size(&v)); + let mut cursor = SketchSlice::new(&raw); + let got = T::deserialize_value(&mut cursor).unwrap(); + assert_eq!(got, v); + } + + #[test] + fn round_trip_integers() { + round_trip(0_i32); + round_trip(i32::MIN); + round_trip(i32::MAX); + round_trip(0_i64); + round_trip(i64::MIN); + round_trip(i64::MAX); + round_trip(0_u32); + round_trip(u32::MAX); + round_trip(0_u64); + round_trip(u64::MAX); + } + + #[test] + fn round_trip_floats() { + round_trip(0.0_f32); + round_trip(-1.5_f32); + round_trip(f32::MIN); + round_trip(f32::MAX); + round_trip(f32::INFINITY); + round_trip(f32::NEG_INFINITY); + round_trip(0.0_f64); + round_trip(-1.5_f64); + round_trip(f64::MIN); + round_trip(f64::MAX); + round_trip(f64::INFINITY); + round_trip(f64::NEG_INFINITY); + } + + #[test] + fn total_cmp_handles_nan_for_floats() { + // Pure NaN comparisons under PartialOrd return None; total_cmp must give a definite + // Ordering. + let nan = f64::NAN; + let one = 1.0_f64; + assert_ne!(::total_cmp(&nan, &one), Ordering::Equal); + assert_eq!(::total_cmp(&nan, &nan), Ordering::Equal); + } + + #[test] + fn total_cmp_for_integers_matches_ord() { + assert_eq!(::total_cmp(&3, &5), Ordering::Less); + assert_eq!(::total_cmp(&5, &5), Ordering::Equal); + assert_eq!(::total_cmp(&7, &5), Ordering::Greater); + } +} diff --git a/datasketches/tests/req_test/accuracy.rs b/datasketches/tests/req_test/accuracy.rs new file mode 100644 index 00000000..3d9e8e7c --- /dev/null +++ b/datasketches/tests/req_test/accuracy.rs @@ -0,0 +1,54 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! End-to-end accuracy checks for ReqSketch. + +use datasketches::error::Error; +use datasketches::req::ReqSketch; +use datasketches::req::SearchCriteria; + +#[test] +fn rank_space_error_is_bounded() -> Result<(), Error> { + let mut sketch = ReqSketch::new(); + let n = 50_000; + + for i in 0..n { + sketch.update(i as f64); + } + + assert_eq!(sketch.n(), n as u64); + + for rank in [0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99] { + let quantile = sketch.quantile(rank, SearchCriteria::Inclusive)?; + let estimated_rank = sketch.rank(&quantile, SearchCriteria::Inclusive)?; + let abs_rank_error = (estimated_rank - rank).abs(); + let max_abs_rank_error = if rank >= 0.9 { 0.01 } else { 0.02 }; + + assert!( + abs_rank_error <= max_abs_rank_error, + "rank {} abs error {:.4} > {:.4}", + rank, + abs_rank_error, + max_abs_rank_error + ); + } + + assert!(!sketch.is_empty()); + Ok(()) +} diff --git a/datasketches/tests/req_test/bounds.rs b/datasketches/tests/req_test/bounds.rs new file mode 100644 index 00000000..7fc14693 --- /dev/null +++ b/datasketches/tests/req_test/bounds.rs @@ -0,0 +1,171 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Rank error bounds and sigma coverage for ReqSketch. + +use datasketches::error::Error; +use datasketches::req::RankAccuracy; +use datasketches::req::ReqSketch; +use datasketches::req::SearchCriteria; + +#[test] +fn bounds_are_nested_and_in_unit_interval() { + let mut sketch = ReqSketch::builder() + .k(12) + .expect("valid k") + .rank_accuracy(RankAccuracy::HighRank) + .build() + .expect("build should succeed"); + + for i in 0..50_000 { + sketch.update(i as f64); + } + + for rank in [0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99, 0.999] { + let bounds: Vec<(f64, f64)> = (1..=3u8) + .map(|sigma| { + ( + sketch.rank_lower_bound(rank, sigma), + sketch.rank_upper_bound(rank, sigma), + ) + }) + .collect(); + + for (lower, upper) in &bounds { + assert!(lower <= upper); + assert!((0.0..=1.0).contains(lower)); + assert!((0.0..=1.0).contains(upper)); + } + + assert!(bounds[1].0 <= bounds[0].0 && bounds[0].1 <= bounds[1].1); + assert!(bounds[2].0 <= bounds[1].0 && bounds[1].1 <= bounds[2].1); + } +} + +#[test] +fn theoretical_error_bounds_cover_uniform_quantiles() -> Result<(), Error> { + let mut sketch = ReqSketch::new(); + let n = 50_000; + + for i in 0..n { + sketch.update(i as f64); + } + + for rank in [ + 0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.4, 0.5, 0.6, 0.7, 0.75, 0.8, 0.85, 0.9, 0.92, + 0.95, 0.97, 0.98, 0.99, 0.995, 0.999, + ] { + let true_quantile = rank * (n - 1) as f64; + let estimated_rank = sketch.rank(&true_quantile, SearchCriteria::Inclusive)?; + let lower = sketch.rank_lower_bound(rank, 3); + let upper = sketch.rank_upper_bound(rank, 3); + assert!( + estimated_rank >= lower && estimated_rank <= upper, + "rank {} estimate {:.6} outside [{:.6}, {:.6}]", + rank, + estimated_rank, + lower, + upper + ); + } + + Ok(()) +} + +#[test] +fn hra_and_lra_bounds_are_tighter_at_their_target_end() -> Result<(), Error> { + for rank in [0.05, 0.25, 0.5, 0.75, 0.95] { + let mut hra = ReqSketch::builder() + .rank_accuracy(RankAccuracy::HighRank) + .build()?; + let mut lra = ReqSketch::builder() + .rank_accuracy(RankAccuracy::LowRank) + .build()?; + + for i in 0..10_000 { + hra.update(i as f64); + lra.update(i as f64); + } + + let hra_error = + (rank - hra.rank_lower_bound(rank, 2)).max(hra.rank_upper_bound(rank, 2) - rank); + let lra_error = + (rank - lra.rank_lower_bound(rank, 2)).max(lra.rank_upper_bound(rank, 2) - rank); + + if rank >= 0.75 { + assert!(hra_error <= lra_error); + } else if rank <= 0.25 { + assert!(lra_error <= hra_error); + } + } + + Ok(()) +} + +#[test] +fn exact_mode_bounds_are_tight() { + let mut sketch = ReqSketch::new(); + + for i in 0..20 { + sketch.update(i as f64); + } + + assert!(!sketch.is_estimation_mode()); + + for rank in [0.1, 0.25, 0.5, 0.75, 0.9] { + let lower = sketch.rank_lower_bound(rank, 2); + let upper = sketch.rank_upper_bound(rank, 2); + assert!((upper - lower) / 2.0 < 0.05); + } +} + +#[test] +fn high_rank_accuracy_matches_tight_thresholds() { + let mut sketch = ReqSketch::new(); + let n = 50_000; + + for i in 0..n { + sketch.update(i as f64); + } + + assert_eq!(sketch.n(), n as u64); + + for rank in [0.5, 0.9, 0.95, 0.99, 0.999] { + let true_quantile = rank * (n - 1) as f64; + let estimated_rank = sketch + .rank(&true_quantile, SearchCriteria::Inclusive) + .expect("rank should succeed"); + let abs_error = (estimated_rank - rank).abs(); + let max_abs_error = if rank >= 0.99 { + 0.005 + } else if rank >= 0.9 { + 0.01 + } else { + 0.02 + }; + + assert!(abs_error <= max_abs_error); + } + + for rank in [0.9, 0.99, 0.999] { + let lower = sketch.rank_lower_bound(rank, 3); + let upper = sketch.rank_upper_bound(rank, 3); + assert!(rank >= lower && rank <= upper); + } +} diff --git a/datasketches/tests/req_test/core.rs b/datasketches/tests/req_test/core.rs new file mode 100644 index 00000000..ee363207 --- /dev/null +++ b/datasketches/tests/req_test/core.rs @@ -0,0 +1,243 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Core ReqSketch construction and update behavior. + +use approx::assert_relative_eq; +use datasketches::error::Error; +use datasketches::req::RankAccuracy; +use datasketches::req::ReqSketch; +use datasketches::req::SearchCriteria; + +#[test] +fn empty_sketch_has_default_state_and_rejects_queries() { + let sketch: ReqSketch = ReqSketch::new(); + + assert_eq!(sketch.k(), 12); + assert!(sketch.is_empty()); + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.n(), 0); + assert_eq!(sketch.num_retained(), 0); + assert!(sketch.min_item().is_none()); + assert!(sketch.max_item().is_none()); + + assert!(sketch.rank(&0.0, SearchCriteria::Inclusive).is_err()); + assert!(sketch.quantile(0.5, SearchCriteria::Inclusive).is_err()); + assert!(sketch.pmf(&[0.0], SearchCriteria::Inclusive).is_err()); + assert!(sketch.cdf(&[0.0], SearchCriteria::Inclusive).is_err()); +} + +#[test] +fn single_value_hra_answers_exactly() { + let mut sketch = ReqSketch::new(); + sketch.update(1.0f32); + + assert!(!sketch.is_empty()); + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.n(), 1); + assert_eq!(sketch.num_retained(), 1); + assert_eq!(sketch.min_item(), Some(&1.0)); + assert_eq!(sketch.max_item(), Some(&1.0)); + + assert_relative_eq!( + sketch + .rank(&1.0, SearchCriteria::Exclusive) + .expect("rank should succeed"), + 0.0 + ); + assert_relative_eq!( + sketch + .rank(&1.0, SearchCriteria::Inclusive) + .expect("rank should succeed"), + 1.0 + ); + assert_relative_eq!( + sketch + .rank(&1.1, SearchCriteria::Exclusive) + .expect("rank should succeed"), + 1.0 + ); + assert_relative_eq!( + sketch + .rank(&f32::INFINITY, SearchCriteria::Inclusive) + .expect("rank should succeed"), + 1.0 + ); + + for rank in [0.0, 0.5, 1.0] { + assert_relative_eq!( + sketch + .quantile(rank, SearchCriteria::Exclusive) + .expect("quantile should succeed"), + 1.0 + ); + } +} + +#[test] +fn single_value_lra_preserves_configuration() { + let mut sketch: ReqSketch = ReqSketch::builder() + .rank_accuracy(RankAccuracy::LowRank) + .build() + .expect("build should succeed"); + sketch.update(1.0f32); + + assert_eq!(sketch.rank_accuracy(), RankAccuracy::LowRank); + assert!(!sketch.is_empty()); + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.n(), 1); + assert_eq!(sketch.num_retained(), 1); +} + +#[test] +fn repeated_values_respect_search_criteria() { + let mut sketch = ReqSketch::new(); + for _ in 0..3 { + sketch.update(1.0f32); + } + for _ in 0..3 { + sketch.update(2.0f32); + } + + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.n(), 6); + assert_eq!(sketch.num_retained(), 6); + + assert_relative_eq!( + sketch + .rank(&1.0, SearchCriteria::Exclusive) + .expect("rank should succeed"), + 0.0 + ); + assert_relative_eq!( + sketch + .rank(&1.0, SearchCriteria::Inclusive) + .expect("rank should succeed"), + 0.5 + ); + assert_relative_eq!( + sketch + .rank(&2.0, SearchCriteria::Exclusive) + .expect("rank should succeed"), + 0.5 + ); + assert_relative_eq!( + sketch + .rank(&2.0, SearchCriteria::Inclusive) + .expect("rank should succeed"), + 1.0 + ); +} + +#[test] +fn estimation_mode_compresses_and_keeps_min_max() { + let mut sketch = ReqSketch::new(); + let n = 100_000; + + for i in 0..n { + sketch.update(i as f32); + } + + assert!(!sketch.is_empty()); + assert!(sketch.is_estimation_mode()); + assert_eq!(sketch.n(), n); + assert!(sketch.num_retained() < n as u32); + assert_eq!(sketch.min_item(), Some(&0.0)); + assert_eq!(sketch.max_item(), Some(&((n - 1) as f32))); + + let r0 = sketch + .rank(&0.0, SearchCriteria::Exclusive) + .expect("rank should succeed"); + let rmid = sketch + .rank(&(n as f32 / 2.0), SearchCriteria::Exclusive) + .expect("rank should succeed"); + let rmax = sketch + .rank(&(n as f32), SearchCriteria::Exclusive) + .expect("rank should succeed"); + + assert!((r0 - 0.0).abs() <= 1e-3); + assert!((rmid - 0.5).abs() <= 0.01); + assert!((rmax - 1.0).abs() <= 1e-3); +} + +#[test] +fn nan_updates_are_silently_skipped_for_f64() { + let mut sketch: ReqSketch = ReqSketch::new(); + sketch.update(f64::NAN); + sketch.update(f64::NAN); + assert!(sketch.is_empty()); + assert_eq!(sketch.n(), 0); + + sketch.update(1.0); + sketch.update(f64::NAN); + sketch.update(2.0); + assert_eq!(sketch.n(), 2); + assert_eq!(sketch.min_item(), Some(&1.0)); + assert_eq!(sketch.max_item(), Some(&2.0)); +} + +#[test] +fn nan_updates_are_silently_skipped_for_f32() { + let mut sketch: ReqSketch = ReqSketch::new(); + sketch.update(f32::NAN); + assert!(sketch.is_empty()); + assert_eq!(sketch.n(), 0); + + sketch.update(5.0f32); + sketch.update(f32::NAN); + assert_eq!(sketch.n(), 1); + assert_eq!( + sketch + .quantile(0.5, SearchCriteria::Inclusive) + .expect("quantile should succeed"), + 5.0f32 + ); +} + +#[test] +fn small_edge_cases_answer_reasonably() -> Result<(), Error> { + let mut single = ReqSketch::new(); + single.update(42.0); + assert_eq!(single.quantile(0.5, SearchCriteria::Inclusive)?, 42.0); + + let mut two_values = ReqSketch::new(); + two_values.update(1.0); + two_values.update(100.0); + let median = two_values.quantile(0.5, SearchCriteria::Inclusive)?; + assert!((1.0..=100.0).contains(&median)); + + let mut duplicates = ReqSketch::new(); + for _ in 0..100 { + duplicates.update(42.0); + } + assert_eq!(duplicates.quantile(0.5, SearchCriteria::Inclusive)?, 42.0); + + Ok(()) +} + +#[test] +fn constructors_validate_k() { + // k must be even and within the supported range; both constructors enforce it. + assert!(ReqSketch::::try_new(0, RankAccuracy::HighRank).is_err()); + assert!(ReqSketch::::try_new(3, RankAccuracy::HighRank).is_err()); // odd + assert!(ReqSketch::::try_new(4096, RankAccuracy::HighRank).is_err()); // too large + assert!(ReqSketch::::try_new(12, RankAccuracy::HighRank).is_ok()); + assert!(ReqSketch::::builder().k(5).is_err()); // odd via builder + assert!(ReqSketch::::builder().k(12).is_ok()); +} diff --git a/datasketches/tests/req_test/main.rs b/datasketches/tests/req_test/main.rs new file mode 100644 index 00000000..a06f9cab --- /dev/null +++ b/datasketches/tests/req_test/main.rs @@ -0,0 +1,54 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::path::PathBuf; + +fn serialization_test_data(sub_dir: &str, name: &str) -> PathBuf { + const SERDE_TESTS_DIR: &str = "tests/serde_tests"; + + let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join(SERDE_TESTS_DIR) + .join(sub_dir) + .join(name); + + if !path.exists() { + panic!( + r#"serialization test data file not found: {} + + Please ensure test data files are present in the repository. Generally, you can + run the following commands from the project root to prepare the test data files + if they are missing: + + $ cargo x prepare-testdata + "#, + path.display(), + ); + } + + path +} + +mod accuracy; +mod bounds; +mod core; +mod merge; +mod property; +mod query; +mod serialization; +mod sorted_view_api; +mod structure; +mod union; diff --git a/datasketches/tests/req_test/merge.rs b/datasketches/tests/req_test/merge.rs new file mode 100644 index 00000000..4837625c --- /dev/null +++ b/datasketches/tests/req_test/merge.rs @@ -0,0 +1,141 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Merge behavior for ReqSketch. + +use datasketches::req::RankAccuracy; +use datasketches::req::ReqSketch; +use datasketches::req::SearchCriteria; + +#[test] +fn merge_into_empty_preserves_source_distribution() { + let mut target: ReqSketch = ReqSketch::builder() + .k(40) + .expect("valid k") + .build() + .expect("build should succeed"); + let mut source: ReqSketch = ReqSketch::builder() + .k(40) + .expect("valid k") + .build() + .expect("build should succeed"); + + for i in 0..1000 { + source.update(i as f32); + } + + target.merge(&source).expect("merge should succeed"); + assert_eq!(target.min_item(), Some(&0.0)); + assert_eq!(target.max_item(), Some(&999.0)); + + let q25 = target + .quantile(0.25, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + let q50 = target + .quantile(0.5, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + let q75 = target + .quantile(0.75, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + let r50 = target + .rank(&500.0, SearchCriteria::Inclusive) + .expect("rank should succeed"); + + assert!((q25 - 250.0).abs() / 250.0 <= 0.01); + assert!((q50 - 500.0).abs() / 500.0 <= 0.01); + assert!((q75 - 750.0).abs() / 750.0 <= 0.01); + assert!((r50 - 0.5).abs() / 0.5 <= 0.01); +} + +#[test] +fn merge_two_ranges_preserves_distribution() { + let mut left: ReqSketch = ReqSketch::builder() + .k(100) + .expect("valid k") + .build() + .expect("build should succeed"); + let mut right: ReqSketch = ReqSketch::builder() + .k(100) + .expect("valid k") + .build() + .expect("build should succeed"); + + for i in 0..1000 { + left.update(i as f32); + } + for i in 1000..2000 { + right.update(i as f32); + } + + left.merge(&right).expect("merge should succeed"); + assert_eq!(left.min_item(), Some(&0.0)); + assert_eq!(left.max_item(), Some(&1999.0)); + + let q25 = left + .quantile(0.25, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + let q50 = left + .quantile(0.5, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + let q75 = left + .quantile(0.75, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + let r50 = left + .rank(&1000.0, SearchCriteria::Inclusive) + .expect("rank should succeed"); + + assert!((q25 - 500.0).abs() / 500.0 <= 0.02); + assert!((q50 - 1000.0).abs() / 1000.0 <= 0.01); + assert!((q75 - 1500.0).abs() / 1500.0 <= 0.01); + assert!((r50 - 0.5).abs() / 0.5 <= 0.01); +} + +#[test] +fn merge_rejects_incompatible_accuracy_modes() { + let mut high_rank = ReqSketch::new(); + let low_rank: ReqSketch = ReqSketch::builder() + .rank_accuracy(RankAccuracy::LowRank) + .build() + .expect("build should succeed"); + + high_rank.update(1.0); + assert!(high_rank.merge(&low_rank).is_err()); +} + +#[test] +fn many_small_merges_preserve_count_bounds_and_median() { + let mut sketch = ReqSketch::new(); + + for batch in 0..100 { + let mut batch_sketch = ReqSketch::new(); + for i in 0..100 { + batch_sketch.update((batch * 100 + i) as f64); + } + sketch.merge(&batch_sketch).expect("merge should succeed"); + } + + assert_eq!(sketch.n(), 10_000); + assert_eq!(sketch.min_item(), Some(&0.0)); + assert_eq!(sketch.max_item(), Some(&9999.0)); + + let median = sketch + .quantile(0.5, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + assert!((median - 4999.5).abs() < 500.0); +} diff --git a/datasketches/tests/req_test/property.rs b/datasketches/tests/req_test/property.rs new file mode 100644 index 00000000..43b654d2 --- /dev/null +++ b/datasketches/tests/req_test/property.rs @@ -0,0 +1,113 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Property-based ReqSketch tests. + +use datasketches::req::ReqSketch; +use datasketches::req::SearchCriteria; +use proptest::prelude::*; + +proptest! { + #[test] + fn prop_quantile_rank_consistency( + values in prop::collection::vec(0.0f64..1000.0, 500..1500), + ) { + let mut sketch = ReqSketch::new(); + for value in values { + sketch.update(value); + } + + // These sizes push the sketch past the compaction threshold, so the + // round-trip exercises the estimation path rather than exact storage. + prop_assume!(sketch.is_estimation_mode()); + + for rank in [0.1, 0.25, 0.5, 0.75, 0.9] { + let quantile = sketch + .quantile(rank, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + let recovered = sketch + .rank(&quantile, SearchCriteria::Inclusive) + .expect("rank should succeed"); + + // The recovered rank must land within the sketch's own 3-sigma rank + // interval for the target rank (plus a small cushion for snapping to a + // stored item). This scales with k and n, unlike a fixed slack, so it + // actually constrains the result instead of always passing. + let lower = sketch.rank_lower_bound(rank, 3) - 0.02; + let upper = sketch.rank_upper_bound(rank, 3) + 0.02; + prop_assert!( + (lower..=upper).contains(&recovered), + "rank {rank} -> quantile {quantile} -> recovered {recovered}, expected within [{lower:.4}, {upper:.4}]" + ); + } + } + + #[test] + fn prop_sketch_bounds(values in prop::collection::vec(-1000.0f64..1000.0, 1..1000)) { + let mut sketch = ReqSketch::new(); + for value in &values { + sketch.update(*value); + } + + if sketch.is_empty() { + return Ok(()); + } + + let true_min = values.iter().fold(f64::INFINITY, |a, &b| a.min(b)); + let true_max = values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)); + + prop_assert_eq!(sketch.min_item(), Some(&true_min)); + prop_assert_eq!(sketch.max_item(), Some(&true_max)); + + for rank in [0.0, 0.25, 0.5, 0.75, 1.0] { + let quantile = sketch + .quantile(rank, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + prop_assert!( + quantile >= true_min && quantile <= true_max, + "quantile {} out of bounds [{}, {}]", + quantile, + true_min, + true_max + ); + } + } + + #[test] + fn prop_rank_monotonicity(values in prop::collection::vec(0.0f64..1000.0, 10..100)) { + let mut sketch = ReqSketch::new(); + for value in values { + sketch.update(value); + } + + if sketch.is_empty() { + return Ok(()); + } + + let mut last_rank = -1.0; + for value in [0.0, 100.0, 200.0, 500.0, 800.0, 1000.0] { + let rank = sketch + .rank(&value, SearchCriteria::Inclusive) + .expect("rank should succeed"); + prop_assert!(rank >= last_rank, "rank {} after {}", rank, last_rank); + prop_assert!((0.0..=1.0).contains(&rank), "rank {} out of bounds", rank); + last_rank = rank; + } + } +} diff --git a/datasketches/tests/req_test/query.rs b/datasketches/tests/req_test/query.rs new file mode 100644 index 00000000..94fe2b16 --- /dev/null +++ b/datasketches/tests/req_test/query.rs @@ -0,0 +1,194 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Rank, quantile, PMF, and CDF behavior for ReqSketch. + +use approx::assert_relative_eq; +use datasketches::error::Error; +use datasketches::req::ReqSketch; +use datasketches::req::SearchCriteria; + +#[test] +fn exact_mode_rank_quantile_pmf_and_cdf_match_reference() { + let mut sketch = ReqSketch::new(); + for i in 1..=10 { + sketch.update(i as f32); + } + + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.n(), 10); + assert_eq!(sketch.num_retained(), 10); + + for (value, expected) in [(1.0, 0.0), (2.0, 0.1), (6.0, 0.5), (9.0, 0.8), (10.0, 0.9)] { + assert_relative_eq!( + sketch + .rank(&value, SearchCriteria::Exclusive) + .expect("rank should succeed"), + expected, + epsilon = 1e-6 + ); + } + + for (value, expected) in [(1.0, 0.1), (2.0, 0.2), (5.0, 0.5), (9.0, 0.9), (10.0, 1.0)] { + assert_relative_eq!( + sketch + .rank(&value, SearchCriteria::Inclusive) + .expect("rank should succeed"), + expected, + epsilon = 1e-6 + ); + } + + for (rank, expected) in [(0.0, 1.0), (0.1, 2.0), (0.5, 6.0), (0.9, 10.0), (1.0, 10.0)] { + assert_relative_eq!( + sketch + .quantile(rank, SearchCriteria::Exclusive) + .expect("quantile should succeed"), + expected, + epsilon = 1e-6 + ); + } + + for (rank, expected) in [(0.0, 1.0), (0.1, 1.0), (0.5, 5.0), (0.9, 9.0), (1.0, 10.0)] { + assert_relative_eq!( + sketch + .quantile(rank, SearchCriteria::Inclusive) + .expect("quantile should succeed"), + expected, + epsilon = 1e-6 + ); + } + + let splits = [2.0, 6.0, 9.0]; + let cdf = sketch + .cdf(&splits, SearchCriteria::Exclusive) + .expect("cdf should succeed"); + assert_relative_eq!(cdf[0], 0.1, epsilon = 1e-6); + assert_relative_eq!(cdf[1], 0.5, epsilon = 1e-6); + assert_relative_eq!(cdf[2], 0.8, epsilon = 1e-6); + assert_relative_eq!(cdf[3], 1.0, epsilon = 1e-6); + + let pmf = sketch + .pmf(&splits, SearchCriteria::Exclusive) + .expect("pmf should succeed"); + assert_relative_eq!(pmf[0], 0.1, epsilon = 1e-6); + assert_relative_eq!(pmf[1], 0.4, epsilon = 1e-6); + assert_relative_eq!(pmf[2], 0.3, epsilon = 1e-6); + assert_relative_eq!(pmf[3], 0.2, epsilon = 1e-6); +} + +#[test] +fn pmf_and_cdf_are_consistent() { + let mut sketch = ReqSketch::new(); + for i in 0..1000 { + sketch.update(i as f64); + } + + let split_points = [100.0, 300.0, 500.0, 700.0, 900.0]; + let pmf = sketch + .pmf(&split_points, SearchCriteria::Inclusive) + .expect("pmf should succeed"); + let cdf = sketch + .cdf(&split_points, SearchCriteria::Inclusive) + .expect("cdf should succeed"); + + assert_relative_eq!(pmf.iter().sum::(), 1.0, epsilon = 1e-10); + + let mut cumulative = 0.0; + for i in 0..pmf.len() { + cumulative += pmf[i]; + assert_relative_eq!(cdf[i], cumulative, epsilon = 1e-10); + } + assert_relative_eq!(cdf[cdf.len() - 1], 1.0, epsilon = 1e-10); +} + +#[test] +fn rank_is_monotonic_and_bounded() { + let mut sketch = ReqSketch::new(); + for i in 0..10_000 { + sketch.update(i as f64); + } + + let test_values: Vec = (0..10_000).step_by(1000).map(|i| i as f64).collect(); + let mut last_rank = 0.0; + + for value in test_values { + let rank = sketch + .rank(&value, SearchCriteria::Inclusive) + .expect("rank should succeed"); + assert!(rank >= last_rank, "ranks should be monotonic"); + assert!((0.0..=1.0).contains(&rank), "rank should be in [0,1]"); + last_rank = rank; + } +} + +#[test] +fn quantiles_are_monotonic() -> Result<(), Error> { + let mut sketch = ReqSketch::new(); + for i in 0..10_000 { + sketch.update(i as f64); + } + + let ranks = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]; + let mut previous = 0.0; + + for rank in ranks { + let quantile = sketch.quantile(rank, SearchCriteria::Inclusive)?; + assert!(quantile >= previous); + previous = quantile; + } + + Ok(()) +} + +#[test] +fn rank_quantile_round_trip_is_consistent() -> Result<(), Error> { + let mut sketch = ReqSketch::new(); + for i in 0..10_000 { + sketch.update(i as f64); + } + + for target_rank in [0.1, 0.25, 0.5, 0.75, 0.9] { + let quantile = sketch.quantile(target_rank, SearchCriteria::Inclusive)?; + let recovered_rank = sketch.rank(&quantile, SearchCriteria::Inclusive)?; + let error = (recovered_rank - target_rank).abs() / target_rank; + assert!(error < 0.2); + } + + Ok(()) +} + +#[test] +fn search_criteria_rank_consistency() -> Result<(), Error> { + let mut sketch = ReqSketch::new(); + for i in 0..1000 { + sketch.update(i as f64); + } + + for value in [100.0, 250.0, 500.0, 750.0] { + let inclusive_rank = sketch.rank(&value, SearchCriteria::Inclusive)?; + let exclusive_rank = sketch.rank(&value, SearchCriteria::Exclusive)?; + + assert!(exclusive_rank <= inclusive_rank); + assert!((0.0..=1.0).contains(&inclusive_rank)); + assert!((0.0..=1.0).contains(&exclusive_rank)); + } + + Ok(()) +} diff --git a/datasketches/tests/req_test/serialization.rs b/datasketches/tests/req_test/serialization.rs new file mode 100644 index 00000000..a7c07632 --- /dev/null +++ b/datasketches/tests/req_test/serialization.rs @@ -0,0 +1,374 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Serialization round-trip and cross-language compatibility tests for ReqSketch. + +use std::fs; +use std::path::PathBuf; + +use crate::serialization_test_data; +use datasketches::req::RankAccuracy; +use datasketches::req::ReqSketch; +use datasketches::req::ReqValue; +use datasketches::req::SearchCriteria; + +// ---------- Rust ↔ Rust round-trip ---------- + +fn round_trip_one(k: u16, ra: RankAccuracy, n: u64, make_item: impl Fn(u64) -> T) +where + T: ReqValue + std::fmt::Debug + PartialEq, +{ + let mut a: ReqSketch = ReqSketch::try_new(k, ra).unwrap(); + for i in 0..n { + a.update(make_item(i)); + } + let bytes = a.serialize(); + let b: ReqSketch = ReqSketch::deserialize(&bytes).unwrap(); + assert_eq!(a.n(), b.n()); + assert_eq!(a.k(), b.k()); + assert_eq!(a.rank_accuracy(), b.rank_accuracy()); + assert_eq!(a.min_item(), b.min_item()); + assert_eq!(a.max_item(), b.max_item()); + assert_eq!(bytes, b.serialize(), "non-stable serialization"); +} + +#[test] +fn round_trip_f64_matrix() { + for &k in &[4u16, 12, 1024] { + for &ra in &[RankAccuracy::HighRank, RankAccuracy::LowRank] { + for &n in &[0u64, 1, 4, 5, 100, 10_000] { + round_trip_one::(k, ra, n, |i| i as f64); + } + } + } +} + +#[test] +fn round_trip_f32_basic() { + for &n in &[0u64, 1, 4, 5, 1000] { + round_trip_one::(12, RankAccuracy::HighRank, n, |i| i as f32); + } +} + +#[test] +fn round_trip_i64_basic() { + for &n in &[0u64, 1, 4, 5, 1000] { + round_trip_one::(12, RankAccuracy::HighRank, n, |i| i as i64); + } +} + +// ---------- Deserialize error paths ---------- +// +// Each test crafts a malformed byte sequence and asserts that deserialize returns +// Err, exercising the validation guards in ReqSketch::deserialize. + +use datasketches::error::ErrorKind; + +#[test] +fn deserialize_truncated_preamble() { + // Less than 8 bytes — can't even read the fixed preamble. + for n in 0..8usize { + let bytes = vec![0u8; n]; + let result = ReqSketch::::deserialize(&bytes); + assert!( + result.is_err(), + "deserialize succeeded with {} bytes (expected error)", + n + ); + } +} + +#[test] +fn deserialize_wrong_family_id() { + // Valid preamble structure but family != 17. + // Flags=4 (IS_EMPTY), k=12 (little-endian: 12, 0). + let bytes = [ + 2u8, // preamble_ints (PREAMBLE_INTS_EXACT) + 1u8, // serial_version + 99u8, // family — wrong (REQ is 17) + 4u8, // flags (IS_EMPTY) + 12u8, 0u8, // k = 12 + 0u8, // num_levels + 0u8, // num_raw_items + ]; + let result = ReqSketch::::deserialize(&bytes); + assert!(result.is_err(), "deserialize accepted wrong family id"); + let err = result.unwrap_err(); + assert_eq!( + err.kind(), + ErrorKind::InvalidData, + "wrong error kind: {:?}", + err.kind() + ); +} + +#[test] +fn deserialize_wrong_serial_version() { + // Serial version != 1 should be rejected. + let bytes = [ + 2u8, 99u8, // serial_version — wrong (REQ uses 1) + 17u8, 4u8, // IS_EMPTY + 12u8, 0u8, 0u8, 0u8, + ]; + let result = ReqSketch::::deserialize(&bytes); + assert!(result.is_err(), "deserialize accepted wrong serial version"); +} + +#[test] +fn deserialize_invalid_preamble_ints() { + // preamble_ints must be 2 (exact) or 4 (estimation). Try 3. + let bytes = [3u8, 1, 17, 4, 12, 0, 0, 0]; + let result = ReqSketch::::deserialize(&bytes); + assert!( + result.is_err(), + "deserialize accepted invalid preamble_ints=3" + ); +} + +#[test] +fn deserialize_rejects_non_empty_zero_levels() { + // Non-empty flags with num_levels=0 used to create a sketch with n=1 but + // no level-0 compactor, causing the next update to panic. + let bytes = [ + 2u8, // PREAMBLE_INTS_EXACT + 1, 17, 8u8, // IS_HIGH_RANK only: not empty, not raw + 12, 0, // k + 0u8, // num_levels=0 is invalid for non-empty sketches + 0u8, + ]; + let result = ReqSketch::::deserialize(&bytes); + assert!( + result.is_err(), + "deserialize accepted non-empty sketch with zero levels" + ); +} + +#[test] +fn deserialize_rejects_inconsistent_raw_items_header() { + // RAW_ITEMS is only valid for one non-empty level with 1..=4 raw items. + let raw_with_no_items = [ + 2u8, 1, 17, 24u8, // IS_HIGH_RANK | RAW_ITEMS + 12, 0, 1u8, // num_levels + 0u8, // invalid raw item count + ]; + assert!( + ReqSketch::::deserialize(&raw_with_no_items).is_err(), + "deserialize accepted raw-items sketch with no raw items" + ); + + let raw_with_two_levels = [ + 4u8, 1, 17, 24u8, // IS_HIGH_RANK | RAW_ITEMS + 12, 0, 2u8, // invalid for raw-items sketches + 1u8, + ]; + assert!( + ReqSketch::::deserialize(&raw_with_two_levels).is_err(), + "deserialize accepted raw-items sketch with multiple levels" + ); +} + +#[test] +fn deserialize_odd_k() { + // k must be even. Try k=11. + let bytes = [ + 2u8, 1, 17, 4u8, // IS_EMPTY + 11u8, 0u8, // k=11 (odd) + 0u8, 0u8, + ]; + let result = ReqSketch::::deserialize(&bytes); + assert!(result.is_err(), "deserialize accepted odd k=11"); +} + +#[test] +fn deserialize_k_out_of_range() { + // k must be in [4, 1024]. Try k=2 (too small). + let bytes_small = [2u8, 1, 17, 4, 2, 0, 0, 0]; + assert!( + ReqSketch::::deserialize(&bytes_small).is_err(), + "accepted k=2" + ); + + // k=2048 (too large): little-endian 2048 = [0x00, 0x08] + let bytes_big = [2u8, 1, 17, 4, 0, 8, 0, 0]; + assert!( + ReqSketch::::deserialize(&bytes_big).is_err(), + "accepted k=2048" + ); +} + +#[test] +fn deserialize_truncated_estimation_mode() { + // preamble_ints=4, num_levels=2 (multi-level), not empty — code will try to read + // n (u64) + min_f32 + max_f32 + compactor preambles, but we provide nothing beyond + // the 8-byte preamble. + // flags=8 (IS_HIGH_RANK only — not empty, not raw). + let bytes = [ + 4u8, // PREAMBLE_INTS_ESTIMATION + 1, 17, 8u8, // IS_HIGH_RANK only (not empty, not raw) + 12, 0, // k + 2u8, // num_levels = 2 (triggers n/min/max read) + 0u8, /* num_raw_items + * no payload — truncated */ + ]; + let result = ReqSketch::::deserialize(&bytes); + assert!( + result.is_err(), + "deserialize accepted truncated estimation-mode bytes" + ); +} + +#[test] +fn deserialize_truncated_raw_items() { + // raw_items=true (FLAG_RAW_ITEMS=0x10), num_raw_items=3, but only 1 f32 follows. + // flags = IS_HIGH_RANK | RAW_ITEMS = 8 | 16 = 24, num_levels=1 + let bytes = [ + 2u8, 1, 17, 24u8, // IS_HIGH_RANK | RAW_ITEMS + 12, 0, 1u8, // num_levels=1 + 3u8, // num_raw_items=3 (but only 1 f32 supplied) + 0u8, 0, 0x80, 0x3f, // 1.0_f32 (only 1 of the 3 promised items) + ]; + let result = ReqSketch::::deserialize(&bytes); + assert!(result.is_err(), "deserialize accepted truncated raw_items"); +} + +// ---------- Deserialize hardening: malformed compactor fields ---------- +// +// A non-empty, non-raw, single-level sketch carries a full 20-byte compactor +// preamble whose `section_size_raw`, `lg_weight`, and `num_items` fields are read +// straight off the wire. Without bounds checks these crafted values either panic +// (arithmetic overflow) or trigger an unbounded allocation in `Compactor::deserialize`. + +/// Builds a non-empty, non-raw, single-level (`num_levels = 1`) REQ sketch image +/// with a fully specified compactor preamble, so an individual field can be made +/// malformed in isolation. With valid inputs the result deserializes successfully +/// (see `single_level_image_is_valid_baseline`). +fn single_level_image( + section_size_raw: f32, + lg_weight: u8, + num_sections: u8, + num_items: u32, + items: &[f32], +) -> Vec { + // Preamble (8 bytes): preamble_ints = 2 (EXACT, since num_levels == 1), + // serial_version = 1, family = 17 (REQ), flags = 8 (IS_HIGH_RANK: not empty, + // not raw), k = 12 (u16 LE), num_levels = 1, num_raw_items = 0. + let mut b = vec![2u8, 1, 17, 8, 12, 0, 1, 0]; + // Compactor preamble (20 bytes). + b.extend_from_slice(&0u64.to_le_bytes()); // state + b.extend_from_slice(§ion_size_raw.to_le_bytes()); + b.push(lg_weight); + b.push(num_sections); + b.extend_from_slice(&0u16.to_le_bytes()); // padding + b.extend_from_slice(&num_items.to_le_bytes()); + for &item in items { + b.extend_from_slice(&item.to_le_bytes()); + } + b +} + +#[test] +fn single_level_image_is_valid_baseline() { + // Control: the builder with well-formed fields round-trips, so the malformed + // variants below isolate exactly one bad field. + let bytes = single_level_image(12.0, 0, 3, 1, &[1.0]); + assert!( + ReqSketch::::deserialize(&bytes).is_ok(), + "baseline single-level image should deserialize" + ); +} + +#[test] +fn deserialize_rejects_out_of_range_section_size() { + // A garbage section_size_raw drives the `nominal_capacity` arithmetic to overflow. + let bytes = single_level_image(1e30, 0, 3, 1, &[1.0]); + assert!( + ReqSketch::::deserialize(&bytes).is_err(), + "deserialize accepted out-of-range section_size_raw" + ); +} + +#[test] +fn deserialize_rejects_oversized_lg_weight() { + // lg_weight >= 64 makes the per-item weight `1u64 << lg_weight` overflow. + let bytes = single_level_image(12.0, 64, 3, 1, &[1.0]); + assert!( + ReqSketch::::deserialize(&bytes).is_err(), + "deserialize accepted lg_weight = 64" + ); +} + +#[test] +fn deserialize_rejects_oversized_compactor_num_items() { + // num_items claims billions of items while only one is supplied: deserialize + // must fail gracefully without attempting a multi-gigabyte allocation. + let bytes = single_level_image(12.0, 0, 3, u32::MAX, &[1.0]); + assert!( + ReqSketch::::deserialize(&bytes).is_err(), + "deserialize accepted oversized num_items" + ); +} + +// ---------- Cross-language compatibility ---------- +// +// Requires fixtures generated by `tools/generate_serialization_test_data.py`. +// If `tests/serde_tests/{cpp,java}_generated_files/` is missing, the +// `serialization_test_data` helper panics with regeneration instructions. + +fn validate_cross_language_fixture(path: PathBuf, expected_n: u64) { + let bytes = + fs::read(&path).unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display())); + let sketch = ReqSketch::::deserialize(&bytes) + .unwrap_or_else(|e| panic!("deserialize failed for {}: {e}", path.display())); + + assert_eq!(sketch.n(), expected_n, "n mismatch on {}", path.display()); + assert_eq!(sketch.k(), 12, "k mismatch on {}", path.display()); + assert_eq!(sketch.rank_accuracy(), RankAccuracy::HighRank); + + if expected_n > 0 { + assert_eq!(sketch.min_item().copied(), Some(1.0_f32)); + assert_eq!(sketch.max_item().copied(), Some(expected_n as f32)); + let _ = sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(); + } + + let serialized = sketch.serialize(); + assert_eq!( + bytes, + serialized, + "byte mismatch on {} — wire format diverges from C++/Java", + path.display() + ); +} + +#[test] +fn cpp_compatibility() { + for n in [0u64, 1, 10, 100, 1000, 10000, 100000, 1000000] { + let path = + serialization_test_data("cpp_generated_files", &format!("req_float_n{n}_cpp.sk")); + validate_cross_language_fixture(path, n); + } +} + +#[test] +fn java_compatibility() { + for n in [0u64, 1, 10, 100, 1000, 10000, 100000, 1000000] { + let path = + serialization_test_data("java_generated_files", &format!("req_float_n{n}_java.sk")); + validate_cross_language_fixture(path, n); + } +} diff --git a/datasketches/tests/req_test/sorted_view_api.rs b/datasketches/tests/req_test/sorted_view_api.rs new file mode 100644 index 00000000..b8803a23 --- /dev/null +++ b/datasketches/tests/req_test/sorted_view_api.rs @@ -0,0 +1,165 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Tests for the user-managed SortedView API (ported from reqsketch-rs #25): +//! distribution queries take `&self`, and `sorted_view()` returns an owned +//! snapshot instead of relying on an internal cache. + +use datasketches::error::ErrorKind; +use datasketches::req::ReqSketch; +use datasketches::req::SearchCriteria; +use datasketches::req::SortedView; + +fn populated_sketch(n: u64) -> ReqSketch { + let mut sketch = ReqSketch::new(); + for i in 0..n { + sketch.update(i as f64); + } + sketch +} + +/// All distribution queries must work through a shared (`&self`) reference. +fn query_through_shared_ref(sketch: &ReqSketch) { + sketch + .quantile(0.5, SearchCriteria::Inclusive) + .expect("quantile"); + sketch + .quantiles(&[0.25, 0.5, 0.75], SearchCriteria::Inclusive) + .expect("quantiles"); + sketch.rank(&50.0, SearchCriteria::Inclusive).expect("rank"); + sketch + .pmf(&[10.0, 50.0], SearchCriteria::Inclusive) + .expect("pmf"); + sketch + .cdf(&[10.0, 50.0], SearchCriteria::Inclusive) + .expect("cdf"); + assert!(!sketch.sorted_view().is_empty()); +} + +#[test] +fn queries_work_through_shared_reference() { + let sketch = populated_sketch(100); + query_through_shared_ref(&sketch); +} + +#[test] +fn sorted_view_is_an_owned_snapshot() { + let mut sketch = populated_sketch(100); + + let view: SortedView = sketch.sorted_view(); + assert_eq!(view.total_weight(), 100); + + // Updating the sketch while the view is alive must compile (owned view) + // and must not affect the snapshot. + for i in 100..200 { + sketch.update(i as f64); + } + assert_eq!(view.total_weight(), 100); + + // A fresh view reflects the new state. + let fresh = sketch.sorted_view(); + assert_eq!(fresh.total_weight(), 200); +} + +#[test] +fn sorted_view_on_empty_sketch_is_an_empty_view() { + let sketch: ReqSketch = ReqSketch::new(); + let view = sketch.sorted_view(); + assert!(view.is_empty()); + assert_eq!(view.len(), 0); + assert_eq!(view.total_weight(), 0); + // Queries on the empty view still report an error. + assert!(view.quantile(0.5, SearchCriteria::Inclusive).is_err()); +} + +#[test] +fn empty_sketch_pmf_cdf_report_error() { + let sketch: ReqSketch = ReqSketch::new(); + assert!(sketch.pmf(&[1.0], SearchCriteria::Inclusive).is_err()); + assert!(sketch.cdf(&[1.0], SearchCriteria::Inclusive).is_err()); +} + +#[test] +fn view_rank_is_primary_query_name() { + let sketch = populated_sketch(10); + let view = sketch.sorted_view(); + let r = view.rank(&5.0, SearchCriteria::Inclusive).expect("rank"); + assert!((r - 0.6).abs() < 1e-10, "rank of 5.0 in 0..10, got {r}"); +} + +#[test] +fn nan_query_items_are_rejected() { + let sketch = populated_sketch(100); + let err = sketch + .rank(&f64::NAN, SearchCriteria::Inclusive) + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::InvalidArgument); + + let view = sketch.sorted_view(); + assert!(view.rank(&f64::NAN, SearchCriteria::Inclusive).is_err()); +} + +#[test] +fn error_precedence_empty_before_invalid_rank() { + // On an empty sketch the emptiness is reported before the out-of-range rank. + let empty: ReqSketch = ReqSketch::new(); + let empty_err = empty.quantile(2.0, SearchCriteria::Inclusive).unwrap_err(); + assert!( + empty_err.message().contains("empty"), + "expected emptiness error, got: {}", + empty_err.message() + ); + + // On a populated sketch the out-of-range rank is reported. + let sketch = populated_sketch(10); + let range_err = sketch.quantile(2.0, SearchCriteria::Inclusive).unwrap_err(); + assert_eq!(range_err.kind(), ErrorKind::InvalidArgument); + assert!( + range_err.message().contains("must be in"), + "expected range error, got: {}", + range_err.message() + ); +} + +#[test] +fn view_is_send_and_sync() { + fn assert_send_sync() {} + assert_send_sync::>(); + assert_send_sync::>(); +} + +#[test] +fn concurrent_readers_share_the_sketch() { + let sketch = std::sync::Arc::new(populated_sketch(1_000)); + let handles: Vec<_> = (0..4) + .map(|i| { + let sketch = std::sync::Arc::clone(&sketch); + std::thread::spawn(move || { + let rank = 0.2 * (i + 1) as f64; + sketch + .quantile(rank, SearchCriteria::Inclusive) + .expect("quantile from shared sketch") + }) + }) + .collect(); + for handle in handles { + let q = handle.join().expect("thread"); + assert!((0.0..1_000.0).contains(&q)); + } +} diff --git a/datasketches/tests/req_test/structure.rs b/datasketches/tests/req_test/structure.rs new file mode 100644 index 00000000..aa30cc3a --- /dev/null +++ b/datasketches/tests/req_test/structure.rs @@ -0,0 +1,77 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Public iterator behavior for ReqSketch. + +use datasketches::req::ReqSketch; + +#[test] +fn iterator_weights_sum_to_n_and_items_are_in_range() { + let mut sketch = ReqSketch::new(); + for i in 0..1000 { + sketch.update(i as f64); + } + + let total_weight: u64 = sketch.iter().map(|(_, weight)| weight).sum(); + assert_eq!(total_weight, sketch.n()); + + for (item, weight) in sketch.iter() { + assert!(weight >= 1); + assert!(item >= *sketch.min_item().expect("non-empty sketch")); + assert!(item <= *sketch.max_item().expect("non-empty sketch")); + } +} + +#[test] +fn small_sketch_iterator_reports_unit_weights() { + // Below the compaction threshold every retained item still has weight 1. + let mut sketch = ReqSketch::new(); + for i in 0..10 { + sketch.update(i as f64); + } + + let items: Vec<(f64, u64)> = sketch.iter().collect(); + assert_eq!(items.len(), 10); + assert!(items.iter().all(|&(_, weight)| weight == 1)); +} + +#[test] +fn empty_sketch_iterator_yields_nothing() { + let sketch: ReqSketch = ReqSketch::new(); + assert_eq!(sketch.iter().count(), 0); +} + +#[test] +fn compaction_promotes_surviving_items_to_higher_weights() { + // After enough updates to trigger compaction, surviving items are promoted up a + // level at double weight, so the maximum item weight exceeds 1. + let mut sketch = ReqSketch::new(); + for i in 0..100_000 { + sketch.update(i as f64); + } + + let max_weight = sketch.iter().map(|(_, weight)| weight).max().unwrap(); + assert!( + max_weight > 1, + "expected promoted items, got max weight {max_weight}" + ); + + // Every weight is a power of two (2^level). + assert!(sketch.iter().all(|(_, weight)| weight.is_power_of_two())); +} diff --git a/datasketches/tests/req_test/union.rs b/datasketches/tests/req_test/union.rs new file mode 100644 index 00000000..b069f7a1 --- /dev/null +++ b/datasketches/tests/req_test/union.rs @@ -0,0 +1,116 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "req")] + +//! Integration tests for ReqUnion. + +use datasketches::req::RankAccuracy; +use datasketches::req::ReqSketch; +use datasketches::req::ReqUnion; +use datasketches::req::SearchCriteria; + +#[test] +fn union_equivalent_to_chained_merge() { + let make = |range: std::ops::Range| -> ReqSketch { + let mut sketch = ReqSketch::new(); + for i in range { + sketch.update(i as f64); + } + sketch + }; + let s1 = make(0..1000); + let s2 = make(1000..2000); + let s3 = make(2000..3000); + + let mut union: ReqUnion = ReqUnion::new(); + union.merge(&s1).expect("merge should succeed"); + union.merge(&s2).expect("merge should succeed"); + union.merge(&s3).expect("merge should succeed"); + let union_result = union.to_sketch(); + + let mut via_merge: ReqSketch = ReqSketch::new(); + via_merge.merge(&s1).expect("merge should succeed"); + via_merge.merge(&s2).expect("merge should succeed"); + via_merge.merge(&s3).expect("merge should succeed"); + + assert_eq!(union_result.n(), via_merge.n()); + assert_eq!(union_result.min_item(), via_merge.min_item()); + assert_eq!(union_result.max_item(), via_merge.max_item()); + + let true_median = 1499.5_f64; + let tolerance = 0.05 * 3000.0; + let q_union = union_result + .quantile(0.5, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + let q_merge = via_merge + .quantile(0.5, SearchCriteria::Inclusive) + .expect("quantile should succeed"); + + assert!((q_union - true_median).abs() <= tolerance); + assert!((q_merge - true_median).abs() <= tolerance); +} + +#[test] +fn empty_union_returns_empty_sketch() { + let union: ReqUnion = ReqUnion::new(); + assert!(union.to_sketch().is_empty()); +} + +#[test] +fn reset_clears_union_state() { + let mut sketch: ReqSketch = ReqSketch::new(); + for i in 0..100 { + sketch.update(i as f64); + } + + let mut union: ReqUnion = ReqUnion::new(); + union.merge(&sketch).expect("merge should succeed"); + assert!(!union.is_empty()); + + union.reset(); + assert!(union.is_empty()); +} + +#[test] +fn try_new_validates_k() { + assert!(ReqUnion::::try_new(3, RankAccuracy::HighRank).is_err()); + assert!(ReqUnion::::try_new(12, RankAccuracy::HighRank).is_ok()); +} + +#[test] +fn empty_union_uses_default_configuration() { + let union: ReqUnion = ReqUnion::new(); + assert_eq!(union.k(), 12); + assert_eq!(union.rank_accuracy(), RankAccuracy::HighRank); +} + +#[test] +fn union_keeps_default_k_when_merging_mismatched_sketch() { + // The union retains its own k even when fed a sketch built with a different k. + let mut other = ReqSketch::::try_new(16, RankAccuracy::HighRank).expect("valid k"); + for i in 0..50 { + other.update(i as f64); + } + + let mut union: ReqUnion = ReqUnion::new(); + union.merge(&other).expect("merge should succeed"); + + let result = union.to_sketch(); + assert_eq!(result.k(), 12); + assert_eq!(result.n(), 50); +} From 1f27cd6d921f9f54365639e0628680fb6022a95e Mon Sep 17 00:00:00 2001 From: Patrick McGleenon Date: Tue, 18 Aug 2026 13:58:19 +0100 Subject: [PATCH 2/2] feat(req): cargo +nightly fmt --all --- datasketches/tests/req_test/serialization.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/datasketches/tests/req_test/serialization.rs b/datasketches/tests/req_test/serialization.rs index a7c07632..ca5893c9 100644 --- a/datasketches/tests/req_test/serialization.rs +++ b/datasketches/tests/req_test/serialization.rs @@ -22,12 +22,13 @@ use std::fs; use std::path::PathBuf; -use crate::serialization_test_data; use datasketches::req::RankAccuracy; use datasketches::req::ReqSketch; use datasketches::req::ReqValue; use datasketches::req::SearchCriteria; +use crate::serialization_test_data; + // ---------- Rust ↔ Rust round-trip ---------- fn round_trip_one(k: u16, ra: RankAccuracy, n: u64, make_item: impl Fn(u64) -> T)