From 17e8b77114ba4f55dfd1d52cf27fe9652164f576 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:05:49 +0800 Subject: [PATCH 01/22] refactor(kll): clarify capacity module boundaries --- .../src/kll/{helper.rs => capacity.rs} | 28 ++++++------------- datasketches/src/kll/mod.rs | 2 +- datasketches/src/kll/serialization.rs | 22 +++++++-------- datasketches/src/kll/sketch.rs | 18 ++++++------ datasketches/src/kll/sorted_view.rs | 12 ++++---- 5 files changed, 37 insertions(+), 45 deletions(-) rename datasketches/src/kll/{helper.rs => capacity.rs} (73%) diff --git a/datasketches/src/kll/helper.rs b/datasketches/src/kll/capacity.rs similarity index 73% rename from datasketches/src/kll/helper.rs rename to datasketches/src/kll/capacity.rs index 6b993aa7..03057aa1 100644 --- a/datasketches/src/kll/helper.rs +++ b/datasketches/src/kll/capacity.rs @@ -49,7 +49,7 @@ const POWERS_OF_THREE: [u64; 31] = [ 205891132094649, ]; -pub(super) fn compute_total_capacity(k: u16, m: u8, num_levels: usize) -> u32 { +pub fn 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); @@ -57,27 +57,27 @@ pub(super) fn compute_total_capacity(k: u16, m: u8, num_levels: usize) -> u32 { total } -pub(super) fn level_capacity(k: u16, num_levels: usize, height: usize, min_wid: u8) -> u32 { +pub fn level_capacity(k: u16, num_levels: usize, height: usize, min_width: 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) + let cap = capacity_at_depth(k, depth as u8); + std::cmp::max(min_width as u32, cap as u32) } -fn int_cap_aux(k: u16, depth: u8) -> u16 { +fn capacity_at_depth(k: u16, depth: u8) -> u16 { if depth > 60 { panic!("depth must be <= 60"); } if depth <= 30 { - return int_cap_aux_aux(k, depth); + return capacity_at_shallow_depth(k, depth); } let half = depth / 2; let rest = depth - half; - let tmp = int_cap_aux_aux(k, half); - int_cap_aux_aux(tmp, rest) + let tmp = capacity_at_shallow_depth(k, half); + capacity_at_shallow_depth(tmp, rest) } -fn int_cap_aux_aux(k: u16, depth: u8) -> u16 { +fn capacity_at_shallow_depth(k: u16, depth: u8) -> u16 { if depth > 30 { panic!("depth must be <= 30"); } @@ -87,13 +87,3 @@ fn int_cap_aux_aux(k: u16, depth: u8) -> u16 { 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..d87da0ea 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -36,7 +36,7 @@ //! assert!(q >= 1.0 && q <= 2.0); //! ``` -mod helper; +mod capacity; mod serialization; mod sketch; mod sorted_view; 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..ec611904 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; @@ -505,7 +504,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 { @@ -692,7 +691,7 @@ impl> KllSketch { } fn capacity(&self) -> usize { - compute_total_capacity(self.k, self.m, self.levels.len()) as usize + total_capacity(self.k, self.m, self.levels.len()) as usize } fn level_offsets(&self) -> Vec { @@ -847,8 +846,11 @@ impl> KllSketch { } 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> { @@ -986,7 +988,7 @@ fn general_compress>( ) -> 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; diff --git a/datasketches/src/kll/sorted_view.rs b/datasketches/src/kll/sorted_view.rs index bcaad640..075a4726 100644 --- a/datasketches/src/kll/sorted_view.rs +++ b/datasketches/src/kll/sorted_view.rs @@ -21,7 +21,7 @@ use super::sketch::KllComparator; use super::sketch::KllItem; #[derive(Debug, Clone)] -pub(super) struct SortedView> { +pub struct SortedView> { comparator: C, entries: Vec>, total_weight: u64, @@ -48,7 +48,7 @@ impl> SortedView { } } - pub(super) fn rank(&self, item: &T, inclusive: bool) -> f64 { + pub fn rank(&self, item: &T, inclusive: bool) -> f64 { if self.entries.is_empty() { return 0.0; } @@ -66,7 +66,7 @@ impl> SortedView { weight as f64 / self.total_weight as f64 } - pub(super) fn quantile(&self, rank: f64, inclusive: bool) -> T { + pub fn quantile(&self, rank: f64, inclusive: bool) -> T { let weight = if inclusive { (rank * self.total_weight as f64).ceil() as u64 } else { @@ -85,7 +85,7 @@ impl> SortedView { self.entries[idx].item.clone() } - pub(super) fn cdf(&self, split_points: &[T], inclusive: bool) -> Vec { + pub fn cdf(&self, split_points: &[T], inclusive: bool) -> Vec { check_split_points(split_points, &self.comparator); let mut ranks = Vec::with_capacity(split_points.len() + 1); for item in split_points { @@ -95,7 +95,7 @@ impl> SortedView { ranks } - pub(super) fn pmf(&self, split_points: &[T], inclusive: bool) -> Vec { + pub 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]; @@ -104,7 +104,7 @@ impl> SortedView { } } -pub(super) fn build_sorted_view>( +pub fn build_sorted_view>( levels: &[Vec], comparator: C, ) -> SortedView { From 28a74d275070b619a2d8843f61afd7cb910b6679 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:06:12 +0800 Subject: [PATCH 02/22] refactor(kll): minimize public configuration API --- datasketches/src/kll/mod.rs | 17 ++++------------- datasketches/src/kll/sketch.rs | 4 ++-- tests-integration/tests/kll_test/sketch.rs | 6 +++--- tests-integration/tests/serde_tests/kll.rs | 3 ++- 4 files changed, 11 insertions(+), 19 deletions(-) diff --git a/datasketches/src/kll/mod.rs b/datasketches/src/kll/mod.rs index d87da0ea..c98a9311 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -46,20 +46,11 @@ 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; - /// 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/sketch.rs b/datasketches/src/kll/sketch.rs index ec611904..f1cf1751 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -125,7 +125,7 @@ impl KllSketch { /// /// # Errors /// - /// Returns an error if `k` is outside [`MIN_K`, `MAX_K`]. + /// Returns an error if `k` is outside `8..=65535`. /// /// # Examples /// @@ -144,7 +144,7 @@ impl> KllSketch { /// /// # Errors /// - /// Returns an error if `k` is outside [`MIN_K`, `MAX_K`]. + /// Returns an error if `k` is outside `8..=65535`. pub fn new_with_comparator(k: u16, comparator: C) -> Result { if !(MIN_K..=MAX_K).contains(&k) { return Err(Error::invalid_argument(format!( diff --git a/tests-integration/tests/kll_test/sketch.rs b/tests-integration/tests/kll_test/sketch.rs index d9acaaa8..18e68c47 100644 --- a/tests-integration/tests/kll_test/sketch.rs +++ b/tests-integration/tests/kll_test/sketch.rs @@ -18,12 +18,12 @@ 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 DEFAULT_K: u16 = 200; +const MIN_K: u16 = 8; +const MAX_K: u16 = u16::MAX; const NUMERIC_NOISE_TOLERANCE: f64 = 1e-6; fn assert_approx_eq(actual: f64, expected: f64, tolerance: f64) { diff --git a/tests-integration/tests/serde_tests/kll.rs b/tests-integration/tests/serde_tests/kll.rs index 70d9249a..fad3eb37 100644 --- a/tests-integration/tests/serde_tests/kll.rs +++ b/tests-integration/tests/serde_tests/kll.rs @@ -28,12 +28,13 @@ use std::cmp::Ordering; use std::fs; use std::path::PathBuf; -use datasketches::kll::DEFAULT_K; use datasketches::kll::KllComparator; use datasketches::kll::KllSketch; use crate::serialization_test_data; +const DEFAULT_K: u16 = 200; + #[derive(Clone, Copy)] struct NumericStringOrder; From f6868d264e0189b4cb61169a8a4dda17e7175af9 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:08:29 +0800 Subject: [PATCH 03/22] refactor(kll): unify item ordering and serialization --- datasketches/src/kll/mod.rs | 8 +- datasketches/src/kll/order.rs | 57 +++++ datasketches/src/kll/sketch.rs | 269 +++------------------ datasketches/src/kll/sorted_view.rs | 31 +-- datasketches/src/kll/value.rs | 121 +++++++++ tests-integration/tests/kll_test/sketch.rs | 6 +- tests-integration/tests/serde_tests/kll.rs | 4 + 7 files changed, 237 insertions(+), 259 deletions(-) create mode 100644 datasketches/src/kll/order.rs create mode 100644 datasketches/src/kll/value.rs diff --git a/datasketches/src/kll/mod.rs b/datasketches/src/kll/mod.rs index c98a9311..cf7679dc 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -37,14 +37,16 @@ //! ``` mod capacity; +mod order; mod serialization; mod sketch; mod sorted_view; +mod value; -pub use self::sketch::KllComparator; -pub use self::sketch::KllItem; +pub use self::order::KllComparator; +pub use self::order::NaturalOrder; pub use self::sketch::KllSketch; -pub use self::sketch::NaturalOrder; +pub use self::value::KllValue; /// Default value of parameter k. const DEFAULT_K: u16 = 200; diff --git a/datasketches/src/kll/order.rs b/datasketches/src/kll/order.rs new file mode 100644 index 00000000..daf2264b --- /dev/null +++ b/datasketches/src/kll/order.rs @@ -0,0 +1,57 @@ +// 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; + +/// Defines the ordering used by a KLL sketch. +/// +/// Accepted values must form a total order. Sketches can be merged only when their comparators are +/// compatible: they must accept the same values and order every pair of accepted values +/// identically. +pub trait KllComparator: Clone { + /// Compares two accepted values. + fn compare(&self, left: &T, right: &T) -> Ordering; + + /// Returns whether `item` belongs to this comparator's ordered domain. + /// + /// Updates with rejected values are ignored. The default accepts every value. + fn accepts(&self, _item: &T) -> bool { + true + } + + /// Returns whether `other` defines the same ordered domain and comparison semantics. + fn is_compatible(&self, other: &Self) -> bool; +} + +/// Uses the value's natural partial ordering and rejects unordered values such as NaN. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct NaturalOrder; + +impl KllComparator for NaturalOrder { + fn compare(&self, left: &T, right: &T) -> Ordering { + left.partial_cmp(right) + .expect("accepted KLL values must be totally ordered") + } + + fn accepts(&self, item: &T) -> bool { + item.partial_cmp(item).is_some() + } + + fn is_compatible(&self, _other: &Self) -> bool { + true + } +} diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index f1cf1751..8153bd09 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -23,6 +23,8 @@ use super::MAX_K; use super::MIN_K; use super::capacity::level_capacity; use super::capacity::total_capacity; +use super::order::KllComparator; +use super::order::NaturalOrder; use super::serialization::DATA_START; use super::serialization::DATA_START_SINGLE_ITEM; use super::serialization::EMPTY_SIZE_BYTES; @@ -35,6 +37,7 @@ use super::serialization::PREAMBLE_INTS_SHORT; use super::serialization::SERIAL_VERSION_1; use super::serialization::SERIAL_VERSION_2; 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; @@ -42,53 +45,6 @@ use crate::codec::assert::insufficient_data; use crate::codec::family::Family; 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. @@ -105,7 +61,7 @@ pub struct KllSketch { max_item: Option, } -impl Default for KllSketch { +impl Default for KllSketch { fn default() -> Self { Self::make( NaturalOrder, @@ -120,7 +76,7 @@ impl Default for KllSketch { } } -impl KllSketch { +impl KllSketch { /// Creates a new sketch with the given value of k. /// /// # Errors @@ -139,7 +95,7 @@ impl KllSketch { } } -impl> KllSketch { +impl> KllSketch { /// Creates a new sketch with the given value of k and ordering policy. /// /// # Errors @@ -207,7 +163,7 @@ impl> KllSketch { /// /// NaN values are ignored for floating-point types. pub fn update(&mut self, item: T) { - if T::is_nan(&item) { + if !self.comparator.accepts(&item) { return; } self.update_min_max(&item); @@ -307,7 +263,7 @@ impl> KllSketch { } } -fn serialized_size>(sketch: &KllSketch) -> usize { +fn serialized_size>(sketch: &KllSketch) -> usize { if sketch.is_empty() { return EMPTY_SIZE_BYTES; } @@ -331,7 +287,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); @@ -403,7 +359,7 @@ fn serialize_with_serde>(sketch: &KllSketch>( +fn deserialize_with_serde>( bytes: &[u8], comparator: C, ) -> Result, Error> { @@ -590,83 +546,36 @@ fn deserialize_with_serde>( Ok(sketch) } -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 { +impl> KllSketch { /// Serializes the sketch to bytes. pub fn serialize(&self) -> Vec { serialize_with_serde(self) } /// Deserializes a sketch using the supplied ordering policy. + /// + /// # Errors + /// + /// Returns `InvalidData` if the image is truncated, malformed, or inconsistent with + /// `comparator`. pub fn deserialize_with_comparator(bytes: &[u8], comparator: C) -> Result { deserialize_with_serde(bytes, comparator) } } -impl KllSketch { +impl KllSketch { /// Deserializes a sketch from bytes. + /// + /// # Errors + /// + /// Returns `InvalidData` if the image is truncated, malformed, or contains values that are not + /// naturally ordered. pub fn deserialize(bytes: &[u8]) -> Result { deserialize_with_serde(bytes, NaturalOrder) } } -impl> KllSketch { +impl> KllSketch { fn make( comparator: C, k: u16, @@ -863,8 +772,10 @@ 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.accepts(min_item) || !self.comparator.accepts(max_item) { + return Err(Error::deserial( + "minimum and maximum items must belong to the comparator's ordered domain", + )); } if self.comparator.compare(min_item, max_item) == Ordering::Greater { return Err(Error::deserial( @@ -894,8 +805,10 @@ impl> KllSketch { } for item in level { - if T::is_nan(item) { - return Err(Error::deserial("retained items must not be NaN")); + if !self.comparator.accepts(item) { + return Err(Error::deserial( + "retained items must belong to the comparator's ordered domain", + )); } if self.comparator.compare(item, min_item) == Ordering::Less || self.comparator.compare(item, max_item) == Ordering::Greater @@ -933,7 +846,7 @@ fn normalized_rank_error(k: u16, pmf: bool) -> f64 { } } -fn downsample(items: Vec, offset: bool, use_up: bool) -> Vec { +fn downsample(items: Vec, offset: bool, use_up: bool) -> Vec { let len = items.len(); debug_assert!(len % 2 == 0, "length must be even"); let offset = usize::from(offset); @@ -958,7 +871,7 @@ fn take_leftover(items: &mut Vec, level: usize, is_level_zero_sorted: bool } } -fn merge_sorted_vec>( +fn merge_sorted_vec>( left: Vec, right: Vec, comparator: &C, @@ -979,7 +892,7 @@ fn merge_sorted_vec>( merged } -fn general_compress>( +fn general_compress>( mut levels_in: Vec>, k: u16, m: u8, @@ -1052,117 +965,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 075a4726..4ab88384 100644 --- a/datasketches/src/kll/sorted_view.rs +++ b/datasketches/src/kll/sorted_view.rs @@ -17,11 +17,10 @@ use std::cmp::Ordering; -use super::sketch::KllComparator; -use super::sketch::KllItem; +use super::order::KllComparator; #[derive(Debug, Clone)] -pub struct SortedView> { +pub struct SortedView> { comparator: C, entries: Vec>, total_weight: u64, @@ -33,7 +32,7 @@ struct Entry { weight: u64, } -impl> SortedView { +impl> SortedView { fn new(mut entries: Vec>, comparator: C) -> Self { entries.sort_by(|a, b| comparator.compare(&a.item, &b.item)); let mut total_weight = 0u64; @@ -104,7 +103,7 @@ impl> SortedView { } } -pub fn build_sorted_view>( +pub fn build_sorted_view>( levels: &[Vec], comparator: C, ) -> SortedView { @@ -125,10 +124,10 @@ pub fn build_sorted_view>( } #[track_caller] -fn check_split_points>(split_points: &[T], comparator: &C) { +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" + split_points.iter().all(|point| comparator.accepts(point)), + "split_points must belong to the comparator's ordered domain" ); for pair in split_points.windows(2) { assert!( @@ -138,11 +137,7 @@ fn check_split_points>(split_points: &[T], compa } } -fn lower_bound>( - entries: &[Entry], - item: &T, - comparator: &C, -) -> usize { +fn lower_bound>(entries: &[Entry], item: &T, comparator: &C) -> usize { let mut left = 0usize; let mut right = entries.len(); while left < right { @@ -156,11 +151,7 @@ fn lower_bound>( left } -fn upper_bound>( - entries: &[Entry], - item: &T, - comparator: &C, -) -> usize { +fn upper_bound>(entries: &[Entry], item: &T, comparator: &C) -> usize { let mut left = 0usize; let mut right = entries.len(); while left < right { @@ -174,7 +165,7 @@ fn upper_bound>( left } -fn lower_bound_by_weight(entries: &[Entry], weight: u64) -> usize { +fn lower_bound_by_weight(entries: &[Entry], weight: u64) -> usize { let mut left = 0usize; let mut right = entries.len(); while left < right { @@ -188,7 +179,7 @@ fn lower_bound_by_weight(entries: &[Entry], weight: u64) -> usize left } -fn upper_bound_by_weight(entries: &[Entry], weight: u64) -> usize { +fn upper_bound_by_weight(entries: &[Entry], weight: u64) -> usize { let mut left = 0usize; let mut right = entries.len(); while left < right { diff --git a/datasketches/src/kll/value.rs b/datasketches/src/kll/value.rs new file mode 100644 index 00000000..e14057ae --- /dev/null +++ b/datasketches/src/kll/value.rs @@ -0,0 +1,121 @@ +// 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 crate::codec::SketchBytes; +use crate::codec::SketchSlice; +use crate::error::Error; + +/// Defines the compact binary representation of a KLL item. +/// +/// This trait is required only for serialization. In-memory KLL operations support any cloneable +/// item type with a [`KllComparator`](crate::kll::KllComparator). The encoded representation must +/// preserve the comparator's 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 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 KllValue 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 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(|_| Error::insufficient_data("i64")) + } +} + +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(|_| 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/tests-integration/tests/kll_test/sketch.rs b/tests-integration/tests/kll_test/sketch.rs index 18e68c47..7fcca7e1 100644 --- a/tests-integration/tests/kll_test/sketch.rs +++ b/tests-integration/tests/kll_test/sketch.rs @@ -47,6 +47,10 @@ impl KllComparator for NumericStringOrder { .unwrap() .cmp(&right.parse::().unwrap()) } + + fn is_compatible(&self, _other: &Self) -> bool { + true + } } #[test] @@ -245,7 +249,7 @@ fn test_out_of_order_split_points_panics() { } #[test] -#[should_panic(expected = "split_points must not contain NaN values")] +#[should_panic(expected = "split_points must belong to the comparator's ordered domain")] fn test_nan_split_point_panics() { let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); sketch.update(0.0); diff --git a/tests-integration/tests/serde_tests/kll.rs b/tests-integration/tests/serde_tests/kll.rs index fad3eb37..aac40997 100644 --- a/tests-integration/tests/serde_tests/kll.rs +++ b/tests-integration/tests/serde_tests/kll.rs @@ -42,6 +42,10 @@ impl KllComparator for NumericStringOrder { fn compare(&self, left: &String, right: &String) -> Ordering { parse_string_value(left).cmp(&parse_string_value(right)) } + + fn is_compatible(&self, _other: &Self) -> bool { + true + } } fn test_f32_file(path: PathBuf, expected_n: usize) { From 1610ba7ed8ce79861935f70296afeac7505e5741 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:11:54 +0800 Subject: [PATCH 04/22] refactor(kll): align quantile query APIs --- datasketches/src/common/mod.rs | 2 + datasketches/src/common/search_criteria.rs | 26 +++ datasketches/src/kll/mod.rs | 5 +- datasketches/src/kll/sketch.rs | 66 +++++--- datasketches/src/kll/sorted_view.rs | 51 +++--- datasketches/src/req/mod.rs | 11 +- tests-integration/tests/kll_test/sketch.rs | 177 ++++++++++++++------- 7 files changed, 231 insertions(+), 107 deletions(-) create mode 100644 datasketches/src/common/search_criteria.rs 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..7f480c73 --- /dev/null +++ b/datasketches/src/common/search_criteria.rs @@ -0,0 +1,26 @@ +// 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, 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, +} diff --git a/datasketches/src/kll/mod.rs b/datasketches/src/kll/mod.rs index cf7679dc..8b8e5631 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -28,11 +28,11 @@ //! # Usage //! //! ```rust -//! # use datasketches::kll::KllSketch; +//! # use datasketches::kll::{KllSketch, SearchCriteria}; //! let mut sketch = KllSketch::::new(200).unwrap(); //! sketch.update(1.0); //! sketch.update(2.0); -//! let q = sketch.quantile(0.5, true).unwrap(); +//! let q = sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(); //! assert!(q >= 1.0 && q <= 2.0); //! ``` @@ -47,6 +47,7 @@ pub use self::order::KllComparator; pub use self::order::NaturalOrder; pub use self::sketch::KllSketch; pub use self::value::KllValue; +pub use crate::common::SearchCriteria; /// Default value of parameter k. const DEFAULT_K: u16 = 200; diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index 8153bd09..ae9045f3 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -43,6 +43,7 @@ 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; /// KLL sketch for estimating quantiles and ranks. @@ -217,49 +218,78 @@ impl> KllSketch { } /// 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 or `item` is outside the comparator's ordered + /// domain. + 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")); + } + if !self.comparator.accepts(item) { + return Err(Error::invalid_argument( + "item must belong to the comparator's ordered domain", + )); } let view = build_sorted_view(&self.levels, self.comparator.clone()); - Some(view.rank(item, inclusive)) + Ok(view.rank(item, criteria)) } /// 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")); + } + if !(0.0..=1.0).contains(&rank) { + return Err(Error::invalid_argument(format!( + "rank must be in [0.0, 1.0], got {rank}" + ))); } - 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)) + Ok(view.quantile(rank, 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, a split point is outside the comparator's ordered + /// domain, 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)) + 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, a split point is outside the comparator's ordered + /// domain, 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)) + view.pmf(split_points, criteria) + } + + /// 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) } } diff --git a/datasketches/src/kll/sorted_view.rs b/datasketches/src/kll/sorted_view.rs index 4ab88384..5c8cd5bb 100644 --- a/datasketches/src/kll/sorted_view.rs +++ b/datasketches/src/kll/sorted_view.rs @@ -18,6 +18,8 @@ use std::cmp::Ordering; use super::order::KllComparator; +use crate::common::SearchCriteria; +use crate::error::Error; #[derive(Debug, Clone)] pub struct SortedView> { @@ -47,12 +49,12 @@ impl> SortedView { } } - pub fn rank(&self, item: &T, inclusive: bool) -> f64 { + pub fn rank(&self, item: &T, criteria: SearchCriteria) -> f64 { if self.entries.is_empty() { return 0.0; } - let idx = if inclusive { + let idx = if criteria == SearchCriteria::Inclusive { upper_bound(&self.entries, item, &self.comparator) } else { lower_bound(&self.entries, item, &self.comparator) @@ -65,14 +67,14 @@ impl> SortedView { weight as f64 / self.total_weight as f64 } - pub fn quantile(&self, rank: f64, inclusive: bool) -> T { - let weight = if inclusive { + pub fn quantile(&self, rank: f64, criteria: SearchCriteria) -> T { + 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 idx = if criteria == SearchCriteria::Inclusive { lower_bound_by_weight(&self.entries, weight) } else { upper_bound_by_weight(&self.entries, weight) @@ -84,22 +86,22 @@ impl> SortedView { self.entries[idx].item.clone() } - pub fn cdf(&self, split_points: &[T], inclusive: bool) -> Vec { - check_split_points(split_points, &self.comparator); + pub fn cdf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { + check_split_points(split_points, &self.comparator)?; 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 fn pmf(&self, split_points: &[T], inclusive: bool) -> Vec { - let mut buckets = self.cdf(split_points, inclusive); + pub fn pmf(&self, split_points: &[T], criteria: SearchCriteria) -> Result, Error> { + let mut buckets = self.cdf(split_points, criteria)?; for i in (1..buckets.len()).rev() { buckets[i] -= buckets[i - 1]; } - buckets + Ok(buckets) } } @@ -123,18 +125,23 @@ pub fn build_sorted_view>( SortedView::new(entries, comparator) } -#[track_caller] -fn check_split_points>(split_points: &[T], comparator: &C) { - assert!( - split_points.iter().all(|point| comparator.accepts(point)), - "split_points must belong to the comparator's ordered domain" - ); +fn check_split_points>( + split_points: &[T], + comparator: &C, +) -> Result<(), Error> { + if !split_points.iter().all(|point| comparator.accepts(point)) { + return Err(Error::invalid_argument( + "split points must belong to the comparator's ordered domain", + )); + } for pair in split_points.windows(2) { - assert!( - comparator.compare(&pair[0], &pair[1]) == Ordering::Less, - "split_points must be unique and monotonically increasing" - ); + if comparator.compare(&pair[0], &pair[1]) != Ordering::Less { + return Err(Error::invalid_argument( + "split points must be unique and monotonically increasing", + )); + } } + Ok(()) } fn lower_bound>(entries: &[Entry], item: &T, comparator: &C) -> usize { diff --git a/datasketches/src/req/mod.rs b/datasketches/src/req/mod.rs index 8714f5ec..41d8f408 100644 --- a/datasketches/src/req/mod.rs +++ b/datasketches/src/req/mod.rs @@ -61,6 +61,7 @@ pub use self::sketch::ReqSketch; pub use self::sorted_view::SortedView; pub use self::value::ReqFloat; pub use self::value::ReqValue; +pub use crate::common::SearchCriteria; /// Default value of `k` if not specified. Roughly 1% relative error at 95% confidence. const DEFAULT_K: u16 = 12; @@ -79,16 +80,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/tests-integration/tests/kll_test/sketch.rs b/tests-integration/tests/kll_test/sketch.rs index 7fcca7e1..afc5090d 100644 --- a/tests-integration/tests/kll_test/sketch.rs +++ b/tests-integration/tests/kll_test/sketch.rs @@ -20,6 +20,7 @@ use std::cmp::Ordering; use datasketches::error::ErrorKind; use datasketches::kll::KllComparator; use datasketches::kll::KllSketch; +use datasketches::kll::SearchCriteria; const DEFAULT_K: u16 = 200; const MIN_K: u16 = 8; @@ -35,7 +36,7 @@ fn assert_approx_eq(actual: f64, expected: f64, tolerance: f64) { } fn rank_eps(sketch: &KllSketch) -> f64 { - sketch.normalized_rank_error(false) + sketch.normalized_rank_error() } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -74,18 +75,20 @@ fn test_empty() { 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()); + assert!(sketch.rank(&0.0, SearchCriteria::Inclusive).is_err()); + assert!(sketch.quantile(0.5, SearchCriteria::Inclusive).is_err()); + assert!(sketch.pmf(&[0.0f32], SearchCriteria::Inclusive).is_err()); + assert!(sketch.cdf(&[0.0f32], SearchCriteria::Inclusive).is_err()); } #[test] -#[should_panic(expected = "rank must be in [0.0, 1.0]")] -fn test_quantile_out_of_range_panics() { +fn test_quantile_out_of_range_returns_error() { let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); sketch.update(0.0); - sketch.quantile(-1.0, true); + let error = sketch + .quantile(-1.0, SearchCriteria::Inclusive) + .unwrap_err(); + assert_eq!(error.kind(), ErrorKind::InvalidArgument); } #[test] @@ -96,12 +99,15 @@ fn test_one_item() { 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.rank(&1.0, SearchCriteria::Exclusive).unwrap(), 0.0); + assert_eq!(sketch.rank(&1.0, SearchCriteria::Inclusive).unwrap(), 1.0); + assert_eq!(sketch.rank(&2.0, SearchCriteria::Exclusive).unwrap(), 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)); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), + 1.0 + ); } #[test] @@ -111,12 +117,18 @@ fn test_duplicate_items_follow_inclusive_and_exclusive_semantics() { 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)); + assert_eq!(sketch.rank(&1.0, SearchCriteria::Exclusive).unwrap(), 0.0); + assert_eq!(sketch.rank(&1.0, SearchCriteria::Inclusive).unwrap(), 0.5); + assert_eq!(sketch.rank(&2.0, SearchCriteria::Exclusive).unwrap(), 0.5); + assert_eq!(sketch.rank(&2.0, SearchCriteria::Inclusive).unwrap(), 1.0); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), + 1.0 + ); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Exclusive).unwrap(), + 2.0 + ); } #[test] @@ -141,15 +153,27 @@ fn test_many_items_exact_mode() { 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.quantile(0.0, SearchCriteria::Inclusive).unwrap(), + 1.0 + ); assert_eq!(sketch.max_item().cloned(), Some(n as f32)); - assert_eq!(sketch.quantile(1.0, true), Some(n as f32)); + assert_eq!( + sketch.quantile(1.0, SearchCriteria::Inclusive).unwrap(), + 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)); + assert_eq!( + sketch.rank(&(i as f32), SearchCriteria::Inclusive).unwrap(), + inclusive_rank + ); let exclusive_rank = (i - 1) as f64 / n as f64; - assert_eq!(sketch.rank(&(i as f32), false), Some(exclusive_rank)); + assert_eq!( + sketch.rank(&(i as f32), SearchCriteria::Exclusive).unwrap(), + exclusive_rank + ); } } @@ -159,10 +183,22 @@ fn test_ten_items_quantiles() { 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)); + assert_eq!( + sketch.quantile(0.0, SearchCriteria::Inclusive).unwrap(), + 1.0 + ); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), + 5.0 + ); + assert_eq!( + sketch.quantile(0.99, SearchCriteria::Inclusive).unwrap(), + 10.0 + ); + assert_eq!( + sketch.quantile(1.0, SearchCriteria::Inclusive).unwrap(), + 10.0 + ); } #[test] @@ -171,11 +207,26 @@ fn test_hundred_items_quantiles() { 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)); + assert_eq!( + sketch.quantile(0.0, SearchCriteria::Inclusive).unwrap(), + 0.0 + ); + assert_eq!( + sketch.quantile(0.01, SearchCriteria::Inclusive).unwrap(), + 0.0 + ); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), + 49.0 + ); + assert_eq!( + sketch.quantile(0.99, SearchCriteria::Inclusive).unwrap(), + 98.0 + ); + assert_eq!( + sketch.quantile(1.0, SearchCriteria::Inclusive).unwrap(), + 99.0 + ); } #[test] @@ -193,7 +244,7 @@ fn test_many_items_estimation_mode_rank_error() { 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(); + let rank = sketch.rank(&(i as f32), SearchCriteria::Exclusive).unwrap(); assert_approx_eq(rank, true_rank, rank_eps); } @@ -210,12 +261,12 @@ fn test_rank_cdf_pmf_consistency() { values.push(i as f32); } - let ranks = sketch.cdf(&values, false).unwrap(); - let pmf = sketch.pmf(&values, false).unwrap(); + let ranks = sketch.cdf(&values, SearchCriteria::Exclusive).unwrap(); + let pmf = sketch.pmf(&values, SearchCriteria::Exclusive).unwrap(); let mut subtotal = 0.0; for i in 0..n { - let rank = sketch.rank(&values[i], false).unwrap(); + let rank = sketch.rank(&values[i], SearchCriteria::Exclusive).unwrap(); assert_eq!(rank, ranks[i]); subtotal += pmf[i]; assert!( @@ -224,12 +275,12 @@ fn test_rank_cdf_pmf_consistency() { ); } - let ranks = sketch.cdf(&values, true).unwrap(); - let pmf = sketch.pmf(&values, true).unwrap(); + let ranks = sketch.cdf(&values, SearchCriteria::Inclusive).unwrap(); + let pmf = sketch.pmf(&values, SearchCriteria::Inclusive).unwrap(); let mut subtotal = 0.0; for i in 0..n { - let rank = sketch.rank(&values[i], true).unwrap(); + let rank = sketch.rank(&values[i], SearchCriteria::Inclusive).unwrap(); assert_eq!(rank, ranks[i]); subtotal += pmf[i]; assert!( @@ -240,21 +291,25 @@ fn test_rank_cdf_pmf_consistency() { } #[test] -#[should_panic(expected = "split_points must be unique and monotonically increasing")] -fn test_out_of_order_split_points_panics() { +fn test_out_of_order_split_points_return_error() { 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); + let error = sketch + .cdf(&split_points, SearchCriteria::Inclusive) + .unwrap_err(); + assert_eq!(error.kind(), ErrorKind::InvalidArgument); } #[test] -#[should_panic(expected = "split_points must belong to the comparator's ordered domain")] -fn test_nan_split_point_panics() { +fn test_nan_split_point_returns_error() { let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); sketch.update(0.0); let split_points = [f32::NAN]; - let _ = sketch.cdf(&split_points, true); + let error = sketch + .cdf(&split_points, SearchCriteria::Inclusive) + .unwrap_err(); + assert_eq!(error.kind(), ErrorKind::InvalidArgument); } #[test] @@ -278,7 +333,7 @@ fn test_merge() { 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 median = sketch1.quantile(0.5, SearchCriteria::Inclusive).unwrap(); let rank_eps = rank_eps(&sketch1); assert_approx_eq(median as f64, n as f64, n as f64 * rank_eps); } @@ -299,14 +354,14 @@ fn test_merge_lower_k() { 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) + sketch1.normalized_rank_error(), + sketch2.normalized_rank_error() ); assert_eq!( - sketch1.normalized_rank_error(true), - sketch2.normalized_rank_error(true) + sketch1.normalized_pmf_error(), + sketch2.normalized_pmf_error() ); - let median = sketch1.quantile(0.5, true).unwrap(); + let median = sketch1.quantile(0.5, SearchCriteria::Inclusive).unwrap(); let rank_eps = rank_eps(&sketch1); assert_approx_eq(median as f64, n as f64, n as f64 * rank_eps); } @@ -320,14 +375,14 @@ fn test_merge_exact_mode_lower_k() { sketch1.update(i as f32); } - let err_before = sketch1.normalized_rank_error(true); + let err_before = sketch1.normalized_pmf_error(); sketch1.merge(&sketch2); - assert_eq!(sketch1.normalized_rank_error(true), err_before); + assert_eq!(sketch1.normalized_pmf_error(), 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 median = sketch1.quantile(0.5, SearchCriteria::Inclusive).unwrap(); let rank_eps = rank_eps(&sketch1); assert_approx_eq(median as f64, (n / 2) as f64, (n as f64 / 2.0) * rank_eps); } @@ -386,7 +441,13 @@ fn test_custom_comparator_roundtrip() { 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")); + assert_eq!( + sketch + .quantile(0.5, SearchCriteria::Inclusive) + .as_deref() + .unwrap(), + "2" + ); let bytes = sketch.serialize(); let decoded = KllSketch::::deserialize_with_comparator( @@ -397,5 +458,11 @@ fn test_custom_comparator_roundtrip() { 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")); + assert_eq!( + decoded + .quantile(0.5, SearchCriteria::Inclusive) + .as_deref() + .unwrap(), + "2" + ); } From c93611448b4ce0499fae178703018af9a9ec0bc7 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:12:38 +0800 Subject: [PATCH 05/22] fix(kll): reject incompatible sketch merges --- datasketches/src/kll/sketch.rs | 31 ++++++++++----- tests-integration/tests/kll_test/sketch.rs | 46 +++++++++++++++++++--- 2 files changed, 62 insertions(+), 15 deletions(-) diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index ae9045f3..9dcecb80 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -184,23 +184,33 @@ 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 sketches use incompatible comparators or their 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.comparator.is_compatible(&other.comparator) { + return Err(Error::invalid_argument( + "cannot merge sketches with incompatible comparators", + )); + } + 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("combined stream weight exceeds u64::MAX"))?; 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()); } @@ -215,6 +225,7 @@ 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. diff --git a/tests-integration/tests/kll_test/sketch.rs b/tests-integration/tests/kll_test/sketch.rs index afc5090d..a52102fd 100644 --- a/tests-integration/tests/kll_test/sketch.rs +++ b/tests-integration/tests/kll_test/sketch.rs @@ -54,6 +54,25 @@ impl KllComparator for NumericStringOrder { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct DirectionalOrder { + descending: bool, +} + +impl KllComparator for DirectionalOrder { + fn compare(&self, left: &i64, right: &i64) -> Ordering { + if self.descending { + right.cmp(left) + } else { + left.cmp(right) + } + } + + fn is_compatible(&self, other: &Self) -> bool { + self == other + } +} + #[test] fn test_k_limits() { let _min = KllSketch::::new(MIN_K).unwrap(); @@ -327,7 +346,7 @@ fn test_merge() { 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); + sketch1.merge(&sketch2).unwrap(); assert!(!sketch1.is_empty()); assert_eq!(sketch1.n(), (2 * n) as u64); @@ -348,7 +367,7 @@ fn test_merge_lower_k() { sketch2.update((2 * n - i - 1) as f32); } - sketch1.merge(&sketch2); + sketch1.merge(&sketch2).unwrap(); assert_eq!(sketch1.n(), (2 * n) as u64); assert_eq!(sketch1.min_item().cloned(), Some(0.0)); @@ -376,7 +395,7 @@ fn test_merge_exact_mode_lower_k() { } let err_before = sketch1.normalized_pmf_error(); - sketch1.merge(&sketch2); + sketch1.merge(&sketch2).unwrap(); assert_eq!(sketch1.normalized_pmf_error(), err_before); assert_eq!(sketch1.n(), n as u64); @@ -393,11 +412,28 @@ fn test_merge_min_max_from_other() { let mut sketch2 = KllSketch::::new(DEFAULT_K).unwrap(); sketch1.update(1.0); sketch2.update(2.0); - sketch2.merge(&sketch1); + sketch2.merge(&sketch1).unwrap(); assert_eq!(sketch2.min_item().cloned(), Some(1.0)); assert_eq!(sketch2.max_item().cloned(), Some(2.0)); } +#[test] +fn test_merge_rejects_incompatible_comparators_without_mutation() { + let mut ascending = + KllSketch::new_with_comparator(200, DirectionalOrder { descending: false }).unwrap(); + let mut descending = + KllSketch::new_with_comparator(200, DirectionalOrder { descending: true }).unwrap(); + ascending.update(1); + descending.update(2); + + let error = ascending.merge(&descending).unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::InvalidArgument); + assert_eq!(ascending.n(), 1); + assert_eq!(ascending.min_item(), Some(&1)); + assert_eq!(ascending.max_item(), Some(&1)); +} + #[test] fn test_merge_min_max_large_other() { let mut sketch1 = KllSketch::::new(DEFAULT_K).unwrap(); @@ -405,7 +441,7 @@ fn test_merge_min_max_large_other() { sketch1.update(i as f32); } let mut sketch2 = KllSketch::::new(DEFAULT_K).unwrap(); - sketch2.merge(&sketch1); + sketch2.merge(&sketch1).unwrap(); assert_eq!(sketch2.min_item().cloned(), Some(0.0)); assert_eq!(sketch2.max_item().cloned(), Some(999_999.0)); } From 20e0707c73bdec60dbf73e3ba0d97d10b5dc2a65 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:16:19 +0800 Subject: [PATCH 06/22] perf(kll): optimize updates and repeated queries --- datasketches/src/kll/mod.rs | 1 + datasketches/src/kll/sketch.rs | 166 ++++++++++------ datasketches/src/kll/sorted_view.rs | 218 ++++++++++++++------- tests-integration/tests/kll_test/sketch.rs | 46 +++++ 4 files changed, 296 insertions(+), 135 deletions(-) diff --git a/datasketches/src/kll/mod.rs b/datasketches/src/kll/mod.rs index 8b8e5631..e65f3758 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -46,6 +46,7 @@ mod value; pub use self::order::KllComparator; pub use self::order::NaturalOrder; pub use self::sketch::KllSketch; +pub use self::sorted_view::SortedView; pub use self::value::KllValue; pub use crate::common::SearchCriteria; diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index 9dcecb80..bae05f43 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -36,6 +36,7 @@ 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; @@ -56,6 +57,8 @@ pub struct KllSketch { m: u8, min_k: u16, n: u64, + num_retained: usize, + capacity: usize, is_level_zero_sorted: bool, levels: Vec>, min_item: Option, @@ -142,7 +145,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. @@ -175,6 +178,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()); @@ -243,8 +248,20 @@ impl> KllSketch { "item must belong to the comparator's ordered domain", )); } - let view = build_sorted_view(&self.levels, self.comparator.clone()); - Ok(view.rank(item, criteria)) + let inclusive = criteria == SearchCriteria::Inclusive; + let mut weight = 0u64; + for (level, items) in self.levels.iter().enumerate() { + let count = items + .iter() + .filter(|retained| match self.comparator.compare(retained, 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. @@ -261,8 +278,18 @@ impl> KllSketch { "rank must be in [0.0, 1.0], got {rank}" ))); } - let view = build_sorted_view(&self.levels, self.comparator.clone()); - Ok(view.quantile(rank, criteria)) + 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. @@ -275,8 +302,7 @@ impl> KllSketch { if self.is_empty() { return Err(Error::invalid_argument("cannot query an empty sketch")); } - let view = build_sorted_view(&self.levels, self.comparator.clone()); - view.cdf(split_points, criteria) + self.sorted_view().cdf(split_points, criteria) } /// Returns the approximate PMF for the given split points. @@ -289,8 +315,18 @@ impl> KllSketch { if self.is_empty() { return Err(Error::invalid_argument("cannot query an empty sketch")); } - let view = build_sorted_view(&self.levels, self.comparator.clone()); - view.pmf(split_points, criteria) + 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, + self.comparator.clone(), + ) } /// Returns the normalized single-sided rank error for the configured k. @@ -627,12 +663,16 @@ 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, @@ -640,12 +680,8 @@ impl> KllSketch { } } - fn capacity(&self) -> usize { - 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"); @@ -704,10 +740,11 @@ 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.num_retained += 1; self.is_level_zero_sorted = false; self.levels[0].push(item); } @@ -718,25 +755,17 @@ 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, + &self.comparator, + rand::random::(), + use_up, + ); if above.is_empty() { above = promoted; } else { @@ -749,6 +778,7 @@ impl> KllSketch { new_level.push(item); } self.levels[level] = new_level; + self.refresh_capacity_state(); } fn find_level_to_compact(&self) -> usize { @@ -793,6 +823,12 @@ impl> KllSketch { self.is_level_zero_sorted, &self.comparator, ); + 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 { @@ -887,7 +923,36 @@ 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, + comparator: &C, + 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_by(|left, right| comparator.compare(left, right)); + } + + 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"); let offset = usize::from(offset); @@ -898,20 +963,11 @@ 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, @@ -957,25 +1013,17 @@ 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, + comparator, + rand::random::(), + use_up, + ); let promoted_len = promoted.len(); if above.is_empty() { above = promoted; diff --git a/datasketches/src/kll/sorted_view.rs b/datasketches/src/kll/sorted_view.rs index 5c8cd5bb..2337a6c3 100644 --- a/datasketches/src/kll/sorted_view.rs +++ b/datasketches/src/kll/sorted_view.rs @@ -21,6 +21,10 @@ use super::order::KllComparator; 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 struct SortedView> { comparator: C, @@ -31,16 +35,15 @@ pub 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)); + fn from_sorted(mut entries: Vec>, comparator: C) -> 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, @@ -49,57 +52,119 @@ impl> SortedView { } } - pub fn rank(&self, item: &T, criteria: SearchCriteria) -> 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() + } + + /// 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 or `item` is outside the comparator's ordered domain. + pub fn rank(&self, item: &T, criteria: SearchCriteria) -> Result { + if self.is_empty() { + return Err(Error::invalid_argument("cannot query an empty view")); + } + if !self.comparator.accepts(item) { + return Err(Error::invalid_argument( + "item must belong to the comparator's ordered domain", + )); } - let idx = if criteria == SearchCriteria::Inclusive { + let index = if criteria == SearchCriteria::Inclusive { upper_bound(&self.entries, item, &self.comparator) } else { lower_bound(&self.entries, item, &self.comparator) }; - 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 fn quantile(&self, rank: f64, criteria: SearchCriteria) -> T { + /// 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 criteria == SearchCriteria::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() } + /// 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, &self.comparator)?; let mut ranks = Vec::with_capacity(split_points.len() + 1); for item in split_points { - ranks.push(self.rank(item, criteria)); + ranks.push(self.rank(item, criteria)?); } ranks.push(1.0); Ok(ranks) } + /// 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 i in (1..buckets.len()).rev() { - buckets[i] -= buckets[i - 1]; + for index in (1..buckets.len()).rev() { + buckets[index] -= buckets[index - 1]; } Ok(buckets) } @@ -107,22 +172,63 @@ impl> SortedView { pub fn build_sorted_view>( levels: &[Vec], + is_level_zero_sorted: bool, comparator: C, ) -> SortedView { - let num_retained: usize = levels.iter().map(|level| level.len()).sum(); - let mut entries = Vec::with_capacity(num_retained); + 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| comparator.compare(&left.item, &right.item)); + } + if !run.is_empty() { + runs.push(run); + } + } - for (level_idx, level) in levels.iter().enumerate() { - let weight = 1u64 << level_idx; - for item in level { - entries.push(Entry { - item: item.clone(), - weight, - }); + 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, &comparator)); + } else { + merged_runs.push(left); + } } + runs = merged_runs; } - SortedView::new(entries, comparator) + SortedView::from_sorted(runs.pop().unwrap_or_default(), comparator) +} + +fn merge_sorted_entries>( + left: Vec>, + right: Vec>, + comparator: &C, +) -> 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 comparator.compare(&left_entry.item, &right_entry.item) == Ordering::Greater { + merged.push(right.next().unwrap()); + } else { + merged.push(left.next().unwrap()); + } + } + merged.extend(left); + merged.extend(right); + merged } fn check_split_points>( @@ -145,57 +251,17 @@ fn check_split_points>( } 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; - } else { - right = mid; - } - } - left + entries.partition_point(|entry| comparator.compare(&entry.item, item) == Ordering::Less) } 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; - } - } - left + entries.partition_point(|entry| comparator.compare(&entry.item, item) != Ordering::Greater) } 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 + entries.partition_point(|entry| entry.cumulative_weight < weight) } 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 + entries.partition_point(|entry| entry.cumulative_weight <= weight) } diff --git a/tests-integration/tests/kll_test/sketch.rs b/tests-integration/tests/kll_test/sketch.rs index a52102fd..5fcd0680 100644 --- a/tests-integration/tests/kll_test/sketch.rs +++ b/tests-integration/tests/kll_test/sketch.rs @@ -309,6 +309,52 @@ fn test_rank_cdf_pmf_consistency() { } } +#[test] +fn test_sorted_view_supports_repeated_and_batch_queries() { + let mut sketch = KllSketch::::new(64).unwrap(); + for item in 0..1_000 { + sketch.update(item as f32); + } + + 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.0); + assert_eq!(view.total_weight(), 1_000); + assert_eq!( + view.quantile(1.0, SearchCriteria::Inclusive).unwrap(), + 999.0 + ); +} + +#[test] +fn test_empty_sorted_view_queries_return_errors() { + let sketch = KllSketch::::new(DEFAULT_K).unwrap(); + let view = sketch.sorted_view(); + + assert!(view.is_empty()); + assert!(view.quantile(0.5, SearchCriteria::Inclusive).is_err()); + assert!(view.rank(&0.0, SearchCriteria::Inclusive).is_err()); +} + #[test] fn test_out_of_order_split_points_return_error() { let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); From 8b5eed1b07c2fd81f4e096ded97ff084b38d619d Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:17:30 +0800 Subject: [PATCH 07/22] bench(kll): cover update query merge and serde --- benchmarks/Cargo.toml | 2 +- benchmarks/kll/merge.rs | 38 ++++++++++++++++++++++++++++++ benchmarks/kll/mod.rs | 22 ++++++++++++++++++ benchmarks/kll/query.rs | 49 +++++++++++++++++++++++++++++++++++++++ benchmarks/kll/serde.rs | 40 ++++++++++++++++++++++++++++++++ benchmarks/kll/support.rs | 42 +++++++++++++++++++++++++++++++++ benchmarks/kll/update.rs | 31 +++++++++++++++++++++++++ benchmarks/main.rs | 1 + 8 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 benchmarks/kll/merge.rs create mode 100644 benchmarks/kll/mod.rs create mode 100644 benchmarks/kll/query.rs create mode 100644 benchmarks/kll/serde.rs create mode 100644 benchmarks/kll/support.rs create mode 100644 benchmarks/kll/update.rs 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..3e501125 --- /dev/null +++ b/benchmarks/kll/query.rs @@ -0,0 +1,49 @@ +// 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::SearchCriteria; +use divan::Bencher; +use divan::black_box; + +use super::support::prepared_sketch; + +#[divan::bench] +fn rank(bencher: Bencher) { + let sketch = prepared_sketch(); + bencher + .bench_local(|| black_box(&sketch).rank(black_box(&500_000.0), 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..28169afc --- /dev/null +++ b/benchmarks/kll/serde.rs @@ -0,0 +1,40 @@ +// 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::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..5f756eac --- /dev/null +++ b/benchmarks/kll/support.rs @@ -0,0 +1,42 @@ +// 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::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(value); + } + 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; From 9402a41171f8fe7df6b761ee3aa99a819b891a5b Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:19:56 +0800 Subject: [PATCH 08/22] test(kll): organize deterministic coverage --- tests-integration/tests/kll_test/core.rs | 86 +++ tests-integration/tests/kll_test/generic.rs | 68 +++ tests-integration/tests/kll_test/main.rs | 5 +- tests-integration/tests/kll_test/merge.rs | 124 +++++ tests-integration/tests/kll_test/query.rs | 173 ++++++ tests-integration/tests/kll_test/sketch.rs | 550 -------------------- tests-integration/tests/serde_tests/kll.rs | 84 +++ 7 files changed, 539 insertions(+), 551 deletions(-) create mode 100644 tests-integration/tests/kll_test/core.rs create mode 100644 tests-integration/tests/kll_test/generic.rs create mode 100644 tests-integration/tests/kll_test/merge.rs create mode 100644 tests-integration/tests/kll_test/query.rs delete mode 100644 tests-integration/tests/kll_test/sketch.rs diff --git a/tests-integration/tests/kll_test/core.rs b/tests-integration/tests/kll_test/core.rs new file mode 100644 index 00000000..24915ab4 --- /dev/null +++ b/tests-integration/tests/kll_test/core.rs @@ -0,0 +1,86 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datasketches::error::ErrorKind; +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 as f32); + } + 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 unordered_updates_are_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); + assert_eq!(sketch.num_retained(), 1); +} + +#[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 as f32); + 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.0)); + assert_eq!(decoded.max_item(), Some(&99_999.0)); +} diff --git a/tests-integration/tests/kll_test/generic.rs b/tests-integration/tests/kll_test/generic.rs new file mode 100644 index 00000000..62b0fe3a --- /dev/null +++ b/tests-integration/tests/kll_test/generic.rs @@ -0,0 +1,68 @@ +// 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::kll::KllComparator; +use datasketches::kll::KllSketch; +use datasketches::kll::SearchCriteria; + +#[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()) + } + + fn is_compatible(&self, _other: &Self) -> bool { + true + } +} + +#[test] +fn custom_comparator_controls_queries_and_survives_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, SearchCriteria::Inclusive).unwrap(), + "2" + ); + + let decoded = KllSketch::::deserialize_with_comparator( + &sketch.serialize(), + NumericStringOrder, + ) + .unwrap(); + assert_eq!(decoded.n(), sketch.n()); + assert_eq!(decoded.num_retained(), sketch.num_retained()); + 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, SearchCriteria::Inclusive).unwrap(), + "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..a20f19a9 --- /dev/null +++ b/tests-integration/tests/kll_test/merge.rs @@ -0,0 +1,124 @@ +// 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::KllComparator; +use datasketches::kll::KllSketch; +use datasketches::kll::SearchCriteria; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct DirectionalOrder { + descending: bool, +} + +impl KllComparator for DirectionalOrder { + fn compare(&self, left: &i64, right: &i64) -> Ordering { + if self.descending { + right.cmp(left) + } else { + left.cmp(right) + } + } + + fn is_compatible(&self, other: &Self) -> bool { + self == other + } +} + +#[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 as f32); + right.update((19_999 - item) as f32); + } + + left.merge(&right).unwrap(); + + assert_eq!(left.n(), 20_000); + assert_eq!(left.min_item(), Some(&0.0)); + assert_eq!(left.max_item(), Some(&19_999.0)); + 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 as f32); + right.update((20_000 - item) as f32); + } + + 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 as f32); + } + 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.0); + second.update(2.0); + + second.merge(&first).unwrap(); + + assert_eq!(second.min_item(), Some(&1.0)); + assert_eq!(second.max_item(), Some(&2.0)); +} + +#[test] +fn merge_rejects_incompatible_comparators_without_mutation() { + let mut ascending = + KllSketch::new_with_comparator(200, DirectionalOrder { descending: false }).unwrap(); + let mut descending = + KllSketch::new_with_comparator(200, DirectionalOrder { descending: true }).unwrap(); + ascending.update(1); + descending.update(2); + + let error = ascending.merge(&descending).unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::InvalidArgument); + assert_eq!(ascending.n(), 1); + assert_eq!(ascending.min_item(), Some(&1)); + assert_eq!(ascending.max_item(), Some(&1)); +} diff --git a/tests-integration/tests/kll_test/query.rs b/tests-integration/tests/kll_test/query.rs new file mode 100644 index 00000000..81343c96 --- /dev/null +++ b/tests-integration/tests/kll_test/query.rs @@ -0,0 +1,173 @@ +// 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::KllSketch; +use datasketches::kll::SearchCriteria; + +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.0, SearchCriteria::Inclusive).is_err()); + assert!(sketch.quantile(0.5, SearchCriteria::Inclusive).is_err()); + assert!(sketch.pmf(&[0.0], SearchCriteria::Inclusive).is_err()); + assert!(sketch.cdf(&[0.0], SearchCriteria::Inclusive).is_err()); + + sketch.update(0.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); + } + for split_points in [&[1.0, 0.0][..], &[f32::NAN][..]] { + let error = sketch + .cdf(split_points, 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.0, 1.0, 2.0, 2.0] { + sketch.update(item); + } + + assert_eq!(sketch.rank(&1.0, SearchCriteria::Exclusive).unwrap(), 0.0); + assert_eq!(sketch.rank(&1.0, SearchCriteria::Inclusive).unwrap(), 0.5); + assert_eq!(sketch.rank(&2.0, SearchCriteria::Exclusive).unwrap(), 0.5); + assert_eq!(sketch.rank(&2.0, SearchCriteria::Inclusive).unwrap(), 1.0); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), + 1.0 + ); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Exclusive).unwrap(), + 2.0 + ); +} + +#[test] +fn exact_mode_queries_match_the_stream() { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + for item in 1..=100 { + sketch.update(item as f32); + } + + assert_eq!( + sketch.quantile(0.0, SearchCriteria::Inclusive).unwrap(), + 1.0 + ); + assert_eq!( + sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), + 50.0 + ); + assert_eq!( + sketch.quantile(1.0, SearchCriteria::Inclusive).unwrap(), + 100.0 + ); + for item in 1..=100 { + assert_eq!( + sketch + .rank(&(item as f32), 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 as f32); + } + + let mut previous_rank = 0.0; + for item in (0..10_000).step_by(100) { + let rank = sketch + .rank(&(item as f32), SearchCriteria::Inclusive) + .unwrap(); + assert!(rank >= previous_rank); + assert!((0.0..=1.0).contains(&rank)); + previous_rank = rank; + } + assert_eq!(sketch.min_item(), Some(&0.0)); + assert_eq!(sketch.max_item(), Some(&9_999.0)); + 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 as f32); + } + let split_points: Vec<_> = (100..10_000).step_by(100).map(|item| item as f32).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 as f32); + } + 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.0); + assert_eq!(view.total_weight(), 1_000); + assert_eq!( + view.quantile(1.0, SearchCriteria::Inclusive).unwrap(), + 999.0 + ); +} diff --git a/tests-integration/tests/kll_test/sketch.rs b/tests-integration/tests/kll_test/sketch.rs deleted file mode 100644 index 5fcd0680..00000000 --- a/tests-integration/tests/kll_test/sketch.rs +++ /dev/null @@ -1,550 +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::KllComparator; -use datasketches::kll::KllSketch; -use datasketches::kll::SearchCriteria; - -const DEFAULT_K: u16 = 200; -const MIN_K: u16 = 8; -const MAX_K: u16 = u16::MAX; -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() -} - -#[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()) - } - - fn is_compatible(&self, _other: &Self) -> bool { - true - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct DirectionalOrder { - descending: bool, -} - -impl KllComparator for DirectionalOrder { - fn compare(&self, left: &i64, right: &i64) -> Ordering { - if self.descending { - right.cmp(left) - } else { - left.cmp(right) - } - } - - fn is_compatible(&self, other: &Self) -> bool { - self == other - } -} - -#[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, SearchCriteria::Inclusive).is_err()); - assert!(sketch.quantile(0.5, SearchCriteria::Inclusive).is_err()); - assert!(sketch.pmf(&[0.0f32], SearchCriteria::Inclusive).is_err()); - assert!(sketch.cdf(&[0.0f32], SearchCriteria::Inclusive).is_err()); -} - -#[test] -fn test_quantile_out_of_range_returns_error() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(0.0); - let error = sketch - .quantile(-1.0, SearchCriteria::Inclusive) - .unwrap_err(); - assert_eq!(error.kind(), ErrorKind::InvalidArgument); -} - -#[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, SearchCriteria::Exclusive).unwrap(), 0.0); - assert_eq!(sketch.rank(&1.0, SearchCriteria::Inclusive).unwrap(), 1.0); - assert_eq!(sketch.rank(&2.0, SearchCriteria::Exclusive).unwrap(), 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, SearchCriteria::Inclusive).unwrap(), - 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, SearchCriteria::Exclusive).unwrap(), 0.0); - assert_eq!(sketch.rank(&1.0, SearchCriteria::Inclusive).unwrap(), 0.5); - assert_eq!(sketch.rank(&2.0, SearchCriteria::Exclusive).unwrap(), 0.5); - assert_eq!(sketch.rank(&2.0, SearchCriteria::Inclusive).unwrap(), 1.0); - assert_eq!( - sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), - 1.0 - ); - assert_eq!( - sketch.quantile(0.5, SearchCriteria::Exclusive).unwrap(), - 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, SearchCriteria::Inclusive).unwrap(), - 1.0 - ); - assert_eq!(sketch.max_item().cloned(), Some(n as f32)); - assert_eq!( - sketch.quantile(1.0, SearchCriteria::Inclusive).unwrap(), - n as f32 - ); - - for i in 1..=n { - let inclusive_rank = i as f64 / n as f64; - assert_eq!( - sketch.rank(&(i as f32), SearchCriteria::Inclusive).unwrap(), - inclusive_rank - ); - let exclusive_rank = (i - 1) as f64 / n as f64; - assert_eq!( - sketch.rank(&(i as f32), SearchCriteria::Exclusive).unwrap(), - 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, SearchCriteria::Inclusive).unwrap(), - 1.0 - ); - assert_eq!( - sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), - 5.0 - ); - assert_eq!( - sketch.quantile(0.99, SearchCriteria::Inclusive).unwrap(), - 10.0 - ); - assert_eq!( - sketch.quantile(1.0, SearchCriteria::Inclusive).unwrap(), - 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, SearchCriteria::Inclusive).unwrap(), - 0.0 - ); - assert_eq!( - sketch.quantile(0.01, SearchCriteria::Inclusive).unwrap(), - 0.0 - ); - assert_eq!( - sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), - 49.0 - ); - assert_eq!( - sketch.quantile(0.99, SearchCriteria::Inclusive).unwrap(), - 98.0 - ); - assert_eq!( - sketch.quantile(1.0, SearchCriteria::Inclusive).unwrap(), - 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), SearchCriteria::Exclusive).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, SearchCriteria::Exclusive).unwrap(); - let pmf = sketch.pmf(&values, SearchCriteria::Exclusive).unwrap(); - - let mut subtotal = 0.0; - for i in 0..n { - let rank = sketch.rank(&values[i], SearchCriteria::Exclusive).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, SearchCriteria::Inclusive).unwrap(); - let pmf = sketch.pmf(&values, SearchCriteria::Inclusive).unwrap(); - - let mut subtotal = 0.0; - for i in 0..n { - let rank = sketch.rank(&values[i], SearchCriteria::Inclusive).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] -fn test_sorted_view_supports_repeated_and_batch_queries() { - let mut sketch = KllSketch::::new(64).unwrap(); - for item in 0..1_000 { - sketch.update(item as f32); - } - - 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.0); - assert_eq!(view.total_weight(), 1_000); - assert_eq!( - view.quantile(1.0, SearchCriteria::Inclusive).unwrap(), - 999.0 - ); -} - -#[test] -fn test_empty_sorted_view_queries_return_errors() { - let sketch = KllSketch::::new(DEFAULT_K).unwrap(); - let view = sketch.sorted_view(); - - assert!(view.is_empty()); - assert!(view.quantile(0.5, SearchCriteria::Inclusive).is_err()); - assert!(view.rank(&0.0, SearchCriteria::Inclusive).is_err()); -} - -#[test] -fn test_out_of_order_split_points_return_error() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(0.0); - let split_points = [1.0, 0.0]; - let error = sketch - .cdf(&split_points, SearchCriteria::Inclusive) - .unwrap_err(); - assert_eq!(error.kind(), ErrorKind::InvalidArgument); -} - -#[test] -fn test_nan_split_point_returns_error() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(0.0); - let split_points = [f32::NAN]; - let error = sketch - .cdf(&split_points, SearchCriteria::Inclusive) - .unwrap_err(); - assert_eq!(error.kind(), ErrorKind::InvalidArgument); -} - -#[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).unwrap(); - - 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, SearchCriteria::Inclusive).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).unwrap(); - - 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(), - sketch2.normalized_rank_error() - ); - assert_eq!( - sketch1.normalized_pmf_error(), - sketch2.normalized_pmf_error() - ); - let median = sketch1.quantile(0.5, SearchCriteria::Inclusive).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_pmf_error(); - sketch1.merge(&sketch2).unwrap(); - assert_eq!(sketch1.normalized_pmf_error(), 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, SearchCriteria::Inclusive).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).unwrap(); - assert_eq!(sketch2.min_item().cloned(), Some(1.0)); - assert_eq!(sketch2.max_item().cloned(), Some(2.0)); -} - -#[test] -fn test_merge_rejects_incompatible_comparators_without_mutation() { - let mut ascending = - KllSketch::new_with_comparator(200, DirectionalOrder { descending: false }).unwrap(); - let mut descending = - KllSketch::new_with_comparator(200, DirectionalOrder { descending: true }).unwrap(); - ascending.update(1); - descending.update(2); - - let error = ascending.merge(&descending).unwrap_err(); - - assert_eq!(error.kind(), ErrorKind::InvalidArgument); - assert_eq!(ascending.n(), 1); - assert_eq!(ascending.min_item(), Some(&1)); - assert_eq!(ascending.max_item(), Some(&1)); -} - -#[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).unwrap(); - 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, SearchCriteria::Inclusive) - .as_deref() - .unwrap(), - "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, SearchCriteria::Inclusive) - .as_deref() - .unwrap(), - "2" - ); -} diff --git a/tests-integration/tests/serde_tests/kll.rs b/tests-integration/tests/serde_tests/kll.rs index aac40997..f8ecbb77 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,8 +29,13 @@ use std::cmp::Ordering; use std::fs; use std::path::PathBuf; +use datasketches::codec::SketchBytes; +use datasketches::codec::SketchSlice; +use datasketches::error::Error; +use datasketches::error::ErrorKind; use datasketches::kll::KllComparator; use datasketches::kll::KllSketch; +use datasketches::kll::KllValue; use crate::serialization_test_data; @@ -48,6 +54,35 @@ impl KllComparator for NumericStringOrder { } } +#[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) @@ -324,6 +359,55 @@ 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(); From a9d61f5591cb0832d538a4c5f6a92b3d1f3e270b Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:20:50 +0800 Subject: [PATCH 09/22] docs(kll): describe compatibility and query model --- CHANGELOG.md | 6 +++++- datasketches/src/kll/mod.rs | 5 ++--- datasketches/src/kll/sketch.rs | 9 ++++++++- 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 339c6a6d..c05956d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,11 @@ All significant changes to this project will be documented in this file. ### 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, including inclusive and exclusive rank, quantile, PMF, and CDF queries; reusable sorted views and batch quantiles; comparator-checked merging; custom item ordering and value encodings; and C++, Java, and Go-compatible serialization. + +### Performance improvements + +* Speed up KLL updates and rank queries by caching retained-capacity state and scanning retained items directly, and make repeated quantile queries reuse sorted levels through an owned `SortedView`. ## v0.5.0 diff --git a/datasketches/src/kll/mod.rs b/datasketches/src/kll/mod.rs index e65f3758..2638dfc9 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -21,9 +21,8 @@ //! 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. //! //! # Usage //! diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index bae05f43..175e93a8 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -166,6 +166,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 !self.comparator.accepts(&item) { return; @@ -743,7 +747,10 @@ impl> KllSketch { if self.num_retained >= self.capacity { self.compress_while_updating(); } - self.n += 1; + self.n = self + .n + .checked_add(1) + .expect("stream weight exceeds u64::MAX"); self.num_retained += 1; self.is_level_zero_sorted = false; self.levels[0].push(item); From 226a538b576e4c7993d1f684a9c72100566d17d4 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 12:22:27 +0800 Subject: [PATCH 10/22] docs(kll): polish public query types --- datasketches/src/kll/sketch.rs | 3 ++- datasketches/src/kll/sorted_view.rs | 3 ++- datasketches/src/lib.rs | 5 +++-- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index 175e93a8..ec620916 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -165,7 +165,8 @@ impl> KllSketch { /// Updates the sketch with a new item. /// - /// NaN values are ignored for floating-point types. + /// Values rejected by the configured comparator are ignored. This includes NaN values when + /// using [`NaturalOrder`]. /// /// # Panics /// diff --git a/datasketches/src/kll/sorted_view.rs b/datasketches/src/kll/sorted_view.rs index 2337a6c3..cd5f6624 100644 --- a/datasketches/src/kll/sorted_view.rs +++ b/datasketches/src/kll/sorted_view.rs @@ -18,6 +18,7 @@ use std::cmp::Ordering; use super::order::KllComparator; +use super::order::NaturalOrder; use crate::common::SearchCriteria; use crate::error::Error; @@ -26,7 +27,7 @@ use crate::error::Error; /// Build one with [`KllSketch::sorted_view`](super::KllSketch::sorted_view) when running repeated /// queries against the same sketch state. #[derive(Debug, Clone)] -pub struct SortedView> { +pub struct SortedView = NaturalOrder> { comparator: C, entries: Vec>, total_weight: u64, 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. From 1c836ac6444119e7f925abfae74bca1277b1eafb Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 13:19:19 +0800 Subject: [PATCH 11/22] docs(kll): simplify changelog entries --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c05956d7..c153d3c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,11 +6,11 @@ All significant changes to this project will be documented in this file. ### New features -* Add KLL sketches behind the `kll` feature, including inclusive and exclusive rank, quantile, PMF, and CDF queries; reusable sorted views and batch quantiles; comparator-checked merging; custom item ordering and value encodings; and C++, Java, and Go-compatible serialization. +* Add KLL sketches behind the `kll` feature, with rank, quantile, PMF, and CDF queries, merging, custom item types, and serialization. ### Performance improvements -* Speed up KLL updates and rank queries by caching retained-capacity state and scanning retained items directly, and make repeated quantile queries reuse sorted levels through an owned `SortedView`. +* Improve KLL update and query performance. ## v0.5.0 From 5b99707cbe1b8192d78d0d5138e8b1ee4ad2db7c Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 13:20:51 +0800 Subject: [PATCH 12/22] refactor: expose search criteria from common only --- benchmarks/kll/query.rs | 2 +- benchmarks/req/query.rs | 2 +- datasketches/src/kll/mod.rs | 5 ++--- datasketches/src/req/mod.rs | 4 +--- datasketches/src/req/sketch.rs | 2 +- datasketches/src/req/sorted_view.rs | 2 +- tests-integration/tests/kll_test/generic.rs | 2 +- tests-integration/tests/kll_test/merge.rs | 2 +- tests-integration/tests/kll_test/query.rs | 2 +- tests-integration/tests/req_test/accuracy.rs | 2 +- tests-integration/tests/req_test/bounds.rs | 2 +- tests-integration/tests/req_test/core.rs | 2 +- tests-integration/tests/req_test/generic.rs | 2 +- tests-integration/tests/req_test/merge.rs | 2 +- tests-integration/tests/req_test/property.rs | 2 +- tests-integration/tests/req_test/query.rs | 2 +- tests-integration/tests/req_test/sorted_view_api.rs | 2 +- tests-integration/tests/serde_tests/req.rs | 2 +- 18 files changed, 19 insertions(+), 22 deletions(-) diff --git a/benchmarks/kll/query.rs b/benchmarks/kll/query.rs index 3e501125..084c836e 100644 --- a/benchmarks/kll/query.rs +++ b/benchmarks/kll/query.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use datasketches::kll::SearchCriteria; +use datasketches::common::SearchCriteria; use divan::Bencher; use divan::black_box; 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/kll/mod.rs b/datasketches/src/kll/mod.rs index 2638dfc9..f39fdac6 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -27,7 +27,8 @@ //! # Usage //! //! ```rust -//! # use datasketches::kll::{KllSketch, SearchCriteria}; +//! # use datasketches::common::SearchCriteria; +//! # use datasketches::kll::KllSketch; //! let mut sketch = KllSketch::::new(200).unwrap(); //! sketch.update(1.0); //! sketch.update(2.0); @@ -47,8 +48,6 @@ pub use self::order::NaturalOrder; pub use self::sketch::KllSketch; pub use self::sorted_view::SortedView; pub use self::value::KllValue; -pub use crate::common::SearchCriteria; - /// Default value of parameter k. const DEFAULT_K: u16 = 200; /// Default value of parameter m. diff --git a/datasketches/src/req/mod.rs b/datasketches/src/req/mod.rs index 41d8f408..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,8 +61,6 @@ pub use self::sketch::ReqSketch; pub use self::sorted_view::SortedView; pub use self::value::ReqFloat; pub use self::value::ReqValue; -pub use crate::common::SearchCriteria; - /// Default value of `k` if not specified. Roughly 1% relative error at 95% confidence. const DEFAULT_K: u16 = 12; /// Minimum allowed value of `k`. 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/tests-integration/tests/kll_test/generic.rs b/tests-integration/tests/kll_test/generic.rs index 62b0fe3a..31e3ac83 100644 --- a/tests-integration/tests/kll_test/generic.rs +++ b/tests-integration/tests/kll_test/generic.rs @@ -17,9 +17,9 @@ use std::cmp::Ordering; +use datasketches::common::SearchCriteria; use datasketches::kll::KllComparator; use datasketches::kll::KllSketch; -use datasketches::kll::SearchCriteria; #[derive(Debug, Clone, Copy, PartialEq, Eq)] struct NumericStringOrder; diff --git a/tests-integration/tests/kll_test/merge.rs b/tests-integration/tests/kll_test/merge.rs index a20f19a9..b7052fb5 100644 --- a/tests-integration/tests/kll_test/merge.rs +++ b/tests-integration/tests/kll_test/merge.rs @@ -17,10 +17,10 @@ use std::cmp::Ordering; +use datasketches::common::SearchCriteria; use datasketches::error::ErrorKind; use datasketches::kll::KllComparator; use datasketches::kll::KllSketch; -use datasketches::kll::SearchCriteria; #[derive(Debug, Clone, Copy, PartialEq, Eq)] struct DirectionalOrder { diff --git a/tests-integration/tests/kll_test/query.rs b/tests-integration/tests/kll_test/query.rs index 81343c96..308183c7 100644 --- a/tests-integration/tests/kll_test/query.rs +++ b/tests-integration/tests/kll_test/query.rs @@ -15,9 +15,9 @@ // specific language governing permissions and limitations // under the License. +use datasketches::common::SearchCriteria; use datasketches::error::ErrorKind; use datasketches::kll::KllSketch; -use datasketches::kll::SearchCriteria; const DEFAULT_K: u16 = 200; const NUMERIC_NOISE_TOLERANCE: f64 = 1e-6; 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/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; From 52949d620cfacbe7bc2d99317cab9344fe135ce8 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 13:40:25 +0800 Subject: [PATCH 13/22] refactor: require explicit search criteria --- datasketches/src/common/search_criteria.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/datasketches/src/common/search_criteria.rs b/datasketches/src/common/search_criteria.rs index 7f480c73..3e98ca65 100644 --- a/datasketches/src/common/search_criteria.rs +++ b/datasketches/src/common/search_criteria.rs @@ -16,10 +16,9 @@ // under the License. /// Selects the rank definition used by rank, quantile, PMF, and CDF queries. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] 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, From 101fa1b6b00c89546b1de8769886f8cbb81d1323 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 13:45:53 +0800 Subject: [PATCH 14/22] docs: record search criteria migration --- CHANGELOG.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c153d3c3..5b3e7fcf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,11 +4,15 @@ 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, with rank, quantile, PMF, and CDF queries, merging, custom item types, and serialization. -### Performance improvements +### Improvements * Improve KLL update and query performance. From 83c1c7f2605081d50313b83d880f9dadac39e3ed Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 13:52:19 +0800 Subject: [PATCH 15/22] refactor(kll): tighten capacity helpers --- datasketches/src/kll/capacity.rs | 64 +++++++++++++++++++------------- 1 file changed, 39 insertions(+), 25 deletions(-) diff --git a/datasketches/src/kll/capacity.rs b/datasketches/src/kll/capacity.rs index 03057aa1..fdbce7fc 100644 --- a/datasketches/src/kll/capacity.rs +++ b/datasketches/src/kll/capacity.rs @@ -49,41 +49,55 @@ const POWERS_OF_THREE: [u64; 31] = [ 205891132094649, ]; -pub fn total_capacity(k: u16, m: u8, num_levels: usize) -> u32 { +const MAX_DEPTH: usize = 60; +const MAX_SHALLOW_DEPTH: usize = POWERS_OF_THREE.len() - 1; + +pub const fn total_capacity(k: u16, minimum_capacity: u8, num_levels: usize) -> u32 { let mut total: u32 = 0; - for level in 0..num_levels { - total += level_capacity(k, num_levels, level, m); + let mut level = 0; + while level < num_levels { + total += level_capacity(k, num_levels, level, minimum_capacity); + level += 1; } total } -pub fn level_capacity(k: u16, num_levels: usize, height: usize, min_width: u8) -> u32 { - assert!(height < num_levels, "height must be < num_levels"); - let depth = num_levels - height - 1; - let cap = capacity_at_depth(k, depth as u8); - std::cmp::max(min_width as u32, cap as u32) +pub const fn level_capacity(k: u16, num_levels: usize, level: usize, minimum_capacity: u8) -> u32 { + assert!( + level < num_levels, + "level index must be less than the number of levels" + ); + 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 + } } -fn capacity_at_depth(k: u16, depth: u8) -> u16 { - if depth > 60 { - panic!("depth must be <= 60"); - } - if depth <= 30 { +const fn capacity_at_depth(k: u16, depth: usize) -> u16 { + assert!(depth <= MAX_DEPTH, "KLL capacity depth must be at most 60"); + if depth <= MAX_SHALLOW_DEPTH { return capacity_at_shallow_depth(k, depth); } - let half = depth / 2; - let rest = depth - half; - let tmp = capacity_at_shallow_depth(k, half); - capacity_at_shallow_depth(tmp, rest) + 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) } -fn capacity_at_shallow_depth(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"); +const fn capacity_at_shallow_depth(k: u16, depth: usize) -> u16 { + assert!( + depth <= MAX_SHALLOW_DEPTH, + "shallow KLL capacity depth must be at most 30" + ); + let twice_k = (k as u64) << 1; + let scaled_capacity = (twice_k << depth) / POWERS_OF_THREE[depth]; + let result = (scaled_capacity + 1) >> 1; + assert!( + result <= k as u64, + "computed level capacity must not exceed k" + ); result as u16 } From b043f02fd6afd73391977ae192a6526e4fc80b16 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 13:54:15 +0800 Subject: [PATCH 16/22] docs(req): clarify section size comment --- datasketches/src/req/compactor.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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, From 75d7d6c36cb0deb72df20674827cadbada64b495 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 14:01:21 +0800 Subject: [PATCH 17/22] refactor(kll): improve invariant diagnostics --- datasketches/src/kll/capacity.rs | 50 +++++---- datasketches/src/kll/sketch.rs | 159 ++++++++++++++++++---------- datasketches/src/kll/sorted_view.rs | 20 ++-- datasketches/src/kll/value.rs | 12 ++- 4 files changed, 153 insertions(+), 88 deletions(-) diff --git a/datasketches/src/kll/capacity.rs b/datasketches/src/kll/capacity.rs index fdbce7fc..e4ad7f1f 100644 --- a/datasketches/src/kll/capacity.rs +++ b/datasketches/src/kll/capacity.rs @@ -52,21 +52,42 @@ const POWERS_OF_THREE: [u64; 31] = [ const MAX_DEPTH: usize = 60; const MAX_SHALLOW_DEPTH: usize = POWERS_OF_THREE.len() - 1; -pub const fn total_capacity(k: u16, minimum_capacity: u8, num_levels: usize) -> u32 { +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; - let mut level = 0; - while level < num_levels { - total += level_capacity(k, num_levels, level, minimum_capacity); - level += 1; + for level in 0..num_levels { + total += level_capacity_unchecked(k, num_levels, level, minimum_capacity); } total } -pub const fn level_capacity(k: u16, num_levels: usize, level: usize, minimum_capacity: u8) -> u32 { +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, - "level index must be less than the number of 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 { @@ -77,7 +98,6 @@ pub const fn level_capacity(k: u16, num_levels: usize, level: usize, minimum_cap } const fn capacity_at_depth(k: u16, depth: usize) -> u16 { - assert!(depth <= MAX_DEPTH, "KLL capacity depth must be at most 60"); if depth <= MAX_SHALLOW_DEPTH { return capacity_at_shallow_depth(k, depth); } @@ -88,16 +108,6 @@ const fn capacity_at_depth(k: u16, depth: usize) -> u16 { } const fn capacity_at_shallow_depth(k: u16, depth: usize) -> u16 { - assert!( - depth <= MAX_SHALLOW_DEPTH, - "shallow KLL capacity depth must be at most 30" - ); - let twice_k = (k as u64) << 1; - let scaled_capacity = (twice_k << depth) / POWERS_OF_THREE[depth]; - let result = (scaled_capacity + 1) >> 1; - assert!( - result <= k as u64, - "computed level capacity must not exceed k" - ); - result as 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/sketch.rs b/datasketches/src/kll/sketch.rs index ec620916..b2d918b3 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -214,10 +214,14 @@ impl> KllSketch { self.m, other.m ))); } - let final_n = self - .n - .checked_add(other.n) - .ok_or_else(|| Error::invalid_argument("combined stream weight exceeds u64::MAX"))?; + 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); @@ -494,12 +498,17 @@ 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, @@ -523,17 +532,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 { @@ -554,21 +560,22 @@ 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 @@ -584,9 +591,17 @@ fn deserialize_with_serde>( 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::deserial(format!( + "insufficient item data: expected at least {min_item_bytes} bytes, got {available_item_bytes}" + ))); } let mut levels = Vec::with_capacity(num_levels); @@ -621,8 +636,11 @@ 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) @@ -688,7 +706,10 @@ impl> KllSketch { fn level_offsets(&self) -> Vec { 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; @@ -748,10 +769,13 @@ impl> KllSketch { if self.num_retained >= self.capacity { self.compress_while_updating(); } - self.n = self - .n - .checked_add(1) - .expect("stream weight exceeds u64::MAX"); + 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); @@ -798,7 +822,10 @@ 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) { @@ -857,9 +884,14 @@ impl> KllSketch { .as_ref() .ok_or_else(|| Error::deserial("non-empty sketch must have a maximum item"))?; - if !self.comparator.accepts(min_item) || !self.comparator.accepts(max_item) { + if !self.comparator.accepts(min_item) { + return Err(Error::deserial( + "serialized minimum item is outside the comparator's ordered domain", + )); + } + if !self.comparator.accepts(max_item) { return Err(Error::deserial( - "minimum and maximum items must belong to the comparator's ordered domain", + "serialized maximum item is outside the comparator's ordered domain", )); } if self.comparator.compare(min_item, max_item) == Ordering::Greater { @@ -873,41 +905,55 @@ 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 self.comparator.compare(&pair[0], &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 { + for (item_index, item) in level.iter().enumerate() { if !self.comparator.accepts(item) { - return Err(Error::deserial( - "retained items must belong to the comparator's ordered domain", - )); + return Err(Error::deserial(format!( + "retained item at level {level_index}, index {item_index} is outside the comparator's ordered domain" + ))); } 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", - )); + 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}" + )) + })?; } } @@ -962,7 +1008,10 @@ fn compact_level>( 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 diff --git a/datasketches/src/kll/sorted_view.rs b/datasketches/src/kll/sorted_view.rs index cd5f6624..2a4b24a4 100644 --- a/datasketches/src/kll/sorted_view.rs +++ b/datasketches/src/kll/sorted_view.rs @@ -236,16 +236,20 @@ fn check_split_points>( split_points: &[T], comparator: &C, ) -> Result<(), Error> { - if !split_points.iter().all(|point| comparator.accepts(point)) { - return Err(Error::invalid_argument( - "split points must belong to the comparator's ordered domain", - )); + if let Some(index) = split_points + .iter() + .position(|point| !comparator.accepts(point)) + { + return Err(Error::invalid_argument(format!( + "split point at index {index} is outside the comparator's ordered domain" + ))); } - for pair in split_points.windows(2) { + for (index, pair) in split_points.windows(2).enumerate() { if comparator.compare(&pair[0], &pair[1]) != Ordering::Less { - return Err(Error::invalid_argument( - "split points must be unique and monotonically increasing", - )); + return Err(Error::invalid_argument(format!( + "split points at indices {index} and {} must be strictly increasing", + index + 1 + ))); } } Ok(()) diff --git a/datasketches/src/kll/value.rs b/datasketches/src/kll/value.rs index e14057ae..0d394dc7 100644 --- a/datasketches/src/kll/value.rs +++ b/datasketches/src/kll/value.rs @@ -108,12 +108,14 @@ impl KllValue for String { 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 available = input.remaining().len(); + let bytes = input.remaining().get(..len).ok_or_else(|| { + Error::deserial(format!( + "insufficient string data: expected {len} bytes, got {available}" + )) + })?; let value = std::str::from_utf8(bytes) - .map_err(|_| Error::deserial("invalid utf-8 string"))? + .map_err(|error| Error::deserial(format!("invalid UTF-8 string: {error}")))? .to_owned(); input.advance(len as u64); Ok(value) From a2a76474c93edc0467f495a6319e0d6ee29848a3 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 14:36:31 +0800 Subject: [PATCH 18/22] refactor(kll): encode ordering in item types --- CHANGELOG.md | 2 +- benchmarks/kll/query.rs | 5 +- benchmarks/kll/serde.rs | 3 +- benchmarks/kll/support.rs | 7 +- datasketches/src/kll/mod.rs | 16 +- datasketches/src/kll/order.rs | 57 ------ datasketches/src/kll/sketch.rs | 186 +++++--------------- datasketches/src/kll/sorted_view.rs | 66 +++---- datasketches/src/kll/value.rs | 127 +++++++++++-- tests-integration/tests/kll_test/core.rs | 37 ++-- tests-integration/tests/kll_test/generic.rs | 66 ++++--- tests-integration/tests/kll_test/merge.rs | 78 ++------ tests-integration/tests/kll_test/query.rs | 89 ++++------ tests-integration/tests/serde_tests/kll.rs | 65 ++++--- 14 files changed, 344 insertions(+), 460 deletions(-) delete mode 100644 datasketches/src/kll/order.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b3e7fcf..67670801 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ All significant changes to this project will be documented in this file. ### New features -* Add KLL sketches behind the `kll` feature, with rank, quantile, PMF, and CDF queries, merging, custom item types, and 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 diff --git a/benchmarks/kll/query.rs b/benchmarks/kll/query.rs index 084c836e..1c61aa08 100644 --- a/benchmarks/kll/query.rs +++ b/benchmarks/kll/query.rs @@ -16,6 +16,7 @@ // under the License. use datasketches::common::SearchCriteria; +use datasketches::kll::KllFloat; use divan::Bencher; use divan::black_box; @@ -24,8 +25,8 @@ use super::support::prepared_sketch; #[divan::bench] fn rank(bencher: Bencher) { let sketch = prepared_sketch(); - bencher - .bench_local(|| black_box(&sketch).rank(black_box(&500_000.0), SearchCriteria::Inclusive)); + let item = KllFloat::::new(500_000.0).unwrap(); + bencher.bench_local(|| black_box(&sketch).rank(black_box(&item), SearchCriteria::Inclusive)); } #[divan::bench] diff --git a/benchmarks/kll/serde.rs b/benchmarks/kll/serde.rs index 28169afc..3882f23d 100644 --- a/benchmarks/kll/serde.rs +++ b/benchmarks/kll/serde.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use datasketches::kll::KllFloat; use datasketches::kll::KllSketch; use divan::Bencher; use divan::black_box; @@ -36,5 +37,5 @@ fn deserialize(bencher: Bencher) { let bytes = prepared_sketch().serialize(); bencher .counter(BytesCount::new(bytes.len())) - .bench_local(|| KllSketch::::deserialize(black_box(&bytes)).unwrap()); + .bench_local(|| KllSketch::>::deserialize(black_box(&bytes)).unwrap()); } diff --git a/benchmarks/kll/support.rs b/benchmarks/kll/support.rs index 5f756eac..20708f19 100644 --- a/benchmarks/kll/support.rs +++ b/benchmarks/kll/support.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use datasketches::kll::KllFloat; use datasketches::kll::KllSketch; use rand::RngExt; use rand::SeedableRng; @@ -29,14 +30,14 @@ pub(super) fn values(len: usize) -> Vec { .collect() } -pub(super) fn build_sketch(values: &[f64]) -> KllSketch { +pub(super) fn build_sketch(values: &[f64]) -> KllSketch> { let mut sketch = KllSketch::new(DEFAULT_K).unwrap(); for &value in values { - sketch.update(value); + sketch.update(KllFloat::::new(value).unwrap()); } sketch } -pub(super) fn prepared_sketch() -> KllSketch { +pub(super) fn prepared_sketch() -> KllSketch> { build_sketch(&values(100_000)) } diff --git a/datasketches/src/kll/mod.rs b/datasketches/src/kll/mod.rs index f39fdac6..469bc7f1 100644 --- a/datasketches/src/kll/mod.rs +++ b/datasketches/src/kll/mod.rs @@ -24,29 +24,31 @@ //! 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 mut sketch = KllSketch::::new(200).unwrap(); +//! sketch.update(1); +//! sketch.update(2); //! let q = sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(); -//! assert!(q >= 1.0 && q <= 2.0); +//! assert!((1..=2).contains(&q)); //! ``` mod capacity; -mod order; mod serialization; mod sketch; mod sorted_view; mod value; -pub use self::order::KllComparator; -pub use self::order::NaturalOrder; pub use self::sketch::KllSketch; pub use self::sorted_view::SortedView; +pub use self::value::KllFloat; pub use self::value::KllValue; /// Default value of parameter k. const DEFAULT_K: u16 = 200; diff --git a/datasketches/src/kll/order.rs b/datasketches/src/kll/order.rs deleted file mode 100644 index daf2264b..00000000 --- a/datasketches/src/kll/order.rs +++ /dev/null @@ -1,57 +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; - -/// Defines the ordering used by a KLL sketch. -/// -/// Accepted values must form a total order. Sketches can be merged only when their comparators are -/// compatible: they must accept the same values and order every pair of accepted values -/// identically. -pub trait KllComparator: Clone { - /// Compares two accepted values. - fn compare(&self, left: &T, right: &T) -> Ordering; - - /// Returns whether `item` belongs to this comparator's ordered domain. - /// - /// Updates with rejected values are ignored. The default accepts every value. - fn accepts(&self, _item: &T) -> bool { - true - } - - /// Returns whether `other` defines the same ordered domain and comparison semantics. - fn is_compatible(&self, other: &Self) -> bool; -} - -/// Uses the value's natural partial ordering and rejects unordered values such as NaN. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct NaturalOrder; - -impl KllComparator for NaturalOrder { - fn compare(&self, left: &T, right: &T) -> Ordering { - left.partial_cmp(right) - .expect("accepted KLL values must be totally ordered") - } - - fn accepts(&self, item: &T) -> bool { - item.partial_cmp(item).is_some() - } - - fn is_compatible(&self, _other: &Self) -> bool { - true - } -} diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index b2d918b3..fc9bf5e0 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -23,8 +23,6 @@ use super::MAX_K; use super::MIN_K; use super::capacity::level_capacity; use super::capacity::total_capacity; -use super::order::KllComparator; -use super::order::NaturalOrder; use super::serialization::DATA_START; use super::serialization::DATA_START_SINGLE_ITEM; use super::serialization::EMPTY_SIZE_BYTES; @@ -51,8 +49,7 @@ use crate::error::Error; /// /// 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, @@ -65,22 +62,13 @@ pub struct KllSketch { 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 @@ -91,36 +79,16 @@ impl KllSketch { /// /// ``` /// # 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 `8..=65535`. - 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. @@ -165,16 +133,10 @@ impl> KllSketch { /// Updates the sketch with a new item. /// - /// Values rejected by the configured comparator are ignored. This includes NaN values when - /// using [`NaturalOrder`]. - /// /// # Panics /// /// Panics if the stream weight would exceed [`u64::MAX`]. pub fn update(&mut self, item: T) { - if !self.comparator.accepts(&item) { - return; - } self.update_min_max(&item); self.internal_update(item); } @@ -196,18 +158,12 @@ impl> KllSketch { /// /// # Errors /// - /// Returns an error if the sketches use incompatible comparators or their combined stream - /// weight exceeds [`u64::MAX`]. - pub fn merge(&mut self, other: &KllSketch) -> Result<(), Error> { + /// 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 Ok(()); } - if !self.comparator.is_compatible(&other.comparator) { - return Err(Error::invalid_argument( - "cannot merge sketches with incompatible comparators", - )); - } if self.m != other.m { return Err(Error::invalid_argument(format!( "cannot merge sketches with different m values: {} and {}", @@ -246,23 +202,17 @@ impl> KllSketch { /// /// # Errors /// - /// Returns an error if the sketch is empty or `item` is outside the comparator's ordered - /// domain. + /// Returns an error if the sketch is empty. pub fn rank(&self, item: &T, criteria: SearchCriteria) -> Result { if self.is_empty() { return Err(Error::invalid_argument("cannot query an empty sketch")); } - if !self.comparator.accepts(item) { - return Err(Error::invalid_argument( - "item must belong to the comparator's ordered domain", - )); - } let inclusive = criteria == SearchCriteria::Inclusive; let mut weight = 0u64; for (level, items) in self.levels.iter().enumerate() { let count = items .iter() - .filter(|retained| match self.comparator.compare(retained, item) { + .filter(|retained| match (*retained).cmp(item) { Ordering::Less => true, Ordering::Equal => inclusive, Ordering::Greater => false, @@ -305,8 +255,8 @@ impl> KllSketch { /// /// # Errors /// - /// Returns an error if the sketch is empty, a split point is outside the comparator's ordered - /// domain, or the split points are not unique and strictly increasing. + /// 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 Err(Error::invalid_argument("cannot query an empty sketch")); @@ -318,8 +268,8 @@ impl> KllSketch { /// /// # Errors /// - /// Returns an error if the sketch is empty, a split point is outside the comparator's ordered - /// domain, or the split points are not unique and strictly increasing. + /// 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 Err(Error::invalid_argument("cannot query an empty sketch")); @@ -330,12 +280,8 @@ impl> KllSketch { /// 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, - self.comparator.clone(), - ) + 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. @@ -349,7 +295,7 @@ impl> KllSketch { } } -fn serialized_size>(sketch: &KllSketch) -> usize { +fn serialized_size(sketch: &KllSketch) -> usize { if sketch.is_empty() { return EMPTY_SIZE_BYTES; } @@ -373,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); @@ -445,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 @@ -511,7 +454,6 @@ fn deserialize_with_serde>( ))); } return Ok(KllSketch::make( - comparator, k, k, 0, @@ -618,7 +560,6 @@ fn deserialize_with_serde>( } let mut sketch = KllSketch::make( - comparator, k, min_k, n, @@ -646,38 +587,25 @@ fn deserialize_with_serde>( 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. - /// - /// # Errors - /// - /// Returns `InvalidData` if the image is truncated, malformed, or inconsistent with - /// `comparator`. - 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 - /// naturally ordered. + /// 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, @@ -689,7 +617,6 @@ impl> KllSketch { 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, @@ -728,11 +655,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()); } } @@ -740,7 +667,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(); @@ -748,12 +675,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()); } } @@ -794,14 +721,13 @@ impl> KllSketch { current, level, self.is_level_zero_sorted, - &self.comparator, 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; @@ -828,7 +754,7 @@ impl> KllSketch { ); } - 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]; @@ -847,17 +773,11 @@ 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(); } @@ -884,17 +804,7 @@ impl> KllSketch { .as_ref() .ok_or_else(|| Error::deserial("non-empty sketch must have a maximum item"))?; - if !self.comparator.accepts(min_item) { - return Err(Error::deserial( - "serialized minimum item is outside the comparator's ordered domain", - )); - } - if !self.comparator.accepts(max_item) { - return Err(Error::deserial( - "serialized maximum item is outside the comparator's ordered domain", - )); - } - 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", )); @@ -922,7 +832,7 @@ impl> KllSketch { let must_be_sorted = level_index > 0 || self.is_level_zero_sorted; if must_be_sorted { for (item_index, pair) in level.windows(2).enumerate() { - if self.comparator.compare(&pair[0], &pair[1]) == Ordering::Greater { + 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 @@ -932,14 +842,7 @@ impl> KllSketch { } for (item_index, item) in level.iter().enumerate() { - if !self.comparator.accepts(item) { - return Err(Error::deserial(format!( - "retained item at level {level_index}, index {item_index} is outside the comparator's ordered domain" - ))); - } - if self.comparator.compare(item, min_item) == Ordering::Less - || self.comparator.compare(item, max_item) == Ordering::Greater - { + 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" ))); @@ -977,11 +880,10 @@ fn normalized_rank_error(k: u16, pmf: bool) -> f64 { } } -fn compact_level>( +fn compact_level( mut items: Vec, level: usize, is_level_zero_sorted: bool, - comparator: &C, offset: bool, use_up: bool, ) -> (Option, Vec) { @@ -993,7 +895,7 @@ fn compact_level>( None }; if level_zero_needs_sorting { - items.sort_unstable_by(|left, right| comparator.compare(left, right)); + items.sort_unstable(); } let mut items = items.into_iter(); @@ -1025,17 +927,13 @@ fn downsample>(items: I, offset: bool, use_up: .collect() } -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()); @@ -1046,12 +944,11 @@ 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(); @@ -1077,7 +974,6 @@ fn general_compress>( current, current_level, is_level_zero_sorted, - comparator, rand::random::(), use_up, ); @@ -1085,7 +981,7 @@ fn general_compress>( 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; diff --git a/datasketches/src/kll/sorted_view.rs b/datasketches/src/kll/sorted_view.rs index 2a4b24a4..b58a69a9 100644 --- a/datasketches/src/kll/sorted_view.rs +++ b/datasketches/src/kll/sorted_view.rs @@ -17,8 +17,6 @@ use std::cmp::Ordering; -use super::order::KllComparator; -use super::order::NaturalOrder; use crate::common::SearchCriteria; use crate::error::Error; @@ -27,8 +25,7 @@ use crate::error::Error; /// Build one with [`KllSketch::sorted_view`](super::KllSketch::sorted_view) when running repeated /// queries against the same sketch state. #[derive(Debug, Clone)] -pub struct SortedView = NaturalOrder> { - comparator: C, +pub struct SortedView { entries: Vec>, total_weight: u64, } @@ -39,15 +36,14 @@ struct Entry { cumulative_weight: u64, } -impl> SortedView { - fn from_sorted(mut entries: Vec>, comparator: C) -> Self { +impl SortedView { + fn from_sorted(mut entries: Vec>) -> Self { let mut total_weight = 0u64; for entry in &mut entries { total_weight += entry.cumulative_weight; entry.cumulative_weight = total_weight; } Self { - comparator, entries, total_weight, } @@ -72,21 +68,15 @@ impl> SortedView { /// /// # Errors /// - /// Returns an error if the view is empty or `item` is outside the comparator's ordered domain. + /// 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")); } - if !self.comparator.accepts(item) { - return Err(Error::invalid_argument( - "item must belong to the comparator's ordered domain", - )); - } - let index = if criteria == SearchCriteria::Inclusive { - upper_bound(&self.entries, item, &self.comparator) + upper_bound(&self.entries, item) } else { - lower_bound(&self.entries, item, &self.comparator) + lower_bound(&self.entries, item) }; if index == 0 { @@ -148,7 +138,7 @@ impl> SortedView { if self.is_empty() { return Err(Error::invalid_argument("cannot query an empty view")); } - check_split_points(split_points, &self.comparator)?; + 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, criteria)?); @@ -171,11 +161,10 @@ impl> SortedView { } } -pub fn build_sorted_view>( +pub fn build_sorted_view( levels: &[Vec], is_level_zero_sorted: bool, - comparator: C, -) -> SortedView { +) -> SortedView { let mut runs = Vec::with_capacity(levels.len()); for (level_index, level) in levels.iter().enumerate() { let weight = 1u64 << level_index; @@ -188,7 +177,7 @@ pub fn build_sorted_view>( }) .collect(); if level_index == 0 && !is_level_zero_sorted { - run.sort_unstable_by(|left, right| comparator.compare(&left.item, &right.item)); + run.sort_unstable_by(|left, right| left.item.cmp(&right.item)); } if !run.is_empty() { runs.push(run); @@ -200,7 +189,7 @@ pub fn build_sorted_view>( 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, &comparator)); + merged_runs.push(merge_sorted_entries(left, right)); } else { merged_runs.push(left); } @@ -208,20 +197,16 @@ pub fn build_sorted_view>( runs = merged_runs; } - SortedView::from_sorted(runs.pop().unwrap_or_default(), comparator) + SortedView::from_sorted(runs.pop().unwrap_or_default()) } -fn merge_sorted_entries>( - left: Vec>, - right: Vec>, - comparator: &C, -) -> Vec> { +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 comparator.compare(&left_entry.item, &right_entry.item) == Ordering::Greater { + if left_entry.item.cmp(&right_entry.item) == Ordering::Greater { merged.push(right.next().unwrap()); } else { merged.push(left.next().unwrap()); @@ -232,20 +217,9 @@ fn merge_sorted_entries>( merged } -fn check_split_points>( - split_points: &[T], - comparator: &C, -) -> Result<(), Error> { - if let Some(index) = split_points - .iter() - .position(|point| !comparator.accepts(point)) - { - return Err(Error::invalid_argument(format!( - "split point at index {index} is outside the comparator's ordered domain" - ))); - } +fn check_split_points(split_points: &[T]) -> Result<(), Error> { for (index, pair) in split_points.windows(2).enumerate() { - if comparator.compare(&pair[0], &pair[1]) != Ordering::Less { + 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 @@ -255,12 +229,12 @@ fn check_split_points>( Ok(()) } -fn lower_bound>(entries: &[Entry], item: &T, comparator: &C) -> usize { - entries.partition_point(|entry| comparator.compare(&entry.item, item) == Ordering::Less) +fn lower_bound(entries: &[Entry], item: &T) -> usize { + entries.partition_point(|entry| entry.item.cmp(item) == Ordering::Less) } -fn upper_bound>(entries: &[Entry], item: &T, comparator: &C) -> usize { - entries.partition_point(|entry| comparator.compare(&entry.item, item) != Ordering::Greater) +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 { diff --git a/datasketches/src/kll/value.rs b/datasketches/src/kll/value.rs index 0d394dc7..60699c9e 100644 --- a/datasketches/src/kll/value.rs +++ b/datasketches/src/kll/value.rs @@ -15,15 +15,110 @@ // 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::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 -/// item type with a [`KllComparator`](crate::kll::KllComparator). The encoded representation must -/// preserve the comparator's ordering across a round trip. +/// 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; @@ -38,39 +133,41 @@ pub trait KllValue: Clone { fn deserialize(input: &mut SketchSlice<'_>) -> Result; } -impl KllValue for f32 { - const MIN_SERIALIZED_SIZE: usize = 4; +impl KllValue for KllFloat { + const MIN_SERIALIZED_SIZE: usize = size_of::(); fn serialized_size(_value: &Self) -> usize { - 4 + size_of::() } fn serialize(value: &Self, bytes: &mut SketchBytes) { - bytes.write_f32_le(*value); + bytes.write_f32_le(value.0); } fn deserialize(input: &mut SketchSlice<'_>) -> Result { - input + let value = input .read_f32_le() - .map_err(|_| Error::insufficient_data("f32")) + .map_err(|_| Error::insufficient_data("f32"))?; + Self::new(value).map_err(|_| Error::deserial("KLL float must not be NaN")) } } -impl KllValue for f64 { - const MIN_SERIALIZED_SIZE: usize = 8; +impl KllValue for KllFloat { + const MIN_SERIALIZED_SIZE: usize = size_of::(); fn serialized_size(_value: &Self) -> usize { - 8 + size_of::() } fn serialize(value: &Self, bytes: &mut SketchBytes) { - bytes.write_f64_le(*value); + bytes.write_f64_le(value.0); } fn deserialize(input: &mut SketchSlice<'_>) -> Result { - input + let value = input .read_f64_le() - .map_err(|_| Error::insufficient_data("f64")) + .map_err(|_| Error::insufficient_data("f64"))?; + Self::new(value).map_err(|_| Error::deserial("KLL float must not be NaN")) } } diff --git a/tests-integration/tests/kll_test/core.rs b/tests-integration/tests/kll_test/core.rs index 24915ab4..c0bcaec0 100644 --- a/tests-integration/tests/kll_test/core.rs +++ b/tests-integration/tests/kll_test/core.rs @@ -16,6 +16,7 @@ // under the License. use datasketches::error::ErrorKind; +use datasketches::kll::KllFloat; use datasketches::kll::KllSketch; const DEFAULT_K: u16 = 200; @@ -24,16 +25,16 @@ const MAX_K: u16 = u16::MAX; #[test] fn k_limits() { - KllSketch::::new(MIN_K).unwrap(); - KllSketch::::new(MAX_K).unwrap(); + KllSketch::::new(MIN_K).unwrap(); + KllSketch::::new(MAX_K).unwrap(); - let error = KllSketch::::new(MIN_K - 1).unwrap_err(); + 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(); + let mut sketch = KllSketch::::new(64).unwrap(); assert!(sketch.is_empty()); assert!(!sketch.is_estimation_mode()); assert_eq!(sketch.n(), 0); @@ -42,7 +43,7 @@ fn empty_and_reset_state() { assert_eq!(sketch.max_item(), None); for item in 0..10_000 { - sketch.update(item as f32); + sketch.update(item); } assert!(sketch.is_estimation_mode()); assert!(sketch.num_retained() > 0); @@ -59,28 +60,28 @@ fn empty_and_reset_state() { } #[test] -fn unordered_updates_are_ignored() { - let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); - sketch.update(f32::NAN); - assert!(sketch.is_empty()); +fn float_adapter_rejects_nan() { + assert_eq!( + KllFloat::::new(f32::NAN).unwrap_err().kind(), + ErrorKind::InvalidArgument + ); - sketch.update(0.0); - sketch.update(f32::NAN); - assert_eq!(sketch.n(), 1); - assert_eq!(sketch.num_retained(), 1); + 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(); + let mut sketch = KllSketch::::new(32).unwrap(); for item in 0..100_000 { - sketch.update(item as f32); + sketch.update(item); assert!(sketch.num_retained() <= sketch.n() as usize); } - let decoded = KllSketch::::deserialize(&sketch.serialize()).unwrap(); + 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.0)); - assert_eq!(decoded.max_item(), Some(&99_999.0)); + 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 index 31e3ac83..14d06677 100644 --- a/tests-integration/tests/kll_test/generic.rs +++ b/tests-integration/tests/kll_test/generic.rs @@ -17,52 +17,68 @@ use std::cmp::Ordering; +use datasketches::codec::SketchBytes; +use datasketches::codec::SketchSlice; use datasketches::common::SearchCriteria; -use datasketches::kll::KllComparator; +use datasketches::error::Error; use datasketches::kll::KllSketch; +use datasketches::kll::KllValue; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct NumericStringOrder; +#[derive(Debug, Clone, PartialEq, Eq)] +struct NumericString(String); -impl KllComparator for NumericStringOrder { - fn compare(&self, left: &String, right: &String) -> Ordering { - left.parse::() +impl Ord for NumericString { + fn cmp(&self, other: &Self) -> Ordering { + self.0 + .parse::() .unwrap() - .cmp(&right.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 is_compatible(&self, _other: &Self) -> bool { - true + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + String::deserialize(input).map(Self) } } #[test] -fn custom_comparator_controls_queries_and_survives_roundtrip() { - let mut sketch = - KllSketch::::new_with_comparator(200, NumericStringOrder) - .unwrap(); +fn custom_item_order_controls_queries_and_survives_roundtrip() { + let mut sketch = KllSketch::::new(200).unwrap(); for item in ["2", "10", "1"] { - sketch.update(item.to_owned()); + sketch.update(NumericString(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.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(), + sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap().0, "2" ); - let decoded = KllSketch::::deserialize_with_comparator( - &sketch.serialize(), - NumericStringOrder, - ) - .unwrap(); + 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(String::as_str), Some("1")); - assert_eq!(decoded.max_item().map(String::as_str), Some("10")); + 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(), + decoded.quantile(0.5, SearchCriteria::Inclusive).unwrap().0, "2" ); } diff --git a/tests-integration/tests/kll_test/merge.rs b/tests-integration/tests/kll_test/merge.rs index b7052fb5..417fa1b1 100644 --- a/tests-integration/tests/kll_test/merge.rs +++ b/tests-integration/tests/kll_test/merge.rs @@ -15,46 +15,23 @@ // specific language governing permissions and limitations // under the License. -use std::cmp::Ordering; - use datasketches::common::SearchCriteria; -use datasketches::error::ErrorKind; -use datasketches::kll::KllComparator; use datasketches::kll::KllSketch; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct DirectionalOrder { - descending: bool, -} - -impl KllComparator for DirectionalOrder { - fn compare(&self, left: &i64, right: &i64) -> Ordering { - if self.descending { - right.cmp(left) - } else { - left.cmp(right) - } - } - - fn is_compatible(&self, other: &Self) -> bool { - self == other - } -} - #[test] fn merge_preserves_weight_extrema_and_query_invariants() { - let mut left = KllSketch::::new(200).unwrap(); - let mut right = KllSketch::::new(200).unwrap(); + let mut left = KllSketch::::new(200).unwrap(); + let mut right = KllSketch::::new(200).unwrap(); for item in 0..10_000 { - left.update(item as f32); - right.update((19_999 - item) as f32); + 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.0)); - assert_eq!(left.max_item(), Some(&19_999.0)); + 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) @@ -64,11 +41,11 @@ fn merge_preserves_weight_extrema_and_query_invariants() { #[test] fn merge_tracks_the_smallest_estimation_k() { - let mut left = KllSketch::::new(256).unwrap(); - let mut right = KllSketch::::new(128).unwrap(); + let mut left = KllSketch::::new(256).unwrap(); + let mut right = KllSketch::::new(128).unwrap(); for item in 0..10_000 { - left.update(item as f32); - right.update((20_000 - item) as f32); + left.update(item); + right.update(20_000 - item); } left.merge(&right).unwrap(); @@ -80,11 +57,11 @@ fn merge_tracks_the_smallest_estimation_k() { #[test] fn merging_an_empty_lower_k_sketch_does_not_change_accuracy() { - let mut sketch = KllSketch::::new(256).unwrap(); + let mut sketch = KllSketch::::new(256).unwrap(); for item in 0..10_000 { - sketch.update(item as f32); + sketch.update(item); } - let empty = KllSketch::::new(128).unwrap(); + let empty = KllSketch::::new(128).unwrap(); let rank_error = sketch.normalized_rank_error(); sketch.merge(&empty).unwrap(); @@ -95,30 +72,13 @@ fn merging_an_empty_lower_k_sketch_does_not_change_accuracy() { #[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.0); - second.update(2.0); + 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.0)); - assert_eq!(second.max_item(), Some(&2.0)); -} - -#[test] -fn merge_rejects_incompatible_comparators_without_mutation() { - let mut ascending = - KllSketch::new_with_comparator(200, DirectionalOrder { descending: false }).unwrap(); - let mut descending = - KllSketch::new_with_comparator(200, DirectionalOrder { descending: true }).unwrap(); - ascending.update(1); - descending.update(2); - - let error = ascending.merge(&descending).unwrap_err(); - - assert_eq!(error.kind(), ErrorKind::InvalidArgument); - assert_eq!(ascending.n(), 1); - assert_eq!(ascending.min_item(), Some(&1)); - assert_eq!(ascending.max_item(), Some(&1)); + 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 index 308183c7..5474df95 100644 --- a/tests-integration/tests/kll_test/query.rs +++ b/tests-integration/tests/kll_test/query.rs @@ -24,72 +24,54 @@ 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.0, SearchCriteria::Inclusive).is_err()); + 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.0], SearchCriteria::Inclusive).is_err()); - assert!(sketch.cdf(&[0.0], SearchCriteria::Inclusive).is_err()); + assert!(sketch.pmf(&[0], SearchCriteria::Inclusive).is_err()); + assert!(sketch.cdf(&[0], SearchCriteria::Inclusive).is_err()); - sketch.update(0.0); + 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); } - for split_points in [&[1.0, 0.0][..], &[f32::NAN][..]] { - let error = sketch - .cdf(split_points, 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.0, 1.0, 2.0, 2.0] { + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); + for item in [1, 1, 2, 2] { sketch.update(item); } - assert_eq!(sketch.rank(&1.0, SearchCriteria::Exclusive).unwrap(), 0.0); - assert_eq!(sketch.rank(&1.0, SearchCriteria::Inclusive).unwrap(), 0.5); - assert_eq!(sketch.rank(&2.0, SearchCriteria::Exclusive).unwrap(), 0.5); - assert_eq!(sketch.rank(&2.0, SearchCriteria::Inclusive).unwrap(), 1.0); - assert_eq!( - sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), - 1.0 - ); - assert_eq!( - sketch.quantile(0.5, SearchCriteria::Exclusive).unwrap(), - 2.0 - ); + 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(); + let mut sketch = KllSketch::::new(DEFAULT_K).unwrap(); for item in 1..=100 { - sketch.update(item as f32); + sketch.update(item); } - assert_eq!( - sketch.quantile(0.0, SearchCriteria::Inclusive).unwrap(), - 1.0 - ); - assert_eq!( - sketch.quantile(0.5, SearchCriteria::Inclusive).unwrap(), - 50.0 - ); + 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.0 + 100 ); for item in 1..=100 { assert_eq!( - sketch - .rank(&(item as f32), SearchCriteria::Inclusive) - .unwrap(), + sketch.rank(&item, SearchCriteria::Inclusive).unwrap(), item as f64 / 100.0 ); } @@ -97,32 +79,30 @@ fn exact_mode_queries_match_the_stream() { #[test] fn estimation_mode_queries_preserve_deterministic_invariants() { - let mut sketch = KllSketch::::new(64).unwrap(); + let mut sketch = KllSketch::::new(64).unwrap(); for item in 0..10_000 { - sketch.update(item as f32); + sketch.update(item); } let mut previous_rank = 0.0; for item in (0..10_000).step_by(100) { - let rank = sketch - .rank(&(item as f32), SearchCriteria::Inclusive) - .unwrap(); + 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.0)); - assert_eq!(sketch.max_item(), Some(&9_999.0)); + 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(); + let mut sketch = KllSketch::::new(64).unwrap(); for item in 0..10_000 { - sketch.update(item as f32); + sketch.update(item); } - let split_points: Vec<_> = (100..10_000).step_by(100).map(|item| item as f32).collect(); + 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(); @@ -139,9 +119,9 @@ fn rank_cdf_and_pmf_are_consistent() { #[test] fn sorted_view_supports_repeated_and_batch_queries() { - let mut sketch = KllSketch::::new(64).unwrap(); + let mut sketch = KllSketch::::new(64).unwrap(); for item in 0..1_000 { - sketch.update(item as f32); + sketch.update(item); } let view = sketch.sorted_view(); let ranks = [0.0, 0.25, 0.5, 0.75, 1.0]; @@ -164,10 +144,7 @@ fn sorted_view_supports_repeated_and_batch_queries() { ); } - sketch.update(2_000.0); + sketch.update(2_000); assert_eq!(view.total_weight(), 1_000); - assert_eq!( - view.quantile(1.0, SearchCriteria::Inclusive).unwrap(), - 999.0 - ); + assert_eq!(view.quantile(1.0, SearchCriteria::Inclusive).unwrap(), 999); } diff --git a/tests-integration/tests/serde_tests/kll.rs b/tests-integration/tests/serde_tests/kll.rs index f8ecbb77..aed2a57b 100644 --- a/tests-integration/tests/serde_tests/kll.rs +++ b/tests-integration/tests/serde_tests/kll.rs @@ -33,7 +33,7 @@ use datasketches::codec::SketchBytes; use datasketches::codec::SketchSlice; use datasketches::error::Error; use datasketches::error::ErrorKind; -use datasketches::kll::KllComparator; +use datasketches::kll::KllFloat; use datasketches::kll::KllSketch; use datasketches::kll::KllValue; @@ -41,16 +41,34 @@ use crate::serialization_test_data; const DEFAULT_K: u16 = 200; -#[derive(Clone, Copy)] -struct NumericStringOrder; +#[derive(Debug, Clone, PartialEq, Eq)] +struct NumericString(String); -impl KllComparator for NumericStringOrder { - fn compare(&self, left: &String, right: &String) -> Ordering { - parse_string_value(left).cmp(&parse_string_value(right)) +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 is_compatible(&self, _other: &Self) -> bool { - true + fn deserialize(input: &mut SketchSlice<'_>) -> Result { + String::deserialize(input).map(Self) } } @@ -85,7 +103,7 @@ impl KllValue for Record { 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()); @@ -113,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() @@ -131,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()); @@ -159,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() @@ -230,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!( @@ -263,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() @@ -410,17 +425,17 @@ fn test_custom_kll_value_roundtrip() { #[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()); } From 043333d8dc3e0a0a769b8c2c78de189c591e08ca Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 14:46:35 +0800 Subject: [PATCH 19/22] refactor(codec): unify truncated input diagnostics --- CHANGELOG.md | 1 + datasketches/src/bloom/sketch.rs | 9 +-- datasketches/src/codec/assert.rs | 2 +- datasketches/src/codec/decode.rs | 58 +++++++++++++++++++ datasketches/src/countmin/sketch.rs | 9 +-- datasketches/src/cpc/sketch.rs | 5 +- datasketches/src/frequencies/serialization.rs | 37 +++++------- datasketches/src/frequencies/sketch.rs | 26 ++++----- datasketches/src/hll/array4.rs | 16 ++--- datasketches/src/hll/array6.rs | 9 +-- datasketches/src/hll/array8.rs | 9 +-- datasketches/src/hll/hash_set.rs | 23 +++----- datasketches/src/hll/list.rs | 17 +++--- datasketches/src/kll/sketch.rs | 31 ++++++---- datasketches/src/kll/value.rs | 19 +++--- datasketches/src/req/compactor.rs | 2 +- datasketches/src/tdigest/sketch.rs | 13 ++--- datasketches/src/thetafamily/theta/sketch.rs | 18 ++---- .../src/thetafamily/tuple/serialization.rs | 9 ++- datasketches/src/thetafamily/tuple/sketch.rs | 9 +-- .../tests/countmin_test/sketch.rs | 4 +- .../tests/serde_tests/frequencies.rs | 5 +- tests-integration/tests/serde_tests/kll.rs | 13 +++++ 23 files changed, 186 insertions(+), 158 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 67670801..67d4e2eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ All significant changes to this project will be documented in this file. ### Improvements * Improve KLL update and query performance. +* Improve truncated-input diagnostics across sketch deserializers. ## v0.5.0 diff --git a/datasketches/src/bloom/sketch.rs b/datasketches/src/bloom/sketch.rs index 6a119e91..747cc692 100644 --- a/datasketches/src/bloom/sketch.rs +++ b/datasketches/src/bloom/sketch.rs @@ -495,12 +495,9 @@ 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() - ))); - } + cursor + .ensure_remaining(payload_bytes) + .map_err(insufficient_data("Bloom filter payload"))?; } let mut bit_array = vec![0u64; num_words].into_boxed_slice(); let num_bits_set = if is_empty { 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/codec/decode.rs b/datasketches/src/codec/decode.rs index bce6e0b8..cabc0806 100644 --- a/datasketches/src/codec/decode.rs +++ b/datasketches/src/codec/decode.rs @@ -48,8 +48,36 @@ impl SketchSlice<'_> { &buf[pos..] } + /// Returns the number of not-yet-read bytes. + pub fn remaining_len(&self) -> usize { + self.remaining().len() + } + + /// Verifies that at least `expected` bytes remain without advancing the cursor. + pub fn ensure_remaining(&self, expected: usize) -> io::Result<()> { + let actual = self.remaining_len(); + if actual < expected { + Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("expected {expected} bytes, got {actual}"), + )) + } else { + Ok(()) + } + } + + /// Returns and consumes the next `len` bytes. + pub fn read_bytes(&mut self, len: usize) -> io::Result<&[u8]> { + self.ensure_remaining(len)?; + let start = self.slice.position() as usize; + let end = start + len; + self.slice.set_position(end as u64); + Ok(&self.slice.get_ref()[start..end]) + } + /// Reads exactly `buf.len()` bytes from the slice into `buf`. pub fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> { + self.ensure_remaining(buf.len())?; self.slice.read_exact(buf) } @@ -179,3 +207,33 @@ impl SketchSlice<'_> { Ok(f64::from_be_bytes(buf)) } } + +#[cfg(test)] +mod tests { + use std::io::ErrorKind; + + use super::SketchSlice; + + #[test] + fn insufficient_reads_report_expected_and_available_bytes_without_advancing() { + let mut input = SketchSlice::new(&[1, 2]); + + let error = input.read_u32_le().unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::UnexpectedEof); + assert_eq!(error.to_string(), "expected 4 bytes, got 2"); + assert_eq!(input.remaining(), &[1, 2]); + } + + #[test] + fn read_bytes_consumes_only_the_requested_bytes() { + let mut input = SketchSlice::new(&[1, 2, 3]); + + assert_eq!(input.read_bytes(2).unwrap(), &[1, 2]); + assert_eq!(input.remaining(), &[3]); + + let error = input.read_bytes(2).unwrap_err(); + assert_eq!(error.to_string(), "expected 2 bytes, got 1"); + assert_eq!(input.remaining(), &[3]); + } +} diff --git a/datasketches/src/countmin/sketch.rs b/datasketches/src/countmin/sketch.rs index 7970c0c9..b749499f 100644 --- a/datasketches/src/countmin/sketch.rs +++ b/datasketches/src/countmin/sketch.rs @@ -430,12 +430,9 @@ 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() - ))); - } + cursor + .ensure_remaining(payload_bytes) + .map_err(insufficient_data("CountMin payload"))?; } let mut sketch = Self::make(num_hashes, num_buckets, seed, expected_seed_hash, entries); diff --git a/datasketches/src/cpc/sketch.rs b/datasketches/src/cpc/sketch.rs index a9db2020..af0cd6eb 100644 --- a/datasketches/src/cpc/sketch.rs +++ b/datasketches/src/cpc/sketch.rs @@ -776,9 +776,8 @@ impl CpcSketch { .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"))?; + .read_bytes(payload_bytes) + .map_err(insufficient_data("CPC compressed payload"))?; 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..669e2859 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,16 @@ 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 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 len = cursor + .read_u32_le() + .map_err(insufficient_data("string item length"))? as usize; + let bytes = cursor + .read_bytes(len) + .map_err(insufficient_data("string item payload"))?; + + std::str::from_utf8(bytes) + .map(str::to_owned) + .map_err(|_| Error::deserial("invalid UTF-8 string payload")) } } @@ -87,11 +80,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..0052cafc 100644 --- a/datasketches/src/frequencies/sketch.rs +++ b/datasketches/src/frequencies/sketch.rs @@ -679,20 +679,17 @@ 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"))?; + cursor + .ensure_remaining(weight_bytes) + .map_err(insufficient_data("frequent item weights"))?; 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 +780,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..e798173d 100644 --- a/datasketches/src/hll/array4.rs +++ b/datasketches/src/hll/array4.rs @@ -356,12 +356,9 @@ 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() - ))); - } + cursor + .ensure_remaining(required_bytes) + .map_err(insufficient_data("HLL4 payload"))?; // Read packed 4-bit byte array let mut data = vec![0u8; num_bytes]; @@ -375,10 +372,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..c5e6be2f 100644 --- a/datasketches/src/hll/array6.rs +++ b/datasketches/src/hll/array6.rs @@ -203,12 +203,9 @@ 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() - ))); - } + cursor + .ensure_remaining(num_bytes) + .map_err(insufficient_data("HLL6 payload"))?; // Read packed byte array from offset HLL_BYTE_ARR_START let mut data = vec![0u8; num_bytes]; diff --git a/datasketches/src/hll/array8.rs b/datasketches/src/hll/array8.rs index 56756400..69e40279 100644 --- a/datasketches/src/hll/array8.rs +++ b/datasketches/src/hll/array8.rs @@ -275,12 +275,9 @@ 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() - ))); - } + cursor + .ensure_remaining(k) + .map_err(insufficient_data("HLL8 payload"))?; // Read byte array from offset HLL_BYTE_ARR_START let mut data = vec![0u8; k]; diff --git a/datasketches/src/hll/hash_set.rs b/datasketches/src/hll/hash_set.rs index 66b4714d..eca44498 100644 --- a/datasketches/src/hll/hash_set.rs +++ b/datasketches/src/hll/hash_set.rs @@ -111,22 +111,18 @@ 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() - ))); - } + cursor + .ensure_remaining(required_bytes) + .map_err(insufficient_data("HLL SET mode coupons"))?; if compact { // Compact mode: only couponCount coupons are stored // 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 +135,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..a1d70862 100644 --- a/datasketches/src/hll/list.rs +++ b/datasketches/src/hll/list.rs @@ -22,6 +22,7 @@ use crate::codec::SketchBytes; use crate::codec::SketchSlice; +use crate::codec::assert::insufficient_data; use crate::codec::family::Family; use crate::error::Error; use crate::hll::Coupon; @@ -99,21 +100,19 @@ 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 { + cursor + .ensure_remaining(required_bytes) + .map_err(insufficient_data("HLL LIST mode coupons"))?; } // 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/sketch.rs b/datasketches/src/kll/sketch.rs index fc9bf5e0..82a7b844 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -447,7 +447,7 @@ fn deserialize_with_serde(bytes: &[u8]) -> Result(bytes: &[u8]) -> Result(bytes: &[u8]) -> Result(bytes: &[u8]) -> Result { fn deserialize(input: &mut SketchSlice<'_>) -> Result { let value = input .read_f32_le() - .map_err(|_| Error::insufficient_data("f32"))?; + .map_err(insufficient_data("KLL f32 item"))?; Self::new(value).map_err(|_| Error::deserial("KLL float must not be NaN")) } } @@ -166,7 +167,7 @@ impl KllValue for KllFloat { fn deserialize(input: &mut SketchSlice<'_>) -> Result { let value = input .read_f64_le() - .map_err(|_| Error::insufficient_data("f64"))?; + .map_err(insufficient_data("KLL f64 item"))?; Self::new(value).map_err(|_| Error::deserial("KLL float must not be NaN")) } } @@ -185,7 +186,7 @@ impl KllValue for i64 { fn deserialize(input: &mut SketchSlice<'_>) -> Result { input .read_i64_le() - .map_err(|_| Error::insufficient_data("i64")) + .map_err(insufficient_data("KLL i64 item")) } } @@ -204,17 +205,13 @@ impl KllValue for String { fn deserialize(input: &mut SketchSlice<'_>) -> Result { let len = input .read_u32_le() - .map_err(|_| Error::insufficient_data("string_len"))? as usize; - let available = input.remaining().len(); - let bytes = input.remaining().get(..len).ok_or_else(|| { - Error::deserial(format!( - "insufficient string data: expected {len} bytes, got {available}" - )) - })?; + .map_err(insufficient_data("KLL string length"))? as usize; + let bytes = input + .read_bytes(len) + .map_err(insufficient_data("KLL string payload"))?; let value = std::str::from_utf8(bytes) .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/req/compactor.rs b/datasketches/src/req/compactor.rs index 8f9cc92d..2daaa38c 100644 --- a/datasketches/src/req/compactor.rs +++ b/datasketches/src/req/compactor.rs @@ -408,7 +408,7 @@ where // a multi-gigabyte reservation before the per-item reads below fail. The buffer // holds at most `remaining` more items (each item is ≥ 1 byte), so cap the // pre-allocation there; `push` still grows the Vec as the validated data needs. - let capacity = (num_items as usize).min(cursor.remaining().len()); + let capacity = (num_items as usize).min(cursor.remaining_len()); let mut items = Vec::with_capacity(capacity); for _ in 0..num_items { items.push(T::deserialize_value(cursor)?); diff --git a/datasketches/src/tdigest/sketch.rs b/datasketches/src/tdigest/sketch.rs index 942c7358..e5795bb7 100644 --- a/datasketches/src/tdigest/sketch.rs +++ b/datasketches/src/tdigest/sketch.rs @@ -673,16 +673,11 @@ 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 payload = cursor + .read_bytes(required_payload_bytes) + .map_err(insufficient_data("TDigest payload"))?; + 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..23290bed 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -739,12 +739,9 @@ 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() - ))); - } + cursor + .ensure_remaining(required_bytes) + .map_err(insufficient_data("Theta entries"))?; let mut entries = Vec::with_capacity(num_entries); for _ in 0..num_entries { let hash = cursor.read_u64_le().map_err(insufficient_data("entries"))?; @@ -969,12 +966,9 @@ 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() - ))); - } + cursor + .ensure_remaining(required_bytes) + .map_err(insufficient_data("Theta compressed entries"))?; // unpack blocks of BLOCK_WIDTH deltas let mut i = 0usize; 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..ef6ad630 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -716,12 +716,9 @@ 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() - ))); - } + cursor + .ensure_remaining(required_hash_bytes) + .map_err(insufficient_data("Tuple entry hashes"))?; let mut retained_entries = Vec::with_capacity(num_entries); for _ in 0..num_entries { let hash = cursor 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/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 aed2a57b..44ba22b1 100644 --- a/tests-integration/tests/serde_tests/kll.rs +++ b/tests-integration/tests/serde_tests/kll.rs @@ -439,3 +439,16 @@ fn test_rejects_truncated_or_trailing_data() { with_trailing_data.push(0); 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" + ); +} From b21c4cbfa79bdf9eac4f8478e282226d67f2d487 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 15:01:42 +0800 Subject: [PATCH 20/22] refactor(codec): keep byte helpers internal --- datasketches/src/codec/decode.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/datasketches/src/codec/decode.rs b/datasketches/src/codec/decode.rs index cabc0806..bcf24aa3 100644 --- a/datasketches/src/codec/decode.rs +++ b/datasketches/src/codec/decode.rs @@ -17,7 +17,6 @@ use std::io; use std::io::Cursor; -use std::io::Read; /// A wrapper around a byte slice that provides methods for reading various types of data from it. pub struct SketchSlice<'a> { @@ -49,12 +48,12 @@ impl SketchSlice<'_> { } /// Returns the number of not-yet-read bytes. - pub fn remaining_len(&self) -> usize { + pub(crate) fn remaining_len(&self) -> usize { self.remaining().len() } /// Verifies that at least `expected` bytes remain without advancing the cursor. - pub fn ensure_remaining(&self, expected: usize) -> io::Result<()> { + pub(crate) fn ensure_remaining(&self, expected: usize) -> io::Result<()> { let actual = self.remaining_len(); if actual < expected { Err(io::Error::new( @@ -67,7 +66,7 @@ impl SketchSlice<'_> { } /// Returns and consumes the next `len` bytes. - pub fn read_bytes(&mut self, len: usize) -> io::Result<&[u8]> { + pub(crate) fn read_bytes(&mut self, len: usize) -> io::Result<&[u8]> { self.ensure_remaining(len)?; let start = self.slice.position() as usize; let end = start + len; @@ -77,8 +76,8 @@ impl SketchSlice<'_> { /// Reads exactly `buf.len()` bytes from the slice into `buf`. pub fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> { - self.ensure_remaining(buf.len())?; - self.slice.read_exact(buf) + buf.copy_from_slice(self.read_bytes(buf.len())?); + Ok(()) } /// Reads a single byte from the slice and returns it as a `u8`. From b4b8302b3b15994762f8c859e08d2133de0568a5 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 15:08:30 +0800 Subject: [PATCH 21/22] refactor(codec): keep cursor abstraction minimal --- datasketches/src/bloom/sketch.rs | 10 ++- datasketches/src/codec/decode.rs | 61 +------------------ datasketches/src/countmin/sketch.rs | 10 ++- datasketches/src/cpc/sketch.rs | 11 +++- datasketches/src/frequencies/serialization.rs | 17 ++++-- datasketches/src/frequencies/sketch.rs | 10 ++- datasketches/src/hll/array4.rs | 10 ++- datasketches/src/hll/array6.rs | 10 ++- datasketches/src/hll/array8.rs | 10 ++- datasketches/src/hll/hash_set.rs | 10 ++- datasketches/src/hll/list.rs | 11 ++-- datasketches/src/kll/sketch.rs | 14 +++-- datasketches/src/kll/value.rs | 13 ++-- datasketches/src/req/compactor.rs | 2 +- datasketches/src/tdigest/sketch.rs | 11 +++- datasketches/src/thetafamily/theta/sketch.rs | 20 ++++-- datasketches/src/thetafamily/tuple/sketch.rs | 10 ++- 17 files changed, 125 insertions(+), 115 deletions(-) diff --git a/datasketches/src/bloom/sketch.rs b/datasketches/src/bloom/sketch.rs index 747cc692..8bf39220 100644 --- a/datasketches/src/bloom/sketch.rs +++ b/datasketches/src/bloom/sketch.rs @@ -495,9 +495,13 @@ impl BloomFilter { .checked_add(1) .and_then(|words| words.checked_mul(size_of::())) .ok_or_else(|| Error::deserial("Bloom filter payload length overflows"))?; - cursor - .ensure_remaining(payload_bytes) - .map_err(insufficient_data("Bloom filter payload"))?; + 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(); let num_bits_set = if is_empty { diff --git a/datasketches/src/codec/decode.rs b/datasketches/src/codec/decode.rs index bcf24aa3..bce6e0b8 100644 --- a/datasketches/src/codec/decode.rs +++ b/datasketches/src/codec/decode.rs @@ -17,6 +17,7 @@ use std::io; use std::io::Cursor; +use std::io::Read; /// A wrapper around a byte slice that provides methods for reading various types of data from it. pub struct SketchSlice<'a> { @@ -47,37 +48,9 @@ impl SketchSlice<'_> { &buf[pos..] } - /// Returns the number of not-yet-read bytes. - pub(crate) fn remaining_len(&self) -> usize { - self.remaining().len() - } - - /// Verifies that at least `expected` bytes remain without advancing the cursor. - pub(crate) fn ensure_remaining(&self, expected: usize) -> io::Result<()> { - let actual = self.remaining_len(); - if actual < expected { - Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - format!("expected {expected} bytes, got {actual}"), - )) - } else { - Ok(()) - } - } - - /// Returns and consumes the next `len` bytes. - pub(crate) fn read_bytes(&mut self, len: usize) -> io::Result<&[u8]> { - self.ensure_remaining(len)?; - let start = self.slice.position() as usize; - let end = start + len; - self.slice.set_position(end as u64); - Ok(&self.slice.get_ref()[start..end]) - } - /// Reads exactly `buf.len()` bytes from the slice into `buf`. pub fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> { - buf.copy_from_slice(self.read_bytes(buf.len())?); - Ok(()) + self.slice.read_exact(buf) } /// Reads a single byte from the slice and returns it as a `u8`. @@ -206,33 +179,3 @@ impl SketchSlice<'_> { Ok(f64::from_be_bytes(buf)) } } - -#[cfg(test)] -mod tests { - use std::io::ErrorKind; - - use super::SketchSlice; - - #[test] - fn insufficient_reads_report_expected_and_available_bytes_without_advancing() { - let mut input = SketchSlice::new(&[1, 2]); - - let error = input.read_u32_le().unwrap_err(); - - assert_eq!(error.kind(), ErrorKind::UnexpectedEof); - assert_eq!(error.to_string(), "expected 4 bytes, got 2"); - assert_eq!(input.remaining(), &[1, 2]); - } - - #[test] - fn read_bytes_consumes_only_the_requested_bytes() { - let mut input = SketchSlice::new(&[1, 2, 3]); - - assert_eq!(input.read_bytes(2).unwrap(), &[1, 2]); - assert_eq!(input.remaining(), &[3]); - - let error = input.read_bytes(2).unwrap_err(); - assert_eq!(error.to_string(), "expected 2 bytes, got 1"); - assert_eq!(input.remaining(), &[3]); - } -} diff --git a/datasketches/src/countmin/sketch.rs b/datasketches/src/countmin/sketch.rs index b749499f..500f4396 100644 --- a/datasketches/src/countmin/sketch.rs +++ b/datasketches/src/countmin/sketch.rs @@ -430,9 +430,13 @@ impl CountMinSketch { let payload_bytes = payload_values .checked_mul(LONG_SIZE_BYTES) .ok_or_else(|| Error::deserial("CountMin payload size overflows"))?; - cursor - .ensure_remaining(payload_bytes) - .map_err(insufficient_data("CountMin payload"))?; + 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}"), + )); + } } let mut sketch = Self::make(num_hashes, num_buckets, seed, expected_seed_hash, entries); diff --git a/datasketches/src/cpc/sketch.rs b/datasketches/src/cpc/sketch.rs index af0cd6eb..2cc4a7a6 100644 --- a/datasketches/src/cpc/sketch.rs +++ b/datasketches/src/cpc/sketch.rs @@ -775,9 +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 - .read_bytes(payload_bytes) - .map_err(insufficient_data("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 669e2859..d5a8b615 100644 --- a/datasketches/src/frequencies/serialization.rs +++ b/datasketches/src/frequencies/serialization.rs @@ -58,13 +58,18 @@ impl FrequentItemValue for String { let len = cursor .read_u32_le() .map_err(insufficient_data("string item length"))? as usize; - let bytes = cursor - .read_bytes(len) - .map_err(insufficient_data("string item payload"))?; - - std::str::from_utf8(bytes) + 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 value = std::str::from_utf8(&cursor.remaining()[..len]) .map(str::to_owned) - .map_err(|_| Error::deserial("invalid UTF-8 string payload")) + .map_err(|_| Error::deserial("invalid UTF-8 string payload"))?; + cursor.advance(len as u64); + Ok(value) } } diff --git a/datasketches/src/frequencies/sketch.rs b/datasketches/src/frequencies/sketch.rs index 0052cafc..2b163f27 100644 --- a/datasketches/src/frequencies/sketch.rs +++ b/datasketches/src/frequencies/sketch.rs @@ -682,9 +682,13 @@ impl FrequentItemsSketch { let weight_bytes = active_items .checked_mul(size_of::()) .ok_or_else(|| Error::deserial("frequent item weight payload length overflows"))?; - cursor - .ensure_remaining(weight_bytes) - .map_err(insufficient_data("frequent item weights"))?; + 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 { diff --git a/datasketches/src/hll/array4.rs b/datasketches/src/hll/array4.rs index e798173d..2380e001 100644 --- a/datasketches/src/hll/array4.rs +++ b/datasketches/src/hll/array4.rs @@ -356,9 +356,13 @@ 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"))?; - cursor - .ensure_remaining(required_bytes) - .map_err(insufficient_data("HLL4 payload"))?; + 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 let mut data = vec![0u8; num_bytes]; diff --git a/datasketches/src/hll/array6.rs b/datasketches/src/hll/array6.rs index c5e6be2f..3654bb80 100644 --- a/datasketches/src/hll/array6.rs +++ b/datasketches/src/hll/array6.rs @@ -203,9 +203,13 @@ impl Array6 { "HLL6 zero count must not exceed k and auxiliary count must be zero", )); } - cursor - .ensure_remaining(num_bytes) - .map_err(insufficient_data("HLL6 payload"))?; + 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 let mut data = vec![0u8; num_bytes]; diff --git a/datasketches/src/hll/array8.rs b/datasketches/src/hll/array8.rs index 69e40279..e28c7241 100644 --- a/datasketches/src/hll/array8.rs +++ b/datasketches/src/hll/array8.rs @@ -275,9 +275,13 @@ impl Array8 { "HLL8 zero count must not exceed k and auxiliary count must be zero", )); } - cursor - .ensure_remaining(k) - .map_err(insufficient_data("HLL8 payload"))?; + 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 let mut data = vec![0u8; k]; diff --git a/datasketches/src/hll/hash_set.rs b/datasketches/src/hll/hash_set.rs index eca44498..dd2f5be5 100644 --- a/datasketches/src/hll/hash_set.rs +++ b/datasketches/src/hll/hash_set.rs @@ -111,9 +111,13 @@ impl HashSet { } let read_count = if compact { coupon_count } else { array_size }; let required_bytes = read_count * size_of::(); - cursor - .ensure_remaining(required_bytes) - .map_err(insufficient_data("HLL SET mode coupons"))?; + 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 { // Compact mode: only couponCount coupons are stored diff --git a/datasketches/src/hll/list.rs b/datasketches/src/hll/list.rs index a1d70862..47370fa5 100644 --- a/datasketches/src/hll/list.rs +++ b/datasketches/src/hll/list.rs @@ -22,7 +22,6 @@ use crate::codec::SketchBytes; use crate::codec::SketchSlice; -use crate::codec::assert::insufficient_data; use crate::codec::family::Family; use crate::error::Error; use crate::hll::Coupon; @@ -101,9 +100,13 @@ impl List { let read_count = if compact { coupon_count } else { array_size }; let required_bytes = read_count * size_of::(); if !empty { - cursor - .ensure_remaining(required_bytes) - .map_err(insufficient_data("HLL LIST mode coupons"))?; + 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. diff --git a/datasketches/src/kll/sketch.rs b/datasketches/src/kll/sketch.rs index 82a7b844..82399cdb 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -447,7 +447,7 @@ fn deserialize_with_serde(bytes: &[u8]) -> Result(bytes: &[u8]) -> Result(bytes: &[u8]) -> Result()) .ok_or_else(|| Error::deserial("Theta entry payload length overflows"))?; - cursor - .ensure_remaining(required_bytes) - .map_err(insufficient_data("Theta entries"))?; + 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 { let hash = cursor.read_u64_le().map_err(insufficient_data("entries"))?; @@ -966,9 +970,13 @@ impl CompactThetaSketch { .and_then(|bits| bits.checked_add(7)) .map(|bits| bits / 8) .ok_or_else(|| Error::deserial("Theta compressed payload length overflows"))?; - cursor - .ensure_remaining(required_bytes) - .map_err(insufficient_data("Theta compressed entries"))?; + 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 let mut i = 0usize; diff --git a/datasketches/src/thetafamily/tuple/sketch.rs b/datasketches/src/thetafamily/tuple/sketch.rs index ef6ad630..2d6b93f4 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -716,9 +716,13 @@ impl CompactTupleSketch { let required_hash_bytes = num_entries .checked_mul(size_of::()) .ok_or_else(|| Error::deserial("Tuple entry payload length overflows"))?; - cursor - .ensure_remaining(required_hash_bytes) - .map_err(insufficient_data("Tuple entry hashes"))?; + 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 { let hash = cursor From 40e2fd0b9b0a02aeaf479d6e7e29005f4a74291b Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 2 Sep 2026 15:30:23 +0800 Subject: [PATCH 22/22] Apply batched suggestions from code review Co-authored-by: tison --- CHANGELOG.md | 1 - 1 file changed, 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 67d4e2eb..42666317 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,7 +14,6 @@ All significant changes to this project will be documented in this file. ### Improvements -* Improve KLL update and query performance. * Improve truncated-input diagnostics across sketch deserializers. ## v0.5.0