From 77dcbff6b2355b6a72837c29cec4dc7c355dea22 Mon Sep 17 00:00:00 2001
From: Kerollmops <clement@meilisearch.com>
Date: Mon, 23 Oct 2023 13:49:54 +0200
Subject: [PATCH] Remove and Insert the DelAdd geo points

---
 .../extract/extract_geo_points.rs             |  1 -
 .../src/update/index_documents/typed_chunk.rs | 28 +++++++++++++------
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/milli/src/update/index_documents/extract/extract_geo_points.rs b/milli/src/update/index_documents/extract/extract_geo_points.rs
index a818bb91c..cc283121e 100644
--- a/milli/src/update/index_documents/extract/extract_geo_points.rs
+++ b/milli/src/update/index_documents/extract/extract_geo_points.rs
@@ -15,7 +15,6 @@ use crate::{FieldId, InternalError, Result};
 /// Returns the generated grenad reader containing the docid as key associated to the (latitude, longitude)
 #[logging_timer::time]
 pub fn extract_geo_points<R: io::Read + io::Seek>(
-    // TODO grenad::Reader<Obkv<FieldId, Obkv<DelAdd, JsonValue>>>
     obkv_documents: grenad::Reader<R>,
     indexer: GrenadParameters,
     primary_key_id: FieldId,
diff --git a/milli/src/update/index_documents/typed_chunk.rs b/milli/src/update/index_documents/typed_chunk.rs
index 0d618ad28..9d4d63f90 100644
--- a/milli/src/update/index_documents/typed_chunk.rs
+++ b/milli/src/update/index_documents/typed_chunk.rs
@@ -279,14 +279,17 @@ pub(crate) fn write_typed_chunk_into_index(
                 // 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)];
-                let xyz_point = lat_lng_to_xyz(&point);
-
-                rtree.insert(GeoPoint::new(xyz_point, (docid, point)));
-                geo_faceted_docids.insert(docid);
+                let deladd_obkv = KvReaderDelAdd::new(value);
+                if let Some(value) = deladd_obkv.get(DelAdd::Deletion) {
+                    let geopoint = extract_geo_point(value, docid);
+                    rtree.remove(&geopoint);
+                    geo_faceted_docids.remove(docid);
+                }
+                if let Some(value) = deladd_obkv.get(DelAdd::Addition) {
+                    let geopoint = extract_geo_point(value, docid);
+                    rtree.insert(geopoint);
+                    geo_faceted_docids.insert(docid);
+                }
             }
             index.put_geo_rtree(wtxn, &rtree)?;
             index.put_geo_faceted_documents_ids(wtxn, &geo_faceted_docids)?;
@@ -368,6 +371,15 @@ pub(crate) fn write_typed_chunk_into_index(
     Ok((RoaringBitmap::new(), is_merged_database))
 }
 
+/// Converts the latitude and longitude back to an xyz GeoPoint.
+fn extract_geo_point(value: &[u8], docid: DocumentId) -> GeoPoint {
+    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)];
+    let xyz_point = lat_lng_to_xyz(&point);
+    GeoPoint::new(xyz_point, (docid, point))
+}
+
 fn merge_word_docids_reader_into_fst(
     word_docids_iter: grenad::Reader<io::Cursor<ClonableMmap>>,
     exact_word_docids_iter: grenad::Reader<io::Cursor<ClonableMmap>>,