ecp_core/search/
index.rs

1use std::path::PathBuf;
2use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
3use std::sync::{Arc, Mutex, RwLock};
4
5use ndarray::Array1;
6use ndarray::Array2;
7
8use dashmap::DashMap;
9use moka::sync::Cache;
10use zarrs::array::Array;
11use zarrs::filesystem::FilesystemStore;
12use zarrs::storage::ReadableWritableListableStorage;
13
14use ordered_float::NotNan;
15
16use crate::build::tree::{BuildConfig, NodeCache as BuildNodeCache, add_data};
17use crate::build::writer::write_info_u32;
18use crate::dtype::{dtype_of_array, read_subset_as_f32};
19use crate::error::{EcpError, Result, ResultExt};
20use crate::format::FORMAT_VERSION;
21use crate::metric::Metric;
22use crate::search::info::{read_format_version, read_info_fields, read_info_u32};
23use crate::search::node::Node;
24
25mod persistence;
26mod query;
27
28use query::{HeapEntry, QueryState};
29
30/// Share of the memory limit kept for open queries' search state. The rest
31/// caches tree nodes read from disk.
32const QUERY_CACHE_MEMORY_FRACTION: f64 = 0.05;
33
34type NodeCache = Cache<(usize, u32), Arc<Node>>;
35type QueryCache = Cache<usize, Arc<Mutex<QueryState>>>;
36
37/// An eCP index opened from disk. Nodes are read on first visit and cached
38/// in memory, along with open queries. Safe to share across threads; a
39/// search and an insert only wait for each other on the same leaf.
40pub struct Index {
41    store: ReadableWritableListableStorage,
42    metric: Metric,
43    is_normalized: bool,
44    /// Node levels below the root; the last one holds the leaves.
45    levels: u32,
46    /// The root node's embeddings, always in memory.
47    root: Array2<f32>,
48    /// Nodes read so far, keyed by `(level, node_id)` with `level` 0-based.
49    nodes: NodeCache,
50    /// Open queries' search state, by query id. A query saved to disk (by
51    /// `shutdown` or eviction) is read back on first use.
52    queries: QueryCache,
53    next_query_id: AtomicUsize,
54    memory_limit_bytes: Option<usize>,
55    /// False after `shutdown`; searches then return nothing.
56    accepting: AtomicBool,
57    /// One lock per leaf node id, held only while that leaf is read from disk
58    /// or appended to, so a search never reads a leaf in the middle of an
59    /// insert's append. zarrs leaves that to its callers.
60    leaf_locks: DashMap<u32, Arc<RwLock<()>>>,
61    /// The id `insert` gives the next item. Only ever increases; it runs
62    /// ahead of `total_items` when a crash stops an insert between taking
63    /// ids and writing the items.
64    next_item_id: Mutex<u32>,
65    /// How many items are stored on disk.
66    total_items: Mutex<u32>,
67}
68
69impl Index {
70    /// Opens the index at `index_path`, reading its `info/*` fields and
71    /// root. Other nodes load on first visit. `memory_limit_bytes` caps the
72    /// node and query data kept in memory; `None` means no cap.
73    pub fn load(index_path: PathBuf, memory_limit_bytes: Option<usize>) -> Result<Self> {
74        if !index_path.exists() {
75            return Err(EcpError::NotFound(format!(
76                "index directory {} does not exist",
77                index_path.display()
78            )));
79        }
80        let store: ReadableWritableListableStorage =
81            Arc::new(FilesystemStore::new(&index_path).store_err("failed to open store")?);
82        Self::load_from_store(store, memory_limit_bytes)
83    }
84
85    /// Loads the index from `store` instead of a path. Otherwise the same as `load`.
86    /// Note: This function is only split out from `load` for unit tests.
87    fn load_from_store(
88        store: ReadableWritableListableStorage,
89        memory_limit_bytes: Option<usize>,
90    ) -> Result<Self> {
91        // Check the version first, so a newer index reports its version
92        let format_version = read_format_version(&store.clone().readable_listable())?;
93
94        let (levels, metric, is_normalized) = read_info_fields(&store.clone().readable_listable())?;
95        // Not normalized is the setting that scores any data correctly
96        let is_normalized = is_normalized.unwrap_or(false);
97
98        let root_array = Array::open(store.clone(), "/index_root/embeddings")
99            .store_err("failed to open index_root/embeddings")?;
100        let root: Array2<f32> = read_subset_as_f32(
101            &root_array,
102            &root_array.subset_all(),
103            "index_root/embeddings",
104        )?;
105
106        let (nodes, queries) = Self::build_caches(memory_limit_bytes, store.clone());
107
108        // Start after the highest saved query id, so no id is reused
109        let next_query_id = persistence::query_ids_on_disk(&store)?
110            .into_iter()
111            .max()
112            .map_or(0, |max_id| max_id + 1);
113
114        let readable = store.clone().readable_listable();
115        let next_item_id = read_info_u32(&readable, "next_item_id")?;
116        let total_items = read_info_u32(&readable, "total_items")?;
117
118        // Add the missing version now that the index has loaded
119        if format_version.is_none() {
120            match write_info_u32(&store, "format_version", FORMAT_VERSION) {
121                Ok(()) => log::warn!("index had no info/format_version, added {FORMAT_VERSION}"),
122                Err(e) => log::warn!(
123                    "index has no info/format_version and adding it failed, \
124                     loading it as version {FORMAT_VERSION}: {e}"
125                ),
126            }
127        }
128
129        log::info!(
130            "index loaded: levels={levels} metric={metric:?} items={total_items} \
131             memory_limit={memory_limit_bytes:?}"
132        );
133
134        Ok(Index {
135            store,
136            metric,
137            is_normalized,
138            levels,
139            root,
140            nodes,
141            queries,
142            next_query_id: AtomicUsize::new(next_query_id),
143            memory_limit_bytes,
144            accepting: AtomicBool::new(true),
145            leaf_locks: DashMap::new(),
146            next_item_id: Mutex::new(next_item_id),
147            total_items: Mutex::new(total_items),
148        })
149    }
150
151    /// Builds the node and query caches, uncapped when `memory_limit_bytes`
152    /// is `None`. A query evicted from its cache is saved to `store`.
153    fn build_caches(
154        memory_limit_bytes: Option<usize>,
155        store: ReadableWritableListableStorage,
156    ) -> (NodeCache, QueryCache) {
157        // One node cache for every level, so one capacity caps the whole tree.
158        let mut nodes_builder = Cache::builder();
159        let mut queries_builder = Cache::builder();
160
161        // With a limit, split it between the caches, sizing each entry in bytes
162        if let Some(limit) = memory_limit_bytes {
163            let query_capacity = (limit as f64 * QUERY_CACHE_MEMORY_FRACTION) as u64;
164            if query_capacity == 0 {
165                log::warn!(
166                    "query cache capacity is 0 bytes for memory_limit_bytes={limit}; \
167                     every query will be evicted to disk immediately and reloaded \
168                     on next use"
169                );
170            }
171            let node_capacity = limit as u64 - query_capacity;
172            nodes_builder = nodes_builder
173                .max_capacity(node_capacity)
174                .weigher(|_key: &(usize, u32), node: &Arc<Node>| node.resident_bytes() as u32);
175            queries_builder = queries_builder.max_capacity(query_capacity).weigher(
176                |_key: &usize, state: &Arc<Mutex<QueryState>>| {
177                    let state = state.lock().unwrap();
178                    (state.query.len() * size_of::<f32>()
179                        + state.tree_pq.len() * size_of::<HeapEntry>()
180                        + state.items.len() * size_of::<(NotNan<f32>, u32)>()
181                        + state.exclude.len() * size_of::<u32>()) as u32
182                },
183            );
184        }
185
186        // moka also notifies when an entry is replaced or removed by hand, so
187        // both listeners ask whether the entry really left to free memory.
188        let nodes = nodes_builder
189            .eviction_listener(|key, node: Arc<Node>, cause| {
190                if cause.was_evicted() {
191                    log::debug!(
192                        "evicted lvl={} node={} cause={cause:?} bytes={}",
193                        key.0,
194                        key.1,
195                        node.resident_bytes()
196                    );
197                }
198            })
199            .build();
200        // An evicted query goes to disk, so it can still be resumed.
201        let queries = queries_builder
202            .eviction_listener(move |query_id, state_arc: Arc<Mutex<QueryState>>, cause| {
203                if !cause.was_evicted() {
204                    return;
205                }
206                log::debug!("evicting query_id={query_id} cause={cause:?}");
207                let state = state_arc.lock().unwrap();
208                // No caller is waiting on an eviction, so a failure here can only be logged.
209                if let Err(e) = persistence::persist_or_erase(&store, *query_id, &state) {
210                    log::error!("failed to persist evicted query_id={query_id}: {e}");
211                }
212            })
213            .build();
214
215        (nodes, queries)
216    }
217
218    /// Changes the memory limit without reloading; `None` removes the cap.
219    /// It rebuilds both caches, so it suits occasional reconfiguration, not
220    /// a call per query.
221    pub fn set_memory_limit_bytes(&mut self, memory_limit_bytes: Option<usize>) {
222        self.memory_limit_bytes = memory_limit_bytes;
223        let (new_nodes, new_queries) = Self::build_caches(memory_limit_bytes, self.store.clone());
224
225        // Move every cached entry over
226        for (key, value) in self.nodes.iter() {
227            new_nodes.insert(*key, value);
228        }
229        for (key, value) in self.queries.iter() {
230            new_queries.insert(*key, value);
231        }
232
233        // Evict now if the new limit is already exceeded
234        new_nodes.run_pending_tasks();
235        new_queries.run_pending_tasks();
236        self.nodes = new_nodes;
237        self.queries = new_queries;
238        log::info!(
239            "memory limit set to {memory_limit_bytes:?}, resident now {} bytes",
240            self.resident_bytes()
241        );
242    }
243
244    /// Gets node `node_id` at 0-based level `lvl`, reading it from disk on a cache miss.
245    /// `node_id` may not exist on disk (only nodes that received a child
246    /// during build do); `embeddings`/`children` return `Ok(None)` for those.
247    fn node_at(&self, lvl: usize, node_id: u32) -> Arc<Node> {
248        let is_leaf = lvl + 1 == self.levels as usize;
249        let child_key = if is_leaf { "item_ids" } else { "node_ids" };
250
251        // Internal Node
252        if !is_leaf {
253            log::trace!(
254                "lvl={lvl} node={node_id} cache {}",
255                if self.nodes.contains_key(&(lvl, node_id)) {
256                    "hit"
257                } else {
258                    "miss, reading from disk"
259                }
260            );
261            return self.nodes.get_with((lvl, node_id), || {
262                Arc::new(Node::new(
263                    self.store.clone().readable_listable(),
264                    format!("/lvl_{}/node_{node_id}", lvl + 1),
265                    child_key.to_string(),
266                ))
267            });
268        }
269
270        // Leaf node already cached
271        if let Some(hit) = self.nodes.get(&(lvl, node_id)) {
272            log::trace!("lvl={lvl} node={node_id} cache hit");
273            return hit;
274        }
275
276        // Acquire lock before reading from disk
277        let lock = self
278            .leaf_locks
279            .entry(node_id)
280            .or_insert_with(|| Arc::new(RwLock::new(())))
281            .clone();
282        let _guard = lock.read().unwrap();
283
284        // Someone else may have populated it while we waited for the lock.
285        if let Some(hit) = self.nodes.get(&(lvl, node_id)) {
286            log::trace!("lvl={lvl} node={node_id} cache hit after waiting for its lock");
287            return hit;
288        }
289        log::trace!("lvl={lvl} node={node_id} cache miss, reading from disk");
290        let node = Arc::new(Node::new(
291            self.store.clone().readable_listable(),
292            format!("/lvl_{}/node_{node_id}", lvl + 1),
293            child_key.to_string(),
294        ));
295        // Read the node's arrays into cache
296        let _ = node.embeddings();
297        let _ = node.children();
298        // Store before the lock drops, so insert's invalidate() can never
299        // land in between and miss this entry.
300        self.nodes.insert((lvl, node_id), node.clone());
301        node
302    }
303
304    /// Returns how many bytes the node cache currently holds.
305    fn resident_bytes(&self) -> usize {
306        self.nodes.weighted_size() as usize
307    }
308
309    /// Stops new searches and saves every open query to disk, so a later
310    /// `Index` can resume it; a query with nothing left to return is erased
311    /// instead. Safe to call more than once. `Err` if any query fails to save.
312    pub fn shutdown(&self) -> Result<()> {
313        self.accepting.store(false, Ordering::SeqCst);
314        let mut total = 0;
315        let mut failed = 0;
316        // Count a failed save and go on, so every open query is tried
317        for (query_id, state_arc) in self.queries.iter() {
318            total += 1;
319            let state = state_arc.lock().unwrap();
320            if let Err(e) = persistence::persist_or_erase(&self.store, *query_id, &state) {
321                log::error!("failed to persist query_id={query_id} on shutdown: {e}");
322                failed += 1;
323            }
324        }
325        log::info!("shutdown: saved {} of {total} open queries", total - failed);
326        if failed > 0 {
327            return Err(EcpError::Store(format!(
328                "failed to persist {failed} of {total} open queries"
329            )));
330        }
331        Ok(())
332    }
333
334    /// Erases every query saved to disk before `cutoff_unix_secs` (seconds
335    /// since the Unix epoch). Returns how many were erased.
336    pub fn cleanup_persisted_queries_older_than(&self, cutoff_unix_secs: u64) -> Result<usize> {
337        let erased = persistence::cleanup_older_than(&self.store, cutoff_unix_secs)?;
338        log::info!("cleanup: erased {erased} persisted queries older than {cutoff_unix_secs}");
339        Ok(erased)
340    }
341
342    /// Adds each row of `embeddings` to its nearest leaf, converted to the
343    /// index's dtype, and returns the new ids in row order. Leaves only grow;
344    /// the tree is never rebalanced. Safe alongside searches and inserts.
345    pub fn insert(&self, embeddings: Array2<f32>) -> Result<std::ops::Range<u32>> {
346        if embeddings.nrows() == 0 {
347            let next = *self.next_item_id.lock().unwrap();
348            return Ok(next..next);
349        }
350
351        // Read the dtype and chunk shape from the root; every node shares them
352        let root_array = Array::open(self.store.clone(), "/index_root/embeddings")
353            .store_err("failed to open index_root/embeddings")?;
354        let embedding_dtype = dtype_of_array(&root_array, "index_root/embeddings")?;
355        let chunk_shape: Vec<u64> = root_array
356            .chunk_shape_usize(&[0, 0])
357            .store_err("failed to read index_root/embeddings chunk shape")?
358            .into_iter()
359            .map(|v| v as u64)
360            .collect();
361
362        // A small cache for this call, since it only holds the nodes on the way
363        // down to each leaf
364        let cache_capacity_bytes = self.memory_limit_bytes.unwrap_or(usize::MAX) as u64 / 20;
365        let node_cache = BuildNodeCache::new(cache_capacity_bytes);
366        let config = BuildConfig {
367            store: &self.store,
368            node_cache: &node_cache,
369            target_level: self.levels,
370            total_levels: self.levels,
371            metric: self.metric,
372            is_normalized: self.is_normalized,
373            chunk_shape: &chunk_shape,
374            embedding_dtype,
375            leaf_locks: Some(&self.leaf_locks),
376        };
377
378        // Take the next ids and save the new next_item_id under one lock, so
379        // concurrent inserts can't save it out of order. Done before the
380        // append, so a crash skips ids instead of reusing them.
381        let start = {
382            let mut next = self.next_item_id.lock().unwrap();
383            let start = *next;
384            *next += embeddings.nrows() as u32;
385            write_info_u32(&self.store, "next_item_id", *next)?;
386            start
387        };
388        let end = start + embeddings.nrows() as u32;
389        let ids = Array1::from_iter(start..end);
390
391        // Append, then drop the stale cached copy of every leaf appended to
392        let touched = add_data(&config, &self.root, &embeddings, &ids)?;
393        for (level, node_id) in touched {
394            // node_at's cache key is 0-based; on-disk level is 1-based.
395            let key = (level as usize - 1, node_id);
396            log::debug!("dropping the cached copy of lvl={} node={node_id}", key.0);
397            self.nodes.invalidate(&key);
398        }
399
400        // Counted only once the vectors are on disk
401        {
402            let mut total = self.total_items.lock().unwrap();
403            *total += embeddings.nrows() as u32;
404            write_info_u32(&self.store, "total_items", *total)?;
405        }
406
407        log::info!(
408            "insert: added {} rows, ids {start}..{end}",
409            embeddings.nrows()
410        );
411        Ok(start..end)
412    }
413}
414
415#[cfg(test)]
416#[path = "utests/fixtures.rs"]
417mod fixtures;
418
419#[cfg(test)]
420#[path = "utests/index.rs"]
421mod tests;
ecp_core/search/index.rs

ecp_core/search/
index.rs
