MRG: add generic support for any type of sketch collection as query or database

.cargo/config.toml

@@ -0,0 +1,2 @@
+[build]
+rustdocflags = ["--document-private-items"]

Cargo.toml

@@ -16,7 +16,7 @@ sourmash = { version = "0.16.0", features = ["branchwater"] }
 serde_json = "1.0.128"
 niffler = "2.4.0"
 log = "0.4.22"
-env_logger = { version = "0.11.5", optional = true }
+env_logger = { version = "0.11.5" }
 simple-error = "0.3.1"
 anyhow = "1.0.89"
 zip = { version = "2.0", default-features = false }

src/fastgather.rs

@@ -33,7 +33,9 @@ pub fn fastgather(
         )
     }
     // get single query sig and minhash
-    let query_sig = query_collection.sig_for_dataset(0)?; // need this for original md5sum
+    let query_sig = query_collection.get_first_sig().expect("no queries!?");
+
+    // @CTB avoid clone?
     let query_sig_ds = query_sig.clone().select(selection)?; // downsample
     let query_mh = match query_sig_ds.minhash() {
         Some(query_mh) => query_mh,
@@ -89,7 +91,17 @@ pub fn fastgather(
     }
 
     if prefetch_output.is_some() {
-        write_prefetch(&query_sig, prefetch_output, &matchlist).ok();
+        let query_filename = query_sig.filename();
+        let query_name = query_sig.name();
+        let query_md5 = query_sig.md5sum();
+        write_prefetch(
+            query_filename,
+            query_name,
+            query_md5,
+            prefetch_output,
+            &matchlist,
+        )
+        .ok();
     }
 
     // run the gather!

src/fastmultigather.rs

@@ -2,7 +2,7 @@
 use anyhow::Result;
 use rayon::prelude::*;
 
-use sourmash::prelude::ToWriter;
+use sourmash::prelude::{Storage, ToWriter};
 use sourmash::{selection::Selection, signature::SigsTrait};
 
 use std::sync::atomic;
@@ -15,13 +15,14 @@ use camino::Utf8Path as PathBuf;
 use std::collections::HashSet;
 use std::fs::File;
 
+use log::trace;
+
 use sourmash::signature::Signature;
 use sourmash::sketch::minhash::KmerMinHash;
 use sourmash::sketch::Sketch;
 
 use crate::utils::{
-    consume_query_by_gather, load_collection, load_sketches, write_prefetch, PrefetchResult,
-    ReportType,
+    consume_query_by_gather, load_collection, write_prefetch, PrefetchResult, ReportType,
 };
 
 pub fn fastmultigather(
@@ -34,6 +35,8 @@ pub fn fastmultigather(
     save_matches: bool,
     create_empty_results: bool,
 ) -> Result<()> {
+    let _ = env_logger::try_init();
+
     // load query collection
     let query_collection = load_collection(
         &query_filepath,
@@ -50,8 +53,7 @@ pub fn fastmultigather(
             1
         }
     }
-    .try_into()
-    .unwrap();
+    .try_into()?;
 
     println!("threshold overlap: {} {}", threshold_hashes, threshold_bp);
 
@@ -63,123 +65,130 @@ pub fn fastmultigather(
         allow_failed_sigpaths,
     )?;
     // load against sketches into memory, downsampling on the way
-    let against = load_sketches(against_collection, selection, ReportType::Against).unwrap();
+    let against = against_collection.load_sketches(selection)?;
 
     // Iterate over all queries => do prefetch and gather!
     let processed_queries = AtomicUsize::new(0);
     let skipped_paths = AtomicUsize::new(0);
     let failed_paths = AtomicUsize::new(0);
 
-    query_collection.par_iter().for_each(|(_idx, record)| {
+    query_collection.par_iter().for_each(|(c, _idx, record)| {
         // increment counter of # of queries. q: could we instead use the _idx from par_iter(), or will it vary based on thread?
         let _i = processed_queries.fetch_add(1, atomic::Ordering::SeqCst);
         // Load query sig (downsampling happens here)
-        match query_collection.sig_from_record(record) {
+        trace!(
+            "fastmultigather query load: from:{} idx:{} loc:{}",
+            c.storage().spec(),
+            _idx,
+            record.internal_location()
+        );
+        match c.sig_from_record(record) {
             Ok(query_sig) => {
                 let name = query_sig.name();
                 let prefix = name.split(' ').next().unwrap_or_default().to_string();
                 let location = PathBuf::new(&prefix).file_name().unwrap();
-                if let Some(query_mh) = query_sig.minhash() {
-                    let mut matching_hashes = if save_matches { Some(Vec::new()) } else { None };
-                    let matchlist: BinaryHeap<PrefetchResult> = against
-                        .iter()
-                        .filter_map(|against| {
-                            let mut mm: Option<PrefetchResult> = None;
-                            if let Ok(overlap) = against.minhash.count_common(query_mh, false) {
-                                if overlap >= threshold_hashes {
-                                    if save_matches {
-                                        if let Ok(intersection) =
-                                            against.minhash.intersection(query_mh)
-                                        {
-                                            matching_hashes
-                                                .as_mut()
-                                                .unwrap()
-                                                .extend(intersection.0);
-                                        }
-                                    }
-                                    let result = PrefetchResult {
-                                        name: against.name.clone(),
-                                        md5sum: against.md5sum.clone(),
-                                        minhash: against.minhash.clone(),
-                                        location: against.location.clone(),
-                                        overlap,
-                                    };
-                                    mm = Some(result);
-                                }
-                            }
-                            mm
-                        })
-                        .collect();
-                    if !matchlist.is_empty() {
-                        let prefetch_output = format!("{}.prefetch.csv", location);
-                        let gather_output = format!("{}.gather.csv", location);
 
-                        // Save initial list of matches to prefetch output
-                        write_prefetch(&query_sig, Some(prefetch_output), &matchlist).ok();
-
-                        // Now, do the gather!
-                        consume_query_by_gather(
-                            query_sig.clone(),
-                            scaled as u64,
-                            matchlist,
-                            threshold_hashes,
-                            Some(gather_output),
-                        )
-                        .ok();
-
-                        // Save matching hashes to .sig file if save_matches is true
-                        if save_matches {
-                            if let Some(hashes) = matching_hashes {
-                                let sig_filename = format!("{}.matches.sig", name);
-                                if let Ok(mut file) = File::create(&sig_filename) {
-                                    let unique_hashes: HashSet<u64> = hashes.into_iter().collect();
-                                    let mut new_mh = KmerMinHash::new(
-                                        query_mh.scaled().try_into().unwrap(),
-                                        query_mh.ksize().try_into().unwrap(),
-                                        query_mh.hash_function().clone(),
-                                        query_mh.seed(),
-                                        false,
-                                        query_mh.num(),
-                                    );
-                                    new_mh
-                                        .add_many(&unique_hashes.into_iter().collect::<Vec<_>>())
-                                        .ok();
-                                    let mut signature = Signature::default();
-                                    signature.push(Sketch::MinHash(new_mh));
-                                    signature.set_filename(&name);
-                                    if let Err(e) = signature.to_writer(&mut file) {
-                                        eprintln!("Error writing signature file: {}", e);
+                let query_filename = query_sig.filename();
+                let query_name = query_sig.name();
+                let query_md5 = query_sig.md5sum();
+
+                let query_mh = query_sig.minhash().expect("cannot get sketch");
+                let mut matching_hashes = if save_matches { Some(Vec::new()) } else { None };
+                let matchlist: BinaryHeap<PrefetchResult> = against
+                    .iter()
+                    .filter_map(|against| {
+                        let mut mm: Option<PrefetchResult> = None;
+                        if let Ok(overlap) = against.minhash.count_common(query_mh, false) {
+                            if overlap >= threshold_hashes {
+                                if save_matches {
+                                    if let Ok(intersection) = against.minhash.intersection(query_mh)
+                                    {
+                                        matching_hashes.as_mut().unwrap().extend(intersection.0);
                                     }
-                                } else {
-                                    eprintln!("Error creating signature file: {}", sig_filename);
                                 }
+                                let result = PrefetchResult {
+                                    name: against.name.clone(),
+                                    md5sum: against.md5sum.clone(),
+                                    minhash: against.minhash.clone(),
+                                    location: against.location.clone(),
+                                    overlap,
+                                };
+                                mm = Some(result);
                             }
                         }
-                    } else {
-                        println!("No matches to '{}'", location);
-                        if create_empty_results {
-                            let prefetch_output = format!("{}.prefetch.csv", location);
-                            let gather_output = format!("{}.gather.csv", location);
-                            // touch output files
-                            match std::fs::File::create(&prefetch_output) {
-                                Ok(_) => {}
-                                Err(e) => {
-                                    eprintln!("Failed to create empty prefetch output: {}", e)
+                        mm
+                    })
+                    .collect();
+
+                if !matchlist.is_empty() {
+                    let prefetch_output = format!("{}.prefetch.csv", location);
+                    let gather_output = format!("{}.gather.csv", location);
+
+                    // Save initial list of matches to prefetch output
+                    write_prefetch(
+                        query_filename,
+                        query_name,
+                        query_md5,
+                        Some(prefetch_output),
+                        &matchlist,
+                    )
+                    .ok();
+
+                    // Now, do the gather!
+                    consume_query_by_gather(
+                        query_sig.clone(),
+                        scaled as u64,
+                        matchlist,
+                        threshold_hashes,
+                        Some(gather_output),
+                    )
+                    .ok();
+
+                    // Save matching hashes to .sig file if save_matches is true
+                    if save_matches {
+                        if let Some(hashes) = matching_hashes {
+                            let sig_filename = format!("{}.matches.sig", name);
+                            if let Ok(mut file) = File::create(&sig_filename) {
+                                let unique_hashes: HashSet<u64> = hashes.into_iter().collect();
+                                let mut new_mh = KmerMinHash::new(
+                                    query_mh.scaled(),
+                                    query_mh.ksize().try_into().unwrap(),
+                                    query_mh.hash_function().clone(),
+                                    query_mh.seed(),
+                                    false,
+                                    query_mh.num(),
+                                );
+                                new_mh
+                                    .add_many(&unique_hashes.into_iter().collect::<Vec<_>>())
+                                    .ok();
+                                let mut signature = Signature::default();
+                                signature.push(Sketch::MinHash(new_mh));
+                                signature.set_filename(&name);
+                                if let Err(e) = signature.to_writer(&mut file) {
+                                    eprintln!("Error writing signature file: {}", e);
                                 }
-                            }
-                            match std::fs::File::create(&gather_output) {
-                                Ok(_) => {}
-                                Err(e) => eprintln!("Failed to create empty gather output: {}", e),
+                            } else {
+                                eprintln!("Error creating signature file: {}", sig_filename);
                             }
                         }
                     }
                 } else {
-                    // different warning here? Could not load sig from record??
-                    eprintln!(
-                        "WARNING: no compatible sketches in path '{}'",
-                        record.internal_location()
-                    );
-                    let _ = skipped_paths.fetch_add(1, atomic::Ordering::SeqCst);
+                    println!("No matches to '{}'", location);
+                    if create_empty_results {
+                        let prefetch_output = format!("{}.prefetch.csv", location);
+                        let gather_output = format!("{}.gather.csv", location);
+                        // touch output files
+                        match std::fs::File::create(&prefetch_output) {
+                            Ok(_) => {}
+                            Err(e) => {
+                                eprintln!("Failed to create empty prefetch output: {}", e)
+                            }
+                        }
+                        match std::fs::File::create(&gather_output) {
+                            Ok(_) => {}
+                            Err(e) => eprintln!("Failed to create empty gather output: {}", e),
+                        }
+                    }
                 }
             }
             Err(_) => {

src/index.rs

@@ -4,6 +4,7 @@ use sourmash::prelude::*;
 use std::path::Path;
 
 use crate::utils::{load_collection, ReportType};
+use sourmash::collection::{Collection, CollectionSet};
 
 pub fn index<P: AsRef<Path>>(
     siglist: String,
@@ -13,24 +14,53 @@ pub fn index<P: AsRef<Path>>(
     allow_failed_sigpaths: bool,
     use_internal_storage: bool,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    println!("Loading siglist");
+    eprintln!("Loading sketches from {}", siglist);
 
-    let collection = load_collection(
+    let multi = match load_collection(
         &siglist,
         selection,
         ReportType::General,
         allow_failed_sigpaths,
-    )?;
+    ) {
+        Ok(multi) => multi,
+        Err(err) => return Err(err.into()),
+    };
+    eprintln!("Found {} sketches total.", multi.len());
 
-    let mut index = RevIndex::create(
-        output.as_ref(),
-        collection.select(selection)?.try_into()?,
-        colors,
-    )?;
+    // Try to convert it into a Collection and then CollectionSet.
+    let collection = match Collection::try_from(multi.clone()) {
+        // conversion worked!
+        Ok(c) => {
+            let cs: CollectionSet = c.select(selection)?.try_into()?;
+            Ok(cs)
+        }
+        // conversion failed; can we/should we load it into memory?
+        Err(_) => {
+            if use_internal_storage {
+                eprintln!("WARNING: loading all sketches into memory in order to index.");
+                eprintln!("See 'index' documentation for details.");
+                let c: Collection = multi.load_all_sigs(selection)?;
+                let cs: CollectionSet = c.try_into()?;
+                Ok(cs)
+            } else {
+                Err(
+                    anyhow::anyhow!("cannot index this type of collection with external storage")
+                        .into(),
+                )
+            }
+        }
+    };
 
-    if use_internal_storage {
-        index.internalize_storage()?;
-    }
+    match collection {
+        Ok(collection) => {
+            eprintln!("Indexing {} sketches.", collection.len());
+            let mut index = RevIndex::create(output.as_ref(), collection, colors)?;
 
-    Ok(())
+            if use_internal_storage {
+                index.internalize_storage()?;
+            }
+            Ok(())
+        }
+        Err(e) => Err(e),
+    }
 }