Skip to main content

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}