Skip to main content

data_beans/sparse_io/
factory.rs

1#![allow(dead_code)]
2
3use std::sync::Arc;
4
5#[cfg(feature = "hdf5")]
6use crate::sparse_backend::hdf5 as sparse_matrix_hdf5;
7use crate::sparse_backend::zarr as sparse_matrix_zarr;
8
9#[cfg(feature = "ndarray")]
10use super::Array2;
11use super::{DMatrix, SparseIo, SparseIoBackend};
12
13/// Returned when the user asks for the HDF5 backend but data-beans was built
14/// without the `hdf5` feature. Keeps the error message consistent across the
15/// (small) handful of factory callsites that take a `SparseIoBackend`.
16#[cfg(not(feature = "hdf5"))]
17fn hdf5_disabled<T>() -> anyhow::Result<T> {
18    anyhow::bail!(
19        "HDF5 backend selected but data-beans was built without the `hdf5` feature. \
20         Reinstall with `--features hdf5` (requires libhdf5) or use the Zarr backend."
21    )
22}
23
24/// Open a sparse matrix io (backend)
25/// * `backend_file`: file path to the sparse matrix
26/// * `backend`: backend type (HDF5 or Zarr)
27pub fn open_sparse_matrix(
28    backend_file: &str,
29    backend: &SparseIoBackend,
30) -> anyhow::Result<Box<dyn SparseIo<IndexIter = Vec<usize>>>> {
31    match backend {
32        SparseIoBackend::Zarr => Ok(Box::new(sparse_matrix_zarr::SparseMtxData::open(
33            backend_file,
34        )?)),
35        #[cfg(feature = "hdf5")]
36        SparseIoBackend::HDF5 => Ok(Box::new(sparse_matrix_hdf5::SparseMtxData::open(
37            backend_file,
38        )?)),
39        #[cfg(not(feature = "hdf5"))]
40        SparseIoBackend::HDF5 => hdf5_disabled(),
41    }
42}
43
44/// Open a sparse matrix, choosing the backend from the file name
45/// (`.h5`, `.zarr`, or `.zarr.zip`) via [`crate::hdf5_io::resolve_backend_file`].
46pub fn open_sparse_matrix_by_path(
47    file_path: &str,
48) -> anyhow::Result<Box<dyn SparseIo<IndexIter = Vec<usize>>>> {
49    let (backend, backend_file) = crate::hdf5_io::resolve_backend_file(file_path, None)?;
50    open_sparse_matrix(&backend_file, &backend)
51}
52
53/// Create a sparse backend from a borrowed triplet slice.
54///
55/// Clones the slice internally — prefer [`create_sparse_from_triplets_owned`]
56/// when you already own the Vec, since that avoids the extra copy.
57pub fn create_sparse_from_triplets(
58    triplets: &[(u64, u64, f32)],
59    mtx_shape: (usize, usize, usize),
60    backend_file: Option<&str>,
61    backend: Option<&SparseIoBackend>,
62) -> anyhow::Result<Box<dyn SparseIo<IndexIter = Vec<usize>>>> {
63    create_sparse_from_triplets_owned(triplets.to_vec(), mtx_shape, backend_file, backend)
64}
65
66/// Create a sparse backend from an owned triplet Vec.
67///
68/// Moves the Vec into the backend sort/write without copying. Use this on
69/// hot paths (e.g. `from-*` converters, merges) to avoid holding two full
70/// copies of the triplet list simultaneously.
71pub fn create_sparse_from_triplets_owned(
72    mut triplets: Vec<(u64, u64, f32)>,
73    mtx_shape: (usize, usize, usize),
74    backend_file: Option<&str>,
75    backend: Option<&SparseIoBackend>,
76) -> anyhow::Result<Box<dyn SparseIo<IndexIter = Vec<usize>>>> {
77    match backend {
78        #[cfg(feature = "hdf5")]
79        Some(SparseIoBackend::HDF5) => {
80            let mut ret = Box::new(sparse_matrix_hdf5::SparseMtxData::new(backend_file)?);
81
82            ret.record_mtx_shape(Some(mtx_shape))?;
83            ret.record_triplets_by_col(&mut triplets)?;
84            ret.record_triplets_by_row(&mut triplets)?;
85            ret.read_column_indptr()?;
86            ret.read_row_indptr()?;
87            Ok(ret)
88        }
89        #[cfg(not(feature = "hdf5"))]
90        Some(SparseIoBackend::HDF5) => hdf5_disabled(),
91
92        Some(SparseIoBackend::Zarr) | None => {
93            let mut ret = Box::new(sparse_matrix_zarr::SparseMtxData::new(backend_file)?);
94            ret.record_mtx_shape(Some(mtx_shape))?;
95            ret.record_triplets_by_col(&mut triplets)?;
96            ret.record_triplets_by_row(&mut triplets)?;
97            ret.read_column_indptr()?;
98            ret.read_row_indptr()?;
99            Ok(ret)
100        }
101    }
102}
103
104/// Create an empty sparse backend ready for streaming CSC writes.
105///
106/// The returned backend has no data yet — the caller is expected to
107/// drive [`SparseIo::begin_streaming_csc`], one or more
108/// [`SparseIo::append_csc_slab`] calls, [`SparseIo::finalize_streaming_csc`],
109/// and finally [`SparseIo::build_csr_from_csc_streaming`].
110pub fn create_sparse_streaming_empty(
111    backend_file: Option<&str>,
112    backend: Option<&SparseIoBackend>,
113) -> anyhow::Result<Box<dyn SparseIo<IndexIter = Vec<usize>>>> {
114    match backend {
115        #[cfg(feature = "hdf5")]
116        Some(SparseIoBackend::HDF5) => Ok(Box::new(sparse_matrix_hdf5::SparseMtxData::new(
117            backend_file,
118        )?)),
119        #[cfg(not(feature = "hdf5"))]
120        Some(SparseIoBackend::HDF5) => hdf5_disabled(),
121        Some(SparseIoBackend::Zarr) | None => Ok(Box::new(sparse_matrix_zarr::SparseMtxData::new(
122            backend_file,
123        )?)),
124    }
125}
126
127/// Create a sparse matrix io (backend) with 10x mtx
128/// * `mtx_file`: file path to the 10x mtx
129/// * `backend_file`: file path to the sparse matrix
130/// * `backend`: backend type (HDF5 or Zarr)
131pub fn create_sparse_from_mtx_file(
132    mtx_file: &str,
133    backend_file: Option<&str>,
134    backend: Option<&SparseIoBackend>,
135) -> anyhow::Result<Box<dyn SparseIo<IndexIter = Vec<usize>>>> {
136    match backend {
137        #[cfg(feature = "hdf5")]
138        Some(SparseIoBackend::HDF5) => Ok(Box::new(
139            sparse_matrix_hdf5::SparseMtxData::from_mtx_file(mtx_file, backend_file, Some(true))?,
140        )),
141        #[cfg(not(feature = "hdf5"))]
142        Some(SparseIoBackend::HDF5) => hdf5_disabled(),
143
144        Some(SparseIoBackend::Zarr) | None => Ok(Box::new(
145            sparse_matrix_zarr::SparseMtxData::from_mtx_file(mtx_file, backend_file, Some(true))?,
146        )),
147    }
148}
149
150#[cfg(feature = "ndarray")]
151/// Create a sparse matrix io (backend) with dense `Array2`
152/// * `data`: data matrix
153/// * `backend_file`: file path to the sparse matrix
154/// * `backend`: backend type (HDF5 or Zarr)
155pub fn create_sparse_from_ndarray(
156    data: &Array2<f32>,
157    backend_file: Option<&str>,
158    backend: Option<&SparseIoBackend>,
159) -> anyhow::Result<Box<dyn SparseIo<IndexIter = Vec<usize>>>> {
160    match backend {
161        #[cfg(feature = "hdf5")]
162        Some(SparseIoBackend::HDF5) => Ok(Box::new(
163            sparse_matrix_hdf5::SparseMtxData::from_ndarray(data, backend_file, Some(true))?,
164        )),
165        #[cfg(not(feature = "hdf5"))]
166        Some(SparseIoBackend::HDF5) => hdf5_disabled(),
167
168        Some(SparseIoBackend::Zarr) | None => Ok(Box::new(
169            sparse_matrix_zarr::SparseMtxData::from_ndarray(data, backend_file, Some(true))?,
170        )),
171    }
172}
173
174/// Create a sparse matrix io (backend) with dense `DMatrix`
175/// * `data`: data matrix
176/// * `backend_file`: file path to the sparse matrix
177/// * `backend`: backend type (HDF5 or Zarr)
178pub fn create_sparse_from_dmatrix(
179    data: &DMatrix<f32>,
180    backend_file: Option<&str>,
181    backend: Option<&SparseIoBackend>,
182) -> anyhow::Result<Box<dyn SparseIo<IndexIter = Vec<usize>>>> {
183    match backend {
184        #[cfg(feature = "hdf5")]
185        Some(SparseIoBackend::HDF5) => Ok(Box::new(
186            sparse_matrix_hdf5::SparseMtxData::from_dmatrix(data, backend_file, Some(true))?,
187        )),
188        #[cfg(not(feature = "hdf5"))]
189        Some(SparseIoBackend::HDF5) => hdf5_disabled(),
190
191        Some(SparseIoBackend::Zarr) | None => Ok(Box::new(
192            sparse_matrix_zarr::SparseMtxData::from_dmatrix(data, backend_file, Some(true))?,
193        )),
194    }
195}
196
197pub fn sparse_io_box_to_arc<T>(
198    boxed: Box<dyn SparseIo<IndexIter = T>>,
199) -> Arc<dyn SparseIo<IndexIter = T>> {
200    Arc::from(boxed)
201}