feat: improve term search with full-text index (#109) (#112)

Cargo.toml

@@ -17,6 +17,7 @@ console = "0.15"
 csv = "1.3"
 derive_more = "0.99"
 env_logger = "0.10"
+fastobo = "0.15.2"
 fastrand = "2.0"
 flate2 = "1.0"
 hpo = "0.8"
@@ -37,6 +38,8 @@ serde_with = { version = "3.3", features=["alloc", "macros", "indexmap_2"], defa
 shellexpand = "3.0"
 strum = { version = "0.25", features = ["strum_macros", "derive"] }
 strum_macros = "0.25"
+tantivy = "0.21.1"
+tempdir = "0.3.7"
 thiserror = "1.0"
 tracing = "0.1"
 tracing-subscriber = "0.3"

src/index.rs

@@ -0,0 +1,191 @@
+//! Full text index for OBO documents using tantivy.
+
+use tantivy::schema::Schema;
+
+/// Encapsulation of a Tantivy index.
+///
+/// The index lives in a directory and will only persist as long as the lifetime of this struct.
+pub struct Index {
+ /// The index directory.
+ #[allow(dead_code)]
+ tmpdir: tempdir::TempDir, // keep around for RAII
+ /// The HPO document.
+ hpo_doc: fastobo::ast::OboDoc,
+ /// The index.
+ index: tantivy::Index,
+ /// The index schema.
+ schema: tantivy::schema::Schema,
+ /// The single reader.
+ reader: tantivy::IndexReader,
+}
+
+/// Convert ident to String.
+fn ident_to_string(ident: &fastobo::ast::Ident) -> String {
+ match ident {
+ fastobo::ast::Ident::Prefixed(val) => format!("{}:{}", val.prefix(), val.local()),
+ fastobo::ast::Ident::Unprefixed(val) => val.as_str().to_string(),
+ fastobo::ast::Ident::Url(val) => val.as_str().to_string(),
+ }
+}
+
+// Code for creating an `Index`.
+impl Index {
+ /// Create a new index from an OBO document.
+ ///
+ /// # Args
+ ///
+ /// * `hpo_doc` - The OBO document to index.
+ ///
+ /// # Errors
+ ///
+ /// In the case that the index cannot be created.
+ pub fn new(hpo_doc: fastobo::ast::OboDoc) -> Result<Self, anyhow::Error> {
+ let schema = Self::build_schema();
+ let tmpdir = tempdir::TempDir::new("viguno")?;
+ let index = tantivy::Index::create_in_dir(tmpdir.path(), schema.clone())?;
+
+ let mut index_writer = index.writer(100_000_000).map_err(|e| {
+ anyhow::anyhow!(
+ "Error creating tantivy index writer for directory {:?}: {}",
+ tmpdir.path(),
+ e
+ )
+ })?;
+ Self::write_hpo_index(&hpo_doc, &schema, &mut index_writer).map_err(|e| {
+ anyhow::anyhow!(
+ "Error writing HPO index for directory {:?}: {}",
+ tmpdir.path(),
+ e
+ )
+ })?;
+
+ let reader = index
+ .reader_builder()
+ .reload_policy(tantivy::ReloadPolicy::Manual)
+ .try_into()?;
+
+ Ok(Self {
+ tmpdir,
+ hpo_doc,
+ index,
+ schema,
+ reader,
+ })
+ }
+
+ /// Build the tantivy schema for the HPO.
+ fn build_schema() -> Schema {
+ use tantivy::schema::{STORED, STRING, TEXT};
+
+ let mut schema_builder = Schema::builder();
+ schema_builder.add_text_field("term_id", STRING | STORED);
+ schema_builder.add_text_field("alt_id", STRING | STORED);
+ schema_builder.add_text_field("name", TEXT | STORED);
+ schema_builder.add_text_field("def", TEXT | STORED);
+ schema_builder.add_text_field("synonym", TEXT | STORED);
+ schema_builder.add_text_field("xref", STRING | STORED);
+ schema_builder.build()
+ }
+
+ /// Index the HPO document.
+ fn write_hpo_index(
+ hpo_doc: &fastobo::ast::OboDoc,
+ schema: &tantivy::schema::Schema,
+ index_writer: &mut tantivy::IndexWriter,
+ ) -> Result<(), anyhow::Error> {
+ for term_frame in hpo_doc
+ .entities()
+ .iter()
+ .filter_map(fastobo::ast::EntityFrame::as_term)
+ {
+ let mut doc = tantivy::Document::default();
+
+ doc.add_field_value(
+ schema.get_field("term_id")?,
+ ident_to_string(term_frame.id().as_inner().as_ref()),
+ );
+
+ for line in term_frame
+ .clauses()
+ .iter()
+ .map(fastobo::ast::Line::as_inner)
+ {
+ match line {
+ fastobo::ast::TermClause::Name(name) => {
+ doc.add_field_value(schema.get_field("name")?, name.as_str());
+ }
+ fastobo::ast::TermClause::AltId(alt_id) => {
+ doc.add_field_value(
+ schema.get_field("alt_id")?,
+ ident_to_string(alt_id).as_str(),
+ );
+ }
+ fastobo::ast::TermClause::Def(def) => {
+ doc.add_field_value(schema.get_field("def")?, def.text().as_str());
+ }
+ fastobo::ast::TermClause::Synonym(synonym) => {
+ doc.add_field_value(
+ schema.get_field("synonym")?,
+ synonym.description().as_str(),
+ );
+ }
+ fastobo::ast::TermClause::Xref(xref) => {
+ doc.add_field_value(schema.get_field("xref")?, ident_to_string(xref.id()));
+ }
+ fastobo::ast::TermClause::IsAnonymous(_)
+ | fastobo::ast::TermClause::Comment(_)
+ | fastobo::ast::TermClause::Namespace(_)
+ | fastobo::ast::TermClause::Subset(_)
+ | fastobo::ast::TermClause::Builtin(_)
+ | fastobo::ast::TermClause::PropertyValue(_)
+ | fastobo::ast::TermClause::IsA(_)
+ | fastobo::ast::TermClause::IntersectionOf(_, _)
+ | fastobo::ast::TermClause::UnionOf(_)
+ | fastobo::ast::TermClause::EquivalentTo(_)
+ | fastobo::ast::TermClause::DisjointFrom(_)
+ | fastobo::ast::TermClause::Relationship(_, _)
+ | fastobo::ast::TermClause::CreatedBy(_)
+ | fastobo::ast::TermClause::CreationDate(_)
+ | fastobo::ast::TermClause::IsObsolete(_)
+ | fastobo::ast::TermClause::ReplacedBy(_)
+ | fastobo::ast::TermClause::Consider(_) => (),
+ }
+ }
+
+ index_writer.add_document(doc).map_err(|e| {
+ anyhow::anyhow!(
+ "Error adding document to tantivy index writer: {}",
+ e.to_string()
+ )
+ })?;
+ }
+
+ index_writer
+ .commit()
+ .map_err(|e| anyhow::anyhow!("Error committing tantivy index writer: {}", e))?;
+ Ok(())
+ }
+}
+
+// Accessor code.
+impl Index {
+ /// Get the HPO document.
+ pub fn hpo_doc(&self) -> &fastobo::ast::OboDoc {
+ &self.hpo_doc
+ }
+
+ /// Get the index.
+ pub fn index(&self) -> &tantivy::Index {
+ &self.index
+ }
+
+ /// Get the schema.
+ pub fn schema(&self) -> &tantivy::schema::Schema {
+ &self.schema
+ }
+
+ /// Get the reader.
+ pub fn reader(&self) -> &tantivy::IndexReader {
+ &self.reader
+ }
+}

src/main.rs

@@ -7,6 +7,7 @@
 pub mod algos;
 pub mod common;
 pub mod convert;
+pub mod index;
 pub mod pbs;
 pub mod query;
 pub mod server;

src/server/actix_server/hpo_genes.rs

@@ -192,13 +192,15 @@ mod test {
  let ncbi_to_hgnc =
  crate::common::hgnc_xlink::load_ncbi_to_hgnc("tests/data/hgnc_xlink.tsv")?;
  let hgnc_to_ncbi = crate::common::hgnc_xlink::inverse_hashmap(&ncbi_to_hgnc);
+ let hpo_doc = fastobo::from_file("tests/data/hpo/hp.obo")?;
  let app = actix_web::test::init_service(
  actix_web::App::new()
  .app_data(actix_web::web::Data::new(crate::server::WebServerData {
  ontology,
  db: None,
  ncbi_to_hgnc,
  hgnc_to_ncbi,
+ full_text_index: crate::index::Index::new(hpo_doc)?,
  }))
  .service(super::handle),
  )

src/server/actix_server/hpo_omims.rs

@@ -219,13 +219,15 @@ mod test {
  let ncbi_to_hgnc =
  crate::common::hgnc_xlink::load_ncbi_to_hgnc("tests/data/hgnc_xlink.tsv")?;
  let hgnc_to_ncbi = crate::common::hgnc_xlink::inverse_hashmap(&ncbi_to_hgnc);
+ let hpo_doc = fastobo::from_file("tests/data/hpo/hp.obo")?;
  let app = actix_web::test::init_service(
  actix_web::App::new()
  .app_data(actix_web::web::Data::new(crate::server::WebServerData {
  ontology,
  db: None,
  ncbi_to_hgnc,
  hgnc_to_ncbi,
+ full_text_index: crate::index::Index::new(hpo_doc)?,
  }))
  .service(super::handle),
  )

src/server/actix_server/hpo_sim/term_gene.rs

@@ -118,6 +118,7 @@ mod test {
  ["meta", "scores"],
  true,
  )?);
+ let hpo_doc = fastobo::from_file("tests/data/hpo/hp.obo")?;
 
  let app = actix_web::test::init_service(
  actix_web::App::new()
@@ -126,6 +127,7 @@ mod test {
  db,
  ncbi_to_hgnc,
  hgnc_to_ncbi,
+ full_text_index: crate::index::Index::new(hpo_doc)?,
  }))
  .service(super::handle),
  )

src/server/actix_server/hpo_sim/term_term.rs

@@ -173,6 +173,7 @@ mod test {
  ["meta", "scores"],
  true,
  )?);
+ let hpo_doc = fastobo::from_file("tests/data/hpo/hp.obo")?;
 
  let app = actix_web::test::init_service(
  actix_web::App::new()
@@ -181,6 +182,7 @@ mod test {
  db,
  ncbi_to_hgnc,
  hgnc_to_ncbi,
+ full_text_index: crate::index::Index::new(hpo_doc)?,
  }))
  .service(super::handle),
  )