diff --git a/src/lexical.rs b/src/lexical.rs index 4b0fd9d..8a593ca 100644 --- a/src/lexical.rs +++ b/src/lexical.rs @@ -1 +1,74 @@ -// implemented in a later task +use crate::model::{AlphaIdEntry, AppError}; +use crate::normalize::normalize; +use tantivy::collector::TopDocs; +use tantivy::query::QueryParser; +use tantivy::schema::{Schema, Value, STORED, TEXT}; +use tantivy::{doc, Index, TantivyDocument}; + +pub struct LexicalIndex { + index: Index, + text_field: tantivy::schema::Field, + idx_field: tantivy::schema::Field, +} + +impl LexicalIndex { + fn schema() -> (Schema, tantivy::schema::Field, tantivy::schema::Field) { + let mut b = Schema::builder(); + let text_field = b.add_text_field("text", TEXT); + let idx_field = b.add_u64_field("idx", STORED); + (b.build(), text_field, idx_field) + } + + pub fn build_in_ram(entries: &[AlphaIdEntry]) -> Result { + let (schema, text_field, idx_field) = Self::schema(); + let index = Index::create_in_ram(schema); + let mut w = index + .writer(50_000_000) + .map_err(|e| AppError::Io(format!("tantivy writer: {e}")))?; + for (i, e) in entries.iter().enumerate() { + w.add_document(doc!(text_field => normalize(&e.text), idx_field => i as u64)) + .map_err(|e| AppError::Io(format!("tantivy add: {e}")))?; + } + w.commit() + .map_err(|e| AppError::Io(format!("tantivy commit: {e}")))?; + Ok(Self { + index, + text_field, + idx_field, + }) + } + + /// Returns (entry_index, bm25_score) descending. + pub fn search(&self, query: &str, k: usize) -> Result, AppError> { + let reader = self + .index + .reader() + .map_err(|e| AppError::Io(format!("tantivy reader: {e}")))?; + let searcher = reader.searcher(); + let norm = normalize(query); + let sanitized = norm.replace(['(', ')', '"', ':'], " "); + // Use OR semantics (default) for recall-favoring behaviour; fall back to + // the raw normalised query on parse error. + let parsed = QueryParser::for_index(&self.index, vec![self.text_field]) + .parse_query(&sanitized) + .or_else(|_| { + let mut qp = QueryParser::for_index(&self.index, vec![self.text_field]); + qp.set_conjunction_by_default(); + qp.parse_query(&norm) + }) + .map_err(|e| AppError::Parse(format!("query: {e}")))?; + let top = searcher + .search(&parsed, &TopDocs::with_limit(k.max(1))) + .map_err(|e| AppError::Io(format!("tantivy search: {e}")))?; + let mut out = Vec::new(); + for (score, addr) in top { + let d: TantivyDocument = searcher + .doc(addr) + .map_err(|e| AppError::Io(format!("tantivy doc: {e}")))?; + if let Some(v) = d.get_first(self.idx_field).and_then(|v| v.as_u64()) { + out.push((v as usize, score)); + } + } + Ok(out) + } +} diff --git a/tests/lexical_tests.rs b/tests/lexical_tests.rs new file mode 100644 index 0000000..f1c2db4 --- /dev/null +++ b/tests/lexical_tests.rs @@ -0,0 +1,23 @@ +use alpha_id::lexical::LexicalIndex; +use alpha_id::model::AlphaIdEntry; + +fn entry(id: &str, code: &str, text: &str) -> AlphaIdEntry { + AlphaIdEntry { + alpha_id: id.into(), valid: true, + icd_primary: code.into(), icd_star: "".into(), icd_addon: "".into(), + icd_primary2: "".into(), orpha: "".into(), text: text.into(), + } +} + +#[test] +fn finds_entry_by_fuzzy_query() { + let entries = vec![ + entry("I1", "E11.9", "Diabetes mellitus Typ 2"), + entry("I2", "I10.90", "Arterielle essentielle Hypertonie"), + ]; + let idx = LexicalIndex::build_in_ram(&entries).unwrap(); + let hits = idx.search("diabetes typ 2", 5).unwrap(); + assert!(!hits.is_empty()); + assert_eq!(hits[0].0, 0); // entry index 0 + assert!(hits[0].1 > 0.0); // bm25 score +}