diff --git a/datasketches/src/theta/jaccard_similarity.rs b/datasketches/src/theta/jaccard_similarity.rs new file mode 100644 index 0000000..db4d86b --- /dev/null +++ b/datasketches/src/theta/jaccard_similarity.rs @@ -0,0 +1,74 @@ +// 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. + +//! Jaccard similarity for Theta sketches. + +use crate::error::Error; +use crate::hash::DEFAULT_UPDATE_SEED; +use crate::theta::ThetaSketchView; +use crate::thetacommon::jaccard_similarity::RawThetaJaccardSimilarity; + +/// Jaccard similarity result for two Theta sketches. +/// +/// The bounds use a 95.4% confidence interval, equivalent to +/- 2 standard deviations. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct JaccardSimilarity { + lower_bound: f64, + estimate: f64, + upper_bound: f64, +} + +impl JaccardSimilarity { + /// Computes the Jaccard similarity index with the default update seed. + pub fn between( + sketch_a: &A, + sketch_b: &B, + ) -> Result { + Self::between_with_seed(sketch_a, sketch_b, DEFAULT_UPDATE_SEED) + } + + /// Computes the Jaccard similarity index with an explicit update seed. + /// + /// Returns an error if a non-empty sketch was built with a different seed. + pub fn between_with_seed( + sketch_a: &A, + sketch_b: &B, + seed: u64, + ) -> Result { + let raw = RawThetaJaccardSimilarity::compute(sketch_a, sketch_b, seed)?; + Ok(Self { + lower_bound: raw.lower_bound, + estimate: raw.estimate, + upper_bound: raw.upper_bound, + }) + } + + /// Returns the approximate lower bound for the Jaccard index. + pub fn lower_bound(&self) -> f64 { + self.lower_bound + } + + /// Returns the estimate of the Jaccard index. + pub fn estimate(&self) -> f64 { + self.estimate + } + + /// Returns the approximate upper bound for the Jaccard index. + pub fn upper_bound(&self) -> f64 { + self.upper_bound + } +} diff --git a/datasketches/src/theta/mod.rs b/datasketches/src/theta/mod.rs index 4b5d28b..96d648a 100644 --- a/datasketches/src/theta/mod.rs +++ b/datasketches/src/theta/mod.rs @@ -43,6 +43,7 @@ mod a_not_b; mod bit_pack; mod hash_table; mod intersection; +mod jaccard_similarity; mod serialization; mod sketch; mod union; @@ -50,6 +51,7 @@ mod union; pub use self::a_not_b::ThetaANotB; pub use self::hash_table::ThetaEntry; pub use self::intersection::ThetaIntersection; +pub use self::jaccard_similarity::JaccardSimilarity; pub use self::sketch::CompactThetaSketch; pub use self::sketch::ThetaSketch; pub use self::sketch::ThetaSketchBuilder; diff --git a/datasketches/src/thetacommon/bounds_binomial_proportions.rs b/datasketches/src/thetacommon/bounds_binomial_proportions.rs new file mode 100644 index 0000000..ba603a0 --- /dev/null +++ b/datasketches/src/thetacommon/bounds_binomial_proportions.rs @@ -0,0 +1,161 @@ +// 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::error::Error; + +/// Computes an approximate lower bound for an unknown binomial proportion. +pub(crate) fn approximate_lower_bound_on_p( + n: u64, + k: u64, + num_std_devs: f64, +) -> Result { + check_inputs(n, k)?; + if n == 0 || k == 0 { + Ok(0.0) + } else if k == 1 { + Ok(exact_lower_bound_on_p_k_eq_1( + n, + delta_of_num_stdevs(num_std_devs), + )) + } else if k == n { + Ok(exact_lower_bound_on_p_k_eq_n( + n, + delta_of_num_stdevs(num_std_devs), + )) + } else { + let x = abramowitz_stegun_formula_26p5p22((n - k) as f64 + 1.0, k as f64, -num_std_devs); + Ok(1.0 - x) + } +} + +/// Computes an approximate upper bound for an unknown binomial proportion. +pub(crate) fn approximate_upper_bound_on_p( + n: u64, + k: u64, + num_std_devs: f64, +) -> Result { + check_inputs(n, k)?; + if n == 0 || k == n { + Ok(1.0) + } else if k == n - 1 { + Ok(exact_upper_bound_on_p_k_eq_minusone( + n, + delta_of_num_stdevs(num_std_devs), + )) + } else if k == 0 { + Ok(exact_upper_bound_on_p_k_eq_zero( + n, + delta_of_num_stdevs(num_std_devs), + )) + } else { + let x = abramowitz_stegun_formula_26p5p22((n - k) as f64, k as f64 + 1.0, num_std_devs); + Ok(1.0 - x) + } +} + +fn check_inputs(n: u64, k: u64) -> Result<(), Error> { + if k > n { + return Err(Error::invalid_argument(format!( + "k cannot exceed n: k={k}, n={n}" + ))); + } + Ok(()) +} + +fn delta_of_num_stdevs(kappa: f64) -> f64 { + normal_cdf(-kappa) +} + +fn normal_cdf(x: f64) -> f64 { + 0.5 * (1.0 + erf(x / 2.0_f64.sqrt())) +} + +fn erf(x: f64) -> f64 { + if x < 0.0 { + -erf_of_nonneg(-x) + } else { + erf_of_nonneg(x) + } +} + +fn erf_of_nonneg(x: f64) -> f64 { + let a1 = 0.0705230784; + let a2 = 0.0422820123; + let a3 = 0.0092705272; + let a4 = 0.0001520143; + let a5 = 0.0002765672; + let a6 = 0.0000430638; + let x2 = x * x; + let x3 = x2 * x; + let x4 = x2 * x2; + let x5 = x2 * x3; + let x6 = x3 * x3; + let sum = 1.0 + (a1 * x) + (a2 * x2) + (a3 * x3) + (a4 * x4) + (a5 * x5) + (a6 * x6); + let sum2 = sum * sum; + let sum4 = sum2 * sum2; + let sum8 = sum4 * sum4; + let sum16 = sum8 * sum8; + 1.0 - (1.0 / sum16) +} + +fn abramowitz_stegun_formula_26p5p22(a: f64, b: f64, yp: f64) -> f64 { + let b2m1 = (2.0 * b) - 1.0; + let a2m1 = (2.0 * a) - 1.0; + let lambda = ((yp * yp) - 3.0) / 6.0; + let htmp = (1.0 / a2m1) + (1.0 / b2m1); + let h = 2.0 / htmp; + let term1 = (yp * (h + lambda).sqrt()) / h; + let term2 = (1.0 / b2m1) - (1.0 / a2m1); + let term3 = (lambda + (5.0 / 6.0)) - (2.0 / (3.0 * h)); + let w = term1 - (term2 * term3); + a / (a + (b * (2.0 * w).exp())) +} + +fn exact_upper_bound_on_p_k_eq_zero(n: u64, delta: f64) -> f64 { + 1.0 - delta.powf(1.0 / n as f64) +} + +fn exact_lower_bound_on_p_k_eq_n(n: u64, delta: f64) -> f64 { + delta.powf(1.0 / n as f64) +} + +fn exact_lower_bound_on_p_k_eq_1(n: u64, delta: f64) -> f64 { + 1.0 - (1.0 - delta).powf(1.0 / n as f64) +} + +fn exact_upper_bound_on_p_k_eq_minusone(n: u64, delta: f64) -> f64 { + (1.0 - delta).powf(1.0 / n as f64) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_invalid_counts() { + assert!(approximate_lower_bound_on_p(1, 2, 2.0).is_err()); + assert!(approximate_upper_bound_on_p(1, 2, 2.0).is_err()); + } + + #[test] + fn computes_exact_edge_cases() { + assert_eq!(approximate_lower_bound_on_p(0, 0, 2.0).unwrap(), 0.0); + assert_eq!(approximate_upper_bound_on_p(0, 0, 2.0).unwrap(), 1.0); + assert_eq!(approximate_lower_bound_on_p(10, 0, 2.0).unwrap(), 0.0); + assert_eq!(approximate_upper_bound_on_p(10, 10, 2.0).unwrap(), 1.0); + } +} diff --git a/datasketches/src/thetacommon/jaccard_similarity.rs b/datasketches/src/thetacommon/jaccard_similarity.rs new file mode 100644 index 0000000..12e4929 --- /dev/null +++ b/datasketches/src/thetacommon/jaccard_similarity.rs @@ -0,0 +1,202 @@ +// 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::common::ResizeFactor; +use crate::error::Error; +use crate::thetacommon::RetainedEntry; +use crate::thetacommon::ThetaFamilySketchView; +use crate::thetacommon::bounds_binomial_proportions; +use crate::thetacommon::constants::MAX_LG_K; +use crate::thetacommon::constants::MAX_THETA; +use crate::thetacommon::constants::MIN_LG_K; +use crate::thetacommon::intersection::IntersectionMergePolicy; +use crate::thetacommon::intersection::IntersectionState; +use crate::thetacommon::union::UnionMergePolicy; +use crate::thetacommon::union::UnionState; + +const NUM_STD_DEVS: f64 = 2.0; + +#[derive(Clone, Copy, Debug, PartialEq)] +pub(crate) struct RawThetaJaccardSimilarity { + pub(crate) lower_bound: f64, + pub(crate) estimate: f64, + pub(crate) upper_bound: f64, +} + +#[derive(Debug)] +struct NoopMergePolicy; + +impl UnionMergePolicy for NoopMergePolicy { + fn merge(&self, _existing: &mut E, _incoming: E) {} +} + +impl IntersectionMergePolicy for NoopMergePolicy { + fn merge(&self, _existing: &mut E, _incoming: E) {} +} + +struct CompactSketchView { + entries: Vec, + theta: u64, + seed_hash: u16, + ordered: bool, + empty: bool, +} + +impl ThetaFamilySketchView for CompactSketchView { + type Entry = E; + + fn seed_hash(&self) -> u16 { + self.seed_hash + } + + fn theta64(&self) -> u64 { + self.theta + } + + fn is_empty(&self) -> bool { + self.empty + } + + fn is_ordered(&self) -> bool { + self.ordered + } + + fn iter(&self) -> impl Iterator + '_ { + self.entries.iter().cloned() + } + + fn num_retained(&self) -> usize { + self.entries.len() + } +} + +impl RawThetaJaccardSimilarity { + pub(crate) fn compute(sketch_a: &A, sketch_b: &B, seed: u64) -> Result + where + A: ThetaFamilySketchView, + B: ThetaFamilySketchView, + A::Entry: Clone, + { + if sketch_a.is_empty() && sketch_b.is_empty() { + return Ok(Self::exact(1.0)); + } + if sketch_a.is_empty() || sketch_b.is_empty() { + return Ok(Self::exact(0.0)); + } + + let mut union = UnionState::new( + union_lg_k(sketch_a.num_retained(), sketch_b.num_retained()), + ResizeFactor::X8, + 1.0, + seed, + NoopMergePolicy, + ); + union.update(sketch_a)?; + union.update(sketch_b)?; + let union = union.to_compact_parts(false); + + if union.entries.len() == sketch_a.num_retained() + && union.entries.len() == sketch_b.num_retained() + && union.theta == sketch_a.theta64() + && union.theta == sketch_b.theta64() + { + return Ok(Self::exact(1.0)); + } + + let union = CompactSketchView { + entries: union.entries, + theta: union.theta, + seed_hash: union.seed_hash, + ordered: union.ordered, + empty: union.empty, + }; + let mut intersection = IntersectionState::new(seed, NoopMergePolicy); + intersection.update(sketch_a)?; + intersection.update(sketch_b)?; + intersection.update(&union)?; + let intersection = intersection.result(false); + + Self::ratio_bounds( + union.num_retained() as u64, + intersection.entries.len() as u64, + union.theta64(), + ) + } + + fn exact(value: f64) -> Self { + Self { + lower_bound: value, + estimate: value, + upper_bound: value, + } + } + + fn ratio_bounds(union_count: u64, intersection_count: u64, theta: u64) -> Result { + if intersection_count > union_count { + return Err(Error::invalid_argument(format!( + "intersection count cannot exceed union count: {intersection_count} > {union_count}" + ))); + } + if union_count == 0 { + return Ok(Self { + lower_bound: 0.0, + estimate: 0.5, + upper_bound: 1.0, + }); + } + + let sampling_probability = theta as f64 / MAX_THETA as f64; + if sampling_probability <= 0.0 || sampling_probability > 1.0 { + return Err(Error::invalid_argument(format!( + "theta must produce a probability in (0.0, 1.0], got {sampling_probability}" + ))); + } + if sampling_probability == 1.0 { + return Ok(Self::exact(intersection_count as f64 / union_count as f64)); + } + + let adjustment = NUM_STD_DEVS * sampling_adjuster(sampling_probability); + Ok(Self { + lower_bound: bounds_binomial_proportions::approximate_lower_bound_on_p( + union_count, + intersection_count, + adjustment, + )?, + estimate: intersection_count as f64 / union_count as f64, + upper_bound: bounds_binomial_proportions::approximate_upper_bound_on_p( + union_count, + intersection_count, + adjustment, + )?, + }) + } +} + +fn sampling_adjuster(sampling_probability: f64) -> f64 { + let adjustment = (1.0 - sampling_probability).sqrt(); + if sampling_probability <= 0.5 { + adjustment + } else { + adjustment + (0.01 * (sampling_probability - 0.5)) + } +} + +fn union_lg_k(left_count: usize, right_count: usize) -> u8 { + let required_capacity = left_count.saturating_add(right_count).max(1); + let lg_k = usize::BITS - (required_capacity - 1).leading_zeros(); + (lg_k as u8).clamp(MIN_LG_K, MAX_LG_K) +} diff --git a/datasketches/src/thetacommon/mod.rs b/datasketches/src/thetacommon/mod.rs index eb8db95..6b666f3 100644 --- a/datasketches/src/thetacommon/mod.rs +++ b/datasketches/src/thetacommon/mod.rs @@ -19,9 +19,11 @@ pub(crate) mod a_not_b; pub(crate) mod binomial_bounds; +pub(crate) mod bounds_binomial_proportions; pub(crate) mod constants; pub(crate) mod hash_table; pub(crate) mod intersection; +pub(crate) mod jaccard_similarity; pub(crate) mod union; /// An entry retained by a Theta sketch family hash table. diff --git a/datasketches/tests/theta_test/jaccard_similarity.rs b/datasketches/tests/theta_test/jaccard_similarity.rs new file mode 100644 index 0000000..6e300ab --- /dev/null +++ b/datasketches/tests/theta_test/jaccard_similarity.rs @@ -0,0 +1,145 @@ +// 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. + +#![cfg(feature = "theta")] + +use datasketches::theta::JaccardSimilarity; +use datasketches::theta::ThetaSketch; +use datasketches::theta::ThetaSketchBuilder; + +fn assert_jaccard_exact(actual: datasketches::theta::JaccardSimilarity, expected: f64) { + assert_eq!(actual.lower_bound(), expected); + assert_eq!(actual.estimate(), expected); + assert_eq!(actual.upper_bound(), expected); +} + +fn assert_close(actual: f64, expected: f64, margin: f64) { + assert!( + (actual - expected).abs() <= margin, + "actual={actual}, expected={expected}, margin={margin}" + ); +} + +fn assert_jaccard_estimate(actual: JaccardSimilarity, expected: f64) { + assert_close(actual.estimate(), expected, 0.01); + assert!(actual.lower_bound() <= actual.estimate()); + assert!(actual.estimate() <= actual.upper_bound()); +} + +fn sketch_with_range(start: u64, count: u64) -> ThetaSketch { + let mut sketch = ThetaSketchBuilder::default().build(); + for value in start..start + count { + sketch.update(value); + } + sketch +} + +fn sketch_with_range_and_seed(start: u64, count: u64, seed: u64) -> ThetaSketch { + let mut sketch = ThetaSketchBuilder::default().seed(seed).build(); + for value in start..start + count { + sketch.update(value); + } + sketch +} + +#[test] +fn test_empty() { + let sketch_a = ThetaSketchBuilder::default().build(); + let sketch_b = ThetaSketchBuilder::default().build(); + + let jaccard = JaccardSimilarity::between(&sketch_a, &sketch_b).unwrap(); + + assert_jaccard_exact(jaccard, 1.0); +} + +#[test] +fn test_same_sketch_exact_mode() { + let sketch = sketch_with_range(0, 1000); + + let jaccard = JaccardSimilarity::between(&sketch, &sketch).unwrap(); + assert_jaccard_exact(jaccard, 1.0); + + let jaccard = JaccardSimilarity::between(&sketch.compact(true), &sketch.compact(true)).unwrap(); + assert_jaccard_exact(jaccard, 1.0); +} + +#[test] +fn test_full_overlap_exact_mode() { + let sketch_a = sketch_with_range(0, 1000); + let sketch_b = sketch_with_range(0, 1000); + + let jaccard = JaccardSimilarity::between(&sketch_a, &sketch_b).unwrap(); + assert_jaccard_exact(jaccard, 1.0); + + let jaccard = + JaccardSimilarity::between(&sketch_a.compact(true), &sketch_b.compact(true)).unwrap(); + assert_jaccard_exact(jaccard, 1.0); +} + +#[test] +fn test_disjoint_exact_mode() { + let sketch_a = sketch_with_range(0, 1000); + let sketch_b = sketch_with_range(1000, 1000); + + let jaccard = JaccardSimilarity::between(&sketch_a, &sketch_b).unwrap(); + assert_jaccard_exact(jaccard, 0.0); + + let jaccard = + JaccardSimilarity::between(&sketch_a.compact(true), &sketch_b.compact(true)).unwrap(); + assert_jaccard_exact(jaccard, 0.0); +} + +#[test] +fn test_half_overlap_estimation_mode() { + let sketch_a = sketch_with_range(0, 10000); + let sketch_b = sketch_with_range(5000, 10000); + + let jaccard = JaccardSimilarity::between(&sketch_a, &sketch_b).unwrap(); + assert_jaccard_estimate(jaccard, 0.33); + + let jaccard = + JaccardSimilarity::between(&sketch_a.compact(true), &sketch_b.compact(true)).unwrap(); + assert_jaccard_estimate(jaccard, 0.33); +} + +#[test] +fn test_half_overlap_estimation_mode_custom_seed() { + let seed = 123; + let sketch_a = sketch_with_range_and_seed(0, 10000, seed); + let sketch_b = sketch_with_range_and_seed(5000, 10000, seed); + + let jaccard = JaccardSimilarity::between_with_seed(&sketch_a, &sketch_b, seed).unwrap(); + assert_jaccard_estimate(jaccard, 0.33); + + let jaccard = JaccardSimilarity::between_with_seed( + &sketch_a.compact(true), + &sketch_b.compact(true), + seed, + ) + .unwrap(); + assert_jaccard_estimate(jaccard, 0.33); +} + +#[test] +fn test_seed_mismatch() { + let mut sketch_a = ThetaSketchBuilder::default().build(); + sketch_a.update(1u64); + let mut sketch_b = ThetaSketchBuilder::default().seed(123).build(); + sketch_b.update(1u64); + + assert!(JaccardSimilarity::between(&sketch_a, &sketch_b).is_err()); +} diff --git a/datasketches/tests/theta_test/main.rs b/datasketches/tests/theta_test/main.rs index f3d4a72..d9d0800 100644 --- a/datasketches/tests/theta_test/main.rs +++ b/datasketches/tests/theta_test/main.rs @@ -17,5 +17,6 @@ mod a_not_b; mod intersection; +mod jaccard_similarity; mod sketch; mod union;