Skip to main content

qdrant_edge/edge/edge_shard/
optimize.rs

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    /// Run shard optimizers in-process and blocking until no more optimization plans are produced.
26    ///
27    /// This is synchronous and does not spawn background optimization workers.
28    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                // Bypass budget in Edge, always allocate the full desired IO for the optimizer.
49                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            // Optimization swapped segments in/out; reflect the new set in the manifest.
76            self.update_segment_manifest()?;
77
78            // Avoid repeating the same plan forever if no optimizer made effective progress.
79            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        // Live read of the vector names so the merge can prune a vector deleted from the config
91        // (whose data still lingers in older segment files) instead of cancelling forever, while
92        // still cancelling on the CreateVectorName race. Safe here for the same reason as the
93        // server wiring: `update` holds the segments read guard across both the segment
94        // application and the config update, so any name a proxy-frozen source segment carries is
95        // already visible in this read. The provider uses a synchronous `parking_lot` read, fine
96        // on this blocking path.
97        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// Tests in this module exercise platform-agnostic optimizer logic but run
154// 5-25x slower on Windows due to filesystem IO. They are marked
155// `#[ignore]` on Windows; Linux and macOS jobs provide full coverage.
156// To execute them locally on Windows, run with `cargo test -- --ignored`.
157#[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        // Delete 250/1000 = 25%, above DEFAULT_DELETED_THRESHOLD (20%)
249        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        // Verify surviving points are queryable with correct vectors
264        assert_points_retrievable_with_vectors(&shard, &[251, 500, 999, 1000]);
265    }
266
267    /// A fresh shard with a single small segment and no deletions should not
268    /// trigger any optimizer.
269    #[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    /// An empty shard (no data at all) should be a no-op.
293    #[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    /// Creating more segments than `default_segment_number` should trigger
309    /// the merge optimizer to reduce the segment count.
310    #[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        // All duplicated segments contained the same point (id=1). After merge,
341        // the exact count should deduplicate across remaining segments.
342        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    /// After a merge optimization, a second run should be a no-op.
349    #[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        // First explicit optimization triggers merge.
369        reopened.optimize().unwrap();
370        let segments_after_first = reopened.info().unwrap().segments_count;
371
372        // Second explicit optimization should be a no-op.
373        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    /// Deleting less than 20% of points (below the vacuum threshold)
387    /// should NOT trigger the vacuum optimizer.
388    #[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        // Delete 5% — below the 20% threshold (DEFAULT_DELETED_THRESHOLD)
404        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        // Surviving points should still have correct vectors
416        assert_points_retrievable_with_vectors(&shard, &[51, 500, 1000]);
417    }
418
419    /// Deleting below the minimum vector count (< 1000 total points)
420    /// should NOT trigger the vacuum optimizer even with a high deletion ratio.
421    #[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        // Only 100 points total (below DEFAULT_VACUUM_MIN_VECTOR_NUMBER=1000)
432        let points = (1..=100).map(point).collect::<Vec<_>>();
433        shard
434            .update(PointOperation(UpsertPoints(PointsList(points))))
435            .unwrap();
436
437        // Delete 50% — above ratio threshold (20%), but total count is below minimum
438        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    /// After vacuum optimization, all non-deleted points should still be
454    /// retrievable and deleted points should be gone.
455    #[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        // Delete points 1..=250 (25%, above DEFAULT_DELETED_THRESHOLD=20%)
471        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        // Verify point count
482        let count = shard.count(CountRequest::new()).unwrap();
483        assert_eq!(count, 750, "should have 750 remaining points after vacuum");
484
485        // Verify deleted points are gone
486        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        // Verify surviving points are accessible with correct vectors
500        assert_points_retrievable_with_vectors(&shard, &[251, 500, 750, 1000]);
501    }
502
503    /// Deleting all points from a segment should be handled gracefully.
504    /// The vacuum optimizer plans the segment for rebuild, but because the
505    /// resulting segment has 0 points, `optimize_all_segments_blocking`
506    /// reports `false` (zero points processed). The shard should still be
507    /// valid and accept new data afterward.
508    #[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        // Delete ALL points
524        let deleted_ids = (1..=1000).map(ExtendedPointId::NumId).collect();
525        shard
526            .update(PointOperation(DeletePoints { ids: deleted_ids }))
527            .unwrap();
528
529        // The vacuum optimizer rebuilds the segment, but since 0 points remain
530        // in the result, `points_optimized == 0` and the function returns false.
531        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 should still be functional — can insert new points
537        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    /// Vacuum at exactly the threshold boundary (20% deleted, 1000 total).
547    /// The threshold check is strictly greater-than, so exactly 20% should
548    /// NOT trigger vacuum.
549    #[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        // Delete exactly 20% (200 out of 1000) — matches DEFAULT_DELETED_THRESHOLD
565        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    /// Just above the vacuum threshold should trigger optimization.
580    #[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        // Delete 201 out of 1000 = 20.1% — just above DEFAULT_DELETED_THRESHOLD (20%)
596        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        // Points 202..=1000 should survive with correct vectors
608        assert_points_retrievable_with_vectors(&shard, &[202, 500, 1000]);
609    }
610
611    /// When there are excess segments AND some have high deletion ratios,
612    /// optimization should handle both (merge + vacuum).
613    #[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        // Insert 1000 points, then delete 250 (25% — above DEFAULT_DELETED_THRESHOLD=20%)
626        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        // Create excess segments
637        multiply_segments(dir.path(), target_count);
638
639        // Explicit optimization (both merge and vacuum should run)
640        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        // The duplicated segments each had 750 surviving points (same IDs).
651        // After merge, the shard should be functional with correct data.
652        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        // Surviving points (251..=1000) should be queryable with correct vectors
659        assert_points_retrievable_with_vectors(&reopened, &[251, 500, 1000]);
660
661        // Second run should be idle
662        let optimized = reopened.optimize().unwrap();
663        assert!(!optimized, "second run should be idle after merge+vacuum");
664    }
665
666    /// Optimized shard should survive a reload and still serve correct data.
667    #[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        // Delete 250 points (25%, above DEFAULT_DELETED_THRESHOLD=20%), then optimize
683        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        // Reload the shard
693        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        // Verify specific points survive reload with correct vectors
699        assert_points_retrievable_with_vectors(&reopened, &[251, 500, 750, 1000]);
700    }
701
702    /// Regression test for the deleted-vector optimizer deadlock on edge.
703    ///
704    /// A vector name deleted from the config can still linger in older segment files: a
705    /// `DeleteVectorName` landing while a segment is proxy-frozen mid-optimization updates the
706    /// config and the proxy, but not the frozen source. This test recreates that persisted state
707    /// directly (config lacks the vector, segments still carry its data) and checks `optimize()`
708    /// prunes the stale data instead of cancelling every merge attempt (and, since edge propagates
709    /// the cancellation, erroring out of `optimize()` forever).
710    #[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        // Create excess segments so the merge optimizer has to rebuild them.
729        multiply_segments(dir.path(), target_count);
730
731        let reopened = EdgeShard::load(dir.path(), None).unwrap();
732
733        // Emulate the raced deletion: remove the vector from the config only, leaving its data in
734        // the on-disk segments, exactly what a mid-optimization `DeleteVectorName` leaves behind.
735        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        // No data loss on the surviving vector.
755        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    /// The optimizers must observe vector-name deletions that land after they were built: the
783    /// live provider re-reads the config on every call rather than snapshotting it at build time.
784    /// A snapshot would wrongly prune a vector created between the snapshot and the merge.
785    #[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    /// Retrieve points by ID and verify each one is present with the correct
828    /// vector value. Every test point was created with vector `[id as f32]`.
829    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    /// Like [`test_config`], but with two dense vectors so one can be deleted.
896    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    /// A point populating both vectors of [`two_vector_config`]: keep = `[id]`, drop = `[-id]`.
916    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    /// Copy the first segment on disk to reach `target_count` total segments.
945    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}