refactor search

This commit is contained in:
mpostma 2021-03-04 11:23:41 +01:00
parent ae3c8af56c
commit 9e2a95b1a3
No known key found for this signature in database
GPG Key ID: CBC8A7C1D7A28C3A
7 changed files with 23 additions and 395 deletions

View File

@ -1,8 +1,6 @@
pub mod search;
mod updates;
pub use search::{SearchQuery, SearchResult, DEFAULT_SEARCH_LIMIT};
use std::fs::create_dir_all;
use std::ops::Deref;
use std::sync::Arc;

View File

@ -1,203 +1,8 @@
use std::collections::{HashSet, BTreeMap};
use std::mem;
use std::time::Instant;
use anyhow::bail;
use either::Either;
use heed::RoTxn;
use meilisearch_tokenizer::{Analyzer, AnalyzerConfig};
use milli::{obkv_to_json, FacetCondition, Index, facet::FacetValue};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::index::{SearchQuery, SearchResult};
use super::Data;
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)]
#[allow(dead_code)]
pub struct SearchQuery {
pub q: Option<String>,
pub offset: Option<usize>,
#[serde(default = "default_search_limit")]
pub limit: usize,
pub attributes_to_retrieve: Option<Vec<String>>,
pub attributes_to_crop: Option<Vec<String>>,
pub crop_length: Option<usize>,
pub attributes_to_highlight: Option<HashSet<String>>,
pub filters: Option<String>,
pub matches: Option<bool>,
pub facet_filters: Option<Value>,
pub facet_distributions: Option<Vec<String>>,
}
impl SearchQuery {
pub fn perform(&self, index: impl AsRef<Index>) -> anyhow::Result<SearchResult> {
let index = index.as_ref();
let before_search = Instant::now();
let rtxn = index.read_txn()?;
let mut search = index.search(&rtxn);
if let Some(ref query) = self.q {
search.query(query);
}
search.limit(self.limit);
search.offset(self.offset.unwrap_or_default());
if let Some(ref facets) = self.facet_filters {
if let Some(facets) = parse_facets(facets, index, &rtxn)? {
search.facet_condition(facets);
}
}
let milli::SearchResult {
documents_ids,
found_words,
candidates,
} = search.execute()?;
let mut documents = Vec::new();
let fields_ids_map = index.fields_ids_map(&rtxn)?;
let displayed_fields_ids = index.displayed_fields_ids(&rtxn)?;
let attributes_to_retrieve_ids = match self.attributes_to_retrieve {
Some(ref attrs) if attrs.iter().any(|f| f == "*") => None,
Some(ref attrs) => attrs
.iter()
.filter_map(|f| fields_ids_map.id(f))
.collect::<Vec<_>>()
.into(),
None => None,
};
let displayed_fields_ids = match (displayed_fields_ids, attributes_to_retrieve_ids) {
(_, Some(ids)) => ids,
(Some(ids), None) => ids,
(None, None) => fields_ids_map.iter().map(|(id, _)| id).collect(),
};
let stop_words = fst::Set::default();
let highlighter = Highlighter::new(&stop_words);
for (_id, obkv) in index.documents(&rtxn, documents_ids)? {
let mut object = obkv_to_json(&displayed_fields_ids, &fields_ids_map, obkv)?;
if let Some(ref attributes_to_highlight) = self.attributes_to_highlight {
highlighter.highlight_record(&mut object, &found_words, attributes_to_highlight);
}
documents.push(object);
}
let nb_hits = candidates.len();
let facet_distributions = match self.facet_distributions {
Some(ref fields) => {
let mut facet_distribution = index.facets_distribution(&rtxn);
if fields.iter().all(|f| f != "*") {
facet_distribution.facets(fields);
}
Some(facet_distribution.candidates(candidates).execute()?)
}
None => None,
};
Ok(SearchResult {
hits: documents,
nb_hits,
query: self.q.clone().unwrap_or_default(),
limit: self.limit,
offset: self.offset.unwrap_or_default(),
processing_time_ms: before_search.elapsed().as_millis(),
facet_distributions,
})
}
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchResult {
pub hits: Vec<Map<String, Value>>,
pub nb_hits: u64,
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<FacetValue, u64>>>,
}
pub struct Highlighter<'a, A> {
analyzer: Analyzer<'a, A>,
}
impl<'a, A: AsRef<[u8]>> Highlighter<'a, A> {
pub fn new(stop_words: &'a fst::Set<A>) -> Self {
let analyzer = Analyzer::new(AnalyzerConfig::default_with_stopwords(stop_words));
Self { analyzer }
}
pub fn highlight_value(&self, value: Value, words_to_highlight: &HashSet<String>) -> 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.contains(token.text());
if to_highlight {
string.push_str("<mark>")
}
string.push_str(word);
if to_highlight {
string.push_str("</mark>")
}
} 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(),
),
}
}
pub fn highlight_record(
&self,
object: &mut Map<String, Value>,
words_to_highlight: &HashSet<String>,
attributes_to_highlight: &HashSet<String>,
) {
// TODO do we need to create a string for element that are not and needs to be highlight?
for (key, value) in object.iter_mut() {
if attributes_to_highlight.contains(key) {
let old_value = mem::take(value);
*value = self.highlight_value(old_value, words_to_highlight);
}
}
}
}
impl Data {
pub async fn search<S: AsRef<str>>(
&self,
@ -303,49 +108,3 @@ impl Data {
//document
}
}
fn parse_facets_array(
txn: &RoTxn,
index: &Index,
arr: &Vec<Value>,
) -> anyhow::Result<Option<FacetCondition>> {
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
),
}
}
FacetCondition::from_array(txn, index, ands)
}
fn parse_facets(
facets: &Value,
index: &Index,
txn: &RoTxn,
) -> anyhow::Result<Option<FacetCondition>> {
match facets {
// Disabled for now
//Value::String(expr) => Ok(Some(FacetCondition::from_str(txn, index, expr)?)),
Value::Array(arr) => parse_facets_array(txn, index, arr),
v => bail!(
"Invalid facet expression, expected Array, found: {:?}",
v
),
}
}

View File

@ -2,25 +2,20 @@ use std::collections::{HashMap, hash_map::Entry};
use std::fs::{File, create_dir_all};
use std::path::{PathBuf, Path};
use std::sync::Arc;
use std::time::Instant;
use anyhow::bail;
use either::Either;
use async_stream::stream;
use chrono::Utc;
use futures::stream::StreamExt;
use heed::{EnvOpenOptions, RoTxn};
use heed::EnvOpenOptions;
use log::info;
use milli::{FacetCondition, Index};
use serde_json::Value;
use thiserror::Error;
use tokio::sync::{mpsc, oneshot, RwLock};
use uuid::Uuid;
use super::update_handler::UpdateHandler;
use crate::data::{SearchQuery, SearchResult};
use crate::index_controller::{IndexMetadata, UpdateMeta, updates::{Processed, Failed, Processing}, UpdateResult as UResult};
use crate::option::IndexerOpts;
use crate::index::{Index, SearchQuery, SearchResult};
pub type Result<T> = std::result::Result<T, IndexError>;
type AsyncMap<K, V> = Arc<RwLock<HashMap<K, V>>>;
@ -49,8 +44,8 @@ pub enum IndexError {
#[async_trait::async_trait]
trait IndexStore {
async fn create_index(&self, uuid: Uuid, primary_key: Option<String>) -> Result<IndexMetadata>;
async fn get_or_create(&self, uuid: Uuid) -> Result<Arc<Index>>;
async fn get(&self, uuid: Uuid) -> Result<Option<Arc<Index>>>;
async fn get_or_create(&self, uuid: Uuid) -> Result<Index>;
async fn get(&self, uuid: Uuid) -> Result<Option<Index>>;
}
impl<S: IndexStore + Sync + Send> IndexActor<S> {
@ -88,7 +83,7 @@ impl<S: IndexStore + Sync + Send> IndexActor<S> {
async fn handle_search(&self, uuid: Uuid, query: SearchQuery, ret: oneshot::Sender<anyhow::Result<SearchResult>>) {
let index = self.store.get(uuid).await.unwrap().unwrap();
tokio::task::spawn_blocking(move || {
let result = perform_search(&index, query);
let result = index.perform_search(query);
ret.send(result)
});
@ -104,138 +99,12 @@ impl<S: IndexStore + Sync + Send> IndexActor<S> {
let uuid = meta.index_uuid().clone();
let index = self.store.get_or_create(uuid).await.unwrap();
let update_handler = self.update_handler.clone();
let result = tokio::task::spawn_blocking(move || update_handler.handle_update(meta, data, index.as_ref())).await;
let result = tokio::task::spawn_blocking(move || update_handler.handle_update(meta, data, index)).await;
let result = result.unwrap();
let _ = ret.send(result);
}
}
fn perform_search(index: &Index, query: SearchQuery) -> anyhow::Result<SearchResult> {
let before_search = Instant::now();
let rtxn = index.read_txn()?;
let mut search = index.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 facets) = query.facet_filters {
if let Some(facets) = parse_facets(facets, index, &rtxn)? {
search.facet_condition(facets);
}
}
let milli::SearchResult {
documents_ids,
found_words,
candidates,
..
} = search.execute()?;
let mut documents = Vec::new();
let fields_ids_map = index.fields_ids_map(&rtxn).unwrap();
let displayed_fields_ids = index.displayed_fields_ids(&rtxn).unwrap();
let attributes_to_retrieve_ids = match query.attributes_to_retrieve {
Some(ref attrs) if attrs.iter().any(|f| f == "*") => None,
Some(ref attrs) => attrs
.iter()
.filter_map(|f| fields_ids_map.id(f))
.collect::<Vec<_>>()
.into(),
None => None,
};
let displayed_fields_ids = match (displayed_fields_ids, attributes_to_retrieve_ids) {
(_, Some(ids)) => ids,
(Some(ids), None) => ids,
(None, None) => fields_ids_map.iter().map(|(id, _)| id).collect(),
};
let stop_words = fst::Set::default();
let highlighter = crate::data::search::Highlighter::new(&stop_words);
for (_id, obkv) in index.documents(&rtxn, documents_ids)? {
let mut object = milli::obkv_to_json(&displayed_fields_ids, &fields_ids_map, obkv).unwrap();
if let Some(ref attributes_to_highlight) = query.attributes_to_highlight {
highlighter.highlight_record(&mut object, &found_words, attributes_to_highlight);
}
documents.push(object);
}
let nb_hits = candidates.len();
let facet_distributions = match query.facet_distributions {
Some(ref fields) => {
let mut facet_distribution = index.facets_distribution(&rtxn);
if fields.iter().all(|f| f != "*") {
facet_distribution.facets(fields);
}
Some(facet_distribution.candidates(candidates).execute()?)
}
None => None,
};
let result = SearchResult {
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 parse_facets_array(
txn: &RoTxn,
index: &Index,
arr: &Vec<Value>,
) -> anyhow::Result<Option<FacetCondition>> {
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
),
}
}
FacetCondition::from_array(txn, index, ands)
}
fn parse_facets(
facets: &Value,
index: &Index,
txn: &RoTxn,
) -> anyhow::Result<Option<FacetCondition>> {
match facets {
// Disabled for now
//Value::String(expr) => Ok(Some(FacetCondition::from_str(txn, index, expr)?)),
Value::Array(arr) => parse_facets_array(txn, index, arr),
v => bail!(
"Invalid facet expression, expected Array, found: {:?}",
v
),
}
}
#[derive(Clone)]
pub struct IndexActorHandle {
sender: mpsc::Sender<IndexMsg>,
@ -276,7 +145,7 @@ impl IndexActorHandle {
struct MapIndexStore {
root: PathBuf,
meta_store: AsyncMap<Uuid, IndexMetadata>,
index_store: AsyncMap<Uuid, Arc<Index>>,
index_store: AsyncMap<Uuid, Index>,
}
#[async_trait::async_trait]
@ -301,17 +170,18 @@ impl IndexStore for MapIndexStore {
create_dir_all(&db_path).expect("can't create db");
let mut options = EnvOpenOptions::new();
options.map_size(4096 * 100_000);
let index = Index::new(options, &db_path)
let index = milli::Index::new(options, &db_path)
.map_err(|e| IndexError::Error(e))?;
let index = Index(Arc::new(index));
Ok(index)
}).await.expect("thread died");
self.index_store.write().await.insert(meta.uuid.clone(), Arc::new(index?));
self.index_store.write().await.insert(meta.uuid.clone(), index?);
Ok(meta)
}
async fn get_or_create(&self, uuid: Uuid) -> Result<Arc<Index>> {
async fn get_or_create(&self, uuid: Uuid) -> Result<Index> {
match self.index_store.write().await.entry(uuid.clone()) {
Entry::Vacant(entry) => {
match self.meta_store.write().await.entry(uuid.clone()) {
@ -327,7 +197,7 @@ impl IndexStore for MapIndexStore {
}
}
async fn get(&self, uuid: Uuid) -> Result<Option<Arc<Index>>> {
async fn get(&self, uuid: Uuid) -> Result<Option<Index>> {
Ok(self.index_store.read().await.get(&uuid).cloned())
}
}

View File

@ -12,7 +12,7 @@ use super::IndexMetadata;
use futures::stream::StreamExt;
use actix_web::web::Payload;
use super::UpdateMeta;
use crate::data::{SearchResult, SearchQuery};
use crate::index::{SearchResult, SearchQuery};
use actix_web::web::Bytes;
pub struct IndexController {

View File

@ -7,7 +7,7 @@ use flate2::read::GzDecoder;
use grenad::CompressionType;
use log::info;
use milli::update::{IndexDocumentsMethod, UpdateBuilder, UpdateFormat};
use milli::Index;
use crate::index::Index;
use rayon::ThreadPool;
use crate::index_controller::updates::{Failed, Processed, Processing};
@ -225,7 +225,7 @@ impl UpdateHandler {
&self,
meta: Processing<UpdateMeta>,
content: File,
index: &Index,
index: Index,
) -> Result<Processed<UpdateMeta, UpdateResult>, Failed<UpdateMeta, String>> {
use UpdateMeta::*;
@ -244,12 +244,12 @@ impl UpdateHandler {
content,
update_builder,
primary_key.as_deref(),
index,
&index,
),
ClearDocuments => self.clear_documents(update_builder, index),
DeleteDocuments => self.delete_documents(content, update_builder, index),
Settings(settings) => self.update_settings(settings, update_builder, index),
Facets(levels) => self.update_facets(levels, update_builder, index),
ClearDocuments => self.clear_documents(update_builder, &index),
DeleteDocuments => self.delete_documents(content, update_builder, &index),
Settings(settings) => self.update_settings(settings, update_builder, &index),
Facets(levels) => self.update_facets(levels, update_builder, &index),
};
match result {

View File

@ -8,6 +8,7 @@ pub mod error;
pub mod helpers;
pub mod option;
pub mod routes;
mod index;
mod index_controller;
pub use option::Opt;

View File

@ -4,11 +4,11 @@ use std::convert::{TryFrom, TryInto};
use actix_web::{get, post, web, HttpResponse};
use serde::Deserialize;
use crate::data::{SearchQuery, DEFAULT_SEARCH_LIMIT};
use crate::error::ResponseError;
use crate::helpers::Authentication;
use crate::routes::IndexParam;
use crate::Data;
use crate::index::{SearchQuery, DEFAULT_SEARCH_LIMIT};
pub fn services(cfg: &mut web::ServiceConfig) {
cfg.service(search_with_post).service(search_with_url_query);