mmap geo index api (#5163)

* define mmap geo index

fix compilation

deleted flags

load new mmap

geo index tests

fix tests

fix build after rebase

add files list

* refactor get_stored_sub_regions output type

* review remanings

* mmap geo index api

* fix after rebase

docs/grpc/docs.md

@@ -822,6 +822,11 @@
 
 
 
+| Field | Type | Label | Description |
+| ----- | ---- | ----- | ----------- |
+| on_disk | [bool](#bool) | optional | If true - store index on disk. |
+
+
 
 
 

docs/redoc/master/openapi.json

@@ -6832,6 +6832,11 @@
         "properties": {
           "type": {
             "$ref": "#/components/schemas/GeoIndexType"
+          },
+          "on_disk": {
+            "description": "If true, store the index on disk. Default: false.",
+            "type": "boolean",
+            "nullable": true
           }
         }
       },

lib/api/src/grpc/conversions.rs

@@ -258,9 +258,11 @@ impl From<segment::data_types::index::FloatIndexParams> for PayloadIndexParams {
 }
 
 impl From<segment::data_types::index::GeoIndexParams> for PayloadIndexParams {
-    fn from(_params: segment::data_types::index::GeoIndexParams) -> Self {
+    fn from(params: segment::data_types::index::GeoIndexParams) -> Self {
         PayloadIndexParams {
-            index_params: Some(IndexParams::GeoIndexParams(GeoIndexParams {})),
+            index_params: Some(IndexParams::GeoIndexParams(GeoIndexParams {
+                on_disk: params.on_disk,
+            })),
         }
     }
 }
@@ -417,9 +419,10 @@ impl TryFrom<FloatIndexParams> for segment::data_types::index::FloatIndexParams
 
 impl TryFrom<GeoIndexParams> for segment::data_types::index::GeoIndexParams {
     type Error = Status;
-    fn try_from(_params: GeoIndexParams) -> Result<Self, Self::Error> {
+    fn try_from(params: GeoIndexParams) -> Result<Self, Self::Error> {
         Ok(segment::data_types::index::GeoIndexParams {
             r#type: GeoIndexType::Geo,
+            on_disk: params.on_disk,
         })
     }
 }

lib/api/src/grpc/proto/collections.proto

@@ -418,6 +418,7 @@ message FloatIndexParams {
 }
 
 message GeoIndexParams {
+  optional bool on_disk = 1; // If true - store index on disk.
 }
 
 message TextIndexParams {

lib/api/src/grpc/qdrant.rs

@@ -693,7 +693,11 @@ pub struct FloatIndexParams {
 #[derive(serde::Serialize)]
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
-pub struct GeoIndexParams {}
+pub struct GeoIndexParams {
+    /// If true - store index on disk.
+    #[prost(bool, optional, tag = "1")]
+    pub on_disk: ::core::option::Option<bool>,
+}
 #[derive(serde::Serialize)]
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]

lib/segment/src/data_types/index.rs

@@ -120,7 +120,6 @@ pub struct GeoIndexParams {
 
     /// If true, store the index on disk. Default: false.
     #[serde(default, skip_serializing_if = "Option::is_none")]
-    #[cfg(any())]
     pub on_disk: Option<bool>,
 }
 

lib/segment/src/index/field_index/field_index_base.rs

@@ -8,7 +8,7 @@ use super::binary_index::BinaryIndexBuilder;
 use super::facet_index::FacetIndex;
 use super::full_text_index::mmap_text_index::FullTextMmapIndexBuilder;
 use super::full_text_index::text_index::{FullTextIndex, FullTextIndexBuilder};
-use super::geo_index::GeoMapIndexBuilder;
+use super::geo_index::{GeoMapIndexBuilder, GeoMapIndexMmapBuilder};
 use super::map_index::{MapIndex, MapIndexBuilder, MapIndexMmapBuilder};
 use super::numeric_index::{
     NumericIndex, NumericIndexBuilder, NumericIndexMmapBuilder, StreamRange,
@@ -421,6 +421,7 @@ pub enum FieldIndexBuilder {
     FloatIndex(NumericIndexBuilder<FloatPayloadType, FloatPayloadType>),
     FloatMmapIndex(NumericIndexMmapBuilder<FloatPayloadType, FloatPayloadType>),
     GeoIndex(GeoMapIndexBuilder),
+    GeoMmapIndex(GeoMapIndexMmapBuilder),
     FullTextIndex(FullTextIndexBuilder),
     FullTextMmapIndex(FullTextMmapIndexBuilder),
     BinaryIndex(BinaryIndexBuilder),
@@ -444,6 +445,7 @@ impl FieldIndexBuilderTrait for FieldIndexBuilder {
             Self::FloatIndex(index) => index.init(),
             Self::FloatMmapIndex(index) => index.init(),
             Self::GeoIndex(index) => index.init(),
+            Self::GeoMmapIndex(index) => index.init(),
             Self::BinaryIndex(index) => index.init(),
             Self::FullTextIndex(index) => index.init(),
             Self::FullTextMmapIndex(builder) => builder.init(),
@@ -465,6 +467,7 @@ impl FieldIndexBuilderTrait for FieldIndexBuilder {
             Self::FloatIndex(index) => index.add_point(id, payload),
             Self::FloatMmapIndex(index) => index.add_point(id, payload),
             Self::GeoIndex(index) => index.add_point(id, payload),
+            Self::GeoMmapIndex(index) => index.add_point(id, payload),
             Self::BinaryIndex(index) => index.add_point(id, payload),
             Self::FullTextIndex(index) => index.add_point(id, payload),
             Self::FullTextMmapIndex(builder) => {
@@ -488,6 +491,7 @@ impl FieldIndexBuilderTrait for FieldIndexBuilder {
             Self::FloatIndex(index) => FieldIndex::FloatIndex(index.finalize()?),
             Self::FloatMmapIndex(index) => FieldIndex::FloatIndex(index.finalize()?),
             Self::GeoIndex(index) => FieldIndex::GeoIndex(index.finalize()?),
+            Self::GeoMmapIndex(index) => FieldIndex::GeoIndex(index.finalize()?),
             Self::BinaryIndex(index) => FieldIndex::BinaryIndex(index.finalize()?),
             Self::FullTextIndex(index) => FieldIndex::FullTextIndex(index.finalize()?),
             Self::FullTextMmapIndex(builder) => FieldIndex::FullTextIndex(builder.finalize()?),

lib/segment/src/index/field_index/index_selector.rs

@@ -5,6 +5,7 @@ use parking_lot::RwLock;
 use rocksdb::DB;
 
 use super::binary_index::BinaryIndex;
+use super::geo_index::{GeoMapIndexBuilder, GeoMapIndexMmapBuilder};
 use super::histogram::Numericable;
 use super::map_index::{MapIndex, MapIndexBuilder, MapIndexKey, MapIndexMmapBuilder};
 use super::mmap_point_to_values::MmapValue;
@@ -62,11 +63,7 @@ impl<'a> IndexSelector<'a> {
             )
             .collect(),
             PayloadSchemaParams::Float(_) => vec![FieldIndex::FloatIndex(self.numeric_new(field)?)],
-            PayloadSchemaParams::Geo(_) => vec![FieldIndex::GeoIndex(GeoMapIndex::new_memory(
-                self.as_rocksdb()?.db.clone(),
-                &field.to_string(),
-                self.as_rocksdb()?.is_appendable,
-            ))],
+            PayloadSchemaParams::Geo(_) => vec![FieldIndex::GeoIndex(self.geo_new(field)?)],
             PayloadSchemaParams::Text(text_index_params) => {
                 vec![FieldIndex::FullTextIndex(
                     self.text_new(field, text_index_params.clone())?,
@@ -126,10 +123,11 @@ impl<'a> IndexSelector<'a> {
                 )]
             }
             PayloadSchemaParams::Geo(_) => {
-                vec![FieldIndexBuilder::GeoIndex(GeoMapIndex::builder(
-                    self.as_rocksdb()?.db.clone(),
-                    &field.to_string(),
-                ))]
+                vec![self.geo_builder(
+                    field,
+                    FieldIndexBuilder::GeoIndex,
+                    FieldIndexBuilder::GeoMmapIndex,
+                )]
             }
             PayloadSchemaParams::Text(text_index_params) => {
                 vec![self.text_builder(field, text_index_params.clone())]
@@ -218,6 +216,33 @@ impl<'a> IndexSelector<'a> {
         }
     }
 
+    fn geo_new(&self, field: &JsonPath) -> OperationResult<GeoMapIndex> {
+        Ok(match self {
+            IndexSelector::RocksDb(IndexSelectorRocksDb { db, is_appendable }) => {
+                GeoMapIndex::new_memory(Arc::clone(db), &field.to_string(), *is_appendable)
+            }
+            IndexSelector::OnDisk(IndexSelectorOnDisk { dir }) => {
+                GeoMapIndex::new_mmap(&map_dir(dir, field))?
+            }
+        })
+    }
+
+    fn geo_builder(
+        &self,
+        field: &JsonPath,
+        make_rocksdb: fn(GeoMapIndexBuilder) -> FieldIndexBuilder,
+        make_mmap: fn(GeoMapIndexMmapBuilder) -> FieldIndexBuilder,
+    ) -> FieldIndexBuilder {
+        match self {
+            IndexSelector::RocksDb(IndexSelectorRocksDb { db, .. }) => {
+                make_rocksdb(GeoMapIndex::builder(Arc::clone(db), &field.to_string()))
+            }
+            IndexSelector::OnDisk(IndexSelectorOnDisk { dir }) => {
+                make_mmap(GeoMapIndex::mmap_builder(&map_dir(dir, field)))
+            }
+        }
+    }
+
     fn text_new(
         &self,
         field: &JsonPath,

lib/segment/src/types.rs

@@ -1281,7 +1281,7 @@ impl PayloadSchemaParams {
             PayloadSchemaParams::Datetime(i) => i.on_disk.unwrap_or_default(),
             PayloadSchemaParams::Uuid(i) => i.on_disk.unwrap_or_default(),
             PayloadSchemaParams::Text(i) => i.on_disk.unwrap_or_default(),
-            PayloadSchemaParams::Geo(_) => false,
+            PayloadSchemaParams::Geo(i) => i.on_disk.unwrap_or_default(),
             PayloadSchemaParams::Bool(_) => false,
         }
     }