From aa7e33b68b1bbfc64cef22a7abc91c980f765f19 Mon Sep 17 00:00:00 2001 From: Ivan Pleshkov Date: Mon, 7 Oct 2024 14:35:17 +0200 Subject: [PATCH] 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 | 5 +++ docs/redoc/master/openapi.json | 5 +++ lib/api/src/grpc/conversions.rs | 9 ++-- lib/api/src/grpc/proto/collections.proto | 1 + lib/api/src/grpc/qdrant.rs | 6 ++- lib/segment/src/data_types/index.rs | 1 - .../src/index/field_index/field_index_base.rs | 6 ++- .../src/index/field_index/index_selector.rs | 43 +++++++++++++++---- lib/segment/src/types.rs | 2 +- 9 files changed, 62 insertions(+), 16 deletions(-) diff --git a/docs/grpc/docs.md b/docs/grpc/docs.md index 93d353a86c..dcfa4e2eb8 100644 --- a/docs/grpc/docs.md +++ b/docs/grpc/docs.md @@ -822,6 +822,11 @@ +| Field | Type | Label | Description | +| ----- | ---- | ----- | ----------- | +| on_disk | [bool](#bool) | optional | If true - store index on disk. | + + diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json index 7d6a111242..2cdd45bc18 100644 --- a/docs/redoc/master/openapi.json +++ b/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 } } }, diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index d867f9b97c..56d62d323e 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -258,9 +258,11 @@ impl From for PayloadIndexParams { } impl From 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 for segment::data_types::index::FloatIndexParams impl TryFrom for segment::data_types::index::GeoIndexParams { type Error = Status; - fn try_from(_params: GeoIndexParams) -> Result { + fn try_from(params: GeoIndexParams) -> Result { Ok(segment::data_types::index::GeoIndexParams { r#type: GeoIndexType::Geo, + on_disk: params.on_disk, }) } } diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index cfc64a7b0a..474735f62c 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/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 { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 6e0db7663f..514b1c99f6 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/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, +} #[derive(serde::Serialize)] #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/lib/segment/src/data_types/index.rs b/lib/segment/src/data_types/index.rs index f030f976f0..91b0b600dc 100644 --- a/lib/segment/src/data_types/index.rs +++ b/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, } diff --git a/lib/segment/src/index/field_index/field_index_base.rs b/lib/segment/src/index/field_index/field_index_base.rs index 9b1278cc51..ef3bd83502 100644 --- a/lib/segment/src/index/field_index/field_index_base.rs +++ b/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), FloatMmapIndex(NumericIndexMmapBuilder), 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()?), diff --git a/lib/segment/src/index/field_index/index_selector.rs b/lib/segment/src/index/field_index/index_selector.rs index aa7c2a2ac8..b8753b5e98 100644 --- a/lib/segment/src/index/field_index/index_selector.rs +++ b/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 { + 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, diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index e8ff21fbc8..7b35e4894c 100644 --- a/lib/segment/src/types.rs +++ b/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, } }