MRG: refactor calculate_gather_stats to disallow repeated downsampling

include/sourmash.h

@@ -30,6 +30,10 @@ enum SourmashErrorCode {
   SOURMASH_ERROR_CODE_NON_EMPTY_MIN_HASH = 106,
   SOURMASH_ERROR_CODE_MISMATCH_NUM = 107,
   SOURMASH_ERROR_CODE_NEEDS_ABUNDANCE_TRACKING = 108,
+  SOURMASH_ERROR_CODE_CANNOT_UPSAMPLE_SCALED = 109,
+  SOURMASH_ERROR_CODE_NO_MIN_HASH_FOUND = 110,
+  SOURMASH_ERROR_CODE_EMPTY_SIGNATURE = 111,
+  SOURMASH_ERROR_CODE_MULTIPLE_SKETCHES_FOUND = 112,
   SOURMASH_ERROR_CODE_INVALID_DNA = 1101,
   SOURMASH_ERROR_CODE_INVALID_PROT = 1102,
   SOURMASH_ERROR_CODE_INVALID_CODON_LENGTH = 1103,

src/core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sourmash"
-version = "0.15.2"
+version = "0.16.0"
 authors = ["Luiz Irber <[email protected]>", "N. Tessa Pierce-Ward <[email protected]>"]
 description = "tools for comparing biological sequences with k-mer sketches"
 repository = "https://github.com/sourmash-bio/sourmash"

src/core/src/collection.rs

@@ -241,6 +241,7 @@ mod test {
     use crate::prelude::Select;
     use crate::selection::Selection;
     use crate::signature::Signature;
+    #[cfg(all(feature = "branchwater", not(target_arch = "wasm32")))]
     use crate::Result;
 
     #[test]

src/core/src/errors.rs

@@ -7,6 +7,9 @@ pub enum SourmashError {
     #[error("internal error: {message:?}")]
     Internal { message: String },
 
+    #[error("new scaled smaller than previous; cannot upsample")]
+    CannotUpsampleScaled,
+
     #[error("must have same num: {n1} != {n2}")]
     MismatchNum { n1: u32, n2: u32 },
 
@@ -28,6 +31,15 @@ pub enum SourmashError {
     #[error("sketch needs abundance for this operation")]
     NeedsAbundanceTracking,
 
+    #[error("Expected a MinHash sketch in this signature")]
+    NoMinHashFound,
+
+    #[error("Empty signature")]
+    EmptySignature,
+
+    #[error("Multiple sketches found, expected one")]
+    MultipleSketchesFound,
+
     #[error("Invalid hash function: {function:?}")]
     InvalidHashFunction { function: String },
 
@@ -104,6 +116,10 @@ pub enum SourmashErrorCode {
     NonEmptyMinHash = 1_06,
     MismatchNum = 1_07,
     NeedsAbundanceTracking = 1_08,
+    CannotUpsampleScaled = 1_09,
+    NoMinHashFound = 1_10,
+    EmptySignature = 1_11,
+    MultipleSketchesFound = 1_12,
     // Input sequence errors
     InvalidDNA = 11_01,
     InvalidProt = 11_02,
@@ -132,6 +148,7 @@ impl SourmashErrorCode {
         match error {
             SourmashError::Internal { .. } => SourmashErrorCode::Internal,
             SourmashError::Panic { .. } => SourmashErrorCode::Panic,
+            SourmashError::CannotUpsampleScaled { .. } => SourmashErrorCode::CannotUpsampleScaled,
             SourmashError::MismatchNum { .. } => SourmashErrorCode::MismatchNum,
             SourmashError::NeedsAbundanceTracking { .. } => {
                 SourmashErrorCode::NeedsAbundanceTracking
@@ -142,6 +159,9 @@ impl SourmashErrorCode {
             SourmashError::MismatchSeed => SourmashErrorCode::MismatchSeed,
             SourmashError::MismatchSignatureType => SourmashErrorCode::MismatchSignatureType,
             SourmashError::NonEmptyMinHash { .. } => SourmashErrorCode::NonEmptyMinHash,
+            SourmashError::NoMinHashFound => SourmashErrorCode::NoMinHashFound,
+            SourmashError::EmptySignature => SourmashErrorCode::EmptySignature,
+            SourmashError::MultipleSketchesFound => SourmashErrorCode::MultipleSketchesFound,
             SourmashError::InvalidDNA { .. } => SourmashErrorCode::InvalidDNA,
             SourmashError::InvalidProt { .. } => SourmashErrorCode::InvalidProt,
             SourmashError::InvalidCodonLength { .. } => SourmashErrorCode::InvalidCodonLength,

src/core/src/index/mod.rs

@@ -27,6 +27,7 @@
 use crate::signature::SigsTrait;
 use crate::sketch::minhash::KmerMinHash;
 use crate::storage::SigStore;
+use crate::Error::CannotUpsampleScaled;
 use crate::Result;
 
 #[derive(TypedBuilder, CopyGetters, Getters, Setters, Serialize, Deserialize, Debug, PartialEq)]
@@ -205,11 +206,10 @@
     }
 }
 
-// note all mh should be selected/downsampled prior to being passed in here.
 #[allow(clippy::too_many_arguments)]
 pub fn calculate_gather_stats(
     orig_query: &KmerMinHash,
-    query: KmerMinHash,
+    remaining_query: KmerMinHash,
     match_sig: SigStore,
     match_size: usize,
     gather_result_rank: usize,
@@ -218,18 +218,31 @@
     calc_abund_stats: bool,
     calc_ani_ci: bool,
     confidence: Option<f64>,
-) -> Result<GatherResult> {
+) -> Result<(GatherResult, (Vec<u64>, u64))> {
     // get match_mh
-    let match_mh = match_sig.minhash().unwrap();
+    let match_mh = match_sig.minhash().expect("cannot retrieve sketch");
+
+    // it's ok to downsample match, but query is often big and repeated,
+    // so we do not allow downsampling of query in this function.
+    if match_mh.scaled() > remaining_query.scaled() {
+        return Err(CannotUpsampleScaled);
+    }
+
+    let match_mh = match_mh
+        .clone()
+        .downsample_scaled(remaining_query.scaled())
+        .expect("cannot downsample match");
 
     // calculate intersection
-    let isect = match_mh.intersection(&query)?;
+    let isect = match_mh
+        .intersection(&remaining_query)
+        .expect("could not do intersection");
     let isect_size = isect.0.len();
     trace!("isect_size: {}", isect_size);
-    trace!("query.size: {}", query.size());
+    trace!("query.size: {}", remaining_query.size());
 
     //bp remaining in subtracted query
-    let remaining_bp = (query.size() - isect_size) * query.scaled() as usize;
+    let remaining_bp = (remaining_query.size() - isect_size) * remaining_query.scaled() as usize;
 
     // stats for this match vs original query
     let (intersect_orig, _) = match_mh.intersection_size(orig_query).unwrap();
@@ -289,7 +302,7 @@
     // If abundance, calculate abund-related metrics (vs current query)
     if calc_abund_stats {
         // take abunds from subtracted query
-        let (abunds, unique_weighted_found) = match match_mh.inflated_abundances(&query) {
+        let (abunds, unique_weighted_found) = match match_mh.inflated_abundances(&remaining_query) {
             Ok((abunds, unique_weighted_found)) => (abunds, unique_weighted_found),
             Err(e) => {
                 return Err(e);
@@ -336,7 +349,7 @@
         .sum_weighted_found(sum_total_weighted_found)
         .total_weighted_hashes(total_weighted_hashes)
         .build();
-    Ok(result)
+    Ok((result, isect))
 }
 
 #[cfg(test)]
@@ -392,7 +405,7 @@
         let gather_result_rank = 0;
         let calc_abund_stats = true;
         let calc_ani_ci = false;
-        let result = calculate_gather_stats(
+        let (result, _isect) = calculate_gather_stats(
             &orig_query,
             query,
             match_sig.into(),
@@ -405,6 +418,7 @@
             None,
         )
         .unwrap();
+
         // first, print all results
         assert_eq!(result.filename(), "match-filename");
         assert_eq!(result.name(), "match-name");

src/core/src/index/revindex/disk_revindex.rs

@@ -1,3 +1,4 @@
+use std::cmp::max;
 use std::hash::{BuildHasher, BuildHasherDefault, Hash, Hasher};
 use std::path::Path;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -372,7 +373,6 @@ impl RevIndexOps for RevIndex {
         let mut query = KmerMinHashBTree::from(orig_query.clone());
         let mut sum_weighted_found = 0;
         let _selection = selection.unwrap_or_else(|| self.collection.selection());
-        let mut orig_query_ds = orig_query.clone();
         let total_weighted_hashes = orig_query.sum_abunds();
 
         // or set this with user --track-abundance?
@@ -393,29 +393,29 @@ impl RevIndexOps for RevIndex {
                 break;
             }
 
-            // this should downsample mh for us
             let match_sig = self.collection.sig_for_dataset(dataset_id)?;
-
-            // get downsampled minhashes for comparison.
             let match_mh = match_sig.minhash().unwrap().clone();
-            query = query.downsample_scaled(match_mh.scaled())?;
-            orig_query_ds = orig_query_ds.downsample_scaled(match_mh.scaled())?;
 
-            // just calculate essentials here
-            let gather_result_rank = matches.len();
+            // make downsampled minhashes
+            let max_scaled = max(match_mh.scaled(), query.scaled());
+
+            let match_mh = match_mh
+                .downsample_scaled(max_scaled)
+                .expect("cannot downsample match");
 
+            query = query
+                .downsample_scaled(max_scaled)
+                .expect("cannot downsample query");
             let query_mh = KmerMinHash::from(query.clone());
 
-            // grab the specific intersection:
-            let isect = match_mh.intersection(&query_mh)?;
-            let mut isect_mh = match_mh.clone();
-            isect_mh.clear();
-            isect_mh.add_many(&isect.0)?;
+            // just calculate essentials here
+            let gather_result_rank = matches.len();
 
+            // grab the specific intersection:
             // Calculate stats
-            let gather_result = calculate_gather_stats(
-                &orig_query_ds,
-                KmerMinHash::from(query.clone()),
+            let (gather_result, isect) = calculate_gather_stats(
+                &orig_query,
+                query_mh,
                 match_sig,
                 match_size,
                 gather_result_rank,
@@ -424,7 +424,14 @@ impl RevIndexOps for RevIndex {
                 calc_abund_stats,
                 calc_ani_ci,
                 ani_confidence_interval_fraction,
-            )?;
+            )
+            .expect("could not calculate gather stats");
+
+            // use intersection from calc_gather_stats to make a KmerMinHash.
+            let mut isect_mh = match_mh.clone();
+            isect_mh.clear();
+            isect_mh.add_many(&isect.0)?;
+
             // keep track of the sum weighted found
             sum_weighted_found = gather_result.sum_weighted_found();
             matches.push(gather_result);

src/core/src/index/revindex/mod.rs

@@ -903,14 +903,16 @@ mod test {
 
         let (counter, query_colors, hash_to_color) = index.prepare_gather_counters(&query);
 
-        let matches_external = index.gather(
-            counter,
-            query_colors,
-            hash_to_color,
-            0,
-            &query,
-            Some(selection.clone()),
-        )?;
+        let matches_external = index
+            .gather(
+                counter,
+                query_colors,
+                hash_to_color,
+                0,
+                &query,
+                Some(selection.clone()),
+            )
+            .expect("failed to gather!");
 
         {
             let mut index = index;

src/core/src/signature.rs

@@ -842,7 +842,7 @@
                 // TODO: also account for LargeMinHash
                 if let Sketch::MinHash(mh) = sketch {
                     if (mh.scaled() as u32) < sel_scaled {
-                        *sketch = Sketch::MinHash(mh.downsample_scaled(sel_scaled as u64)?);
+                        *sketch = Sketch::MinHash(mh.clone().downsample_scaled(sel_scaled as u64)?);
                     }
                 }
             }
@@ -887,6 +887,28 @@
     }
 }
 
+impl TryInto<KmerMinHash> for Signature {
+    type Error = Error;
+
+    fn try_into(self) -> Result<KmerMinHash, Error> {
+        match self.signatures.len() {
+            1 => self
+                .signatures
+                .into_iter()
+                .find_map(|sk| {
+                    if let Sketch::MinHash(mh) = sk {
+                        Some(mh)
+                    } else {
+                        None
+                    }
+                })
+                .ok_or_else(|| Error::NoMinHashFound),
+            0 => Err(Error::EmptySignature),
+            _ => Err(Error::MultipleSketchesFound),
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::fs::File;