MeiliSearch/meilisearch-http/src/index/search.rs
2021-06-16 17:15:56 +02:00

486 lines
14 KiB
Rust

use std::borrow::Cow;
use std::collections::{BTreeMap, HashSet};
use std::time::Instant;
use anyhow::bail;
use either::Either;
use heed::RoTxn;
use indexmap::IndexMap;
use itertools::Itertools;
use meilisearch_tokenizer::{Analyzer, AnalyzerConfig};
use milli::{FilterCondition, FieldId, FieldsIdsMap, MatchingWords};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::Index;
pub type Document = IndexMap<String, Value>;
pub const DEFAULT_SEARCH_LIMIT: usize = 20;
const fn default_search_limit() -> usize {
DEFAULT_SEARCH_LIMIT
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct SearchQuery {
pub q: Option<String>,
pub offset: Option<usize>,
#[serde(default = "default_search_limit")]
pub limit: usize,
pub attributes_to_retrieve: Option<HashSet<String>>,
pub attributes_to_crop: Option<HashSet<String>>,
pub crop_length: Option<usize>,
pub attributes_to_highlight: Option<HashSet<String>>,
pub matches: Option<bool>,
pub filter: Option<Value>,
pub facet_distributions: Option<Vec<String>>,
}
#[derive(Debug, Clone, Serialize)]
pub struct SearchHit {
#[serde(flatten)]
pub document: Document,
#[serde(rename = "_formatted", skip_serializing_if = "Document::is_empty")]
pub formatted: Document,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchResult {
pub hits: Vec<SearchHit>,
pub nb_hits: u64,
pub exhaustive_nb_hits: bool,
pub query: String,
pub limit: usize,
pub offset: usize,
pub processing_time_ms: u128,
#[serde(skip_serializing_if = "Option::is_none")]
pub facet_distributions: Option<BTreeMap<String, BTreeMap<String, u64>>>,
}
impl Index {
pub fn perform_search(&self, query: SearchQuery) -> anyhow::Result<SearchResult> {
let before_search = Instant::now();
let rtxn = self.read_txn()?;
let mut search = self.search(&rtxn);
if let Some(ref query) = query.q {
search.query(query);
}
search.limit(query.limit);
search.offset(query.offset.unwrap_or_default());
if let Some(ref filter) = query.filter {
if let Some(facets) = parse_facets(filter, self, &rtxn)? {
search.filter(facets);
}
}
let milli::SearchResult {
documents_ids,
matching_words,
candidates,
..
} = search.execute()?;
let mut documents = Vec::new();
let fields_ids_map = self.fields_ids_map(&rtxn).unwrap();
let displayed_ids = self
.displayed_fields_ids(&rtxn)?
.map(|fields| fields.into_iter().collect::<HashSet<_>>())
.unwrap_or_else(|| fields_ids_map.iter().map(|(id, _)| id).collect());
let fids = |attrs: &HashSet<String>| {
let mut ids = HashSet::new();
for attr in attrs {
if attr == "*" {
ids = displayed_ids.clone();
break;
}
if let Some(id) = fields_ids_map.id(attr) {
ids.insert(id);
}
}
ids
};
let to_retrieve_ids = query
.attributes_to_retrieve
.as_ref()
.map(fids)
.unwrap_or_else(|| displayed_ids.clone());
let to_highlight_ids = query
.attributes_to_highlight
.as_ref()
.map(fids)
.unwrap_or_default();
let to_crop_ids = query
.attributes_to_crop
.as_ref()
.map(fids)
.unwrap_or_default();
// The attributes to retrieve are:
// - the ones explicitly marked as to retrieve that are also in the displayed attributes
let all_attributes: Vec<_> = to_retrieve_ids
.intersection(&displayed_ids)
.cloned()
.sorted()
.collect();
// The formatted attributes are:
// - The one in either highlighted attributes or cropped attributes if there are attributes
// to retrieve
// - All the attributes to retrieve if there are either highlighted or cropped attributes
// the request specified that all attributes are to retrieve (i.e attributes to retrieve is
// empty in the query)
let all_formatted = if query.attributes_to_retrieve.is_none() {
if query.attributes_to_highlight.is_some() || query.attributes_to_crop.is_some() {
Cow::Borrowed(&all_attributes)
} else {
Cow::Owned(Vec::new())
}
} else {
let attrs = (&to_crop_ids | &to_highlight_ids)
.intersection(&displayed_ids)
.cloned()
.collect::<Vec<_>>();
Cow::Owned(attrs)
};
let stop_words = fst::Set::default();
let highlighter =
Highlighter::new(&stop_words, (String::from("<em>"), String::from("</em>")));
for (_id, obkv) in self.documents(&rtxn, documents_ids)? {
let document = make_document(&all_attributes, &fields_ids_map, obkv)?;
let formatted = compute_formatted(
&fields_ids_map,
obkv,
&highlighter,
&matching_words,
all_formatted.as_ref().as_slice(),
&to_highlight_ids,
)?;
let hit = SearchHit {
document,
formatted,
};
documents.push(hit);
}
let nb_hits = candidates.len();
let facet_distributions = match query.facet_distributions {
Some(ref fields) => {
let mut facet_distribution = self.facets_distribution(&rtxn);
if fields.iter().all(|f| f != "*") {
facet_distribution.facets(fields);
}
Some(facet_distribution.candidates(candidates).execute()?)
}
None => None,
};
let result = SearchResult {
exhaustive_nb_hits: false, // not implemented yet
hits: documents,
nb_hits,
query: query.q.clone().unwrap_or_default(),
limit: query.limit,
offset: query.offset.unwrap_or_default(),
processing_time_ms: before_search.elapsed().as_millis(),
facet_distributions,
};
Ok(result)
}
}
fn make_document(
attributes_to_retrieve: &[FieldId],
field_ids_map: &FieldsIdsMap,
obkv: obkv::KvReader,
) -> anyhow::Result<Document> {
let mut document = Document::new();
for attr in attributes_to_retrieve {
if let Some(value) = obkv.get(*attr) {
let value = serde_json::from_slice(value)?;
// This unwrap must be safe since we got the ids from the fields_ids_map just
// before.
let key = field_ids_map
.name(*attr)
.expect("Missing field name")
.to_string();
document.insert(key, value);
}
}
Ok(document)
}
fn compute_formatted<A: AsRef<[u8]>>(
field_ids_map: &FieldsIdsMap,
obkv: obkv::KvReader,
highlighter: &Highlighter<A>,
matching_words: &impl Matcher,
all_formatted: &[FieldId],
to_highlight_ids: &HashSet<FieldId>,
) -> anyhow::Result<Document> {
let mut document = Document::new();
for field in all_formatted {
if let Some(value) = obkv.get(*field) {
let mut value: Value = serde_json::from_slice(value)?;
if to_highlight_ids.contains(field) {
value = highlighter.highlight_value(value, matching_words);
}
// This unwrap must be safe since we got the ids from the fields_ids_map just
// before.
let key = field_ids_map
.name(*field)
.expect("Missing field name")
.to_string();
document.insert(key, value);
}
}
Ok(document)
}
/// trait to allow unit testing of `compute_formated`
trait Matcher {
fn matches(&self, w: &str) -> bool;
}
#[cfg(test)]
impl Matcher for HashSet<String> {
fn matches(&self, w: &str) -> bool {
self.contains(w)
}
}
impl Matcher for MatchingWords {
fn matches(&self, w: &str) -> bool {
self.matching_bytes(w).is_some()
}
}
struct Highlighter<'a, A> {
analyzer: Analyzer<'a, A>,
marks: (String, String),
}
impl<'a, A: AsRef<[u8]>> Highlighter<'a, A> {
pub fn new(stop_words: &'a fst::Set<A>, marks: (String, String)) -> Self {
let mut config = AnalyzerConfig::default();
config.stop_words(stop_words);
let analyzer = Analyzer::new(config);
Self { analyzer, marks }
}
fn highlight_value(&self, value: Value, words_to_highlight: &impl Matcher) -> Value {
match value {
Value::Null => Value::Null,
Value::Bool(boolean) => Value::Bool(boolean),
Value::Number(number) => Value::Number(number),
Value::String(old_string) => {
let mut string = String::new();
let analyzed = self.analyzer.analyze(&old_string);
for (word, token) in analyzed.reconstruct() {
if token.is_word() {
let to_highlight = words_to_highlight.matches(token.text());
if to_highlight {
string.push_str(&self.marks.0)
}
string.push_str(word);
if to_highlight {
string.push_str(&self.marks.1)
}
} else {
string.push_str(word);
}
}
Value::String(string)
}
Value::Array(values) => Value::Array(
values
.into_iter()
.map(|v| self.highlight_value(v, words_to_highlight))
.collect(),
),
Value::Object(object) => Value::Object(
object
.into_iter()
.map(|(k, v)| (k, self.highlight_value(v, words_to_highlight)))
.collect(),
),
}
}
}
fn parse_facets(
facets: &Value,
index: &Index,
txn: &RoTxn,
) -> anyhow::Result<Option<FilterCondition>> {
match facets {
Value::String(expr) => Ok(Some(FilterCondition::from_str(txn, index, expr)?)),
Value::Array(arr) => parse_facets_array(txn, index, arr),
v => bail!("Invalid facet expression, expected Array, found: {:?}", v),
}
}
fn parse_facets_array(
txn: &RoTxn,
index: &Index,
arr: &[Value],
) -> anyhow::Result<Option<FilterCondition>> {
let mut ands = Vec::new();
for value in arr {
match value {
Value::String(s) => ands.push(Either::Right(s.clone())),
Value::Array(arr) => {
let mut ors = Vec::new();
for value in arr {
match value {
Value::String(s) => ors.push(s.clone()),
v => bail!("Invalid facet expression, expected String, found: {:?}", v),
}
}
ands.push(Either::Left(ors));
}
v => bail!(
"Invalid facet expression, expected String or [String], found: {:?}",
v
),
}
}
Ok(FilterCondition::from_array(txn, &index.0, ands)?)
}
#[cfg(test)]
mod test {
use std::iter::FromIterator;
use super::*;
#[test]
fn no_formatted() {
let stop_words = fst::Set::default();
let highlighter =
Highlighter::new(&stop_words, (String::from("<em>"), String::from("</em>")));
let mut fields = FieldsIdsMap::new();
let id = fields.insert("test").unwrap();
let mut buf = Vec::new();
let mut obkv = obkv::KvWriter::new(&mut buf);
obkv.insert(id, Value::String("hello".into()).to_string().as_bytes())
.unwrap();
obkv.finish().unwrap();
let obkv = obkv::KvReader::new(&buf);
let all_formatted = Vec::new();
let to_highlight_ids = HashSet::new();
let matching_words = MatchingWords::default();
let value = compute_formatted(
&fields,
obkv,
&highlighter,
&matching_words,
&all_formatted,
&to_highlight_ids,
)
.unwrap();
assert!(value.is_empty());
}
#[test]
fn formatted_no_highlight() {
let stop_words = fst::Set::default();
let highlighter =
Highlighter::new(&stop_words, (String::from("<em>"), String::from("</em>")));
let mut fields = FieldsIdsMap::new();
let id = fields.insert("test").unwrap();
let mut buf = Vec::new();
let mut obkv = obkv::KvWriter::new(&mut buf);
obkv.insert(id, Value::String("hello".into()).to_string().as_bytes())
.unwrap();
obkv.finish().unwrap();
let obkv = obkv::KvReader::new(&buf);
let all_formatted = vec![id];
let to_highlight_ids = HashSet::new();
let matching_words = MatchingWords::default();
let value = compute_formatted(
&fields,
obkv,
&highlighter,
&matching_words,
&all_formatted,
&to_highlight_ids,
)
.unwrap();
assert_eq!(value["test"], "hello");
}
#[test]
fn formatted_with_highlight() {
let stop_words = fst::Set::default();
let highlighter =
Highlighter::new(&stop_words, (String::from("<em>"), String::from("</em>")));
let mut fields = FieldsIdsMap::new();
let id = fields.insert("test").unwrap();
let mut buf = Vec::new();
let mut obkv = obkv::KvWriter::new(&mut buf);
obkv.insert(id, Value::String("hello".into()).to_string().as_bytes())
.unwrap();
obkv.finish().unwrap();
let obkv = obkv::KvReader::new(&buf);
let all_formatted = vec![id];
let to_highlight_ids = HashSet::from_iter(Some(id));
let matching_words = HashSet::from_iter(Some(String::from("hello")));
let value = compute_formatted(
&fields,
obkv,
&highlighter,
&matching_words,
&all_formatted,
&to_highlight_ids,
)
.unwrap();
assert_eq!(value["test"], "<em>hello</em>");
}
}