This is an automated email from the ASF dual-hosted git repository. tisonkun pushed a commit to branch codex/kll-api-and-performance in repository https://gitbox.apache.org/repos/asf/datasketches-rust.git
commit 20e0707c73bdec60dbf73e3ba0d97d10b5dc2a65 Author: tison <[email protected]> AuthorDate: Wed Sep 2 12:16:19 2026 +0800 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 8b8e563..e65f375 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 9dcecb8..bae05f4 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<T, C = NaturalOrder> { m: u8, min_k: u16, n: u64, + num_retained: usize, + capacity: usize, is_level_zero_sorted: bool, levels: Vec<Vec<T>>, min_item: Option<T>, @@ -142,7 +145,7 @@ impl<T: Clone, C: KllComparator<T>> KllSketch<T, C> { /// 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { "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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { "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<Vec<T>, Error> { + self.sorted_view().quantiles(ranks, criteria) } /// Returns the approximate CDF for the given split points. @@ -275,8 +302,7 @@ impl<T: Clone, C: KllComparator<T>> KllSketch<T, C> { 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { 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<T, C> { + 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { max_item: Option<T>, 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { } } - fn capacity(&self) -> usize { - total_capacity(self.k, self.m, self.levels.len()) as usize - } - fn level_offsets(&self) -> Vec<u32> { - 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { } 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { 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::<bool>(), use_up); + let (leftover, promoted) = compact_level( + current, + level, + self.is_level_zero_sorted, + &self.comparator, + rand::random::<bool>(), + use_up, + ); if above.is_empty() { above = promoted; } else { @@ -749,6 +778,7 @@ impl<T: Clone, C: KllComparator<T>> KllSketch<T, C> { 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { 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<T: Clone>(items: Vec<T>, offset: bool, use_up: bool) -> Vec<T> { +fn compact_level<T, C: KllComparator<T>>( + mut items: Vec<T>, + level: usize, + is_level_zero_sorted: bool, + comparator: &C, + offset: bool, + use_up: bool, +) -> (Option<T>, Vec<T>) { + 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<T, I: ExactSizeIterator<Item = T>>(items: I, offset: bool, use_up: bool) -> Vec<T> { let len = items.len(); debug_assert!(len % 2 == 0, "length must be even"); let offset = usize::from(offset); @@ -898,20 +963,11 @@ fn downsample<T: Clone>(items: Vec<T>, offset: bool, use_up: bool) -> Vec<T> { }; items - .into_iter() .enumerate() .filter_map(|(idx, item)| if idx % 2 == parity { Some(item) } else { None }) .collect() } -fn take_leftover<T>(items: &mut Vec<T>, 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<T: Clone, C: KllComparator<T>>( left: Vec<T>, right: Vec<T>, @@ -957,25 +1013,17 @@ fn general_compress<T: Clone, C: KllComparator<T>>( 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::<bool>(), use_up); + let (leftover, promoted) = compact_level( + current, + current_level, + is_level_zero_sorted, + comparator, + rand::random::<bool>(), + 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 5c8cd5b..2337a6c 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<T: Clone, C: KllComparator<T>> { comparator: C, @@ -31,16 +35,15 @@ pub struct SortedView<T: Clone, C: KllComparator<T>> { #[derive(Debug, Clone)] struct Entry<T> { item: T, - weight: u64, + cumulative_weight: u64, } impl<T: Clone, C: KllComparator<T>> SortedView<T, C> { - fn new(mut entries: Vec<Entry<T>>, comparator: C) -> Self { - entries.sort_by(|a, b| comparator.compare(&a.item, &b.item)); + fn from_sorted(mut entries: Vec<Entry<T>>, 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<T: Clone, C: KllComparator<T>> SortedView<T, C> { } } - 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<f64, Error> { + 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<T, Error> { + 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<Vec<T>, 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<Vec<f64>, 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<Vec<f64>, 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<T: Clone, C: KllComparator<T>> SortedView<T, C> { pub fn build_sorted_view<T: Clone, C: KllComparator<T>>( levels: &[Vec<T>], + is_level_zero_sorted: bool, comparator: C, ) -> SortedView<T, C> { - 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<T, C: KllComparator<T>>( + left: Vec<Entry<T>>, + right: Vec<Entry<T>>, + comparator: &C, +) -> Vec<Entry<T>> { + 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<T, C: KllComparator<T>>( @@ -145,57 +251,17 @@ fn check_split_points<T, C: KllComparator<T>>( } fn lower_bound<T, C: KllComparator<T>>(entries: &[Entry<T>], 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<T, C: KllComparator<T>>(entries: &[Entry<T>], 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<T>(entries: &[Entry<T>], 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<T>(entries: &[Entry<T>], 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 a52102f..5fcd068 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::<f32>::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::<f32>::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::<f32>::new(DEFAULT_K).unwrap(); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
