1#![allow(dead_code, unused_imports)]
2
3#[cfg(feature = "tensor")]
4pub use legume_numeric::candle_core::Tensor;
5pub use nalgebra::DMatrix;
6pub use nalgebra_sparse::{csc::CscMatrix, csr::CsrMatrix};
7#[cfg(feature = "ndarray")]
8pub use ndarray::prelude::*;
9
10pub const MAX_ROW_NAME_IDX: usize = 3;
11pub const MAX_COLUMN_NAME_IDX: usize = 10;
12pub const COLUMN_SEP: &str = "@";
13pub const ROW_SEP: &str = "_";
14
15use super::helpers::*;
16use super::index_audit;
17use super::meta::Metadata;
18use anyhow::Context;
19
20use crate::column_subset::{distinct_within, stream_csc_selection, RowMap};
21use crate::sparse_data_visitors::styled_progress_bar;
22use clap::ValueEnum;
23use indicatif::ParallelProgressIterator;
24use legume_numeric::matrix::mtx_io::*;
25use legume_numeric::matrix::traits::*;
26use legume_numeric::matrix::utils::byte_budget_intervals;
27use log::info;
28use rayon::prelude::*;
29use rustc_hash::FxHashMap as HashMap;
30use std::ops::Range;
31
32#[cfg(test)]
33mod tests;
34
35#[derive(ValueEnum, Clone, Debug, PartialEq)]
36#[clap(rename_all = "lowercase")]
37pub enum SparseIoBackend {
38 Zarr,
39 HDF5,
40}
41
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
46pub enum CsKey {
47 CscData,
48 CscIndices,
49 CscIndptr,
50 CsrData,
51 CsrIndices,
52 CsrIndptr,
53}
54
55const SLAB_NNZ: usize = 1 << 20;
59
60fn slab_end(
67 triplets: &[(u64, u64, f32)],
68 start: usize,
69 slab_nnz: usize,
70 n_major: usize,
71 major: impl Fn(&(u64, u64, f32)) -> u64,
72) -> (usize, u64) {
73 debug_assert!(slab_nnz > 0);
74 let nnz = triplets.len();
75 let mut end = (start + slab_nnz).min(nnz);
76 while end < nnz && major(&triplets[end]) == major(&triplets[end - 1]) {
77 end += 1;
78 }
79 let band_end = if end == nnz {
80 n_major as u64
81 } else {
82 major(&triplets[end])
83 };
84 (end, band_end)
85}
86
87fn check_triplet_bounds(
92 label: &str,
93 triplets: &[(u64, u64, f32)],
94 nrow: usize,
95 ncol: usize,
96) -> anyhow::Result<()> {
97 let (nrow_u, ncol_u) = (nrow as u64, ncol as u64);
98 if let Some(k) = triplets
99 .par_iter()
100 .position_first(|&(r, c, _)| r >= nrow_u || c >= ncol_u)
101 {
102 let (r, c, v) = triplets[k];
103 anyhow::bail!(
104 "{label}: triplet {k} = (row {r} (0x{r:016x}), column {c} (0x{c:016x}), {v}) \
105 is outside the {nrow} x {ncol} matrix"
106 );
107 }
108 Ok(())
109}
110
111fn replace_backend_content<S>(
125 data: &mut S,
126 stage_tag: &str,
127 build: impl FnOnce(&S, &mut dyn SparseIo<IndexIter = Vec<usize>>) -> anyhow::Result<()>,
128) -> anyhow::Result<()>
129where
130 S: SparseIo + ?Sized,
131{
132 let final_path = data.get_backend_file_name().to_string();
133 anyhow::ensure!(
134 !final_path.ends_with(".zip"),
135 "{final_path} is a zip archive; convert it to a directory \
136 backend first (data-beans convert)"
137 );
138 let name = std::path::Path::new(&final_path)
141 .file_name()
142 .map(|n| n.to_string_lossy().into_owned())
143 .unwrap_or_default();
144 let temp_path = crate::zarr_io::in_stage(&final_path, &format!("{name}.{stage_tag}"));
145 crate::zarr_io::make_stage(&temp_path)?;
146 if std::path::Path::new(&temp_path).exists() {
147 crate::sparse_io::remove_backend_path(&temp_path)?;
148 }
149
150 {
151 let backend_kind = data.backend_type();
152 let mut out =
153 crate::sparse_io::create_sparse_streaming_empty(Some(&temp_path), Some(&backend_kind))?;
154 build(data, &mut *out)?;
155 out.set_metadata(&data.metadata())?;
156 }
157
158 data.remove_backend_file()?;
159 std::fs::rename(&temp_path, &final_path)?;
160 crate::zarr_io::leave_stage(&temp_path);
161 data.reopen_backend()?;
162 info!("registered new data to {}", data.get_backend_file_name());
163 Ok(())
164}
165
166pub(crate) fn build_csr_from_csc_banded<S>(data: &mut S, band_bytes: usize) -> anyhow::Result<()>
171where
172 S: SparseIo + ?Sized,
173{
174 let nrow = data
175 .num_rows()
176 .ok_or_else(|| anyhow::anyhow!("nrow not set before build_csr_from_csc_streaming"))?;
177 let ncol = data
178 .num_columns()
179 .ok_or_else(|| anyhow::anyhow!("ncol not set before build_csr_from_csc_streaming"))?;
180 let nnz = data
181 .num_non_zeros()
182 .ok_or_else(|| anyhow::anyhow!("nnz not set before build_csr_from_csc_streaming"))?;
183
184 if nnz == 0 {
185 data.cs_create(CsKey::CsrData, 0)?;
186 data.cs_create(CsKey::CsrIndices, 0)?;
187 data.cs_create(CsKey::CsrIndptr, nrow + 1)?;
188 let zeros = vec![0u64; nrow + 1];
189 data.cs_write_u64(CsKey::CsrIndptr, 0, &zeros)?;
190 data.read_row_indptr()?;
191 return Ok(());
192 }
193
194 const COL_BLOCK: usize = 1024;
195 let n_col_blocks = ncol.div_ceil(COL_BLOCK);
196 let bar1 = styled_progress_bar(n_col_blocks as u64, "transpose count");
197 let mut row_counts = vec![0u64; nrow];
198 let mut col_lo = 0usize;
199 while col_lo < ncol {
200 let col_hi = (col_lo + COL_BLOCK).min(ncol);
201 let cols: S::IndexIter = (col_lo..col_hi).collect();
202 let (_, _, triplets) = data.read_triplets_by_columns(cols)?;
203 for t in &triplets {
204 let slot = row_counts.get_mut(t.0 as usize).ok_or_else(|| {
205 anyhow::anyhow!(
206 "build_csr_from_csc_streaming: row index {} (0x{:016x}) outside \
207 the {nrow}-row matrix",
208 t.0,
209 t.0
210 )
211 })?;
212 *slot += 1;
213 }
214 col_lo = col_hi;
215 bar1.inc(1);
216 }
217 bar1.finish_and_clear();
218
219 let mut rowptr = vec![0u64; nrow + 1];
220 let mut acc = 0u64;
221 for i in 0..nrow {
222 rowptr[i] = acc;
223 acc += row_counts[i];
224 }
225 rowptr[nrow] = acc;
226 anyhow::ensure!(
227 acc == nnz as u64,
228 "build_csr_from_csc_streaming: counted {acc} entries in the CSC arrays, \
229 but {nnz} are declared"
230 );
231
232 data.cs_create(CsKey::CsrData, nnz)?;
233 data.cs_create(CsKey::CsrIndices, nnz)?;
234 data.cs_create(CsKey::CsrIndptr, nrow + 1)?;
235 data.cs_write_u64(CsKey::CsrIndptr, 0, &rowptr)?;
236
237 let bands: Vec<(usize, usize)> =
241 byte_budget_intervals(&row_counts, band_bytes, PRELOAD_BYTES_PER_NNZ)
242 .into_iter()
243 .filter(|&(lo, hi)| rowptr[hi] > rowptr[lo])
244 .collect();
245
246 let bar2 = styled_progress_bar(bands.len() as u64, "transpose scatter");
247 for (band_lo, band_hi) in bands {
248 let band_nnz_start = rowptr[band_lo];
249 let band_nnz = (rowptr[band_hi] - band_nnz_start) as usize;
250
251 let mut out_indices = vec![0u64; band_nnz];
252 let mut out_values = vec![0f32; band_nnz];
253 let mut cursor = vec![0u64; band_hi - band_lo];
254
255 let mut col_lo = 0usize;
256 while col_lo < ncol {
257 let col_hi = (col_lo + COL_BLOCK).min(ncol);
258 let cols: S::IndexIter = (col_lo..col_hi).collect();
259 let (_, _, triplets) = data.read_triplets_by_columns(cols)?;
260 for &(row_i, col_j_local, x) in &triplets {
261 let row_i_us = row_i as usize;
262 if row_i_us >= band_lo && row_i_us < band_hi {
263 let band_idx = row_i_us - band_lo;
264 let col_j_global = col_j_local + col_lo as u64;
269 let offset_in_band =
270 (rowptr[band_lo + band_idx] - band_nnz_start) + cursor[band_idx];
271 out_indices[offset_in_band as usize] = col_j_global;
272 out_values[offset_in_band as usize] = x;
273 cursor[band_idx] += 1;
274 }
275 }
276 col_lo = col_hi;
277 }
278
279 data.cs_write_u64(CsKey::CsrIndices, band_nnz_start, &out_indices)?;
280 data.cs_write_f32(CsKey::CsrData, band_nnz_start, &out_values)?;
281
282 bar2.inc(1);
283 }
284 bar2.finish_and_clear();
285
286 data.read_row_indptr()?;
287 Ok(())
288}
289
290pub trait SparseIo: Sync + Send {
291 type IndexIter: IntoIterator<Item = usize> + FromIterator<usize>;
292
293 #[cfg(feature = "ndarray")]
298 fn read_columns_ndarray(&self, columns: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
302 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
303 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
304 }
305
306 #[cfg(feature = "tensor")]
307 fn read_columns_tensor(&self, columns: Self::IndexIter) -> anyhow::Result<Tensor> {
311 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
312 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
313 }
314
315 fn read_columns_dmatrix(&self, columns: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
319 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
320 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
321 }
322
323 fn read_columns_csr(&self, columns: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
327 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
328 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
329 }
330
331 fn read_columns_csc(&self, columns: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
335 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
336 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
337 }
338
339 fn csc_column_arrays(&self) -> Option<(&[u64], &[u64], &[f32])> {
345 None
346 }
347
348 #[cfg(feature = "ndarray")]
349 fn read_rows_ndarray(&self, rows: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
353 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
354 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
355 }
356
357 #[cfg(feature = "tensor")]
358 fn read_rows_tensor(&self, rows: Self::IndexIter) -> anyhow::Result<Tensor> {
362 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
363 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
364 }
365
366 fn read_rows_dmatrix(&self, rows: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
370 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
371 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
372 }
373
374 fn read_rows_csr(&self, rows: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
378 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
379 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
380 }
381
382 fn read_rows_csc(&self, rows: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
386 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
387 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
388 }
389
390 fn import_mtx_file(&mut self, mtx_file: &str, index_by_row: bool) -> anyhow::Result<()> {
400 let (mut mtx_triplets, mtx_shape) = read_mtx_triplets(mtx_file)?;
401 info!("read mtx file: {}", mtx_file);
402 if mtx_triplets.is_empty() {
403 return Err(anyhow::anyhow!("No data in mtx file"));
404 }
405 self.record_mtx_shape(Some(mtx_shape))?;
406 info!("recording the column index");
407 self.record_triplets_by_col(&mut mtx_triplets)?;
408 if index_by_row {
409 info!("recording the row index");
410 self.record_triplets_by_row(&mut mtx_triplets)?;
411 }
412 Ok(())
413 }
414
415 fn import_dmatrix_by_row(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
422 let (nrow, ncol) = matrix.shape();
423 let mut mtx_triplets = dmatrix_to_triplets(matrix);
424 let mtx_shape = (nrow, ncol, mtx_triplets.len());
425 self.record_mtx_shape(Some(mtx_shape))?;
426 self.record_triplets_by_row(&mut mtx_triplets)
427 }
428
429 fn import_dmatrix_by_col(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
432 let (nrow, ncol) = matrix.shape();
433 let mut mtx_triplets = dmatrix_to_triplets(matrix);
434 let mtx_shape = (nrow, ncol, mtx_triplets.len());
435 self.record_mtx_shape(Some(mtx_shape))?;
436 self.record_triplets_by_col(&mut mtx_triplets)
437 }
438
439 #[cfg(feature = "ndarray")]
444 fn import_ndarray_by_row(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
447 let nrow = array.shape()[0];
448 let ncol = array.shape()[1];
449
450 let mut mtx_triplets = ndarray_to_triplets(array);
452
453 let nnz = mtx_triplets.len();
454 let mtx_shape = (nrow, ncol, nnz);
455 self.record_mtx_shape(Some(mtx_shape))?;
456
457 self.record_triplets_by_row(&mut mtx_triplets)
460 }
461
462 #[cfg(feature = "ndarray")]
463 fn import_ndarray_by_col(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
466 let nrow = array.shape()[0];
467 let ncol = array.shape()[1];
468
469 let mut mtx_triplets = ndarray_to_triplets(array);
471
472 let nnz = mtx_triplets.len();
473 let mtx_shape = (nrow, ncol, nnz);
474 self.record_mtx_shape(Some(mtx_shape))?;
475
476 self.record_triplets_by_col(&mut mtx_triplets)
479 }
480
481 #[allow(clippy::type_complexity)]
489 fn read_triplets_by_rows(
490 &self,
491 rows: Self::IndexIter,
492 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
493
494 #[allow(clippy::type_complexity)]
498 fn read_triplets_by_columns(
499 &self,
500 columns: Self::IndexIter,
501 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
502
503 #[allow(clippy::type_complexity)]
507 fn read_triplets_by_single_column(
508 &self,
509 col: usize,
510 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
511
512 fn to_mtx_file(&self, mtx_file: &str) -> anyhow::Result<()>;
515
516 fn num_rows(&self) -> Option<usize>;
518
519 fn num_columns(&self) -> Option<usize>;
521
522 fn num_non_zeros(&self) -> Option<usize>;
524
525 fn reopen_backend(&mut self) -> anyhow::Result<()>;
530
531 #[doc(hidden)]
535 fn note_streamed_nnz(&mut self, n: u64);
543
544 #[doc(hidden)]
546 fn streamed_nnz(&self) -> u64;
547
548 #[doc(hidden)]
550 fn reset_streamed_nnz(&mut self);
551
552 fn column_indptr(&self) -> &[u64];
557
558 fn row_indptr(&self) -> &[u64];
561
562 fn column_nnz(&self, col: usize) -> Option<u64> {
568 let indptr = self.column_indptr();
569 let hi = *indptr.get(col + 1)?;
570 let lo = *indptr.get(col)?;
571 hi.checked_sub(lo)
572 }
573
574 fn register_row_names_file(&mut self, row_name_file: &str);
577
578 fn register_column_names_file(&mut self, column_name_file: &str);
581
582 fn register_row_names_vec(&mut self, rows: &[Box<str>]);
585
586 fn register_column_names_vec(&mut self, columns: &[Box<str>]);
589
590 fn register_names_file(
596 &mut self,
597 key: &str,
598 name_file: &str,
599 name_columns: Range<usize>,
600 name_sep: &str,
601 ) -> anyhow::Result<()>;
602
603 fn register_names_vec(&mut self, key: &str, names: &[Box<str>]) -> anyhow::Result<()>;
607
608 fn row_names(&self) -> anyhow::Result<Vec<Box<str>>>;
609
610 fn column_names(&self) -> anyhow::Result<Vec<Box<str>>>;
611
612 fn retrieve_registered_names(&self, key: &str) -> anyhow::Result<Vec<Box<str>>>;
615
616 fn metadata(&self) -> Metadata;
624
625 fn set_metadata(&mut self, meta: &Metadata) -> anyhow::Result<()>;
630
631 fn meta(&self, key: &str) -> Option<String> {
633 self.metadata().remove(key)
634 }
635
636 fn set_meta(&mut self, key: &str, value: &str) -> anyhow::Result<()> {
638 let mut meta = self.metadata();
639 meta.insert(key.to_string(), value.to_string());
640 self.set_metadata(&meta)
641 }
642
643 fn subset_columns_rows(
660 &mut self,
661 columns: Option<&Vec<usize>>,
662 rows: Option<&Vec<usize>>,
663 ) -> anyhow::Result<()> {
664 let ncol_data = self
665 .num_columns()
666 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
667 let nrow_data = self
668 .num_rows()
669 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
670
671 if let Some(cols) = columns {
672 anyhow::ensure!(!cols.is_empty(), "subset: empty column selection");
673 distinct_within(cols, ncol_data, "column")?;
674 }
675 let row_map = match rows {
676 Some(rs) => {
677 anyhow::ensure!(!rs.is_empty(), "subset: empty row selection");
678 Some(RowMap::keeping(rs, nrow_data)?)
679 }
680 None => None,
681 };
682
683 let take_names = |sel: Option<&Vec<usize>>,
685 names: Vec<Box<str>>,
686 what: &str|
687 -> anyhow::Result<Vec<Box<str>>> {
688 let Some(sel) = sel else {
689 return Ok(names);
690 };
691 sel.iter()
692 .map(|&i| {
693 names.get(i).cloned().ok_or_else(|| {
694 anyhow::anyhow!("subset: {what} {i} has no name ({} named)", names.len())
695 })
696 })
697 .collect()
698 };
699 let new_col_names = take_names(columns, self.column_names()?, "column")?;
700 let new_row_names = take_names(rows, self.row_names()?, "row")?;
701
702 let every_column: Vec<usize>;
703 let old_cols: &[usize] = match columns {
704 Some(cols) => cols,
705 None => {
706 every_column = (0..ncol_data).collect();
707 &every_column
708 }
709 };
710 replace_backend_content(self, "subset_tmp", |data, out| {
711 stream_csc_selection(data, old_cols, row_map.as_ref(), out)?;
712 out.register_row_names_vec(&new_row_names);
713 out.register_column_names_vec(&new_col_names);
714 Ok(())
715 })
716 }
717
718 fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()> {
728 let Some(ncol) = self.num_columns() else {
729 self.clean_preloaded_columns();
730 self.clean_preloaded_rows();
731 return Ok(());
732 };
733 let nrow_data = self
734 .num_rows()
735 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
736
737 let row_map = RowMap::by_names(&self.row_names()?, row_names_order, nrow_data)?;
738
739 let new_col_names = self.column_names()?;
740 let old_cols: Vec<usize> = (0..ncol).collect();
741
742 info!("remapping triplets ...");
743 replace_backend_content(self, "reorder_tmp", |data, out| {
744 stream_csc_selection(data, &old_cols, Some(&row_map), out)?;
745 out.register_row_names_vec(row_names_order);
746 out.register_column_names_vec(&new_col_names);
747 Ok(())
748 })
749 }
750
751 fn remove_backend_file(&self) -> anyhow::Result<()>;
753
754 fn initialize_backend(&mut self) -> anyhow::Result<()>;
756
757 fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()>;
758
759 fn record_triplets_by_row(
762 &mut self,
763 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
764 ) -> anyhow::Result<()> {
765 let nrow = self.num_rows().expect("should have `nrow`");
766 let ncol = self.num_columns().expect("should have `ncol`");
767 let nnz = row_col_val_triplets.len();
768
769 if nnz == 0 {
770 let csr_rowptr = vec![0u64; nrow + 1];
771 return self.record_csr_dataset_backend(&[], &[], &csr_rowptr);
772 }
773
774 check_triplet_bounds("record_triplets_by_row", row_col_val_triplets, nrow, ncol)?;
775
776 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (row, col));
780
781 self.begin_streaming_csr((nrow, ncol, nnz))?;
782
783 let mut local_rowptr: Vec<u64> = Vec::new();
784 let mut cols: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
785 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
786
787 let mut start = 0_usize;
788 let mut row_offset = 0_u64;
789 while (row_offset as usize) < nrow {
790 let (end, band_end_row) =
791 slab_end(row_col_val_triplets, start, SLAB_NNZ, nrow, |t| t.0);
792
793 local_rowptr.clear();
794 cols.clear();
795 vals.clear();
796 let mut i = start;
797 for row in row_offset..band_end_row {
798 local_rowptr.push((i - start) as u64);
799 while i < end && row_col_val_triplets[i].0 == row {
800 cols.push(row_col_val_triplets[i].1);
801 vals.push(row_col_val_triplets[i].2);
802 i += 1;
803 }
804 }
805 if i != end {
806 anyhow::bail!(
807 "record_triplets_by_row: band walk left triplet {:?} at {i}",
808 row_col_val_triplets[i]
809 );
810 }
811
812 self.append_csr_slab(row_offset, start as u64, &local_rowptr, &cols, &vals)?;
813 start = end;
814 row_offset = band_end_row;
815 }
816
817 self.finalize_streaming_csr()
818 }
819
820 fn record_triplets_by_col(
827 &mut self,
828 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
829 ) -> anyhow::Result<()> {
830 let nrow = self.num_rows().expect("should have `nrow`");
831 let ncol = self.num_columns().expect("should have `ncol`");
832 let nnz = row_col_val_triplets.len();
833
834 if nnz == 0 {
835 let csc_colptr = vec![0u64; ncol + 1];
836 return self.record_csc_dataset_backend(&[], &[], &csc_colptr);
837 }
838
839 check_triplet_bounds("record_triplets_by_col", row_col_val_triplets, nrow, ncol)?;
840
841 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (col, row));
843
844 self.begin_streaming_csc((nrow, ncol, nnz))?;
845
846 let mut local_colptr: Vec<u64> = Vec::new();
847 let mut rows: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
848 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
849
850 let mut start = 0_usize;
851 let mut col_offset = 0_u64;
852 while (col_offset as usize) < ncol {
853 let (end, band_end_col) =
854 slab_end(row_col_val_triplets, start, SLAB_NNZ, ncol, |t| t.1);
855
856 local_colptr.clear();
857 rows.clear();
858 vals.clear();
859 let mut i = start;
860 for col in col_offset..band_end_col {
861 local_colptr.push((i - start) as u64);
862 while i < end && row_col_val_triplets[i].1 == col {
863 rows.push(row_col_val_triplets[i].0);
864 vals.push(row_col_val_triplets[i].2);
865 i += 1;
866 }
867 }
868 if i != end {
869 anyhow::bail!(
870 "record_triplets_by_col: band walk left triplet {:?} at {i}",
871 row_col_val_triplets[i]
872 );
873 }
874
875 self.append_csc_slab(col_offset, start as u64, &local_colptr, &rows, &vals)?;
876 start = end;
877 col_offset = band_end_col;
878 }
879
880 self.finalize_streaming_csc()
881 }
882
883 fn record_csr_dataset_backend(
892 &mut self,
893 csr_cols: &[u64],
894 csr_vals: &[f32],
895 csr_rowptr: &[u64],
896 ) -> anyhow::Result<()>;
897
898 fn record_csc_dataset_backend(
908 &mut self,
909 csc_rows: &[u64],
910 csc_vals: &[f32],
911 csc_colptr: &[u64],
912 ) -> anyhow::Result<()>;
913
914 fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()>;
917
918 fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()>;
921
922 fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()>;
925
926 fn begin_streaming_csc(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
931 self.reset_streamed_nnz();
935 let (_, ncol, nnz) = shape;
936 self.record_mtx_shape(Some(shape))?;
937 self.cs_create(CsKey::CscData, nnz)?;
938 self.cs_create(CsKey::CscIndices, nnz)?;
939 self.cs_create(CsKey::CscIndptr, ncol + 1)?;
940 Ok(())
941 }
942
943 fn append_csc_slab(
952 &mut self,
953 col_offset: u64,
954 nnz_offset: u64,
955 local_colptr: &[u64],
956 row_indices: &[u64],
957 values: &[f32],
958 ) -> anyhow::Result<()> {
959 anyhow::ensure!(
965 row_indices.len() == values.len(),
966 "append_csc_slab: {} row indices vs {} values",
967 row_indices.len(),
968 values.len()
969 );
970 anyhow::ensure!(
971 local_colptr.first().copied() == Some(0) || local_colptr.is_empty(),
972 "append_csc_slab: local_colptr must start at 0"
973 );
974 anyhow::ensure!(
975 local_colptr.windows(2).all(|w| w[0] <= w[1]),
976 "append_csc_slab: local_colptr must be monotone non-decreasing"
977 );
978 if let Some(&last) = local_colptr.last() {
979 anyhow::ensure!(
980 last <= values.len() as u64,
981 "append_csc_slab: colptr claims {last} entries, slab holds {}",
982 values.len()
983 );
984 }
985 if let Some(nrow) = self.num_rows() {
986 if let Some(&bad) = row_indices.iter().find(|&&r| r >= nrow as u64) {
987 anyhow::bail!("append_csc_slab: row index {bad} outside the {nrow}-row matrix");
988 }
989 }
990 for (c, &start) in local_colptr.iter().enumerate() {
993 let end = local_colptr
994 .get(c + 1)
995 .copied()
996 .unwrap_or(values.len() as u64) as usize;
997 anyhow::ensure!(
998 row_indices[start as usize..end]
999 .windows(2)
1000 .all(|w| w[0] < w[1]),
1001 "append_csc_slab: rows within column {} of this band must be \
1002 strictly ascending — repeated rows usually mean duplicate \
1003 (row, col) coordinates in the source (an MTX with repeated \
1004 entries, or a union remap folding rows together)",
1005 col_offset as usize + c
1006 );
1007 }
1008
1009 let shifted: Vec<u64> = local_colptr.iter().map(|&p| p + nnz_offset).collect();
1010 self.cs_write_u64(CsKey::CscIndptr, col_offset, &shifted)?;
1011 self.cs_write_u64(CsKey::CscIndices, nnz_offset, row_indices)?;
1012 self.cs_write_f32(CsKey::CscData, nnz_offset, values)?;
1013 self.note_streamed_nnz(values.len() as u64);
1014 Ok(())
1015 }
1016
1017 fn finalize_streaming_csc(&mut self) -> anyhow::Result<()> {
1020 let ncol = self
1021 .num_columns()
1022 .ok_or_else(|| anyhow::anyhow!("ncol not set before finalize_streaming_csc"))?;
1023 let nnz = self
1024 .num_non_zeros()
1025 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csc"))?;
1026 self.cs_write_u64(CsKey::CscIndptr, ncol as u64, &[nnz as u64])?;
1027 self.read_column_indptr()?;
1028
1029 index_audit::check_indptr("finalize_streaming_csc", self.column_indptr(), ncol, nnz)
1038 .context(
1039 "the slabs were appended with a gap or overlap in their nnz offsets, \
1040 or the first slab was never appended",
1041 )?;
1042 let appended = self.streamed_nnz();
1048 anyhow::ensure!(
1049 appended == nnz as u64,
1050 "finalize_streaming_csc: {appended} entries appended but {nnz} \
1051 declared — the difference reads back as fill values wearing real \
1052 entries' positions"
1053 );
1054 Ok(())
1055 }
1056
1057 fn begin_streaming_csr(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
1060 self.reset_streamed_nnz();
1061 let (nrow, _, nnz) = shape;
1062 self.record_mtx_shape(Some(shape))?;
1063 self.cs_create(CsKey::CsrData, nnz)?;
1064 self.cs_create(CsKey::CsrIndices, nnz)?;
1065 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1066 Ok(())
1067 }
1068
1069 fn append_csr_slab(
1079 &mut self,
1080 row_offset: u64,
1081 nnz_offset: u64,
1082 local_rowptr: &[u64],
1083 col_indices: &[u64],
1084 values: &[f32],
1085 ) -> anyhow::Result<()> {
1086 anyhow::ensure!(
1087 col_indices.len() == values.len(),
1088 "append_csr_slab: {} column indices vs {} values",
1089 col_indices.len(),
1090 values.len()
1091 );
1092 anyhow::ensure!(
1093 local_rowptr.first().copied() == Some(0) || local_rowptr.is_empty(),
1094 "append_csr_slab: local_rowptr must start at 0"
1095 );
1096 anyhow::ensure!(
1097 local_rowptr.windows(2).all(|w| w[0] <= w[1]),
1098 "append_csr_slab: local_rowptr must be monotone non-decreasing"
1099 );
1100 if let Some(&last) = local_rowptr.last() {
1101 anyhow::ensure!(
1102 last <= values.len() as u64,
1103 "append_csr_slab: rowptr claims {last} entries, slab holds {}",
1104 values.len()
1105 );
1106 }
1107 if let Some(ncol) = self.num_columns() {
1108 if let Some(&bad) = col_indices.iter().find(|&&c| c >= ncol as u64) {
1109 anyhow::bail!(
1110 "append_csr_slab: column index {bad} outside the {ncol}-column matrix"
1111 );
1112 }
1113 }
1114 for (r, &start) in local_rowptr.iter().enumerate() {
1115 let end = local_rowptr
1116 .get(r + 1)
1117 .copied()
1118 .unwrap_or(values.len() as u64) as usize;
1119 anyhow::ensure!(
1120 col_indices[start as usize..end]
1121 .windows(2)
1122 .all(|w| w[0] < w[1]),
1123 "append_csr_slab: columns within row {} of this band must be \
1124 strictly ascending — repeated columns usually mean duplicate \
1125 (row, col) coordinates in the source",
1126 row_offset as usize + r
1127 );
1128 }
1129
1130 let shifted: Vec<u64> = local_rowptr.iter().map(|&p| p + nnz_offset).collect();
1131 self.cs_write_u64(CsKey::CsrIndptr, row_offset, &shifted)?;
1132 self.cs_write_u64(CsKey::CsrIndices, nnz_offset, col_indices)?;
1133 self.cs_write_f32(CsKey::CsrData, nnz_offset, values)?;
1134 self.note_streamed_nnz(values.len() as u64);
1135 Ok(())
1136 }
1137
1138 fn finalize_streaming_csr(&mut self) -> anyhow::Result<()> {
1142 let nrow = self
1143 .num_rows()
1144 .ok_or_else(|| anyhow::anyhow!("nrow not set before finalize_streaming_csr"))?;
1145 let nnz = self
1146 .num_non_zeros()
1147 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csr"))?;
1148 self.cs_write_u64(CsKey::CsrIndptr, nrow as u64, &[nnz as u64])?;
1149 self.read_row_indptr()?;
1150
1151 index_audit::check_indptr("finalize_streaming_csr", self.row_indptr(), nrow, nnz).context(
1153 "the slabs were appended with a gap or overlap in their nnz offsets, \
1154 or the first slab was never appended",
1155 )?;
1156
1157 let appended = self.streamed_nnz();
1158 anyhow::ensure!(
1159 appended == nnz as u64,
1160 "finalize_streaming_csr: {appended} entries appended but {nnz} \
1161 declared — the slabs did not cover the matrix"
1162 );
1163 Ok(())
1164 }
1165
1166 fn build_csr_from_csc_streaming(&mut self) -> anyhow::Result<()> {
1173 let wanted = self
1174 .num_non_zeros()
1175 .unwrap_or(0)
1176 .saturating_mul(PRELOAD_BYTES_PER_NNZ);
1177 let (band_bytes, _reservation) = reserve_transpose_band(wanted);
1179 build_csr_from_csc_banded(self, band_bytes)
1180 }
1181
1182 fn read_row_indptr(&mut self) -> anyhow::Result<()>;
1184
1185 fn read_column_indptr(&mut self) -> anyhow::Result<()>;
1187
1188 fn preload_columns(&mut self) -> anyhow::Result<()>;
1190
1191 fn clean_preloaded_columns(&mut self);
1193
1194 fn preload_rows(&mut self) -> anyhow::Result<()>;
1196
1197 fn clean_preloaded_rows(&mut self);
1199
1200 fn get_backend_file_name(&self) -> &str;
1202
1203 fn backend_type(&self) -> SparseIoBackend;
1205}