diff --git a/diskann-benchmark-core/src/search/api.rs b/diskann-benchmark-core/src/search/api.rs index 811f02b0ee..64eb3744c5 100644 --- a/diskann-benchmark-core/src/search/api.rs +++ b/diskann-benchmark-core/src/search/api.rs @@ -517,7 +517,7 @@ where T: Search, { let mut lengths = Vec::with_capacity(range.len()); - let mut ids = Matrix::new(views::Init(T::Id::default), range.len(), num_ids.into()); + let mut ids = Matrix::try_new(views::Init(T::Id::default), range.len(), num_ids.into())?; let mut latencies = Vec::::with_capacity(range.len()); let mut outputs = Vec::::with_capacity(range.len()); diff --git a/diskann-benchmark/src/exhaustive/minmax.rs b/diskann-benchmark/src/exhaustive/minmax.rs index 444d4d6be9..7544223c74 100644 --- a/diskann-benchmark/src/exhaustive/minmax.rs +++ b/diskann-benchmark/src/exhaustive/minmax.rs @@ -389,7 +389,7 @@ mod imp { // The APIs below should correctly handle these variables. let output_dim = quantizer.output_dim(); let bytes = Data::::canonical_bytes(output_dim); - let mut data = diskann_utils::views::Matrix::new(0, input.nrows(), bytes); + let mut data = diskann_utils::views::Matrix::try_new(0, input.nrows(), bytes)?; // Compress the data. // diff --git a/diskann-benchmark/src/exhaustive/product.rs b/diskann-benchmark/src/exhaustive/product.rs index 43538117dd..676653a069 100644 --- a/diskann-benchmark/src/exhaustive/product.rs +++ b/diskann-benchmark/src/exhaustive/product.rs @@ -332,8 +332,11 @@ mod imp { quantizer: diskann_providers::model::pq::FixedChunkPQTable, progress: &ProgressBar, ) -> anyhow::Result { - let mut data = - diskann_utils::views::Matrix::new(0, input.nrows(), quantizer.get_num_chunks()); + let mut data = diskann_utils::views::Matrix::try_new( + 0, + input.nrows(), + quantizer.get_num_chunks(), + )?; // Compress the data. #[expect(clippy::disallowed_methods)] diff --git a/diskann-benchmark/src/exhaustive/spherical.rs b/diskann-benchmark/src/exhaustive/spherical.rs index 1bf3c09349..1fecf198a7 100644 --- a/diskann-benchmark/src/exhaustive/spherical.rs +++ b/diskann-benchmark/src/exhaustive/spherical.rs @@ -416,7 +416,7 @@ mod imp { let output_dim = plan.quantizer().output_dim(); let bytes = diskann_quantization::spherical::DataRef::::canonical_bytes(output_dim); - let mut data = diskann_utils::views::Matrix::new(0, input.nrows(), bytes); + let mut data = diskann_utils::views::Matrix::try_new(0, input.nrows(), bytes)?; // Compress the data. #[expect(clippy::disallowed_methods)] diff --git a/diskann-benchmark/src/utils/datafiles.rs b/diskann-benchmark/src/utils/datafiles.rs index abfe06a7de..5bf6950400 100644 --- a/diskann-benchmark/src/utils/datafiles.rs +++ b/diskann-benchmark/src/utils/datafiles.rs @@ -71,21 +71,21 @@ impl ConvertingLoad for f32 { ))] fn converting_load(path: BinFile<'_>, data_type: DataType) -> anyhow::Result> { #[inline(never)] - fn convert(from: diskann_utils::views::MatrixView) -> Matrix + fn convert(from: diskann_utils::views::MatrixView) -> anyhow::Result> where U: Default + Clone + From, T: Copy, { - let mut to = Matrix::new(U::default(), from.nrows(), from.ncols()); + let mut to = Matrix::try_new(U::default(), from.nrows(), from.ncols())?; std::iter::zip(to.as_mut_slice().iter_mut(), from.as_slice().iter()) .for_each(|(t, f)| *t = (*f).into()); - to + Ok(to) } match data_type { DataType::Float32 => load_dataset::(path), - DataType::Float16 => Ok(convert(load_dataset::(path)?.as_view())), - DataType::UInt8 => Ok(convert(load_dataset::(path)?.as_view())), - DataType::Int8 => Ok(convert(load_dataset::(path)?.as_view())), + DataType::Float16 => convert(load_dataset::(path)?.as_view()), + DataType::UInt8 => convert(load_dataset::(path)?.as_view()), + DataType::Int8 => convert(load_dataset::(path)?.as_view()), _ => Err(anyhow::anyhow!( "data type {:?} is not supported for loading `f32` data", data_type diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index 5092ba2e7f..b4532a7a4e 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -555,7 +555,7 @@ where ensure_vertex_loaded(&mut accessor.scratch.vertex_provider, &candidate_ids)?; - let mut candidate_vectors = Matrix::new(0.0f32, candidate_ids.len(), query_f32.len()); + let mut candidate_vectors = Matrix::try_new(0.0f32, candidate_ids.len(), query_f32.len())?; let mut candidate_distances = Vec::with_capacity(candidate_ids.len()); let mut associated_data = Vec::with_capacity(candidate_ids.len()); diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs b/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs index a1e4b81e2f..7ccbc74bae 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs @@ -434,7 +434,10 @@ where let candidates: Vec> = candidates.collect(); let candidate_count = candidates.len(); let store: &FullPrecisionStore = accessor.as_full_precision(); - let mut vectors = Matrix::new(0.0f32, candidate_count, query.len()); + let mut vectors = match Matrix::try_new(0.0f32, candidate_count, query.len()) { + Ok(vectors) => vectors, + Err(error) => return std::future::ready(Err(error.into())), + }; let mut ids = Vec::with_capacity(candidate_count); let mut distances = Vec::with_capacity(candidate_count); diff --git a/diskann-quantization/src/product/train.rs b/diskann-quantization/src/product/train.rs index 946b18a244..5e92ac293b 100644 --- a/diskann-quantization/src/product/train.rs +++ b/diskann-quantization/src/product/train.rs @@ -169,7 +169,15 @@ impl TrainQuantizer for LightPQTrainingParameters { // Allocate scratch data structures. let norms: Vec = view.rows().map(square_norm).collect(); let transpose = BlockTransposed::::from_strided(view); - let mut centers = Matrix::new(0.0, trainer.ncenters, range.len()); + let mut centers = + Matrix::try_new(0.0, trainer.ncenters, range.len()).map_err(|err| { + PQTrainingError { + chunk: i, + of: schema.len(), + dim: range.len(), + kind: PQTrainingErrorKind::InternalError(Box::new(err)), + } + })?; // Construct the random number generator seeded by the PQ chunk. let mut rng = rng_builder.build_boxed_rng(i); diff --git a/diskann/src/error/ann_error.rs b/diskann/src/error/ann_error.rs index bfd211281d..d810f8e8a6 100644 --- a/diskann/src/error/ann_error.rs +++ b/diskann/src/error/ann_error.rs @@ -262,6 +262,7 @@ convert_error!(TryFromIntError); convert_error!(TryFromSliceError); convert_error!(diskann_utils::io::ReadBinError); convert_error!(diskann_utils::io::SaveBinError); +convert_error!(diskann_utils::views::LayoutError); convert_error!(diskann_utils::views::TryFromErrorLight); // Convert from `mpsc::SendError` to `ANNError` diff --git a/diskann/src/graph/start_point.rs b/diskann/src/graph/start_point.rs index c829b9dcfe..15d3b1064c 100644 --- a/diskann/src/graph/start_point.rs +++ b/diskann/src/graph/start_point.rs @@ -67,6 +67,9 @@ pub enum StartPointError { #[error("Error getting row_id {} from training data matrix view", row_id)] MatrixRowError { row_id: usize }, + + #[error(transparent)] + MatrixLayout(#[from] diskann_utils::views::LayoutError), } impl StartPointStrategy { @@ -98,7 +101,7 @@ impl StartPointStrategy { let indices = rand::seq::index::sample(&mut rng, train_data.nrows(), nsamples.get()); - let mut points = Matrix::new(T::default(), nsamples.get(), train_data.ncols()); + let mut points = Matrix::try_new(T::default(), nsamples.get(), train_data.ncols())?; std::iter::zip(points.row_iter_mut(), indices).for_each(|(dst, src)| { dst.copy_from_slice(train_data.row(src)); }); @@ -115,7 +118,7 @@ impl StartPointStrategy { } => { let mut rng = StdRng::seed_from_u64(*seed); let dim = train_data.ncols(); - let mut points = Matrix::new(T::default(), nsamples.get(), dim); + let mut points = Matrix::try_new(T::default(), nsamples.get(), dim)?; points.row_iter_mut().for_each(|row| { row.copy_from_slice(&WithApproximateNorm::with_approximate_norm( dim, *norm, &mut rng,