1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
// ============================================================================
// Point Lookup - Single-value equality operations
// ============================================================================
use vibesql_types::SqlValue;
use super::{
index_metadata::{acquire_btree_lock, IndexData},
value_normalization::{normalize_cow, normalize_for_comparison},
};
impl IndexData {
/// Optimized single-key lookup - avoids Vec allocation for common single-column index case.
///
/// # Arguments
/// * `key` - Single key value to look up (will be normalized for consistent comparison)
///
/// # Returns
/// Owned vector of row indices if key exists, None otherwise
///
/// # Performance
/// This method avoids the Vec<SqlValue> allocation that `get()` requires, making it
/// significantly faster for single-column index lookups (the common case for primary keys).
#[inline]
pub fn get_single(&self, key: &SqlValue) -> Option<Vec<usize>> {
// Use Cow to avoid cloning non-numeric values
let normalized_key = normalize_cow(key);
match self {
IndexData::InMemory { data } => {
// Create a temporary slice for lookup without Vec allocation
let key_slice: &[SqlValue] = std::slice::from_ref(normalized_key.as_ref());
data.get(key_slice).cloned()
}
IndexData::DiskBacked { btree, .. } => {
// For disk-backed, we still need a Vec for the API, but this is less common
let key_vec = vec![normalized_key.into_owned()];
match acquire_btree_lock(btree) {
Ok(guard) => match guard.lookup(&key_vec) {
Ok(row_ids) if !row_ids.is_empty() => Some(row_ids),
Ok(_) => None,
Err(e) => {
log::warn!("BTreeIndex lookup failed in get_single: {}", e);
None
}
},
Err(e) => {
log::warn!("BTreeIndex lock acquisition failed in get_single: {}", e);
None
}
}
}
IndexData::IVFFlat { .. } | IndexData::Hnsw { .. } => None,
}
}
/// Optimized single-key existence check - avoids Vec allocation.
///
/// # Arguments
/// * `key` - Single key value to check (will be normalized for consistent comparison)
///
/// # Returns
/// true if key exists, false otherwise
///
/// # Performance
/// This method avoids the Vec<SqlValue> allocation that `contains_key()` requires.
#[inline]
pub fn contains_key_single(&self, key: &SqlValue) -> bool {
let normalized_key = normalize_cow(key);
match self {
IndexData::InMemory { data, .. } => {
let key_slice: &[SqlValue] = std::slice::from_ref(normalized_key.as_ref());
data.contains_key(key_slice)
}
IndexData::DiskBacked { btree, .. } => {
let key_vec = vec![normalized_key.into_owned()];
match acquire_btree_lock(btree) {
Ok(guard) => match guard.lookup(&key_vec) {
Ok(row_ids) => !row_ids.is_empty(),
Err(e) => {
log::warn!("BTreeIndex lookup failed in contains_key_single: {}", e);
false
}
},
Err(e) => {
log::warn!(
"BTreeIndex lock acquisition failed in contains_key_single: {}",
e
);
false
}
}
}
IndexData::IVFFlat { .. } | IndexData::Hnsw { .. } => false,
}
}
/// Lookup exact key in the index
///
/// # Arguments
/// * `key` - Key to look up (values will be normalized for consistent comparison)
///
/// # Returns
/// Owned vector of row indices if key exists, None otherwise
///
/// # Note
/// This is the primary point-lookup API for index queries.
/// Returns owned data to support both in-memory (cloned) and disk-backed (loaded) indexes.
/// Values are normalized (e.g., Integer -> Double) to match insertion-time normalization.
///
/// # Performance
/// For single-column indexes, prefer `get_single()` which avoids Vec allocation.
pub fn get(&self, key: &[SqlValue]) -> Option<Vec<usize>> {
// Fast path for single-key lookups (common case)
if key.len() == 1 {
return self.get_single(&key[0]);
}
// Multi-column key path: normalize all key values
let normalized_key: Vec<SqlValue> = key.iter().map(normalize_for_comparison).collect();
match self {
IndexData::InMemory { data } => {
data.get(&normalized_key).cloned()
}
IndexData::DiskBacked { btree, .. } => {
// Safely acquire lock and perform lookup
match acquire_btree_lock(btree) {
Ok(guard) => {
match guard.lookup(&normalized_key) {
Ok(row_ids) if !row_ids.is_empty() => Some(row_ids),
Ok(_) => None, // Empty result means key not found
Err(e) => {
log::warn!("BTreeIndex lookup failed in get: {}", e);
None
}
}
}
Err(e) => {
log::warn!("BTreeIndex lock acquisition failed in get: {}", e);
None
}
}
}
IndexData::IVFFlat { .. } => {
// IVFFlat indexes don't support point lookups - use search() method instead
None
}
IndexData::Hnsw { .. } => {
// HNSW indexes don't support point lookups - use search() method instead
None
}
}
}
/// Check if a key exists in the index
///
/// # Arguments
/// * `key` - Key to check (values will be normalized for consistent comparison)
///
/// # Returns
/// true if key exists, false otherwise
///
/// # Note
/// Used primarily for UNIQUE constraint validation.
/// Values are normalized (e.g., Integer -> Double) to match insertion-time normalization.
///
/// # Performance
/// For single-column indexes, prefer `contains_key_single()` which avoids Vec allocation.
pub fn contains_key(&self, key: &[SqlValue]) -> bool {
// Fast path for single-key lookups (common case)
if key.len() == 1 {
return self.contains_key_single(&key[0]);
}
// Multi-column key path: normalize all key values
let normalized_key: Vec<SqlValue> = key.iter().map(normalize_for_comparison).collect();
match self {
IndexData::InMemory { data, .. } => data.contains_key(&normalized_key),
IndexData::DiskBacked { btree, .. } => {
// Safely acquire lock and check if key exists
match acquire_btree_lock(btree) {
Ok(guard) => match guard.lookup(&normalized_key) {
Ok(row_ids) => !row_ids.is_empty(),
Err(e) => {
log::warn!("BTreeIndex lookup failed in contains_key: {}", e);
false
}
},
Err(e) => {
log::warn!("BTreeIndex lock acquisition failed in contains_key: {}", e);
false
}
}
}
IndexData::IVFFlat { .. } => {
// IVFFlat indexes don't support contains_key - use search() method instead
false
}
IndexData::Hnsw { .. } => {
// HNSW indexes don't support contains_key - use search() method instead
false
}
}
}
/// Lookup multiple values in the index (for IN predicates)
///
/// # Arguments
/// * `values` - List of values to look up
///
/// # Returns
/// Vector of row indices that match any of the values
///
/// # Performance
/// Uses Cow normalization to avoid cloning non-numeric values, and slice::from_ref
/// to avoid Vec allocation per key lookup.
pub fn multi_lookup(&self, values: &[SqlValue]) -> Vec<usize> {
match self {
IndexData::InMemory { data } => {
// Deduplicate values to avoid returning duplicate rows
// For example, WHERE a IN (10, 10, 20) should only look up 10 once
let mut unique_values: Vec<&SqlValue> = values.iter().collect();
unique_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
unique_values.dedup();
let mut matching_row_indices = Vec::new();
for value in unique_values {
// Use Cow to avoid cloning non-numeric values
let normalized_value = normalize_cow(value);
// Use slice::from_ref to avoid Vec allocation per key
let search_key: &[SqlValue] = std::slice::from_ref(normalized_value.as_ref());
if let Some(row_indices) = data.get(search_key) {
matching_row_indices.extend(row_indices);
}
}
// Return row indices in the order they were collected from BTreeMap
// For IN predicates, we collect results for each value in the order
// specified. We should NOT sort by row index as that would destroy
// the semantic ordering of the results.
matching_row_indices
}
IndexData::DiskBacked { btree, .. } => {
// Deduplicate values to avoid returning duplicate rows
let mut unique_values: Vec<&SqlValue> = values.iter().collect();
unique_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
unique_values.dedup();
// For disk-backed, we need Vec<Vec<SqlValue>> for the API
// but we still use normalize_cow to reduce cloning
let keys: Vec<Vec<SqlValue>> =
unique_values.iter().map(|v| vec![normalize_cow(v).into_owned()]).collect();
// Safely acquire lock and call BTreeIndex::multi_lookup
match acquire_btree_lock(btree) {
Ok(guard) => guard.multi_lookup(&keys).unwrap_or_else(|_| vec![]),
Err(e) => {
// Log error and return empty result set
log::warn!("BTreeIndex lock acquisition failed in multi_lookup: {}", e);
vec![]
}
}
}
IndexData::IVFFlat { .. } => {
// IVFFlat indexes don't support multi_lookup - use search() method instead
vec![]
}
IndexData::Hnsw { .. } => {
// HNSW indexes don't support multi_lookup - use search() method instead
vec![]
}
}
}
/// Get an iterator over all key-value pairs in the index
///
/// # Returns
/// Iterator yielding owned (key, row_indices) pairs
///
/// # Note
/// For in-memory indexes, iteration is in sorted key order (BTreeMap ordering).
/// This method enables index scanning operations without exposing internal data structures.
/// Returns owned data to support both in-memory (cloned) and disk-backed (loaded) indexes.
///
/// **Note**: For disk-backed indexes, this requires a full B+ tree scan and is expensive.
/// Currently returns an empty iterator as the BTreeIndex doesn't expose key-level iteration.
/// Most use cases should use `values()` for full scans or `multi_lookup()` for specific keys.
pub fn iter(&self) -> Box<dyn Iterator<Item = (Vec<SqlValue>, Vec<usize>)> + '_> {
match self {
IndexData::InMemory { data } => {
Box::new(data.iter().map(move |(k, v)| (k.clone(), v.clone())))
}
IndexData::DiskBacked { .. } => {
// BTreeIndex doesn't currently expose an API for iterating over (key, row_ids)
// pairs This would require adding a scan API that preserves key
// groupings For now, return empty iterator since this method is
// rarely used Callers should use values() for full scans or
// lookup()/multi_lookup() for point queries
log::warn!("DiskBacked iter() is not yet implemented - use values() instead");
Box::new(std::iter::empty())
}
IndexData::IVFFlat { .. } => {
// IVFFlat indexes don't support iteration - use search() method instead
Box::new(std::iter::empty())
}
IndexData::Hnsw { .. } => {
// HNSW indexes don't support iteration - use search() method instead
Box::new(std::iter::empty())
}
}
}
/// Get an iterator over all row index vectors in the index
///
/// # Returns
/// Iterator yielding owned row index vectors
///
/// # Note
/// This method is used for full index scans where we need all row indices
/// regardless of the key values. Returns owned data to support both in-memory
/// (cloned) and disk-backed (loaded from disk) indexes.
pub fn values(&self) -> Box<dyn Iterator<Item = Vec<usize>> + '_> {
match self {
IndexData::InMemory { data } => Box::new(data.values().cloned()),
IndexData::DiskBacked { btree, .. } => {
// Perform a full range scan to get all values
// Use range_scan with no bounds to scan entire index
match acquire_btree_lock(btree) {
Ok(guard) => {
match guard.range_scan(None, None, true, true) {
Ok(all_row_ids) => {
// Group row_ids by their appearance (BTree returns them in key
// order) For full scan, we just
// need all row IDs, so wrap in a single Vec
Box::new(std::iter::once(all_row_ids))
}
Err(e) => {
log::warn!("BTreeIndex range_scan failed in values: {}", e);
Box::new(std::iter::empty())
}
}
}
Err(e) => {
log::warn!("BTreeIndex lock acquisition failed in values: {}", e);
Box::new(std::iter::empty())
}
}
}
IndexData::IVFFlat { index } => {
// Return all row IDs stored in the IVFFlat index
let all_row_ids: Vec<usize> = index.all_row_ids();
Box::new(std::iter::once(all_row_ids))
}
IndexData::Hnsw { index } => {
// Return all row IDs stored in the HNSW index
let all_row_ids: Vec<usize> = index.all_row_ids();
Box::new(std::iter::once(all_row_ids))
}
}
}
}