1use std::sync::Arc;
2use std::sync::atomic::AtomicBool;
3
4use crate::common::budget::ResourceBudget;
5use crate::common::progress_tracker::new_progress_tracker;
6use crate::segment::common::operation_error::{OperationError, OperationResult};
7use crate::segment::types::HnswGlobalConfig;
8use crate::shard::files::SEGMENTS_PATH;
9use crate::shard::optimizers::config::{
10 DEFAULT_DELETED_THRESHOLD, DEFAULT_VACUUM_MIN_VECTOR_NUMBER, LiveVectorNamesProvider,
11 TEMP_SEGMENTS_PATH,
12};
13use crate::shard::optimizers::config_mismatch_optimizer::ConfigMismatchOptimizer;
14use crate::shard::optimizers::indexing_optimizer::IndexingOptimizer;
15use crate::shard::optimizers::merge_optimizer::MergeOptimizer;
16use crate::shard::optimizers::segment_optimizer::{
17 Optimizer, max_num_indexing_threads, plan_optimizations,
18};
19use crate::shard::optimizers::vacuum_optimizer::VacuumOptimizer;
20use uuid::Uuid;
21
22use crate::edge::EdgeShard;
23
24impl EdgeShard {
25 pub fn optimize(&self) -> OperationResult<bool> {
29 let optimizers = self.build_blocking_optimizers();
30 let stopped = AtomicBool::new(false);
31 let mut optimized_any = false;
32
33 loop {
34 let planned = {
35 let segments = self.segments.read();
36 plan_optimizations(&segments, &optimizers)
37 };
38
39 if planned.is_empty() {
40 return Ok(optimized_any);
41 }
42
43 let mut optimized_in_iteration = false;
44
45 for (optimizer, segment_ids) in planned {
46 let num_indexing_threads = optimizer.num_indexing_threads();
47 let desired_io = num_indexing_threads;
48 let budget = ResourceBudget::new(num_indexing_threads, desired_io);
50 let permit = budget.try_acquire(0, desired_io).ok_or_else(|| {
51 OperationError::service_error(format!(
52 "failed to acquire resource permit for {} optimizer",
53 optimizer.name(),
54 ))
55 })?;
56
57 let (_, progress) = new_progress_tracker();
58 let points_optimized = optimizer.as_ref().optimize(
59 self.segments.clone(),
60 segment_ids,
61 Uuid::new_v4(),
62 permit,
63 budget,
64 &stopped,
65 progress,
66 Box::new(|| ()),
67 )?;
68
69 if points_optimized > 0 {
70 optimized_in_iteration = true;
71 optimized_any = true;
72 }
73 }
74
75 self.update_segment_manifest()?;
77
78 if !optimized_in_iteration {
80 return Ok(optimized_any);
81 }
82 }
83 }
84
85 fn build_blocking_optimizers(&self) -> Vec<Arc<Optimizer>> {
86 let segments_path = self.path.join(SEGMENTS_PATH);
87 let temp_segments_path = self.path.join(TEMP_SEGMENTS_PATH);
88
89 let cfg = self.config();
90 let live_vector_names = {
98 let config = Arc::clone(&self.config);
99 LiveVectorNamesProvider::new(move || config.read().vector_names())
100 };
101 let segment_optimizer_config = cfg
102 .segment_optimizer_config()
103 .with_live_vector_names(live_vector_names);
104 let global_hnsw_config = cfg.hnsw_config();
105 let optimizers_config = cfg.optimizers();
106 let hnsw_global_config = HnswGlobalConfig::default();
107 let num_indexing_threads = max_num_indexing_threads(&segment_optimizer_config);
108 let threshold_config = cfg.optimizer_thresholds(num_indexing_threads);
109 let default_segments_number = optimizers_config.get_number_segments();
110
111 vec![
112 Arc::new(MergeOptimizer::new(
113 default_segments_number,
114 threshold_config,
115 segments_path.clone(),
116 temp_segments_path.clone(),
117 segment_optimizer_config.clone(),
118 hnsw_global_config.clone(),
119 )),
120 Arc::new(IndexingOptimizer::new(
121 default_segments_number,
122 threshold_config,
123 segments_path.clone(),
124 temp_segments_path.clone(),
125 segment_optimizer_config.clone(),
126 hnsw_global_config.clone(),
127 )),
128 Arc::new(VacuumOptimizer::new(
129 optimizers_config
130 .deleted_threshold
131 .unwrap_or(DEFAULT_DELETED_THRESHOLD),
132 optimizers_config
133 .vacuum_min_vector_number
134 .unwrap_or(DEFAULT_VACUUM_MIN_VECTOR_NUMBER),
135 threshold_config,
136 segments_path.clone(),
137 temp_segments_path.clone(),
138 segment_optimizer_config.clone(),
139 hnsw_global_config.clone(),
140 )),
141 Arc::new(ConfigMismatchOptimizer::new(
142 threshold_config,
143 segments_path,
144 temp_segments_path,
145 segment_optimizer_config,
146 global_hnsw_config,
147 hnsw_global_config,
148 )),
149 ]
150 }
151}
152
153#[cfg(test)]
158mod tests {
159 #![expect(clippy::wildcard_enum_match_arm, reason = "test code")]
160
161 use std::collections::HashMap;
162 use std::path::Path;
163
164 use fs_err as fs;
165 use crate::segment::data_types::vectors::{VectorInternal, VectorStructInternal};
166 use crate::segment::types::{Distance, ExtendedPointId, WithPayloadInterface, WithVector};
167 use crate::shard::operations::CollectionUpdateOperations::{PointOperation, VectorNameOperation};
168 use crate::shard::operations::VectorNameOperations;
169 use crate::shard::operations::point_ops::PointInsertOperationsInternal::PointsList;
170 use crate::shard::operations::point_ops::PointOperations::{DeletePoints, UpsertPoints};
171 use crate::shard::operations::point_ops::{PointStructPersisted, VectorStructPersisted};
172 use crate::shard::operations::vector_name_ops::DeleteVectorName;
173 use crate::shard::optimizers::config::default_segment_number;
174 use uuid::Uuid;
175
176 use crate::edge::config::vectors::EdgeVectorParams;
177 use crate::edge::{CountRequest, EdgeConfig, EdgeShard, RetrieveRequestBuilder};
178
179 const VECTOR_NAME: &str = "edge-test-vector";
180
181 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
182 #[test]
183 fn exact_count_deduplicates_across_segments() {
184 let dir = tempfile::Builder::new()
185 .prefix("edge-exact-count-dedup")
186 .tempdir()
187 .unwrap();
188
189 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
190 shard
191 .update(PointOperation(UpsertPoints(PointsList(vec![point(1)]))))
192 .unwrap();
193 drop(shard);
194
195 duplicate_single_segment(dir.path());
196
197 let reopened = EdgeShard::load(dir.path(), None).unwrap();
198 assert_eq!(reopened.info().unwrap().segments_count, 2);
199
200 let count = reopened.count(CountRequest::new()).unwrap();
201 assert_eq!(
202 count, 1,
203 "exact count should deduplicate point ids across segments"
204 );
205 }
206
207 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
208 #[test]
209 fn does_not_force_merge_all_segments_into_one() {
210 let dir = tempfile::Builder::new()
211 .prefix("edge-opt-do-not-force-one")
212 .tempdir()
213 .unwrap();
214
215 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
216 shard
217 .update(PointOperation(UpsertPoints(PointsList(vec![point(1)]))))
218 .unwrap();
219 drop(shard);
220
221 duplicate_single_segment(dir.path());
222
223 let reopened = EdgeShard::load(dir.path(), None).unwrap();
224 assert_eq!(reopened.info().unwrap().segments_count, 2);
225
226 let optimized = reopened.optimize().unwrap();
227 assert!(!optimized, "optimizer should not force-merge all segments");
228 assert_eq!(reopened.info().unwrap().segments_count, 2);
229
230 assert_points_retrievable_with_vectors(&reopened, &[1]);
231 }
232
233 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
234 #[test]
235 fn vacuum_optimizer_runs_in_blocking_mode_until_idle() {
236 let dir = tempfile::Builder::new()
237 .prefix("edge-opt-vacuum")
238 .tempdir()
239 .unwrap();
240
241 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
242
243 let points = (1..=1000).map(point).collect::<Vec<_>>();
244 shard
245 .update(PointOperation(UpsertPoints(PointsList(points))))
246 .unwrap();
247
248 let deleted_ids = (1..=250).map(ExtendedPointId::NumId).collect::<Vec<_>>();
250 shard
251 .update(PointOperation(DeletePoints { ids: deleted_ids }))
252 .unwrap();
253
254 let optimized = shard.optimize().unwrap();
255 assert!(optimized, "vacuum candidate should be optimized");
256
257 let optimized_again = shard.optimize().unwrap();
258 assert!(
259 !optimized_again,
260 "second run should be idle after blocking optimization"
261 );
262
263 assert_points_retrievable_with_vectors(&shard, &[251, 500, 999, 1000]);
265 }
266
267 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
270 #[test]
271 fn no_op_on_single_segment_without_deletions() {
272 let dir = tempfile::Builder::new()
273 .prefix("edge-opt-noop-single")
274 .tempdir()
275 .unwrap();
276
277 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
278
279 let points = (1..=100).map(point).collect::<Vec<_>>();
280 shard
281 .update(PointOperation(UpsertPoints(PointsList(points))))
282 .unwrap();
283
284 let optimized = shard.optimize().unwrap();
285 assert!(!optimized, "single clean segment should not be optimized");
286 assert_eq!(shard.info().unwrap().points_count, 100);
287 assert_eq!(shard.info().unwrap().segments_count, 1);
288
289 assert_points_retrievable_with_vectors(&shard, &[1, 50, 100]);
290 }
291
292 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
294 #[test]
295 fn no_op_on_empty_shard() {
296 let dir = tempfile::Builder::new()
297 .prefix("edge-opt-noop-empty")
298 .tempdir()
299 .unwrap();
300
301 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
302
303 let optimized = shard.optimize().unwrap();
304 assert!(!optimized, "empty shard should not trigger optimization");
305 assert_eq!(shard.info().unwrap().points_count, 0);
306 }
307
308 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
311 #[test]
312 fn merge_reduces_excess_segments() {
313 let target_count = default_segment_number() + 6;
314
315 let dir = tempfile::Builder::new()
316 .prefix("edge-opt-merge-excess")
317 .tempdir()
318 .unwrap();
319
320 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
321 shard
322 .update(PointOperation(UpsertPoints(PointsList(vec![point(1)]))))
323 .unwrap();
324 drop(shard);
325
326 multiply_segments(dir.path(), target_count);
327
328 let reopened = EdgeShard::load(dir.path(), None).unwrap();
329 reopened.optimize().unwrap();
330 let info = reopened.info().unwrap();
331 assert!(
332 info.segments_count <= default_segment_number() + 1,
333 "segments should be reduced after merge: got {} segments, \
334 expected at most {} (default_segment_number={}, +1 for appendable)",
335 info.segments_count,
336 default_segment_number() + 1,
337 default_segment_number(),
338 );
339
340 let count = reopened.count(CountRequest::new()).unwrap();
343 assert_eq!(count, 1, "merged shard should have exactly one point");
344
345 assert_points_retrievable_with_vectors(&reopened, &[1]);
346 }
347
348 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
350 #[test]
351 fn optimization_is_idempotent_after_merge() {
352 let target_count = default_segment_number() + 6;
353
354 let dir = tempfile::Builder::new()
355 .prefix("edge-opt-merge-idempotent")
356 .tempdir()
357 .unwrap();
358
359 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
360 shard
361 .update(PointOperation(UpsertPoints(PointsList(vec![point(1)]))))
362 .unwrap();
363 drop(shard);
364
365 multiply_segments(dir.path(), target_count);
366
367 let reopened = EdgeShard::load(dir.path(), None).unwrap();
368 reopened.optimize().unwrap();
370 let segments_after_first = reopened.info().unwrap().segments_count;
371
372 let optimized = reopened.optimize().unwrap();
374 assert!(
375 !optimized,
376 "second optimization run should be idle after merge"
377 );
378 assert_eq!(
379 reopened.info().unwrap().segments_count,
380 segments_after_first
381 );
382
383 assert_points_retrievable_with_vectors(&reopened, &[1]);
384 }
385
386 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
389 #[test]
390 fn vacuum_below_threshold_is_noop() {
391 let dir = tempfile::Builder::new()
392 .prefix("edge-opt-vacuum-below")
393 .tempdir()
394 .unwrap();
395
396 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
397
398 let points = (1..=1000).map(point).collect::<Vec<_>>();
399 shard
400 .update(PointOperation(UpsertPoints(PointsList(points))))
401 .unwrap();
402
403 let deleted_ids = (1..=50).map(ExtendedPointId::NumId).collect::<Vec<_>>();
405 shard
406 .update(PointOperation(DeletePoints { ids: deleted_ids }))
407 .unwrap();
408
409 let optimized = shard.optimize().unwrap();
410 assert!(
411 !optimized,
412 "5% deletion should not trigger vacuum (threshold is 20%)"
413 );
414
415 assert_points_retrievable_with_vectors(&shard, &[51, 500, 1000]);
417 }
418
419 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
422 #[test]
423 fn vacuum_below_min_vector_count_is_noop() {
424 let dir = tempfile::Builder::new()
425 .prefix("edge-opt-vacuum-min-vecs")
426 .tempdir()
427 .unwrap();
428
429 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
430
431 let points = (1..=100).map(point).collect::<Vec<_>>();
433 shard
434 .update(PointOperation(UpsertPoints(PointsList(points))))
435 .unwrap();
436
437 let deleted_ids = (1..=50).map(ExtendedPointId::NumId).collect::<Vec<_>>();
439 shard
440 .update(PointOperation(DeletePoints { ids: deleted_ids }))
441 .unwrap();
442
443 let optimized = shard.optimize().unwrap();
444 assert!(
445 !optimized,
446 "high deletion ratio with only 100 total points should not trigger vacuum \
447 (min_vectors_number=1000)"
448 );
449
450 assert_points_retrievable_with_vectors(&shard, &[51, 75, 100]);
451 }
452
453 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
456 #[test]
457 fn vacuum_preserves_remaining_points() {
458 let dir = tempfile::Builder::new()
459 .prefix("edge-opt-vacuum-data")
460 .tempdir()
461 .unwrap();
462
463 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
464
465 let points = (1..=1000).map(point).collect::<Vec<_>>();
466 shard
467 .update(PointOperation(UpsertPoints(PointsList(points))))
468 .unwrap();
469
470 let deleted_ids = (1..=250).map(ExtendedPointId::NumId).collect::<Vec<_>>();
472 shard
473 .update(PointOperation(DeletePoints {
474 ids: deleted_ids.clone(),
475 }))
476 .unwrap();
477
478 let optimized = shard.optimize().unwrap();
479 assert!(optimized, "25% deletion should trigger vacuum");
480
481 let count = shard.count(CountRequest::new()).unwrap();
483 assert_eq!(count, 750, "should have 750 remaining points after vacuum");
484
485 let deleted_results = shard
487 .retrieve(
488 RetrieveRequestBuilder::new(deleted_ids.clone())
489 .with_payload(WithPayloadInterface::Bool(false))
490 .with_vector(WithVector::Bool(false))
491 .build(),
492 )
493 .unwrap();
494 assert!(
495 deleted_results.is_empty(),
496 "deleted points should not be retrievable"
497 );
498
499 assert_points_retrievable_with_vectors(&shard, &[251, 500, 750, 1000]);
501 }
502
503 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
509 #[test]
510 fn vacuum_after_all_points_deleted() {
511 let dir = tempfile::Builder::new()
512 .prefix("edge-opt-vacuum-all-deleted")
513 .tempdir()
514 .unwrap();
515
516 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
517
518 let points = (1..=1000).map(point).collect::<Vec<_>>();
519 shard
520 .update(PointOperation(UpsertPoints(PointsList(points))))
521 .unwrap();
522
523 let deleted_ids = (1..=1000).map(ExtendedPointId::NumId).collect();
525 shard
526 .update(PointOperation(DeletePoints { ids: deleted_ids }))
527 .unwrap();
528
529 let _optimized = shard.optimize().unwrap();
532
533 let count = shard.count(CountRequest::new()).unwrap();
534 assert_eq!(count, 0, "all points should be gone after vacuum");
535
536 shard
538 .update(PointOperation(UpsertPoints(PointsList(vec![point(9999)]))))
539 .unwrap();
540 let count = shard.count(CountRequest::new()).unwrap();
541 assert_eq!(count, 1, "shard should accept new points after full vacuum");
542
543 assert_points_retrievable_with_vectors(&shard, &[9999]);
544 }
545
546 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
550 #[test]
551 fn vacuum_at_exact_threshold_boundary_is_noop() {
552 let dir = tempfile::Builder::new()
553 .prefix("edge-opt-vacuum-boundary")
554 .tempdir()
555 .unwrap();
556
557 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
558
559 let points = (1..=1000).map(point).collect::<Vec<_>>();
560 shard
561 .update(PointOperation(UpsertPoints(PointsList(points))))
562 .unwrap();
563
564 let deleted_ids = (1..=200).map(ExtendedPointId::NumId).collect();
566 shard
567 .update(PointOperation(DeletePoints { ids: deleted_ids }))
568 .unwrap();
569
570 let optimized = shard.optimize().unwrap();
571 assert!(
572 !optimized,
573 "exactly 20% deletion (not strictly greater) should not trigger vacuum"
574 );
575
576 assert_points_retrievable_with_vectors(&shard, &[201, 500, 1000]);
577 }
578
579 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
581 #[test]
582 fn vacuum_just_above_threshold_triggers() {
583 let dir = tempfile::Builder::new()
584 .prefix("edge-opt-vacuum-above")
585 .tempdir()
586 .unwrap();
587
588 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
589
590 let points = (1..=1000).map(point).collect::<Vec<_>>();
591 shard
592 .update(PointOperation(UpsertPoints(PointsList(points))))
593 .unwrap();
594
595 let deleted_ids = (1..=201).map(ExtendedPointId::NumId).collect();
597 shard
598 .update(PointOperation(DeletePoints { ids: deleted_ids }))
599 .unwrap();
600
601 let optimized = shard.optimize().unwrap();
602 assert!(
603 optimized,
604 "20.1% deletion should trigger vacuum (threshold is >20%)"
605 );
606
607 assert_points_retrievable_with_vectors(&shard, &[202, 500, 1000]);
609 }
610
611 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
614 #[test]
615 fn merge_and_vacuum_cooperate() {
616 let target_count = default_segment_number() + 6;
617
618 let dir = tempfile::Builder::new()
619 .prefix("edge-opt-merge-vacuum")
620 .tempdir()
621 .unwrap();
622
623 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
624
625 let points = (1..=1000).map(point).collect::<Vec<_>>();
627 shard
628 .update(PointOperation(UpsertPoints(PointsList(points))))
629 .unwrap();
630 let deleted_ids = (1..=250).map(ExtendedPointId::NumId).collect();
631 shard
632 .update(PointOperation(DeletePoints { ids: deleted_ids }))
633 .unwrap();
634 drop(shard);
635
636 multiply_segments(dir.path(), target_count);
638
639 let reopened = EdgeShard::load(dir.path(), None).unwrap();
641 reopened.optimize().unwrap();
642
643 let info = reopened.info().unwrap();
644 assert!(
645 info.segments_count <= default_segment_number() + 1,
646 "excess segments should be merged: got {}",
647 info.segments_count,
648 );
649
650 let count = reopened.count(CountRequest::new()).unwrap();
653 assert_eq!(
654 count, 750,
655 "merged shard should preserve surviving points without double-counting"
656 );
657
658 assert_points_retrievable_with_vectors(&reopened, &[251, 500, 1000]);
660
661 let optimized = reopened.optimize().unwrap();
663 assert!(!optimized, "second run should be idle after merge+vacuum");
664 }
665
666 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
668 #[test]
669 fn data_survives_optimize_and_reload() {
670 let dir = tempfile::Builder::new()
671 .prefix("edge-opt-reload")
672 .tempdir()
673 .unwrap();
674
675 let shard = EdgeShard::new(dir.path(), test_config()).unwrap();
676
677 let points = (1..=1000).map(point).collect::<Vec<_>>();
678 shard
679 .update(PointOperation(UpsertPoints(PointsList(points))))
680 .unwrap();
681
682 let deleted_ids = (1..=250).map(ExtendedPointId::NumId).collect();
684 shard
685 .update(PointOperation(DeletePoints { ids: deleted_ids }))
686 .unwrap();
687
688 let optimized = shard.optimize().unwrap();
689 assert!(optimized);
690 drop(shard);
691
692 let reopened = EdgeShard::load(dir.path(), None).unwrap();
694
695 let count = reopened.count(CountRequest::new()).unwrap();
696 assert_eq!(count, 750, "point count should be preserved across reload");
697
698 assert_points_retrievable_with_vectors(&reopened, &[251, 500, 750, 1000]);
700 }
701
702 #[cfg_attr(target_os = "windows", ignore = "slow on Windows, not OS-specific")]
711 #[test]
712 fn optimize_prunes_vector_deleted_from_config() {
713 let target_count = default_segment_number() + 6;
714
715 let dir = tempfile::Builder::new()
716 .prefix("edge-opt-prune-deleted-vector")
717 .tempdir()
718 .unwrap();
719
720 let shard = EdgeShard::new(dir.path(), two_vector_config()).unwrap();
721 shard
722 .update(PointOperation(UpsertPoints(PointsList(vec![
723 two_vector_point(1),
724 ]))))
725 .unwrap();
726 drop(shard);
727
728 multiply_segments(dir.path(), target_count);
730
731 let reopened = EdgeShard::load(dir.path(), None).unwrap();
732
733 reopened
736 .config
737 .write(|cfg| {
738 cfg.vectors.remove(DROP_VECTOR_NAME);
739 })
740 .unwrap();
741
742 let optimized = reopened
743 .optimize()
744 .expect("optimize must prune the deleted vector data, not cancel the merge");
745 assert!(optimized, "excess segments should have been merged");
746
747 let info = reopened.info().unwrap();
748 assert!(
749 info.segments_count <= default_segment_number() + 1,
750 "segments should be reduced after merge: got {} segments",
751 info.segments_count,
752 );
753
754 let results = reopened
756 .retrieve(
757 RetrieveRequestBuilder::new(vec![ExtendedPointId::NumId(1)])
758 .with_payload(WithPayloadInterface::Bool(false))
759 .with_vector(WithVector::Bool(true))
760 .build(),
761 )
762 .unwrap();
763 assert_eq!(results.len(), 1, "point should survive the merge");
764 let vectors = match results[0]
765 .vector
766 .as_ref()
767 .expect("vector should be present")
768 {
769 VectorStructInternal::Named(named) => named,
770 other => panic!("expected Named vectors, got {other:?}"),
771 };
772 let keep = match vectors
773 .get(KEEP_VECTOR_NAME)
774 .expect("surviving vector should be present")
775 {
776 VectorInternal::Dense(v) => v,
777 other => panic!("expected Dense vector, got {other:?}"),
778 };
779 assert_eq!(keep, &vec![1.0], "surviving vector value mismatch");
780 }
781
782 #[test]
786 fn optimizers_read_live_vector_names() {
787 let dir = tempfile::Builder::new()
788 .prefix("edge-opt-live-vector-names")
789 .tempdir()
790 .unwrap();
791
792 let shard = EdgeShard::new(dir.path(), two_vector_config()).unwrap();
793 let optimizers = shard.build_blocking_optimizers();
794
795 for optimizer in &optimizers {
796 let names = optimizer
797 .segment_optimizer_config()
798 .live_vector_names()
799 .expect("live vector names must be wired into edge optimizers");
800 assert!(
801 names.contains(KEEP_VECTOR_NAME) && names.contains(DROP_VECTOR_NAME),
802 "both vector names should be live before the delete: {names:?}",
803 );
804 }
805
806 shard
807 .update(VectorNameOperation(VectorNameOperations::DeleteVectorName(
808 DeleteVectorName {
809 vector_name: DROP_VECTOR_NAME.to_string(),
810 },
811 )))
812 .unwrap();
813
814 for optimizer in &optimizers {
815 let names = optimizer
816 .segment_optimizer_config()
817 .live_vector_names()
818 .expect("live vector names must be wired into edge optimizers");
819 assert!(names.contains(KEEP_VECTOR_NAME));
820 assert!(
821 !names.contains(DROP_VECTOR_NAME),
822 "a deletion after optimizer build must be visible to the live provider",
823 );
824 }
825 }
826
827 fn assert_points_retrievable_with_vectors(shard: &EdgeShard, ids: &[u64]) {
830 let point_ids = ids
831 .iter()
832 .map(|id| ExtendedPointId::NumId(*id))
833 .collect::<Vec<_>>();
834 let results = shard
835 .retrieve(
836 RetrieveRequestBuilder::new(point_ids)
837 .with_payload(WithPayloadInterface::Bool(false))
838 .with_vector(WithVector::Bool(true))
839 .build(),
840 )
841 .unwrap();
842 assert_eq!(
843 results.len(),
844 ids.len(),
845 "expected {} retrievable points, got {}",
846 ids.len(),
847 results.len(),
848 );
849 for (result, &expected_id) in results.iter().zip(ids) {
850 assert_eq!(result.id, ExtendedPointId::NumId(expected_id));
851 let vectors = match result.vector.as_ref().expect("vector should be present") {
852 VectorStructInternal::Named(named) => named,
853 other => panic!("expected Named vectors, got {other:?}"),
854 };
855 let vec = match vectors.get(VECTOR_NAME).expect("vector name should exist") {
856 VectorInternal::Dense(v) => v,
857 other => panic!("expected Dense vector, got {other:?}"),
858 };
859 assert_eq!(
860 vec,
861 &vec![expected_id as f32],
862 "vector value mismatch for point {expected_id}"
863 );
864 }
865 }
866
867 fn test_config() -> EdgeConfig {
868 EdgeConfig {
869 on_disk_payload: Some(false),
870 vectors: HashMap::from([(
871 VECTOR_NAME.to_string(),
872 EdgeVectorParams {
873 size: 1,
874 distance: Distance::Dot,
875 quantization_config: None,
876 multivector_config: None,
877 datatype: None,
878 on_disk: None,
879 hnsw_config: None,
880 },
881 )]),
882 sparse_vectors: HashMap::new(),
883 hnsw_config: None,
884 quantization_config: None,
885 optimizers: None,
886 wal_options: None,
887 max_search_threads: None,
888 search_pool_core: None,
889 }
890 }
891
892 const KEEP_VECTOR_NAME: &str = "edge-keep-vector";
893 const DROP_VECTOR_NAME: &str = "edge-drop-vector";
894
895 fn two_vector_config() -> EdgeConfig {
897 let params = EdgeVectorParams {
898 size: 1,
899 distance: Distance::Dot,
900 quantization_config: None,
901 multivector_config: None,
902 datatype: None,
903 on_disk: None,
904 hnsw_config: None,
905 };
906 EdgeConfig {
907 vectors: HashMap::from([
908 (KEEP_VECTOR_NAME.to_string(), params.clone()),
909 (DROP_VECTOR_NAME.to_string(), params),
910 ]),
911 ..test_config()
912 }
913 }
914
915 fn two_vector_point(id: u64) -> PointStructPersisted {
917 PointStructPersisted {
918 id: ExtendedPointId::NumId(id),
919 vector: VectorStructPersisted::from(VectorStructInternal::Named(HashMap::from([
920 (
921 KEEP_VECTOR_NAME.to_string(),
922 VectorInternal::from(vec![id as f32]),
923 ),
924 (
925 DROP_VECTOR_NAME.to_string(),
926 VectorInternal::from(vec![-(id as f32)]),
927 ),
928 ]))),
929 payload: None,
930 }
931 }
932
933 fn point(id: u64) -> PointStructPersisted {
934 PointStructPersisted {
935 id: ExtendedPointId::NumId(id),
936 vector: VectorStructPersisted::from(VectorStructInternal::Named(HashMap::from([(
937 VECTOR_NAME.to_string(),
938 VectorInternal::from(vec![id as f32]),
939 )]))),
940 payload: None,
941 }
942 }
943
944 fn multiply_segments(shard_dir: &Path, target_count: usize) {
946 let segments_path = shard_dir.join("segments");
947 let segment_dirs = fs::read_dir(&segments_path)
948 .unwrap()
949 .filter_map(Result::ok)
950 .map(|entry| entry.path())
951 .filter(|path| path.is_dir())
952 .collect::<Vec<_>>();
953 assert!(!segment_dirs.is_empty(), "need at least one source segment");
954
955 let source = &segment_dirs[0];
956 let current_count = segment_dirs.len();
957 for _ in current_count..target_count {
958 let target = segments_path.join(Uuid::new_v4().to_string());
959 copy_dir_recursive(source, &target);
960 }
961 }
962
963 fn duplicate_single_segment(shard_dir: &Path) {
964 let segments_path = shard_dir.join("segments");
965 let segment_dirs = fs::read_dir(&segments_path)
966 .unwrap()
967 .filter_map(Result::ok)
968 .map(|entry| entry.path())
969 .filter(|path| path.is_dir())
970 .collect::<Vec<_>>();
971 assert_eq!(segment_dirs.len(), 1, "expected exactly one source segment");
972
973 let source = &segment_dirs[0];
974 let target = segments_path.join(Uuid::new_v4().to_string());
975 copy_dir_recursive(source, &target);
976 }
977
978 fn copy_dir_recursive(from: &Path, to: &Path) {
979 fs::create_dir_all(to).unwrap();
980 for entry in fs::read_dir(from).unwrap().filter_map(Result::ok) {
981 let from_path = entry.path();
982 let to_path = to.join(entry.file_name());
983 if entry.file_type().unwrap().is_dir() {
984 copy_dir_recursive(&from_path, &to_path);
985 } else {
986 fs::copy(&from_path, &to_path).unwrap();
987 }
988 }
989 }
990}