ecp_core/search/
node.rs

1use std::sync::OnceLock;
2
3use ndarray::{Array1, Array2};
4use zarrs::array::{Array, ArrayCreateError};
5use zarrs::storage::ReadableListableStorage;
6
7use crate::dtype::read_subset_as_f32;
8use crate::error::{EcpError, Result, ResultExt};
9
10/// One node's embeddings and child ids, read from the store on first use and
11/// then kept in memory.
12pub struct Node {
13    store: ReadableListableStorage,
14    /// The node's group in the store, such as `/lvl_2/node_7`.
15    pub group_path: String,
16    /// `node_ids` for an internal node, `item_ids` for a leaf.
17    pub child_key: String,
18    embeddings: OnceLock<Result<Option<Array2<f32>>>>,
19    children: OnceLock<Result<Option<Array1<u32>>>>,
20}
21
22/// Opens the array at path. Ok(None) if there is no array there,
23/// Err if the store failed
24fn open_if_present(
25    store: &ReadableListableStorage,
26    path: &str,
27) -> Result<Option<Array<dyn zarrs::storage::ReadableListableStorageTraits>>> {
28    match Array::open(store.clone(), path) {
29        Ok(array) => Ok(Some(array)),
30        Err(ArrayCreateError::MissingMetadata) => Ok(None),
31        Err(e) => Err(EcpError::Store(format!("failed to open {path}: {e}"))),
32    }
33}
34
35impl Node {
36    /// Creates a `Node` for the group at `group_path`. Nothing is read until
37    /// `embeddings` or `children` is called.
38    pub fn new(store: ReadableListableStorage, group_path: String, child_key: String) -> Self {
39        Node {
40            store,
41            group_path,
42            child_key,
43            embeddings: OnceLock::new(),
44            children: OnceLock::new(),
45        }
46    }
47
48    /// Returns the node's embeddings as f32, reading them on the first call.
49    /// Ok(None) if the node was never written, Err if the read failed.
50    /// Both results are cached.
51    pub fn embeddings(&self) -> Result<Option<&Array2<f32>>> {
52        let result = self.embeddings.get_or_init(|| {
53            let embeddings_path = format!("{}/embeddings", self.group_path);
54            match open_if_present(&self.store, &embeddings_path)? {
55                Some(array) => Ok(Some(read_subset_as_f32(
56                    &array,
57                    &array.subset_all(),
58                    &embeddings_path,
59                )?)),
60                None => Ok(None),
61            }
62        });
63        match result {
64            Ok(opt) => Ok(opt.as_ref()),
65            Err(e) => Err(e.clone()),
66        }
67    }
68
69    /// Returns the node's child ids as u32, reading them on the first call.
70    /// `Ok(None)` means the node was never written, `Err` if the read failed.
71    /// Both results are cached.
72    pub fn children(&self) -> Result<Option<&Array1<u32>>> {
73        let result = self.children.get_or_init(|| {
74            let ids_path = format!("{}/{}", self.group_path, self.child_key);
75            match open_if_present(&self.store, &ids_path)? {
76                Some(array) => {
77                    let ids = array
78                        .retrieve_array_subset::<Array1<u32>>(&array.subset_all())
79                        .store_err_with(|| format!("failed to retrieve {ids_path}"))?;
80                    Ok(Some(ids))
81                }
82                None => Ok(None),
83            }
84        });
85        match result {
86            Ok(opt) => Ok(opt.as_ref()),
87            Err(e) => Err(e.clone()),
88        }
89    }
90
91    /// Checks whether `embeddings` or `children` has read data. Always false
92    /// for a node missing on disk or one whose read failed.
93    pub fn is_loaded(&self) -> bool {
94        self.embeddings
95            .get()
96            .is_some_and(|r| matches!(r, Ok(Some(_))))
97            || self
98                .children
99                .get()
100                .is_some_and(|r| matches!(r, Ok(Some(_))))
101    }
102
103    /// Returns the bytes of embeddings and child ids this node holds in memory,
104    /// which the node cache counts against its limit.
105    pub fn resident_bytes(&self) -> usize {
106        let emb_bytes = self
107            .embeddings
108            .get()
109            .and_then(|r| r.as_ref().ok())
110            .and_then(Option::as_ref)
111            .map_or(0, |e| e.len() * size_of::<f32>());
112        let child_bytes = self
113            .children
114            .get()
115            .and_then(|r| r.as_ref().ok())
116            .and_then(Option::as_ref)
117            .map_or(0, |c| c.len() * size_of::<u32>());
118        emb_bytes + child_bytes
119    }
120}
121
122#[cfg(test)]
123#[path = "utests/node.rs"]
124mod tests;
ecp_core/search/node.rs

ecp_core/search/
node.rs
