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 c93611448b4ce0499fae178703018af9a9ec0bc7 Author: tison <[email protected]> AuthorDate: Wed Sep 2 12:12:38 2026 +0800 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 ae9045f..9dcecb8 100644 --- a/datasketches/src/kll/sketch.rs +++ b/datasketches/src/kll/sketch.rs @@ -184,23 +184,33 @@ impl<T: Clone, C: KllComparator<T>> KllSketch<T, C> { /// Merges another sketch into this one. /// - /// # Panics + /// # Errors /// - /// Panics if the sketches have incompatible parameters. - pub fn merge(&mut self, other: &KllSketch<T, C>) { + /// Returns an error if the sketches use incompatible comparators or their combined stream + /// weight exceeds [`u64::MAX`]. + pub fn merge(&mut self, other: &KllSketch<T, C>) -> 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<T: Clone, C: KllComparator<T>> KllSketch<T, C> { } 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 afc5090..a52102f 100644 --- a/tests-integration/tests/kll_test/sketch.rs +++ b/tests-integration/tests/kll_test/sketch.rs @@ -54,6 +54,25 @@ impl KllComparator<String> for NumericStringOrder { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct DirectionalOrder { + descending: bool, +} + +impl KllComparator<i64> 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::<f32>::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::<f32>::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::<f32>::new(DEFAULT_K).unwrap(); @@ -405,7 +441,7 @@ fn test_merge_min_max_large_other() { sketch1.update(i as f32); } let mut sketch2 = KllSketch::<f32>::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)); } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
