diff --git a/src/command_line/search/basic_search_count.rs b/src/command_line/search/basic_search_count.rs new file mode 100644 index 0000000..1e6dd09 --- /dev/null +++ b/src/command_line/search/basic_search_count.rs @@ -0,0 +1,40 @@ +use crate::command_line::prelude::*; +use crate::search::basic_search::basic_search_count; + +pub const NAME: &str = "basic-search-count"; + +pub fn command() -> Command { + Command::new(NAME) + .arg( + Arg::new("index") + .required(true) + .long("index") + .short('i') + .num_args(1), + ) + .arg( + Arg::new("query") + .required(true) + .long("query") + .short('q') + .num_args(1), + ) +} + +pub fn action(matches: &ArgMatches) -> eyre::Result<()> { + let index_path = matches + .get_one::("index") + .ok_or(eyre::eyre!("Failed to extract index path"))?; + let query = matches + .get_one::("query") + .ok_or(eyre::eyre!("Failed to extract query"))?; + + let index = open_index(index_path)?; + let reader = index.reader()?; + let searcher = reader.searcher(); + let result_count = basic_search_count(&searcher, query)?; + + log::info!("{:#?}", result_count); + + Ok(()) +} diff --git a/src/command_line/search/mod.rs b/src/command_line/search/mod.rs index dbfbc6a..add6c99 100644 --- a/src/command_line/search/mod.rs +++ b/src/command_line/search/mod.rs @@ -1,4 +1,5 @@ pub mod basic_search; +pub mod basic_search_count; pub mod cli_structure_search; pub mod identity_search; pub mod similarity_search; diff --git a/src/main.rs b/src/main.rs index f728d27..3d10a3f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -24,6 +24,7 @@ async fn main() -> eyre::Result<()> { .subcommand(command_line::pubchem::fetch_pubchem::command()) .subcommand(command_line::pubchem::stream_pubchem_sdf::command()) .subcommand(command_line::search::basic_search::command()) + .subcommand(command_line::search::basic_search_count::command()) .subcommand(command_line::search::substructure_search::command()) .subcommand(command_line::search::superstructure_search::command()) .subcommand(command_line::search::similarity_search::command()) @@ -59,6 +60,9 @@ async fn main() -> eyre::Result<()> { (command_line::search::basic_search::NAME, matches) => { command_line::search::basic_search::action(matches) } + (command_line::search::basic_search_count::NAME, matches) => { + command_line::search::basic_search_count::action(matches) + } (command_line::search::substructure_search::NAME, matches) => { command_line::search::substructure_search::action(matches) } diff --git a/src/search/basic_search.rs b/src/search/basic_search.rs index ca72ca3..da42732 100644 --- a/src/search/basic_search.rs +++ b/src/search/basic_search.rs @@ -1,6 +1,7 @@ use crate::search::sort_docs; use rayon::prelude::*; -use tantivy::{collector::TopDocs, query::QueryParser, DocAddress, Searcher}; +use tantivy::collector::{Count, TopDocs}; +use tantivy::{query::QueryParser, DocAddress, Searcher}; #[allow(clippy::ptr_arg)] pub fn basic_search( @@ -20,3 +21,13 @@ pub fn basic_search( sort_docs(&mut final_results); Ok(final_results) } + +#[allow(clippy::ptr_arg)] +pub fn basic_search_count(searcher: &Searcher, query: &String) -> eyre::Result { + let index = searcher.index(); + let query_parser = QueryParser::for_index(index, vec![]); + let query = query_parser.parse_query(query)?; + let result_count = searcher.search(&query, &Count)?; + + Ok(result_count) +}