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::optimizers::config::{
DEFAULT_DELETED_THRESHOLD, DEFAULT_VACUUM_MIN_VECTOR_NUMBER, 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, SEGMENTS_PATH};
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;
}
}
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 segment_optimizer_config = cfg.segment_optimizer_config();
let global_hnsw_config = cfg.hnsw_config;
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 = cfg.optimizers.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(
cfg.optimizers
.deleted_threshold
.unwrap_or(DEFAULT_DELETED_THRESHOLD),
cfg.optimizers
.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::count::CountRequestInternal;
use crate::shard::operations::CollectionUpdateOperations::PointOperation;
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::optimizers::config::default_segment_number;
use uuid::Uuid;
use crate::edge::config::vectors::EdgeVectorParams;
use crate::edge::{EdgeConfig, EdgeShard};
const VECTOR_NAME: &str = "edge-test-vector";
#[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().segments_count, 2);
let optimized = reopened.optimize().unwrap();
assert!(!optimized, "optimizer should not force-merge all segments");
assert_eq!(reopened.info().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().points_count, 100);
assert_eq!(shard.info().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().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();
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(CountRequestInternal {
filter: None,
exact: true,
})
.unwrap();
assert!(count >= 1, "shard should still have data after merge");
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().segments_count;
let optimized = reopened.optimize().unwrap();
assert!(
!optimized,
"second optimization run should be idle after merge"
);
assert_eq!(reopened.info().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(CountRequestInternal {
filter: None,
exact: true,
})
.unwrap();
assert_eq!(count, 750, "should have 750 remaining points after vacuum");
let deleted_results = shard
.retrieve(
&deleted_ids,
Some(WithPayloadInterface::Bool(false)),
Some(WithVector::Bool(false)),
)
.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(CountRequestInternal {
filter: None,
exact: true,
})
.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(CountRequestInternal {
filter: None,
exact: true,
})
.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();
assert!(
info.segments_count <= default_segment_number() + 1,
"excess segments should be merged: got {}",
info.segments_count,
);
let count = reopened
.count(CountRequestInternal {
filter: None,
exact: true,
})
.unwrap();
assert!(
count >= 750,
"merged shard should preserve surviving points"
);
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(CountRequestInternal {
filter: None,
exact: true,
})
.unwrap();
assert_eq!(count, 750, "point count should be preserved across reload");
assert_points_retrievable_with_vectors(&reopened, &[251, 500, 750, 1000]);
}
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(
&point_ids,
Some(WithPayloadInterface::Bool(false)),
Some(WithVector::Bool(true)),
)
.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: 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: Default::default(),
quantization_config: None,
optimizers: Default::default(),
wal_options: 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();
}
}
}
}