68 lines
1.8 KiB
Rust
68 lines
1.8 KiB
Rust
use anyhow::Result;
|
|
|
|
use crate::index::{LocalIndex, SearchHit};
|
|
|
|
pub struct EngineOutput {
|
|
pub hits: Vec<SearchHit>,
|
|
pub total: Option<u64>,
|
|
}
|
|
|
|
pub trait SearchEngine: Send + Sync {
|
|
fn search(&self, text: &str, budget: usize, only_shared: bool) -> Result<EngineOutput>;
|
|
fn doc_count(&self) -> u64;
|
|
}
|
|
|
|
pub struct TantivyEngine {
|
|
pub index: LocalIndex,
|
|
pub min_coverage: f64,
|
|
}
|
|
|
|
impl SearchEngine for TantivyEngine {
|
|
fn search(&self, text: &str, budget: usize, only_shared: bool) -> Result<EngineOutput> {
|
|
let (hits, total) =
|
|
self.index
|
|
.search_with_coverage(text, budget, only_shared, self.min_coverage)?;
|
|
Ok(EngineOutput {
|
|
hits,
|
|
total: Some(total),
|
|
})
|
|
}
|
|
|
|
fn doc_count(&self) -> u64 {
|
|
self.index.doc_count()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::index::Collection;
|
|
use crate::message::EXPOSURE_FULL;
|
|
use std::fs;
|
|
|
|
#[test]
|
|
fn tantivy_engine_reports_hits_and_total() {
|
|
let temp = tempfile::tempdir().unwrap();
|
|
let dir = temp.path().join("corpus");
|
|
fs::create_dir_all(&dir).unwrap();
|
|
fs::write(dir.join("a.txt"), "rust ownership guide").unwrap();
|
|
let engine = TantivyEngine {
|
|
index: LocalIndex::open(&temp.path().join("index")).unwrap(),
|
|
min_coverage: 0.4,
|
|
};
|
|
engine
|
|
.index
|
|
.add_collection(&Collection {
|
|
name: "docs".to_string(),
|
|
path: dir.display().to_string(),
|
|
shared: true,
|
|
exposure: EXPOSURE_FULL.to_string(),
|
|
})
|
|
.unwrap();
|
|
let output = engine.search("rust", 5, false).unwrap();
|
|
assert_eq!(output.hits.len(), 1);
|
|
assert_eq!(output.total, Some(1));
|
|
assert_eq!(engine.doc_count(), 1);
|
|
}
|
|
}
|