ecp_core/build/
source.rs

1use std::path::Path;
2use std::sync::Arc;
3
4use ndarray::{Array2, s};
5use rust_hdf5::{DatatypeMessage, H5Dataset, H5File, Hdf5Error};
6use zarrs::array::{Array, ArrayCreateError, ArraySubset};
7use zarrs::filesystem::FilesystemStore;
8use zarrs::storage::ReadableListableStorage;
9
10use crate::dtype::{EmbeddingDtype, dtype_of_array, read_subset_as_f32};
11use crate::error::{EcpError, Result, ResultExt};
12
13/// Turns an HDF5 failure into an [`EcpError`], `context` naming what was
14/// being read. A missing dataset maps to [`EcpError::NotFound`]; anything
15/// else maps to [`EcpError::Store`].
16fn hdf5_err(e: Hdf5Error, context: &str) -> EcpError {
17    if matches!(&e, Hdf5Error::NotFound(_)) {
18        EcpError::NotFound(format!("{context}: {e}"))
19    } else {
20        EcpError::Store(format!("{context}: {e}"))
21    }
22}
23
24/// Opens the zarr array at `path` in `store`. A missing array maps to
25/// [`EcpError::NotFound`]; any other failure maps to [`EcpError::Store`].
26fn open_zarr_array(
27    store: &ReadableListableStorage,
28    path: &str,
29) -> Result<Array<dyn zarrs::storage::ReadableListableStorageTraits>> {
30    Array::open(store.clone(), path).map_err(|e| match e {
31        ArrayCreateError::MissingMetadata => EcpError::NotFound(format!("no array at {path}")),
32        other => EcpError::Store(format!("failed to open {path}: {other}")),
33    })
34}
35
36/// The embeddings a build reads, from an `.h5` or `.zarr` file or an
37/// `Array2<f32>` already in memory. Files may store f32, f16, uint8 or int8,
38/// but `read_vecs` always returns f32 and only reads the rows asked for.
39pub enum EmbeddingsSource {
40    Hdf5(H5Dataset),
41    Zarr {
42        store: ReadableListableStorage,
43        path: String,
44    },
45    /// Already in memory; `read_vecs` just copies a slice.
46    Memory(Array2<f32>),
47}
48
49/// Maps an HDF5 dataset's dtype to its `EmbeddingDtype`.
50fn hdf5_dtype(dataset: &H5Dataset) -> Result<EmbeddingDtype> {
51    match dataset
52        .datatype()
53        .map_err(|e| hdf5_err(e, "failed to read HDF5 dataset datatype"))?
54    {
55        DatatypeMessage::FloatingPoint { size: 2, .. } => Ok(EmbeddingDtype::F16),
56        DatatypeMessage::FloatingPoint { size: 4, .. } => Ok(EmbeddingDtype::F32),
57        DatatypeMessage::FixedPoint {
58            size: 1,
59            signed: false,
60            ..
61        } => Ok(EmbeddingDtype::UInt8),
62        DatatypeMessage::FixedPoint {
63            size: 1,
64            signed: true,
65            ..
66        } => Ok(EmbeddingDtype::Int8),
67        other => Err(EcpError::InvalidInput(format!(
68            "unsupported embeddings dtype: {other:?} (use float32, float16, uint8 or int8)"
69        ))),
70    }
71}
72
73impl EmbeddingsSource {
74    /// Opens the dataset (`.h5`) or array (`.zarr`) called `name` in the file
75    /// at `path`, picking the format from `path`'s extension.
76    pub fn open(path: &Path, name: &str) -> Result<Self> {
77        if !path.exists() {
78            return Err(EcpError::NotFound(format!(
79                "{} does not exist",
80                path.display()
81            )));
82        }
83        match path.extension().and_then(|ext| ext.to_str()) {
84            Some("h5") => {
85                let file =
86                    H5File::open(path).map_err(|e| hdf5_err(e, "failed to open HDF5 file"))?;
87                let dataset = file
88                    .dataset(name)
89                    .map_err(|e| hdf5_err(e, "failed to open HDF5 dataset"))?;
90                Ok(EmbeddingsSource::Hdf5(dataset))
91            }
92            Some("zarr") => {
93                let store: ReadableListableStorage =
94                    Arc::new(FilesystemStore::new(path).store_err("failed to open zarr store")?);
95                Ok(EmbeddingsSource::Zarr {
96                    store,
97                    path: format!("/{name}"),
98                })
99            }
100            other => Err(EcpError::InvalidInput(format!(
101                "unsupported embeddings file format: {other:?} (use \"h5\" or \"zarr\")"
102            ))),
103        }
104    }
105
106    /// Creates an embeddings source from the zarr array at `path` in `store`.
107    pub fn from_zarr(store: ReadableListableStorage, path: String) -> Self {
108        EmbeddingsSource::Zarr { store, path }
109    }
110
111    /// Returns the source's size as `(total_items, dim)`.
112    pub fn shape(&self) -> Result<(usize, usize)> {
113        Ok(match self {
114            EmbeddingsSource::Hdf5(dataset) => {
115                let shape = dataset.shape();
116                (shape[0], shape[1])
117            }
118            EmbeddingsSource::Zarr { store, path } => {
119                let array = open_zarr_array(store, path)?;
120                let shape = array.shape();
121                (shape[0] as usize, shape[1] as usize)
122            }
123            EmbeddingsSource::Memory(embeddings) => (embeddings.nrows(), embeddings.ncols()),
124        })
125    }
126
127    /// Returns how many vectors one on-disk chunk holds, or `fallback` if the
128    /// source isn't chunked.
129    pub fn natural_chunk_vecs(&self, fallback: usize) -> Result<usize> {
130        Ok(match self {
131            EmbeddingsSource::Hdf5(dataset) => {
132                dataset.chunk_dims().map(|dims| dims[0]).unwrap_or_else(|| {
133                    log::warn!(
134                        "HDF5 dataset has no chunking; using the fallback batch size of \
135                     {fallback} rows"
136                    );
137                    fallback
138                })
139            }
140            EmbeddingsSource::Zarr { store, path } => {
141                let array = open_zarr_array(store, path)?;
142                array
143                    .chunk_shape_usize(&[0, 0])
144                    .store_err("failed to read zarr chunk shape")?[0]
145            }
146            EmbeddingsSource::Memory(_) => fallback,
147        })
148    }
149
150    /// Picks a read batch size that covers whole chunks. Rounds `memory_floor_vecs`,
151    /// the batch size the memory budget allows, up to whole chunks, so no chunk is
152    /// split and decoded twice. `fallback` works as in `natural_chunk_vecs`.
153    pub fn chunk_aligned_batch_vecs(
154        &self,
155        memory_floor_vecs: usize,
156        fallback: usize,
157    ) -> Result<usize> {
158        let chunk_vecs = self.natural_chunk_vecs(fallback)?.max(1);
159        Ok(memory_floor_vecs.max(1).div_ceil(chunk_vecs) * chunk_vecs)
160    }
161
162    /// Returns the dtype the embeddings are stored as, always `F32` for `Memory`.
163    pub fn native_dtype(&self) -> Result<EmbeddingDtype> {
164        match self {
165            EmbeddingsSource::Hdf5(dataset) => hdf5_dtype(dataset),
166            EmbeddingsSource::Zarr { store, path } => {
167                let array = open_zarr_array(store, path)?;
168                dtype_of_array(&array, path)
169            }
170            EmbeddingsSource::Memory(_) => Ok(EmbeddingDtype::F32),
171        }
172    }
173
174    /// Reads vectors `start..end` as f32.
175    pub fn read_vecs(&self, start: usize, end: usize) -> Result<Array2<f32>> {
176        match self {
177            EmbeddingsSource::Hdf5(dataset) => {
178                let dim = dataset.shape()[1];
179                let rows = end - start;
180                // f16 widens to f32 inside the HDF5 crate, but it refuses to
181                // read an integer dataset as f32, so integers are read at
182                // their own width and widened here.
183                let flat: Vec<f32> = match hdf5_dtype(dataset)? {
184                    EmbeddingDtype::F16 | EmbeddingDtype::F32 => dataset
185                        .read_numeric_slice_as::<f32>(&[start, 0], &[rows, dim])
186                        .map_err(|e| hdf5_err(e, "failed to read HDF5 vec range"))?,
187                    EmbeddingDtype::UInt8 => dataset
188                        .read_numeric_slice_as::<u8>(&[start, 0], &[rows, dim])
189                        .map_err(|e| hdf5_err(e, "failed to read HDF5 vec range"))?
190                        .into_iter()
191                        .map(|v| v as f32)
192                        .collect(),
193                    EmbeddingDtype::Int8 => dataset
194                        .read_numeric_slice_as::<i8>(&[start, 0], &[rows, dim])
195                        .map_err(|e| hdf5_err(e, "failed to read HDF5 vec range"))?
196                        .into_iter()
197                        .map(|v| v as f32)
198                        .collect(),
199                };
200                Ok(Array2::from_shape_vec((rows, dim), flat)
201                    .expect("HDF5 vec range didn't match its declared shape"))
202            }
203            EmbeddingsSource::Zarr { store, path } => {
204                let array = open_zarr_array(store, path)?;
205                let dim = array.shape()[1];
206                let subset = ArraySubset::new_with_ranges(&[start as u64..end as u64, 0..dim]);
207                read_subset_as_f32(&array, &subset, path)
208            }
209            EmbeddingsSource::Memory(embeddings) => {
210                Ok(embeddings.slice(s![start..end, ..]).to_owned())
211            }
212        }
213    }
214}
215
216#[cfg(test)]
217#[path = "utests/source.rs"]
218mod tests;
ecp_core/build/source.rs

ecp_core/build/
source.rs
