diskann_label_filter/set/traits.rs
1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::borrow::Cow;
7use std::hash::Hash;
8
9use diskann::ANNResult;
10use diskann_utils::future::AsyncFriendly;
11
12/// A simple abstraction over a mathematical set, parameterized by the element type.
13///
14/// This trait lets you plug in different concrete set backends, for example a roaring bitmap
15/// for integer sets, without changing call sites. Each concrete `Set` only accepts its own
16/// element type.
17///
18/// Requirements:
19/// - `Clone` and `Default` so callers can duplicate and construct sets easily
20/// - `IntoIterator<Item = Element>` so callers can iterate elements directly
21pub trait Set<Element>: Clone + Default + IntoIterator<Item = Element> {
22 /// Create an empty set, equivalent to `Default::default`.
23 fn empty_set() -> Self;
24
25 /// Return the intersection of `self` and `other`.
26 fn intersection(&self, other: &Self) -> Self;
27
28 /// Return the union of `self` and `other`.
29 fn union(&self, other: &Self) -> Self;
30
31 /// Insert a value, return false if the element already exists
32 /// and true otherwise.
33 /// Return ANNError if there was an error in the operation.
34 fn insert(&mut self, value: &Element) -> ANNResult<bool>;
35
36 /// Remove a value, return true if the value was present,
37 /// false otherwise.
38 /// Return ANNError if there was an underlying error.
39 fn remove(&mut self, value: &Element) -> ANNResult<bool>;
40
41 /// Return true if `value` is a member of the set.
42 /// Return ANNError if the operation failed.
43 fn contains(&self, value: &Element) -> ANNResult<bool>;
44
45 /// Remove all elements from the set.
46 /// Return ANNError if the operation failed.
47 fn clear(&mut self) -> ANNResult<()>;
48
49 /// Return the number of elements in the set.
50 /// Return ANNError if the operation failed.
51 fn len(&self) -> ANNResult<usize>;
52
53 /// Return true if the set is empty.
54 /// Return ANNError if the operation failed.
55 fn is_empty(&self) -> ANNResult<bool>;
56}
57
58/// Provider for sets that may live in memory or in a storage layer.
59///
60/// A `SetProvider` owns or manages many sets, addressed by a key,
61/// for example vector identifiers to label sets. Each implementation
62/// chooses a concrete set type that fits its backend and usage.
63///
64/// Type parameters:
65/// - `Key` is the lookup key for an individual set, for example a vector id
66/// - `Value` is the element type stored inside each set
67pub trait SetProvider<Key, Value>: AsyncFriendly
68where
69 Key: Clone + Eq + Hash,
70 Value: Clone + Eq + Hash,
71{
72 /// Concrete set type managed by this provider.
73 type S: Set<Value> + AsyncFriendly;
74
75 /// Get the set for `key`.
76 ///
77 /// Implementations may borrow internally and return `Cow::Borrowed`,
78 /// or materialize an owned set and return `Cow::Owned`.
79 /// Returns:
80 /// Ok(Some(Set)) of the key if the key exists
81 /// Ok(None) if the key doesn't exist
82 /// ANNError if there was an error, say in the underlying store.
83 fn get(&'_ self, key: &Key) -> ANNResult<Option<Cow<'_, Self::S>>>;
84
85 /// Number of keys managed by this provider.
86 /// Returns ANNError if there was an underlying error while retrieving
87 /// the result (for instance, in the backing store)
88 fn count(&self) -> ANNResult<usize>;
89
90 /// Check if a key exists.
91 /// Returns true if the key exists, false otherwise,
92 /// and ANNError if an error was encountered while
93 /// processing the operation.
94 fn exists(&self, key: &Key) -> ANNResult<bool>;
95
96 /// Insert `value` into the set for `key`.
97 ///
98 /// If the key is missing, create an empty set first, then insert.
99 /// Returns true if the value was inserted, false if the *value*
100 /// already exists, has nothing to do with the key. This is because
101 /// we always expect multiple insert calls with the same key, so it
102 /// is not important whether the key exists before insert or not.
103 ///
104 /// Returns ANNError if there was an underlying error in the operation.
105 fn insert(&mut self, key: &Key, value: &Value) -> ANNResult<bool>;
106
107 /// Insert values for a key
108 /// If the key is missing, it will create a new entry.
109 /// Returns true if the all the values were inserted, false
110 /// if any of them already exist.
111 ///
112 /// Returns ANNError if there is an underlying failure.
113 /// Currently, there is no way to indicate the insertion status
114 /// of an individual value.
115 fn insert_values(&mut self, key: &Key, values: &[Value]) -> ANNResult<bool>;
116
117 /// Delete the entire set for `key`.
118 ///
119 /// Returns true if the key is present. false, otherwise, and ANNError
120 /// if an error was encountered during the operation.
121 fn delete(&mut self, key: &Key) -> ANNResult<bool>;
122
123 /// Remove `value` from the set for `key`.
124 ///
125 /// Returns true if the value was present and removed. false otherwise.
126 /// ANNError is returned if the operation fails.
127 fn delete_from_set(&mut self, key: &Key, value: &Value) -> ANNResult<bool>;
128
129 ///Clear all the keys and values in the set.
130 ///
131 /// Returns error if the operation failedd
132 fn clear(&mut self) -> ANNResult<()>;
133}