WIP: simplify RevIndex in Python & expand functionality (#3578)

Note: PR into #3545

This PR tackles #1939,
fixing & customizing & simplifying the `index.RevIndex` Python container
wrapping `MemRevIndex` on the Rust side.

A few notes -

* fixes ksize calculation for templating/selection for protein, hp, and
dayhoff sketches;
* fixes Jaccard calculation inside `MemRevIndex` search;
* adds hooks in `sourmash gather` (CLI) and `Index.prefetch(...)`
(Python layer) to test out `RevIndex` as a replacement index class &
`CounterGather` implementation;
* simplifies the Python layer and removes unused Rust code for
`MemRevIndex`;

Overall things seem to work well enough, with a few hiccups around
`scaled` and sketch `location`.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: ctb

include/sourmash.h

@@ -388,20 +388,9 @@ const SourmashSearchResult *const *revindex_gather(const SourmashRevIndex *ptr,
 
 uint64_t revindex_len(const SourmashRevIndex *ptr);
 
-SourmashRevIndex *revindex_new_with_paths(const SourmashStr *const *search_sigs_ptr,
-                                          uintptr_t insigs,
-                                          const SourmashKmerMinHash *template_ptr,
-                                          uintptr_t threshold,
-                                          const SourmashKmerMinHash *const *queries_ptr,
-                                          uintptr_t inqueries,
-                                          bool keep_sigs);
-
 SourmashRevIndex *revindex_new_with_sigs(const SourmashSignature *const *search_sigs_ptr,
                                          uintptr_t insigs,
-                                         const SourmashKmerMinHash *template_ptr,
-                                         uintptr_t threshold,
-                                         const SourmashKmerMinHash *const *queries_ptr,
-                                         uintptr_t inqueries);
+                                         const SourmashKmerMinHash *template_ptr);
 
 ScaledType revindex_scaled(const SourmashRevIndex *ptr);
 

src/core/src/ffi/index/revindex.rs

@@ -1,17 +1,14 @@
 use std::slice;
 
-use camino::Utf8PathBuf as PathBuf;
-
 use crate::encodings::*;
 use crate::ffi::index::SourmashSearchResult;
 use crate::ffi::minhash::SourmashKmerMinHash;
 use crate::ffi::signature::SourmashSignature;
-use crate::ffi::utils::{ForeignObject, SourmashStr};
+use crate::ffi::utils::ForeignObject;
 use crate::index::revindex::mem_revindex;
 use crate::index::Index;
 use crate::prelude::*;
 use crate::signature::{Signature, SigsTrait};
-use crate::sketch::minhash::KmerMinHash;
 use crate::sketch::Sketch;
 use crate::ScaledType;
 
@@ -38,7 +35,7 @@ fn from_template(template: &Sketch) -> Selection {
     let adj_ksize: u32 = match moltype {
         HashFunctions::Murmur64Dna => ksize,
         HashFunctions::Murmur64Protein => ksize / 3,
-        HashFunctions::Murmur64Dayhoff => ksize /  3,
+        HashFunctions::Murmur64Dayhoff => ksize / 3,
         HashFunctions::Murmur64Hp => ksize / 3,
         HashFunctions::Murmur64Skipm1n3 => ksize,
         HashFunctions::Murmur64Skipm2n3 => ksize,
@@ -52,6 +49,7 @@ fn from_template(template: &Sketch) -> Selection {
         .build()
 }
 
+/*
 ffi_fn! {
 unsafe fn revindex_new_with_paths(
     search_sigs_ptr: *const *const SourmashStr,
@@ -105,15 +103,13 @@ unsafe fn revindex_new_with_paths(
     Ok(SourmashRevIndex::from_rust(revindex))
 }
 }
+*/
 
 ffi_fn! {
 unsafe fn revindex_new_with_sigs(
     search_sigs_ptr: *const *const SourmashSignature,
     insigs: usize,
     template_ptr: *const SourmashKmerMinHash,
-    threshold: usize,
-    queries_ptr: *const *const SourmashKmerMinHash,
-    inqueries: usize,
 ) -> Result<*mut SourmashRevIndex> {
     let search_sigs: Vec<Signature> = {
         assert!(!search_sigs_ptr.is_null());
@@ -130,21 +126,8 @@ unsafe fn revindex_new_with_sigs(
         Sketch::MinHash(SourmashKmerMinHash::as_rust(template_ptr).clone())
     };
 
-    let queries_vec: Vec<KmerMinHash>;
-    let queries: Option<&[KmerMinHash]> = if queries_ptr.is_null() {
-        None
-    } else {
-        queries_vec = slice::from_raw_parts(queries_ptr, inqueries)
-            .iter()
-            .map(|mh_ptr|
-            // TODO: avoid this clone
-          SourmashKmerMinHash::as_rust(*mh_ptr).clone())
-            .collect();
-        Some(queries_vec.as_ref())
-    };
-
     let selection = from_template(&template);
-    let revindex = mem_revindex::RevIndex::new_with_sigs(search_sigs, &selection, threshold, queries)?;
+    let revindex = mem_revindex::RevIndex::new_with_sigs(search_sigs, &selection, 0, None)?;
     Ok(SourmashRevIndex::from_rust(revindex))
 }
 }

src/core/src/index/revindex/mem_revindex.rs

@@ -17,6 +17,7 @@ use crate::signature::{Signature, SigsTrait};
 use crate::sketch::minhash::KmerMinHash;
 use crate::sketch::Sketch;
 use crate::Result;
+use crate::ScaledType;
 
 pub struct RevIndex {
     linear: LinearIndex,
@@ -236,6 +237,14 @@ impl RevIndex {
         self.linear.template().clone()
     }
 
+    pub fn scaled(&self) -> ScaledType {
+        if let Sketch::MinHash(mh) = self.linear.template() {
+            mh.clone().scaled() // @CTB avoid clone
+        } else {
+            unimplemented!()
+        }
+    }
+
     // TODO: mh should be a sketch, or even a sig...
     pub(crate) fn find_signatures(
         &self,
@@ -244,10 +253,24 @@ impl RevIndex {
         containment: bool,
         _ignore_scaled: bool,
     ) -> Result<Vec<(f64, Signature, String)>> {
+        let index_scaled = self.scaled();
+        let query_scaled = mh.scaled();
+
+        // @CTB avoid clones?
+        let query_mh = {
+            if query_scaled < index_scaled {
+                mh.clone()
+                    .downsample_scaled(index_scaled)
+                    .expect("cannot downsample query")
+            } else {
+                mh.clone()
+            }
+        };
+
         // TODO: proper threshold calculation
-        let threshold: usize = (threshold * (mh.size() as f64)) as _;
+        let threshold: usize = (threshold * (query_mh.size() as f64)) as _;
 
-        let counter = self.counter_for_query(mh);
+        let counter = self.counter_for_query(&query_mh);
 
         debug!(
             "number of matching signatures for hashes: {}",
@@ -256,7 +279,9 @@ impl RevIndex {
 
         let mut results = vec![];
         for (dataset_id, size) in counter.most_common() {
-            let match_size = if size >= threshold { size } else { break };
+            if size < threshold {
+                break;
+            };
 
             let match_sig = self.linear.sig_for_dataset(dataset_id)?;
             let match_path = self
@@ -266,16 +291,19 @@ impl RevIndex {
                 .internal_location();
 
             let mut match_mh = None;
+            //@CTB use something other than mh here?
             if let Some(Sketch::MinHash(mh)) = match_sig.select_sketch(self.linear.template()) {
                 match_mh = Some(mh);
             }
             let match_mh = match_mh.unwrap();
 
             if size >= threshold {
                 let score = if containment {
-                    size as f64 / mh.size() as f64
+                    size as f64 / query_mh.size() as f64
                 } else {
-                    mh.jaccard(match_mh).expect("cannot calculate Jaccard")
+                    query_mh
+                        .jaccard(match_mh)
+                        .expect("cannot calculate Jaccard")
                 };
                 let filename = match_path.to_string();
                 let mut sig: Signature = match_sig.clone().into();

src/sourmash/commands.py

@@ -942,8 +942,9 @@ def gather(args):
 
     if args.linear:  # force linear traversal?
         databases = [LazyLinearIndex(db) for db in databases]
-    else:
+    elif 0:
         # @CTB foo revindex
+        print("XXX NOTE: using RevIndex for ZipFileLinearIndex")
         from sourmash.index.revindex import RevIndex
         from sourmash.index import ZipFileLinearIndex
 

src/sourmash/index/__init__.py

@@ -320,14 +320,23 @@ def counter_gather(self, query, threshold_bp, **kwargs):
             # tada!
             return counter
         else:
-            from .revindex import RevIndex
+            print("XXX NOTE: Using RevIndex CounterGather")
+            from .revindex import RevIndex_CounterGather, RevIndex
 
             revindex = RevIndex(template=prefetch_query.minhash)
+            cg = RevIndex_CounterGather(
+                prefetch_query, revindex, threshold_bp, allow_insert=True
+            )
 
+            n_added = 0
             for result in self.prefetch(prefetch_query, threshold_bp, **kwargs):
-                revindex.insert(result.signature)
+                cg.add(result.signature, location=result.location)
+                n_added += 1
+
+            if n_added == 0:
+                raise ValueError("no signatures to count")
 
-            return revindex.counter_gather(prefetch_query, threshold_bp)
+            return cg
 
     @abstractmethod
     def select(