refactor search

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;

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<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
-        ),
-    }
-}
-

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<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_or_create(&self, uuid: Uuid) -> Result<Index>;
-    async fn get(&self, uuid: Uuid) -> Result<Option<Arc<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())
     }
 }

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 {

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<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),
+            ClearDocuments => self.clear_documents(update_builder, &index),
-            DeleteDocuments => self.delete_documents(content, update_builder, index),
+            DeleteDocuments => self.delete_documents(content, update_builder, &index),
-            Settings(settings) => self.update_settings(settings, update_builder, index),
+            Settings(settings) => self.update_settings(settings, update_builder, &index),
-            Facets(levels) => self.update_facets(levels, update_builder, index),
+            Facets(levels) => self.update_facets(levels, update_builder, &index),
         };
 
         match result {

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;

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);