meilisearch/milli/src/search/mod.rs

229 lines
6.7 KiB
Rust
Raw Normal View History

use std::fmt;
use levenshtein_automata::{LevenshteinAutomatonBuilder as LevBuilder, DFA};
use once_cell::sync::Lazy;
use roaring::bitmap::RoaringBitmap;
pub use self::facet::{FacetDistribution, Filter, DEFAULT_VALUES_PER_FACET};
pub use self::new::matches::{FormatOptions, MatchBounds, Matcher, MatcherBuilder, MatchingWords};
use self::new::PartialSearchResult;
2023-03-23 16:35:53 +08:00
use crate::{
execute_search, AscDesc, DefaultSearchLogger, DocumentId, Index, Result, SearchContext,
};
// Building these factories is not free.
static LEVDIST0: Lazy<LevBuilder> = Lazy::new(|| LevBuilder::new(0, true));
static LEVDIST1: Lazy<LevBuilder> = Lazy::new(|| LevBuilder::new(1, true));
static LEVDIST2: Lazy<LevBuilder> = Lazy::new(|| LevBuilder::new(2, true));
pub mod facet;
2022-03-16 00:28:57 +08:00
mod fst_utils;
2023-02-21 16:40:41 +08:00
pub mod new;
pub struct Search<'a> {
query: Option<String>,
// this should be linked to the String in the query
filter: Option<Filter<'a>>,
offset: usize,
limit: usize,
2021-08-23 17:37:18 +08:00
sort_criteria: Option<Vec<AscDesc>>,
2023-04-13 19:45:34 +08:00
geo_strategy: new::GeoSortStrategy,
terms_matching_strategy: TermsMatchingStrategy,
2021-04-14 01:10:58 +08:00
words_limit: usize,
exhaustive_number_hits: bool,
rtxn: &'a heed::RoTxn<'a>,
index: &'a Index,
}
impl<'a> Search<'a> {
pub fn new(rtxn: &'a heed::RoTxn, index: &'a Index) -> Search<'a> {
Search {
query: None,
filter: None,
offset: 0,
limit: 20,
2021-08-23 17:37:18 +08:00
sort_criteria: None,
2023-04-13 19:45:34 +08:00
geo_strategy: new::GeoSortStrategy::default(),
terms_matching_strategy: TermsMatchingStrategy::default(),
exhaustive_number_hits: false,
2021-04-14 01:10:58 +08:00
words_limit: 10,
rtxn,
index,
}
}
pub fn query(&mut self, query: impl Into<String>) -> &mut Search<'a> {
self.query = Some(query.into());
self
}
pub fn offset(&mut self, offset: usize) -> &mut Search<'a> {
self.offset = offset;
self
}
pub fn limit(&mut self, limit: usize) -> &mut Search<'a> {
self.limit = limit;
self
}
2021-08-23 17:37:18 +08:00
pub fn sort_criteria(&mut self, criteria: Vec<AscDesc>) -> &mut Search<'a> {
self.sort_criteria = Some(criteria);
self
}
pub fn terms_matching_strategy(&mut self, value: TermsMatchingStrategy) -> &mut Search<'a> {
self.terms_matching_strategy = value;
self
}
2021-04-14 01:10:58 +08:00
pub fn words_limit(&mut self, value: usize) -> &mut Search<'a> {
self.words_limit = value;
self
}
pub fn filter(&mut self, condition: Filter<'a>) -> &mut Search<'a> {
self.filter = Some(condition);
self
}
2023-04-13 19:45:34 +08:00
#[cfg(test)]
pub fn geo_sort_strategy(&mut self, strategy: new::GeoSortStrategy) -> &mut Search<'a> {
self.geo_strategy = strategy;
self
}
2022-10-17 20:41:57 +08:00
/// Force the search to exhastivelly compute the number of candidates,
/// this will increase the search time but allows finite pagination.
pub fn exhaustive_number_hits(&mut self, exhaustive_number_hits: bool) -> &mut Search<'a> {
self.exhaustive_number_hits = exhaustive_number_hits;
self
}
pub fn execute(&self) -> Result<SearchResult> {
let mut ctx = SearchContext::new(self.index, self.rtxn);
let PartialSearchResult { located_query_terms, candidates, documents_ids } =
execute_search(
&mut ctx,
&self.query,
self.terms_matching_strategy,
self.exhaustive_number_hits,
&self.filter,
&self.sort_criteria,
2023-04-13 19:45:34 +08:00
self.geo_strategy,
self.offset,
self.limit,
Some(self.words_limit),
&mut DefaultSearchLogger,
&mut DefaultSearchLogger,
)?;
// consume context and located_query_terms to build MatchingWords.
let matching_words = match located_query_terms {
Some(located_query_terms) => MatchingWords::new(ctx, located_query_terms),
None => MatchingWords::default(),
};
Ok(SearchResult { matching_words, candidates, documents_ids })
}
}
impl fmt::Debug for Search<'_> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let Search {
query,
filter,
offset,
limit,
2021-08-23 17:37:18 +08:00
sort_criteria,
2023-04-13 19:45:34 +08:00
geo_strategy: _,
terms_matching_strategy,
2021-04-14 01:10:58 +08:00
words_limit,
exhaustive_number_hits,
rtxn: _,
index: _,
} = self;
f.debug_struct("Search")
.field("query", query)
.field("filter", filter)
.field("offset", offset)
.field("limit", limit)
2021-08-23 17:37:18 +08:00
.field("sort_criteria", sort_criteria)
.field("terms_matching_strategy", terms_matching_strategy)
.field("exhaustive_number_hits", exhaustive_number_hits)
2021-04-14 01:10:58 +08:00
.field("words_limit", words_limit)
.finish()
}
}
#[derive(Default, Debug)]
pub struct SearchResult {
pub matching_words: MatchingWords,
pub candidates: RoaringBitmap,
// TODO those documents ids should be associated with their criteria scores.
pub documents_ids: Vec<DocumentId>,
}
2022-08-18 23:36:08 +08:00
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TermsMatchingStrategy {
// remove last word first
Last,
// all words are mandatory
All,
}
impl Default for TermsMatchingStrategy {
fn default() -> Self {
Self::Last
}
}
2022-01-21 01:35:11 +08:00
fn get_first(s: &str) -> &str {
match s.chars().next() {
Some(c) => &s[..c.len_utf8()],
2022-01-25 17:06:27 +08:00
None => panic!("unexpected empty query"),
2022-01-21 01:35:11 +08:00
}
}
pub fn build_dfa(word: &str, typos: u8, is_prefix: bool) -> DFA {
let lev = match typos {
0 => &LEVDIST0,
1 => &LEVDIST1,
_ => &LEVDIST2,
};
if is_prefix {
lev.build_prefix_dfa(word)
} else {
lev.build_dfa(word)
}
}
2022-03-31 15:54:49 +08:00
#[cfg(test)]
mod test {
use super::*;
2022-03-31 20:06:23 +08:00
use crate::index::tests::TempIndex;
2022-03-31 15:54:49 +08:00
#[cfg(feature = "default")]
2023-02-02 01:39:54 +08:00
#[test]
fn test_kanji_language_detection() {
let index = TempIndex::new();
index
.add_documents(documents!([
{ "id": 0, "title": "The quick (\"brown\") fox can't jump 32.3 feet, right? Brr, it's 29.3°F!" },
{ "id": 1, "title": "東京のお寿司。" },
{ "id": 2, "title": "הַשּׁוּעָל הַמָּהִיר (״הַחוּם״) לֹא יָכוֹל לִקְפֹּץ 9.94 מֶטְרִים, נָכוֹן? ברר, 1.5°C- בַּחוּץ!" }
]))
.unwrap();
let txn = index.write_txn().unwrap();
let mut search = Search::new(&txn, &index);
search.query("東京");
let SearchResult { documents_ids, .. } = search.execute().unwrap();
assert_eq!(documents_ids, vec![1]);
}
2022-03-31 15:54:49 +08:00
}