ecp_core/build/
tree.rs

1use std::sync::{Arc, RwLock};
2
3use dashmap::DashMap;
4use moka::sync::Cache;
5use ndarray::{Array1, Array2, Axis, s};
6use rayon::prelude::*;
7use zarrs::storage::{ReadableListableStorage, ReadableWritableListableStorage};
8
9use crate::build::assign::determine_node_assignments;
10use crate::build::builder::TRACKED_MEMORY_FRACTION;
11use crate::build::source::EmbeddingsSource;
12use crate::build::writer::{append_node_batch, build_group};
13use crate::dtype::EmbeddingDtype;
14use crate::error::{EcpError, Result};
15use crate::metric::Metric;
16use crate::search::Node;
17
18/// A cached node's `(representatives, children)`.
19type CachedNode = Arc<(Array2<f32>, Array1<u32>)>;
20
21/// Caches internal nodes' `(representatives, children)` during one build or insert,
22/// so routing each batch down the tree doesn't re-read them from disk.
23pub(crate) struct NodeCache {
24    /// Keyed by group path, such as `/lvl_1/node_3`.
25    cache: Cache<String, CachedNode>,
26}
27
28impl NodeCache {
29    /// Returns the bytes an entry takes, which the cache counts against its capacity.
30    fn entry_bytes(entry: &(Array2<f32>, Array1<u32>)) -> u32 {
31        (entry.0.len() * size_of::<f32>() + entry.1.len() * size_of::<u32>()) as u32
32    }
33
34    /// Creates an empty cache holding at most `max_capacity_bytes`.
35    pub(crate) fn new(max_capacity_bytes: u64) -> Self {
36        NodeCache {
37            cache: Cache::builder()
38                .max_capacity(max_capacity_bytes)
39                .weigher(|_, v: &CachedNode| Self::entry_bytes(v))
40                .build(),
41        }
42    }
43
44    /// Gets a node's representatives and children, reading from disk only on a cache miss.
45    fn get_or_read(
46        &self,
47        store: &ReadableWritableListableStorage,
48        group_path: &str,
49    ) -> Result<CachedNode> {
50        self.cache
51            .try_get_with(group_path.to_string(), || {
52                let read_store: ReadableListableStorage = store.clone().readable_listable();
53                let node = Node::new(read_store, group_path.to_string(), "node_ids".to_string());
54                // Either the build never wrote this node, or an insert routed
55                // into a branch the build left empty.
56                let missing = || {
57                    EcpError::Corrupt(format!(
58                        "{group_path} has no embeddings or children on disk"
59                    ))
60                };
61                let representatives = node.embeddings()?.ok_or_else(missing)?.clone();
62                let child_ids = node.children()?.ok_or_else(missing)?.clone();
63                Ok::<CachedNode, EcpError>(Arc::new((representatives, child_ids)))
64            })
65            .map_err(|e: Arc<EcpError>| (*e).clone())
66    }
67}
68
69/// The settings `add_data` passes unchanged down the tree.
70pub(crate) struct BuildConfig<'a> {
71    pub(crate) store: &'a ReadableWritableListableStorage,
72    pub(crate) node_cache: &'a NodeCache,
73    /// The 1-based level written to; levels above it are only routed through.
74    pub(crate) target_level: u32,
75    pub(crate) total_levels: u32,
76    pub(crate) metric: Metric,
77    pub(crate) is_normalized: bool,
78    pub(crate) chunk_shape: &'a [u64],
79    pub(crate) embedding_dtype: EmbeddingDtype,
80
81    /// Per-leaf locks, taken around each write. `None` when nothing else can
82    /// read the index yet, as during a build.
83    pub(crate) leaf_locks: Option<&'a DashMap<u32, Arc<RwLock<()>>>>,
84}
85
86/// Calls `process(i)` for each `i` in `0..count` in parallel and collects the results.
87/// Set `on_caller_thread` to run them one at a time on this thread instead. Do so when
88/// `process` may wait on a leaf lock, since that can deadlock rayon's pool.
89fn fan_out<F>(count: usize, on_caller_thread: bool, process: F) -> Result<Vec<(u32, u32)>>
90where
91    F: Fn(usize) -> Result<Vec<(u32, u32)>> + Sync + Send,
92{
93    let per_index: Vec<Vec<(u32, u32)>> = if on_caller_thread {
94        (0..count).map(process).collect::<Result<_>>()?
95    } else {
96        (0..count)
97            .into_par_iter()
98            .map(process)
99            .collect::<Result<_>>()?
100    };
101    Ok(per_index.into_iter().flatten().collect())
102}
103
104/// Routes a batch of vectors that belong under node `node_idx` at `level` down
105/// to `config.target_level`, and appends each to its nearest node there.
106/// Returns the on-disk `(level, node_idx)` of every node written to.
107fn route_batch_to_node(
108    config: &BuildConfig,
109    level: u32,
110    node_idx: u32,
111    data_embeddings: &Array2<f32>,
112    data_ids: &Array1<u32>,
113) -> Result<Vec<(u32, u32)>> {
114    let group_path = format!("/lvl_{level}/node_{node_idx}");
115
116    // At the target level, append the batch (under the leaf's lock, if any)
117    if level == config.target_level {
118        let child_key = if level == config.total_levels {
119            "item_ids"
120        } else {
121            "node_ids"
122        };
123        let write = || {
124            append_node_batch(
125                config.store,
126                &group_path,
127                child_key,
128                data_embeddings,
129                data_ids,
130                config.chunk_shape,
131                config.embedding_dtype,
132            )
133        };
134        match config.leaf_locks {
135            Some(locks) => {
136                let lock = locks
137                    .entry(node_idx)
138                    .or_insert_with(|| Arc::new(RwLock::new(())))
139                    .clone();
140                let _guard = lock.write().unwrap();
141                write()?;
142            }
143            None => write()?,
144        }
145        return Ok(vec![(level, node_idx)]);
146    }
147
148    // Above the target level, split the batch by nearest child and recurse into each
149    let entry = config.node_cache.get_or_read(config.store, &group_path)?;
150    let (representatives, child_ids) = (&entry.0, &entry.1);
151
152    let (offsets, assignment) = determine_node_assignments(
153        representatives,
154        data_embeddings,
155        config.metric,
156        config.is_normalized,
157    );
158
159    fan_out(
160        representatives.nrows(),
161        config.leaf_locks.is_some(),
162        |child| {
163            let start = offsets[child] as usize;
164            let end = offsets[child + 1] as usize;
165            if start == end {
166                return Ok(Vec::new());
167            }
168            let vec_indices: Vec<usize> = assignment
169                .slice(s![start..end])
170                .iter()
171                .map(|&i| i as usize)
172                .collect();
173            let child_embeddings = data_embeddings.select(Axis(0), &vec_indices);
174            let child_ids_batch = Array1::from_iter(vec_indices.iter().map(|&i| data_ids[i]));
175            route_batch_to_node(
176                config,
177                level + 1,
178                child_ids[child],
179                &child_embeddings,
180                &child_ids_batch,
181            )
182        },
183    )
184}
185
186/// Adds vectors to the tree at `config.target_level`. Each row of `data_embeddings`
187/// goes down from the root to its nearest node there, along with its id from
188/// `data_ids`. Returns the on-disk `(level, node_idx)` of every node written to.
189pub(crate) fn add_data(
190    config: &BuildConfig,
191    root_embeddings: &Array2<f32>,
192    data_embeddings: &Array2<f32>,
193    data_ids: &Array1<u32>,
194) -> Result<Vec<(u32, u32)>> {
195    let (offsets, assignment) = determine_node_assignments(
196        root_embeddings,
197        data_embeddings,
198        config.metric,
199        config.is_normalized,
200    );
201
202    fan_out(
203        root_embeddings.nrows(),
204        config.leaf_locks.is_some(),
205        |root_node| {
206            let start = offsets[root_node] as usize;
207            let end = offsets[root_node + 1] as usize;
208            if start == end {
209                return Ok(Vec::new());
210            }
211            let vec_indices: Vec<usize> = assignment
212                .slice(s![start..end])
213                .iter()
214                .map(|&i| i as usize)
215                .collect();
216            let node_embeddings = data_embeddings.select(Axis(0), &vec_indices);
217            let node_ids = Array1::from_iter(vec_indices.iter().map(|&i| data_ids[i]));
218            route_batch_to_node(config, 1, root_node as u32, &node_embeddings, &node_ids)
219        },
220    )
221}
222
223/// `build_tree`'s parameters, bundled to keep its own signature manageable.
224#[derive(Clone, Copy)]
225pub struct BuildTreeArgs<'a> {
226    pub store: &'a ReadableWritableListableStorage,
227    pub root_embeddings: &'a Array2<f32>,
228    /// Every representative, starting with the root's. The non-leaf levels are
229    /// built from these.
230    pub representatives: &'a EmbeddingsSource,
231    /// The items the leaves hold.
232    pub dataset: &'a EmbeddingsSource,
233    pub total_levels: u32,
234    pub metric: Metric,
235    pub is_normalized: bool,
236    /// The chunk size assumed for a source that isn't chunked.
237    pub fallback_batch_vecs: usize,
238    pub chunk_shape: &'a [u64],
239    pub embedding_dtype: EmbeddingDtype,
240    pub memory_limit_bytes: usize,
241}
242
243/// Builds the tree one level at a time from the top, appending each input
244/// vector to its nearest node on that level. Level `l` reads the first
245/// `ns^(l+1)` representatives, where `ns` is the number of root entries; the
246/// leaf level reads `dataset`.
247///
248/// Example, `root_embeddings.nrows() = ns = 100`, `total_levels = 3`,
249/// `representatives.shape().0 = R = 1_000_000`:
250///
251/// ```text
252/// target_level=1: reads first ns^2 = 10_000 of `representatives`
253/// target_level=2: reads first ns^3 = 1_000_000 of `representatives` (all of R)
254/// target_level=3 (== total_levels): reads all of `dataset`
255/// ```
256pub fn build_tree(args: &BuildTreeArgs) -> Result<()> {
257    let BuildTreeArgs {
258        store,
259        root_embeddings,
260        representatives,
261        dataset,
262        total_levels,
263        metric,
264        is_normalized,
265        fallback_batch_vecs,
266        chunk_shape,
267        embedding_dtype,
268        memory_limit_bytes,
269    } = *args;
270
271    let node_size = root_embeddings.nrows() as u64;
272    let tracked_budget = (memory_limit_bytes as f64 * TRACKED_MEMORY_FRACTION) as usize;
273    let bytes_per_vec = (root_embeddings.ncols() * size_of::<f32>()).max(1);
274
275    // The node cache gets what every non-leaf level needs, up to 3/4 of the
276    // budget; batches get the rest.
277    let total_cacheable_bytes: usize = (1..total_levels)
278        .map(|l| node_size.pow(l + 1) as usize * bytes_per_vec)
279        .sum();
280    let cache_capacity = total_cacheable_bytes.min(tracked_budget * 3 / 4);
281    let batch_share = tracked_budget.saturating_sub(cache_capacity);
282    if batch_share < bytes_per_vec {
283        log::warn!(
284            "build memory budget leaves {batch_share} bytes for batching against \
285             {bytes_per_vec} bytes per vector; batches will collapse to about \
286             one vector each, which will be slow"
287        );
288    }
289    let memory_floor_vecs = (batch_share / bytes_per_vec).max(1);
290    let node_cache = NodeCache::new(cache_capacity as u64);
291
292    for target_level in 1..=total_levels {
293        build_group(store, &format!("/lvl_{target_level}"))?;
294
295        // Non-leaf levels are built from representatives, the leaf level from the dataset
296        let source = if target_level == total_levels {
297            dataset
298        } else {
299            representatives
300        };
301        let (total_vec_count, _dim) = source.shape()?;
302        let vec_count = if target_level == total_levels {
303            total_vec_count
304        } else {
305            (node_size.pow(target_level + 1) as usize).min(total_vec_count)
306        };
307        let batch_vecs = source.chunk_aligned_batch_vecs(memory_floor_vecs, fallback_batch_vecs)?;
308
309        let config = BuildConfig {
310            store,
311            node_cache: &node_cache,
312            target_level,
313            total_levels,
314            metric,
315            is_normalized,
316            chunk_shape,
317            embedding_dtype,
318            leaf_locks: None,
319        };
320
321        // Row indices are item ids on the leaf level, and representative ids
322        // (the next level's node ids) above it
323
324        let mut start = 0;
325        while start < vec_count {
326            let end = (start + batch_vecs).min(vec_count);
327            log::debug!("target_level={target_level}: processing vecs {start}..{end}");
328            let batch_embeddings = source.read_vecs(start, end)?;
329            let batch_ids: Array1<u32> = (start as u32..end as u32).collect();
330            add_data(&config, root_embeddings, &batch_embeddings, &batch_ids)?;
331            start = end;
332        }
333    }
334    Ok(())
335}
336
337#[cfg(test)]
338#[path = "utests/tree.rs"]
339mod tests;
ecp_core/build/tree.rs

ecp_core/build/
tree.rs
