diskann_benchmark_core/streaming/graph/
inplace_delete.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::{ops::Range, sync::Arc};
7
8use diskann::{
9    ANNResult,
10    graph::{self, glue},
11    provider,
12};
13
14use crate::build::{Build, ids::ToIdSized};
15
16/// A built-in helper for running and benchmarking
17/// [`inplace_delete`](graph::DiskANNIndex::inplace_delete).
18///
19/// This is intended to be used in conjunction with [`crate::build::build`] and
20/// [`crate::build::build_tracked`].
21pub struct InplaceDelete<DP, S>
22where
23    DP: provider::DataProvider,
24{
25    index: Arc<graph::DiskANNIndex<DP>>,
26    strategy: S,
27    num_to_replace: usize,
28    inplace_delete_method: graph::InplaceDeleteMethod,
29    to_id: Box<dyn ToIdSized<DP::ExternalId>>,
30}
31
32impl<DP, S> InplaceDelete<DP, S>
33where
34    DP: provider::DataProvider,
35{
36    /// Construct a new [`InplaceDelete`] build stage.
37    ///
38    /// This [`Build`] object will run for all Ids provided by `to_id`, invoking
39    /// [`diskann::graph::DiskANNIndex::inplace_delete`] on each ID.
40    ///
41    /// Arguments `num_to_replace` and `inplace_delete_method` are passed directly to the method
42    /// on [`diskann::graph::DiskANNIndex`].
43    pub fn new(
44        index: Arc<graph::DiskANNIndex<DP>>,
45        strategy: S,
46        num_to_replace: usize,
47        inplace_delete_method: graph::InplaceDeleteMethod,
48        to_id: impl ToIdSized<DP::ExternalId> + 'static,
49    ) -> Arc<Self> {
50        Arc::new(Self {
51            index,
52            strategy,
53            num_to_replace,
54            inplace_delete_method,
55            to_id: Box::new(to_id),
56        })
57    }
58}
59
60impl<DP, S> Build for InplaceDelete<DP, S>
61where
62    DP: provider::DataProvider<Context: Default> + provider::Delete,
63    S: glue::InplaceDeleteStrategy<DP> + Clone,
64{
65    type Output = ();
66
67    fn num_data(&self) -> usize {
68        self.to_id.len()
69    }
70
71    async fn build(&self, range: Range<usize>) -> ANNResult<Self::Output> {
72        for i in range {
73            let context = DP::Context::default();
74            self.index
75                .inplace_delete(
76                    self.strategy.clone(),
77                    &context,
78                    &self.to_id.to_id(i)?,
79                    self.num_to_replace,
80                    self.inplace_delete_method,
81                )
82                .await?;
83        }
84        Ok(())
85    }
86}
87
88///////////
89// Tests //
90///////////
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    use std::num::NonZeroUsize;
97
98    use diskann::{graph::test::provider, provider::Delete, utils::ONE};
99
100    use crate::{build, streaming::graph::test};
101
102    #[test]
103    fn test_inplace_delete() {
104        let (index, num_points) = test::build_test_index();
105        let points_to_delete = [10, 20, 30, 40];
106
107        let rt = crate::tokio::runtime(2).unwrap();
108        let _ = build::build(
109            InplaceDelete::new(
110                index.clone(),
111                provider::Strategy::new(),
112                4,
113                graph::InplaceDeleteMethod::TwoHopAndOneHop,
114                build::ids::Slice::new(points_to_delete.into()),
115            ),
116            build::Parallelism::dynamic(ONE, NonZeroUsize::new(2).unwrap()),
117            &rt,
118        )
119        .unwrap();
120
121        let num_points: u32 = num_points.try_into().unwrap();
122        let ctx = provider::Context::new();
123        for i in 0..num_points {
124            let is_deleted = rt
125                .block_on(index.provider().status_by_external_id(&ctx, &i))
126                .unwrap()
127                .is_deleted();
128            if points_to_delete.contains(&i) {
129                assert!(is_deleted, "expected {i} to be deleted");
130            } else {
131                assert!(!is_deleted, "expected {i} to not be deleted");
132            }
133        }
134    }
135}
diskann_benchmark_core/streaming/graph/inplace_delete.rs

diskann_benchmark_core/streaming/graph/
inplace_delete.rs
