upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-07-30 22:51:00 -05:00
Code Issues Packages Projects Releases Wiki Activity

Strict mode max collection vector size (#5501)

Browse Source 
* Strict mode config: Max collection size

* api specs

* Add tests + set/update payload check

* Improve function names and add comments

* rename config to separate vectors and payload

* fix tests

* Adjust configs docs

* add benchmark

* improve performance by caching shard info

* add bench for size_info() and fix tests

* Also limit the batch-size for vector updates (#5508)

* Also limit the batch-size for vector updates

* clippy

* add lost commit

* Load cache on collection initialization

* add unit type to parameter name

* fix renaming in test

* clearer error message

* fix test

* review remarks

* remove unused function for now

---------

Co-authored-by: Arnaud Gourlay <arnaud.gourlay@gmail.com>
This commit is contained in:
Jojii
2024-12-05 17:28:09 +01:00
committed by GitHub
parent 3b8a4180f8
commit 42fd2e27e2
32 changed files with 629 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
docs/grpc/docs.md
View File

@@ -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 | |

7
docs/redoc/master/openapi.json
View File

@@ -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
}
}
},

6
lib/api/src/grpc/conversions.rs
View File

@@ -1632,6 +1632,9 @@ impl From<StrictModeConfig> 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<segment::types::StrictModeConfig> 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),
}
}
}

1
lib/api/src/grpc/proto/collections.proto
View File

@@ -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 {

2
lib/api/src/grpc/qdrant.rs
View File

@@ -454,6 +454,8 @@ pub struct StrictModeConfig {
pub search_max_oversampling: ::core::option::Option<f32>,
#[prost(uint64, optional, tag = "9")]
pub upsert_max_batchsize: ::core::option::Option<u64>,
#[prost(uint64, optional, tag = "10")]
pub max_collection_vector_size_bytes: ::core::option::Option<u64>,
}
#[derive(validator::Validate)]
#[derive(serde::Serialize)]

14
lib/api/src/rest/schema.rs
View File

@@ -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
}
}

38
lib/collection/src/collection/mod.rs
View File

@@ -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<dyn Fn(ShardTransfer) + Send + Sync>;
@@ -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<RwLock<ShardHolder>>) -> 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<LocalDataStats> {
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;

7
lib/collection/src/collection_manager/holders/proxy_segment.rs
View File

@@ -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();

65
lib/collection/src/common/local_data_stats.rs Normal file
View File

@@ -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;
}
}

1
lib/collection/src/common/mod.rs
View File

@@ -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;

5
lib/collection/src/operations/verification/discovery.rs
View File

@@ -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(())

24
lib/collection/src/operations/verification/mod.rs
View File

@@ -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<T: PartialOrd + Display>(
}
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

10
lib/collection/src/operations/verification/query.rs
View File

@@ -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?;
}
}

5
lib/collection/src/operations/verification/search.rs
View File

@@ -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(())
}

81
lib/collection/src/operations/verification/update.rs
View File

@@ -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<usize> {
None
}
fn indexed_filter_read(&self) -> Option<&Filter> {
None
}
fn indexed_filter_write(&self) -> Option<&Filter> {
None
}
fn request_exact(&self) -> Option<bool> {
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(())
}

30
lib/collection/src/shards/replica_set/mod.rs
View File

@@ -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

10
lib/collection/src/shards/shard_holder/mod.rs
View File

@@ -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)]

23
lib/segment/Cargo.toml
View File

@@ -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

76
lib/segment/benches/segment_info.rs Normal file
View File

@@ -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<String, Value> = 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);

4
lib/segment/src/entry/entry_point.rs
View File

@@ -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;

31
lib/segment/src/segment/entry.rs
View File

@@ -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
}

6
lib/segment/src/types.rs
View File

@@ -707,6 +707,10 @@ pub struct StrictModeConfig {
/// Max batchsize when upserting
#[serde(skip_serializing_if = "Option::is_none")]
pub upsert_max_batchsize: Option<usize>,
/// Max size of a collections vector storage in bytes
#[serde(skip_serializing_if = "Option::is_none")]
pub max_collection_vector_size_bytes: Option<usize>,
}
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);
}

4
lib/storage/src/content_manager/collection_verification.rs
View File

@@ -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 {

3
lib/storage/src/content_manager/conversions.rs
View File

@@ -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),
}
}

10
src/actix/api/update_api.rs
View File

@@ -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<UpdateParam>,
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,

72
src/common/points.rs
View File

@@ -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
}
}
}
}

9
src/tonic/api/points_api.rs
View File

@@ -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<Response<PointsOperationResponse>, 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,

26
src/tonic/api/points_common.rs
View File

@@ -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<TableOfContent>,
toc_provider: impl CheckedTocProvider,
upsert_points: UpsertPoints,
clock_tag: Option<ClockTag>,
shard_selection: Option<ShardId>,
@@ -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<TableOfContent>,
toc_provider: impl CheckedTocProvider,
update_point_vectors: UpdatePointVectors,
clock_tag: Option<ClockTag>,
shard_selection: Option<ShardId>,
@@ -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,

4
src/tonic/api/points_internal_api.rs
View File

@@ -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,

142
tests/openapi/test_strictmode.py
View File

@@ -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"

1
tests/storage/aliases/data.json Normal file
View File

@@ -0,0 +1 @@
{}

1
tests/storage/raft_state.json Normal file
View File

@@ -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}