weavatrix_search_vector/
metadata.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3/// Deterministically ordered metadata value attached to a vector key.
4#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
5#[non_exhaustive]
6pub enum MetadataValue {
7    Bool(bool),
8    I64(i64),
9    U64(u64),
10    Text(String),
11    Bytes(Vec<u8>),
12}
13
14/// Metadata fields for one vector.
15#[derive(Debug, Clone, Default, PartialEq, Eq)]
16pub struct Metadata {
17    fields: BTreeMap<String, MetadataValue>,
18}
19
20impl Metadata {
21    #[must_use]
22    pub const fn new() -> Self {
23        Self {
24            fields: BTreeMap::new(),
25        }
26    }
27
28    pub fn insert(
29        &mut self,
30        field: impl Into<String>,
31        value: MetadataValue,
32    ) -> Option<MetadataValue> {
33        self.fields.insert(field.into(), value)
34    }
35
36    #[must_use]
37    pub fn get(&self, field: &str) -> Option<&MetadataValue> {
38        self.fields.get(field)
39    }
40
41    pub fn remove(&mut self, field: &str) -> Option<MetadataValue> {
42        self.fields.remove(field)
43    }
44
45    #[must_use]
46    pub fn len(&self) -> usize {
47        self.fields.len()
48    }
49
50    #[must_use]
51    pub fn is_empty(&self) -> bool {
52        self.fields.is_empty()
53    }
54
55    #[must_use]
56    pub fn iter(&self) -> impl ExactSizeIterator<Item = (&str, &MetadataValue)> {
57        self.fields
58            .iter()
59            .map(|(field, value)| (field.as_str(), value))
60    }
61}
62
63/// Composable deterministic predicate over vector metadata.
64#[derive(Debug, Clone, PartialEq, Eq)]
65#[non_exhaustive]
66pub enum MetadataFilter {
67    All,
68    Exists(String),
69    Equal {
70        field: String,
71        value: MetadataValue,
72    },
73    I64Range {
74        field: String,
75        minimum: Option<i64>,
76        maximum: Option<i64>,
77    },
78    U64Range {
79        field: String,
80        minimum: Option<u64>,
81        maximum: Option<u64>,
82    },
83    TextPrefix {
84        field: String,
85        prefix: String,
86    },
87    And(Vec<Self>),
88    Or(Vec<Self>),
89    Not(Box<Self>),
90}
91
92impl MetadataFilter {
93    #[must_use]
94    pub fn matches(&self, metadata: &Metadata) -> bool {
95        match self {
96            Self::All => true,
97            Self::Exists(field) => metadata.get(field).is_some(),
98            Self::Equal { field, value } => metadata.get(field) == Some(value),
99            Self::I64Range {
100                field,
101                minimum,
102                maximum,
103            } => matches!(
104                metadata.get(field),
105                Some(MetadataValue::I64(value))
106                    if minimum.is_none_or(|minimum| *value >= minimum)
107                        && maximum.is_none_or(|maximum| *value <= maximum)
108            ),
109            Self::U64Range {
110                field,
111                minimum,
112                maximum,
113            } => matches!(
114                metadata.get(field),
115                Some(MetadataValue::U64(value))
116                    if minimum.is_none_or(|minimum| *value >= minimum)
117                        && maximum.is_none_or(|maximum| *value <= maximum)
118            ),
119            Self::TextPrefix { field, prefix } => matches!(
120                metadata.get(field),
121                Some(MetadataValue::Text(value)) if value.starts_with(prefix)
122            ),
123            Self::And(filters) => filters.iter().all(|filter| filter.matches(metadata)),
124            Self::Or(filters) => filters.iter().any(|filter| filter.matches(metadata)),
125            Self::Not(filter) => !filter.matches(metadata),
126        }
127    }
128}
129
130/// Key-addressed metadata used to filter owned or memory-mapped indexes.
131#[derive(Debug, Clone, Default)]
132pub struct MetadataIndex {
133    records: BTreeMap<u64, Metadata>,
134}
135
136impl MetadataIndex {
137    #[must_use]
138    pub const fn new() -> Self {
139        Self {
140            records: BTreeMap::new(),
141        }
142    }
143
144    pub fn insert(&mut self, key: u64, metadata: Metadata) -> Option<Metadata> {
145        self.records.insert(key, metadata)
146    }
147
148    pub fn remove(&mut self, key: u64) -> Option<Metadata> {
149        self.records.remove(&key)
150    }
151
152    #[must_use]
153    pub fn get(&self, key: u64) -> Option<&Metadata> {
154        self.records.get(&key)
155    }
156
157    #[must_use]
158    pub fn len(&self) -> usize {
159        self.records.len()
160    }
161
162    #[must_use]
163    pub fn is_empty(&self) -> bool {
164        self.records.is_empty()
165    }
166
167    #[must_use]
168    pub fn iter(&self) -> impl ExactSizeIterator<Item = (u64, &Metadata)> {
169        self.records.iter().map(|(key, metadata)| (*key, metadata))
170    }
171
172    #[must_use]
173    pub fn matches(&self, key: u64, filter: &MetadataFilter) -> bool {
174        self.records
175            .get(&key)
176            .is_some_and(|metadata| filter.matches(metadata))
177    }
178
179    #[must_use]
180    pub fn matching_keys(&self, filter: &MetadataFilter) -> BTreeSet<u64> {
181        self.records
182            .iter()
183            .filter_map(|(key, metadata)| filter.matches(metadata).then_some(*key))
184            .collect()
185    }
186}
weavatrix_search_vector/metadata.rs

weavatrix_search_vector/
metadata.rs
