diff --git a/src/fastgather.rs b/src/fastgather.rs index 46512025..08a1ebbd 100644 --- a/src/fastgather.rs +++ b/src/fastgather.rs @@ -1,4 +1,5 @@ /// fastgather: Run gather with a query against a list of files. +use pyo3::Python; use anyhow::Result; use sourmash::prelude::Select; use sourmash::selection::Selection; @@ -10,6 +11,7 @@ use crate::utils::{ #[allow(clippy::too_many_arguments)] pub fn fastgather( + py: Python, query_filepath: String, against_filepath: String, threshold_bp: usize, @@ -94,6 +96,7 @@ pub fn fastgather( // run the gather! consume_query_by_gather( + py, query_sig, scaled as u64, matchlist, diff --git a/src/fastmultigather.rs b/src/fastmultigather.rs index 608fc215..e58b6aac 100644 --- a/src/fastmultigather.rs +++ b/src/fastmultigather.rs @@ -1,6 +1,7 @@ /// fastmultigather: Run gather for multiple queries against a list of files. use anyhow::Result; use rayon::prelude::*; +use pyo3::Python; use sourmash::selection::Selection; @@ -17,6 +18,7 @@ use crate::utils::{ }; pub fn fastmultigather( + py: Python, query_filepath: String, against_filepath: String, threshold_bp: usize, @@ -98,6 +100,7 @@ pub fn fastmultigather( // Now, do the gather! consume_query_by_gather( + py, query_sig.clone(), scaled as u64, matchlist, diff --git a/src/lib.rs b/src/lib.rs index e8d22f2a..4567e03f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,6 +23,7 @@ use camino::Utf8PathBuf as PathBuf; #[pyfunction] fn do_manysearch( + py: Python, querylist_path: String, siglist_path: String, threshold: f64, @@ -73,6 +74,7 @@ fn do_manysearch( #[pyfunction] #[allow(clippy::too_many_arguments)] fn do_fastgather( + py: Python, query_filename: String, siglist_path: String, threshold_bp: usize, @@ -86,6 +88,7 @@ fn do_fastgather( let allow_failed_sigpaths = true; match fastgather::fastgather( + py, query_filename, siglist_path, threshold_bp, @@ -105,6 +108,7 @@ fn do_fastgather( #[pyfunction] fn do_fastmultigather( + py: Python, query_filenames: String, siglist_path: String, threshold_bp: usize, @@ -138,6 +142,7 @@ fn do_fastmultigather( bail!("output path specified, but not running fastmultigather against a rocksdb. See issue #239"); } match fastmultigather::fastmultigather( + py, query_filenames, siglist_path, threshold_bp, @@ -173,6 +178,7 @@ fn set_global_thread_pool(num_threads: usize) -> PyResult { #[pyfunction] fn do_index( + py: Python, siglist: String, ksize: u8, scaled: usize, @@ -206,6 +212,7 @@ fn do_check(index: String, quick: bool) -> anyhow::Result { #[pyfunction] #[allow(clippy::too_many_arguments)] fn do_multisearch( + py: Python, querylist_path: String, siglist_path: String, threshold: f64, @@ -238,6 +245,7 @@ fn do_multisearch( #[pyfunction] #[allow(clippy::too_many_arguments)] fn do_pairwise( + py: Python, siglist_path: String, threshold: f64, ksize: u8, @@ -268,6 +276,7 @@ fn do_pairwise( #[pyfunction] fn do_manysketch( + py: Python, filelist: String, param_str: String, output: String, @@ -285,6 +294,7 @@ fn do_manysketch( #[pyfunction] fn do_cluster( + py: Python, pairwise_csv: String, output_clusters: String, similarity_column: String, diff --git a/src/utils.rs b/src/utils.rs index 89779133..ae123abb 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,5 +1,6 @@ /// Utility functions for sourmash_plugin_branchwater. use rayon::prelude::*; +use pyo3::Python; use sourmash::encodings::HashFunctions; use sourmash::selection::Select; @@ -926,6 +927,7 @@ pub fn branchwater_calculate_gather_stats( /// removing matches in 'matchlist' from 'query'. pub fn consume_query_by_gather( + py: Python, query: SigStore, scaled: u64, matchlist: BinaryHeap, @@ -994,6 +996,7 @@ pub fn consume_query_by_gather( ); while !matching_sketches.is_empty() { + py.check_signals()?; let best_element = matching_sketches.peek().unwrap(); query_mh = query_mh.downsample_scaled(best_element.minhash.scaled())?;