ecp_core/build/
writer.rs

1use std::sync::Arc;
2
3use half::f16;
4use ndarray::{Array1, Array2};
5use zarrs::array::codec::ZstdCodec;
6use zarrs::array::data_type::{bool, float16, float32, int8, string, uint8, uint32};
7use zarrs::array::{
8    Array, ArrayBuilder, ArrayCreateError, ArraySubset, BytesToBytesCodecTraits, DataType,
9    FillValue, FillValueMetadata,
10};
11use zarrs::group::GroupBuilder;
12use zarrs::storage::{ReadableWritableListableStorage, ReadableWritableListableStorageTraits};
13
14use crate::dtype::EmbeddingDtype;
15use crate::error::{EcpError, Result, ResultExt};
16use crate::format::FORMAT_VERSION;
17use crate::metric::Metric;
18
19/// Compresses every written array with zstd at level 3. Chunks are sized for
20/// I/O rather than a node's data, so most are largely empty padding, which
21/// zstd shrinks to almost nothing on disk.
22pub(super) fn compressor() -> Vec<Arc<dyn BytesToBytesCodecTraits>> {
23    vec![Arc::new(ZstdCodec::new(3, false))]
24}
25
26/// Registers `path` as a Zarr group, so a generic Zarr reader (such as
27/// zarr-python) can list and open it as a real node, not just a directory
28/// two array paths happen to share.
29pub(crate) fn build_group(store: &ReadableWritableListableStorage, path: &str) -> Result<()> {
30    GroupBuilder::default()
31        .build(store.clone(), path)
32        .store_err_with(|| format!("failed to build {path} group"))?
33        .store_metadata()
34        .store_err_with(|| format!("failed to store {path} group metadata"))
35}
36
37/// Creates the embeddings array at `path` with `shape`, stored as `dtype`,
38/// zero-filled and zstd-compressed, and returns it for writing.
39pub(crate) fn build_embeddings_array(
40    store: &ReadableWritableListableStorage,
41    path: &str,
42    shape: Vec<u64>,
43    chunk_shape: &[u64],
44    dtype: EmbeddingDtype,
45) -> Result<Array<dyn ReadableWritableListableStorageTraits>> {
46    let (data_type, fill_value): (DataType, FillValue) = match dtype {
47        EmbeddingDtype::F32 => (float32(), 0.0f32.into()),
48        EmbeddingDtype::F16 => (float16(), f16::from_f32(0.0).into()),
49        EmbeddingDtype::UInt8 => (uint8(), 0u8.into()),
50        EmbeddingDtype::Int8 => (int8(), 0i8.into()),
51    };
52    let mut builder = ArrayBuilder::new(shape, chunk_shape.to_vec(), data_type, fill_value);
53    builder.bytes_to_bytes_codecs(compressor());
54    let array = builder
55        .build(store.clone(), path)
56        .store_err_with(|| format!("failed to build {path} array"))?;
57    array
58        .store_metadata()
59        .store_err_with(|| format!("failed to store {path} metadata"))?;
60    Ok(array)
61}
62
63/// Writes `embeddings` into the `subset` region of `array`, converted to
64/// `dtype`. Integer dtypes truncate fractions and clamp out-of-range values.
65/// `context` names the array in error messages.
66pub(crate) fn store_embeddings_subset(
67    array: &Array<dyn ReadableWritableListableStorageTraits>,
68    subset: &ArraySubset,
69    embeddings: &Array2<f32>,
70    dtype: EmbeddingDtype,
71    context: &str,
72) -> Result<()> {
73    let context = format!("failed to store {context}");
74    match dtype {
75        EmbeddingDtype::F32 => array
76            .store_array_subset(subset, embeddings)
77            .store_err(&context),
78        EmbeddingDtype::F16 => array
79            .store_array_subset(subset, embeddings.mapv(f16::from_f32))
80            .store_err(&context),
81        EmbeddingDtype::UInt8 => array
82            .store_array_subset(subset, embeddings.mapv(|x| x as u8))
83            .store_err(&context),
84        EmbeddingDtype::Int8 => array
85            .store_array_subset(subset, embeddings.mapv(|x| x as i8))
86            .store_err(&context),
87    }
88}
89
90/// Appends `embeddings` (stored as `dtype`) to the array at `embeddings_path`
91/// and `ids` to the one at `ids_path`, creating both on the first call.
92/// `chunk_shape` only applies on creation.
93// The ids arrays' `&[a..b]` is intended, one range for their single dimension.
94#[allow(clippy::single_range_in_vec_init)]
95pub fn zarrs_append(
96    store: &ReadableWritableListableStorage,
97    embeddings_path: &str,
98    ids_path: &str,
99    embeddings: &Array2<f32>,
100    ids: &Array1<u32>,
101    chunk_shape: &[u64],
102    dtype: EmbeddingDtype,
103) -> Result<()> {
104    match Array::open(store.clone(), embeddings_path) {
105        // Grow both arrays, then write into the new rows
106        Ok(mut array) => {
107            let existing_vecs = array.shape()[0];
108            let dim = array.shape()[1];
109            let new_vecs = existing_vecs + embeddings.nrows() as u64;
110            array
111                .set_shape(vec![new_vecs, dim])
112                .store_err("failed to grow embeddings array")?;
113            array
114                .store_metadata()
115                .store_err("failed to store embeddings metadata")?;
116            let subset = ArraySubset::new_with_ranges(&[existing_vecs..new_vecs, 0..dim]);
117            store_embeddings_subset(&array, &subset, embeddings, dtype, embeddings_path)?;
118
119            let mut ids_array = Array::open(store.clone(), ids_path)
120                .store_err("ids array missing alongside embeddings")?;
121            let existing_ids = ids_array.shape()[0];
122            let new_ids = existing_ids + ids.len() as u64;
123            ids_array
124                .set_shape(vec![new_ids])
125                .store_err("failed to grow ids array")?;
126            ids_array
127                .store_metadata()
128                .store_err("failed to store ids metadata")?;
129            ids_array
130                .store_array_subset(&ArraySubset::new_with_ranges(&[existing_ids..new_ids]), ids)
131                .store_err("failed to append ids")?;
132        }
133        // On the first call, create both arrays
134        Err(ArrayCreateError::MissingMetadata) => {
135            let dim = embeddings.ncols() as u64;
136            let emb_shape = vec![embeddings.nrows() as u64, dim];
137            let subset = ArraySubset::new_with_ranges(&[0..emb_shape[0], 0..dim]);
138            let emb_array =
139                build_embeddings_array(store, embeddings_path, emb_shape, chunk_shape, dtype)?;
140            store_embeddings_subset(&emb_array, &subset, embeddings, dtype, embeddings_path)?;
141
142            // Same chunk-count as the embeddings array, so a given
143            // chunk index lines up across both.
144            let ids_shape = vec![ids.len() as u64];
145            let mut ids_builder =
146                ArrayBuilder::new(ids_shape.clone(), vec![chunk_shape[0]], uint32(), 0u32);
147            ids_builder.bytes_to_bytes_codecs(compressor());
148            let ids_array = ids_builder
149                .build(store.clone(), ids_path)
150                .store_err("failed to build ids array")?;
151            ids_array
152                .store_metadata()
153                .store_err("failed to store ids metadata")?;
154            ids_array
155                .store_array_subset(&ArraySubset::new_with_ranges(&[0..ids_shape[0]]), ids)
156                .store_err("failed to store ids")?;
157        }
158        Err(e) => {
159            return Err(EcpError::Store(format!(
160                "failed to open {embeddings_path}: {e}"
161            )));
162        }
163    }
164    Ok(())
165}
166
167/// Writes `info/format_version`, `info/levels`, `info/metric`, and `info/is_normalized`.
168pub fn write_index_info(
169    store: &ReadableWritableListableStorage,
170    levels: u32,
171    metric: Metric,
172    is_normalized: bool,
173) -> Result<()> {
174    write_info_u32(store, "format_version", FORMAT_VERSION)?;
175
176    // Zarr has no bare-scalar type; each of these is a rank-0 array.
177    let scalar_shape: Vec<u64> = vec![];
178
179    let field = ArrayBuilder::new(scalar_shape.clone(), scalar_shape.clone(), uint32(), 0u32)
180        .build(store.clone(), "/info/levels")
181        .store_err("failed to build info/levels array")?;
182    field
183        .store_metadata()
184        .store_err("failed to store info/levels metadata")?;
185    field
186        .store_chunk(&[], vec![levels])
187        .store_err("failed to store info/levels chunk")?;
188
189    let field = ArrayBuilder::new(scalar_shape.clone(), scalar_shape.clone(), string(), "")
190        .build(store.clone(), "/info/metric")
191        .store_err("failed to build info/metric array")?;
192    field
193        .store_metadata()
194        .store_err("failed to store info/metric metadata")?;
195    field
196        .store_chunk(&[], vec![metric.as_str().to_string()])
197        .store_err("failed to store info/metric chunk")?;
198
199    let field = ArrayBuilder::new(
200        scalar_shape.clone(),
201        scalar_shape,
202        bool(),
203        FillValueMetadata::Bool(false),
204    )
205    .build(store.clone(), "/info/is_normalized")
206    .store_err("failed to build info/is_normalized array")?;
207    field
208        .store_metadata()
209        .store_err("failed to store info/is_normalized metadata")?;
210    field
211        .store_chunk(&[], vec![is_normalized])
212        .store_err("failed to store info/is_normalized chunk")?;
213    Ok(())
214}
215
216/// Writes `info/{name}` as a `u32` scalar, replacing any existing value.
217pub fn write_info_u32(
218    store: &ReadableWritableListableStorage,
219    name: &str,
220    value: u32,
221) -> Result<()> {
222    let scalar_shape: Vec<u64> = vec![];
223    let path = format!("/info/{name}");
224
225    let field = ArrayBuilder::new(scalar_shape.clone(), scalar_shape, uint32(), 0u32)
226        .build(store.clone(), &path)
227        .store_err_with(|| format!("failed to build {path} array"))?;
228    field
229        .store_metadata()
230        .store_err_with(|| format!("failed to store {path} metadata"))?;
231    field
232        .store_chunk(&[], vec![value])
233        .store_err_with(|| format!("failed to store {path} chunk"))?;
234    Ok(())
235}
236
237/// Writes `root_embeddings`, the root node's representatives, to
238/// `index_root/embeddings` as `dtype`.
239pub fn write_index_root(
240    store: &ReadableWritableListableStorage,
241    root_embeddings: &Array2<f32>,
242    chunk_shape: &[u64],
243    dtype: EmbeddingDtype,
244) -> Result<()> {
245    let shape = vec![
246        root_embeddings.nrows() as u64,
247        root_embeddings.ncols() as u64,
248    ];
249    let subset = ArraySubset::new_with_ranges(&[0..shape[0], 0..shape[1]]);
250    let path = "/index_root/embeddings";
251    let array = build_embeddings_array(store, path, shape, chunk_shape, dtype)?;
252    store_embeddings_subset(&array, &subset, root_embeddings, dtype, path)
253}
254
255/// Appends `embeddings` and their `children` ids to the node at `group_path`,
256/// creating its arrays on the first call; the ids go in `child_key`
257/// (`node_ids` or `item_ids`). `border` is left at its fill value, never populated.
258pub fn append_node_batch(
259    store: &ReadableWritableListableStorage,
260    group_path: &str,
261    child_key: &str,
262    embeddings: &Array2<f32>,
263    children: &Array1<u32>,
264    chunk_shape: &[u64],
265    dtype: EmbeddingDtype,
266) -> Result<()> {
267    let embeddings_path = format!("{group_path}/embeddings");
268    let children_path = format!("{group_path}/{child_key}");
269    let is_new = matches!(
270        Array::open(store.clone(), &embeddings_path),
271        Err(ArrayCreateError::MissingMetadata)
272    );
273
274    zarrs_append(
275        store,
276        &embeddings_path,
277        &children_path,
278        embeddings,
279        children,
280        chunk_shape,
281        dtype,
282    )?;
283
284    if is_new {
285        build_group(store, group_path)?;
286        let border_shape = vec![2u64];
287        let border_array = ArrayBuilder::new(border_shape.clone(), border_shape, float32(), 0.0f32)
288            .build(store.clone(), &format!("{group_path}/border"))
289            .store_err("failed to build border array")?;
290        border_array
291            .store_metadata()
292            .store_err("failed to store border metadata")?;
293    }
294    Ok(())
295}
296
297#[cfg(test)]
298#[path = "utests/writer.rs"]
299mod tests;
ecp_core/build/writer.rs

ecp_core/build/
writer.rs
