diff --git a/CHANGELOG.md b/CHANGELOG.md index 546d83de..f65b2678 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,9 +4,17 @@ All significant changes to this project will be documented in this file. ## Unreleased +### Breaking changes + +* Move `SearchCriteria` from `req` to `common` and remove its `Default` implementation. Import `datasketches::common::SearchCriteria` and explicitly choose `Inclusive` or `Exclusive` for each query. + ### New features -* Add KLL sketches behind the `kll` feature, including rank, quantile, PMF, and CDF queries, custom item ordering, merging, and C++/Java-compatible serialization. +* Add KLL sketches behind the `kll` feature, with rank, quantile, PMF, and CDF queries, merging, totally ordered custom item types, a `KllFloat` adapter for non-NaN floating-point values, and serialization. + +### Improvements + +* Improve truncated-input diagnostics across sketch deserializers. ### Bug fixes diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 115320ab..3bbbda5c 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -23,7 +23,7 @@ edition.workspace = true rust-version.workspace = true [dev-dependencies] -datasketches = { workspace = true, features = ["cpc", "req", "tdigest"] } +datasketches = { workspace = true, features = ["cpc", "kll", "req", "tdigest"] } divan = { workspace = true } rand = { workspace = true } diff --git a/benchmarks/kll/merge.rs b/benchmarks/kll/merge.rs new file mode 100644 index 00000000..c35a1afb --- /dev/null +++ b/benchmarks/kll/merge.rs @@ -0,0 +1,38 @@ +// 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 divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use super::support::build_sketch; +use super::support::values; + +#[divan::bench] +fn merge(bencher: Bencher) { + let values = values(200_000); + let left = build_sketch(&values[..100_000]); + let right = build_sketch(&values[100_000..]); + + bencher + .counter(ItemsCount::new(values.len())) + .with_inputs(|| left.clone()) + .bench_local_values(|mut left| { + left.merge(black_box(&right)).unwrap(); + black_box(left) + }); +} diff --git a/benchmarks/kll/mod.rs b/benchmarks/kll/mod.rs new file mode 100644 index 00000000..4ff11ff9 --- /dev/null +++ b/benchmarks/kll/mod.rs @@ -0,0 +1,22 @@ +// 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. + +mod merge; +mod query; +mod serde; +mod support; +mod update; diff --git a/benchmarks/kll/query.rs b/benchmarks/kll/query.rs new file mode 100644 index 00000000..1c61aa08 --- /dev/null +++ b/benchmarks/kll/query.rs @@ -0,0 +1,50 @@ +// 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 datasketches::common::SearchCriteria; +use datasketches::kll::KllFloat; +use divan::Bencher; +use divan::black_box; + +use super::support::prepared_sketch; + +#[divan::bench] +fn rank(bencher: Bencher) { + let sketch = prepared_sketch(); + let item = KllFloat::::new(500_000.0).unwrap(); + bencher.bench_local(|| black_box(&sketch).rank(black_box(&item), SearchCriteria::Inclusive)); +} + +#[divan::bench] +fn quantile(bencher: Bencher) { + let sketch = prepared_sketch(); + bencher.bench_local(|| black_box(&sketch).quantile(black_box(0.5), SearchCriteria::Inclusive)); +} + +#[divan::bench] +fn sorted_view_quantile(bencher: Bencher) { + let view = prepared_sketch().sorted_view(); + bencher.bench_local(|| black_box(&view).quantile(black_box(0.5), SearchCriteria::Inclusive)); +} + +#[divan::bench] +fn batch_quantiles(bencher: Bencher) { + let sketch = prepared_sketch(); + let ranks = [0.01, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99]; + bencher + .bench_local(|| black_box(&sketch).quantiles(black_box(&ranks), SearchCriteria::Inclusive)); +} diff --git a/benchmarks/kll/serde.rs b/benchmarks/kll/serde.rs new file mode 100644 index 00000000..3882f23d --- /dev/null +++ b/benchmarks/kll/serde.rs @@ -0,0 +1,41 @@ +// 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 datasketches::kll::KllFloat; +use datasketches::kll::KllSketch; +use divan::Bencher; +use divan::black_box; +use divan::counter::BytesCount; + +use super::support::prepared_sketch; + +#[divan::bench] +fn serialize(bencher: Bencher) { + let sketch = prepared_sketch(); + let bytes = sketch.serialize(); + bencher + .counter(BytesCount::new(bytes.len())) + .bench_local(|| black_box(&sketch).serialize()); +} + +#[divan::bench] +fn deserialize(bencher: Bencher) { + let bytes = prepared_sketch().serialize(); + bencher + .counter(BytesCount::new(bytes.len())) + .bench_local(|| KllSketch::>::deserialize(black_box(&bytes)).unwrap()); +} diff --git a/benchmarks/kll/support.rs b/benchmarks/kll/support.rs new file mode 100644 index 00000000..20708f19 --- /dev/null +++ b/benchmarks/kll/support.rs @@ -0,0 +1,43 @@ +// 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 datasketches::kll::KllFloat; +use datasketches::kll::KllSketch; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; + +pub(super) const DEFAULT_K: u16 = 200; + +pub(super) fn values(len: usize) -> Vec { + let mut rng = StdRng::seed_from_u64(42); + (0..len) + .map(|_| rng.random_range(0.0..1_000_000.0)) + .collect() +} + +pub(super) fn build_sketch(values: &[f64]) -> KllSketch> { + let mut sketch = KllSketch::new(DEFAULT_K).unwrap(); + for &value in values { + sketch.update(KllFloat::::new(value).unwrap()); + } + sketch +} + +pub(super) fn prepared_sketch() -> KllSketch> { + build_sketch(&values(100_000)) +} diff --git a/benchmarks/kll/update.rs b/benchmarks/kll/update.rs new file mode 100644 index 00000000..6a3c865e --- /dev/null +++ b/benchmarks/kll/update.rs @@ -0,0 +1,31 @@ +// 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 divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use super::support::build_sketch; +use super::support::values; + +#[divan::bench(args = [1_000, 10_000, 100_000])] +fn update(bencher: Bencher, len: usize) { + let values = values(len); + bencher + .counter(ItemsCount::new(len)) + .bench_local(|| build_sketch(black_box(&values))); +} diff --git a/benchmarks/main.rs b/benchmarks/main.rs index 4f8ed78c..11c99621 100644 --- a/benchmarks/main.rs +++ b/benchmarks/main.rs @@ -21,6 +21,7 @@ use divan::AllocProfiler; static ALLOC: AllocProfiler = AllocProfiler::system(); mod cpc; +mod kll; mod req; mod tdigest; diff --git a/benchmarks/req/query.rs b/benchmarks/req/query.rs index 916f5e88..b52dfc52 100644 --- a/benchmarks/req/query.rs +++ b/benchmarks/req/query.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. +use datasketches::common::SearchCriteria; use datasketches::req::ReqFloat; -use datasketches::req::SearchCriteria; use divan::Bencher; use divan::black_box; diff --git a/datasketches/src/bloom/sketch.rs b/datasketches/src/bloom/sketch.rs index 6a119e91..8bf39220 100644 --- a/datasketches/src/bloom/sketch.rs +++ b/datasketches/src/bloom/sketch.rs @@ -495,11 +495,12 @@ impl BloomFilter { .checked_add(1) .and_then(|words| words.checked_mul(size_of::())) .ok_or_else(|| Error::deserial("Bloom filter payload length overflows"))?; - if payload_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "Bloom filter payload requires {payload_bytes} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < payload_bytes { + return Err(Error::insufficient_data_of( + "Bloom filter payload", + format_args!("expected {payload_bytes} bytes, got {available_bytes}"), + )); } } let mut bit_array = vec![0u64; num_words].into_boxed_slice(); diff --git a/datasketches/src/codec/assert.rs b/datasketches/src/codec/assert.rs index 71d18ec7..523d43df 100644 --- a/datasketches/src/codec/assert.rs +++ b/datasketches/src/codec/assert.rs @@ -21,7 +21,7 @@ use std::ops::RangeBounds; use crate::error::Error; pub fn insufficient_data(tag: &'static str) -> impl FnOnce(std::io::Error) -> Error { - move |_| Error::insufficient_data(tag) + move |error| Error::insufficient_data_of(tag, error) } pub fn ensure_serial_version_is(expected: u8, actual: u8) -> Result<(), Error> { diff --git a/datasketches/src/common/mod.rs b/datasketches/src/common/mod.rs index 6d4c6c6e..918c57b5 100644 --- a/datasketches/src/common/mod.rs +++ b/datasketches/src/common/mod.rs @@ -19,8 +19,10 @@ mod num_std_dev; mod resize; +mod search_criteria; pub use self::num_std_dev::NumStdDev; pub use self::resize::ResizeFactor; +pub use self::search_criteria::SearchCriteria; #[cfg(any(feature = "cpc", feature = "hll"))] pub(crate) mod inv_pow2; diff --git a/datasketches/src/common/search_criteria.rs b/datasketches/src/common/search_criteria.rs new file mode 100644 index 00000000..3e98ca65 --- /dev/null +++ b/datasketches/src/common/search_criteria.rs @@ -0,0 +1,25 @@ +// 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. + +/// Selects the rank definition used by rank, quantile, PMF, and CDF queries. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SearchCriteria { + /// Define rank as the fraction of values less than or equal to the boundary. + Inclusive, + /// Define rank as the fraction of values strictly less than the boundary. + Exclusive, +} diff --git a/datasketches/src/countmin/sketch.rs b/datasketches/src/countmin/sketch.rs index 7970c0c9..500f4396 100644 --- a/datasketches/src/countmin/sketch.rs +++ b/datasketches/src/countmin/sketch.rs @@ -430,11 +430,12 @@ impl CountMinSketch { let payload_bytes = payload_values .checked_mul(LONG_SIZE_BYTES) .ok_or_else(|| Error::deserial("CountMin payload size overflows"))?; - if payload_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "CountMin payload requires {payload_bytes} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < payload_bytes { + return Err(Error::insufficient_data_of( + "CountMin payload", + format_args!("expected {payload_bytes} bytes, got {available_bytes}"), + )); } } diff --git a/datasketches/src/cpc/sketch.rs b/datasketches/src/cpc/sketch.rs index a9db2020..2cc4a7a6 100644 --- a/datasketches/src/cpc/sketch.rs +++ b/datasketches/src/cpc/sketch.rs @@ -775,10 +775,14 @@ impl CpcSketch { let payload_bytes = window_data_bytes .checked_add(table_data_bytes) .ok_or_else(|| Error::deserial("CPC payload length overflows"))?; - let payload = cursor - .remaining() - .get(..payload_bytes) - .ok_or_else(|| Error::deserial("insufficient data for CPC compressed payload"))?; + let available_bytes = cursor.remaining().len(); + if available_bytes < payload_bytes { + return Err(Error::insufficient_data_of( + "CPC compressed payload", + format_args!("expected {payload_bytes} bytes, got {available_bytes}"), + )); + } + let payload = &cursor.remaining()[..payload_bytes]; let (window_data, table_data) = payload.split_at(window_data_bytes); let (table, window) = match flavor { Flavor::Empty => (PairTable::new(2, lg_k + 6), vec![]), diff --git a/datasketches/src/frequencies/serialization.rs b/datasketches/src/frequencies/serialization.rs index 5ef3bf2b..d5a8b615 100644 --- a/datasketches/src/frequencies/serialization.rs +++ b/datasketches/src/frequencies/serialization.rs @@ -19,6 +19,7 @@ use std::hash::Hash; use crate::codec::SketchBytes; use crate::codec::SketchSlice; +use crate::codec::assert::insufficient_data; use crate::error::Error; /// Serialization version. @@ -54,24 +55,21 @@ impl FrequentItemValue for String { } fn deserialize_value(cursor: &mut SketchSlice<'_>) -> Result { - let len = cursor.read_u32_le().map_err(|_| { - Error::insufficient_data("failed to read string item length".to_string()) - })? as usize; - - let remaining = cursor.remaining().len(); - if len > remaining { - return Err(Error::insufficient_data(format!( - "string item length ({len}) exceeds the remaining {remaining} bytes" - ))); + let len = cursor + .read_u32_le() + .map_err(insufficient_data("string item length"))? as usize; + let available_bytes = cursor.remaining().len(); + if available_bytes < len { + return Err(Error::insufficient_data_of( + "string item payload", + format_args!("expected {len} bytes, got {available_bytes}"), + )); } - - let mut slice = vec![0; len]; - cursor.read_exact(&mut slice).map_err(|_| { - Error::insufficient_data("failed to read string item bytes".to_string()) - })?; - - String::from_utf8(slice) - .map_err(|_| Error::deserial("invalid UTF-8 string payload".to_string())) + let value = std::str::from_utf8(&cursor.remaining()[..len]) + .map(str::to_owned) + .map_err(|_| Error::deserial("invalid UTF-8 string payload"))?; + cursor.advance(len as u64); + Ok(value) } } @@ -87,11 +85,9 @@ macro_rules! impl_primitive { } fn deserialize_value(cursor: &mut SketchSlice<'_>) -> Result { - cursor.$read().map_err(|_| { - Error::insufficient_data( - concat!("failed to read ", stringify!($name), " item bytes").to_string(), - ) - }) + cursor + .$read() + .map_err(insufficient_data(concat!(stringify!($name), " item"))) } } }; diff --git a/datasketches/src/frequencies/sketch.rs b/datasketches/src/frequencies/sketch.rs index 31383e91..2b163f27 100644 --- a/datasketches/src/frequencies/sketch.rs +++ b/datasketches/src/frequencies/sketch.rs @@ -679,20 +679,21 @@ impl FrequentItemsSketch { // Each active item has an eight-byte weight before its encoded key. Check // that lower bound before trusting the count for `Vec` preallocation. - let weight_bytes = active_items.checked_mul(size_of::()); - if !weight_bytes.is_some_and(|needed| needed <= cursor.remaining().len()) { - return Err(Error::insufficient_data(format!( - "active_items ({active_items}) exceeds the remaining {} bytes", - cursor.remaining().len() - ))); + let weight_bytes = active_items + .checked_mul(size_of::()) + .ok_or_else(|| Error::deserial("frequent item weight payload length overflows"))?; + let available_bytes = cursor.remaining().len(); + if available_bytes < weight_bytes { + return Err(Error::insufficient_data_of( + "frequent item weights", + format_args!("expected {weight_bytes} bytes, got {available_bytes}"), + )); } let mut values = Vec::with_capacity(active_items); for i in 0..active_items { - values.push(cursor.read_u64_le().map_err(|_| { - Error::insufficient_data(format!( - "expected {active_items} weights, failed at index {i}" - )) + values.push(cursor.read_u64_le().map_err(|error| { + Error::insufficient_data_of("frequent item weight", error).with_context("index", i) })?); } @@ -783,11 +784,8 @@ impl FrequentItemsSketch { Self::deserialize_inner(bytes, |mut cursor, num_items| { let mut items = Vec::with_capacity(num_items); for i in 0..num_items { - let item = T::deserialize_value(&mut cursor).map_err(|_| { - Error::insufficient_data(format!( - "expected {num_items} items, failed to read item at index {i}" - )) - })?; + let item = T::deserialize_value(&mut cursor) + .map_err(|error| error.with_context("item index", i))?; items.push(item); } Ok(items) diff --git a/datasketches/src/hll/array4.rs b/datasketches/src/hll/array4.rs index a1175608..2380e001 100644 --- a/datasketches/src/hll/array4.rs +++ b/datasketches/src/hll/array4.rs @@ -356,11 +356,12 @@ impl Array4 { .checked_mul(COUPON_SIZE_BYTES) .and_then(|aux_bytes| num_bytes.checked_add(aux_bytes)) .ok_or_else(|| Error::deserial("HLL4 payload length overflows"))?; - if required_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "HLL4 payload requires {required_bytes} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < required_bytes { + return Err(Error::insufficient_data_of( + "HLL4 payload", + format_args!("expected {required_bytes} bytes, got {available_bytes}"), + )); } // Read packed 4-bit byte array @@ -375,10 +376,9 @@ impl Array4 { let mut aux = AuxMap::new(lg_config_k); let mut decoded_count = 0; for i in 0..aux_slots { - let coupon = cursor.read_u32_le().map_err(|_| { - Error::insufficient_data(format!( - "expected {aux_slots} HLL4 auxiliary slots, failed at index {i}", - )) + let coupon = cursor.read_u32_le().map_err(|error| { + Error::insufficient_data_of("HLL4 auxiliary slot", error) + .with_context("index", i) })?; let coupon = Coupon(coupon); if coupon.is_empty() && !compact { diff --git a/datasketches/src/hll/array6.rs b/datasketches/src/hll/array6.rs index 70c519fc..3654bb80 100644 --- a/datasketches/src/hll/array6.rs +++ b/datasketches/src/hll/array6.rs @@ -203,11 +203,12 @@ impl Array6 { "HLL6 zero count must not exceed k and auxiliary count must be zero", )); } - if num_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "HLL6 payload requires {num_bytes} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < num_bytes { + return Err(Error::insufficient_data_of( + "HLL6 payload", + format_args!("expected {num_bytes} bytes, got {available_bytes}"), + )); } // Read packed byte array from offset HLL_BYTE_ARR_START diff --git a/datasketches/src/hll/array8.rs b/datasketches/src/hll/array8.rs index 56756400..e28c7241 100644 --- a/datasketches/src/hll/array8.rs +++ b/datasketches/src/hll/array8.rs @@ -275,11 +275,12 @@ impl Array8 { "HLL8 zero count must not exceed k and auxiliary count must be zero", )); } - if k > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "HLL8 payload requires {k} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < k { + return Err(Error::insufficient_data_of( + "HLL8 payload", + format_args!("expected {k} bytes, got {available_bytes}"), + )); } // Read byte array from offset HLL_BYTE_ARR_START diff --git a/datasketches/src/hll/hash_set.rs b/datasketches/src/hll/hash_set.rs index 66b4714d..dd2f5be5 100644 --- a/datasketches/src/hll/hash_set.rs +++ b/datasketches/src/hll/hash_set.rs @@ -111,11 +111,12 @@ impl HashSet { } let read_count = if compact { coupon_count } else { array_size }; let required_bytes = read_count * size_of::(); - if required_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "SET mode coupons require {required_bytes} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < required_bytes { + return Err(Error::insufficient_data_of( + "HLL SET mode coupons", + format_args!("expected {required_bytes} bytes, got {available_bytes}"), + )); } if compact { @@ -123,10 +124,9 @@ impl HashSet { // Create a new hash set and insert coupons one by one let mut hash_set = HashSet::new(lg_arr); for i in 0..coupon_count { - let coupon = cursor.read_u32_le().map_err(|_| { - Error::insufficient_data(format!( - "expected {coupon_count} coupons, failed at index {i}" - )) + let coupon = cursor.read_u32_le().map_err(|error| { + Error::insufficient_data_of("HLL SET mode coupon", error) + .with_context("index", i) })?; hash_set.update(Coupon(coupon)); } @@ -139,10 +139,9 @@ impl HashSet { // Read entire hash table including empty slots let mut coupons = vec![Coupon::EMPTY; array_size]; for (i, coupon) in coupons.iter_mut().enumerate() { - let raw = cursor.read_u32_le().map_err(|_| { - Error::insufficient_data(format!( - "expected {array_size} coupons, failed at index {i}" - )) + let raw = cursor.read_u32_le().map_err(|error| { + Error::insufficient_data_of("HLL SET mode coupon", error) + .with_context("index", i) })?; *coupon = Coupon(raw); } diff --git a/datasketches/src/hll/list.rs b/datasketches/src/hll/list.rs index afed99a0..47370fa5 100644 --- a/datasketches/src/hll/list.rs +++ b/datasketches/src/hll/list.rs @@ -99,21 +99,23 @@ impl List { } let read_count = if compact { coupon_count } else { array_size }; let required_bytes = read_count * size_of::(); - if !empty && required_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "LIST mode coupons require {required_bytes} bytes, got {}", - cursor.remaining().len() - ))); + if !empty { + let available_bytes = cursor.remaining().len(); + if available_bytes < required_bytes { + return Err(Error::insufficient_data_of( + "HLL LIST mode coupons", + format_args!("expected {required_bytes} bytes, got {available_bytes}"), + )); + } } // Read coupons into the front of the full-sized array; remaining slots stay Coupon::EMPTY. let mut coupons = vec![Coupon::EMPTY; array_size]; if !empty && coupon_count > 0 { for (i, coupon) in coupons.iter_mut().take(read_count).enumerate() { - let raw = cursor.read_u32_le().map_err(|_| { - Error::insufficient_data(format!( - "expect {coupon_count} coupons, failed at index {i}" - )) + let raw = cursor.read_u32_le().map_err(|error| { + Error::insufficient_data_of("HLL LIST mode coupon", error) + .with_context("index", i) })?; *coupon = Coupon(raw); } diff --git a/datasketches/src/kll/capacity.rs b/datasketches/src/kll/capacity.rs new file mode 100644 index 00000000..e4ad7f1f --- /dev/null +++ b/datasketches/src/kll/capacity.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. + +const POWERS_OF_THREE: [u64; 31] = [ + 1, + 3, + 9, + 27, + 81, + 243, + 729, + 2187, + 6561, + 19683, + 59049, + 177147, + 531441, + 1594323, + 4782969, + 14348907, + 43046721, + 129140163, + 387420489, + 1162261467, + 3486784401, + 10460353203, + 31381059609, + 94143178827, + 282429536481, + 847288609443, + 2541865828329, + 7625597484987, + 22876792454961, + 68630377364883, + 205891132094649, +]; + +const MAX_DEPTH: usize = 60; +const MAX_SHALLOW_DEPTH: usize = POWERS_OF_THREE.len() - 1; + +pub fn total_capacity(k: u16, minimum_capacity: u8, num_levels: usize) -> u32 { + validate_inputs(k, minimum_capacity, num_levels); + let mut total: u32 = 0; + for level in 0..num_levels { + total += level_capacity_unchecked(k, num_levels, level, minimum_capacity); + } + total +} + +pub fn level_capacity(k: u16, num_levels: usize, level: usize, minimum_capacity: u8) -> u32 { + validate_inputs(k, minimum_capacity, num_levels); + assert!( + level < num_levels, + "KLL level index must be in [0, {num_levels}), got {level}" + ); + level_capacity_unchecked(k, num_levels, level, minimum_capacity) +} + +fn validate_inputs(k: u16, minimum_capacity: u8, num_levels: usize) { + assert!( + (1..=MAX_DEPTH + 1).contains(&num_levels), + "KLL number of levels must be in [1, {}], got {num_levels}", + MAX_DEPTH + 1 + ); + assert!( + minimum_capacity as u16 <= k, + "KLL minimum level capacity must not exceed k: minimum capacity {minimum_capacity}, k {k}" + ); +} + +const fn level_capacity_unchecked( + k: u16, + num_levels: usize, + level: usize, + minimum_capacity: u8, +) -> u32 { + let depth = num_levels - level - 1; + let capacity = capacity_at_depth(k, depth) as u32; + if capacity < minimum_capacity as u32 { + minimum_capacity as u32 + } else { + capacity + } +} + +const fn capacity_at_depth(k: u16, depth: usize) -> u16 { + if depth <= MAX_SHALLOW_DEPTH { + return capacity_at_shallow_depth(k, depth); + } + let first_depth = depth / 2; + let remaining_depth = depth - first_depth; + let intermediate_capacity = capacity_at_shallow_depth(k, first_depth); + capacity_at_shallow_depth(intermediate_capacity, remaining_depth) +} + +const fn capacity_at_shallow_depth(k: u16, depth: usize) -> u16 { + let scaled_capacity = ((k as u64) << (depth + 1)) / POWERS_OF_THREE[depth]; + ((scaled_capacity + 1) >> 1) as u16 +} diff --git a/datasketches/src/kll/helper.rs b/datasketches/src/kll/helper.rs deleted file mode 100644 index 6b993aa7..00000000 --- a/datasketches/src/kll/helper.rs +++ /dev/null @@ -1,99 +0,0 @@ -// 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. - -const POWERS_OF_THREE: [u64; 31] = [ - 1, - 3, - 9, - 27, - 81, - 243, - 729, - 2187, - 6561, - 19683, - 59049, - 177147, - 531441, - 1594323, - 4782969, - 14348907, - 43046721, - 129140163, - 387420489, - 1162261467, - 3486784401, - 10460353203, - 31381059609, - 94143178827, - 282429536481, - 847288609443, - 2541865828329, - 7625597484987, - 22876792454961, - 68630377364883, - 205891132094649, -]; - -pub(super) fn compute_total_capacity(k: u16, m: u8, num_levels: usize) -> u32 { - let mut total: u32 = 0; - for level in 0..num_levels { - total += level_capacity(k, num_levels, level, m); - } - total -} - -pub(super) fn level_capacity(k: u16, num_levels: usize, height: usize, min_wid: u8) -> u32 { - assert!(height < num_levels, "height must be < num_levels"); - let depth = num_levels - height - 1; - let cap = int_cap_aux(k, depth as u8); - std::cmp::max(min_wid as u32, cap as u32) -} - -fn int_cap_aux(k: u16, depth: u8) -> u16 { - if depth > 60 { - panic!("depth must be <= 60"); - } - if depth <= 30 { - return int_cap_aux_aux(k, depth); - } - let half = depth / 2; - let rest = depth - half; - let tmp = int_cap_aux_aux(k, half); - int_cap_aux_aux(tmp, rest) -} - -fn int_cap_aux_aux(k: u16, depth: u8) -> u16 { - if depth > 30 { - panic!("depth must be <= 30"); - } - let twok = (k as u64) << 1; - let tmp = (twok << depth) / POWERS_OF_THREE[depth as usize]; - let result = (tmp + 1) >> 1; - assert!(result <= k as u64, "capacity result exceeds k"); - result as u16 -} - -pub(super) fn sum_the_sample_weights(level_sizes: &[usize]) -> u64 { - let mut total = 0u64; - let mut weight = 1u64; - for &size in level_sizes { - total += weight * size as u64; - weight <<= 1; - } - total -} diff --git a/datasketches/src/kll/mod.rs b/datasketches/src/kll/mod.rs index 47cfd23f..469bc7f1 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -21,45 +21,40 @@ //! near-optimal accuracy per retained item. It supports one-pass updates, //! approximate quantiles, ranks, PMF, and CDF queries. //! -//! This implementation follows Apache DataSketches semantics (Java KllSketch -//! / KllPreambleUtil, C++ kll_sketch) and uses the same binary serialization -//! format as those implementations. +//! This implementation follows Apache DataSketches semantics and uses the compact binary +//! serialization format shared by the Java, C++, and Go implementations. +//! +//! Items must implement [`Ord`]. Wrap `f32` or `f64` values in [`KllFloat`], which rejects NaN and +//! provides their ordinary numerical order. Custom ordering should be expressed with a newtype +//! that implements [`Ord`], keeping the ordering semantics part of the item type. //! //! # Usage //! //! ```rust +//! # use datasketches::common::SearchCriteria; //! # use datasketches::kll::KllSketch; -//! let mut sketch = KllSketch::::new(200).unwrap(); -//! sketch.update(1.0); -//! sketch.update(2.0); -//! let q = sketch.quantile(0.5, true).unwrap(); -//! assert!(q >= 1.0 && q <= 2.0); +//! let mut sketch = KllSketch::::new(200).unwrap(); +//! sketch.update(1); +//! sketch.update(2); +//! let q = sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(); +//! assert!((1..=2).contains(&q)); //! ``` -mod helper; +mod capacity; mod serialization; mod sketch; mod sorted_view; +mod value; -pub use self::sketch::KllComparator; -pub use self::sketch::KllItem; pub use self::sketch::KllSketch; -pub use self::sketch::NaturalOrder; - -/// KLL sketch specialized for `f64`. -pub type KllSketchF64 = KllSketch; -/// KLL sketch specialized for `f32`. -pub type KllSketchF32 = KllSketch; -/// KLL sketch specialized for `i64`. -pub type KllSketchI64 = KllSketch; -/// KLL sketch specialized for `String`. -pub type KllSketchString = KllSketch; - +pub use self::sorted_view::SortedView; +pub use self::value::KllFloat; +pub use self::value::KllValue; /// Default value of parameter k. -pub const DEFAULT_K: u16 = 200; +const DEFAULT_K: u16 = 200; /// Default value of parameter m. -pub const DEFAULT_M: u8 = 8; +const DEFAULT_M: u8 = 8; /// Minimum value of parameter k. -pub const MIN_K: u16 = DEFAULT_M as u16; +const MIN_K: u16 = DEFAULT_M as u16; /// Maximum value of parameter k. -pub const MAX_K: u16 = u16::MAX; +const MAX_K: u16 = u16::MAX; diff --git a/datasketches/src/kll/serialization.rs b/datasketches/src/kll/serialization.rs index 3ba4a791..41e585c3 100644 --- a/datasketches/src/kll/serialization.rs +++ b/datasketches/src/kll/serialization.rs @@ -23,28 +23,28 @@ //! intentionally outside this module's scope. /// Serialization version for empty or full sketches (KllPreambleUtil.SERIAL_VERSION_EMPTY_FULL). -pub(super) const SERIAL_VERSION_1: u8 = 1; +pub const SERIAL_VERSION_1: u8 = 1; /// Serialization version for single-item sketches (KllPreambleUtil.SERIAL_VERSION_SINGLE). -pub(super) const SERIAL_VERSION_2: u8 = 2; +pub const SERIAL_VERSION_2: u8 = 2; /// Preamble ints for empty and single-item sketches (KllPreambleUtil.PREAMBLE_INTS_EMPTY_SINGLE). -pub(super) const PREAMBLE_INTS_SHORT: u8 = 2; +pub const PREAMBLE_INTS_SHORT: u8 = 2; /// Preamble ints for sketches with more than one item (KllPreambleUtil.PREAMBLE_INTS_FULL). -pub(super) const PREAMBLE_INTS_FULL: u8 = 5; +pub const PREAMBLE_INTS_FULL: u8 = 5; /// Flag indicating the sketch is empty (KllPreambleUtil.EMPTY_BIT_MASK). -pub(super) const FLAG_EMPTY: u8 = 1 << 0; +pub const FLAG_EMPTY: u8 = 1 << 0; /// Flag indicating level zero is sorted (KllPreambleUtil.LEVEL_ZERO_SORTED_BIT_MASK). -pub(super) const FLAG_LEVEL_ZERO_SORTED: u8 = 1 << 1; +pub const FLAG_LEVEL_ZERO_SORTED: u8 = 1 << 1; /// Flag indicating the sketch has a single item (KllPreambleUtil.SINGLE_ITEM_BIT_MASK). -pub(super) const FLAG_SINGLE_ITEM: u8 = 1 << 2; +pub const FLAG_SINGLE_ITEM: u8 = 1 << 2; /// Serialized size for an empty sketch in bytes (KllPreambleUtil.DATA_START_ADR_SINGLE_ITEM). -pub(super) const EMPTY_SIZE_BYTES: usize = 8; +pub const EMPTY_SIZE_BYTES: usize = 8; /// Data offset for single-item sketches (KllPreambleUtil.DATA_START_ADR_SINGLE_ITEM). -pub(super) const DATA_START_SINGLE_ITEM: usize = 8; +pub const DATA_START_SINGLE_ITEM: usize = 8; /// Data offset for sketches with more than one item (KllPreambleUtil.DATA_START_ADR). -pub(super) const DATA_START: usize = 20; +pub const DATA_START: usize = 20; /// Maximum level count supported by the KLL capacity calculation. -pub(super) const MAX_NUM_LEVELS: usize = 61; +pub const MAX_NUM_LEVELS: usize = 61; diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index 8248fe22..82399cdb 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -21,9 +21,8 @@ use super::DEFAULT_K; use super::DEFAULT_M; use super::MAX_K; use super::MIN_K; -use super::helper::compute_total_capacity; -use super::helper::level_capacity; -use super::helper::sum_the_sample_weights; +use super::capacity::level_capacity; +use super::capacity::total_capacity; use super::serialization::DATA_START; use super::serialization::DATA_START_SINGLE_ITEM; use super::serialization::EMPTY_SIZE_BYTES; @@ -35,133 +34,61 @@ use super::serialization::PREAMBLE_INTS_FULL; use super::serialization::PREAMBLE_INTS_SHORT; use super::serialization::SERIAL_VERSION_1; use super::serialization::SERIAL_VERSION_2; +use super::sorted_view::SortedView; use super::sorted_view::build_sorted_view; +use super::value::KllValue; use crate::codec::SketchBytes; use crate::codec::SketchSlice; use crate::codec::assert::ensure_serial_version_is; use crate::codec::assert::insufficient_data; use crate::codec::family::Family; +use crate::common::SearchCriteria; use crate::error::Error; -/// Trait implemented by item types supported by [`KllSketch`]. -/// -/// Implementations must provide a total ordering via `cmp`. -/// For floating-point types, ensure `cmp` handles NaN consistently and `is_nan` -/// returns true for values that should be ignored by updates. -pub trait KllItem: Clone { - /// Compare two items. - fn cmp(a: &Self, b: &Self) -> Ordering; - - /// Returns true if the item is NaN. - fn is_nan(_value: &Self) -> bool { - false - } -} - -/// Ordering policy used by a [`KllSketch`]. -/// -/// A sketch and every sketch merged into it must use equivalent ordering policies. -pub trait KllComparator: Clone { - /// Compare two items. - fn compare(&self, left: &T, right: &T) -> Ordering; -} - -/// Uses the natural ordering supplied by [`KllItem::cmp`]. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct NaturalOrder; - -impl KllComparator for NaturalOrder { - fn compare(&self, left: &T, right: &T) -> Ordering { - T::cmp(left, right) - } -} - -trait KllSerde: KllItem { - /// Minimum serialized size in bytes for one item. - const MIN_SERIALIZED_SIZE: usize; - - /// Serialized size in bytes. - fn serialized_size(value: &Self) -> usize; - - /// Serialize a single item into the buffer. - fn serialize(value: &Self, bytes: &mut SketchBytes); - - /// Deserialize a single item from the input. - fn deserialize(input: &mut SketchSlice<'_>) -> Result; -} - /// KLL sketch for estimating quantiles and ranks. /// /// See the [kll module level documentation](crate::kll) for more. #[derive(Debug, Clone, PartialEq)] -pub struct KllSketch { - comparator: C, +pub struct KllSketch { k: u16, m: u8, min_k: u16, n: u64, + num_retained: usize, + capacity: usize, is_level_zero_sorted: bool, levels: Vec>, min_item: Option, max_item: Option, } -impl Default for KllSketch { +impl Default for KllSketch { fn default() -> Self { - Self::make( - NaturalOrder, - DEFAULT_K, - DEFAULT_K, - 0, - vec![Vec::new()], - None, - None, - false, - ) + Self::make(DEFAULT_K, DEFAULT_K, 0, vec![Vec::new()], None, None, false) } } -impl KllSketch { +impl KllSketch { /// Creates a new sketch with the given value of k. /// /// # Errors /// - /// Returns an error if `k` is outside [`MIN_K`, `MAX_K`]. + /// Returns an error if `k` is outside `8..=65535`. /// /// # Examples /// /// ``` /// # use datasketches::kll::KllSketch; - /// let sketch = KllSketch::::new(200).unwrap(); + /// let sketch = KllSketch::::new(200).unwrap(); /// assert_eq!(sketch.k(), 200); /// ``` pub fn new(k: u16) -> Result { - Self::new_with_comparator(k, NaturalOrder) - } -} - -impl> KllSketch { - /// Creates a new sketch with the given value of k and ordering policy. - /// - /// # Errors - /// - /// Returns an error if `k` is outside [`MIN_K`, `MAX_K`]. - pub fn new_with_comparator(k: u16, comparator: C) -> Result { if !(MIN_K..=MAX_K).contains(&k) { return Err(Error::invalid_argument(format!( "k must be in [{MIN_K}, {MAX_K}], got {k}" ))); } - Ok(Self::make( - comparator, - k, - k, - 0, - vec![Vec::new()], - None, - None, - false, - )) + Ok(Self::make(k, k, 0, vec![Vec::new()], None, None, false)) } /// Returns parameter k used to configure this sketch. @@ -186,7 +113,7 @@ impl> KllSketch { /// Returns the number of retained items. pub fn num_retained(&self) -> usize { - self.levels.iter().map(|level| level.len()).sum() + self.num_retained } /// Returns true if the sketch is in estimation mode. @@ -206,11 +133,10 @@ impl> KllSketch { /// Updates the sketch with a new item. /// - /// NaN values are ignored for floating-point types. + /// # Panics + /// + /// Panics if the stream weight would exceed [`u64::MAX`]. pub fn update(&mut self, item: T) { - if T::is_nan(&item) { - return; - } self.update_min_max(&item); self.internal_update(item); } @@ -219,6 +145,8 @@ impl> KllSketch { pub fn reset(&mut self) { self.min_k = self.k; self.n = 0; + self.num_retained = 0; + self.capacity = total_capacity(self.k, self.m, 1) as usize; self.is_level_zero_sorted = false; self.levels.clear(); self.levels.push(Vec::new()); @@ -228,23 +156,31 @@ impl> KllSketch { /// Merges another sketch into this one. /// - /// # Panics + /// # Errors /// - /// Panics if the sketches have incompatible parameters. - pub fn merge(&mut self, other: &KllSketch) { + /// Returns an error if the combined stream weight exceeds [`u64::MAX`]. + pub fn merge(&mut self, other: &KllSketch) -> Result<(), Error> { if other.is_empty() { - return; + return Ok(()); } - assert_eq!( - self.m, other.m, - "incompatible m values: {} and {}", - self.m, other.m - ); + if self.m != other.m { + return Err(Error::invalid_argument(format!( + "cannot merge sketches with different m values: {} and {}", + self.m, other.m + ))); + } + let final_n = self.n.checked_add(other.n).ok_or_else(|| { + Error::invalid_argument(format!( + "combined stream weight exceeds {}: left {}, right {}", + u64::MAX, + self.n, + other.n + )) + })?; self.update_min_max_from_other(other); - let final_n = self.n + other.n; for item in &other.levels[0] { self.internal_update(item.clone()); } @@ -259,56 +195,107 @@ impl> KllSketch { } debug_assert_eq!(self.total_weight(), self.n, "total weight does not match n"); + Ok(()) } /// Returns the normalized rank of the given item. - pub fn rank(&self, item: &T, inclusive: bool) -> Option { + /// + /// # Errors + /// + /// Returns an error if the sketch is empty. + pub fn rank(&self, item: &T, criteria: SearchCriteria) -> Result { if self.is_empty() { - return None; + return Err(Error::invalid_argument("cannot query an empty sketch")); } - let view = build_sorted_view(&self.levels, self.comparator.clone()); - Some(view.rank(item, inclusive)) + let inclusive = criteria == SearchCriteria::Inclusive; + let mut weight = 0u64; + for (level, items) in self.levels.iter().enumerate() { + let count = items + .iter() + .filter(|retained| match (*retained).cmp(item) { + Ordering::Less => true, + Ordering::Equal => inclusive, + Ordering::Greater => false, + }) + .count() as u64; + weight += count << level; + } + Ok(weight as f64 / self.n as f64) } /// Returns the quantile for the given normalized rank. /// - /// # Panics + /// # Errors /// - /// Panics if rank is not in [0.0, 1.0]. - pub fn quantile(&self, rank: f64, inclusive: bool) -> Option { + /// Returns an error if the sketch is empty or `rank` is outside `[0.0, 1.0]`. + pub fn quantile(&self, rank: f64, criteria: SearchCriteria) -> Result { if self.is_empty() { - return None; + return Err(Error::invalid_argument("cannot query an empty sketch")); } - assert!((0.0..=1.0).contains(&rank), "rank must be in [0.0, 1.0]"); - let view = build_sorted_view(&self.levels, self.comparator.clone()); - Some(view.quantile(rank, inclusive)) + if !(0.0..=1.0).contains(&rank) { + return Err(Error::invalid_argument(format!( + "rank must be in [0.0, 1.0], got {rank}" + ))); + } + self.sorted_view().quantile(rank, criteria) + } + + /// Returns approximate quantiles for the given normalized ranks. + /// + /// The sorted view is built once for the whole batch. + /// + /// # Errors + /// + /// Returns an error if the sketch is empty or any rank is outside `[0.0, 1.0]`. + pub fn quantiles(&self, ranks: &[f64], criteria: SearchCriteria) -> Result, Error> { + self.sorted_view().quantiles(ranks, criteria) } /// Returns the approximate CDF for the given split points. - pub fn cdf(&self, split_points: &[T], inclusive: bool) -> Option> { + /// + /// # Errors + /// + /// Returns an error if the sketch is empty or the split points are not unique and strictly + /// increasing. + pub fn cdf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { if self.is_empty() { - return None; + return Err(Error::invalid_argument("cannot query an empty sketch")); } - let view = build_sorted_view(&self.levels, self.comparator.clone()); - Some(view.cdf(split_points, inclusive)) + self.sorted_view().cdf(split_points, criteria) } /// Returns the approximate PMF for the given split points. - pub fn pmf(&self, split_points: &[T], inclusive: bool) -> Option> { + /// + /// # Errors + /// + /// Returns an error if the sketch is empty or the split points are not unique and strictly + /// increasing. + pub fn pmf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { if self.is_empty() { - return None; + return Err(Error::invalid_argument("cannot query an empty sketch")); } - let view = build_sorted_view(&self.levels, self.comparator.clone()); - Some(view.pmf(split_points, inclusive)) + self.sorted_view().pmf(split_points, criteria) + } + + /// Returns an owned, sorted snapshot of the current sketch state. + /// + /// The view can be reused for repeated queries while this sketch continues to receive updates. + pub fn sorted_view(&self) -> SortedView { + build_sorted_view(&self.levels, self.is_level_zero_sorted) + } + + /// Returns the normalized single-sided rank error for the configured k. + pub fn normalized_rank_error(&self) -> f64 { + normalized_rank_error(self.min_k, false) } - /// Returns normalized rank error for the configured k. - pub fn normalized_rank_error(&self, pmf: bool) -> f64 { - normalized_rank_error(self.min_k, pmf) + /// Returns the normalized double-sided rank error for PMF queries for the configured k. + pub fn normalized_pmf_error(&self) -> f64 { + normalized_rank_error(self.min_k, true) } } -fn serialized_size>(sketch: &KllSketch) -> usize { +fn serialized_size(sketch: &KllSketch) -> usize { if sketch.is_empty() { return EMPTY_SIZE_BYTES; } @@ -332,7 +319,7 @@ fn serialized_size>(sketch: &KllSketch) - size } -fn serialize_with_serde>(sketch: &KllSketch) -> Vec { +fn serialize_with_serde(sketch: &KllSketch) -> Vec { let size = serialized_size(sketch); let mut bytes = SketchBytes::with_capacity(size); @@ -404,10 +391,7 @@ fn serialize_with_serde>(sketch: &KllSketch>( - bytes: &[u8], - comparator: C, -) -> Result, Error> { +fn deserialize_with_serde(bytes: &[u8]) -> Result, Error> { let mut cursor = SketchSlice::new(bytes); let preamble_ints = cursor @@ -457,15 +441,19 @@ fn deserialize_with_serde>( ensure_serial_version_is(expected_version, serial_version)?; if !(MIN_K..=MAX_K).contains(&k) { - return Err(Error::deserial(format!("k out of range: {k}"))); + return Err(Error::deserial(format!( + "k must be in [{MIN_K}, {MAX_K}], got {k}" + ))); } if is_empty { - if !cursor.remaining().is_empty() { - return Err(Error::deserial("unexpected trailing data")); + let trailing_bytes = cursor.remaining().len(); + if trailing_bytes != 0 { + return Err(Error::deserial(format!( + "expected end of KLL image, found {trailing_bytes} trailing bytes" + ))); } return Ok(KllSketch::make( - comparator, k, k, 0, @@ -486,17 +474,14 @@ fn deserialize_with_serde>( (n, min_k, num_levels as usize) }; - if num_levels == 0 { - return Err(Error::deserial("num_levels must be > 0")); - } - if num_levels > MAX_NUM_LEVELS { + if !(1..=MAX_NUM_LEVELS).contains(&num_levels) { return Err(Error::deserial(format!( - "num_levels must be at most {MAX_NUM_LEVELS}, got {num_levels}" + "num_levels must be in [1, {MAX_NUM_LEVELS}], got {num_levels}" ))); } if !is_single_item && n < 2 { return Err(Error::deserial(format!( - "full sketch must have n >= 2, got {n}" + "full sketch n must be at least 2, got {n}" ))); } if min_k < MIN_K || min_k > k { @@ -505,7 +490,7 @@ fn deserialize_with_serde>( ))); } - let capacity = compute_total_capacity(k, m, num_levels); + let capacity = total_capacity(k, m, num_levels); let mut level_offsets = Vec::with_capacity(num_levels + 1); if !is_single_item { for _ in 0..num_levels { @@ -517,47 +502,67 @@ fn deserialize_with_serde>( } level_offsets.push(capacity); - if level_offsets.is_empty() { - return Err(Error::deserial("levels array is empty")); - } if level_offsets[0] > capacity { - return Err(Error::deserial("levels[0] exceeds capacity")); + return Err(Error::deserial(format!( + "first level offset must not exceed capacity {capacity}, got {}", + level_offsets[0] + ))); } - for window in level_offsets.windows(2) { + for (index, window) in level_offsets.windows(2).enumerate() { if window[1] < window[0] { - return Err(Error::deserial("levels array must be non-decreasing")); + return Err(Error::deserial(format!( + "level offsets must be nondecreasing: offset[{index}] is {}, offset[{}] is {}", + window[0], + index + 1, + window[1] + ))); } } - let last = *level_offsets.last().unwrap(); - if last != capacity { - return Err(Error::deserial("levels last offset must equal capacity")); - } let min_item = if is_single_item { None } else { - Some(T::deserialize(&mut cursor)?) + Some( + T::deserialize(&mut cursor) + .map_err(|error| error.with_context("KLL item", "minimum"))?, + ) }; let max_item = if is_single_item { None } else { - Some(T::deserialize(&mut cursor)?) + Some( + T::deserialize(&mut cursor) + .map_err(|error| error.with_context("KLL item", "maximum"))?, + ) }; let num_retained = (level_offsets[num_levels] - level_offsets[0]) as usize; let min_item_bytes = num_retained .checked_mul(T::MIN_SERIALIZED_SIZE) - .ok_or_else(|| Error::deserial("retained item size overflow"))?; - if cursor.remaining().len() < min_item_bytes { - return Err(Error::insufficient_data("items")); + .ok_or_else(|| { + Error::deserial(format!( + "minimum serialized size overflows usize: {num_retained} retained items, {} bytes per item", + T::MIN_SERIALIZED_SIZE + )) + })?; + let available_item_bytes = cursor.remaining().len(); + if available_item_bytes < min_item_bytes { + return Err(Error::insufficient_data_of( + "KLL item payload", + format_args!("expected {min_item_bytes} bytes, got {available_item_bytes}"), + )); } let mut levels = Vec::with_capacity(num_levels); for level in 0..num_levels { let size = (level_offsets[level + 1] - level_offsets[level]) as usize; let mut items = Vec::with_capacity(size); - for _ in 0..size { - items.push(T::deserialize(&mut cursor)?); + for index in 0..size { + items.push(T::deserialize(&mut cursor).map_err(|error| { + error + .with_context("KLL level", level) + .with_context("item index", index) + })?); } levels.push(items); } @@ -566,7 +571,6 @@ fn deserialize_with_serde>( } let mut sketch = KllSketch::make( - comparator, k, min_k, n, @@ -584,92 +588,35 @@ fn deserialize_with_serde>( } sketch.validate_deserialized_state()?; - if !cursor.remaining().is_empty() { - return Err(Error::deserial("unexpected trailing data")); + let trailing_bytes = cursor.remaining().len(); + if trailing_bytes != 0 { + return Err(Error::deserial(format!( + "expected end of KLL image, found {trailing_bytes} trailing bytes" + ))); } Ok(sketch) } -impl> KllSketch { +impl KllSketch { /// Serializes the sketch to bytes. pub fn serialize(&self) -> Vec { serialize_with_serde(self) } - /// Deserializes a sketch using the supplied ordering policy. - pub fn deserialize_with_comparator(bytes: &[u8], comparator: C) -> Result { - deserialize_with_serde(bytes, comparator) - } -} - -impl KllSketch { - /// Deserializes a sketch from bytes. - pub fn deserialize(bytes: &[u8]) -> Result { - deserialize_with_serde(bytes, NaturalOrder) - } -} - -impl> KllSketch { - /// Serializes the sketch to bytes. - pub fn serialize(&self) -> Vec { - serialize_with_serde(self) - } - - /// Deserializes a sketch using the supplied ordering policy. - pub fn deserialize_with_comparator(bytes: &[u8], comparator: C) -> Result { - deserialize_with_serde(bytes, comparator) - } -} - -impl KllSketch { - /// Deserializes a sketch from bytes. - pub fn deserialize(bytes: &[u8]) -> Result { - deserialize_with_serde(bytes, NaturalOrder) - } -} - -impl> KllSketch { - /// Serializes the sketch to bytes. - pub fn serialize(&self) -> Vec { - serialize_with_serde(self) - } - - /// Deserializes a sketch using the supplied ordering policy. - pub fn deserialize_with_comparator(bytes: &[u8], comparator: C) -> Result { - deserialize_with_serde(bytes, comparator) - } -} - -impl KllSketch { - /// Deserializes a sketch from bytes. - pub fn deserialize(bytes: &[u8]) -> Result { - deserialize_with_serde(bytes, NaturalOrder) - } -} - -impl> KllSketch { - /// Serializes the sketch to bytes. - pub fn serialize(&self) -> Vec { - serialize_with_serde(self) - } - - /// Deserializes a sketch using the supplied ordering policy. - pub fn deserialize_with_comparator(bytes: &[u8], comparator: C) -> Result { - deserialize_with_serde(bytes, comparator) - } -} - -impl KllSketch { /// Deserializes a sketch from bytes. + /// + /// # Errors + /// + /// Returns `InvalidData` if the image is truncated, malformed, or contains values that are not + /// totally ordered. pub fn deserialize(bytes: &[u8]) -> Result { - deserialize_with_serde(bytes, NaturalOrder) + deserialize_with_serde(bytes) } } -impl> KllSketch { +impl KllSketch { fn make( - comparator: C, k: u16, min_k: u16, n: u64, @@ -678,12 +625,15 @@ impl> KllSketch { max_item: Option, is_level_zero_sorted: bool, ) -> Self { + let num_retained = levels.iter().map(Vec::len).sum(); + let capacity = total_capacity(k, DEFAULT_M, levels.len()) as usize; Self { - comparator, k, m: DEFAULT_M, min_k, n, + num_retained, + capacity, is_level_zero_sorted, levels, min_item, @@ -691,14 +641,13 @@ impl> KllSketch { } } - fn capacity(&self) -> usize { - compute_total_capacity(self.k, self.m, self.levels.len()) as usize - } - fn level_offsets(&self) -> Vec { - let capacity = self.capacity() as u32; + let capacity = self.capacity as u32; let retained = self.num_retained() as u32; - assert!(capacity >= retained, "capacity must be >= retained"); + assert!( + capacity >= retained, + "KLL retained item count must not exceed capacity: retained {retained}, capacity {capacity}" + ); let mut offsets = Vec::with_capacity(self.levels.len() + 1); let mut offset = capacity - retained; @@ -717,11 +666,11 @@ impl> KllSketch { self.max_item = Some(item.clone()); } Some(min) => { - if self.comparator.compare(item, min) == Ordering::Less { + if item.cmp(min) == Ordering::Less { self.min_item = Some(item.clone()); } if let Some(max) = &self.max_item { - if self.comparator.compare(max, item) == Ordering::Less { + if max.cmp(item) == Ordering::Less { self.max_item = Some(item.clone()); } } @@ -729,7 +678,7 @@ impl> KllSketch { } } - fn update_min_max_from_other(&mut self, other: &KllSketch) { + fn update_min_max_from_other(&mut self, other: &KllSketch) { match (&self.min_item, &self.max_item) { (None, None) => { self.min_item = other.min_item.clone(); @@ -737,12 +686,12 @@ impl> KllSketch { } (Some(min), Some(max)) => { if let Some(other_min) = &other.min_item { - if self.comparator.compare(other_min, min) == Ordering::Less { + if other_min.cmp(min) == Ordering::Less { self.min_item = Some(other_min.clone()); } } if let Some(other_max) = &other.max_item { - if self.comparator.compare(max, other_max) == Ordering::Less { + if max.cmp(other_max) == Ordering::Less { self.max_item = Some(other_max.clone()); } } @@ -755,10 +704,17 @@ impl> KllSketch { } fn internal_update(&mut self, item: T) { - if self.num_retained() >= self.capacity() { + if self.num_retained >= self.capacity { self.compress_while_updating(); } - self.n += 1; + self.n = self.n.checked_add(1).unwrap_or_else(|| { + panic!( + "cannot update KLL sketch: stream weight is {}, maximum is {}", + self.n, + u64::MAX + ) + }); + self.num_retained += 1; self.is_level_zero_sorted = false; self.levels[0].push(item); } @@ -769,29 +725,20 @@ impl> KllSketch { self.levels.push(Vec::new()); } - let mut current = std::mem::take(&mut self.levels[level]); + let current = std::mem::take(&mut self.levels[level]); let mut above = std::mem::take(&mut self.levels[level + 1]); - - let odd = current.len() % 2 == 1; - let mut leftover = None; - if odd { - leftover = Some(take_leftover( - &mut current, - level, - self.is_level_zero_sorted, - )); - } - - if level == 0 && !self.is_level_zero_sorted { - current.sort_by(|left, right| self.comparator.compare(left, right)); - } - let use_up = above.is_empty(); - let promoted = downsample(current, rand::random::(), use_up); + let (leftover, promoted) = compact_level( + current, + level, + self.is_level_zero_sorted, + rand::random::(), + use_up, + ); if above.is_empty() { above = promoted; } else { - above = merge_sorted_vec(promoted, above, &self.comparator); + above = merge_sorted_vec(promoted, above); } self.levels[level + 1] = above; @@ -800,6 +747,7 @@ impl> KllSketch { new_level.push(item); } self.levels[level] = new_level; + self.refresh_capacity_state(); } fn find_level_to_compact(&self) -> usize { @@ -811,10 +759,13 @@ impl> KllSketch { return level; } } - panic!("no level to compact"); + panic!( + "KLL sketch has {}/{} retained items but no level reached its compaction capacity (k {}, m {}, levels {num_levels})", + self.num_retained, self.capacity, self.k, self.m + ); } - fn merge_higher_levels(&mut self, other: &KllSketch) { + fn merge_higher_levels(&mut self, other: &KllSketch) { let provisional_levels = self.levels.len().max(other.levels.len()); let mut self_levels = std::mem::take(&mut self.levels); let mut work_levels = vec![Vec::new(); provisional_levels]; @@ -833,22 +784,25 @@ impl> KllSketch { } else if right.is_empty() { left } else { - merge_sorted_vec(left, right, &self.comparator) + merge_sorted_vec(left, right) }; } - self.levels = general_compress( - work_levels, - self.k, - self.m, - self.is_level_zero_sorted, - &self.comparator, - ); + self.levels = general_compress(work_levels, self.k, self.m, self.is_level_zero_sorted); + self.refresh_capacity_state(); + } + + fn refresh_capacity_state(&mut self) { + self.num_retained = self.levels.iter().map(Vec::len).sum(); + self.capacity = total_capacity(self.k, self.m, self.levels.len()) as usize; } fn total_weight(&self) -> u64 { - let sizes: Vec = self.levels.iter().map(|level| level.len()).collect(); - sum_the_sample_weights(&sizes) + self.levels + .iter() + .enumerate() + .map(|(level, items)| (items.len() as u64) << level) + .sum() } fn validate_deserialized_state(&self) -> Result<(), Error> { @@ -861,10 +815,7 @@ impl> KllSketch { .as_ref() .ok_or_else(|| Error::deserial("non-empty sketch must have a maximum item"))?; - if T::is_nan(min_item) || T::is_nan(max_item) { - return Err(Error::deserial("minimum and maximum items must not be NaN")); - } - if self.comparator.compare(min_item, max_item) == Ordering::Greater { + if min_item.cmp(max_item) == Ordering::Greater { return Err(Error::deserial( "minimum item must not be greater than maximum item", )); @@ -875,39 +826,48 @@ impl> KllSketch { for (level_index, level) in self.levels.iter().enumerate() { let level_total = level_weight .checked_mul(level.len() as u64) - .ok_or_else(|| Error::deserial("sample weight overflow"))?; + .ok_or_else(|| { + Error::deserial(format!( + "sample weight overflows u64 at level {level_index}: weight {level_weight}, retained items {}", + level.len() + )) + })?; total_weight = total_weight .checked_add(level_total) - .ok_or_else(|| Error::deserial("total sample weight overflow"))?; + .ok_or_else(|| { + Error::deserial(format!( + "total sample weight overflows u64 at level {level_index}: accumulated {total_weight}, level contribution {level_total}" + )) + })?; let must_be_sorted = level_index > 0 || self.is_level_zero_sorted; - if must_be_sorted - && level - .windows(2) - .any(|pair| self.comparator.compare(&pair[0], &pair[1]) == Ordering::Greater) - { - return Err(Error::deserial(format!( - "level {level_index} must be sorted" - ))); + if must_be_sorted { + for (item_index, pair) in level.windows(2).enumerate() { + if pair[0].cmp(&pair[1]) == Ordering::Greater { + return Err(Error::deserial(format!( + "level {level_index} must be sorted: item at index {item_index} is greater than item at index {}", + item_index + 1 + ))); + } + } } - for item in level { - if T::is_nan(item) { - return Err(Error::deserial("retained items must not be NaN")); - } - if self.comparator.compare(item, min_item) == Ordering::Less - || self.comparator.compare(item, max_item) == Ordering::Greater - { - return Err(Error::deserial( - "retained items must be within the minimum and maximum", - )); + for (item_index, item) in level.iter().enumerate() { + if item.cmp(min_item) == Ordering::Less || item.cmp(max_item) == Ordering::Greater { + return Err(Error::deserial(format!( + "retained item at level {level_index}, index {item_index} is outside the serialized minimum and maximum" + ))); } } if level_index + 1 < self.levels.len() { level_weight = level_weight .checked_mul(2) - .ok_or_else(|| Error::deserial("level weight overflow"))?; + .ok_or_else(|| { + Error::deserial(format!( + "level weight overflows u64 after level {level_index}: current weight {level_weight}" + )) + })?; } } @@ -931,9 +891,40 @@ fn normalized_rank_error(k: u16, pmf: bool) -> f64 { } } -fn downsample(items: Vec, offset: bool, use_up: bool) -> Vec { +fn compact_level( + mut items: Vec, + level: usize, + is_level_zero_sorted: bool, + offset: bool, + use_up: bool, +) -> (Option, Vec) { + let odd = items.len() % 2 == 1; + let level_zero_needs_sorting = level == 0 && !is_level_zero_sorted; + let leftover = if odd && level_zero_needs_sorting { + items.pop() + } else { + None + }; + if level_zero_needs_sorting { + items.sort_unstable(); + } + + let mut items = items.into_iter(); + let leftover = if odd && !level_zero_needs_sorting { + items.next() + } else { + leftover + }; + let promoted = downsample(items, offset, use_up); + (leftover, promoted) +} + +fn downsample>(items: I, offset: bool, use_up: bool) -> Vec { let len = items.len(); - debug_assert!(len % 2 == 0, "length must be even"); + debug_assert!( + len % 2 == 0, + "KLL compaction requires an even item count, got {len}" + ); let offset = usize::from(offset); let parity = if use_up { (len - 1 - offset) % 2 @@ -942,31 +933,18 @@ fn downsample(items: Vec, offset: bool, use_up: bool) -> Vec { }; items - .into_iter() .enumerate() .filter_map(|(idx, item)| if idx % 2 == parity { Some(item) } else { None }) .collect() } -fn take_leftover(items: &mut Vec, level: usize, is_level_zero_sorted: bool) -> T { - if level == 0 && !is_level_zero_sorted { - items.pop().expect("odd level must not be empty") - } else { - items.remove(0) - } -} - -fn merge_sorted_vec>( - left: Vec, - right: Vec, - comparator: &C, -) -> Vec { +fn merge_sorted_vec(left: Vec, right: Vec) -> Vec { let mut merged = Vec::with_capacity(left.len() + right.len()); let mut left_iter = left.into_iter().peekable(); let mut right_iter = right.into_iter().peekable(); while let (Some(l), Some(r)) = (left_iter.peek(), right_iter.peek()) { - if comparator.compare(l, r) == Ordering::Less { + if l.cmp(r) == Ordering::Less { merged.push(left_iter.next().unwrap()); } else { merged.push(right_iter.next().unwrap()); @@ -977,16 +955,15 @@ fn merge_sorted_vec>( merged } -fn general_compress>( +fn general_compress( mut levels_in: Vec>, k: u16, m: u8, is_level_zero_sorted: bool, - comparator: &C, ) -> Vec> { let mut current_num_levels = levels_in.len(); let mut current_item_count: usize = levels_in.iter().map(|level| level.len()).sum(); - let mut target_item_count = compute_total_capacity(k, m, current_num_levels) as usize; + let mut target_item_count = total_capacity(k, m, current_num_levels) as usize; let mut levels_out = Vec::with_capacity(current_num_levels + 1); let mut current_level = 0usize; @@ -1001,30 +978,21 @@ fn general_compress>( if current_item_count < target_item_count || raw_pop < cap { levels_out.push(std::mem::take(&mut levels_in[current_level])); } else { - let mut current = std::mem::take(&mut levels_in[current_level]); + let current = std::mem::take(&mut levels_in[current_level]); let mut above = std::mem::take(&mut levels_in[current_level + 1]); - - let odd = current.len() % 2 == 1; - let mut leftover = None; - if odd { - leftover = Some(take_leftover( - &mut current, - current_level, - is_level_zero_sorted, - )); - } - - if current_level == 0 && !is_level_zero_sorted { - current.sort_by(|left, right| comparator.compare(left, right)); - } - let use_up = above.is_empty(); - let promoted = downsample(current, rand::random::(), use_up); + let (leftover, promoted) = compact_level( + current, + current_level, + is_level_zero_sorted, + rand::random::(), + use_up, + ); let promoted_len = promoted.len(); if above.is_empty() { above = promoted; } else { - above = merge_sorted_vec(promoted, above, comparator); + above = merge_sorted_vec(promoted, above); } levels_in[current_level + 1] = above; @@ -1050,117 +1018,3 @@ fn general_compress>( levels_out.truncate(current_num_levels); levels_out } - -impl KllItem for f32 { - fn cmp(a: &Self, b: &Self) -> Ordering { - a.partial_cmp(b).unwrap_or(Ordering::Greater) - } - - fn is_nan(value: &Self) -> bool { - value.is_nan() - } -} - -impl KllSerde for f32 { - const MIN_SERIALIZED_SIZE: usize = 4; - - fn serialized_size(_value: &Self) -> usize { - 4 - } - - fn serialize(value: &Self, bytes: &mut SketchBytes) { - bytes.write_f32_le(*value); - } - - fn deserialize(input: &mut SketchSlice<'_>) -> Result { - input - .read_f32_le() - .map_err(|_| Error::insufficient_data("f32")) - } -} - -impl KllItem for f64 { - fn cmp(a: &Self, b: &Self) -> Ordering { - a.partial_cmp(b).unwrap_or(Ordering::Greater) - } - - fn is_nan(value: &Self) -> bool { - value.is_nan() - } -} - -impl KllSerde for f64 { - const MIN_SERIALIZED_SIZE: usize = 8; - - fn serialized_size(_value: &Self) -> usize { - 8 - } - - fn serialize(value: &Self, bytes: &mut SketchBytes) { - bytes.write_f64_le(*value); - } - - fn deserialize(input: &mut SketchSlice<'_>) -> Result { - input - .read_f64_le() - .map_err(|_| Error::insufficient_data("f64")) - } -} - -impl KllItem for i64 { - fn cmp(a: &Self, b: &Self) -> Ordering { - a.cmp(b) - } -} - -impl KllSerde for i64 { - const MIN_SERIALIZED_SIZE: usize = 8; - - fn serialized_size(_value: &Self) -> usize { - 8 - } - - fn serialize(value: &Self, bytes: &mut SketchBytes) { - bytes.write_i64_le(*value); - } - - fn deserialize(input: &mut SketchSlice<'_>) -> Result { - input - .read_i64_le() - .map_err(|_| Error::insufficient_data("i64")) - } -} - -impl KllItem for String { - fn cmp(a: &Self, b: &Self) -> Ordering { - a.cmp(b) - } -} - -impl KllSerde for String { - const MIN_SERIALIZED_SIZE: usize = 4; - - fn serialized_size(value: &Self) -> usize { - 4 + value.len() - } - - fn serialize(value: &Self, bytes: &mut SketchBytes) { - bytes.write_u32_le(value.len() as u32); - bytes.write(value.as_bytes()); - } - - fn deserialize(input: &mut SketchSlice<'_>) -> Result { - let len = input - .read_u32_le() - .map_err(|_| Error::insufficient_data("string_len"))? as usize; - let bytes = input - .remaining() - .get(..len) - .ok_or_else(|| Error::insufficient_data("string_bytes"))?; - let value = std::str::from_utf8(bytes) - .map_err(|_| Error::deserial("invalid utf-8 string"))? - .to_owned(); - input.advance(len as u64); - Ok(value) - } -} diff --git a/datasketches/src/kll/sorted_view.rs b/datasketches/src/kll/sorted_view.rs index bcaad640..b58a69a9 100644 --- a/datasketches/src/kll/sorted_view.rs +++ b/datasketches/src/kll/sorted_view.rs @@ -17,12 +17,15 @@ use std::cmp::Ordering; -use super::sketch::KllComparator; -use super::sketch::KllItem; +use crate::common::SearchCriteria; +use crate::error::Error; +/// An owned, sorted snapshot of a KLL sketch. +/// +/// Build one with [`KllSketch::sorted_view`](super::KllSketch::sorted_view) when running repeated +/// queries against the same sketch state. #[derive(Debug, Clone)] -pub(super) struct SortedView> { - comparator: C, +pub struct SortedView { entries: Vec>, total_weight: u64, } @@ -30,174 +33,214 @@ pub(super) struct SortedView> { #[derive(Debug, Clone)] struct Entry { item: T, - weight: u64, + cumulative_weight: u64, } -impl> SortedView { - fn new(mut entries: Vec>, comparator: C) -> Self { - entries.sort_by(|a, b| comparator.compare(&a.item, &b.item)); +impl SortedView { + fn from_sorted(mut entries: Vec>) -> Self { let mut total_weight = 0u64; for entry in &mut entries { - total_weight += entry.weight; - entry.weight = total_weight; + total_weight += entry.cumulative_weight; + entry.cumulative_weight = total_weight; } Self { - comparator, entries, total_weight, } } - pub(super) fn rank(&self, item: &T, inclusive: bool) -> f64 { - if self.entries.is_empty() { - return 0.0; - } + /// Returns whether the view contains no retained items. + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } - let idx = if inclusive { - upper_bound(&self.entries, item, &self.comparator) + /// Returns the number of retained items in the view. + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Returns the total stream weight represented by the view. + pub fn total_weight(&self) -> u64 { + self.total_weight + } + + /// Returns the approximate normalized rank of `item`. + /// + /// # Errors + /// + /// Returns an error if the view is empty. + pub fn rank(&self, item: &T, criteria: SearchCriteria) -> Result { + if self.is_empty() { + return Err(Error::invalid_argument("cannot query an empty view")); + } + let index = if criteria == SearchCriteria::Inclusive { + upper_bound(&self.entries, item) } else { - lower_bound(&self.entries, item, &self.comparator) + lower_bound(&self.entries, item) }; - if idx == 0 { - return 0.0; + if index == 0 { + return Ok(0.0); } - let weight = self.entries[idx - 1].weight; - weight as f64 / self.total_weight as f64 + Ok(self.entries[index - 1].cumulative_weight as f64 / self.total_weight as f64) } - pub(super) fn quantile(&self, rank: f64, inclusive: bool) -> T { - let weight = if inclusive { + /// Returns the approximate quantile for `rank`. + /// + /// # Errors + /// + /// Returns an error if the view is empty or `rank` is outside `[0.0, 1.0]`. + pub fn quantile(&self, rank: f64, criteria: SearchCriteria) -> Result { + if self.is_empty() { + return Err(Error::invalid_argument("cannot query an empty view")); + } + if !(0.0..=1.0).contains(&rank) { + return Err(Error::invalid_argument(format!( + "rank must be in [0.0, 1.0], got {rank}" + ))); + } + + let weight = if criteria == SearchCriteria::Inclusive { (rank * self.total_weight as f64).ceil() as u64 } else { (rank * self.total_weight as f64) as u64 }; - - let idx = if inclusive { + let index = if criteria == SearchCriteria::Inclusive { lower_bound_by_weight(&self.entries, weight) } else { upper_bound_by_weight(&self.entries, weight) }; - if idx >= self.entries.len() { - return self.entries[self.entries.len() - 1].item.clone(); + Ok(self.entries[index.min(self.entries.len() - 1)].item.clone()) + } + + /// Returns approximate quantiles for all `ranks`. + /// + /// # Errors + /// + /// Returns an error if the view is empty or any rank is outside `[0.0, 1.0]`. + pub fn quantiles(&self, ranks: &[f64], criteria: SearchCriteria) -> Result, Error> { + if self.is_empty() { + return Err(Error::invalid_argument("cannot query an empty view")); } - self.entries[idx].item.clone() + ranks + .iter() + .map(|&rank| self.quantile(rank, criteria)) + .collect() } - pub(super) fn cdf(&self, split_points: &[T], inclusive: bool) -> Vec { - check_split_points(split_points, &self.comparator); + /// Returns the approximate cumulative distribution over `split_points`. + /// + /// # Errors + /// + /// Returns an error if the view is empty or the split points are invalid. + pub fn cdf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { + if self.is_empty() { + return Err(Error::invalid_argument("cannot query an empty view")); + } + check_split_points(split_points)?; let mut ranks = Vec::with_capacity(split_points.len() + 1); for item in split_points { - ranks.push(self.rank(item, inclusive)); + ranks.push(self.rank(item, criteria)?); } ranks.push(1.0); - ranks + Ok(ranks) } - pub(super) fn pmf(&self, split_points: &[T], inclusive: bool) -> Vec { - let mut buckets = self.cdf(split_points, inclusive); - for i in (1..buckets.len()).rev() { - buckets[i] -= buckets[i - 1]; + /// Returns the approximate probability mass over `split_points`. + /// + /// # Errors + /// + /// Returns an error if the view is empty or the split points are invalid. + pub fn pmf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { + let mut buckets = self.cdf(split_points, criteria)?; + for index in (1..buckets.len()).rev() { + buckets[index] -= buckets[index - 1]; } - buckets + Ok(buckets) } } -pub(super) fn build_sorted_view>( +pub fn build_sorted_view( levels: &[Vec], - comparator: C, -) -> SortedView { - let num_retained: usize = levels.iter().map(|level| level.len()).sum(); - let mut entries = Vec::with_capacity(num_retained); - - for (level_idx, level) in levels.iter().enumerate() { - let weight = 1u64 << level_idx; - for item in level { - entries.push(Entry { - item: item.clone(), - weight, - }); + is_level_zero_sorted: bool, +) -> SortedView { + let mut runs = Vec::with_capacity(levels.len()); + for (level_index, level) in levels.iter().enumerate() { + let weight = 1u64 << level_index; + let mut run: Vec<_> = level + .iter() + .cloned() + .map(|item| Entry { + item, + cumulative_weight: weight, + }) + .collect(); + if level_index == 0 && !is_level_zero_sorted { + run.sort_unstable_by(|left, right| left.item.cmp(&right.item)); + } + if !run.is_empty() { + runs.push(run); } } - SortedView::new(entries, comparator) -} - -#[track_caller] -fn check_split_points>(split_points: &[T], comparator: &C) { - assert!( - split_points.iter().all(|point| !T::is_nan(point)), - "split_points must not contain NaN values" - ); - for pair in split_points.windows(2) { - assert!( - comparator.compare(&pair[0], &pair[1]) == Ordering::Less, - "split_points must be unique and monotonically increasing" - ); + while runs.len() > 1 { + let mut merged_runs = Vec::with_capacity(runs.len().div_ceil(2)); + let mut iter = runs.into_iter(); + while let Some(left) = iter.next() { + if let Some(right) = iter.next() { + merged_runs.push(merge_sorted_entries(left, right)); + } else { + merged_runs.push(left); + } + } + runs = merged_runs; } + + SortedView::from_sorted(runs.pop().unwrap_or_default()) } -fn lower_bound>( - entries: &[Entry], - item: &T, - comparator: &C, -) -> usize { - let mut left = 0usize; - let mut right = entries.len(); - while left < right { - let mid = left + (right - left) / 2; - if comparator.compare(&entries[mid].item, item) == Ordering::Less { - left = mid + 1; +fn merge_sorted_entries(left: Vec>, right: Vec>) -> Vec> { + let mut merged = Vec::with_capacity(left.len() + right.len()); + let mut left = left.into_iter().peekable(); + let mut right = right.into_iter().peekable(); + + while let (Some(left_entry), Some(right_entry)) = (left.peek(), right.peek()) { + if left_entry.item.cmp(&right_entry.item) == Ordering::Greater { + merged.push(right.next().unwrap()); } else { - right = mid; + merged.push(left.next().unwrap()); } } - left + merged.extend(left); + merged.extend(right); + merged } -fn upper_bound>( - entries: &[Entry], - item: &T, - comparator: &C, -) -> usize { - let mut left = 0usize; - let mut right = entries.len(); - while left < right { - let mid = left + (right - left) / 2; - if comparator.compare(&entries[mid].item, item) == Ordering::Greater { - right = mid; - } else { - left = mid + 1; +fn check_split_points(split_points: &[T]) -> Result<(), Error> { + for (index, pair) in split_points.windows(2).enumerate() { + if pair[0].cmp(&pair[1]) != Ordering::Less { + return Err(Error::invalid_argument(format!( + "split points at indices {index} and {} must be strictly increasing", + index + 1 + ))); } } - left + Ok(()) } -fn lower_bound_by_weight(entries: &[Entry], weight: u64) -> usize { - let mut left = 0usize; - let mut right = entries.len(); - while left < right { - let mid = left + (right - left) / 2; - if entries[mid].weight < weight { - left = mid + 1; - } else { - right = mid; - } - } - left +fn lower_bound(entries: &[Entry], item: &T) -> usize { + entries.partition_point(|entry| entry.item.cmp(item) == Ordering::Less) } -fn upper_bound_by_weight(entries: &[Entry], weight: u64) -> usize { - let mut left = 0usize; - let mut right = entries.len(); - while left < right { - let mid = left + (right - left) / 2; - if entries[mid].weight > weight { - right = mid; - } else { - left = mid + 1; - } - } - left +fn upper_bound(entries: &[Entry], item: &T) -> usize { + entries.partition_point(|entry| entry.item.cmp(item) != Ordering::Greater) +} + +fn lower_bound_by_weight(entries: &[Entry], weight: u64) -> usize { + entries.partition_point(|entry| entry.cumulative_weight < weight) +} + +fn upper_bound_by_weight(entries: &[Entry], weight: u64) -> usize { + entries.partition_point(|entry| entry.cumulative_weight <= weight) } diff --git a/datasketches/src/kll/value.rs b/datasketches/src/kll/value.rs new file mode 100644 index 00000000..b69f3695 --- /dev/null +++ b/datasketches/src/kll/value.rs @@ -0,0 +1,222 @@ +// 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::cmp::Ordering; +use std::fmt; +use std::mem::size_of; +use std::ops::Deref; + +use crate::codec::SketchBytes; +use crate::codec::SketchSlice; +use crate::codec::assert::insufficient_data; +use crate::error::Error; + +/// A non-NaN floating-point adapter for [`KllSketch`](crate::kll::KllSketch). +/// +/// KLL requires a totally ordered item domain, while primitive floats are unordered in the +/// presence of NaN. Construction therefore rejects NaN. Other values retain their numerical +/// order: signed zeros compare equal and infinities are allowed. +/// +/// The inner float is available through [`into_inner`](Self::into_inner) or immutable +/// dereferencing. +#[repr(transparent)] +#[derive(Clone, Copy, PartialEq, PartialOrd)] +pub struct KllFloat(T); + +impl KllFloat { + /// Returns the wrapped floating-point value. + #[inline(always)] + pub fn into_inner(self) -> T { + self.0 + } +} + +impl Deref for KllFloat { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl fmt::Debug for KllFloat { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl fmt::Display for KllFloat { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl KllFloat { + /// Creates a non-NaN KLL value. + /// + /// # Errors + /// + /// Returns an error if `value` is NaN. + #[inline(always)] + pub fn new(value: f32) -> Result { + if value.is_nan() { + Err(Error::invalid_argument("KLL float must not be NaN")) + } else { + Ok(Self(value)) + } + } +} + +impl Eq for KllFloat {} + +impl Ord for KllFloat { + #[inline(always)] + fn cmp(&self, other: &Self) -> Ordering { + self.0.partial_cmp(&other.0).unwrap() + } +} + +impl KllFloat { + /// Creates a non-NaN KLL value. + /// + /// # Errors + /// + /// Returns an error if `value` is NaN. + #[inline(always)] + pub fn new(value: f64) -> Result { + if value.is_nan() { + Err(Error::invalid_argument("KLL float must not be NaN")) + } else { + Ok(Self(value)) + } + } +} + +impl Eq for KllFloat {} + +impl Ord for KllFloat { + #[inline(always)] + fn cmp(&self, other: &Self) -> Ordering { + self.0.partial_cmp(&other.0).unwrap() + } +} + +/// Defines the compact binary representation of a KLL item. +/// +/// This trait is required only for serialization. In-memory KLL operations support any cloneable, +/// totally ordered item type. The encoded representation must preserve that ordering across a +/// round trip. +pub trait KllValue: Clone { + /// Minimum number of bytes required to encode one value. + const MIN_SERIALIZED_SIZE: usize; + + /// Returns the number of bytes required to encode `value`. + fn serialized_size(value: &Self) -> usize; + + /// Serializes `value` into `bytes`. + fn serialize(value: &Self, bytes: &mut SketchBytes); + + /// Deserializes one value from `input`. + fn deserialize(input: &mut SketchSlice<'_>) -> Result; +} + +impl KllValue for KllFloat { + const MIN_SERIALIZED_SIZE: usize = size_of::(); + + fn serialized_size(_value: &Self) -> usize { + size_of::() + } + + fn serialize(value: &Self, bytes: &mut SketchBytes) { + bytes.write_f32_le(value.0); + } + + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + let value = input + .read_f32_le() + .map_err(insufficient_data("KLL f32 item"))?; + Self::new(value).map_err(|_| Error::deserial("KLL float must not be NaN")) + } +} + +impl KllValue for KllFloat { + const MIN_SERIALIZED_SIZE: usize = size_of::(); + + fn serialized_size(_value: &Self) -> usize { + size_of::() + } + + fn serialize(value: &Self, bytes: &mut SketchBytes) { + bytes.write_f64_le(value.0); + } + + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + let value = input + .read_f64_le() + .map_err(insufficient_data("KLL f64 item"))?; + Self::new(value).map_err(|_| Error::deserial("KLL float must not be NaN")) + } +} + +impl KllValue for i64 { + const MIN_SERIALIZED_SIZE: usize = 8; + + fn serialized_size(_value: &Self) -> usize { + 8 + } + + fn serialize(value: &Self, bytes: &mut SketchBytes) { + bytes.write_i64_le(*value); + } + + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + input + .read_i64_le() + .map_err(insufficient_data("KLL i64 item")) + } +} + +impl KllValue for String { + const MIN_SERIALIZED_SIZE: usize = 4; + + fn serialized_size(value: &Self) -> usize { + 4 + value.len() + } + + fn serialize(value: &Self, bytes: &mut SketchBytes) { + bytes.write_u32_le(value.len() as u32); + bytes.write(value.as_bytes()); + } + + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + let len = input + .read_u32_le() + .map_err(insufficient_data("KLL string length"))? as usize; + let available_bytes = input.remaining().len(); + if available_bytes < len { + return Err(Error::insufficient_data_of( + "KLL string payload", + format_args!("expected {len} bytes, got {available_bytes}"), + )); + } + let value = std::str::from_utf8(&input.remaining()[..len]) + .map_err(|error| Error::deserial(format!("invalid UTF-8 string: {error}")))? + .to_owned(); + input.advance(len as u64); + Ok(value) + } +} diff --git a/datasketches/src/lib.rs b/datasketches/src/lib.rs index 3dac470a..b4ff0164 100644 --- a/datasketches/src/lib.rs +++ b/datasketches/src/lib.rs @@ -39,8 +39,9 @@ //! * Use `countmin` for point-frequency estimates and `frequencies` for discovering heavy hitters. //! * Use `hll` for fast distinct counts, `cpc` for compact serialized distinct counts, or `theta` //! when set operations are required. -//! * Use `req` or `tdigest` for ranks and quantiles. REQ targets configurable high- or low-rank -//! accuracy; T-Digest emphasizes distribution tails. +//! * Use `kll`, `req`, or `tdigest` for ranks and quantiles. KLL provides strong general-purpose +//! rank accuracy, REQ targets configurable high- or low-rank accuracy, and T-Digest emphasizes +//! distribution tails. //! * Use `tuple` when retained Theta keys need application-defined summaries. //! //! See each module's documentation for accuracy, memory, serialization, and update examples. diff --git a/datasketches/src/req/compactor.rs b/datasketches/src/req/compactor.rs index a81acf70..8f9cc92d 100644 --- a/datasketches/src/req/compactor.rs +++ b/datasketches/src/req/compactor.rs @@ -52,7 +52,7 @@ pub struct Compactor { /// Whether this compactor is configured for high rank accuracy rank_accuracy: RankAccuracy, - /// Raw section size (may be fractional) + /// Raw section size (maybe fractional) section_size_raw: f32, /// Random bit for compaction coin: bool, diff --git a/datasketches/src/req/mod.rs b/datasketches/src/req/mod.rs index 8714f5ec..ce5a65fe 100644 --- a/datasketches/src/req/mod.rs +++ b/datasketches/src/req/mod.rs @@ -35,9 +35,9 @@ //! # Example //! //! ``` +//! use datasketches::common::SearchCriteria; //! use datasketches::req::ReqFloat; //! use datasketches::req::ReqSketch; -//! use datasketches::req::SearchCriteria; //! //! let mut sketch = ReqSketch::default(); //! for value in [1.0, 2.0, 3.0] { @@ -61,7 +61,6 @@ pub use self::sketch::ReqSketch; pub use self::sorted_view::SortedView; pub use self::value::ReqFloat; pub use self::value::ReqValue; - /// Default value of `k` if not specified. Roughly 1% relative error at 95% confidence. const DEFAULT_K: u16 = 12; /// Minimum allowed value of `k`. @@ -79,16 +78,6 @@ pub enum RankAccuracy { LowRank, } -/// Selects the rank definition used by rank, quantile, PMF, and CDF queries. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum SearchCriteria { - /// Define rank as the fraction of values less than or equal to the boundary. - #[default] - Inclusive, - /// Define rank as the fraction of values strictly less than the boundary. - Exclusive, -} - /// Number of sections in a newly created compactor. The section count and size /// determine its capacity and compaction range; the count doubles as its state grows. const INITIAL_SECTIONS_PER_COMPACTOR: u8 = 3; diff --git a/datasketches/src/req/sketch.rs b/datasketches/src/req/sketch.rs index 5826ce31..89a3a70c 100644 --- a/datasketches/src/req/sketch.rs +++ b/datasketches/src/req/sketch.rs @@ -22,13 +22,13 @@ use crate::codec::SketchSlice; use crate::codec::assert::insufficient_data; use crate::codec::family::Family; use crate::common::NumStdDev; +use crate::common::SearchCriteria; use crate::error::Error; use crate::req::DEFAULT_K; use crate::req::INITIAL_SECTIONS_PER_COMPACTOR; use crate::req::MAX_K; use crate::req::MIN_K; use crate::req::RankAccuracy; -use crate::req::SearchCriteria; use crate::req::compactor::Compactor; use crate::req::iter::ReqSketchIterator; use crate::req::serialization::FLAG_IS_EMPTY; diff --git a/datasketches/src/req/sorted_view.rs b/datasketches/src/req/sorted_view.rs index d1313acc..8affbbb7 100644 --- a/datasketches/src/req/sorted_view.rs +++ b/datasketches/src/req/sorted_view.rs @@ -17,8 +17,8 @@ //! Sorted view implementation for efficient quantile queries. +use crate::common::SearchCriteria; use crate::error::Error; -use crate::req::SearchCriteria; /// An owned, sorted snapshot of a [`ReqSketch`](crate::req::ReqSketch). /// diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index bedd640d..d496a9c5 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -682,16 +682,16 @@ impl TDigestMut { let required_payload_bytes = centroid_payload_bytes .checked_add(buffered_payload_bytes) .ok_or_else(|| Error::deserial("TDigest payload size exceeds the supported size"))?; - let remaining = cursor.remaining(); - if remaining.len() < required_payload_bytes { - return Err(Error::insufficient_data(format!( - "TDigest payload requires {required_payload_bytes} bytes, got {}", - remaining.len() - ))); - } // Check the whole payload once so fixed-width records can be decoded without per-field I/O. - let (centroid_payload, buffered_payload) = - remaining[..required_payload_bytes].split_at(centroid_payload_bytes); + let available_bytes = cursor.remaining().len(); + if available_bytes < required_payload_bytes { + return Err(Error::insufficient_data_of( + "TDigest payload", + format_args!("expected {required_payload_bytes} bytes, got {available_bytes}"), + )); + } + let payload = &cursor.remaining()[..required_payload_bytes]; + let (centroid_payload, buffered_payload) = payload.split_at(centroid_payload_bytes); let stored_centroids = num_centroids.checked_add(num_buffered).ok_or_else(|| { Error::deserial("num_centroids and num_buffered exceed the supported size") })?; diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index 7fb9ea0b..c293a775 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -739,11 +739,12 @@ impl CompactThetaSketch { let required_bytes = num_entries .checked_mul(size_of::()) .ok_or_else(|| Error::deserial("Theta entry payload length overflows"))?; - if required_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "Theta entries require {required_bytes} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < required_bytes { + return Err(Error::insufficient_data_of( + "Theta entries", + format_args!("expected {required_bytes} bytes, got {available_bytes}"), + )); } let mut entries = Vec::with_capacity(num_entries); for _ in 0..num_entries { @@ -969,11 +970,12 @@ impl CompactThetaSketch { .and_then(|bits| bits.checked_add(7)) .map(|bits| bits / 8) .ok_or_else(|| Error::deserial("Theta compressed payload length overflows"))?; - if required_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "Theta compressed entries require {required_bytes} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < required_bytes { + return Err(Error::insufficient_data_of( + "Theta compressed entries", + format_args!("expected {required_bytes} bytes, got {available_bytes}"), + )); } // unpack blocks of BLOCK_WIDTH deltas diff --git a/datasketches/src/thetafamily/tuple/serialization.rs b/datasketches/src/thetafamily/tuple/serialization.rs index b077f0c4..1f4cfeeb 100644 --- a/datasketches/src/thetafamily/tuple/serialization.rs +++ b/datasketches/src/thetafamily/tuple/serialization.rs @@ -27,6 +27,7 @@ use crate::codec::SketchBytes; use crate::codec::SketchSlice; +use crate::codec::assert::insufficient_data; use crate::error::Error; /// Current serial version written by this implementation. @@ -73,11 +74,9 @@ macro_rules! impl_primitive_summary { } fn deserialize_value(cursor: &mut SketchSlice<'_>) -> Result { - cursor.$read().map_err(|_| { - Error::insufficient_data( - concat!("failed to read ", stringify!($name), " summary bytes").to_string(), - ) - }) + cursor + .$read() + .map_err(insufficient_data(concat!(stringify!($name), " summary"))) } } }; diff --git a/datasketches/src/thetafamily/tuple/sketch.rs b/datasketches/src/thetafamily/tuple/sketch.rs index 0fddc985..2d6b93f4 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -716,11 +716,12 @@ impl CompactTupleSketch { let required_hash_bytes = num_entries .checked_mul(size_of::()) .ok_or_else(|| Error::deserial("Tuple entry payload length overflows"))?; - if required_hash_bytes > cursor.remaining().len() { - return Err(Error::insufficient_data(format!( - "Tuple entry hashes require at least {required_hash_bytes} bytes, got {}", - cursor.remaining().len() - ))); + let available_bytes = cursor.remaining().len(); + if available_bytes < required_hash_bytes { + return Err(Error::insufficient_data_of( + "Tuple entry hashes", + format_args!("expected {required_hash_bytes} bytes, got {available_bytes}"), + )); } let mut retained_entries = Vec::with_capacity(num_entries); for _ in 0..num_entries { diff --git a/tests-integration/tests/countmin_test/sketch.rs b/tests-integration/tests/countmin_test/sketch.rs index 42f8530e..f35eb78a 100644 --- a/tests-integration/tests/countmin_test/sketch.rs +++ b/tests-integration/tests/countmin_test/sketch.rs @@ -267,7 +267,9 @@ fn test_truncated_non_empty_payload_is_rejected_before_table_allocation() { let error = CountMinSketch::::deserialize(&bytes).unwrap_err(); assert_eq!(error.kind(), ErrorKind::InvalidData); - assert!(error.message().contains("payload requires")); + assert!(error.message().contains("CountMin payload")); + assert!(error.message().contains("expected")); + assert!(error.message().contains("got")); } #[test] diff --git a/tests-integration/tests/kll_test/core.rs b/tests-integration/tests/kll_test/core.rs new file mode 100644 index 00000000..c0bcaec0 --- /dev/null +++ b/tests-integration/tests/kll_test/core.rs @@ -0,0 +1,87 @@ +// 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 datasketches::error::ErrorKind; +use datasketches::kll::KllFloat; +use datasketches::kll::KllSketch; + +const DEFAULT_K: u16 = 200; +const MIN_K: u16 = 8; +const MAX_K: u16 = u16::MAX; + +#[test] +fn k_limits() { + KllSketch::::new(MIN_K).unwrap(); + KllSketch::::new(MAX_K).unwrap(); + + let error = KllSketch::::new(MIN_K - 1).unwrap_err(); + assert_eq!(error.kind(), ErrorKind::InvalidArgument); +} + +#[test] +fn empty_and_reset_state() { + let mut sketch = KllSketch::::new(64).unwrap(); + assert!(sketch.is_empty()); + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.n(), 0); + assert_eq!(sketch.num_retained(), 0); + assert_eq!(sketch.min_item(), None); + assert_eq!(sketch.max_item(), None); + + for item in 0..10_000 { + sketch.update(item); + } + assert!(sketch.is_estimation_mode()); + assert!(sketch.num_retained() > 0); + + sketch.reset(); + assert_eq!(sketch.k(), 64); + assert_eq!(sketch.min_k(), 64); + assert!(sketch.is_empty()); + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.n(), 0); + assert_eq!(sketch.num_retained(), 0); + assert_eq!(sketch.min_item(), None); + assert_eq!(sketch.max_item(), None); +} + +#[test] +fn float_adapter_rejects_nan() { + assert_eq!( + KllFloat::::new(f32::NAN).unwrap_err().kind(), + ErrorKind::InvalidArgument + ); + + let mut sketch = KllSketch::new(DEFAULT_K).unwrap(); + sketch.update(KllFloat::::new(0.0).unwrap()); + assert_eq!(sketch.min_item().map(|value| **value), Some(0.0)); +} + +#[test] +fn retained_count_stays_consistent_through_compaction_and_roundtrip() { + let mut sketch = KllSketch::::new(32).unwrap(); + for item in 0..100_000 { + sketch.update(item); + assert!(sketch.num_retained() <= sketch.n() as usize); + } + + let decoded = KllSketch::::deserialize(&sketch.serialize()).unwrap(); + assert_eq!(decoded.n(), sketch.n()); + assert_eq!(decoded.num_retained(), sketch.num_retained()); + assert_eq!(decoded.min_item(), Some(&0)); + assert_eq!(decoded.max_item(), Some(&99_999)); +} diff --git a/tests-integration/tests/kll_test/generic.rs b/tests-integration/tests/kll_test/generic.rs new file mode 100644 index 00000000..14d06677 --- /dev/null +++ b/tests-integration/tests/kll_test/generic.rs @@ -0,0 +1,84 @@ +// 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::cmp::Ordering; + +use datasketches::codec::SketchBytes; +use datasketches::codec::SketchSlice; +use datasketches::common::SearchCriteria; +use datasketches::error::Error; +use datasketches::kll::KllSketch; +use datasketches::kll::KllValue; + +#[derive(Debug, Clone, PartialEq, Eq)] +struct NumericString(String); + +impl Ord for NumericString { + fn cmp(&self, other: &Self) -> Ordering { + self.0 + .parse::() + .unwrap() + .cmp(&other.0.parse::().unwrap()) + } +} + +impl PartialOrd for NumericString { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl KllValue for NumericString { + const MIN_SERIALIZED_SIZE: usize = String::MIN_SERIALIZED_SIZE; + + fn serialized_size(value: &Self) -> usize { + String::serialized_size(&value.0) + } + + fn serialize(value: &Self, bytes: &mut SketchBytes) { + String::serialize(&value.0, bytes); + } + + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + String::deserialize(input).map(Self) + } +} + +#[test] +fn custom_item_order_controls_queries_and_survives_roundtrip() { + let mut sketch = KllSketch::::new(200).unwrap(); + for item in ["2", "10", "1"] { + sketch.update(NumericString(item.to_owned())); + } + + assert_eq!(sketch.min_item().map(|item| item.0.as_str()), Some("1")); + assert_eq!(sketch.max_item().map(|item| item.0.as_str()), Some("10")); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap().0, + "2" + ); + + let decoded = KllSketch::::deserialize(&sketch.serialize()).unwrap(); + assert_eq!(decoded.n(), sketch.n()); + assert_eq!(decoded.num_retained(), sketch.num_retained()); + assert_eq!(decoded.min_item().map(|item| item.0.as_str()), Some("1")); + assert_eq!(decoded.max_item().map(|item| item.0.as_str()), Some("10")); + assert_eq!( + decoded.quantile(0.5, SearchCriteria::Inclusive).unwrap().0, + "2" + ); +} diff --git a/tests-integration/tests/kll_test/main.rs b/tests-integration/tests/kll_test/main.rs index 825a6281..26441bc2 100644 --- a/tests-integration/tests/kll_test/main.rs +++ b/tests-integration/tests/kll_test/main.rs @@ -15,4 +15,7 @@ // specific language governing permissions and limitations // under the License. -mod sketch; +mod core; +mod generic; +mod merge; +mod query; diff --git a/tests-integration/tests/kll_test/merge.rs b/tests-integration/tests/kll_test/merge.rs new file mode 100644 index 00000000..417fa1b1 --- /dev/null +++ b/tests-integration/tests/kll_test/merge.rs @@ -0,0 +1,84 @@ +// 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 datasketches::common::SearchCriteria; +use datasketches::kll::KllSketch; + +#[test] +fn merge_preserves_weight_extrema_and_query_invariants() { + let mut left = KllSketch::::new(200).unwrap(); + let mut right = KllSketch::::new(200).unwrap(); + for item in 0..10_000 { + left.update(item); + right.update(19_999 - item); + } + + left.merge(&right).unwrap(); + + assert_eq!(left.n(), 20_000); + assert_eq!(left.min_item(), Some(&0)); + assert_eq!(left.max_item(), Some(&19_999)); + assert_eq!(left.sorted_view().total_weight(), left.n()); + let quantiles = left + .quantiles(&[0.0, 0.25, 0.5, 0.75, 1.0], SearchCriteria::Inclusive) + .unwrap(); + assert!(quantiles.windows(2).all(|pair| pair[0] <= pair[1])); +} + +#[test] +fn merge_tracks_the_smallest_estimation_k() { + let mut left = KllSketch::::new(256).unwrap(); + let mut right = KllSketch::::new(128).unwrap(); + for item in 0..10_000 { + left.update(item); + right.update(20_000 - item); + } + + left.merge(&right).unwrap(); + + assert_eq!(left.min_k(), right.min_k()); + assert_eq!(left.normalized_rank_error(), right.normalized_rank_error()); + assert_eq!(left.normalized_pmf_error(), right.normalized_pmf_error()); +} + +#[test] +fn merging_an_empty_lower_k_sketch_does_not_change_accuracy() { + let mut sketch = KllSketch::::new(256).unwrap(); + for item in 0..10_000 { + sketch.update(item); + } + let empty = KllSketch::::new(128).unwrap(); + let rank_error = sketch.normalized_rank_error(); + + sketch.merge(&empty).unwrap(); + + assert_eq!(sketch.n(), 10_000); + assert_eq!(sketch.normalized_rank_error(), rank_error); +} + +#[test] +fn merge_updates_extrema_from_either_side() { + let mut first = KllSketch::::new(200).unwrap(); + let mut second = KllSketch::::new(200).unwrap(); + first.update(1); + second.update(2); + + second.merge(&first).unwrap(); + + assert_eq!(second.min_item(), Some(&1)); + assert_eq!(second.max_item(), Some(&2)); +} diff --git a/tests-integration/tests/kll_test/query.rs b/tests-integration/tests/kll_test/query.rs new file mode 100644 index 00000000..5474df95 --- /dev/null +++ b/tests-integration/tests/kll_test/query.rs @@ -0,0 +1,150 @@ +// 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 datasketches::common::SearchCriteria; +use datasketches::error::ErrorKind; +use datasketches::kll::KllSketch; + +const DEFAULT_K: u16 = 200; +const NUMERIC_NOISE_TOLERANCE: f64 = 1e-6; + +#[test] +fn empty_and_invalid_queries_return_errors() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + assert!(sketch.rank(&0, SearchCriteria::Inclusive).is_err()); + assert!(sketch.quantile(0.5, SearchCriteria::Inclusive).is_err()); + assert!(sketch.pmf(&[0], SearchCriteria::Inclusive).is_err()); + assert!(sketch.cdf(&[0], SearchCriteria::Inclusive).is_err()); + + sketch.update(0); + for rank in [-1.0, f64::NAN, 1.1] { + let error = sketch + .quantile(rank, SearchCriteria::Inclusive) + .unwrap_err(); + assert_eq!(error.kind(), ErrorKind::InvalidArgument); + } + let error = sketch.cdf(&[1, 0], SearchCriteria::Inclusive).unwrap_err(); + assert_eq!(error.kind(), ErrorKind::InvalidArgument); +} + +#[test] +fn inclusive_and_exclusive_semantics_cover_duplicates() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + for item in [1, 1, 2, 2] { + sketch.update(item); + } + + assert_eq!(sketch.rank(&1, SearchCriteria::Exclusive).unwrap(), 0.0); + assert_eq!(sketch.rank(&1, SearchCriteria::Inclusive).unwrap(), 0.5); + assert_eq!(sketch.rank(&2, SearchCriteria::Exclusive).unwrap(), 0.5); + assert_eq!(sketch.rank(&2, SearchCriteria::Inclusive).unwrap(), 1.0); + assert_eq!(sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), 1); + assert_eq!(sketch.quantile(0.5, SearchCriteria::Exclusive).unwrap(), 2); +} + +#[test] +fn exact_mode_queries_match_the_stream() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + for item in 1..=100 { + sketch.update(item); + } + + assert_eq!(sketch.quantile(0.0, SearchCriteria::Inclusive).unwrap(), 1); + assert_eq!(sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), 50); + assert_eq!( + sketch.quantile(1.0, SearchCriteria::Inclusive).unwrap(), + 100 + ); + for item in 1..=100 { + assert_eq!( + sketch.rank(&item, SearchCriteria::Inclusive).unwrap(), + item as f64 / 100.0 + ); + } +} + +#[test] +fn estimation_mode_queries_preserve_deterministic_invariants() { + let mut sketch = KllSketch::::new(64).unwrap(); + for item in 0..10_000 { + sketch.update(item); + } + + let mut previous_rank = 0.0; + for item in (0..10_000).step_by(100) { + let rank = sketch.rank(&item, SearchCriteria::Inclusive).unwrap(); + assert!(rank >= previous_rank); + assert!((0.0..=1.0).contains(&rank)); + previous_rank = rank; + } + assert_eq!(sketch.min_item(), Some(&0)); + assert_eq!(sketch.max_item(), Some(&9_999)); + assert!(sketch.normalized_rank_error() < sketch.normalized_pmf_error()); +} + +#[test] +fn rank_cdf_and_pmf_are_consistent() { + let mut sketch = KllSketch::::new(64).unwrap(); + for item in 0..10_000 { + sketch.update(item); + } + let split_points: Vec<_> = (100..10_000).step_by(100).collect(); + + for criteria in [SearchCriteria::Inclusive, SearchCriteria::Exclusive] { + let cdf = sketch.cdf(&split_points, criteria).unwrap(); + let pmf = sketch.pmf(&split_points, criteria).unwrap(); + let mut subtotal = 0.0; + for (index, split_point) in split_points.iter().enumerate() { + subtotal += pmf[index]; + assert!((cdf[index] - subtotal).abs() <= NUMERIC_NOISE_TOLERANCE); + assert_eq!(cdf[index], sketch.rank(split_point, criteria).unwrap()); + } + assert!((pmf.iter().sum::() - 1.0).abs() <= NUMERIC_NOISE_TOLERANCE); + } +} + +#[test] +fn sorted_view_supports_repeated_and_batch_queries() { + let mut sketch = KllSketch::::new(64).unwrap(); + for item in 0..1_000 { + sketch.update(item); + } + let view = sketch.sorted_view(); + let ranks = [0.0, 0.25, 0.5, 0.75, 1.0]; + let quantiles = sketch.quantiles(&ranks, SearchCriteria::Inclusive).unwrap(); + + assert_eq!(view.len(), sketch.num_retained()); + assert_eq!(view.total_weight(), sketch.n()); + assert_eq!( + view.quantiles(&ranks, SearchCriteria::Inclusive).unwrap(), + quantiles + ); + for (&rank, quantile) in ranks.iter().zip(&quantiles) { + assert_eq!( + view.quantile(rank, SearchCriteria::Inclusive).unwrap(), + *quantile + ); + assert_eq!( + view.rank(quantile, SearchCriteria::Inclusive).unwrap(), + sketch.rank(quantile, SearchCriteria::Inclusive).unwrap() + ); + } + + sketch.update(2_000); + assert_eq!(view.total_weight(), 1_000); + assert_eq!(view.quantile(1.0, SearchCriteria::Inclusive).unwrap(), 999); +} diff --git a/tests-integration/tests/kll_test/sketch.rs b/tests-integration/tests/kll_test/sketch.rs deleted file mode 100644 index d9acaaa8..00000000 --- a/tests-integration/tests/kll_test/sketch.rs +++ /dev/null @@ -1,397 +0,0 @@ -// 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::cmp::Ordering; - -use datasketches::error::ErrorKind; -use datasketches::kll::DEFAULT_K; -use datasketches::kll::KllComparator; -use datasketches::kll::KllSketch; -use datasketches::kll::MAX_K; -use datasketches::kll::MIN_K; - -const NUMERIC_NOISE_TOLERANCE: f64 = 1e-6; - -fn assert_approx_eq(actual: f64, expected: f64, tolerance: f64) { - let delta = (actual - expected).abs(); - assert!( - delta <= tolerance, - "expected {expected} +/- {tolerance}, got {actual}" - ); -} - -fn rank_eps(sketch: &KllSketch) -> f64 { - sketch.normalized_rank_error(false) -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct NumericStringOrder; - -impl KllComparator for NumericStringOrder { - fn compare(&self, left: &String, right: &String) -> Ordering { - left.parse::() - .unwrap() - .cmp(&right.parse::().unwrap()) - } -} - -#[test] -fn test_k_limits() { - let _min = KllSketch::::new(MIN_K).unwrap(); - let _max = KllSketch::::new(MAX_K).unwrap(); -} - -#[test] -fn test_k_too_small_returns_error() { - let error = KllSketch::::new(MIN_K - 1).unwrap_err(); - assert_eq!(error.kind(), ErrorKind::InvalidArgument); -} - -#[test] -fn test_empty() { - let sketch = KllSketch::::new(DEFAULT_K).unwrap(); - 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, true).is_none()); - assert!(sketch.quantile(0.5, true).is_none()); - assert!(sketch.pmf(&[0.0f32], true).is_none()); - assert!(sketch.cdf(&[0.0f32], true).is_none()); -} - -#[test] -#[should_panic(expected = "rank must be in [0.0, 1.0]")] -fn test_quantile_out_of_range_panics() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(0.0); - sketch.quantile(-1.0, true); -} - -#[test] -fn test_one_item() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(1.0); - assert!(!sketch.is_empty()); - assert!(!sketch.is_estimation_mode()); - assert_eq!(sketch.n(), 1); - assert_eq!(sketch.num_retained(), 1); - assert_eq!(sketch.rank(&1.0, false), Some(0.0)); - assert_eq!(sketch.rank(&1.0, true), Some(1.0)); - assert_eq!(sketch.rank(&2.0, false), Some(1.0)); - assert_eq!(sketch.min_item().cloned(), Some(1.0)); - assert_eq!(sketch.max_item().cloned(), Some(1.0)); - assert_eq!(sketch.quantile(0.5, true), Some(1.0)); -} - -#[test] -fn test_duplicate_items_follow_inclusive_and_exclusive_semantics() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - for item in [1.0, 1.0, 2.0, 2.0] { - sketch.update(item); - } - - assert_eq!(sketch.rank(&1.0, false), Some(0.0)); - assert_eq!(sketch.rank(&1.0, true), Some(0.5)); - assert_eq!(sketch.rank(&2.0, false), Some(0.5)); - assert_eq!(sketch.rank(&2.0, true), Some(1.0)); - assert_eq!(sketch.quantile(0.5, true), Some(1.0)); - assert_eq!(sketch.quantile(0.5, false), Some(2.0)); -} - -#[test] -fn test_nan_is_ignored() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(f32::NAN); - assert!(sketch.is_empty()); - sketch.update(0.0); - sketch.update(f32::NAN); - assert_eq!(sketch.n(), 1); -} - -#[test] -fn test_many_items_exact_mode() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - let n = DEFAULT_K as usize; - for i in 1..=n { - sketch.update(i as f32); - assert_eq!(sketch.n(), i as u64); - } - assert!(!sketch.is_empty()); - assert!(!sketch.is_estimation_mode()); - assert_eq!(sketch.num_retained(), n); - assert_eq!(sketch.min_item().cloned(), Some(1.0)); - assert_eq!(sketch.quantile(0.0, true), Some(1.0)); - assert_eq!(sketch.max_item().cloned(), Some(n as f32)); - assert_eq!(sketch.quantile(1.0, true), Some(n as f32)); - - for i in 1..=n { - let inclusive_rank = i as f64 / n as f64; - assert_eq!(sketch.rank(&(i as f32), true), Some(inclusive_rank)); - let exclusive_rank = (i - 1) as f64 / n as f64; - assert_eq!(sketch.rank(&(i as f32), false), Some(exclusive_rank)); - } -} - -#[test] -fn test_ten_items_quantiles() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - for i in 1..=10 { - sketch.update(i as f32); - } - assert_eq!(sketch.quantile(0.0, true), Some(1.0)); - assert_eq!(sketch.quantile(0.5, true), Some(5.0)); - assert_eq!(sketch.quantile(0.99, true), Some(10.0)); - assert_eq!(sketch.quantile(1.0, true), Some(10.0)); -} - -#[test] -fn test_hundred_items_quantiles() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - for i in 0..100 { - sketch.update(i as f32); - } - assert_eq!(sketch.quantile(0.0, true), Some(0.0)); - assert_eq!(sketch.quantile(0.01, true), Some(0.0)); - assert_eq!(sketch.quantile(0.5, true), Some(49.0)); - assert_eq!(sketch.quantile(0.99, true), Some(98.0)); - assert_eq!(sketch.quantile(1.0, true), Some(99.0)); -} - -#[test] -fn test_many_items_estimation_mode_rank_error() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - let n = 10_000; - for i in 0..n { - sketch.update(i as f32); - } - assert!(!sketch.is_empty()); - assert!(sketch.is_estimation_mode()); - assert_eq!(sketch.min_item().cloned(), Some(0.0)); - assert_eq!(sketch.max_item().cloned(), Some((n - 1) as f32)); - - let rank_eps = rank_eps(&sketch); - for i in (0..n).step_by(10) { - let true_rank = i as f64 / n as f64; - let rank = sketch.rank(&(i as f32), false).unwrap(); - assert_approx_eq(rank, true_rank, rank_eps); - } - - assert!(sketch.num_retained() > 0); -} - -#[test] -fn test_rank_cdf_pmf_consistency() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - let n = 200; - let mut values = Vec::with_capacity(n); - for i in 0..n { - sketch.update(i as f32); - values.push(i as f32); - } - - let ranks = sketch.cdf(&values, false).unwrap(); - let pmf = sketch.pmf(&values, false).unwrap(); - - let mut subtotal = 0.0; - for i in 0..n { - let rank = sketch.rank(&values[i], false).unwrap(); - assert_eq!(rank, ranks[i]); - subtotal += pmf[i]; - assert!( - (ranks[i] - subtotal).abs() <= NUMERIC_NOISE_TOLERANCE, - "cdf vs pmf mismatch at index {i}" - ); - } - - let ranks = sketch.cdf(&values, true).unwrap(); - let pmf = sketch.pmf(&values, true).unwrap(); - - let mut subtotal = 0.0; - for i in 0..n { - let rank = sketch.rank(&values[i], true).unwrap(); - assert_eq!(rank, ranks[i]); - subtotal += pmf[i]; - assert!( - (ranks[i] - subtotal).abs() <= NUMERIC_NOISE_TOLERANCE, - "cdf vs pmf mismatch at index {i}" - ); - } -} - -#[test] -#[should_panic(expected = "split_points must be unique and monotonically increasing")] -fn test_out_of_order_split_points_panics() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(0.0); - let split_points = [1.0, 0.0]; - let _ = sketch.cdf(&split_points, true); -} - -#[test] -#[should_panic(expected = "split_points must not contain NaN values")] -fn test_nan_split_point_panics() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(0.0); - let split_points = [f32::NAN]; - let _ = sketch.cdf(&split_points, true); -} - -#[test] -fn test_merge() { - let mut sketch1 = KllSketch::::new(DEFAULT_K).unwrap(); - let mut sketch2 = KllSketch::::new(DEFAULT_K).unwrap(); - let n = 10_000; - for i in 0..n { - sketch1.update(i as f32); - sketch2.update((2 * n - i - 1) as f32); - } - - assert_eq!(sketch1.min_item().cloned(), Some(0.0)); - assert_eq!(sketch1.max_item().cloned(), Some((n - 1) as f32)); - assert_eq!(sketch2.min_item().cloned(), Some(n as f32)); - assert_eq!(sketch2.max_item().cloned(), Some((2 * n - 1) as f32)); - - sketch1.merge(&sketch2); - - assert!(!sketch1.is_empty()); - assert_eq!(sketch1.n(), (2 * n) as u64); - assert_eq!(sketch1.min_item().cloned(), Some(0.0)); - assert_eq!(sketch1.max_item().cloned(), Some((2 * n - 1) as f32)); - let median = sketch1.quantile(0.5, true).unwrap(); - let rank_eps = rank_eps(&sketch1); - assert_approx_eq(median as f64, n as f64, n as f64 * rank_eps); -} - -#[test] -fn test_merge_lower_k() { - let mut sketch1 = KllSketch::::new(256).unwrap(); - let mut sketch2 = KllSketch::::new(128).unwrap(); - let n = 10_000; - for i in 0..n { - sketch1.update(i as f32); - sketch2.update((2 * n - i - 1) as f32); - } - - sketch1.merge(&sketch2); - - assert_eq!(sketch1.n(), (2 * n) as u64); - assert_eq!(sketch1.min_item().cloned(), Some(0.0)); - assert_eq!(sketch1.max_item().cloned(), Some((2 * n - 1) as f32)); - assert_eq!( - sketch1.normalized_rank_error(false), - sketch2.normalized_rank_error(false) - ); - assert_eq!( - sketch1.normalized_rank_error(true), - sketch2.normalized_rank_error(true) - ); - let median = sketch1.quantile(0.5, true).unwrap(); - let rank_eps = rank_eps(&sketch1); - assert_approx_eq(median as f64, n as f64, n as f64 * rank_eps); -} - -#[test] -fn test_merge_exact_mode_lower_k() { - let mut sketch1 = KllSketch::::new(256).unwrap(); - let sketch2 = KllSketch::::new(128).unwrap(); - let n = 10_000; - for i in 0..n { - sketch1.update(i as f32); - } - - let err_before = sketch1.normalized_rank_error(true); - sketch1.merge(&sketch2); - assert_eq!(sketch1.normalized_rank_error(true), err_before); - - assert_eq!(sketch1.n(), n as u64); - assert_eq!(sketch1.min_item().cloned(), Some(0.0)); - assert_eq!(sketch1.max_item().cloned(), Some((n - 1) as f32)); - let median = sketch1.quantile(0.5, true).unwrap(); - let rank_eps = rank_eps(&sketch1); - assert_approx_eq(median as f64, (n / 2) as f64, (n as f64 / 2.0) * rank_eps); -} - -#[test] -fn test_merge_min_max_from_other() { - let mut sketch1 = KllSketch::::new(DEFAULT_K).unwrap(); - let mut sketch2 = KllSketch::::new(DEFAULT_K).unwrap(); - sketch1.update(1.0); - sketch2.update(2.0); - sketch2.merge(&sketch1); - assert_eq!(sketch2.min_item().cloned(), Some(1.0)); - assert_eq!(sketch2.max_item().cloned(), Some(2.0)); -} - -#[test] -fn test_merge_min_max_large_other() { - let mut sketch1 = KllSketch::::new(DEFAULT_K).unwrap(); - for i in 0..1_000_000 { - sketch1.update(i as f32); - } - let mut sketch2 = KllSketch::::new(DEFAULT_K).unwrap(); - sketch2.merge(&sketch1); - assert_eq!(sketch2.min_item().cloned(), Some(0.0)); - assert_eq!(sketch2.max_item().cloned(), Some(999_999.0)); -} - -#[test] -fn test_reset_retains_configuration() { - let mut sketch = KllSketch::::new(64).unwrap(); - for i in 0..10_000 { - sketch.update(i as f32); - } - assert!(sketch.is_estimation_mode()); - - sketch.reset(); - - assert_eq!(sketch.k(), 64); - assert_eq!(sketch.min_k(), 64); - assert!(sketch.is_empty()); - assert!(!sketch.is_estimation_mode()); - assert_eq!(sketch.n(), 0); - assert_eq!(sketch.num_retained(), 0); - assert_eq!(sketch.min_item(), None); - assert_eq!(sketch.max_item(), None); -} - -#[test] -fn test_custom_comparator_roundtrip() { - let mut sketch = - KllSketch::::new_with_comparator(200, NumericStringOrder) - .unwrap(); - for item in ["2", "10", "1"] { - sketch.update(item.to_owned()); - } - - assert_eq!(sketch.min_item().map(String::as_str), Some("1")); - assert_eq!(sketch.max_item().map(String::as_str), Some("10")); - assert_eq!(sketch.quantile(0.5, true).as_deref(), Some("2")); - - let bytes = sketch.serialize(); - let decoded = KllSketch::::deserialize_with_comparator( - &bytes, - NumericStringOrder, - ) - .unwrap(); - assert_eq!(decoded.n(), sketch.n()); - assert_eq!(decoded.min_item().map(String::as_str), Some("1")); - assert_eq!(decoded.max_item().map(String::as_str), Some("10")); - assert_eq!(decoded.quantile(0.5, true).as_deref(), Some("2")); -} diff --git a/tests-integration/tests/req_test/accuracy.rs b/tests-integration/tests/req_test/accuracy.rs index d68f8c21..76c318cc 100644 --- a/tests-integration/tests/req_test/accuracy.rs +++ b/tests-integration/tests/req_test/accuracy.rs @@ -17,9 +17,9 @@ //! End-to-end accuracy checks for ReqSketch. +use datasketches::common::SearchCriteria; use datasketches::error::Error; use datasketches::req::ReqSketch; -use datasketches::req::SearchCriteria; use googletest::assert_that; use googletest::prelude::le; diff --git a/tests-integration/tests/req_test/bounds.rs b/tests-integration/tests/req_test/bounds.rs index 5caab948..d6ee3a77 100644 --- a/tests-integration/tests/req_test/bounds.rs +++ b/tests-integration/tests/req_test/bounds.rs @@ -18,10 +18,10 @@ //! Rank error bounds and sigma coverage for ReqSketch. use datasketches::common::NumStdDev; +use datasketches::common::SearchCriteria; use datasketches::error::Error; use datasketches::req::RankAccuracy; use datasketches::req::ReqSketch; -use datasketches::req::SearchCriteria; use googletest::assert_that; use googletest::prelude::all; use googletest::prelude::ge; diff --git a/tests-integration/tests/req_test/core.rs b/tests-integration/tests/req_test/core.rs index 0b435950..737b4892 100644 --- a/tests-integration/tests/req_test/core.rs +++ b/tests-integration/tests/req_test/core.rs @@ -17,11 +17,11 @@ //! Core ReqSketch construction and update behavior. +use datasketches::common::SearchCriteria; use datasketches::error::Error; use datasketches::error::ErrorKind; use datasketches::req::RankAccuracy; use datasketches::req::ReqSketch; -use datasketches::req::SearchCriteria; use googletest::assert_that; use googletest::prelude::all; use googletest::prelude::anything; diff --git a/tests-integration/tests/req_test/generic.rs b/tests-integration/tests/req_test/generic.rs index 26da8587..a405e815 100644 --- a/tests-integration/tests/req_test/generic.rs +++ b/tests-integration/tests/req_test/generic.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. +use datasketches::common::SearchCriteria; use datasketches::req::ReqSketch; -use datasketches::req::SearchCriteria; #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] struct Reading(i32); diff --git a/tests-integration/tests/req_test/merge.rs b/tests-integration/tests/req_test/merge.rs index 6f1b9caa..d4c80e29 100644 --- a/tests-integration/tests/req_test/merge.rs +++ b/tests-integration/tests/req_test/merge.rs @@ -17,9 +17,9 @@ //! Merge behavior for ReqSketch. +use datasketches::common::SearchCriteria; use datasketches::req::RankAccuracy; use datasketches::req::ReqSketch; -use datasketches::req::SearchCriteria; use googletest::assert_that; use googletest::prelude::anything; use googletest::prelude::err; diff --git a/tests-integration/tests/req_test/property.rs b/tests-integration/tests/req_test/property.rs index 9df8c2c1..cf7a8bd5 100644 --- a/tests-integration/tests/req_test/property.rs +++ b/tests-integration/tests/req_test/property.rs @@ -18,8 +18,8 @@ //! Property-based ReqSketch tests. use datasketches::common::NumStdDev; +use datasketches::common::SearchCriteria; use datasketches::req::ReqSketch; -use datasketches::req::SearchCriteria; use quickcheck::Gen; use quickcheck::QuickCheck; use quickcheck::TestResult; diff --git a/tests-integration/tests/req_test/query.rs b/tests-integration/tests/req_test/query.rs index 0404deca..4fb2e62b 100644 --- a/tests-integration/tests/req_test/query.rs +++ b/tests-integration/tests/req_test/query.rs @@ -17,9 +17,9 @@ //! Rank, quantile, PMF, and CDF behavior for ReqSketch. +use datasketches::common::SearchCriteria; use datasketches::error::Error; use datasketches::req::ReqSketch; -use datasketches::req::SearchCriteria; use googletest::assert_that; use googletest::prelude::all; use googletest::prelude::ge; diff --git a/tests-integration/tests/req_test/sorted_view_api.rs b/tests-integration/tests/req_test/sorted_view_api.rs index 6882f2a1..59c4e6dc 100644 --- a/tests-integration/tests/req_test/sorted_view_api.rs +++ b/tests-integration/tests/req_test/sorted_view_api.rs @@ -19,9 +19,9 @@ //! distribution queries take `&self`, and `sorted_view()` returns an owned //! snapshot instead of relying on an internal cache. +use datasketches::common::SearchCriteria; use datasketches::error::ErrorKind; use datasketches::req::ReqSketch; -use datasketches::req::SearchCriteria; use datasketches::req::SortedView; use googletest::assert_that; use googletest::prelude::all; diff --git a/tests-integration/tests/serde_tests/frequencies.rs b/tests-integration/tests/serde_tests/frequencies.rs index 71bd62cb..a1c8f22e 100644 --- a/tests-integration/tests/serde_tests/frequencies.rs +++ b/tests-integration/tests/serde_tests/frequencies.rs @@ -85,7 +85,10 @@ fn test_string_deserialize_rejects_length_larger_than_input() { let error = String::deserialize_value(&mut cursor).unwrap_err(); assert_eq!(error.kind(), ErrorKind::InvalidData); - assert_that!(error.message(), contains_substring("exceeds the remaining")); + assert_that!( + error.message(), + contains_substring("expected 1024 bytes, got 0") + ); } #[test] diff --git a/tests-integration/tests/serde_tests/kll.rs b/tests-integration/tests/serde_tests/kll.rs index 70d9249a..44ba22b1 100644 --- a/tests-integration/tests/serde_tests/kll.rs +++ b/tests-integration/tests/serde_tests/kll.rs @@ -20,6 +20,7 @@ //! These tests verify binary compatibility with Apache DataSketches implementations: //! - Java (datasketches-java) //! - C++ (datasketches-cpp) +//! - Go (datasketches-go) //! //! Test data is generated by the reference implementations and stored in: //! `tests/serde_tests/{java_generated_files,cpp_generated_files}/`. @@ -28,24 +29,81 @@ use std::cmp::Ordering; use std::fs; use std::path::PathBuf; -use datasketches::kll::DEFAULT_K; -use datasketches::kll::KllComparator; +use datasketches::codec::SketchBytes; +use datasketches::codec::SketchSlice; +use datasketches::error::Error; +use datasketches::error::ErrorKind; +use datasketches::kll::KllFloat; use datasketches::kll::KllSketch; +use datasketches::kll::KllValue; use crate::serialization_test_data; -#[derive(Clone, Copy)] -struct NumericStringOrder; +const DEFAULT_K: u16 = 200; -impl KllComparator for NumericStringOrder { - fn compare(&self, left: &String, right: &String) -> Ordering { - parse_string_value(left).cmp(&parse_string_value(right)) +#[derive(Debug, Clone, PartialEq, Eq)] +struct NumericString(String); + +impl Ord for NumericString { + fn cmp(&self, other: &Self) -> Ordering { + parse_string_value(&self.0).cmp(&parse_string_value(&other.0)) + } +} + +impl PartialOrd for NumericString { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl KllValue for NumericString { + const MIN_SERIALIZED_SIZE: usize = String::MIN_SERIALIZED_SIZE; + + fn serialized_size(value: &Self) -> usize { + String::serialized_size(&value.0) + } + + fn serialize(value: &Self, bytes: &mut SketchBytes) { + String::serialize(&value.0, bytes); + } + + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + String::deserialize(input).map(Self) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +struct Record { + key: i64, + category: u16, +} + +impl KllValue for Record { + const MIN_SERIALIZED_SIZE: usize = 10; + + fn serialized_size(_value: &Self) -> usize { + 10 + } + + fn serialize(value: &Self, bytes: &mut SketchBytes) { + bytes.write_i64_le(value.key); + bytes.write_u16_le(value.category); + } + + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + let key = input + .read_i64_le() + .map_err(|_| Error::new(ErrorKind::InvalidData, "missing record key"))?; + let category = input + .read_u16_le() + .map_err(|_| Error::new(ErrorKind::InvalidData, "missing record category"))?; + Ok(Self { key, category }) } } fn test_f32_file(path: PathBuf, expected_n: usize) { let bytes = fs::read(&path).unwrap(); - let sketch = KllSketch::::deserialize(&bytes) + let sketch = KllSketch::>::deserialize(&bytes) .unwrap_or_else(|error| panic!("{}: {error}", path.display())); assert_eq!(sketch.k(), DEFAULT_K, "wrong k in {}", path.display()); @@ -73,13 +131,13 @@ fn test_f32_file(path: PathBuf, expected_n: usize) { assert!(sketch.max_item().is_none(), "max should be None"); } else { assert_eq!( - sketch.min_item().cloned(), + sketch.min_item().map(|value| **value), Some(1.0), "min item mismatch in {}", path.display() ); assert_eq!( - sketch.max_item().cloned(), + sketch.max_item().map(|value| **value), Some(expected_n as f32), "max item mismatch in {}", path.display() @@ -91,7 +149,7 @@ fn test_f32_file(path: PathBuf, expected_n: usize) { fn test_f64_file(path: PathBuf, expected_n: usize) { let bytes = fs::read(&path).unwrap(); - let sketch = KllSketch::::deserialize(&bytes) + let sketch = KllSketch::>::deserialize(&bytes) .unwrap_or_else(|error| panic!("{}: {error}", path.display())); assert_eq!(sketch.k(), DEFAULT_K, "wrong k in {}", path.display()); @@ -119,13 +177,13 @@ fn test_f64_file(path: PathBuf, expected_n: usize) { assert!(sketch.max_item().is_none(), "max should be None"); } else { assert_eq!( - sketch.min_item().cloned(), + sketch.min_item().map(|value| **value), Some(1.0), "min item mismatch in {}", path.display() ); assert_eq!( - sketch.max_item().cloned(), + sketch.max_item().map(|value| **value), Some(expected_n as f64), "max item mismatch in {}", path.display() @@ -190,11 +248,8 @@ fn parse_string_value(value: &str) -> u64 { fn test_string_file(path: PathBuf, expected_n: usize) { let bytes = fs::read(&path).unwrap(); - let sketch = KllSketch::::deserialize_with_comparator( - &bytes, - NumericStringOrder, - ) - .unwrap_or_else(|error| panic!("{}: {error}", path.display())); + let sketch = KllSketch::::deserialize(&bytes) + .unwrap_or_else(|error| panic!("{}: {error}", path.display())); assert_eq!(sketch.k(), DEFAULT_K, "wrong k in {}", path.display()); assert_eq!( @@ -223,13 +278,13 @@ fn test_string_file(path: PathBuf, expected_n: usize) { let min_item = sketch.min_item().expect("missing min item"); let max_item = sketch.max_item().expect("missing max item"); assert_eq!( - parse_string_value(min_item), + parse_string_value(&min_item.0), 1, "min item mismatch in {}", path.display() ); assert_eq!( - parse_string_value(max_item), + parse_string_value(&max_item.0), expected_n as u64, "max item mismatch in {}", path.display() @@ -319,19 +374,81 @@ fn test_cpp_kll_string_compatibility() { } } +#[test] +fn test_go_kll_float_compatibility() { + for n in [0, 1, 10, 100, 1000, 10000, 100000, 1000000] { + let path = serialization_test_data("go_generated_files", &format!("kll_float_n{n}_go.sk")); + test_f32_file(path, n); + } +} + +#[test] +fn test_go_kll_double_compatibility() { + for n in [0, 1, 10, 100, 1000, 10000, 100000, 1000000] { + let path = serialization_test_data("go_generated_files", &format!("kll_double_n{n}_go.sk")); + test_f64_file(path, n); + } +} + +#[test] +fn test_go_kll_long_compatibility() { + for n in [0, 1, 10, 100, 1000, 10000, 100000, 1000000] { + let path = serialization_test_data("go_generated_files", &format!("kll_long_n{n}_go.sk")); + test_i64_file(path, n); + } +} + +#[test] +fn test_go_kll_string_compatibility() { + for n in [0, 1, 10, 100, 1000, 10000, 100000, 1000000] { + let path = serialization_test_data("go_generated_files", &format!("kll_string_n{n}_go.sk")); + test_string_file(path, n); + } +} + +#[test] +fn test_custom_kll_value_roundtrip() { + let mut sketch = KllSketch::::new(64).unwrap(); + for key in 0..1_000 { + sketch.update(Record { + key, + category: (key % 7) as u16, + }); + } + + let decoded = KllSketch::::deserialize(&sketch.serialize()).unwrap(); + + assert_eq!(decoded, sketch); + assert_eq!(decoded.n(), 1_000); + assert_eq!(decoded.num_retained(), sketch.num_retained()); +} + #[test] fn test_rejects_truncated_or_trailing_data() { - let mut sketch = KllSketch::::default(); + let mut sketch = KllSketch::>::default(); for value in 0..1_000 { - sketch.update(value as f32); + sketch.update(KllFloat::::new(value as f32).unwrap()); } let bytes = sketch.serialize(); for length in [7, 15, bytes.len() - 1] { - assert!(KllSketch::::deserialize(&bytes[..length]).is_err()); + assert!(KllSketch::>::deserialize(&bytes[..length]).is_err()); } let mut with_trailing_data = bytes; with_trailing_data.push(0); - assert!(KllSketch::::deserialize(&with_trailing_data).is_err()); + assert!(KllSketch::>::deserialize(&with_trailing_data).is_err()); +} + +#[test] +fn test_string_value_reports_the_truncated_field_and_byte_counts() { + let mut input = SketchSlice::new(&[5, 0, 0, 0, b'a']); + + let error = String::deserialize(&mut input).unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::InvalidData); + assert_eq!( + error.message(), + "insufficient data (KLL string payload): expected 5 bytes, got 1" + ); } diff --git a/tests-integration/tests/serde_tests/req.rs b/tests-integration/tests/serde_tests/req.rs index b097ed05..397f894d 100644 --- a/tests-integration/tests/serde_tests/req.rs +++ b/tests-integration/tests/serde_tests/req.rs @@ -20,11 +20,11 @@ use std::fs; use std::path::PathBuf; +use datasketches::common::SearchCriteria; use datasketches::req::RankAccuracy; use datasketches::req::ReqFloat; use datasketches::req::ReqSketch; use datasketches::req::ReqValue; -use datasketches::req::SearchCriteria; use googletest::assert_that; use googletest::prelude::anything; use googletest::prelude::err;