Merge #322

322: Geosearch r=ManyTheFish a=irevoire

This PR introduces [basic geo-search functionalities](https://github.com/meilisearch/specifications/pull/59), it makes the engine able to index, filter and, sort by geo-point. We decided to use [the rstar library](https://docs.rs/rstar) and to save the points in [an RTree](https://docs.rs/rstar/0.9.1/rstar/struct.RTree.html) that we de/serialize in the index database [by using serde](https://serde.rs/) with [bincode](https://docs.rs/bincode). This is not an efficient way to query this tree as it will consume a lot of CPU and memory when a search is made, but at least it is an easy first way to do so.

### What we will have to do on the indexing part:
 - [x] Index the `_geo` fields from the documents.
   - [x] Create a new module with an extractor in the `extract` module that takes the `obkv_documents` and retrieves the latitude and longitude coordinates, outputting them in a `grenad::Reader` for further process.
   - [x] Call the extractor in the `extract::extract_documents_data` function and send the result to the `TypedChunk` module.
   - [x] Get the `grenad::Reader` in the `typed_chunk::write_typed_chunk_into_index` function and store all the points in the `rtree`
- [x] Delete the documents from the `RTree` when deleting documents from the database. All this can be done in the `delete_documents.rs` file by getting the data structure and removing the points from it, inserting it back after the modification.
- [x] Clearing the `RTree` entirely when we clear the documents from the database, everything happens in the `clear_documents.rs` file.
- [x] save a Roaring bitmap of all documents containing the `_geo` field

### What we will have to do on the query part:
- [x] Filter the documents at a certain distance around a point, this is done by [collecting the documents from the searched point](https://docs.rs/rstar/0.9.1/rstar/struct.RTree.html#method.nearest_neighbor_iter) while they are in range.
  - [x] We must introduce new `geoLowerThan` and `geoGreaterThan` variants to the `Operator` filter enum.
  - [x] Implement the `negative` method on both variants where the `geoGreaterThan` variant is implemented by executing the `geoLowerThan` and removing the results found from the whole list of geo faceted documents.
  - [x] Add the `_geoRadius` function in the pest parser.
- [x] Introduce a `_geo` ascending ranking function that takes a point in parameter, ~~this function must keep the iterator on the `RTree` and make it peekable~~ This was not possible for now, we had to collect the whole iterator. Only the documents that are part of the candidates must be sent too!
  - [x] This ascending ranking rule will only be active if the search is set up with the `_geoPoint` parameter that indicates the center point of the ascending ranking rule.

-----------

- On Meilisearch part: We must introduce a new concept, returning the documents with a new `_geoDistance` field when it passed by the `_geo` ranking rule, this has never been done before. We could maybe just do it afterward when the documents have been retrieved from the database, computing the distance from the `_geoPoint` and all of the documents to be returned.

Co-authored-by: Irevoire <tamo@meilisearch.com>
Co-authored-by: cvermand <33010418+bidoubiwa@users.noreply.github.com>
Co-authored-by: Tamo <tamo@meilisearch.com>

milli/src/update/clear_documents.rs

@@ -48,6 +48,8 @@ impl<'t, 'u, 'i> ClearDocuments<'t, 'u, 'i> {
         self.index.put_external_documents_ids(self.wtxn, &ExternalDocumentsIds::default())?;
         self.index.put_documents_ids(self.wtxn, &RoaringBitmap::default())?;
         self.index.put_field_distribution(self.wtxn, &FieldDistribution::default())?;
+        self.index.delete_geo_rtree(self.wtxn)?;
+        self.index.delete_geo_faceted_documents_ids(self.wtxn)?;
 
         // We clean all the faceted documents ids.
         let empty = RoaringBitmap::default();
@@ -93,7 +95,7 @@ mod tests {
         let content = &br#"[
             { "id": 0, "name": "kevin", "age": 20 },
             { "id": 1, "name": "kevina" },
-            { "id": 2, "name": "benoit", "country": "France" }
+            { "id": 2, "name": "benoit", "country": "France", "_geo": { "lng": 42, "lat": 35 } }
         ]"#[..];
         let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
         builder.update_format(UpdateFormat::Json);
@@ -107,13 +109,15 @@ mod tests {
 
         let rtxn = index.read_txn().unwrap();
 
-        assert_eq!(index.fields_ids_map(&rtxn).unwrap().len(), 4);
+        assert_eq!(index.fields_ids_map(&rtxn).unwrap().len(), 5);
 
         assert!(index.words_fst(&rtxn).unwrap().is_empty());
         assert!(index.words_prefixes_fst(&rtxn).unwrap().is_empty());
         assert!(index.external_documents_ids(&rtxn).unwrap().is_empty());
         assert!(index.documents_ids(&rtxn).unwrap().is_empty());
         assert!(index.field_distribution(&rtxn).unwrap().is_empty());
+        assert!(index.geo_rtree(&rtxn).unwrap().is_none());
+        assert!(index.geo_faceted_documents_ids(&rtxn).unwrap().is_empty());
 
         assert!(index.word_docids.is_empty(&rtxn).unwrap());
         assert!(index.word_prefix_docids.is_empty(&rtxn).unwrap());

milli/src/update/delete_documents.rs

@@ -380,6 +380,24 @@ impl<'t, 'u, 'i> DeleteDocuments<'t, 'u, 'i> {
 
         drop(iter);
 
+        if let Some(mut rtree) = self.index.geo_rtree(self.wtxn)? {
+            let mut geo_faceted_doc_ids = self.index.geo_faceted_documents_ids(self.wtxn)?;
+
+            let (points_to_remove, docids_to_remove): (Vec<_>, RoaringBitmap) = rtree
+                .iter()
+                .filter(|&point| self.documents_ids.contains(point.data))
+                .cloned()
+                .map(|point| (point, point.data))
+                .unzip();
+            points_to_remove.iter().for_each(|point| {
+                rtree.remove(&point);
+            });
+            geo_faceted_doc_ids -= docids_to_remove;
+
+            self.index.put_geo_rtree(self.wtxn, &rtree)?;
+            self.index.put_geo_faceted_documents_ids(self.wtxn, &geo_faceted_doc_ids)?;
+        }
+
         // We delete the documents ids that are under the facet field id values.
         remove_docids_from_facet_field_id_number_docids(
             self.wtxn,
@@ -542,6 +560,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashSet;
+
     use big_s::S;
     use heed::EnvOpenOptions;
     use maplit::hashset;
@@ -657,4 +677,95 @@ mod tests {
 
         wtxn.commit().unwrap();
     }
+
+    #[test]
+    fn delete_documents_with_geo_points() {
+        let path = tempfile::tempdir().unwrap();
+        let mut options = EnvOpenOptions::new();
+        options.map_size(10 * 1024 * 1024); // 10 MB
+        let index = Index::new(options, &path).unwrap();
+
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index, 0);
+        builder.set_primary_key(S("id"));
+        builder.set_filterable_fields(hashset!(S("_geo")));
+        builder.set_sortable_fields(hashset!(S("_geo")));
+        builder.execute(|_, _| ()).unwrap();
+
+        let content = &r#"[
+            {"id":"1","city":"Lille",             "_geo": { "lat": 50.629973371633746, "lng": 3.0569447399419570 } },
+            {"id":"2","city":"Mons-en-Barœul",    "_geo": { "lat": 50.641586120121050, "lng": 3.1106593480348670 } },
+            {"id":"3","city":"Hellemmes",         "_geo": { "lat": 50.631220965518080, "lng": 3.1106399673339933 } },
+            {"id":"4","city":"Villeneuve-d'Ascq", "_geo": { "lat": 50.622468098014565, "lng": 3.1476425513437140 } },
+            {"id":"5","city":"Hem",               "_geo": { "lat": 50.655250871381355, "lng": 3.1897297266244130 } },
+            {"id":"6","city":"Roubaix",           "_geo": { "lat": 50.692473451896710, "lng": 3.1763326737747650 } },
+            {"id":"7","city":"Tourcoing",         "_geo": { "lat": 50.726397466736480, "lng": 3.1541653659578670 } },
+            {"id":"8","city":"Mouscron",          "_geo": { "lat": 50.745325554908610, "lng": 3.2206407854429853 } },
+            {"id":"9","city":"Tournai",           "_geo": { "lat": 50.605342528602630, "lng": 3.3758586941351414 } },
+            {"id":"10","city":"Ghent",            "_geo": { "lat": 51.053777403679035, "lng": 3.6957733119926930 } },
+            {"id":"11","city":"Brussels",         "_geo": { "lat": 50.846640974544690, "lng": 4.3370663564281840 } },
+            {"id":"12","city":"Charleroi",        "_geo": { "lat": 50.409570138889480, "lng": 4.4347354315085520 } },
+            {"id":"13","city":"Mons",             "_geo": { "lat": 50.450294178855420, "lng": 3.9623722870904690 } },
+            {"id":"14","city":"Valenciennes",     "_geo": { "lat": 50.351817774473545, "lng": 3.5326283646928800 } },
+            {"id":"15","city":"Arras",            "_geo": { "lat": 50.284487528579950, "lng": 2.7637515844478160 } },
+            {"id":"16","city":"Cambrai",          "_geo": { "lat": 50.179340577906700, "lng": 3.2189409952502930 } },
+            {"id":"17","city":"Bapaume",          "_geo": { "lat": 50.111276127236400, "lng": 2.8547894666083120 } },
+            {"id":"18","city":"Amiens",           "_geo": { "lat": 49.931472529669996, "lng": 2.2710499758317080 } },
+            {"id":"19","city":"Compiègne",        "_geo": { "lat": 49.444980887725656, "lng": 2.7913841281529015 } },
+            {"id":"20","city":"Paris",            "_geo": { "lat": 48.902100060895480, "lng": 2.3708400867406930 } }
+        ]"#[..];
+        let external_ids_to_delete = ["5", "6", "7", "12", "17", "19"];
+
+        let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
+        builder.update_format(UpdateFormat::Json);
+        builder.execute(content.as_bytes(), |_, _| ()).unwrap();
+
+        let external_document_ids = index.external_documents_ids(&wtxn).unwrap();
+        let ids_to_delete: Vec<u32> = external_ids_to_delete
+            .iter()
+            .map(|id| external_document_ids.get(id.as_bytes()).unwrap())
+            .collect();
+
+        // Delete some documents.
+        let mut builder = DeleteDocuments::new(&mut wtxn, &index, 1).unwrap();
+        external_ids_to_delete.iter().for_each(|id| drop(builder.delete_external_id(id)));
+        builder.execute().unwrap();
+
+        wtxn.commit().unwrap();
+
+        let rtxn = index.read_txn().unwrap();
+        let rtree = index.geo_rtree(&rtxn).unwrap().unwrap();
+        let geo_faceted_doc_ids = index.geo_faceted_documents_ids(&rtxn).unwrap();
+
+        let all_geo_ids = rtree.iter().map(|point| point.data).collect::<Vec<_>>();
+        let all_geo_documents = index
+            .documents(&rtxn, all_geo_ids.iter().copied())
+            .unwrap()
+            .iter()
+            .map(|(id, _)| *id)
+            .collect::<HashSet<_>>();
+
+        let all_geo_faceted_ids = geo_faceted_doc_ids.iter().collect::<Vec<_>>();
+        let all_geo_faceted_documents = index
+            .documents(&rtxn, all_geo_faceted_ids.iter().copied())
+            .unwrap()
+            .iter()
+            .map(|(id, _)| *id)
+            .collect::<HashSet<_>>();
+
+        assert_eq!(
+            all_geo_documents, all_geo_faceted_documents,
+            "There is an inconsistency between the geo_faceted database and the rtree"
+        );
+
+        for id in all_geo_documents.iter() {
+            assert!(!ids_to_delete.contains(&id), "The document {} was supposed to be deleted", id);
+        }
+
+        assert_eq!(
+            all_geo_ids.len(),
+            all_geo_documents.len(),
+            "We deleted documents that were not supposed to be deleted"
+        );
+    }
 }

milli/src/update/index_documents/extract/extract_geo_points.rs

@@ -0,0 +1,44 @@
+use std::fs::File;
+use std::io;
+
+use concat_arrays::concat_arrays;
+use serde_json::Value;
+
+use super::helpers::{create_writer, writer_into_reader, GrenadParameters};
+use crate::{FieldId, InternalError, Result, UserError};
+
+/// Extracts the geographical coordinates contained in each document under the `_geo` field.
+///
+/// Returns the generated grenad reader containing the docid as key associated to the (latitude, longitude)
+pub fn extract_geo_points<R: io::Read>(
+    mut obkv_documents: grenad::Reader<R>,
+    indexer: GrenadParameters,
+    primary_key_id: FieldId,
+    geo_field_id: FieldId,
+) -> Result<grenad::Reader<File>> {
+    let mut writer = tempfile::tempfile().and_then(|file| {
+        create_writer(indexer.chunk_compression_type, indexer.chunk_compression_level, file)
+    })?;
+
+    while let Some((docid_bytes, value)) = obkv_documents.next()? {
+        let obkv = obkv::KvReader::new(value);
+        let point: Value = match obkv.get(geo_field_id) {
+            Some(point) => serde_json::from_slice(point).map_err(InternalError::SerdeJson)?,
+            None => continue,
+        };
+
+        if let Some((lat, lng)) = point["lat"].as_f64().zip(point["lng"].as_f64()) {
+            // this will create an array of 16 bytes (two 8 bytes floats)
+            let bytes: [u8; 16] = concat_arrays![lat.to_ne_bytes(), lng.to_ne_bytes()];
+            writer.insert(docid_bytes, bytes)?;
+        } else {
+            // All document must have a primary key so we can unwrap safely here
+            let primary_key = obkv.get(primary_key_id).unwrap();
+            let primary_key =
+                serde_json::from_slice(primary_key).map_err(InternalError::SerdeJson)?;
+            Err(UserError::InvalidGeoField { document_id: primary_key, object: point })?
+        }
+    }
+
+    Ok(writer_into_reader(writer)?)
+}

milli/src/update/index_documents/extract/mod.rs

@@ -3,6 +3,7 @@ mod extract_facet_number_docids;
 mod extract_facet_string_docids;
 mod extract_fid_docid_facet_values;
 mod extract_fid_word_count_docids;
+mod extract_geo_points;
 mod extract_word_docids;
 mod extract_word_level_position_docids;
 mod extract_word_pair_proximity_docids;
@@ -19,6 +20,7 @@ use self::extract_facet_number_docids::extract_facet_number_docids;
 use self::extract_facet_string_docids::extract_facet_string_docids;
 use self::extract_fid_docid_facet_values::extract_fid_docid_facet_values;
 use self::extract_fid_word_count_docids::extract_fid_word_count_docids;
+use self::extract_geo_points::extract_geo_points;
 use self::extract_word_docids::extract_word_docids;
 use self::extract_word_level_position_docids::extract_word_level_position_docids;
 use self::extract_word_pair_proximity_docids::extract_word_pair_proximity_docids;
@@ -37,6 +39,8 @@ pub(crate) fn data_from_obkv_documents(
     lmdb_writer_sx: Sender<Result<TypedChunk>>,
     searchable_fields: Option<HashSet<FieldId>>,
     faceted_fields: HashSet<FieldId>,
+    primary_key_id: FieldId,
+    geo_field_id: Option<FieldId>,
     stop_words: Option<fst::Set<&[u8]>>,
 ) -> Result<()> {
     let result: Result<(Vec<_>, (Vec<_>, Vec<_>))> = obkv_chunks
@@ -48,6 +52,8 @@ pub(crate) fn data_from_obkv_documents(
                 lmdb_writer_sx.clone(),
                 &searchable_fields,
                 &faceted_fields,
+                primary_key_id,
+                geo_field_id,
                 &stop_words,
             )
         })
@@ -168,6 +174,8 @@ fn extract_documents_data(
     lmdb_writer_sx: Sender<Result<TypedChunk>>,
     searchable_fields: &Option<HashSet<FieldId>>,
     faceted_fields: &HashSet<FieldId>,
+    primary_key_id: FieldId,
+    geo_field_id: Option<FieldId>,
     stop_words: &Option<fst::Set<&[u8]>>,
 ) -> Result<(
     grenad::Reader<CursorClonableMmap>,
@@ -177,6 +185,19 @@ fn extract_documents_data(
 
     let _ = lmdb_writer_sx.send(Ok(TypedChunk::Documents(documents_chunk.clone())));
 
+    if let Some(geo_field_id) = geo_field_id {
+        let documents_chunk_cloned = documents_chunk.clone();
+        let lmdb_writer_sx_cloned = lmdb_writer_sx.clone();
+        rayon::spawn(move || {
+            let result =
+                extract_geo_points(documents_chunk_cloned, indexer, primary_key_id, geo_field_id);
+            let _ = match result {
+                Ok(geo_points) => lmdb_writer_sx_cloned.send(Ok(TypedChunk::GeoPoints(geo_points))),
+                Err(error) => lmdb_writer_sx_cloned.send(Err(error)),
+            };
+        });
+    }
+
     let (docid_word_positions_chunk, docid_fid_facet_values_chunks): (Result<_>, Result<_>) =
         rayon::join(
             || {

milli/src/update/index_documents/mod.rs

@@ -228,11 +228,27 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
             Receiver<Result<TypedChunk>>,
         ) = crossbeam_channel::unbounded();
 
+        // get the primary key field id
+        let primary_key_id = fields_ids_map.id(&primary_key).unwrap();
+
         // get searchable fields for word databases
         let searchable_fields =
             self.index.searchable_fields_ids(self.wtxn)?.map(HashSet::from_iter);
         // get filterable fields for facet databases
         let faceted_fields = self.index.faceted_fields_ids(self.wtxn)?;
+        // get the fid of the `_geo` field.
+        let geo_field_id = match self.index.fields_ids_map(self.wtxn)?.id("_geo") {
+            Some(gfid) => {
+                let is_sortable = self.index.sortable_fields_ids(self.wtxn)?.contains(&gfid);
+                let is_filterable = self.index.filterable_fields_ids(self.wtxn)?.contains(&gfid);
+                if is_sortable || is_filterable {
+                    Some(gfid)
+                } else {
+                    None
+                }
+            }
+            None => None,
+        };
 
         let stop_words = self.index.stop_words(self.wtxn)?;
         // let stop_words = stop_words.as_ref();
@@ -261,6 +277,8 @@ impl<'t, 'u, 'i, 'a> IndexDocuments<'t, 'u, 'i, 'a> {
                     lmdb_writer_sx.clone(),
                     searchable_fields,
                     faceted_fields,
+                    primary_key_id,
+                    geo_field_id,
                     stop_words,
                 )
             });
@@ -876,12 +894,12 @@ mod tests {
         // First we send 3 documents with an id for only one of them.
         let mut wtxn = index.write_txn().unwrap();
         let documents = &r#"[
-          { "id": 2,    "title": "Pride and Prejudice",                    "author": "Jane Austin",              "genre": "romance",    "price": 3.5 },
+          { "id": 2,    "title": "Pride and Prejudice",                    "author": "Jane Austin",              "genre": "romance",    "price": 3.5, "_geo": { "lat": 12, "lng": 42 } },
           { "id": 456,  "title": "Le Petit Prince",                        "author": "Antoine de Saint-Exupéry", "genre": "adventure" , "price": 10.0 },
           { "id": 1,    "title": "Alice In Wonderland",                    "author": "Lewis Carroll",            "genre": "fantasy",    "price": 25.99 },
           { "id": 1344, "title": "The Hobbit",                             "author": "J. R. R. Tolkien",         "genre": "fantasy" },
           { "id": 4,    "title": "Harry Potter and the Half-Blood Prince", "author": "J. K. Rowling",            "genre": "fantasy" },
-          { "id": 42,   "title": "The Hitchhiker's Guide to the Galaxy",   "author": "Douglas Adams" }
+          { "id": 42,   "title": "The Hitchhiker's Guide to the Galaxy",   "author": "Douglas Adams", "_geo": { "lat": 35, "lng": 23 } }
         ]"#[..];
         let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
         builder.update_format(UpdateFormat::Json);
@@ -917,7 +935,7 @@ mod tests {
             { "objectId": 123, "title": "Pride and Prejudice", "comment": "A great book" },
             { "objectId": 456, "title": "Le Petit Prince",     "comment": "A french book" },
             { "objectId": 1,   "title": "Alice In Wonderland", "comment": "A weird book" },
-            { "objectId": 30,  "title": "Hamlet" }
+            { "objectId": 30,  "title": "Hamlet", "_geo": { "lat": 12, "lng": 89 } }
         ]"#[..];
         let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
         builder.update_format(UpdateFormat::Json);
@@ -934,7 +952,7 @@ mod tests {
         assert!(external_documents_ids.get("30").is_none());
 
         let content = &br#"[
-            { "objectId": 30, "title": "Hamlet" }
+            { "objectId": 30,  "title": "Hamlet", "_geo": { "lat": 12, "lng": 89 } }
         ]"#[..];
         let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
         builder.update_format(UpdateFormat::Json);
@@ -944,7 +962,7 @@ mod tests {
         assert!(external_documents_ids.get("30").is_some());
 
         let content = &br#"[
-            { "objectId": 30, "title": "Hamlet" }
+            { "objectId": 30,  "title": "Hamlet", "_geo": { "lat": 12, "lng": 89 } }
         ]"#[..];
         let mut builder = IndexDocuments::new(&mut wtxn, &index, 0);
         builder.update_format(UpdateFormat::Json);

milli/src/update/index_documents/typed_chunk.rs

@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::convert::TryInto;
 use std::fs::File;
 
 use heed::types::ByteSlice;
@@ -6,11 +7,12 @@ use heed::{BytesDecode, RwTxn};
 use roaring::RoaringBitmap;
 
 use super::helpers::{
-    roaring_bitmap_from_u32s_array, serialize_roaring_bitmap, valid_lmdb_key, CursorClonableMmap,
+    self, roaring_bitmap_from_u32s_array, serialize_roaring_bitmap, valid_lmdb_key,
+    CursorClonableMmap,
 };
 use crate::heed_codec::facet::{decode_prefix_string, encode_prefix_string};
 use crate::update::index_documents::helpers::into_clonable_grenad;
-use crate::{BoRoaringBitmapCodec, CboRoaringBitmapCodec, Index, Result};
+use crate::{BoRoaringBitmapCodec, CboRoaringBitmapCodec, DocumentId, GeoPoint, Index, Result};
 
 pub(crate) enum TypedChunk {
     DocidWordPositions(grenad::Reader<CursorClonableMmap>),
@@ -24,6 +26,7 @@ pub(crate) enum TypedChunk {
     WordPairProximityDocids(grenad::Reader<File>),
     FieldIdFacetStringDocids(grenad::Reader<File>),
     FieldIdFacetNumberDocids(grenad::Reader<File>),
+    GeoPoints(grenad::Reader<File>),
 }
 
 /// Write typed chunk in the corresponding LMDB database of the provided index.
@@ -177,6 +180,24 @@ pub(crate) fn write_typed_chunk_into_index(
             )?;
             is_merged_database = true;
         }
+        TypedChunk::GeoPoints(mut geo_points) => {
+            let mut rtree = index.geo_rtree(wtxn)?.unwrap_or_default();
+            let mut geo_faceted_docids = index.geo_faceted_documents_ids(wtxn)?;
+
+            while let Some((key, value)) = geo_points.next()? {
+                // convert the key back to a u32 (4 bytes)
+                let docid = key.try_into().map(DocumentId::from_be_bytes).unwrap();
+
+                // convert the latitude and longitude back to a f64 (8 bytes)
+                let (lat, tail) = helpers::try_split_array_at::<u8, 8>(value).unwrap();
+                let (lng, _) = helpers::try_split_array_at::<u8, 8>(tail).unwrap();
+                let point = [f64::from_ne_bytes(lat), f64::from_ne_bytes(lng)];
+                rtree.insert(GeoPoint::new(point, docid));
+                geo_faceted_docids.insert(docid);
+            }
+            index.put_geo_rtree(wtxn, &rtree)?;
+            index.put_geo_faceted_documents_ids(wtxn, &geo_faceted_docids)?;
+        }
     }
 
     Ok((RoaringBitmap::new(), is_merged_database))