use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use crate::common::budget::ResourceBudget;
use crate::common::progress_tracker::new_progress_tracker;
use crate::segment::common::operation_error::{OperationError, OperationResult};
use crate::segment::types::HnswGlobalConfig;
use crate::shard::files::SEGMENTS_PATH;
use crate::shard::optimizers::config::{
DEFAULT_DELETED_THRESHOLD, DEFAULT_VACUUM_MIN_VECTOR_NUMBER, LiveVectorNamesProvider,
TEMP_SEGMENTS_PATH,
};
use crate::shard::optimizers::config_mismatch_optimizer::ConfigMismatchOptimizer;
use crate::shard::optimizers::indexing_optimizer::IndexingOptimizer;
use crate::shard::optimizers::merge_optimizer::MergeOptimizer;
use crate::shard::optimizers::segment_optimizer::{
Optimizer, max_num_indexing_threads, plan_optimizations,
};
use crate::shard::optimizers::vacuum_optimizer::VacuumOptimizer;
use uuid::Uuid;
use crate::edge::EdgeShard;
impl EdgeShard {
pub fn optimize(&self) -> OperationResult<bool> {
let optimizers = self.build_blocking_optimizers();
let stopped = AtomicBool::new(false);
let mut optimized_any = false;
loop {
let planned = {
let segments = self.segments.read();
plan_optimizations(&segments, &optimizers)
};
if planned.is_empty() {
return Ok(optimized_any);
}
let mut optimized_in_iteration = false;
for (optimizer, segment_ids) in planned {
let num_indexing_threads = optimizer.num_indexing_threads();
let desired_io = num_indexing_threads;
let budget = ResourceBudget::new(num_indexing_threads, desired_io);
let permit = budget.try_acquire(0, desired_io).ok_or_else(|| {
OperationError::service_error(format!(
"failed to acquire resource permit for {} optimizer",
optimizer.name(),
))
})?;
let (_, progress) = new_progress_tracker();
let points_optimized = optimizer.as_ref().optimize(
self.segments.clone(),
segment_ids,
Uuid::new_v4(),
permit,
budget,
&stopped,
progress,
Box::new(|| ()),
)?;
if points_optimized > 0 {
optimized_in_iteration = true;
optimized_any = true;
}
}
self.update_segment_manifest()?;
if !optimized_in_iteration {
return Ok(optimized_any);
}
}
}
fn build_blocking_optimizers(&self) -> Vec<Arc<Optimizer>> {
let segments_path = self.path.join(SEGMENTS_PATH);
let temp_segments_path = self.path.join(TEMP_SEGMENTS_PATH);
let cfg = self.config();
let live_vector_names = {
let config = Arc::clone(&self.config);
LiveVectorNamesProvider::new(move || config.read().vector_names())
};
let segment_optimizer_config = cfg
.segment_optimizer_config()
.with_live_vector_names(live_vector_names);
let global_hnsw_config = cfg.hnsw_config();
let optimizers_config = cfg.optimizers();
let hnsw_global_config = HnswGlobalConfig::default();
let num_indexing_threads = max_num_indexing_threads(&segment_optimizer_config);
let threshold_config = cfg.optimizer_thresholds(num_indexing_threads);
let default_segments_number = optimizers_config.get_number_segments();
vec![
Arc::new(MergeOptimizer::new(
default_segments_number,
threshold_config,
segments_path.clone(),
temp_segments_path.clone(),
segment_optimizer_config.clone(),
hnsw_global_config.clone(),
)),
Arc::new(IndexingOptimizer::new(
default_segments_number,
threshold_config,
segments_path.clone(),
temp_segments_path.clone(),
segment_optimizer_config.clone(),
hnsw_global_config.clone(),
)),
Arc::new(VacuumOptimizer::new(
optimizers_config
.deleted_threshold
.unwrap_or(DEFAULT_DELETED_THRESHOLD),
optimizers_config
.vacuum_min_vector_number
.unwrap_or(DEFAULT_VACUUM_MIN_VECTOR_NUMBER),
threshold_config,
segments_path.clone(),
temp_segments_path.clone(),
segment_optimizer_config.clone(),
hnsw_global_config.clone(),
)),
Arc::new(ConfigMismatchOptimizer::new(
threshold_config,
segments_path,
temp_segments_path,
segment_optimizer_config,
global_hnsw_config,
hnsw_global_config,
)),
]
}
}
#[cfg(test)]
mod tests {
#![expect(clippy::wildcard_enum_match_arm, reason = "test code")]
use std::collections::HashMap;
use std::path::Path;
use fs_err as fs;
use crate::segment::data_types::vectors::{VectorInternal, VectorStructInternal};
use crate::segment::types::{Distance, ExtendedPointId, WithPayloadInterface, WithVector};
use crate::shard::operations::CollectionUpdateOperations::{PointOperation, VectorNameOperation};
use crate::shard::operations::VectorNameOperations;
use crate::shard::operations::point_ops::PointInsertOperationsInternal::PointsList;
use crate::shard::operations::point_ops::PointOperations::{DeletePoints, UpsertPoints};
use crate::shard::operations::point_ops::{PointStructPersisted, VectorStructPersisted};
use crate::shard::operations::vector_name_ops::DeleteVectorName;
use crate::shard::optimizers::config::default_segment_number;
use uuid::Uuid;
use crate::edge::config::vectors::EdgeVectorParams;
use crate::edge::{CountRequest, EdgeConfig, EdgeShard, RetrieveRequestBuilder};
const VECTOR_NAME: &str = "edge-test-vector";
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn exact_count_deduplicates_across_segments() {
let dir = tempfile::Builder::new()
.prefix("edge-exact-count-dedup")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
shard
.update(PointOperation(UpsertPoints(PointsList(vec![point(1)]))))
.unwrap();
drop(shard);
duplicate_single_segment(dir.path());
let reopened = EdgeShard::load(dir.path(), None).unwrap();
assert_eq!(reopened.info().unwrap().segments_count, 2);
let count = reopened.count(CountRequest::new()).unwrap();
assert_eq!(
count, 1,
"exact count should deduplicate point ids across segments"
);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn does_not_force_merge_all_segments_into_one() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-do-not-force-one")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
shard
.update(PointOperation(UpsertPoints(PointsList(vec![point(1)]))))
.unwrap();
drop(shard);
duplicate_single_segment(dir.path());
let reopened = EdgeShard::load(dir.path(), None).unwrap();
assert_eq!(reopened.info().unwrap().segments_count, 2);
let optimized = reopened.optimize().unwrap();
assert!(!optimized, "optimizer should not force-merge all segments");
assert_eq!(reopened.info().unwrap().segments_count, 2);
assert_points_retrievable_with_vectors(&reopened, &[1]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn vacuum_optimizer_runs_in_blocking_mode_until_idle() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-vacuum")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=1000).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=250).map(ExtendedPointId::NumId).collect::<Vec<_>>();
shard
.update(PointOperation(DeletePoints { ids: deleted_ids }))
.unwrap();
let optimized = shard.optimize().unwrap();
assert!(optimized, "vacuum candidate should be optimized");
let optimized_again = shard.optimize().unwrap();
assert!(
!optimized_again,
"second run should be idle after blocking optimization"
);
assert_points_retrievable_with_vectors(&shard, &[251, 500, 999, 1000]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn no_op_on_single_segment_without_deletions() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-noop-single")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=100).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let optimized = shard.optimize().unwrap();
assert!(!optimized, "single clean segment should not be optimized");
assert_eq!(shard.info().unwrap().points_count, 100);
assert_eq!(shard.info().unwrap().segments_count, 1);
assert_points_retrievable_with_vectors(&shard, &[1, 50, 100]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn no_op_on_empty_shard() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-noop-empty")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let optimized = shard.optimize().unwrap();
assert!(!optimized, "empty shard should not trigger optimization");
assert_eq!(shard.info().unwrap().points_count, 0);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn merge_reduces_excess_segments() {
let target_count = default_segment_number() + 6;
let dir = tempfile::Builder::new()
.prefix("edge-opt-merge-excess")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
shard
.update(PointOperation(UpsertPoints(PointsList(vec![point(1)]))))
.unwrap();
drop(shard);
multiply_segments(dir.path(), target_count);
let reopened = EdgeShard::load(dir.path(), None).unwrap();
reopened.optimize().unwrap();
let info = reopened.info().unwrap();
assert!(
info.segments_count <= default_segment_number() + 1,
"segments should be reduced after merge: got {} segments, \
expected at most {} (default_segment_number={}, +1 for appendable)",
info.segments_count,
default_segment_number() + 1,
default_segment_number(),
);
let count = reopened.count(CountRequest::new()).unwrap();
assert_eq!(count, 1, "merged shard should have exactly one point");
assert_points_retrievable_with_vectors(&reopened, &[1]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn optimization_is_idempotent_after_merge() {
let target_count = default_segment_number() + 6;
let dir = tempfile::Builder::new()
.prefix("edge-opt-merge-idempotent")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
shard
.update(PointOperation(UpsertPoints(PointsList(vec![point(1)]))))
.unwrap();
drop(shard);
multiply_segments(dir.path(), target_count);
let reopened = EdgeShard::load(dir.path(), None).unwrap();
reopened.optimize().unwrap();
let segments_after_first = reopened.info().unwrap().segments_count;
let optimized = reopened.optimize().unwrap();
assert!(
!optimized,
"second optimization run should be idle after merge"
);
assert_eq!(
reopened.info().unwrap().segments_count,
segments_after_first
);
assert_points_retrievable_with_vectors(&reopened, &[1]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn vacuum_below_threshold_is_noop() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-vacuum-below")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=1000).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=50).map(ExtendedPointId::NumId).collect::<Vec<_>>();
shard
.update(PointOperation(DeletePoints { ids: deleted_ids }))
.unwrap();
let optimized = shard.optimize().unwrap();
assert!(
!optimized,
"5% deletion should not trigger vacuum (threshold is 20%)"
);
assert_points_retrievable_with_vectors(&shard, &[51, 500, 1000]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn vacuum_below_min_vector_count_is_noop() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-vacuum-min-vecs")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=100).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=50).map(ExtendedPointId::NumId).collect::<Vec<_>>();
shard
.update(PointOperation(DeletePoints { ids: deleted_ids }))
.unwrap();
let optimized = shard.optimize().unwrap();
assert!(
!optimized,
"high deletion ratio with only 100 total points should not trigger vacuum \
(min_vectors_number=1000)"
);
assert_points_retrievable_with_vectors(&shard, &[51, 75, 100]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn vacuum_preserves_remaining_points() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-vacuum-data")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=1000).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=250).map(ExtendedPointId::NumId).collect::<Vec<_>>();
shard
.update(PointOperation(DeletePoints {
ids: deleted_ids.clone(),
}))
.unwrap();
let optimized = shard.optimize().unwrap();
assert!(optimized, "25% deletion should trigger vacuum");
let count = shard.count(CountRequest::new()).unwrap();
assert_eq!(count, 750, "should have 750 remaining points after vacuum");
let deleted_results = shard
.retrieve(
RetrieveRequestBuilder::new(deleted_ids.clone())
.with_payload(WithPayloadInterface::Bool(false))
.with_vector(WithVector::Bool(false))
.build(),
)
.unwrap();
assert!(
deleted_results.is_empty(),
"deleted points should not be retrievable"
);
assert_points_retrievable_with_vectors(&shard, &[251, 500, 750, 1000]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn vacuum_after_all_points_deleted() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-vacuum-all-deleted")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=1000).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=1000).map(ExtendedPointId::NumId).collect();
shard
.update(PointOperation(DeletePoints { ids: deleted_ids }))
.unwrap();
let _optimized = shard.optimize().unwrap();
let count = shard.count(CountRequest::new()).unwrap();
assert_eq!(count, 0, "all points should be gone after vacuum");
shard
.update(PointOperation(UpsertPoints(PointsList(vec![point(9999)]))))
.unwrap();
let count = shard.count(CountRequest::new()).unwrap();
assert_eq!(count, 1, "shard should accept new points after full vacuum");
assert_points_retrievable_with_vectors(&shard, &[9999]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn vacuum_at_exact_threshold_boundary_is_noop() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-vacuum-boundary")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=1000).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=200).map(ExtendedPointId::NumId).collect();
shard
.update(PointOperation(DeletePoints { ids: deleted_ids }))
.unwrap();
let optimized = shard.optimize().unwrap();
assert!(
!optimized,
"exactly 20% deletion (not strictly greater) should not trigger vacuum"
);
assert_points_retrievable_with_vectors(&shard, &[201, 500, 1000]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn vacuum_just_above_threshold_triggers() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-vacuum-above")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=1000).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=201).map(ExtendedPointId::NumId).collect();
shard
.update(PointOperation(DeletePoints { ids: deleted_ids }))
.unwrap();
let optimized = shard.optimize().unwrap();
assert!(
optimized,
"20.1% deletion should trigger vacuum (threshold is >20%)"
);
assert_points_retrievable_with_vectors(&shard, &[202, 500, 1000]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn merge_and_vacuum_cooperate() {
let target_count = default_segment_number() + 6;
let dir = tempfile::Builder::new()
.prefix("edge-opt-merge-vacuum")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=1000).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=250).map(ExtendedPointId::NumId).collect();
shard
.update(PointOperation(DeletePoints { ids: deleted_ids }))
.unwrap();
drop(shard);
multiply_segments(dir.path(), target_count);
let reopened = EdgeShard::load(dir.path(), None).unwrap();
reopened.optimize().unwrap();
let info = reopened.info().unwrap();
assert!(
info.segments_count <= default_segment_number() + 1,
"excess segments should be merged: got {}",
info.segments_count,
);
let count = reopened.count(CountRequest::new()).unwrap();
assert_eq!(
count, 750,
"merged shard should preserve surviving points without double-counting"
);
assert_points_retrievable_with_vectors(&reopened, &[251, 500, 1000]);
let optimized = reopened.optimize().unwrap();
assert!(!optimized, "second run should be idle after merge+vacuum");
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn data_survives_optimize_and_reload() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-reload")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
let points = (1..=1000).map(point).collect::<Vec<_>>();
shard
.update(PointOperation(UpsertPoints(PointsList(points))))
.unwrap();
let deleted_ids = (1..=250).map(ExtendedPointId::NumId).collect();
shard
.update(PointOperation(DeletePoints { ids: deleted_ids }))
.unwrap();
let optimized = shard.optimize().unwrap();
assert!(optimized);
drop(shard);
let reopened = EdgeShard::load(dir.path(), None).unwrap();
let count = reopened.count(CountRequest::new()).unwrap();
assert_eq!(count, 750, "point count should be preserved across reload");
assert_points_retrievable_with_vectors(&reopened, &[251, 500, 750, 1000]);
}
#[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
#[test]
fn optimize_prunes_vector_deleted_from_config() {
let target_count = default_segment_number() + 6;
let dir = tempfile::Builder::new()
.prefix("edge-opt-prune-deleted-vector")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), two_vector_config()).unwrap();
shard
.update(PointOperation(UpsertPoints(PointsList(vec![
two_vector_point(1),
]))))
.unwrap();
drop(shard);
multiply_segments(dir.path(), target_count);
let reopened = EdgeShard::load(dir.path(), None).unwrap();
reopened
.config
.write(|cfg| {
cfg.vectors.remove(DROP_VECTOR_NAME);
})
.unwrap();
let optimized = reopened
.optimize()
.expect("optimize must prune the deleted vector data, not cancel the merge");
assert!(optimized, "excess segments should have been merged");
let info = reopened.info().unwrap();
assert!(
info.segments_count <= default_segment_number() + 1,
"segments should be reduced after merge: got {} segments",
info.segments_count,
);
let results = reopened
.retrieve(
RetrieveRequestBuilder::new(vec![ExtendedPointId::NumId(1)])
.with_payload(WithPayloadInterface::Bool(false))
.with_vector(WithVector::Bool(true))
.build(),
)
.unwrap();
assert_eq!(results.len(), 1, "point should survive the merge");
let vectors = match results[0]
.vector
.as_ref()
.expect("vector should be present")
{
VectorStructInternal::Named(named) => named,
other => panic!("expected Named vectors, got {other:?}"),
};
let keep = match vectors
.get(KEEP_VECTOR_NAME)
.expect("surviving vector should be present")
{
VectorInternal::Dense(v) => v,
other => panic!("expected Dense vector, got {other:?}"),
};
assert_eq!(keep, &vec![1.0], "surviving vector value mismatch");
}
#[test]
fn optimizers_read_live_vector_names() {
let dir = tempfile::Builder::new()
.prefix("edge-opt-live-vector-names")
.tempdir()
.unwrap();
let shard = EdgeShard::new(dir.path(), two_vector_config()).unwrap();
let optimizers = shard.build_blocking_optimizers();
for optimizer in &optimizers {
let names = optimizer
.segment_optimizer_config()
.live_vector_names()
.expect("live vector names must be wired into edge optimizers");
assert!(
names.contains(KEEP_VECTOR_NAME) && names.contains(DROP_VECTOR_NAME),
"both vector names should be live before the delete: {names:?}",
);
}
shard
.update(VectorNameOperation(VectorNameOperations::DeleteVectorName(
DeleteVectorName {
vector_name: DROP_VECTOR_NAME.to_string(),
},
)))
.unwrap();
for optimizer in &optimizers {
let names = optimizer
.segment_optimizer_config()
.live_vector_names()
.expect("live vector names must be wired into edge optimizers");
assert!(names.contains(KEEP_VECTOR_NAME));
assert!(
!names.contains(DROP_VECTOR_NAME),
"a deletion after optimizer build must be visible to the live provider",
);
}
}
fn assert_points_retrievable_with_vectors(shard: &EdgeShard, ids: &[u64]) {
let point_ids = ids
.iter()
.map(|id| ExtendedPointId::NumId(*id))
.collect::<Vec<_>>();
let results = shard
.retrieve(
RetrieveRequestBuilder::new(point_ids)
.with_payload(WithPayloadInterface::Bool(false))
.with_vector(WithVector::Bool(true))
.build(),
)
.unwrap();
assert_eq!(
results.len(),
ids.len(),
"expected {} retrievable points, got {}",
ids.len(),
results.len(),
);
for (result, &expected_id) in results.iter().zip(ids) {
assert_eq!(result.id, ExtendedPointId::NumId(expected_id));
let vectors = match result.vector.as_ref().expect("vector should be present") {
VectorStructInternal::Named(named) => named,
other => panic!("expected Named vectors, got {other:?}"),
};
let vec = match vectors.get(VECTOR_NAME).expect("vector name should exist") {
VectorInternal::Dense(v) => v,
other => panic!("expected Dense vector, got {other:?}"),
};
assert_eq!(
vec,
&vec![expected_id as f32],
"vector value mismatch for point {expected_id}"
);
}
}
fn test_config() -> EdgeConfig {
EdgeConfig {
on_disk_payload: Some(false),
vectors: HashMap::from([(
VECTOR_NAME.to_string(),
EdgeVectorParams {
size: 1,
distance: Distance::Dot,
quantization_config: None,
multivector_config: None,
datatype: None,
on_disk: None,
hnsw_config: None,
},
)]),
sparse_vectors: HashMap::new(),
hnsw_config: None,
quantization_config: None,
optimizers: None,
wal_options: None,
max_search_threads: None,
search_pool_core: None,
}
}
const KEEP_VECTOR_NAME: &str = "edge-keep-vector";
const DROP_VECTOR_NAME: &str = "edge-drop-vector";
fn two_vector_config() -> EdgeConfig {
let params = EdgeVectorParams {
size: 1,
distance: Distance::Dot,
quantization_config: None,
multivector_config: None,
datatype: None,
on_disk: None,
hnsw_config: None,
};
EdgeConfig {
vectors: HashMap::from([
(KEEP_VECTOR_NAME.to_string(), params.clone()),
(DROP_VECTOR_NAME.to_string(), params),
]),
..test_config()
}
}
fn two_vector_point(id: u64) -> PointStructPersisted {
PointStructPersisted {
id: ExtendedPointId::NumId(id),
vector: VectorStructPersisted::from(VectorStructInternal::Named(HashMap::from([
(
KEEP_VECTOR_NAME.to_string(),
VectorInternal::from(vec![id as f32]),
),
(
DROP_VECTOR_NAME.to_string(),
VectorInternal::from(vec![-(id as f32)]),
),
]))),
payload: None,
}
}
fn point(id: u64) -> PointStructPersisted {
PointStructPersisted {
id: ExtendedPointId::NumId(id),
vector: VectorStructPersisted::from(VectorStructInternal::Named(HashMap::from([(
VECTOR_NAME.to_string(),
VectorInternal::from(vec![id as f32]),
)]))),
payload: None,
}
}
fn multiply_segments(shard_dir: &Path, target_count: usize) {
let segments_path = shard_dir.join("segments");
let segment_dirs = fs::read_dir(&segments_path)
.unwrap()
.filter_map(Result::ok)
.map(|entry| entry.path())
.filter(|path| path.is_dir())
.collect::<Vec<_>>();
assert!(!segment_dirs.is_empty(), "need at least one source segment");
let source = &segment_dirs[0];
let current_count = segment_dirs.len();
for _ in current_count..target_count {
let target = segments_path.join(Uuid::new_v4().to_string());
copy_dir_recursive(source, &target);
}
}
fn duplicate_single_segment(shard_dir: &Path) {
let segments_path = shard_dir.join("segments");
let segment_dirs = fs::read_dir(&segments_path)
.unwrap()
.filter_map(Result::ok)
.map(|entry| entry.path())
.filter(|path| path.is_dir())
.collect::<Vec<_>>();
assert_eq!(segment_dirs.len(), 1, "expected exactly one source segment");
let source = &segment_dirs[0];
let target = segments_path.join(Uuid::new_v4().to_string());
copy_dir_recursive(source, &target);
}
fn copy_dir_recursive(from: &Path, to: &Path) {
fs::create_dir_all(to).unwrap();
for entry in fs::read_dir(from).unwrap().filter_map(Result::ok) {
let from_path = entry.path();
let to_path = to.join(entry.file_name());
if entry.file_type().unwrap().is_dir() {
copy_dir_recursive(&from_path, &to_path);
} else {
fs::copy(&from_path, &to_path).unwrap();
}
}
}
}