ecp_core/build/
representatives.rs

1use ndarray::{Array1, Array2, Axis};
2use rand::seq::index::sample;
3use zarrs::storage::ReadableWritableListableStorage;
4
5use crate::build::builder::TRACKED_MEMORY_FRACTION;
6use crate::build::source::EmbeddingsSource;
7use crate::build::writer::zarrs_append;
8use crate::dtype::EmbeddingDtype;
9use crate::error::Result;
10
11/// How `select_representatives` picks representatives, the items that every
12/// other item is clustered around.
13#[derive(Debug, Clone, Copy)]
14pub enum RepresentativeStrategy {
15    /// Items spaced `target_cluster_items` apart, starting at item 0.
16    Offset,
17    /// A uniform random sample, as many as `Offset` picks.
18    Random,
19}
20
21/// Picks which items become representatives and returns their ids, sorted ascending.
22/// Number of representatives = `total_items / target_cluster_items`, rounded up.
23pub fn select_representative_ids(
24    total_items: usize,
25    target_cluster_items: usize,
26    strategy: RepresentativeStrategy,
27) -> Array1<u32> {
28    match strategy {
29        RepresentativeStrategy::Offset => (0..total_items as u32)
30            .step_by(target_cluster_items)
31            .collect(),
32        RepresentativeStrategy::Random => {
33            let total_clusters = total_items.div_ceil(target_cluster_items);
34            let mut ids: Vec<u32> = sample(&mut rand::rng(), total_items, total_clusters)
35                .into_iter()
36                .map(|i| i as u32)
37                .collect();
38            ids.sort_unstable();
39            Array1::from_vec(ids)
40        }
41    }
42}
43
44/// Checks whether `count` f32 embeddings of length `dim` fit under the memory limit.
45pub fn fits_in_memory(count: usize, dim: usize, memory_limit_bytes: usize) -> bool {
46    count.saturating_mul(dim).saturating_mul(size_of::<f32>()) <= memory_limit_bytes
47}
48
49/// Whether `build` reads the representatives from memory, when they were
50/// small enough to keep, or from disk.
51pub enum Representatives {
52    InMemory {
53        embeddings: Array2<f32>,
54        ids: Array1<u32>,
55    },
56    PersistedOnly,
57}
58
59/// Saves the chosen representatives into the index. Copies the rows of
60/// `source` listed in `selected_ids`, which must be sorted ascending. Reads
61/// whole chunks and skips any chunk with no selected row.
62pub fn collect_representatives(
63    store: &ReadableWritableListableStorage,
64    source: &EmbeddingsSource,
65    selected_ids: &Array1<u32>,
66    fallback_batch_vecs: usize,
67    memory_limit_bytes: usize,
68    chunk_shape: &[u64],
69    dtype: EmbeddingDtype,
70) -> Result<()> {
71    let (total_items, dim) = source.shape()?;
72    let tracked_budget = (memory_limit_bytes as f64 * TRACKED_MEMORY_FRACTION) as usize;
73    let bytes_per_vec = (dim * size_of::<f32>()).max(1);
74    let memory_floor_vecs = (tracked_budget / bytes_per_vec).max(1);
75    let batch_vecs = source.chunk_aligned_batch_vecs(memory_floor_vecs, fallback_batch_vecs)?;
76    let selected: Vec<u32> = selected_ids.to_vec();
77
78    let mut start = 0;
79    while start < total_items {
80        let end = (start + batch_vecs).min(total_items);
81
82        // Selected ids in this batch
83        let first = selected.partition_point(|&id| (id as usize) < start);
84        let in_range = &selected[first..];
85        let matched_ids: Vec<u32> = in_range
86            .iter()
87            .take_while(|&&id| (id as usize) < end)
88            .copied()
89            .collect();
90        if matched_ids.is_empty() {
91            start = end;
92            continue;
93        }
94        log::debug!(
95            "processing batch vecs {start}..{end} ({} matched representatives)",
96            matched_ids.len()
97        );
98
99        // Read the batch and append the selected rows
100        let batch = source.read_vecs(start, end)?;
101        let matched_vec_indices: Vec<usize> =
102            matched_ids.iter().map(|&id| id as usize - start).collect();
103        let matched_embeddings = batch.select(Axis(0), &matched_vec_indices);
104        let matched_ids_array = Array1::from_vec(matched_ids);
105        zarrs_append(
106            store,
107            "/rep_embeddings",
108            "/rep_item_ids",
109            &matched_embeddings,
110            &matched_ids_array,
111            chunk_shape,
112            dtype,
113        )?;
114
115        start = end;
116    }
117    Ok(())
118}
119
120#[cfg(test)]
121#[path = "utests/representatives.rs"]
122mod tests;
ecp_core/build/representatives.rs

ecp_core/build/
representatives.rs
