From 9e2a95b1a372ee2b8909807e16207fa8ed989b2f Mon Sep 17 00:00:00 2001 From: mpostma Date: Thu, 4 Mar 2021 11:23:41 +0100 Subject: [PATCH] refactor search --- src/data/mod.rs | 2 - src/data/search.rs | 243 +----------------- .../actor_index_controller/index_actor.rs | 154 +---------- .../actor_index_controller/mod.rs | 2 +- .../actor_index_controller/update_handler.rs | 14 +- src/lib.rs | 1 + src/routes/search.rs | 2 +- 7 files changed, 23 insertions(+), 395 deletions(-) diff --git a/src/data/mod.rs b/src/data/mod.rs index aa601a670..58acb105a 100644 --- a/src/data/mod.rs +++ b/src/data/mod.rs @@ -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; diff --git a/src/data/search.rs b/src/data/search.rs index 6ab792073..d6bf8438e 100644 --- a/src/data/search.rs +++ b/src/data/search.rs @@ -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, - pub offset: Option, - #[serde(default = "default_search_limit")] - pub limit: usize, - pub attributes_to_retrieve: Option>, - pub attributes_to_crop: Option>, - pub crop_length: Option, - pub attributes_to_highlight: Option>, - pub filters: Option, - pub matches: Option, - pub facet_filters: Option, - pub facet_distributions: Option>, -} - -impl SearchQuery { - pub fn perform(&self, index: impl AsRef) -> anyhow::Result { - 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::>() - .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>, - 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>>, -} - -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) -> Self { - let analyzer = Analyzer::new(AnalyzerConfig::default_with_stopwords(stop_words)); - Self { analyzer } - } - - pub fn highlight_value(&self, value: Value, words_to_highlight: &HashSet) -> 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("") - } - string.push_str(word); - if to_highlight { - string.push_str("") - } - } 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, - words_to_highlight: &HashSet, - attributes_to_highlight: &HashSet, - ) { - // 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>( &self, @@ -303,49 +108,3 @@ impl Data { //document } } - -fn parse_facets_array( - txn: &RoTxn, - index: &Index, - arr: &Vec, -) -> anyhow::Result> { - 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> { - 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 - ), - } -} - diff --git a/src/index_controller/actor_index_controller/index_actor.rs b/src/index_controller/actor_index_controller/index_actor.rs index b45b683bf..27d206c06 100644 --- a/src/index_controller/actor_index_controller/index_actor.rs +++ b/src/index_controller/actor_index_controller/index_actor.rs @@ -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 = std::result::Result; type AsyncMap = Arc>>; @@ -49,8 +44,8 @@ pub enum IndexError { #[async_trait::async_trait] trait IndexStore { async fn create_index(&self, uuid: Uuid, primary_key: Option) -> Result; - async fn get_or_create(&self, uuid: Uuid) -> Result>; - async fn get(&self, uuid: Uuid) -> Result>>; + async fn get_or_create(&self, uuid: Uuid) -> Result; + async fn get(&self, uuid: Uuid) -> Result>; } impl IndexActor { @@ -88,7 +83,7 @@ impl IndexActor { async fn handle_search(&self, uuid: Uuid, query: SearchQuery, ret: oneshot::Sender>) { 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 IndexActor { 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 { - 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::>() - .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, -) -> anyhow::Result> { - 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> { - 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, @@ -276,7 +145,7 @@ impl IndexActorHandle { struct MapIndexStore { root: PathBuf, meta_store: AsyncMap, - index_store: AsyncMap>, + index_store: AsyncMap, } #[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> { + async fn get_or_create(&self, uuid: Uuid) -> Result { 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>> { + async fn get(&self, uuid: Uuid) -> Result> { Ok(self.index_store.read().await.get(&uuid).cloned()) } } diff --git a/src/index_controller/actor_index_controller/mod.rs b/src/index_controller/actor_index_controller/mod.rs index 882b20963..b893b62e4 100644 --- a/src/index_controller/actor_index_controller/mod.rs +++ b/src/index_controller/actor_index_controller/mod.rs @@ -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 { diff --git a/src/index_controller/actor_index_controller/update_handler.rs b/src/index_controller/actor_index_controller/update_handler.rs index d9ac2f866..c42a532ea 100644 --- a/src/index_controller/actor_index_controller/update_handler.rs +++ b/src/index_controller/actor_index_controller/update_handler.rs @@ -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, content: File, - index: &Index, + index: Index, ) -> Result, Failed> { 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 { diff --git a/src/lib.rs b/src/lib.rs index 65c4e8eb7..32c8e4266 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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; diff --git a/src/routes/search.rs b/src/routes/search.rs index 1f4218555..eb03cf94d 100644 --- a/src/routes/search.rs +++ b/src/routes/search.rs @@ -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);