diff --git a/milli/Cargo.toml b/milli/Cargo.toml index c6fe3ea95..d2767afd4 100644 --- a/milli/Cargo.toml +++ b/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" diff --git a/milli/src/index.rs b/milli/src/index.rs index f3a2a3e05..d2b4598d3 100644 --- a/milli/src/index.rs +++ b/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>( + &self, + wtxn: &mut RwTxn, + rtree: &RTree, + ) -> heed::Result<()> { + self.main.put::<_, Str, SerdeBincode>>(wtxn, main_key::GEO_RTREE_KEY, rtree) + } + + pub(crate) fn delete_geo_rtree(&self, wtxn: &mut RwTxn) -> heed::Result { + self.main.delete::<_, Str>(wtxn, main_key::GEO_RTREE_KEY) + } + + pub fn geo_rtree<'t>(&self, rtxn: &'t RoTxn) -> Result>> { + match self + .main + .get::<_, Str, SerdeBincode>>(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 diff --git a/milli/src/lib.rs b/milli/src/lib.rs index af811fe08..2a55b6f3a 100644 --- a/milli/src/lib.rs +++ b/milli/src/lib.rs @@ -51,6 +51,7 @@ pub type DocumentId = u32; pub type FieldId = u16; pub type Position = u32; pub type FieldDistribution = BTreeMap; +pub type GeoPoint = rstar::primitives::GeomWithData<[f64; 2], DocumentId>; /// Transform a raw obkv store into a JSON Object. pub fn obkv_to_json(