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 17e8b77114ba4f55dfd1d52cf27fe9652164f576 Author: tison <[email protected]> AuthorDate: Wed Sep 2 12:05:49 2026 +0800 refactor(kll): clarify capacity module boundaries --- datasketches/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(-) 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 6b993aa..03057aa 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 47cfd23..d87da0e 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 3ba4a79..41e585c 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 8248fe2..ec61190 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<T: KllSerde, C: KllComparator<T>>( ))); } - 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<T: KllItem, C: KllComparator<T>> KllSketch<T, C> { } 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<u32> { @@ -847,8 +846,11 @@ impl<T: KllItem, C: KllComparator<T>> KllSketch<T, C> { } fn total_weight(&self) -> u64 { - let sizes: Vec<usize> = 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<T: KllItem, C: KllComparator<T>>( ) -> Vec<Vec<T>> { 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 bcaad64..075a472 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<T: KllItem, C: KllComparator<T>> { +pub struct SortedView<T: KllItem, C: KllComparator<T>> { comparator: C, entries: Vec<Entry<T>>, total_weight: u64, @@ -48,7 +48,7 @@ impl<T: KllItem, C: KllComparator<T>> SortedView<T, C> { } } - 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<T: KllItem, C: KllComparator<T>> SortedView<T, C> { 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<T: KllItem, C: KllComparator<T>> SortedView<T, C> { self.entries[idx].item.clone() } - pub(super) fn cdf(&self, split_points: &[T], inclusive: bool) -> Vec<f64> { + pub fn cdf(&self, split_points: &[T], inclusive: bool) -> Vec<f64> { 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<T: KllItem, C: KllComparator<T>> SortedView<T, C> { ranks } - pub(super) fn pmf(&self, split_points: &[T], inclusive: bool) -> Vec<f64> { + pub fn pmf(&self, split_points: &[T], inclusive: bool) -> Vec<f64> { 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<T: KllItem, C: KllComparator<T>> SortedView<T, C> { } } -pub(super) fn build_sorted_view<T: KllItem, C: KllComparator<T>>( +pub fn build_sorted_view<T: KllItem, C: KllComparator<T>>( levels: &[Vec<T>], comparator: C, ) -> SortedView<T, C> { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
