ff

src/bloom_filter.rs

@@ -11,6 +11,7 @@ use std::mem::size_of;
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicU32, Ordering};
 
+mod bloom_test;
 // A thread-safe bloom filter.
 pub struct BloomFilter {
     bits: Vec<AtomicU32>,
@@ -85,13 +86,10 @@ impl BloomFilter {
             hash_builder_seeds.push(seeds);
         }
 
-        let mut bits = Vec::new();
         let number_of_u32 = size_in_bytes / size_of::<AtomicU32>();
-        bits.reserve_exact(number_of_u32);
-        for _ in 0..number_of_u32 {
-            bits.push(AtomicU32::new(0));
-        }
-
+        let bits: Vec<AtomicU32> = std::iter::repeat_with(|| AtomicU32::new(0))
+            .take(number_of_u32)
+            .collect();
         Self {
             bits,
             hash_builder_seeds,
@@ -138,8 +136,7 @@ impl BloomFilter {
         }
 
         let number_of_elements = stream.read_u64::<LittleEndian>()?;
-        let mut bits = Vec::new();
-        bits.reserve_exact(number_of_elements as usize);
+        let mut bits = Vec::with_capacity(number_of_elements as usize);
         for _ in 0..number_of_elements {
             bits.push(AtomicU32::new(stream.read_u32::<NativeEndian>()?));
         }
@@ -220,8 +217,7 @@ impl BloomFilter {
                 return false;
             }
         }
-
-        return true;
+        true
     }
 
     pub fn contains(&self, s: &VecDeque<&str>) -> bool {
@@ -237,7 +233,7 @@ impl BloomFilter {
         } else {
             log::info!("Creating new bloom filter...");
             let mut bloom_filter_size: usize = config.size_in_bytes;
-            if bloom_filter_size <= 0 {
+            if bloom_filter_size == 0 {
                 bloom_filter_size = BloomFilter::suggest_size_in_bytes(
                     config.estimated_doc_count,
                     config.desired_false_positive_rate,

src/bloom_filter/bloom_test.rs

@@ -0,0 +1,56 @@
+#[cfg(test)]
+mod tests {
+    use super::super::BloomFilter;
+
+    // n: number of items in filter. p: false positive rate
+    // m: number of bits in filter. k: number of hashers
+    // n = ceil(m / (-k / log(1 - exp(log(p) / k))))
+    // p = pow(1 - exp(-k / (m / n)), k)
+    // m = ceil((n * log(p)) / log(1 / pow(2, log(2))));
+    // k = round((m / n) * log(2));
+
+    #[test]
+    fn bloom_optimal_hasher_number() {
+        let size_in_bytes = 1_000_000_000;
+        let expected_elements = 1_000_000_000;
+        assert_eq!(
+            BloomFilter::optimal_number_of_hashers(size_in_bytes, expected_elements),
+            6
+        );
+        assert_eq!(
+            BloomFilter::optimal_number_of_hashers(1_000_000, 500_000),
+            12
+        )
+    }
+    #[test]
+    fn bloom_test_prob_of_false_positive() {
+        // calculated from https://hur.st/bloomfilter/
+        let size_in_bytes = 1_000_000_000;
+        let expected_elements = 1_000_000_000;
+        let num_hashers = 8;
+        assert_eq!(
+            BloomFilter::prob_of_false_positive(size_in_bytes, expected_elements, num_hashers),
+            0.025_491_740_593_406_025 as f64
+        );
+        assert_eq!(
+            BloomFilter::prob_of_false_positive(1_048_576, 524288, 2),
+            0.013_806_979_447_406_826 as f64
+        )
+    }
+
+    #[test]
+    fn bloom_suggest_size() {
+        // it's hard to derive this exactly since the algorithm is doing closest power of 2
+        // instead of exact theoretical optimum
+        let expected_elements = 1_000_000;
+        let desired_false_positive_rate = 0.0001 as f64;
+        let theoretical_optimum = ((expected_elements as f64 * desired_false_positive_rate.ln())
+            / f64::ln(1.0 / 2.0f64.powf(2.0f64.ln())))
+        .ceil()
+        .div_euclid(8f64) as usize;
+        let suggested_size =
+            BloomFilter::suggest_size_in_bytes(expected_elements, desired_false_positive_rate);
+        assert_eq!(suggested_size, 4_194_304);
+        assert_eq!(suggested_size, theoretical_optimum.next_power_of_two())
+    }
+}

src/deduper.rs

@@ -43,33 +43,32 @@ pub fn run(config: DeduperConfig) -> Result<u32, u32> {
         let failed_shard_count_ref = failed_shard_count_ref.clone();
         threadpool.execute(move || {
             let result = write_attributes(path, work_dirs, dedupe, bloom_filter);
-            match result {
-                Ok(_) => {}
-                Err(e) => {
-                    log::error!("Failed to process {:?}: {}", p, e);
-                    failed_shard_count_ref.fetch_add(1, Ordering::Relaxed);
-                }
+            if let Err(e) = result {
+                log::error!("Failed to process {:?}: {}", p, e);
+                failed_shard_count_ref.fetch_add(1, Ordering::Relaxed);
             }
         });
     }
     threadpool.join();
 
     let bloom_filter_file = PathBuf::from(&config.bloom_filter.file);
     log::info!("Writing bloom filter to {:?}...", config.bloom_filter.file);
-    bloom_filter.write_to_file(&bloom_filter_file).unwrap();
-    log::info!("Bloom filter written.");
-
-    let failure_count = failed_shard_count_ref.fetch_add(0, Ordering::Relaxed);
-    match failure_count {
-        0 => {
-            log::info!("Done!");
-            return Ok(failure_count);
-        }
-        _ => {
-            log::error!("{} shards failed to process.", failure_count);
-            return Err(failure_count);
+    match bloom_filter.write_to_file(&bloom_filter_file) {
+        Ok(_) => log::info!("Bloom filter written."),
+        Err(e) => {
+            log::error!("Write failed: {}", e);
+            panic!("Failed to write bloom filter");
         }
     }
+
+    let failure_count = failed_shard_count_ref.load(Ordering::Relaxed);
+    if failure_count == 0 {
+        log::info!("Done!");
+        Ok(failure_count)
+    } else {
+        log::error!("{} shards failed to process.", failure_count);
+        Err(failure_count)
+    }
 }
 
 // Write attributes for the documents in the given file:
@@ -87,12 +86,8 @@ fn write_attributes(
     };
 
     let attrs_location = {
-        let mut attr_prefix = "/attributes/".to_owned();
-        attr_prefix.push_str(&dedupe_config.name);
-        attr_prefix.push_str("/");
-        docs_location
-            .to_owned()
-            .replace("/documents/", &attr_prefix)
+        let attr_prefix = format!("/attributes/{}/", &dedupe_config.name);
+        docs_location.replace("/documents/", &attr_prefix)
     };
     let local_output = cache.prepare_output(&attrs_location)?;
     if local_output.exists() {
@@ -133,10 +128,9 @@ fn write_attributes(
             GzEncoder::new(tmp_output, Compression::default()),
         );
 
-        let mut line_number = 0;
-        for line in reader.lines() {
-            match line {
-                Ok(_) => {}
+        for (line_number, line) in reader.lines().enumerate() {
+            let line = match line {
+                Ok(line) => line,
                 Err(e) => {
                     log::error!(
                         "Error reading line {} of {}: {}",
@@ -146,45 +140,39 @@ fn write_attributes(
                     );
                     break;
                 }
-            }
-            line_number += 1;
-            let line = line?;
+            };
             let data: Value = serde_json::from_str(&line)?;
             let mut attributes = json!({});
 
-            match dedupe_config.documents {
-                Some(ref cfg) => {
-                    let document_key = {
-                        let mut finder = jsonpath_rust::JsonPathFinder::from_str("{}", &cfg.key)
-                            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
-                            .unwrap();
-                        finder.set_json(Box::new(data.clone()));
-                        finder
-                            .find()
-                            .as_array()
-                            .unwrap()
-                            .get(0)
-                            .unwrap()
-                            .as_str()
-                            .unwrap()
-                            .to_string()
-                    };
+            if let Some(ref cfg) = dedupe_config.documents {
+                let document_key = {
+                    let mut finder = jsonpath_rust::JsonPathFinder::from_str("{}", &cfg.key)
+                        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
+                        .unwrap();
+                    finder.set_json(Box::new(data.clone()));
+                    finder
+                        .find()
+                        .as_array()
+                        .unwrap()
+                        .get(0)
+                        .unwrap()
+                        .as_str()
+                        .unwrap()
+                        .to_string()
+                };
 
-                    if dedupe_config.skip_empty.unwrap_or(false) && document_key.trim().is_empty() {
-                        // skip empty documents if dedupe_config.skip_empty is true
-                        // and the document key is empty after trimming (i.e., removing whitespace)
-                        continue;
-                    } else {
-                        let mut dedupe_key = VecDeque::with_capacity(1);
-                        dedupe_key.push_back(document_key.as_str());
-                        if bloom_filter.contains(&dedupe_key) {
-                            attributes[&cfg.attribute_name] = Value::Bool(true);
-                        } else if !bloom_filter.read_only {
-                            bloom_filter.insert(&dedupe_key);
-                        }
+                if dedupe_config.skip_empty.unwrap_or(false) && document_key.trim().is_empty() {
+                    // skip empty documents if dedupe_config.skip_empty is true
+                    // and the document key is empty after trimming (i.e., removing whitespace)
+                    continue;
+                } else {
+                    let dedupe_key = VecDeque::from([document_key.as_str()]);
+                    if bloom_filter.contains(&dedupe_key) {
+                        attributes[&cfg.attribute_name] = Value::Bool(true);
+                    } else if !bloom_filter.read_only {
+                        bloom_filter.insert(&dedupe_key);
                     }
                 }
-                None => {}
             }
             match dedupe_config.paragraphs {
                 None => {}
@@ -193,7 +181,7 @@ fn write_attributes(
                     let text = data["text"].as_str().unwrap();
                     let text_length = text.len();
                     let mut offset = 0;
-                    let paragraphs = text.split("\n");
+                    let paragraphs = text.split('\n');
                     let mut duplicate_paragraph_spans = Vec::new();
                     for p in paragraphs {
                         let par_start = offset;
@@ -208,13 +196,12 @@ fn write_attributes(
                             // and the paragraph is empty after trimming (i.e., removing whitespace)
                             continue;
                         } else {
-                            let mut dedupe_key = VecDeque::with_capacity(1);
-                            dedupe_key.push_back(p);
+                            let dedupe_key = VecDeque::from([p]);
                             if bloom_filter.contains(&dedupe_key) {
                                 let span = vec![
                                     Value::Number(par_start.into()),
                                     Value::Number(par_end.into()),
-                                    Value::Number(1.into()),
+                                    Value::from(1),
                                 ];
                                 // add span to duplicate_paragraph_spans
                                 duplicate_paragraph_spans.push(Value::Array(span));

src/lib.rs

@@ -15,26 +15,26 @@ use std::env;
 fn deduper_entrypoint(config_str: &str) -> PyResult<()> {
     let config: DeduperConfig = DeduperConfig::parse_from_string(config_str).unwrap();
 
-    match deduper::run(config) {
-        Ok(_) => Ok(()),
-        Err(cnt) => Err(exceptions::PyRuntimeError::new_err(format!(
+    if let Err(cnt) = deduper::run(config) {
+        return Err(exceptions::PyRuntimeError::new_err(format!(
             "Failed with {} errors",
             cnt
-        ))),
+        )));
     }
+    Ok(())
 }
 
 #[pyfunction]
 fn mixer_entrypoint(config_str: &str) -> PyResult<()> {
     //Result<u32, PyErr> {
     let config: MixerConfig = MixerConfig::parse_from_string(config_str).unwrap();
-    match mixer::run(config) {
-        Ok(_) => Ok(()),
-        Err(cnt) => Err(exceptions::PyRuntimeError::new_err(format!(
+    if let Err(cnt) = mixer::run(config) {
+        return Err(exceptions::PyRuntimeError::new_err(format!(
             "Failed with {} errors",
             cnt
-        ))),
+        )));
     }
+    Ok(())
 }
 
 // A Python module implemented in Rust. The name of this function must match

src/mixer.rs

@@ -12,8 +12,7 @@ pub fn run(config: MixerConfig) -> Result<u32, u32> {
     let shards = Shard::split_streams(&config.streams).unwrap();
 
     let threadpool = ThreadPool::new(config.processes);
-    let failed_shard_count = AtomicU32::new(0);
-    let failed_shard_count_ref = Arc::new(failed_shard_count);
+    let failed_shard_count_ref = Arc::new(AtomicU32::new(0));
     for shard in shards {
         let output_path = Path::new(&config.work_dir.output.clone()).join(&shard.output);
         if output_path.exists() {
@@ -26,27 +25,21 @@ pub fn run(config: MixerConfig) -> Result<u32, u32> {
 
         threadpool.execute(move || {
             log::info!("Building output {:?}...", shard.output);
-            match shard.clone().process(work_dirs) {
-                Ok(_) => {}
-                Err(e) => {
-                    log::error!("Error processing {:?}: {}", shard.output, e);
-                    failed_shard_count_ref.fetch_add(1, Ordering::Relaxed);
-                }
+            if let Err(e) = shard.clone().process(work_dirs) {
+                log::error!("Error processing {:?}: {}", shard.output, e);
+                failed_shard_count_ref.fetch_add(1, Ordering::Relaxed);
             }
         });
     }
     threadpool.join();
 
-    let failure_count = failed_shard_count_ref.fetch_add(0, Ordering::Relaxed);
-    match failure_count {
-        0 => {
-            log::info!("Done!");
-            return Ok(failure_count);
-        }
-        _ => {
-            log::error!("{} shards failed to process.", failure_count);
-            return Err(failure_count);
-        }
+    let failure_count = failed_shard_count_ref.load(Ordering::Relaxed);
+    if failure_count == 0 {
+        log::info!("Done!");
+        Ok(failure_count)
+    } else {
+        log::error!("{} shards failed to process.", failure_count);
+        Err(failure_count)
     }
 }