diff --git a/docs/grpc/docs.md b/docs/grpc/docs.md index 3a547b7f0b..8a759e6a23 100644 --- a/docs/grpc/docs.md +++ b/docs/grpc/docs.md @@ -1431,6 +1431,7 @@ Note: 1kB = 1 vector of size 256. | | search_allow_exact | [bool](#bool) | optional | | | search_max_oversampling | [float](#float) | optional | | | upsert_max_batchsize | [uint64](#uint64) | optional | | +| max_collection_vector_size_bytes | [uint64](#uint64) | optional | | diff --git a/docs/redoc/master/openapi.json b/docs/redoc/master/openapi.json index 53f7305cad..b17b517b32 100644 --- a/docs/redoc/master/openapi.json +++ b/docs/redoc/master/openapi.json @@ -7291,6 +7291,13 @@ "format": "uint", "minimum": 0, "nullable": true + }, + "max_collection_vector_size_bytes": { + "description": "Max size of a collections vector storage in bytes", + "type": "integer", + "format": "uint", + "minimum": 0, + "nullable": true } } }, diff --git a/lib/api/src/grpc/conversions.rs b/lib/api/src/grpc/conversions.rs index be81b2f236..4537e65e5a 100644 --- a/lib/api/src/grpc/conversions.rs +++ b/lib/api/src/grpc/conversions.rs @@ -1632,6 +1632,9 @@ impl From for segment::types::StrictModeConfig { search_allow_exact: value.search_allow_exact, search_max_oversampling: value.search_max_oversampling.map(f64::from), upsert_max_batchsize: value.upsert_max_batchsize.map(|i| i as usize), + max_collection_vector_size_bytes: value + .max_collection_vector_size_bytes + .map(|i| i as usize), } } } @@ -1648,6 +1651,9 @@ impl From for StrictModeConfig { search_allow_exact: value.search_allow_exact, search_max_oversampling: value.search_max_oversampling.map(|i| i as f32), upsert_max_batchsize: value.upsert_max_batchsize.map(|i| i as u64), + max_collection_vector_size_bytes: value + .max_collection_vector_size_bytes + .map(|i| i as u64), } } } diff --git a/lib/api/src/grpc/proto/collections.proto b/lib/api/src/grpc/proto/collections.proto index 848b1fe496..fd10cdd2d6 100644 --- a/lib/api/src/grpc/proto/collections.proto +++ b/lib/api/src/grpc/proto/collections.proto @@ -321,6 +321,7 @@ message StrictModeConfig { optional bool search_allow_exact = 7; optional float search_max_oversampling = 8; optional uint64 upsert_max_batchsize = 9; + optional uint64 max_collection_vector_size_bytes = 10; } message CreateCollection { diff --git a/lib/api/src/grpc/qdrant.rs b/lib/api/src/grpc/qdrant.rs index 15215e7a5c..1e338e9469 100644 --- a/lib/api/src/grpc/qdrant.rs +++ b/lib/api/src/grpc/qdrant.rs @@ -454,6 +454,8 @@ pub struct StrictModeConfig { pub search_max_oversampling: ::core::option::Option, #[prost(uint64, optional, tag = "9")] pub upsert_max_batchsize: ::core::option::Option, + #[prost(uint64, optional, tag = "10")] + pub max_collection_vector_size_bytes: ::core::option::Option, } #[derive(validator::Validate)] #[derive(serde::Serialize)] diff --git a/lib/api/src/rest/schema.rs b/lib/api/src/rest/schema.rs index dbffc843fc..033b58b7ea 100644 --- a/lib/api/src/rest/schema.rs +++ b/lib/api/src/rest/schema.rs @@ -1036,3 +1036,17 @@ impl Validate for PointInsertOperations { } } } + +impl PointInsertOperations { + /// Amount of vectors in the operation request. + pub fn len(&self) -> usize { + match self { + PointInsertOperations::PointsBatch(batch) => batch.batch.ids.len(), + PointInsertOperations::PointsList(list) => list.points.len(), + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} diff --git a/lib/collection/src/collection/mod.rs b/lib/collection/src/collection/mod.rs index b524ddb0cf..751d99d5c6 100644 --- a/lib/collection/src/collection/mod.rs +++ b/lib/collection/src/collection/mod.rs @@ -28,6 +28,7 @@ use tokio::sync::{Mutex, RwLock, RwLockWriteGuard}; use crate::collection::payload_index_schema::PayloadIndexSchema; use crate::collection_state::{ShardInfo, State}; use crate::common::is_ready::IsReady; +use crate::common::local_data_stats::{LocalDataStats, LocalDataStatsCache}; use crate::config::CollectionConfigInternal; use crate::operations::config_diff::{DiffConfig, OptimizersConfigDiff}; use crate::operations::shared_storage_config::SharedStorageConfig; @@ -80,6 +81,8 @@ pub struct Collection { // Search runtime handle. search_runtime: Handle, optimizer_cpu_budget: CpuBudget, + // Cached stats over all local shards used in strict mode, may be outdated + local_stats_cache: LocalDataStatsCache, } pub type RequestShardTransfer = Arc; @@ -150,6 +153,10 @@ impl Collection { let locked_shard_holder = Arc::new(LockedShardHolder::new(shard_holder)); + let local_stats_cache = LocalDataStatsCache::new_with_values( + Self::calculate_segment_stats(&locked_shard_holder).await, + ); + // Once the config is persisted - the collection is considered to be successfully created. CollectionVersion::save(path)?; collection_config.save(path)?; @@ -176,6 +183,7 @@ impl Collection { update_runtime: update_runtime.unwrap_or_else(Handle::current), search_runtime: search_runtime.unwrap_or_else(Handle::current), optimizer_cpu_budget, + local_stats_cache, }) } @@ -263,6 +271,10 @@ impl Collection { let locked_shard_holder = Arc::new(LockedShardHolder::new(shard_holder)); + let local_stats_cache = LocalDataStatsCache::new_with_values( + Self::calculate_segment_stats(&locked_shard_holder).await, + ); + Self { id: collection_id.clone(), shards_holder: locked_shard_holder, @@ -285,6 +297,7 @@ impl Collection { update_runtime: update_runtime.unwrap_or_else(Handle::current), search_runtime: search_runtime.unwrap_or_else(Handle::current), optimizer_cpu_budget, + local_stats_cache, } } @@ -784,6 +797,31 @@ impl Collection { pub async fn trigger_optimizers(&self) { self.shards_holder.read().await.trigger_optimizers().await; } + + async fn calculate_segment_stats(shards_holder: &Arc>) -> LocalDataStats { + let shard_lock = shards_holder.read().await; + shard_lock.calculate_local_segments_stats().await + } + + /// Checks and performs a cache update for local data statistics if needed. + /// Returns `Some(..)` with the new values if a cache update has been performed and `None` otherwise. + async fn check_and_update_local_size_stats(&self) -> Option { + if self.local_stats_cache.check_need_update_and_increment() { + let new_stats = Self::calculate_segment_stats(&self.shards_holder).await; + self.local_stats_cache.update(new_stats); + return Some(new_stats); + } + + None + } + + /// Returns the estimated local vector storage size for this collection, cached and auto-updated. + pub async fn estimated_local_vector_storage_size(&self) -> usize { + if let Some(shard_stats) = self.check_and_update_local_size_stats().await { + return shard_stats.vector_storage_size; + } + self.local_stats_cache.get_vector_storage() + } } struct CollectionVersion; diff --git a/lib/collection/src/collection_manager/holders/proxy_segment.rs b/lib/collection/src/collection_manager/holders/proxy_segment.rs index 5e62d73bbb..8d36003480 100644 --- a/lib/collection/src/collection_manager/holders/proxy_segment.rs +++ b/lib/collection/src/collection_manager/holders/proxy_segment.rs @@ -944,9 +944,14 @@ impl SegmentEntry for ProxySegment { SegmentType::Special } + fn size_info(&self) -> SegmentInfo { + // To reduce code complexity for estimations, we use `.info()` directly here. + self.info() + } + fn info(&self) -> SegmentInfo { let wrapped_info = self.wrapped_segment.get().read().info(); - let write_info = self.write_segment.get().read().info(); + let write_info = self.write_segment.get().read().size_info(); // Only fields set by `size_info()` needed! let vector_name_count = self.config().vector_data.len() + self.config().sparse_vector_data.len(); diff --git a/lib/collection/src/common/local_data_stats.rs b/lib/collection/src/common/local_data_stats.rs new file mode 100644 index 0000000000..773fcebb30 --- /dev/null +++ b/lib/collection/src/common/local_data_stats.rs @@ -0,0 +1,65 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +/// Amount of requests that have to be done until the cached data gets updated. +const UPDATE_INTERVAL: usize = 32; + +/// A cache for `LocalDataStats` utilizing `AtomicUsize` for better performance. +#[derive(Default)] +pub(crate) struct LocalDataStatsCache { + vector_storage_size: AtomicUsize, + + request_counter: AtomicUsize, +} + +impl LocalDataStatsCache { + pub fn new_with_values(stats: LocalDataStats) -> Self { + let LocalDataStats { + vector_storage_size, + } = stats; + let vector_storage_size = AtomicUsize::new(vector_storage_size); + Self { + vector_storage_size, + request_counter: AtomicUsize::new(1), // Prevent same data getting loaded a second time when doing the first request. + } + } + + /// Checks whether the cache needs to be updated. + /// For performance reasons, this also assumes a cached value gets read afterwards and brings the + /// Update counter one tick closer to the next update. + pub fn check_need_update_and_increment(&self) -> bool { + let req_counter = self.request_counter.fetch_add(1, Ordering::Relaxed); + req_counter % UPDATE_INTERVAL == 0 + } + + /// Sets all cache values to `new_stats`. + pub fn update(&self, new_stats: LocalDataStats) { + let LocalDataStats { + vector_storage_size, + } = new_stats; + + self.vector_storage_size + .store(vector_storage_size, Ordering::Relaxed); + } + + /// Returns cached vector storage size estimation. + pub fn get_vector_storage(&self) -> usize { + self.vector_storage_size.load(Ordering::Relaxed) + } +} + +/// Statistics for local data, like the size of vector storage. +#[derive(Clone, Copy, Default)] +pub struct LocalDataStats { + /// Estimated amount of vector storage size. + pub vector_storage_size: usize, +} + +impl LocalDataStats { + pub fn accumulate_from(&mut self, other: &Self) { + let LocalDataStats { + vector_storage_size, + } = other; + + self.vector_storage_size += vector_storage_size; + } +} diff --git a/lib/collection/src/common/mod.rs b/lib/collection/src/common/mod.rs index bab1e72654..0f202b7cdf 100644 --- a/lib/collection/src/common/mod.rs +++ b/lib/collection/src/common/mod.rs @@ -3,6 +3,7 @@ pub mod eta_calculator; pub mod fetch_vectors; pub mod file_utils; pub mod is_ready; +pub mod local_data_stats; pub mod retrieve_request_trait; pub mod sha_256; pub mod snapshot_stream; diff --git a/lib/collection/src/operations/verification/discovery.rs b/lib/collection/src/operations/verification/discovery.rs index 6ef291f65c..d77ca40d7d 100644 --- a/lib/collection/src/operations/verification/discovery.rs +++ b/lib/collection/src/operations/verification/discovery.rs @@ -27,14 +27,15 @@ impl StrictModeVerification for DiscoverRequestInternal { } impl StrictModeVerification for DiscoverRequestBatch { - fn check_strict_mode( + async fn check_strict_mode( &self, collection: &Collection, strict_mode_config: &StrictModeConfig, ) -> Result<(), CollectionError> { for i in self.searches.iter() { i.discover_request - .check_strict_mode(collection, strict_mode_config)?; + .check_strict_mode(collection, strict_mode_config) + .await?; } Ok(()) diff --git a/lib/collection/src/operations/verification/mod.rs b/lib/collection/src/operations/verification/mod.rs index 454d4da125..925f007f99 100644 --- a/lib/collection/src/operations/verification/mod.rs +++ b/lib/collection/src/operations/verification/mod.rs @@ -35,7 +35,8 @@ pub struct VerificationPass { /// This trait ignores the `enabled` parameter in `StrictModeConfig`. pub trait StrictModeVerification { /// Implementing this method allows adding a custom check for request specific values. - fn check_custom( + #[allow(async_fn_in_trait)] + async fn check_custom( &self, _collection: &Collection, _strict_mode_config: &StrictModeConfig, @@ -86,13 +87,14 @@ pub trait StrictModeVerification { } /// Checks search parameters. - fn check_search_params( + #[allow(async_fn_in_trait)] + async fn check_search_params( &self, collection: &Collection, strict_mode_config: &StrictModeConfig, ) -> Result<(), CollectionError> { if let Some(search_params) = self.request_search_params() { - search_params.check_strict_mode(collection, strict_mode_config)?; + Box::pin(search_params.check_strict_mode(collection, strict_mode_config)).await?; } Ok(()) } @@ -140,16 +142,18 @@ pub trait StrictModeVerification { /// Does the verification of all configured parameters. Only implement this function if you know what /// you are doing. In most cases implementing `check_custom` is sufficient. - fn check_strict_mode( + #[allow(async_fn_in_trait)] + async fn check_strict_mode( &self, collection: &Collection, strict_mode_config: &StrictModeConfig, ) -> Result<(), CollectionError> { - self.check_custom(collection, strict_mode_config)?; + self.check_custom(collection, strict_mode_config).await?; self.check_request_query_limit(strict_mode_config)?; self.check_request_filter(collection, strict_mode_config)?; self.check_request_exact(strict_mode_config)?; - self.check_search_params(collection, strict_mode_config)?; + self.check_search_params(collection, strict_mode_config) + .await?; Ok(()) } } @@ -196,7 +200,7 @@ pub(crate) fn check_limit_opt( } impl StrictModeVerification for SearchParams { - fn check_custom( + async fn check_custom( &self, _collection: &Collection, strict_mode_config: &StrictModeConfig, @@ -338,6 +342,7 @@ mod test { let strict_mode_config = collection.strict_mode_config().await.unwrap(); let error = request .check_strict_mode(collection, &strict_mode_config) + .await .expect_err("Expected strict mode error but got Ok() value"); if !matches!(error, CollectionError::StrictMode { .. }) { panic!("Expected strict mode error but got {error:#}"); @@ -349,7 +354,9 @@ mod test { collection: &Collection, ) { let strict_mode_config = collection.strict_mode_config().await.unwrap(); - let res = request.check_strict_mode(collection, &strict_mode_config); + let res = request + .check_strict_mode(collection, &strict_mode_config) + .await; if let Err(CollectionError::StrictMode { description }) = res { panic!("Strict mode check should've passed but failed with error: {description:?}"); } else if res.is_err() { @@ -401,6 +408,7 @@ mod test { search_allow_exact: Some(false), search_max_oversampling: Some(0.2), upsert_max_batchsize: None, + max_collection_vector_size_bytes: None, }; fixture_collection(&strict_mode_config).await diff --git a/lib/collection/src/operations/verification/query.rs b/lib/collection/src/operations/verification/query.rs index 190a1aa46e..e8a6be18a0 100644 --- a/lib/collection/src/operations/verification/query.rs +++ b/lib/collection/src/operations/verification/query.rs @@ -8,14 +8,16 @@ use crate::operations::universal_query::collection_query::{ }; impl StrictModeVerification for QueryRequestInternal { - fn check_custom( + async fn check_custom( &self, collection: &Collection, strict_mode_config: &StrictModeConfig, ) -> Result<(), crate::operations::types::CollectionError> { if let Some(prefetch) = &self.prefetch { for prefetch in prefetch { - prefetch.check_strict_mode(collection, strict_mode_config)?; + prefetch + .check_strict_mode(collection, strict_mode_config) + .await?; } } @@ -44,7 +46,7 @@ impl StrictModeVerification for QueryRequestInternal { } impl StrictModeVerification for Prefetch { - fn check_custom( + async fn check_custom( &self, collection: &Collection, strict_mode_config: &StrictModeConfig, @@ -52,7 +54,7 @@ impl StrictModeVerification for Prefetch { // Prefetch.prefetch is of type Prefetch (recursive type) if let Some(prefetch) = &self.prefetch { for prefetch in prefetch { - prefetch.check_strict_mode(collection, strict_mode_config)?; + Box::pin(prefetch.check_strict_mode(collection, strict_mode_config)).await?; } } diff --git a/lib/collection/src/operations/verification/search.rs b/lib/collection/src/operations/verification/search.rs index bf6504f3ce..febd6a80f7 100644 --- a/lib/collection/src/operations/verification/search.rs +++ b/lib/collection/src/operations/verification/search.rs @@ -50,7 +50,7 @@ impl StrictModeVerification for CoreSearchRequest { } impl StrictModeVerification for SearchRequestBatch { - fn check_strict_mode( + async fn check_strict_mode( &self, collection: &Collection, strict_mode_config: &StrictModeConfig, @@ -58,7 +58,8 @@ impl StrictModeVerification for SearchRequestBatch { for search_request in &self.searches { search_request .search_request - .check_strict_mode(collection, strict_mode_config)?; + .check_strict_mode(collection, strict_mode_config) + .await?; } Ok(()) } diff --git a/lib/collection/src/operations/verification/update.rs b/lib/collection/src/operations/verification/update.rs index 29d1a11371..1f0d5ca6f7 100644 --- a/lib/collection/src/operations/verification/update.rs +++ b/lib/collection/src/operations/verification/update.rs @@ -1,4 +1,4 @@ -use api::rest::PointInsertOperations; +use api::rest::{PointInsertOperations, UpdateVectors}; use segment::types::{Filter, StrictModeConfig}; use super::{check_limit_opt, StrictModeVerification}; @@ -56,6 +56,18 @@ impl StrictModeVerification for DeleteVectors { } impl StrictModeVerification for SetPayload { + // TODO: Payload storage size limit + /* + async fn check_custom( + &self, + collection: &Collection, + strict_mode_config: &StrictModeConfig, + ) -> Result<(), CollectionError> { + check_collection_vector_size_limit(collection, strict_mode_config).await?; + Ok(()) + } + */ + fn indexed_filter_write(&self) -> Option<&Filter> { self.filter.as_ref() } @@ -100,22 +112,19 @@ impl StrictModeVerification for DeletePayload { } impl StrictModeVerification for PointInsertOperations { - fn check_custom( + async fn check_custom( &self, - _collection: &Collection, + collection: &Collection, strict_mode_config: &StrictModeConfig, ) -> Result<(), CollectionError> { - let len = match self { - PointInsertOperations::PointsBatch(batch) => batch.batch.ids.len(), - PointInsertOperations::PointsList(list) => list.points.len(), - }; - check_limit_opt( - Some(len), + Some(self.len()), strict_mode_config.upsert_max_batchsize, "upsert limit", )?; + check_collection_vector_size_limit(collection, strict_mode_config).await?; + Ok(()) } @@ -139,3 +148,57 @@ impl StrictModeVerification for PointInsertOperations { None } } + +impl StrictModeVerification for UpdateVectors { + async fn check_custom( + &self, + collection: &Collection, + strict_mode_config: &StrictModeConfig, + ) -> Result<(), CollectionError> { + check_limit_opt( + Some(self.points.len()), + strict_mode_config.upsert_max_batchsize, + "update limit", + )?; + + check_collection_vector_size_limit(collection, strict_mode_config).await?; + Ok(()) + } + + fn query_limit(&self) -> Option { + None + } + + fn indexed_filter_read(&self) -> Option<&Filter> { + None + } + + fn indexed_filter_write(&self) -> Option<&Filter> { + None + } + + fn request_exact(&self) -> Option { + None + } + + fn request_search_params(&self) -> Option<&segment::types::SearchParams> { + None + } +} + +async fn check_collection_vector_size_limit( + collection: &Collection, + strict_mode_config: &StrictModeConfig, +) -> Result<(), CollectionError> { + if let Some(max_vec_storage_size_bytes) = strict_mode_config.max_collection_vector_size_bytes { + let vec_storage_size_bytes = collection.estimated_local_vector_storage_size().await; + if vec_storage_size_bytes >= max_vec_storage_size_bytes { + let size_in_mb = max_vec_storage_size_bytes as f32 / (1024.0 * 1024.0); + return Err(CollectionError::bad_request(format!( + "Max vector storage size limit of {size_in_mb}MB reached!", + ))); + } + } + + Ok(()) +} diff --git a/lib/collection/src/shards/replica_set/mod.rs b/lib/collection/src/shards/replica_set/mod.rs index 57c08f6a1e..b75d420e24 100644 --- a/lib/collection/src/shards/replica_set/mod.rs +++ b/lib/collection/src/shards/replica_set/mod.rs @@ -26,6 +26,7 @@ use super::remote_shard::RemoteShard; use super::transfer::ShardTransfer; use super::CollectionId; use crate::collection::payload_index_schema::PayloadIndexSchema; +use crate::common::local_data_stats::LocalDataStats; use crate::common::snapshots_manager::SnapshotStorageManager; use crate::config::CollectionConfigInternal; use crate::operations::point_ops::{self}; @@ -1010,6 +1011,35 @@ impl ShardReplicaSet { shard.trigger_optimizers(); true } + + /// Returns the estimated size of all locally stored vectors in bytes. + /// Locks and iterates over all segments. + /// Cache this value in performance critical scenarios! + pub(crate) async fn calculate_local_shards_stats(&self) -> LocalDataStats { + self.local + .read() + .await + .as_ref() + .map(|i| match i { + Shard::Local(local) => { + let mut total_vector_size = 0; + + for segment in local.segments.read().iter() { + let size_info = segment.1.get().read().size_info(); + total_vector_size += size_info.vectors_size_bytes; + } + + LocalDataStats { + vector_storage_size: total_vector_size, + } + } + Shard::Proxy(_) + | Shard::ForwardProxy(_) + | Shard::QueueProxy(_) + | Shard::Dummy(_) => LocalDataStats::default(), + }) + .unwrap_or_default() + } } /// Represents a replica set state diff --git a/lib/collection/src/shards/shard_holder/mod.rs b/lib/collection/src/shards/shard_holder/mod.rs index 67a6543d40..d58d5ae5c3 100644 --- a/lib/collection/src/shards/shard_holder/mod.rs +++ b/lib/collection/src/shards/shard_holder/mod.rs @@ -23,6 +23,7 @@ use super::resharding::tasks_pool::ReshardTasksPool; use super::resharding::{ReshardStage, ReshardState}; use super::transfer::transfer_tasks_pool::TransferTasksPool; use crate::collection::payload_index_schema::PayloadIndexSchema; +use crate::common::local_data_stats::LocalDataStats; use crate::common::snapshot_stream::SnapshotStream; use crate::config::{CollectionConfigInternal, ShardingMethod}; use crate::hash_ring::HashRingRouter; @@ -1117,6 +1118,15 @@ impl ShardHolder { } Ok(()) } + + /// Queries and accumulates the statistics for local data, uncached. + pub async fn calculate_local_segments_stats(&self) -> LocalDataStats { + let mut stats = LocalDataStats::default(); + for shard in self.shards.iter() { + stats.accumulate_from(&shard.1.calculate_local_shards_stats().await) + } + stats + } } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/lib/segment/Cargo.toml b/lib/segment/Cargo.toml index 652ba5c392..725161fbc5 100644 --- a/lib/segment/Cargo.toml +++ b/lib/segment/Cargo.toml @@ -12,7 +12,11 @@ edition = "2021" workspace = true [features] -multiling-chinese = ["charabia/chinese-segmentation", "charabia/chinese-normalization", "charabia/chinese-normalization-pinyin"] +multiling-chinese = [ + "charabia/chinese-segmentation", + "charabia/chinese-normalization", + "charabia/chinese-normalization-pinyin", +] multiling-japanese = ["charabia/japanese"] multiling-korean = ["charabia/korean"] testing = ["common/testing", "sparse/testing"] @@ -44,7 +48,10 @@ tempfile = { workspace = true } parking_lot = { workspace = true } rayon = { workspace = true } itertools = { workspace = true } -rocksdb = { version = "0.22.0", default-features = false, features = ["snappy", "lz4"] } +rocksdb = { version = "0.22.0", default-features = false, features = [ + "snappy", + "lz4", +] } uuid = { workspace = true } bincode = "1.3" serde = { workspace = true } @@ -77,7 +84,7 @@ chrono = { workspace = true } smol_str = { version = "0.3.2", features = ["serde"] } fnv = { workspace = true } indexmap = { workspace = true } -ahash = { workspace = true } +ahash = { workspace = true } http = "1.0.0" sha2 = { workspace = true } smallvec = "1.13.2" @@ -88,7 +95,11 @@ zerocopy = { workspace = true } lazy_static = "1.5.0" sysinfo = "0.32" -charabia = { version = "0.9.2", default-features = false, features = ["greek", "hebrew", "thai"] } +charabia = { version = "0.9.2", default-features = false, features = [ + "greek", + "hebrew", + "thai", +] } blob_store = { path = "../blob_store" } @@ -178,3 +189,7 @@ harness = false [[bench]] name = "scorer_mmap" harness = false + +[[bench]] +name = "segment_info" +harness = false diff --git a/lib/segment/benches/segment_info.rs b/lib/segment/benches/segment_info.rs new file mode 100644 index 0000000000..47fa9bab90 --- /dev/null +++ b/lib/segment/benches/segment_info.rs @@ -0,0 +1,76 @@ +use std::collections::HashMap; + +use criterion::{criterion_group, criterion_main, Criterion}; +use segment::data_types::vectors::{only_default_vector, DEFAULT_VECTOR_NAME}; +use segment::entry::entry_point::SegmentEntry; +use segment::json_path::JsonPath; +use segment::segment_constructor::build_segment; +use segment::types::{ + Distance, Indexes, Payload, PayloadFieldSchema, PayloadSchemaType, SegmentConfig, + VectorDataConfig, VectorStorageType, +}; +use serde_json::{Map, Value}; +use tempfile::Builder; + +pub fn criterion_benchmark(c: &mut Criterion) { + let dir = Builder::new().prefix("segment_dir").tempdir().unwrap(); + let dim = 400; + let config = SegmentConfig { + vector_data: HashMap::from([( + DEFAULT_VECTOR_NAME.to_owned(), + VectorDataConfig { + size: dim, + distance: Distance::Dot, + storage_type: VectorStorageType::Memory, + index: Indexes::Plain {}, + quantization_config: None, + multivector_config: None, + datatype: None, + }, + )]), + sparse_vector_data: Default::default(), + payload_storage_type: Default::default(), + }; + let mut segment = build_segment(dir.path(), &config, true).unwrap(); + + let vector = vec![0.1f32; 400]; + + let mut payload: Map = Map::default(); + + for i in 0..3 { + let key = format!("key{i}"); + payload.insert(key.clone(), "value".to_string().into()); + segment + .create_field_index( + 100, + &JsonPath::new(&key), + Some(&PayloadFieldSchema::FieldType(PayloadSchemaType::Keyword)), + ) + .unwrap(); + } + let payload = Payload::from(payload); + + for id in 0..100000u64 { + segment + .upsert_point(100, id.into(), only_default_vector(&vector)) + .unwrap(); + segment + .set_payload(100, id.into(), &payload, &None) + .unwrap(); + } + + c.bench_function("segment-info", |b| { + b.iter(|| { + let _ = segment.info(); + }) + }); + + c.bench_function("segment-size-info", |b| { + b.iter(|| { + let _ = segment.size_info(); + }) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/lib/segment/src/entry/entry_point.rs b/lib/segment/src/entry/entry_point.rs index 62c908d902..53756ba544 100644 --- a/lib/segment/src/entry/entry_point.rs +++ b/lib/segment/src/entry/entry_point.rs @@ -209,6 +209,10 @@ pub trait SegmentEntry { /// Get current stats of the segment fn info(&self) -> SegmentInfo; + /// Get size related stats of the segment. + /// This returns `SegmentInfo` with some non size-related data (like `schema`) unset to improve performance. + fn size_info(&self) -> SegmentInfo; + /// Get segment configuration fn config(&self) -> &SegmentConfig; diff --git a/lib/segment/src/segment/entry.rs b/lib/segment/src/segment/entry.rs index c31e165bbb..e98d67b992 100644 --- a/lib/segment/src/segment/entry.rs +++ b/lib/segment/src/segment/entry.rs @@ -434,17 +434,7 @@ impl SegmentEntry for Segment { self.segment_type } - fn info(&self) -> SegmentInfo { - let payload_index = self.payload_index.borrow(); - let schema = payload_index - .indexed_fields() - .into_iter() - .map(|(key, index_schema)| { - let points_count = payload_index.indexed_points(&key); - (key, PayloadIndexInfo::new(index_schema, points_count)) - }) - .collect(); - + fn size_info(&self) -> SegmentInfo { let num_vectors = self .vector_data .values() @@ -513,11 +503,28 @@ impl SegmentEntry for Segment { ram_usage_bytes: 0, // ToDo: Implement disk_usage_bytes: 0, // ToDo: Implement is_appendable: self.appendable_flag, - index_schema: schema, + index_schema: HashMap::new(), vector_data: vector_data_info, } } + fn info(&self) -> SegmentInfo { + let payload_index = self.payload_index.borrow(); + let schema = payload_index + .indexed_fields() + .into_iter() + .map(|(key, index_schema)| { + let points_count = payload_index.indexed_points(&key); + (key, PayloadIndexInfo::new(index_schema, points_count)) + }) + .collect(); + + let mut info = self.size_info(); + info.index_schema = schema; + + info + } + fn config(&self) -> &SegmentConfig { &self.segment_config } diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs index 309187721d..e50a841c3d 100644 --- a/lib/segment/src/types.rs +++ b/lib/segment/src/types.rs @@ -707,6 +707,10 @@ pub struct StrictModeConfig { /// Max batchsize when upserting #[serde(skip_serializing_if = "Option::is_none")] pub upsert_max_batchsize: Option, + + /// Max size of a collections vector storage in bytes + #[serde(skip_serializing_if = "Option::is_none")] + pub max_collection_vector_size_bytes: Option, } impl Eq for StrictModeConfig {} @@ -724,6 +728,7 @@ impl Hash for StrictModeConfig { // We skip hashing this field because we cannot reliably hash a float search_max_oversampling: _, upsert_max_batchsize, + max_collection_vector_size_bytes, } = self; ( enabled, @@ -734,6 +739,7 @@ impl Hash for StrictModeConfig { search_max_hnsw_ef, search_allow_exact, upsert_max_batchsize, + max_collection_vector_size_bytes, ) .hash(state); } diff --git a/lib/storage/src/content_manager/collection_verification.rs b/lib/storage/src/content_manager/collection_verification.rs index ee2f666127..babd91c68a 100644 --- a/lib/storage/src/content_manager/collection_verification.rs +++ b/lib/storage/src/content_manager/collection_verification.rs @@ -31,7 +31,9 @@ where if let Some(strict_mode_config) = &collection.strict_mode_config().await { if strict_mode_config.enabled.unwrap_or_default() { for request in requests { - request.check_strict_mode(&collection, strict_mode_config)?; + request + .check_strict_mode(&collection, strict_mode_config) + .await?; } if let Some(timeout) = timeout { diff --git a/lib/storage/src/content_manager/conversions.rs b/lib/storage/src/content_manager/conversions.rs index 31569b84dd..40252f7697 100644 --- a/lib/storage/src/content_manager/conversions.rs +++ b/lib/storage/src/content_manager/conversions.rs @@ -86,6 +86,9 @@ pub fn strict_mode_from_api(value: api::grpc::qdrant::StrictModeConfig) -> Stric search_allow_exact: value.search_allow_exact, search_max_oversampling: value.search_max_oversampling.map(f64::from), upsert_max_batchsize: value.upsert_max_batchsize.map(|i| i as usize), + max_collection_vector_size_bytes: value + .max_collection_vector_size_bytes + .map(|i| i as usize), } } diff --git a/src/actix/api/update_api.rs b/src/actix/api/update_api.rs index c387f11331..f5765f37b7 100644 --- a/src/actix/api/update_api.rs +++ b/src/actix/api/update_api.rs @@ -7,7 +7,6 @@ use collection::operations::payload_ops::{DeletePayload, SetPayload}; use collection::operations::point_ops::{PointsSelector, WriteOrdering}; use collection::operations::types::UpdateResult; use collection::operations::vector_ops::DeleteVectors; -use collection::operations::verification::new_unchecked_verification_pass; use schemars::JsonSchema; use segment::json_path::JsonPath; use serde::{Deserialize, Serialize}; @@ -106,13 +105,16 @@ async fn update_vectors( params: Query, ActixAccess(access): ActixAccess, ) -> impl Responder { - // Nothing to verify here. - let pass = new_unchecked_verification_pass(); - let operation = operation.into_inner(); let wait = params.wait.unwrap_or(false); let ordering = params.ordering.unwrap_or_default(); + let pass = + match check_strict_mode(&operation, None, &collection.name, &dispatcher, &access).await { + Ok(pass) => pass, + Err(err) => return process_response_error(err, Instant::now(), None), + }; + helpers::time(do_update_vectors( dispatcher.toc(&access, &pass).clone(), collection.into_inner().name, diff --git a/src/common/points.rs b/src/common/points.rs index 295a0eaa12..dbe8df053c 100644 --- a/src/common/points.rs +++ b/src/common/points.rs @@ -163,34 +163,60 @@ impl StrictModeVerification for UpdateOperation { None } - fn check_strict_mode( + async fn check_strict_mode( &self, collection: &Collection, strict_mode_config: &StrictModeConfig, ) -> Result<(), CollectionError> { match self { - UpdateOperation::Delete(delete_op) => delete_op - .delete - .check_strict_mode(collection, strict_mode_config), - UpdateOperation::SetPayload(set_payload) => set_payload - .set_payload - .check_strict_mode(collection, strict_mode_config), - UpdateOperation::OverwritePayload(overwrite_payload) => overwrite_payload - .overwrite_payload - .check_strict_mode(collection, strict_mode_config), - UpdateOperation::DeletePayload(delete_payload) => delete_payload - .delete_payload - .check_strict_mode(collection, strict_mode_config), - UpdateOperation::ClearPayload(clear_payload) => clear_payload - .clear_payload - .check_strict_mode(collection, strict_mode_config), - UpdateOperation::DeleteVectors(delete_op) => delete_op - .delete_vectors - .check_strict_mode(collection, strict_mode_config), - UpdateOperation::Upsert(upsert_op) => upsert_op - .upsert - .check_strict_mode(collection, strict_mode_config), - UpdateOperation::UpdateVectors(_) => Ok(()), + UpdateOperation::Delete(delete_op) => { + delete_op + .delete + .check_strict_mode(collection, strict_mode_config) + .await + } + UpdateOperation::SetPayload(set_payload) => { + set_payload + .set_payload + .check_strict_mode(collection, strict_mode_config) + .await + } + UpdateOperation::OverwritePayload(overwrite_payload) => { + overwrite_payload + .overwrite_payload + .check_strict_mode(collection, strict_mode_config) + .await + } + UpdateOperation::DeletePayload(delete_payload) => { + delete_payload + .delete_payload + .check_strict_mode(collection, strict_mode_config) + .await + } + UpdateOperation::ClearPayload(clear_payload) => { + clear_payload + .clear_payload + .check_strict_mode(collection, strict_mode_config) + .await + } + UpdateOperation::DeleteVectors(delete_op) => { + delete_op + .delete_vectors + .check_strict_mode(collection, strict_mode_config) + .await + } + UpdateOperation::Upsert(upsert_op) => { + upsert_op + .upsert + .check_strict_mode(collection, strict_mode_config) + .await + } + UpdateOperation::UpdateVectors(update) => { + update + .update_vectors + .check_strict_mode(collection, strict_mode_config) + .await + } } } } diff --git a/src/tonic/api/points_api.rs b/src/tonic/api/points_api.rs index f4e4f1b0c4..9b0cf93e36 100644 --- a/src/tonic/api/points_api.rs +++ b/src/tonic/api/points_api.rs @@ -16,7 +16,6 @@ use api::grpc::qdrant::{ UpdatePointVectors, UpsertPoints, }; use collection::operations::types::CoreSearchRequest; -use collection::operations::verification::new_unchecked_verification_pass; use common::counter::hardware_accumulator::HwMeasurementAcc; use storage::content_manager::toc::request_hw_counter::RequestHwCounter; use storage::dispatcher::Dispatcher; @@ -66,13 +65,10 @@ impl Points for PointsService { ) -> Result, Status> { validate(request.get_ref())?; - // Nothing to verify here. - let pass = new_unchecked_verification_pass(); - let access = extract_access(&mut request); upsert( - self.dispatcher.toc(&access, &pass).clone(), + StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, None, @@ -122,12 +118,11 @@ impl Points for PointsService { validate(request.get_ref())?; // Nothing to verify here. - let pass = new_unchecked_verification_pass(); let access = extract_access(&mut request); update_vectors( - self.dispatcher.toc(&access, &pass).clone(), + StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, None, diff --git a/src/tonic/api/points_common.rs b/src/tonic/api/points_common.rs index fc9b321542..5cfe751047 100644 --- a/src/tonic/api/points_common.rs +++ b/src/tonic/api/points_common.rs @@ -37,7 +37,6 @@ use collection::operations::types::{ RecommendExample, ScrollRequestInternal, }; use collection::operations::vector_ops::DeleteVectors; -use collection::operations::verification::new_unchecked_verification_pass; use collection::operations::{ClockTag, CollectionUpdateOperations, OperationWithClockTag}; use collection::shards::shard::ShardId; use common::counter::hardware_accumulator::HwMeasurementAcc; @@ -112,7 +111,7 @@ pub(crate) fn convert_shard_selector_for_read( } pub async fn upsert( - toc: Arc, + toc_provider: impl CheckedTocProvider, upsert_points: UpsertPoints, clock_tag: Option, shard_selection: Option, @@ -132,9 +131,14 @@ pub async fn upsert( points: points?, shard_key: shard_key_selector.map(ShardKeySelector::from), }); + + let toc = toc_provider + .check_strict_mode(&operation, &collection_name, None, &access) + .await?; + let timing = Instant::now(); let result = do_upsert_points( - toc, + toc.clone(), collection_name, operation, clock_tag, @@ -245,7 +249,7 @@ pub async fn delete( } pub async fn update_vectors( - toc: Arc, + toc_provider: impl CheckedTocProvider, update_point_vectors: UpdatePointVectors, clock_tag: Option, shard_selection: Option, @@ -278,9 +282,13 @@ pub async fn update_vectors( shard_key: shard_key_selector.map(ShardKeySelector::from), }; + let toc = toc_provider + .check_strict_mode(&operation, &collection_name, None, &access) + .await?; + let timing = Instant::now(); let result = do_update_vectors( - toc, + toc.clone(), collection_name, operation, clock_tag, @@ -555,10 +563,8 @@ pub async fn update_batch( points, shard_key_selector, }) => { - // We don't need strict mode checks for upsert! - let toc = dispatcher.toc(&access, &new_unchecked_verification_pass()); upsert( - toc.clone(), + StrictModeCheckedTocProvider::new(dispatcher), UpsertPoints { collection_name, wait, @@ -687,10 +693,8 @@ pub async fn update_batch( shard_key_selector, }, ) => { - // We don't need strict mode checks for vector updates! - let toc = dispatcher.toc(&access, &new_unchecked_verification_pass()); update_vectors( - toc.clone(), + StrictModeCheckedTocProvider::new(dispatcher), UpdatePointVectors { collection_name, wait, diff --git a/src/tonic/api/points_internal_api.rs b/src/tonic/api/points_internal_api.rs index 167ce523b2..5414a085d4 100644 --- a/src/tonic/api/points_internal_api.rs +++ b/src/tonic/api/points_internal_api.rs @@ -187,7 +187,7 @@ impl PointsInternal for PointsInternalService { upsert_points.ok_or_else(|| Status::invalid_argument("UpsertPoints is missing"))?; upsert( - self.toc.clone(), + UncheckedTocProvider::new_unchecked(&self.toc), upsert_points, clock_tag.map(Into::into), shard_id, @@ -237,7 +237,7 @@ impl PointsInternal for PointsInternalService { .ok_or_else(|| Status::invalid_argument("UpdateVectors is missing"))?; update_vectors( - self.toc.clone(), + UncheckedTocProvider::new_unchecked(&self.toc), update_point_vectors, clock_tag.map(Into::into), shard_id, diff --git a/tests/openapi/test_strictmode.py b/tests/openapi/test_strictmode.py index c903da8115..a3974e7eac 100644 --- a/tests/openapi/test_strictmode.py +++ b/tests/openapi/test_strictmode.py @@ -1,4 +1,5 @@ import pytest +import random from .conftest import collection_name from .helpers.collection_setup import basic_collection_setup, drop_collection @@ -461,3 +462,144 @@ def test_strict_mode_update_many_upsert_max_batch_size(collection_name): assert "upsert" in search_fail.json()['status']['error'] assert not search_fail.ok + +def test_strict_mode_update_vectors_max_batch_size(collection_name): + def search_request(): + return request_with_validation( + api='/collections/{collection_name}/points/vectors', + method="PUT", + path_params={'collection_name': collection_name}, + body={ + "points": [ + { + "id": 1, + "vector": [1, 2, 3, 5], + }, + { + "id": 2, + "vector": [1, 2, 3, 5], + }, + { + "id": 3, + "vector": [1, 2, 3, 5], + }, + { + "id": 4, + "vector": [1, 2, 3, 5], + }, + ] + } + ) + + search_request().raise_for_status() + + set_strict_mode(collection_name, { + "enabled": True, + "upsert_max_batchsize": 4, + }) + + search_request().raise_for_status() + + set_strict_mode(collection_name, { + "enabled": True, + "upsert_max_batchsize": 3, + }) + + search_fail = search_request() + + assert "update limit" in search_fail.json()['status']['error'] + assert not search_fail.ok + + +def test_strict_mode_max_collection_size_upsert(collection_name): + basic_collection_setup(collection_name=collection_name) # Clear collection to not depend on other tests + + def upsert_points(ids: list[int]): + length = len(ids) + payloads = [{} for _ in range(length)] + vectors = [[1, 2, 3, 5] for _ in range(length)] + return request_with_validation( + api='/collections/{collection_name}/points', + method="PUT", + path_params={'collection_name': collection_name}, + body={ + "batch": { + "ids": ids, + "payloads": payloads, + "vectors": vectors + } + } + ) + + # Overwriting the same points to trigger cache refreshing + for _ in range(32): + upsert_points([1, 2, 3, 4, 5]).raise_for_status() + + set_strict_mode(collection_name, { + "enabled": True, + "max_collection_vector_size_bytes": 240, + }) + + for _ in range(32): + upsert_points([6, 7, 8, 9, 10]).raise_for_status() + + # Max limit has been reached and one of the next requests must fail. Due to cache it might not be the first call! + for _ in range(32): + failed_upsert = upsert_points([12, 13, 14, 15, 16]) + if failed_upsert.ok: + continue + assert "Max vector storage size" in failed_upsert.json()['status']['error'] + assert not failed_upsert.ok + return + + assert False, "Upserting should have failed but didn't" + + +def test_strict_mode_max_collection_size_upsert_batch(collection_name): + basic_collection_setup(collection_name=collection_name) # Clear collection to not depend on other tests + + def upsert_points(ids: list[int]): + length = len(ids) + payloads = [{} for _ in range(length)] + vectors = [[1, 2, 3, 5] for _ in range(length)] + return request_with_validation( + api='/collections/{collection_name}/points/batch', + method="POST", + path_params={'collection_name': collection_name}, + body={ + "operations": [ + { + "upsert": { + "batch": { + "ids": ids, + "payloads": payloads, + "vectors": vectors + } + } + } + ] + } + ) + + for _ in range(32): + upsert_points([1, 2, 3, 4, 5]).raise_for_status() + + set_strict_mode(collection_name, { + "enabled": True, + "max_collection_vector_size_bytes": 240, + }) + + for _ in range(32): + upsert_points([6, 7, 8, 9, 10]).raise_for_status() + + # Max limit has been reached and one of the next requests must fail. Due to cache it might not be the first call! + for _ in range(32): + failed_upsert = upsert_points([12, 13, 14, 15, 16]) + if failed_upsert.ok: + continue + assert "Max vector storage size" in failed_upsert.json()['status']['error'] + assert not failed_upsert.ok + return + + assert False, "Upserting should have failed but didn't" + diff --git a/tests/storage/aliases/data.json b/tests/storage/aliases/data.json new file mode 100644 index 0000000000..9e26dfeeb6 --- /dev/null +++ b/tests/storage/aliases/data.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/tests/storage/raft_state.json b/tests/storage/raft_state.json new file mode 100644 index 0000000000..365ca4f946 --- /dev/null +++ b/tests/storage/raft_state.json @@ -0,0 +1 @@ +{"state":{"hard_state":{"term":0,"vote":0,"commit":0},"conf_state":{"voters":[2947243872377711],"learners":[],"voters_outgoing":[],"learners_next":[],"auto_leave":false}},"latest_snapshot_meta":{"term":0,"index":0},"apply_progress_queue":null,"first_voter":2947243872377711,"peer_address_by_id":{},"peer_metadata_by_id":{},"this_peer_id":2947243872377711} \ No newline at end of file