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