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::sparse_data_visitors::styled_progress_bar;
21use clap::ValueEnum;
22use indicatif::ParallelProgressIterator;
23use legume_numeric::matrix::mtx_io::*;
24use legume_numeric::matrix::traits::*;
25use log::info;
26use rayon::prelude::*;
27use rustc_hash::FxHashMap as HashMap;
28use std::ops::Range;
29use std::sync::{Arc, Mutex};
30
31#[cfg(test)]
32mod tests;
33
34#[derive(ValueEnum, Clone, Debug, PartialEq)]
35#[clap(rename_all = "lowercase")]
36pub enum SparseIoBackend {
37 Zarr,
38 HDF5,
39}
40
41#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45pub enum CsKey {
46 CscData,
47 CscIndices,
48 CscIndptr,
49 CsrData,
50 CsrIndices,
51 CsrIndptr,
52}
53
54const SLAB_NNZ: usize = 1 << 20;
58
59fn slab_end(
66 triplets: &[(u64, u64, f32)],
67 start: usize,
68 slab_nnz: usize,
69 n_major: usize,
70 major: impl Fn(&(u64, u64, f32)) -> u64,
71) -> (usize, u64) {
72 debug_assert!(slab_nnz > 0);
73 let nnz = triplets.len();
74 let mut end = (start + slab_nnz).min(nnz);
75 while end < nnz && major(&triplets[end]) == major(&triplets[end - 1]) {
76 end += 1;
77 }
78 let band_end = if end == nnz {
79 n_major as u64
80 } else {
81 major(&triplets[end])
82 };
83 (end, band_end)
84}
85
86fn check_triplet_bounds(
91 label: &str,
92 triplets: &[(u64, u64, f32)],
93 nrow: usize,
94 ncol: usize,
95) -> anyhow::Result<()> {
96 let (nrow_u, ncol_u) = (nrow as u64, ncol as u64);
97 if let Some(k) = triplets
98 .par_iter()
99 .position_first(|&(r, c, _)| r >= nrow_u || c >= ncol_u)
100 {
101 let (r, c, v) = triplets[k];
102 anyhow::bail!(
103 "{label}: triplet {k} = (row {r} (0x{r:016x}), column {c} (0x{c:016x}), {v}) \
104 is outside the {nrow} x {ncol} matrix"
105 );
106 }
107 Ok(())
108}
109
110pub trait SparseIo: Sync + Send {
111 type IndexIter: IntoIterator<Item = usize> + FromIterator<usize>;
112
113 #[cfg(feature = "ndarray")]
118 fn read_columns_ndarray(&self, columns: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
122 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
123 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
124 }
125
126 #[cfg(feature = "tensor")]
127 fn read_columns_tensor(&self, columns: Self::IndexIter) -> anyhow::Result<Tensor> {
131 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
132 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
133 }
134
135 fn read_columns_dmatrix(&self, columns: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
139 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
140 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
141 }
142
143 fn read_columns_csr(&self, columns: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
147 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
148 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
149 }
150
151 fn read_columns_csc(&self, columns: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
155 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
156 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
157 }
158
159 fn csc_column_arrays(&self) -> Option<(&[u64], &[u64], &[f32])> {
165 None
166 }
167
168 #[cfg(feature = "ndarray")]
169 fn read_rows_ndarray(&self, rows: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
173 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
174 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
175 }
176
177 #[cfg(feature = "tensor")]
178 fn read_rows_tensor(&self, rows: Self::IndexIter) -> anyhow::Result<Tensor> {
182 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
183 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
184 }
185
186 fn read_rows_dmatrix(&self, rows: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
190 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
191 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
192 }
193
194 fn read_rows_csr(&self, rows: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
198 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
199 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
200 }
201
202 fn read_rows_csc(&self, rows: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
206 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
207 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
208 }
209
210 fn import_mtx_file(&mut self, mtx_file: &str, index_by_row: bool) -> anyhow::Result<()> {
220 let (mut mtx_triplets, mtx_shape) = read_mtx_triplets(mtx_file)?;
221 info!("read mtx file: {}", mtx_file);
222 if mtx_triplets.is_empty() {
223 return Err(anyhow::anyhow!("No data in mtx file"));
224 }
225 self.record_mtx_shape(Some(mtx_shape))?;
226 info!("recording the column index");
227 self.record_triplets_by_col(&mut mtx_triplets)?;
228 if index_by_row {
229 info!("recording the row index");
230 self.record_triplets_by_row(&mut mtx_triplets)?;
231 }
232 Ok(())
233 }
234
235 fn import_dmatrix_by_row(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
242 let (nrow, ncol) = matrix.shape();
243 let mut mtx_triplets = dmatrix_to_triplets(matrix);
244 let mtx_shape = (nrow, ncol, mtx_triplets.len());
245 self.record_mtx_shape(Some(mtx_shape))?;
246 self.record_triplets_by_row(&mut mtx_triplets)
247 }
248
249 fn import_dmatrix_by_col(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
252 let (nrow, ncol) = matrix.shape();
253 let mut mtx_triplets = dmatrix_to_triplets(matrix);
254 let mtx_shape = (nrow, ncol, mtx_triplets.len());
255 self.record_mtx_shape(Some(mtx_shape))?;
256 self.record_triplets_by_col(&mut mtx_triplets)
257 }
258
259 #[cfg(feature = "ndarray")]
264 fn import_ndarray_by_row(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
267 let nrow = array.shape()[0];
268 let ncol = array.shape()[1];
269
270 let mut mtx_triplets = ndarray_to_triplets(array);
272
273 let nnz = mtx_triplets.len();
274 let mtx_shape = (nrow, ncol, nnz);
275 self.record_mtx_shape(Some(mtx_shape))?;
276
277 self.record_triplets_by_row(&mut mtx_triplets)
280 }
281
282 #[cfg(feature = "ndarray")]
283 fn import_ndarray_by_col(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
286 let nrow = array.shape()[0];
287 let ncol = array.shape()[1];
288
289 let mut mtx_triplets = ndarray_to_triplets(array);
291
292 let nnz = mtx_triplets.len();
293 let mtx_shape = (nrow, ncol, nnz);
294 self.record_mtx_shape(Some(mtx_shape))?;
295
296 self.record_triplets_by_col(&mut mtx_triplets)
299 }
300
301 #[allow(clippy::type_complexity)]
309 fn read_triplets_by_rows(
310 &self,
311 rows: Self::IndexIter,
312 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
313
314 #[allow(clippy::type_complexity)]
318 fn read_triplets_by_columns(
319 &self,
320 columns: Self::IndexIter,
321 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
322
323 #[allow(clippy::type_complexity)]
327 fn read_triplets_by_single_column(
328 &self,
329 col: usize,
330 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
331
332 fn to_mtx_file(&self, mtx_file: &str) -> anyhow::Result<()>;
335
336 fn num_rows(&self) -> Option<usize>;
338
339 fn num_columns(&self) -> Option<usize>;
341
342 fn num_non_zeros(&self) -> Option<usize>;
344
345 fn reopen_backend(&mut self) -> anyhow::Result<()>;
350
351 #[doc(hidden)]
355 fn note_streamed_nnz(&mut self, n: u64);
363
364 #[doc(hidden)]
366 fn streamed_nnz(&self) -> u64;
367
368 #[doc(hidden)]
370 fn reset_streamed_nnz(&mut self);
371
372 fn column_indptr(&self) -> &[u64];
377
378 fn row_indptr(&self) -> &[u64];
381
382 fn column_nnz(&self, col: usize) -> Option<u64> {
388 let indptr = self.column_indptr();
389 let hi = *indptr.get(col + 1)?;
390 let lo = *indptr.get(col)?;
391 hi.checked_sub(lo)
392 }
393
394 fn register_row_names_file(&mut self, row_name_file: &str);
397
398 fn register_column_names_file(&mut self, column_name_file: &str);
401
402 fn register_row_names_vec(&mut self, rows: &[Box<str>]);
405
406 fn register_column_names_vec(&mut self, columns: &[Box<str>]);
409
410 fn register_names_file(
416 &mut self,
417 key: &str,
418 name_file: &str,
419 name_columns: Range<usize>,
420 name_sep: &str,
421 ) -> anyhow::Result<()>;
422
423 fn register_names_vec(&mut self, key: &str, names: &[Box<str>]) -> anyhow::Result<()>;
427
428 fn row_names(&self) -> anyhow::Result<Vec<Box<str>>>;
429
430 fn column_names(&self) -> anyhow::Result<Vec<Box<str>>>;
431
432 fn retrieve_registered_names(&self, key: &str) -> anyhow::Result<Vec<Box<str>>>;
435
436 fn metadata(&self) -> Metadata;
444
445 fn set_metadata(&mut self, meta: &Metadata) -> anyhow::Result<()>;
450
451 fn meta(&self, key: &str) -> Option<String> {
453 self.metadata().remove(key)
454 }
455
456 fn set_meta(&mut self, key: &str, value: &str) -> anyhow::Result<()> {
458 let mut meta = self.metadata();
459 meta.insert(key.to_string(), value.to_string());
460 self.set_metadata(&meta)
461 }
462
463 fn subset_columns_rows(
471 &mut self,
472 columns: Option<&Vec<usize>>,
473 rows: Option<&Vec<usize>>,
474 ) -> anyhow::Result<()> {
475 let ncol_data = self
476 .num_columns()
477 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
478 let nrow_data = self
479 .num_rows()
480 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
481
482 let distinct = |sel: &[usize], what: &str| -> anyhow::Result<()> {
491 anyhow::ensure!(!sel.is_empty(), "subset: empty {what} selection");
492 let mut seen = sel.to_vec();
493 seen.sort_unstable();
494 seen.dedup();
495 anyhow::ensure!(
496 seen.len() == sel.len(),
497 "subset: the {what} selection repeats an index ({} of {} are distinct)",
498 seen.len(),
499 sel.len()
500 );
501 Ok(())
502 };
503 if let Some(cols) = columns {
504 distinct(cols, "column")?;
505 }
506 if let Some(rs) = rows {
507 distinct(rs, "row")?;
508 }
509
510 let (old2new_cols, new_col_names) =
515 take_subset_indices_names_if_needed(columns, Some(ncol_data), self.column_names()?);
516 let (old2new_rows, new_row_names) =
517 take_subset_indices_names_if_needed(rows, Some(nrow_data), self.row_names()?);
518 let (new_ncol, new_nrow) = (new_col_names.len(), new_row_names.len());
519 anyhow::ensure!(new_ncol > 0, "subset: no column survived the selection");
520 anyhow::ensure!(new_nrow > 0, "subset: no row survived the selection");
521
522 let mut cols_new_order: Vec<(u64, u64)> =
524 old2new_cols.iter().map(|(&o, &n)| (n, o)).collect();
525 cols_new_order.sort_unstable();
526
527 let mut row_map: Vec<Option<u64>> = vec![None; nrow_data];
532 for (&old, &new) in &old2new_rows {
533 row_map[old as usize] = Some(new);
534 }
535 let monotone_rows = row_map.iter().flatten().is_sorted_by(|a, b| a < b);
536
537 let full_rows = rows.is_none();
545 let per_col_nnz: Vec<u64> = if full_rows {
546 cols_new_order
547 .iter()
548 .map(|&(_, old)| {
549 self.column_nnz(old as usize)
550 .ok_or_else(|| anyhow::anyhow!("subset: no indptr for column {old}"))
551 })
552 .collect::<anyhow::Result<_>>()?
553 } else {
554 let mut counts = vec![0u64; cols_new_order.len()];
559 let coarse = legume_numeric::matrix::utils::generate_minibatch_intervals(
560 cols_new_order.len(),
561 0,
562 Some(8192),
563 );
564 for (lb, ub) in coarse {
565 let old_cols: Vec<usize> = cols_new_order[lb..ub]
566 .iter()
567 .map(|&(_, o)| o as usize)
568 .collect();
569 let (_, _, triplets) =
570 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
571 for (i, c_local, _) in triplets {
572 if row_map[i as usize].is_some() {
573 counts[lb + c_local as usize] += 1;
574 }
575 }
576 }
577 counts
578 };
579 let new_nnz: u64 = per_col_nnz.iter().sum();
580
581 let final_path = self.get_backend_file_name().to_string();
599 anyhow::ensure!(
600 !final_path.ends_with(".zip"),
601 "subset: {final_path} is a zip archive; convert it to a directory \
602 backend first (data-beans convert)"
603 );
604 let temp_path = format!("{final_path}.subset_tmp");
605 if std::path::Path::new(&temp_path).exists() {
606 crate::sparse_io::remove_backend_path(&temp_path)?;
607 }
608
609 {
610 let backend_kind = self.backend_type();
611 let mut out = crate::sparse_io::create_sparse_streaming_empty(
612 Some(&temp_path),
613 Some(&backend_kind),
614 )?;
615 out.begin_streaming_csc((new_nrow, new_ncol, new_nnz as usize))?;
616
617 let blocks = legume_numeric::matrix::utils::byte_budget_intervals(
620 &per_col_nnz,
621 crate::sparse_io::SLAB_BUDGET_BYTES,
622 crate::sparse_io::TRIPLET_BYTES,
623 );
624
625 let mut nnz_offset = 0u64;
626 for (lb, ub) in blocks {
627 let old_cols: Vec<usize> = cols_new_order[lb..ub]
631 .iter()
632 .map(|&(_, o)| o as usize)
633 .collect();
634 let (_, _, triplets) =
635 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
636
637 let n_block = ub - lb;
638 let mut per_col: Vec<Vec<(u64, f32)>> = vec![Vec::new(); n_block];
639 for (i, c_local, x) in triplets {
640 if let Some(new_row) = row_map[i as usize] {
641 per_col[c_local as usize].push((new_row, x));
642 }
643 }
644 let mut local_colptr = Vec::with_capacity(n_block);
645 let mut row_indices = Vec::new();
646 let mut values = Vec::new();
647 for entries in &mut per_col {
648 if !monotone_rows {
649 entries.sort_unstable_by_key(|&(r, _)| r);
652 }
653 local_colptr.push(row_indices.len() as u64);
654 for &(r, x) in entries.iter() {
655 row_indices.push(r);
656 values.push(x);
657 }
658 }
659 out.append_csc_slab(lb as u64, nnz_offset, &local_colptr, &row_indices, &values)?;
660 nnz_offset += values.len() as u64;
661 }
662
663 out.finalize_streaming_csc()?;
664 out.build_csr_from_csc_streaming()?;
665 out.register_row_names_vec(&new_row_names);
666 out.register_column_names_vec(&new_col_names);
667 out.set_metadata(&self.metadata())?;
668 }
669
670 self.remove_backend_file()?;
675 std::fs::rename(&temp_path, &final_path)?;
676 self.reopen_backend()?;
677 self.clean_preloaded_columns();
678 self.clean_preloaded_rows();
679 info!("registered new data to {}", self.get_backend_file_name());
680 Ok(())
681 }
682
683 fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()> {
686 let meta = self.metadata();
688 let new_col_names = self.column_names()?.clone();
689 let name2new = build_name2index_map(row_names_order);
690
691 let block_size = 100;
692
693 let old2new: HashMap<u64, u64> = self
694 .row_names()?
695 .into_par_iter()
696 .enumerate()
697 .filter_map(|(idx_old, name)| {
698 name2new
699 .get(&name)
700 .map(|&idx_new| (idx_old as u64, idx_new as u64))
701 })
702 .collect();
703
704 if let Some(ncol) = self.num_columns() {
705 let arc_triplets = Arc::new(Mutex::new(vec![]));
710
711 let nblock = ncol.div_ceil(block_size);
712
713 info!("remapping triplets ...");
714
715 (0..nblock)
716 .into_par_iter()
717 .progress_with(styled_progress_bar(nblock as u64, "blocks"))
718 .map(|b| {
719 let lb = (b * block_size) as u64;
720 let ub = ((b + 1) * block_size).min(ncol) as u64;
721 (lb, ub)
722 })
723 .for_each(|(lb, ub)| {
724 let (_, _, _triplets_b) = self
725 .read_triplets_by_columns(((lb as usize)..(ub as usize)).collect())
726 .unwrap();
727
728 let _triplets_b = _triplets_b.into_iter().filter_map(|(i, j_loc, x)| {
729 let j_glob = j_loc + lb;
730 old2new.get(&i).map(|&i_new| (i_new, j_glob, x))
731 });
732
733 {
734 let mut triplets = arc_triplets.lock().unwrap();
735 triplets.extend(_triplets_b);
736 }
737 });
738
739 self.remove_backend_file()?;
743
744 self.initialize_backend()?;
748
749 {
751 let mut row_col_val_triplets =
752 arc_triplets.lock().expect("failed to lock triplets");
753
754 let nnz = row_col_val_triplets.len();
755 debug_assert!(row_col_val_triplets.len() <= nnz); let new_nrow = row_names_order.len();
757 let mtx_shape = (new_nrow, ncol, nnz);
758
759 info!("sorting triplets ...");
760
761 self.record_mtx_shape(Some(mtx_shape))?;
762 self.record_triplets_by_col(&mut row_col_val_triplets)?;
763 self.record_triplets_by_row(&mut row_col_val_triplets)?;
764 }
765 self.read_column_indptr()?;
766 self.read_row_indptr()?;
767
768 self.register_row_names_vec(row_names_order);
769 self.register_column_names_vec(&new_col_names);
770 self.set_metadata(&meta)?;
771 info!("registered new data to {}", self.get_backend_file_name());
772 }
773
774 self.clean_preloaded_columns();
775 self.clean_preloaded_rows();
776 Ok(())
777 }
778 fn remove_backend_file(&self) -> anyhow::Result<()>;
782
783 fn initialize_backend(&mut self) -> anyhow::Result<()>;
785
786 fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()>;
787
788 fn record_triplets_by_row(
791 &mut self,
792 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
793 ) -> anyhow::Result<()> {
794 let nrow = self.num_rows().expect("should have `nrow`");
795 let ncol = self.num_columns().expect("should have `ncol`");
796 let nnz = row_col_val_triplets.len();
797
798 if nnz == 0 {
799 let csr_rowptr = vec![0u64; nrow + 1];
800 return self.record_csr_dataset_backend(&[], &[], &csr_rowptr);
801 }
802
803 check_triplet_bounds("record_triplets_by_row", row_col_val_triplets, nrow, ncol)?;
804
805 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (row, col));
809
810 self.begin_streaming_csr((nrow, ncol, nnz))?;
811
812 let mut local_rowptr: Vec<u64> = Vec::new();
813 let mut cols: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
814 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
815
816 let mut start = 0_usize;
817 let mut row_offset = 0_u64;
818 while (row_offset as usize) < nrow {
819 let (end, band_end_row) =
820 slab_end(row_col_val_triplets, start, SLAB_NNZ, nrow, |t| t.0);
821
822 local_rowptr.clear();
823 cols.clear();
824 vals.clear();
825 let mut i = start;
826 for row in row_offset..band_end_row {
827 local_rowptr.push((i - start) as u64);
828 while i < end && row_col_val_triplets[i].0 == row {
829 cols.push(row_col_val_triplets[i].1);
830 vals.push(row_col_val_triplets[i].2);
831 i += 1;
832 }
833 }
834 if i != end {
835 anyhow::bail!(
836 "record_triplets_by_row: band walk left triplet {:?} at {i}",
837 row_col_val_triplets[i]
838 );
839 }
840
841 self.append_csr_slab(row_offset, start as u64, &local_rowptr, &cols, &vals)?;
842 start = end;
843 row_offset = band_end_row;
844 }
845
846 self.finalize_streaming_csr()
847 }
848
849 fn record_triplets_by_col(
856 &mut self,
857 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
858 ) -> anyhow::Result<()> {
859 let nrow = self.num_rows().expect("should have `nrow`");
860 let ncol = self.num_columns().expect("should have `ncol`");
861 let nnz = row_col_val_triplets.len();
862
863 if nnz == 0 {
864 let csc_colptr = vec![0u64; ncol + 1];
865 return self.record_csc_dataset_backend(&[], &[], &csc_colptr);
866 }
867
868 check_triplet_bounds("record_triplets_by_col", row_col_val_triplets, nrow, ncol)?;
869
870 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (col, row));
872
873 self.begin_streaming_csc((nrow, ncol, nnz))?;
874
875 let mut local_colptr: Vec<u64> = Vec::new();
876 let mut rows: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
877 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
878
879 let mut start = 0_usize;
880 let mut col_offset = 0_u64;
881 while (col_offset as usize) < ncol {
882 let (end, band_end_col) =
883 slab_end(row_col_val_triplets, start, SLAB_NNZ, ncol, |t| t.1);
884
885 local_colptr.clear();
886 rows.clear();
887 vals.clear();
888 let mut i = start;
889 for col in col_offset..band_end_col {
890 local_colptr.push((i - start) as u64);
891 while i < end && row_col_val_triplets[i].1 == col {
892 rows.push(row_col_val_triplets[i].0);
893 vals.push(row_col_val_triplets[i].2);
894 i += 1;
895 }
896 }
897 if i != end {
898 anyhow::bail!(
899 "record_triplets_by_col: band walk left triplet {:?} at {i}",
900 row_col_val_triplets[i]
901 );
902 }
903
904 self.append_csc_slab(col_offset, start as u64, &local_colptr, &rows, &vals)?;
905 start = end;
906 col_offset = band_end_col;
907 }
908
909 self.finalize_streaming_csc()
910 }
911
912 fn record_csr_dataset_backend(
921 &mut self,
922 csr_cols: &[u64],
923 csr_vals: &[f32],
924 csr_rowptr: &[u64],
925 ) -> anyhow::Result<()>;
926
927 fn record_csc_dataset_backend(
937 &mut self,
938 csc_rows: &[u64],
939 csc_vals: &[f32],
940 csc_colptr: &[u64],
941 ) -> anyhow::Result<()>;
942
943 fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()>;
946
947 fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()>;
950
951 fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()>;
954
955 fn begin_streaming_csc(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
960 self.reset_streamed_nnz();
964 let (_, ncol, nnz) = shape;
965 self.record_mtx_shape(Some(shape))?;
966 self.cs_create(CsKey::CscData, nnz)?;
967 self.cs_create(CsKey::CscIndices, nnz)?;
968 self.cs_create(CsKey::CscIndptr, ncol + 1)?;
969 Ok(())
970 }
971
972 fn append_csc_slab(
981 &mut self,
982 col_offset: u64,
983 nnz_offset: u64,
984 local_colptr: &[u64],
985 row_indices: &[u64],
986 values: &[f32],
987 ) -> anyhow::Result<()> {
988 anyhow::ensure!(
994 row_indices.len() == values.len(),
995 "append_csc_slab: {} row indices vs {} values",
996 row_indices.len(),
997 values.len()
998 );
999 anyhow::ensure!(
1000 local_colptr.first().copied() == Some(0) || local_colptr.is_empty(),
1001 "append_csc_slab: local_colptr must start at 0"
1002 );
1003 anyhow::ensure!(
1004 local_colptr.windows(2).all(|w| w[0] <= w[1]),
1005 "append_csc_slab: local_colptr must be monotone non-decreasing"
1006 );
1007 if let Some(&last) = local_colptr.last() {
1008 anyhow::ensure!(
1009 last <= values.len() as u64,
1010 "append_csc_slab: colptr claims {last} entries, slab holds {}",
1011 values.len()
1012 );
1013 }
1014 if let Some(nrow) = self.num_rows() {
1015 if let Some(&bad) = row_indices.iter().find(|&&r| r >= nrow as u64) {
1016 anyhow::bail!("append_csc_slab: row index {bad} outside the {nrow}-row matrix");
1017 }
1018 }
1019 for (c, &start) in local_colptr.iter().enumerate() {
1022 let end = local_colptr
1023 .get(c + 1)
1024 .copied()
1025 .unwrap_or(values.len() as u64) as usize;
1026 anyhow::ensure!(
1027 row_indices[start as usize..end]
1028 .windows(2)
1029 .all(|w| w[0] < w[1]),
1030 "append_csc_slab: rows within column {} of this band must be \
1031 strictly ascending — repeated rows usually mean duplicate \
1032 (row, col) coordinates in the source (an MTX with repeated \
1033 entries, or a union remap folding rows together)",
1034 col_offset as usize + c
1035 );
1036 }
1037
1038 let shifted: Vec<u64> = local_colptr.iter().map(|&p| p + nnz_offset).collect();
1039 self.cs_write_u64(CsKey::CscIndptr, col_offset, &shifted)?;
1040 self.cs_write_u64(CsKey::CscIndices, nnz_offset, row_indices)?;
1041 self.cs_write_f32(CsKey::CscData, nnz_offset, values)?;
1042 self.note_streamed_nnz(values.len() as u64);
1043 Ok(())
1044 }
1045
1046 fn finalize_streaming_csc(&mut self) -> anyhow::Result<()> {
1049 let ncol = self
1050 .num_columns()
1051 .ok_or_else(|| anyhow::anyhow!("ncol not set before finalize_streaming_csc"))?;
1052 let nnz = self
1053 .num_non_zeros()
1054 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csc"))?;
1055 self.cs_write_u64(CsKey::CscIndptr, ncol as u64, &[nnz as u64])?;
1056 self.read_column_indptr()?;
1057
1058 index_audit::check_indptr("finalize_streaming_csc", self.column_indptr(), ncol, nnz)
1067 .context(
1068 "the slabs were appended with a gap or overlap in their nnz offsets, \
1069 or the first slab was never appended",
1070 )?;
1071 let appended = self.streamed_nnz();
1077 anyhow::ensure!(
1078 appended == nnz as u64,
1079 "finalize_streaming_csc: {appended} entries appended but {nnz} \
1080 declared — the difference reads back as fill values wearing real \
1081 entries' positions"
1082 );
1083 Ok(())
1084 }
1085
1086 fn begin_streaming_csr(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
1089 self.reset_streamed_nnz();
1090 let (nrow, _, nnz) = shape;
1091 self.record_mtx_shape(Some(shape))?;
1092 self.cs_create(CsKey::CsrData, nnz)?;
1093 self.cs_create(CsKey::CsrIndices, nnz)?;
1094 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1095 Ok(())
1096 }
1097
1098 fn append_csr_slab(
1108 &mut self,
1109 row_offset: u64,
1110 nnz_offset: u64,
1111 local_rowptr: &[u64],
1112 col_indices: &[u64],
1113 values: &[f32],
1114 ) -> anyhow::Result<()> {
1115 anyhow::ensure!(
1116 col_indices.len() == values.len(),
1117 "append_csr_slab: {} column indices vs {} values",
1118 col_indices.len(),
1119 values.len()
1120 );
1121 anyhow::ensure!(
1122 local_rowptr.first().copied() == Some(0) || local_rowptr.is_empty(),
1123 "append_csr_slab: local_rowptr must start at 0"
1124 );
1125 anyhow::ensure!(
1126 local_rowptr.windows(2).all(|w| w[0] <= w[1]),
1127 "append_csr_slab: local_rowptr must be monotone non-decreasing"
1128 );
1129 if let Some(&last) = local_rowptr.last() {
1130 anyhow::ensure!(
1131 last <= values.len() as u64,
1132 "append_csr_slab: rowptr claims {last} entries, slab holds {}",
1133 values.len()
1134 );
1135 }
1136 if let Some(ncol) = self.num_columns() {
1137 if let Some(&bad) = col_indices.iter().find(|&&c| c >= ncol as u64) {
1138 anyhow::bail!(
1139 "append_csr_slab: column index {bad} outside the {ncol}-column matrix"
1140 );
1141 }
1142 }
1143 for (r, &start) in local_rowptr.iter().enumerate() {
1144 let end = local_rowptr
1145 .get(r + 1)
1146 .copied()
1147 .unwrap_or(values.len() as u64) as usize;
1148 anyhow::ensure!(
1149 col_indices[start as usize..end]
1150 .windows(2)
1151 .all(|w| w[0] < w[1]),
1152 "append_csr_slab: columns within row {} of this band must be \
1153 strictly ascending — repeated columns usually mean duplicate \
1154 (row, col) coordinates in the source",
1155 row_offset as usize + r
1156 );
1157 }
1158
1159 let shifted: Vec<u64> = local_rowptr.iter().map(|&p| p + nnz_offset).collect();
1160 self.cs_write_u64(CsKey::CsrIndptr, row_offset, &shifted)?;
1161 self.cs_write_u64(CsKey::CsrIndices, nnz_offset, col_indices)?;
1162 self.cs_write_f32(CsKey::CsrData, nnz_offset, values)?;
1163 self.note_streamed_nnz(values.len() as u64);
1164 Ok(())
1165 }
1166
1167 fn finalize_streaming_csr(&mut self) -> anyhow::Result<()> {
1171 let nrow = self
1172 .num_rows()
1173 .ok_or_else(|| anyhow::anyhow!("nrow not set before finalize_streaming_csr"))?;
1174 let nnz = self
1175 .num_non_zeros()
1176 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csr"))?;
1177 self.cs_write_u64(CsKey::CsrIndptr, nrow as u64, &[nnz as u64])?;
1178 self.read_row_indptr()?;
1179
1180 index_audit::check_indptr("finalize_streaming_csr", self.row_indptr(), nrow, nnz).context(
1182 "the slabs were appended with a gap or overlap in their nnz offsets, \
1183 or the first slab was never appended",
1184 )?;
1185
1186 let appended = self.streamed_nnz();
1187 anyhow::ensure!(
1188 appended == nnz as u64,
1189 "finalize_streaming_csr: {appended} entries appended but {nnz} \
1190 declared — the slabs did not cover the matrix"
1191 );
1192 Ok(())
1193 }
1194
1195 fn build_csr_from_csc_streaming(&mut self) -> anyhow::Result<()> {
1199 let nrow = self
1200 .num_rows()
1201 .ok_or_else(|| anyhow::anyhow!("nrow not set before build_csr_from_csc_streaming"))?;
1202 let ncol = self
1203 .num_columns()
1204 .ok_or_else(|| anyhow::anyhow!("ncol not set before build_csr_from_csc_streaming"))?;
1205 let nnz = self
1206 .num_non_zeros()
1207 .ok_or_else(|| anyhow::anyhow!("nnz not set before build_csr_from_csc_streaming"))?;
1208
1209 if nnz == 0 {
1210 self.cs_create(CsKey::CsrData, 0)?;
1211 self.cs_create(CsKey::CsrIndices, 0)?;
1212 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1213 let zeros = vec![0u64; nrow + 1];
1214 self.cs_write_u64(CsKey::CsrIndptr, 0, &zeros)?;
1215 self.read_row_indptr()?;
1216 return Ok(());
1217 }
1218
1219 const COL_BLOCK: usize = 1024;
1220 let n_col_blocks = ncol.div_ceil(COL_BLOCK);
1221 let bar1 = styled_progress_bar(n_col_blocks as u64, "transpose count");
1222 let mut row_counts = vec![0u64; nrow];
1223 let mut col_lo = 0usize;
1224 while col_lo < ncol {
1225 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1226 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1227 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1228 for t in &triplets {
1229 let slot = row_counts.get_mut(t.0 as usize).ok_or_else(|| {
1230 anyhow::anyhow!(
1231 "build_csr_from_csc_streaming: row index {} (0x{:016x}) outside \
1232 the {nrow}-row matrix",
1233 t.0,
1234 t.0
1235 )
1236 })?;
1237 *slot += 1;
1238 }
1239 col_lo = col_hi;
1240 bar1.inc(1);
1241 }
1242 bar1.finish_and_clear();
1243
1244 let mut rowptr = vec![0u64; nrow + 1];
1245 let mut acc = 0u64;
1246 for i in 0..nrow {
1247 rowptr[i] = acc;
1248 acc += row_counts[i];
1249 }
1250 rowptr[nrow] = acc;
1251 anyhow::ensure!(
1252 acc == nnz as u64,
1253 "build_csr_from_csc_streaming: counted {acc} entries in the CSC arrays, \
1254 but {nnz} are declared"
1255 );
1256
1257 self.cs_create(CsKey::CsrData, nnz)?;
1258 self.cs_create(CsKey::CsrIndices, nnz)?;
1259 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1260 self.cs_write_u64(CsKey::CsrIndptr, 0, &rowptr)?;
1261
1262 const TRANSPOSE_BAND_BYTES: usize = 256 * 1024 * 1024;
1266 let avg_density = nnz.div_ceil(nrow.max(1));
1267 let band_rows = (TRANSPOSE_BAND_BYTES / (12 * avg_density.max(1)))
1268 .max(1)
1269 .min(nrow);
1270 let n_bands = nrow.div_ceil(band_rows);
1271
1272 let bar2 = styled_progress_bar(n_bands as u64, "transpose scatter");
1273 let mut band_lo = 0usize;
1274 while band_lo < nrow {
1275 let band_hi = (band_lo + band_rows).min(nrow);
1276 let band_nnz_start = rowptr[band_lo];
1277 let band_nnz_end = rowptr[band_hi];
1278 let band_nnz = (band_nnz_end - band_nnz_start) as usize;
1279
1280 if band_nnz == 0 {
1281 band_lo = band_hi;
1282 bar2.inc(1);
1283 continue;
1284 }
1285
1286 let mut out_indices = vec![0u64; band_nnz];
1287 let mut out_values = vec![0f32; band_nnz];
1288 let mut cursor = vec![0u64; band_hi - band_lo];
1289
1290 let mut col_lo = 0usize;
1291 while col_lo < ncol {
1292 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1293 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1294 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1295 for &(row_i, col_j_local, x) in &triplets {
1296 let row_i_us = row_i as usize;
1297 if row_i_us >= band_lo && row_i_us < band_hi {
1298 let band_idx = row_i_us - band_lo;
1299 let col_j_global = col_j_local + col_lo as u64;
1304 let offset_in_band =
1305 (rowptr[band_lo + band_idx] - band_nnz_start) + cursor[band_idx];
1306 out_indices[offset_in_band as usize] = col_j_global;
1307 out_values[offset_in_band as usize] = x;
1308 cursor[band_idx] += 1;
1309 }
1310 }
1311 col_lo = col_hi;
1312 }
1313
1314 self.cs_write_u64(CsKey::CsrIndices, band_nnz_start, &out_indices)?;
1315 self.cs_write_f32(CsKey::CsrData, band_nnz_start, &out_values)?;
1316
1317 band_lo = band_hi;
1318 bar2.inc(1);
1319 }
1320 bar2.finish_and_clear();
1321
1322 self.read_row_indptr()?;
1323 Ok(())
1324 }
1325
1326 fn read_row_indptr(&mut self) -> anyhow::Result<()>;
1328
1329 fn read_column_indptr(&mut self) -> anyhow::Result<()>;
1331
1332 fn preload_columns(&mut self) -> anyhow::Result<()>;
1334
1335 fn clean_preloaded_columns(&mut self);
1337
1338 fn preload_rows(&mut self) -> anyhow::Result<()>;
1340
1341 fn clean_preloaded_rows(&mut self);
1343
1344 fn get_backend_file_name(&self) -> &str;
1346
1347 fn backend_type(&self) -> SparseIoBackend;
1349}