Skip to main content

dedup/
lib.rs

1#![allow(
2    clippy::module_name_repetitions,
3    clippy::must_use_candidate,
4    clippy::missing_errors_doc,
5)]
6#![warn(missing_docs)]
7#![warn(clippy::pedantic)]
8#![forbid(unsafe_code)]
9//! # dedup  -  High-performance dataset deduplication for ML training data
10//!
11//! A `MinHash` + LSH implementation for finding near-duplicate documents
12//! in massive datasets. Designed for streaming operation to handle
13//! billions of documents without loading all into memory.
14//!
15//! ## Quick Start
16//!
17//! ```rust
18//! use dedup::{Config, DedupTransformer};
19//!
20//! let config = Config::default()
21//!     .with_similarity_threshold(0.85)
22//!     .with_num_bands(16);
23//!
24//! let dedup = DedupTransformer::new(config).unwrap();
25//! ```
26//!
27//! ## Architecture
28//!
29//! ```text
30//! ┌─────────────┐     ┌──────────────┐     ┌─────────────┐     ┌─────────────┐
31//! │   Shingle   │────▶│   MinHash    │────▶│  LSH Bands  │────▶│   Deduplicate│
32//! │  (k-grams)  │     │  (fast hash) │     │  (buckets)  │     │   (filter)   │
33//! └─────────────┘     └──────────────┘     └─────────────┘     └─────────────┘
34//! ```
35//!
36//! ## `MinHash` + LSH Theory
37//!
38//! - **Shingling**: Convert documents to sets of k-grams (overlapping subsequences)
39//! - **MinHash**: Compress document to a small signature while preserving Jaccard similarity
40//! - **LSH**: Band signatures such that similar documents collide in at least one bucket
41//! - **Threshold**: Documents with estimated Jaccard ≥ threshold are considered duplicates
42
43#![warn(missing_docs, clippy::pedantic)]
44
45#![cfg_attr(
46    not(test),
47    deny(
48        clippy::unwrap_used,
49        clippy::expect_used,
50        clippy::todo,
51        clippy::unimplemented,
52        clippy::panic
53    )
54)]
55#![allow(
56    clippy::cast_possible_truncation,
57    clippy::cast_possible_wrap,
58    clippy::module_name_repetitions,
59    clippy::needless_pass_by_value,
60    clippy::must_use_candidate,
61    clippy::return_self_not_must_use,
62    clippy::unnecessary_literal_bound,
63    clippy::doc_markdown,
64    clippy::cast_precision_loss
65)]
66
67mod cluster;
68mod config;
69mod error;
70mod lsh;
71mod minhash;
72pub mod shingle;
73mod fast_hash;
74mod transform;
75
76pub use config::Config;
77pub use error::{Error, Result};
78pub use cluster::DuplicateCluster;
79pub use lsh::LshIndex;
80pub use minhash::{exact_jaccard_similarity, expected_error, MinHashSignature, MinHasher};
81pub use shingle::ShingleIterator;
82pub use fast_hash::{hash_bytes, FastHasher};
83pub use transform::{DedupTransformer, StatefulDedupTransform};
84
85/// Re-export tenshift types for convenience.
86pub mod tenshift {
87    pub use tenshift_core::sample::Sample;
88    pub use tenshift_core::transform::{Transform, TransformResult};
89}
90
91/// Default number of hash functions (signature size).
92pub const DEFAULT_SIGNATURE_SIZE: usize = 128;
93
94/// Default shingle size in bytes/characters.
95pub const DEFAULT_SHINGLE_SIZE: usize = 5;
96
97/// Default number of LSH bands.
98pub const DEFAULT_NUM_BANDS: usize = 16;
99
100/// Default similarity threshold for considering documents as duplicates.
101pub const DEFAULT_SIMILARITY_THRESHOLD: f64 = 0.9;
102
103/// Compute the number of rows per band given signature size and num bands.
104///
105/// Returns `None` if the signature size is not evenly divisible by num bands.
106#[must_use]
107pub const fn compute_rows_per_band(signature_size: usize, num_bands: usize) -> Option<usize> {
108    if num_bands == 0 || signature_size % num_bands != 0 {
109        return None;
110    }
111    Some(signature_size / num_bands)
112}
113
114/// Estimate the false positive rate given LSH parameters.
115///
116/// The probability that two documents with similarity `s` will be
117/// marked as candidates for comparison.
118#[must_use]
119pub fn candidate_probability(similarity: f64, num_bands: usize, rows_per_band: usize) -> f64 {
120    if similarity <= 0.0 {
121        return 0.0;
122    }
123    if similarity >= 1.0 {
124        return 1.0;
125    }
126    // P(at least one band matches) = 1 - P(no bands match)
127    // P(one band matches) = s^r where r = rows_per_band
128    let band_match_prob = similarity.powf(rows_per_band as f64);
129    1.0 - (1.0 - band_match_prob).powf(num_bands as f64)
130}
131
132/// Find the optimal LSH parameters for a given similarity threshold.
133///
134/// Returns `(num_bands, rows_per_band)` that maximizes the S-curve
135/// steepness around the threshold.
136#[must_use]
137pub fn optimize_lsh_params(
138    signature_size: usize,
139    target_threshold: f64,
140) -> (usize, usize) {
141    // The optimal point is where s^r ≈ 1/b for threshold s
142    // This gives us roughly b * s^r = 1 expected matches at threshold.
143    //
144    // Enumerate every band count that EXACTLY divides `signature_size` (not a
145    // hardcoded {4,8,16,32,64} shortlist). A shortlist returns an INVALID
146    // (bands, rows) whenever `signature_size` divides none of its entries
147    // (e.g. 100: 100/8=12 but 8*12=96 != 100), so the caller would build an
148    // LSH index whose bands don't tile the signature. Iterating real divisors
149    // guarantees `bands * rows == signature_size` for the returned pair, and
150    // the discrimination score naturally rejects the degenerate 1-band /
151    // 1-row divisors.
152    if signature_size == 0 {
153        return (1, 0);
154    }
155    // Only consider band counts up to a sane cap. Realistic LSH configs never
156    // use more than a few hundred bands, and iterating `1..=signature_size`
157    // would hang for an adversarial `signature_size` near `usize::MAX`. The
158    // cap keeps the search O(MAX_CANDIDATE_BANDS); a valid result is still
159    // guaranteed because `num_bands == 1` (rows == signature_size) always
160    // divides and is always in range.
161    const MAX_CANDIDATE_BANDS: usize = 1024;
162    let search_limit = signature_size.min(MAX_CANDIDATE_BANDS);
163    let mut best: Option<(f64, usize, usize)> = None;
164
165    for num_bands in 1..=search_limit {
166        if signature_size % num_bands != 0 {
167            continue;
168        }
169        let rows_per_band = signature_size / num_bands;
170
171        // Score by how steep the curve is at the threshold
172        let p_at_threshold = candidate_probability(target_threshold, num_bands, rows_per_band);
173        let p_below = candidate_probability(target_threshold * 0.9, num_bands, rows_per_band);
174        let p_above =
175            candidate_probability((target_threshold * 1.1).min(1.0), num_bands, rows_per_band);
176
177        // We want high discrimination: low p_below, high p_above
178        let score = p_above - p_below - (p_at_threshold - 0.5).abs() * 0.5;
179
180        if best.map_or(true, |(best_score, _, _)| score > best_score) {
181            best = Some((score, num_bands, rows_per_band));
182        }
183    }
184
185    // `signature_size >= 1` always has the divisor 1, so `best` is `Some`.
186    let (_, bands, rows) = best.unwrap_or((0.0, 1, signature_size));
187    (bands, rows)
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    #[test]
195    fn test_compute_rows_per_band() {
196        assert_eq!(compute_rows_per_band(128, 16), Some(8));
197        assert_eq!(compute_rows_per_band(128, 8), Some(16));
198        assert_eq!(compute_rows_per_band(128, 0), None);
199        assert_eq!(compute_rows_per_band(128, 3), None);
200    }
201
202    #[test]
203    fn test_candidate_probability_bounds() {
204        assert_eq!(candidate_probability(0.0, 16, 8), 0.0);
205        assert_eq!(candidate_probability(1.0, 16, 8), 1.0);
206    }
207
208    #[test]
209    fn test_candidate_probability_increases_with_similarity() {
210        let p1 = candidate_probability(0.5, 16, 8);
211        let p2 = candidate_probability(0.8, 16, 8);
212        let p3 = candidate_probability(0.95, 16, 8);
213        
214        assert!(p1 < p2, "probability should increase with similarity");
215        assert!(p2 < p3, "probability should increase with similarity");
216    }
217
218    #[test]
219    fn test_optimize_lsh_params_produces_valid_params() {
220        let (bands, rows) = optimize_lsh_params(128, 0.9);
221        assert!(bands > 0);
222        assert!(rows > 0);
223        assert_eq!(bands * rows, 128);
224    }
225
226    #[test]
227    fn test_optimize_lsh_params_valid_when_indivisible_by_hardcoded_bands() {
228        // Regression: the old code only tried band counts {4,8,16,32,64} (all
229        // multiples of 4). Any signature_size NOT divisible by 4 matched none,
230        // so it returned the untouched seed (8, size/8) whose product != size:
231        // e.g. size=6 -> (8, 0), size=50 -> (8, 6) [48 != 50]. Now every result
232        // must exactly tile the signature.
233        for size in [1_usize, 2, 3, 5, 6, 7, 10, 25, 30, 50, 100, 127, 200] {
234            let (bands, rows) = optimize_lsh_params(size, 0.85);
235            assert!(bands >= 1, "size {size}: bands must be >= 1, got {bands}");
236            assert!(rows >= 1, "size {size}: rows must be >= 1, got {rows}");
237            assert_eq!(
238                bands * rows,
239                size,
240                "size {size}: returned ({bands}, {rows}) must tile the signature exactly"
241            );
242        }
243    }
244}
245
246/// Compile-checks the README quick-start example as a doctest.
247#[cfg(doctest)]
248#[doc = include_str!("../README.md")]
249struct ReadmeExamples;