EXP: implement some intermediate APIs on top of Rust using pyo3

.cargo/config.toml

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

Cargo.toml

@@ -12,7 +12,8 @@ crate-type = ["cdylib"]
 pyo3 = { version = "0.22.2", features = ["extension-module", "anyhow"] }
 rayon = "1.10.0"
 serde = { version = "1.0.208", features = ["derive"] }
-sourmash = { version = "0.15.0", features = ["branchwater"] }
+sourmash = { git = "https://github.com/sourmash-bio/sourmash.git", branch = "more_rs_updates", features = ["branchwater"] }
+#sourmash = { version = "0.15.0", features = ["branchwater"] }
 serde_json = "1.0.125"
 niffler = "2.4.0"
 log = "0.4.22"

pyproject.toml

@@ -21,6 +21,9 @@ authors = [
 requires = ["maturin>=1.4.0,<2"]
 build-backend = "maturin"
 
+[project.entry-points."sourmash.load_from"]
+collection_reader = "sourmash_plugin_branchwater:load_collection"
+
 [project.entry-points."sourmash.cli_script"]
 manysearch = "sourmash_plugin_branchwater:Branchwater_Manysearch"
 multisearch = "sourmash_plugin_branchwater:Branchwater_Multisearch"

src/branch_api.rs

@@ -0,0 +1,165 @@
+/// Lower-level Python API implementation for sourmash_plugin_branchwater
+use pyo3::prelude::*;
+
+use crate::utils::build_selection;
+use crate::utils::load_collection;
+use crate::utils::ReportType;
+use crate::utils::multicollection::MultiCollection;
+use sourmash::collection::Collection;
+use sourmash::manifest::{Manifest, Record};
+use pyo3::types::{IntoPyDict, PyDict, PyList};
+
+#[pyclass]
+pub struct BranchRecord {
+    record: Record,
+}
+
+#[pymethods]
+impl BranchRecord {
+    pub fn get_name(&self) -> PyResult<String> {
+        Ok(self.record.name().clone())
+    }
+
+    #[getter]
+    pub fn get_as_row<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
+        let dict = {
+            let key_vals: Vec<(&str, PyObject)> = vec![
+                ("ksize", self.record.ksize().to_object(py)),
+                ("moltype", self.record.moltype().to_string().to_object(py)),
+                ("scaled", self.record.scaled().to_object(py)),
+                ("num", self.record.num().to_object(py)),
+                ("with_abundance", self.record.with_abundance().to_object(py)),
+                ("n_hashes", self.record.n_hashes().to_object(py)),
+            ];
+            key_vals.into_py_dict_bound(py)
+        };
+        Ok(dict)
+    }
+}
+
+/*
+impl<T, I> IntoPyDict for I
+where
+    T: PyDictItem
+    I: IntoIterator<Item = T>
+fn into_py_dict(self, py: Python<'_>) -> Bound<'_, PyDict> {
+    let dict = PyDict::new(py);
+    for item in self {
+        dict.set_item(item.key(), item.value())
+            .expect("Failed to set_item on dict");
+    }
+    dict
+}
+}
+*/
+
+#[pyclass]
+pub struct BranchManifest {
+    manifest: Manifest,
+}
+
+#[pymethods]
+impl BranchManifest {
+    pub fn __len__(&self) -> PyResult<usize> {
+        Ok(self.manifest.len())
+    }
+    pub fn _check_row_values(&self) -> PyResult<bool> {
+        Ok(true)
+    }
+    #[getter]
+    pub fn get_rows<'py>(&self, py: Python<'py>) -> PyResult<Vec<Bound<'py, PyDict>>> {
+        let res: Vec<_> = self.manifest.iter().map(|x| { BranchRecord {
+            record: x.clone(),
+        }.get_as_row(py).unwrap()
+        }).collect();
+
+        Ok(res)
+    }
+}
+
+#[pyclass]
+pub struct BranchCollection {
+    #[pyo3(get)]
+    pub location: String,
+
+    #[pyo3(get)]
+    pub is_database: bool,
+
+    #[pyo3(get)]
+    pub has_manifest: bool,
+
+    collection: MultiCollection,
+}
+
+#[pymethods]
+impl BranchCollection {
+    pub fn __len__(&self) -> PyResult<usize> {
+        Ok(self.collection.len())
+    }
+
+    #[getter]
+    pub fn get_manifest(&self) -> PyResult<Py<BranchManifest>> {
+        let manifest: Manifest = self.collection.manifest().clone();
+        let obj =
+            Python::with_gil(|py| Py::new(py, BranchManifest { manifest: manifest }).unwrap());
+        Ok(obj)
+    }
+    pub fn get_first_record(&self) -> PyResult<Py<BranchRecord>> {
+        let records: Vec<_> = self.collection.iter().collect();
+        let first_record = records.first().unwrap().1;
+
+        // @CTB: can I turn this into something automatic?
+        let obj = Python::with_gil(|py| {
+            Py::new(
+                py,
+                BranchRecord {
+                    record: first_record.clone(),
+                },
+            )
+            .unwrap()
+        });
+        Ok(obj)
+    }
+
+    #[getter]
+    pub fn get_rows(&self) -> PyResult<Vec<BranchRecord>> {
+        let records: Vec<_> = self.collection.iter().collect();
+
+        let obj = records
+                .iter()
+                .map(|x| {
+                    BranchRecord {
+                        record: x.1.clone(),
+                    }
+                })
+            .collect();
+
+        // @CTB: this does the GIL grabbing as needed?
+        Ok(obj)
+    }
+}
+
+#[pyfunction]
+pub fn api_load_collection(
+    location: String,
+    ksize: u8,
+    scaled: usize,
+    moltype: String,
+) -> PyResult<Py<BranchCollection>> {
+    let selection = build_selection(ksize, scaled, &moltype);
+
+    let collection = load_collection(&location, &selection, ReportType::Query, true).unwrap();
+    let obj = Python::with_gil(|py| {
+        Py::new(
+            py,
+            BranchCollection {
+                location: location,
+                collection,
+                is_database: false,
+                has_manifest: true,
+            },
+        )
+        .unwrap()
+    });
+    Ok(obj)
+}

src/fastgather.rs

@@ -33,7 +33,7 @@ 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().unwrap();
     let query_sig_ds = query_sig.clone().select(selection)?; // downsample
     let query_mh = match query_sig_ds.minhash() {
         Some(query_mh) => query_mh,

src/fastmultigather.rs

@@ -69,11 +69,11 @@ pub fn fastmultigather(
     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) {
+        match c.sig_from_record(record) {
             Ok(query_sig) => {
                 let name = query_sig.name();
                 let prefix = name.split(' ').next().unwrap_or_default().to_string();
@@ -133,7 +133,7 @@ pub fn fastmultigather(
                                 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.scaled(),
                                         query_mh.ksize().try_into().unwrap(),
                                         query_mh.hash_function().clone(),
                                         query_mh.seed(),

src/lib.rs

@@ -1,9 +1,14 @@
-/// Python interface Rust code for sourmash_plugin_branchwater.
+//! Rust-to-Pyton interface code for sourmash_plugin_branchwater, using pyo3.
+//!
+//! If you're using Rust, you're probably most interested in
+//! [utils](utils/index.html)
+
 use pyo3::prelude::*;
 
 #[macro_use]
 extern crate simple_error;
 
+mod branch_api;
 mod utils;
 use crate::utils::build_selection;
 use crate::utils::is_revindex_database;
@@ -106,6 +111,7 @@ fn do_fastgather(
 }
 
 #[pyfunction]
+#[allow(clippy::too_many_arguments)]
 #[pyo3(signature = (query_filenames, siglist_path, threshold_bp, ksize, scaled, moltype, output_path=None, save_matches=false))]
 fn do_fastmultigather(
     query_filenames: String,
@@ -322,8 +328,11 @@ fn do_cluster(
     }
 }
 
+/// Module interface for the `sourmash_plugin_branchwater` extension module.
+
 #[pymodule]
 fn sourmash_plugin_branchwater(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // top level 'scripts' commands
     m.add_function(wrap_pyfunction!(do_manysearch, m)?)?;
     m.add_function(wrap_pyfunction!(do_fastgather, m)?)?;
     m.add_function(wrap_pyfunction!(do_fastmultigather, m)?)?;
@@ -334,5 +343,10 @@ fn sourmash_plugin_branchwater(_py: Python, m: &Bound<'_, PyModule>) -> PyResult
     m.add_function(wrap_pyfunction!(do_multisearch, m)?)?;
     m.add_function(wrap_pyfunction!(do_pairwise, m)?)?;
     m.add_function(wrap_pyfunction!(do_cluster, m)?)?;
+
+    // lower level API stuff
+    m.add_class::<branch_api::BranchCollection>()?;
+    m.add_function(wrap_pyfunction!(branch_api::api_load_collection, m)?)?;
+
     Ok(())
 }