create a new db with getters and setters

milli/Cargo.toml

@@ -27,6 +27,7 @@ once_cell = "1.5.2"
 ordered-float = "2.1.1"
 rayon = "1.5.0"
 roaring = "0.6.6"
+rstar = { version = "0.9.1", features = ["serde"] }
 serde = { version = "1.0.123", features = ["derive"] }
 serde_json = { version = "1.0.62", features = ["preserve_order"] }
 slice-group-by = "0.2.6"

milli/src/index.rs

@@ -8,6 +8,7 @@ use heed::flags::Flags;
 use heed::types::*;
 use heed::{Database, PolyDatabase, RoTxn, RwTxn};
 use roaring::RoaringBitmap;
+use rstar::RTree;
 
 use crate::error::{InternalError, UserError};
 use crate::fields_ids_map::FieldsIdsMap;
@@ -18,8 +19,8 @@ use crate::heed_codec::facet::{
 use crate::{
     default_criteria, BEU32StrCodec, BoRoaringBitmapCodec, CboRoaringBitmapCodec, Criterion,
     DocumentId, ExternalDocumentsIds, FacetDistribution, FieldDistribution, FieldId,
-    FieldIdWordCountCodec, ObkvCodec, Result, RoaringBitmapCodec, RoaringBitmapLenCodec, Search,
-    StrLevelPositionCodec, StrStrU8Codec, BEU32,
+    FieldIdWordCountCodec, GeoPoint, ObkvCodec, Result, RoaringBitmapCodec, RoaringBitmapLenCodec,
+    Search, StrLevelPositionCodec, StrStrU8Codec, BEU32,
 };
 
 pub mod main_key {
@@ -31,6 +32,7 @@ pub mod main_key {
     pub const SORTABLE_FIELDS_KEY: &str = "sortable-fields";
     pub const FIELD_DISTRIBUTION_KEY: &str = "fields-distribution";
     pub const FIELDS_IDS_MAP_KEY: &str = "fields-ids-map";
+    pub const GEO_RTREE_KEY: &str = "geo";
     pub const HARD_EXTERNAL_DOCUMENTS_IDS_KEY: &str = "hard-external-documents-ids";
     pub const NUMBER_FACETED_DOCUMENTS_IDS_PREFIX: &str = "number-faceted-documents-ids";
     pub const PRIMARY_KEY_KEY: &str = "primary-key";
@@ -294,6 +296,30 @@ impl Index {
             .unwrap_or_default())
     }
 
+    /* geo rtree */
+
+    pub(crate) fn put_geo_rtree<A: AsRef<[u8]>>(
+        &self,
+        wtxn: &mut RwTxn,
+        rtree: &RTree<GeoPoint>,
+    ) -> heed::Result<()> {
+        self.main.put::<_, Str, SerdeBincode<RTree<GeoPoint>>>(wtxn, main_key::GEO_RTREE_KEY, rtree)
+    }
+
+    pub(crate) fn delete_geo_rtree(&self, wtxn: &mut RwTxn) -> heed::Result<bool> {
+        self.main.delete::<_, Str>(wtxn, main_key::GEO_RTREE_KEY)
+    }
+
+    pub fn geo_rtree<'t>(&self, rtxn: &'t RoTxn) -> Result<Option<RTree<GeoPoint>>> {
+        match self
+            .main
+            .get::<_, Str, SerdeBincode<RTree<GeoPoint>>>(rtxn, main_key::GEO_RTREE_KEY)?
+        {
+            Some(rtree) => Ok(Some(rtree)),
+            None => Ok(None),
+        }
+    }
+
     /* field distribution */
 
     /// Writes the field distribution which associates every field name with

milli/src/lib.rs

@@ -51,6 +51,7 @@ pub type DocumentId = u32;
 pub type FieldId = u16;
 pub type Position = u32;
 pub type FieldDistribution = BTreeMap<String, u64>;
+pub type GeoPoint = rstar::primitives::GeomWithData<[f64; 2], DocumentId>;
 
 /// Transform a raw obkv store into a JSON Object.
 pub fn obkv_to_json(