add the geosearch to Meilisearch

meilisearch-lib/Cargo.toml

@@ -27,12 +27,13 @@ heed = { git = "https://github.com/Kerollmops/heed", tag = "v0.12.1" }
 http = "0.2.4"
 indexmap = { version = "1.7.0", features = ["serde-1"] }
 itertools = "0.10.1"
+lazy_static = "1.4.0"
 log = "0.4.14"
 main_error = "0.1.1"
 meilisearch-error = { path = "../meilisearch-error" }
 meilisearch-tokenizer = { git = "https://github.com/meilisearch/tokenizer.git", tag = "v0.2.5" }
 memmap = "0.7.0"
-milli = { git = "https://github.com/meilisearch/milli.git", rev = "6de1b41" }
+milli = { git = "https://github.com/meilisearch/milli.git", version = "0.16.0" }
 mime = "0.3.16"
 num_cpus = "1.13.0"
 once_cell = "1.8.0"

meilisearch-lib/src/error.rs

@@ -35,13 +35,11 @@ impl ErrorCode for MilliError<'_> {
             milli::Error::UserError(ref error) => {
                 match error {
                     // TODO: wait for spec for new error codes.
-                    | UserError::SerdeJson(_)
+                    UserError::SerdeJson(_)
                     | UserError::MaxDatabaseSizeReached
-                    | UserError::InvalidCriterionName { .. }
                     | UserError::InvalidDocumentId { .. }
                     | UserError::InvalidStoreFile
                     | UserError::NoSpaceLeftOnDevice
-                    | UserError::InvalidAscDescSyntax { .. }
                     | UserError::DocumentLimitReached => Code::Internal,
                     UserError::AttributeLimitReached => Code::MaxFieldsLimitExceeded,
                     UserError::InvalidFilter(_) => Code::Filter,
@@ -54,7 +52,10 @@ impl ErrorCode for MilliError<'_> {
                     UserError::SortRankingRuleMissing => Code::Sort,
                     UserError::UnknownInternalDocumentId { .. } => Code::DocumentNotFound,
                     UserError::InvalidFacetsDistribution { .. } => Code::BadRequest,
-                    UserError::InvalidSortableAttribute { .. } => Code::Sort,
+                    UserError::InvalidSortableAttribute { .. }
+                    | UserError::InvalidReservedSortName { .. } => Code::Sort,
+                    UserError::CriterionError(_) => Code::InvalidRankingRule,
+                    UserError::InvalidGeoField { .. } => Code::InvalidGeoField,
                 }
             }
         }

meilisearch-lib/src/index/search.rs

@@ -6,9 +6,12 @@ use either::Either;
 use heed::RoTxn;
 use indexmap::IndexMap;
 use meilisearch_tokenizer::{Analyzer, AnalyzerConfig, Token};
-use milli::{AscDesc, FieldId, FieldsIdsMap, FilterCondition, MatchingWords, UserError};
+use milli::{
+    AscDesc, AscDescError, FieldId, FieldsIdsMap, FilterCondition, MatchingWords, UserError,
+};
+use regex::Regex;
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
+use serde_json::{json, Value};
 
 use crate::index::error::FacetError;
 use crate::index::IndexError;
@@ -110,12 +113,16 @@ impl Index {
         if let Some(ref sort) = query.sort {
             let sort = match sort.iter().map(|s| AscDesc::from_str(s)).collect() {
                 Ok(sorts) => sorts,
-                Err(UserError::InvalidAscDescSyntax { name }) => {
+                Err(AscDescError::InvalidSyntax { name }) => {
                     return Err(IndexError::Milli(
                         UserError::InvalidSortName { name }.into(),
                     ))
                 }
-                Err(err) => return Err(IndexError::Milli(err.into())),
+                Err(AscDescError::ReservedKeyword { name }) => {
+                    return Err(IndexError::Milli(
+                        UserError::InvalidReservedSortName { name }.into(),
+                    ))
+                }
             };
 
             search.sort_criteria(sort);
@@ -193,7 +200,7 @@ impl Index {
         let documents_iter = self.documents(&rtxn, documents_ids)?;
 
         for (_id, obkv) in documents_iter {
-            let document = make_document(&to_retrieve_ids, &fields_ids_map, obkv)?;
+            let mut document = make_document(&to_retrieve_ids, &fields_ids_map, obkv)?;
 
             let matches_info = query
                 .matches
@@ -207,6 +214,10 @@ impl Index {
                 &formatted_options,
             )?;
 
+            if let Some(sort) = query.sort.as_ref() {
+                insert_geo_distance(sort, &mut document);
+            }
+
             let hit = SearchHit {
                 document,
                 formatted,
@@ -247,6 +258,25 @@ impl Index {
     }
 }
 
+fn insert_geo_distance(sorts: &[String], document: &mut Document) {
+    lazy_static::lazy_static! {
+        static ref GEO_REGEX: Regex =
+            Regex::new(r"_geoPoint\(\s*([[:digit:].\-]+)\s*,\s*([[:digit:].\-]+)\s*\)").unwrap();
+    };
+    if let Some(capture_group) = sorts.iter().find_map(|sort| GEO_REGEX.captures(sort)) {
+        // TODO: TAMO: milli encountered an internal error, what do we want to do?
+        let base = [
+            capture_group[1].parse().unwrap(),
+            capture_group[2].parse().unwrap(),
+        ];
+        let geo_point = &document.get("_geo").unwrap_or(&json!(null));
+        if let Some((lat, lng)) = geo_point["lat"].as_f64().zip(geo_point["lng"].as_f64()) {
+            let distance = milli::distance_between_two_points(&base, &[lat, lng]);
+            document.insert("_geoDistance".to_string(), json!(distance.round() as usize));
+        }
+    }
+}
+
 fn compute_matches<A: AsRef<[u8]>>(
     matcher: &impl Matcher,
     document: &Document,
@@ -1332,4 +1362,65 @@ mod test {
             r##"{"about": [MatchInfo { start: 0, length: 6 }, MatchInfo { start: 31, length: 7 }, MatchInfo { start: 191, length: 7 }, MatchInfo { start: 225, length: 7 }, MatchInfo { start: 233, length: 6 }], "color": [MatchInfo { start: 0, length: 3 }]}"##
         );
     }
+
+    #[test]
+    fn test_insert_geo_distance() {
+        let value: Document = serde_json::from_str(
+            r#"{
+      "_geo": {
+        "lat": 50.629973371633746,
+        "lng": 3.0569447399419567
+      },
+      "city": "Lille",
+      "id": "1"
+    }"#,
+        )
+        .unwrap();
+
+        let sorters = &["_geoPoint(50.629973371633746,3.0569447399419567):desc".to_string()];
+        let mut document = value.clone();
+        insert_geo_distance(sorters, &mut document);
+        assert_eq!(document.get("_geoDistance"), Some(&json!(0)));
+
+        let sorters = &["_geoPoint(50.629973371633746, 3.0569447399419567):asc".to_string()];
+        let mut document = value.clone();
+        insert_geo_distance(sorters, &mut document);
+        assert_eq!(document.get("_geoDistance"), Some(&json!(0)));
+
+        let sorters =
+            &["_geoPoint(   50.629973371633746   ,  3.0569447399419567   ):desc".to_string()];
+        let mut document = value.clone();
+        insert_geo_distance(sorters, &mut document);
+        assert_eq!(document.get("_geoDistance"), Some(&json!(0)));
+
+        let sorters = &[
+            "prix:asc",
+            "villeneuve:desc",
+            "_geoPoint(50.629973371633746, 3.0569447399419567):asc",
+            "ubu:asc",
+        ]
+        .map(|s| s.to_string());
+        let mut document = value.clone();
+        insert_geo_distance(sorters, &mut document);
+        assert_eq!(document.get("_geoDistance"), Some(&json!(0)));
+
+        // only the first geoPoint is used to compute the distance
+        let sorters = &[
+            "chien:desc",
+            "_geoPoint(50.629973371633746, 3.0569447399419567):asc",
+            "pangolin:desc",
+            "_geoPoint(100.0, -80.0):asc",
+            "chat:asc",
+        ]
+        .map(|s| s.to_string());
+        let mut document = value.clone();
+        insert_geo_distance(sorters, &mut document);
+        assert_eq!(document.get("_geoDistance"), Some(&json!(0)));
+
+        // there was no _geoPoint so nothing is inserted in the document
+        let sorters = &["chien:asc".to_string()];
+        let mut document = value;
+        insert_geo_distance(sorters, &mut document);
+        assert_eq!(document.get("_geoDistance"), None);
+    }
 }