ecp_core/build/
builder.rs

1use std::path::Path;
2use std::sync::Arc;
3
4use ndarray::{Array1, Array2, s};
5use zarrs::filesystem::FilesystemStore;
6use zarrs::storage::ReadableWritableListableStorage;
7
8use crate::build::representatives::{
9    RepresentativeStrategy, Representatives, collect_representatives, fits_in_memory,
10    select_representative_ids,
11};
12use crate::build::source::EmbeddingsSource;
13use crate::build::tree::{BuildTreeArgs, build_tree};
14use crate::build::writer::{
15    build_group, write_index_info, write_index_root, write_info_u32, zarrs_append,
16};
17use crate::dtype::EmbeddingDtype;
18use crate::error::{EcpError, Result, ResultExt};
19use crate::metric::Metric;
20
21/// Picks the dtype the index stores embeddings as, which is `embedding_dtype`
22/// if the caller set one, else `native`, the source file's own dtype. Warns
23/// when `embedding_dtype` can't hold every value `native` can.
24fn resolve_dtype(
25    embedding_dtype: Option<EmbeddingDtype>,
26    native: EmbeddingDtype,
27) -> EmbeddingDtype {
28    let resolved = embedding_dtype.unwrap_or(native);
29    if resolved.narrows(native) {
30        log::warn!(
31            "writing embeddings as {resolved:?} narrows the source's {native:?} values and loses information; \
32             integer targets additionally clamp out-of-range values and truncate fractions"
33        );
34    }
35    resolved
36}
37
38/// Default chunk size for the representative arrays, in bytes.
39pub const DEFAULT_REP_CHUNK_BYTES: usize = 8 * 1024 * 1024;
40
41/// Default chunk size for the tree nodes, in bytes.
42pub const DEFAULT_NODE_CHUNK_BYTES: usize = 512 * 1024;
43
44/// Share of the memory limit a build spends on each batch of vectors and on
45/// its node cache. The rest is left for memory this crate doesn't track.
46pub(crate) const TRACKED_MEMORY_FRACTION: f64 = 0.8;
47
48/// The chunk sizes a build writes, in bytes.
49#[derive(Debug, Clone, Copy)]
50pub struct ChunkSizes {
51    /// Chunk size for `/rep_embeddings` and `/rep_item_ids`.
52    pub rep_chunk_bytes: usize,
53    /// Chunk size for the root and every tree node.
54    pub node_chunk_bytes: usize,
55}
56
57impl Default for ChunkSizes {
58    fn default() -> Self {
59        ChunkSizes {
60            rep_chunk_bytes: DEFAULT_REP_CHUNK_BYTES,
61            node_chunk_bytes: DEFAULT_NODE_CHUNK_BYTES,
62        }
63    }
64}
65
66/// Returns how many vectors of length `dim` fit in `chunk_bytes` when stored
67/// as `dtype`. Errors for a dimension of 0, or if one vector is wider than
68/// `chunk_bytes`.
69pub fn chunk_rows(dim: usize, dtype: EmbeddingDtype, chunk_bytes: usize) -> Result<u64> {
70    if dim == 0 {
71        return Err(EcpError::InvalidInput(
72            "the embedding dimension must be at least 1".to_string(),
73        ));
74    }
75    let bytes_per_vec = dim.saturating_mul(dtype.bytes());
76    if bytes_per_vec > chunk_bytes {
77        return Err(EcpError::InvalidInput(format!(
78            "one vector of dimension {dim} stored as {dtype:?} takes {bytes_per_vec} bytes, \
79             wider than the {chunk_bytes}-byte chunk"
80        )));
81    }
82    Ok((chunk_bytes / bytes_per_vec) as u64)
83}
84
85/// Picks the tree's fan-out, the number of children each node (the root
86/// included) gets, so the bottom level has room for every cluster.
87/// Fan-out = `total_clusters^(1/levels)`, rounded up.
88fn node_size_for(total_clusters: usize, levels: u32) -> Result<usize> {
89    if levels == 0 {
90        return Err(EcpError::InvalidInput(
91            "levels must be at least 1".to_string(),
92        ));
93    }
94    Ok((total_clusters as f64).powf(1.0 / levels as f64).ceil() as usize)
95}
96
97/// Builds one index in three steps: `create`, then `select_representatives`
98/// or `select_representatives_custom`, then `build`. One build per `Builder`.
99pub struct Builder {
100    store: ReadableWritableListableStorage,
101    levels: u32,
102    metric: Metric,
103    is_normalized: bool,
104    memory_limit_bytes: usize,
105    embedding_dtype: Option<EmbeddingDtype>,
106    chunks: ChunkSizes,
107
108    // Set by either select_representatives method
109    rep_chunk_shape: Vec<u64>,
110    representatives: Option<Representatives>,
111    node_size: usize,
112    resolved_dtype: EmbeddingDtype,
113}
114
115impl Builder {
116    /// Starts a build in `store` instead of a path. Otherwise the same as [`Self::create`].
117    pub fn new(
118        store: ReadableWritableListableStorage,
119        levels: u32,
120        metric: Metric,
121        is_normalized: bool,
122        memory_limit_bytes: usize,
123        embedding_dtype: Option<EmbeddingDtype>,
124        chunks: ChunkSizes,
125    ) -> Result<Self> {
126        if levels == 0 {
127            return Err(EcpError::InvalidInput(
128                "levels must be at least 1".to_string(),
129            ));
130        }
131        build_group(&store, "/")?;
132        build_group(&store, "/info")?;
133        write_index_info(&store, levels, metric, is_normalized)?;
134        Ok(Builder {
135            store,
136            levels,
137            metric,
138            is_normalized,
139            memory_limit_bytes,
140            embedding_dtype,
141            chunks,
142            rep_chunk_shape: Vec::new(),
143            representatives: None,
144            node_size: 0,
145            resolved_dtype: EmbeddingDtype::F32,
146        })
147    }
148
149    /// Creates an index at `index_path`. `levels` is the number of node levels
150    /// below the root. Set `is_normalized` only if every embedding is unit-length,
151    /// and leave `embedding_dtype` as `None` to keep each source's own dtype.
152    pub fn create(
153        index_path: &Path,
154        levels: u32,
155        metric: Metric,
156        is_normalized: bool,
157        memory_limit_bytes: usize,
158        embedding_dtype: Option<EmbeddingDtype>,
159        chunks: ChunkSizes,
160    ) -> Result<Self> {
161        log::info!(
162            "creating index at {}: levels={levels} metric={metric:?} is_normalized={is_normalized}",
163            index_path.display()
164        );
165        let store: ReadableWritableListableStorage =
166            Arc::new(FilesystemStore::new(index_path).store_err("failed to create store")?);
167        Self::new(
168            store,
169            levels,
170            metric,
171            is_normalized,
172            memory_limit_bytes,
173            embedding_dtype,
174            chunks,
175        )
176    }
177
178    /// Picks the representatives the tree is built from, chosen by
179    /// `strategy`. `target_cluster_items` is the average cluster size it aims for, and
180    /// `fallback_batch_vecs` is the chunk size assumed when `source` isn't chunked.
181    pub fn select_representatives(
182        &mut self,
183        source: &EmbeddingsSource,
184        target_cluster_items: usize,
185        strategy: RepresentativeStrategy,
186        fallback_batch_vecs: usize,
187    ) -> Result<()> {
188        let (total_items, dim) = source.shape()?;
189        self.resolved_dtype = resolve_dtype(self.embedding_dtype, source.native_dtype()?);
190        log::info!(
191            "representatives will be stored as {:?}",
192            self.resolved_dtype
193        );
194        self.rep_chunk_shape = vec![
195            chunk_rows(dim, self.resolved_dtype, self.chunks.rep_chunk_bytes)?,
196            dim as u64,
197        ];
198
199        let selected_ids = select_representative_ids(total_items, target_cluster_items, strategy);
200        log::info!(
201            "selected {} representatives via {strategy:?} from {total_items} items (target_cluster_items={target_cluster_items})",
202            selected_ids.len()
203        );
204        self.node_size = node_size_for(selected_ids.len(), self.levels)?;
205        collect_representatives(
206            &self.store,
207            source,
208            &selected_ids,
209            fallback_batch_vecs,
210            self.memory_limit_bytes,
211            &self.rep_chunk_shape,
212            self.resolved_dtype,
213        )?;
214        self.representatives = Some(Representatives::PersistedOnly);
215        Ok(())
216    }
217
218    /// Uses caller-chosen representatives instead of picking them, such as ones
219    /// from an external clustering step. `ids[i]` is the id of row `i` of `embeddings`.
220    pub fn select_representatives_custom(
221        &mut self,
222        ids: Array1<u32>,
223        embeddings: Array2<f32>,
224    ) -> Result<()> {
225        if ids.len() != embeddings.nrows() {
226            return Err(EcpError::InvalidInput(format!(
227                "ids and embeddings must have the same length ({} ids, {} embeddings rows)",
228                ids.len(),
229                embeddings.nrows()
230            )));
231        }
232        let dim = embeddings.ncols();
233        self.resolved_dtype = resolve_dtype(self.embedding_dtype, EmbeddingDtype::F32);
234        log::info!(
235            "representatives will be stored as {:?}",
236            self.resolved_dtype
237        );
238        self.rep_chunk_shape = vec![
239            chunk_rows(dim, self.resolved_dtype, self.chunks.rep_chunk_bytes)?,
240            dim as u64,
241        ];
242        zarrs_append(
243            &self.store,
244            "/rep_embeddings",
245            "/rep_item_ids",
246            &embeddings,
247            &ids,
248            &self.rep_chunk_shape,
249            self.resolved_dtype,
250        )?;
251
252        self.node_size = node_size_for(ids.len(), self.levels)?;
253        self.representatives = Some(if fits_in_memory(ids.len(), dim, self.memory_limit_bytes) {
254            Representatives::InMemory { embeddings, ids }
255        } else {
256            Representatives::PersistedOnly
257        });
258        Ok(())
259    }
260
261    /// Builds the tree from `dataset`, whose rows get item ids 0, 1, 2, ...
262    /// in order. `fallback_batch_vecs` works as in `select_representatives`.
263    pub fn build(&mut self, dataset: &EmbeddingsSource, fallback_batch_vecs: usize) -> Result<()> {
264        log::info!(
265            "building tree: {} levels, metric={:?}",
266            self.levels,
267            self.metric
268        );
269        // Ids 0..total_items, so both counters start at total_items
270        let (total_items, dim) = dataset.shape()?;
271        write_info_u32(&self.store, "total_items", total_items as u32)?;
272        write_info_u32(&self.store, "next_item_id", total_items as u32)?;
273
274        let node_chunk_shape = vec![
275            chunk_rows(dim, self.resolved_dtype, self.chunks.node_chunk_bytes)?,
276            dim as u64,
277        ];
278        log::info!(
279            "node arrays chunk by {} vecs of {:?}",
280            node_chunk_shape[0],
281            self.resolved_dtype
282        );
283
284        let representatives = self.representatives.take().ok_or_else(|| {
285            EcpError::Usage("call select_representatives before build".to_string())
286        })?;
287
288        // Root holds the first node_size representatives
289        let (root_embeddings, representatives_source) = match representatives {
290            Representatives::InMemory { embeddings, .. } => {
291                let root = embeddings.slice(s![..self.node_size, ..]).to_owned();
292                (root, EmbeddingsSource::Memory(embeddings))
293            }
294            Representatives::PersistedOnly => {
295                let source = EmbeddingsSource::from_zarr(
296                    self.store.clone().readable_listable(),
297                    "/rep_embeddings".to_string(),
298                );
299                let root = source.read_vecs(0, self.node_size)?;
300                (root, source)
301            }
302        };
303
304        // Write root, then every level below it
305        build_group(&self.store, "/index_root")?;
306        write_index_root(
307            &self.store,
308            &root_embeddings,
309            &node_chunk_shape,
310            self.resolved_dtype,
311        )?;
312        build_tree(&BuildTreeArgs {
313            store: &self.store,
314            root_embeddings: &root_embeddings,
315            representatives: &representatives_source,
316            dataset,
317            total_levels: self.levels,
318            metric: self.metric,
319            is_normalized: self.is_normalized,
320            fallback_batch_vecs,
321            chunk_shape: &node_chunk_shape,
322            embedding_dtype: self.resolved_dtype,
323            memory_limit_bytes: self.memory_limit_bytes,
324        })
325    }
326}
327
328#[cfg(test)]
329#[path = "utests/builder.rs"]
330mod tests;
ecp_core/build/builder.rs

ecp_core/build/
builder.rs
