search first iteration

src/data.rs

@@ -1,17 +1,53 @@
+use std::borrow::Cow;
+use std::collections::HashSet;
+use std::fs::create_dir_all;
+use std::mem;
 use std::ops::Deref;
 use std::sync::Arc;
-use std::fs::create_dir_all;
+use std::time::Instant;
 
 use async_compression::tokio_02::write::GzipEncoder;
 use futures_util::stream::StreamExt;
 use tokio::io::AsyncWriteExt;
-use milli::Index;
+use milli::{Index, SearchResult as Results, obkv_to_json};
 use milli::update::{IndexDocumentsMethod, UpdateFormat};
 use sha2::Digest;
+use serde_json::{Value, Map};
+use serde::{Deserialize, Serialize};
+use meilisearch_tokenizer::{Analyzer, AnalyzerConfig};
 
 use crate::option::Opt;
 use crate::updates::{UpdateQueue, UpdateMeta, UpdateStatus, UpdateMetaProgress};
 
+const DEFAULT_SEARCH_LIMIT: usize = 20;
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub struct SearchQuery {
+    q: Option<String>,
+    offset: Option<usize>,
+    limit: Option<usize>,
+    attributes_to_retrieve: Option<Vec<String>>,
+    attributes_to_crop: Option<Vec<String>>,
+    crop_length: Option<usize>,
+    attributes_to_highlight: Option<Vec<String>>,
+    filters: Option<String>,
+    matches: Option<bool>,
+    facet_filters: Option<Value>,
+    facets_distribution: Option<Vec<String>>,
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SearchResult {
+    hits: Vec<Map<String, Value>>,
+    nb_hits: usize,
+    query: String,
+    limit: usize,
+    offset: usize,
+    processing_time_ms: u128,
+}
+
 #[derive(Clone)]
 pub struct Data {
     inner: Arc<DataInner>,
@@ -81,8 +117,9 @@ impl Data {
         Ok(Data { inner })
     }
 
-    pub async fn add_documents<B, E>(
+    pub async fn add_documents<B, E, S>(
         &self,
+        _index: S,
         method: IndexDocumentsMethod,
         format: UpdateFormat,
         mut stream: impl futures::Stream<Item=Result<B, E>> + Unpin,
@@ -90,6 +127,7 @@ impl Data {
     where
         B: Deref<Target = [u8]>,
         E: std::error::Error + Send + Sync + 'static,
+        S: AsRef<str>,
     {
         let file = tokio::task::spawn_blocking(tempfile::tempfile).await?;
         let file = tokio::fs::File::from_std(file?);
@@ -115,6 +153,60 @@ impl Data {
         Ok(UpdateStatus::Pending { update_id, meta })
     }
 
+    pub fn search<S: AsRef<str>>(&self, _index: S, search_query: SearchQuery) -> anyhow::Result<SearchResult> {
+        let start =  Instant::now();
+        let index = &self.indexes;
+        let rtxn = index.read_txn()?;
+
+        let mut search = index.search(&rtxn);
+        if let Some(query) = &search_query.q {
+            search.query(query);
+        }
+
+        if let Some(offset) = search_query.offset {
+            search.offset(offset);
+        }
+
+        let limit = search_query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT);
+            search.limit(limit);
+
+        let Results { found_words, documents_ids, nb_hits, .. } = search.execute().unwrap();
+
+        let fields_ids_map = index.fields_ids_map(&rtxn).unwrap();
+
+        let displayed_fields = match index.displayed_fields(&rtxn).unwrap() {
+            Some(fields) => Cow::Borrowed(fields),
+            None => Cow::Owned(fields_ids_map.iter().map(|(id, _)| id).collect()),
+        };
+
+        let attributes_to_highlight = match search_query.attributes_to_highlight {
+            Some(fields) => fields.iter().map(ToOwned::to_owned).collect(),
+            None => HashSet::new(),
+        };
+
+        let stop_words = fst::Set::default();
+        let highlighter = Highlighter::new(&stop_words);
+        let mut documents = Vec::new();
+        for (_id, obkv) in index.documents(&rtxn, documents_ids).unwrap() {
+            let mut object = obkv_to_json(&displayed_fields, &fields_ids_map, obkv).unwrap();
+            highlighter.highlight_record(&mut object, &found_words, &attributes_to_highlight);
+            documents.push(object);
+        }
+
+        let processing_time_ms = start.elapsed().as_millis();
+
+        let result = SearchResult {
+            hits: documents,
+            nb_hits,
+            query: search_query.q.unwrap_or_default(),
+            offset: search_query.offset.unwrap_or(0),
+            limit,
+            processing_time_ms,
+        };
+
+        Ok(result)
+    }
+
     #[inline]
     pub fn http_payload_size_limit(&self) -> usize {
         self.options.http_payload_size_limit.get_bytes() as usize
@@ -125,3 +217,62 @@ impl Data {
         &self.api_keys
     }
 }
+
+struct Highlighter<'a, A> {
+    analyzer: Analyzer<'a, A>,
+}
+
+impl<'a, A: AsRef<[u8]>> Highlighter<'a, A> {
+    fn new(stop_words: &'a fst::Set<A>) -> Self {
+        let analyzer = Analyzer::new(AnalyzerConfig::default_with_stopwords(stop_words));
+        Self { analyzer }
+    }
+
+    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())
+            },
+        }
+    }
+
+    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);
+            }
+        }
+    }
+}

src/routes/document.rs

@@ -117,11 +117,12 @@ async fn update_multiple_documents(
 async fn add_documents_json(
     data: web::Data<Data>,
     path: web::Path<IndexParam>,
-    params: web::Query<UpdateDocumentsQuery>,
+    _params: web::Query<UpdateDocumentsQuery>,
     body: Payload,
 ) -> Result<HttpResponse, ResponseError> {
     let addition_result = data
         .add_documents(
+            &path.index_uid,
             IndexDocumentsMethod::UpdateDocuments,
             UpdateFormat::Json,
             body

src/routes/search.rs

@@ -1,31 +1,31 @@
 use actix_web::{get, post, web, HttpResponse};
-use serde::{Deserialize, Serialize};
-use serde_json::Value;
+use log::error;
 
 use crate::error::ResponseError;
 use crate::helpers::Authentication;
 use crate::routes::IndexParam;
 use crate::Data;
+use crate::data::SearchQuery;
 
 pub fn services(cfg: &mut web::ServiceConfig) {
     cfg.service(search_with_post).service(search_with_url_query);
 }
 
-#[derive(Serialize, Deserialize)]
-#[serde(rename_all = "camelCase", deny_unknown_fields)]
-pub struct SearchQuery {
-    q: Option<String>,
-    offset: Option<usize>,
-    limit: Option<usize>,
-    attributes_to_retrieve: Option<String>,
-    attributes_to_crop: Option<String>,
-    crop_length: Option<usize>,
-    attributes_to_highlight: Option<String>,
-    filters: Option<String>,
-    matches: Option<bool>,
-    facet_filters: Option<String>,
-    facets_distribution: Option<String>,
-}
+//#[derive(Serialize, Deserialize)]
+//#[serde(rename_all = "camelCase", deny_unknown_fields)]
+//pub struct SearchQuery {
+    //q: Option<String>,
+    //offset: Option<usize>,
+    //limit: Option<usize>,
+    //attributes_to_retrieve: Option<String>,
+    //attributes_to_crop: Option<String>,
+    //crop_length: Option<usize>,
+    //attributes_to_highlight: Option<String>,
+    //filters: Option<String>,
+    //matches: Option<bool>,
+    //facet_filters: Option<String>,
+    //facets_distribution: Option<String>,
+//}
 
 #[get("/indexes/{index_uid}/search", wrap = "Authentication::Public")]
 async fn search_with_url_query(
@@ -36,27 +36,21 @@ async fn search_with_url_query(
     todo!()
 }
 
-#[derive(Deserialize)]
-#[serde(rename_all = "camelCase", deny_unknown_fields)]
-pub struct SearchQueryPost {
-    _q: Option<String>,
-    _offset: Option<usize>,
-    _limit: Option<usize>,
-    _attributes_to_retrieve: Option<Vec<String>>,
-    _attributes_to_crop: Option<Vec<String>>,
-    _crop_length: Option<usize>,
-    _attributes_to_highlight: Option<Vec<String>>,
-    _filters: Option<String>,
-    _matches: Option<bool>,
-    _facet_filters: Option<Value>,
-    _facets_distribution: Option<Vec<String>>,
-}
-
 #[post("/indexes/{index_uid}/search", wrap = "Authentication::Public")]
 async fn search_with_post(
-    _data: web::Data<Data>,
-    _path: web::Path<IndexParam>,
-    _params: web::Json<SearchQueryPost>,
+    data: web::Data<Data>,
+    path: web::Path<IndexParam>,
+    params: web::Json<SearchQuery>,
 ) -> Result<HttpResponse, ResponseError> {
-    todo!()
+    let search_result = data.search(&path.index_uid, params.into_inner());
+    match search_result {
+        Ok(docs) => {
+            let docs = serde_json::to_string(&docs).unwrap();
+            Ok(HttpResponse::Ok().body(docs))
+        }
+        Err(e) => {
+            error!("{}", e);
+            todo!()
+        }
+    }
 }