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::*;
16
17use crate::sparse_data_visitors::styled_progress_bar;
18use clap::ValueEnum;
19use indicatif::ParallelProgressIterator;
20use legume_numeric::matrix::mtx_io::*;
21use legume_numeric::matrix::traits::*;
22use log::info;
23use rayon::prelude::*;
24use rustc_hash::FxHashMap as HashMap;
25use std::ops::Range;
26use std::sync::{Arc, Mutex};
27
28#[cfg(test)]
29mod tests;
30
31#[derive(ValueEnum, Clone, Debug, PartialEq)]
32#[clap(rename_all = "lowercase")]
33pub enum SparseIoBackend {
34 Zarr,
35 HDF5,
36}
37
38#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42pub enum CsKey {
43 CscData,
44 CscIndices,
45 CscIndptr,
46 CsrData,
47 CsrIndices,
48 CsrIndptr,
49}
50
51const SLAB_NNZ: usize = 1 << 20;
55
56fn slab_end(
63 triplets: &[(u64, u64, f32)],
64 start: usize,
65 slab_nnz: usize,
66 n_major: usize,
67 major: impl Fn(&(u64, u64, f32)) -> u64,
68) -> (usize, u64) {
69 debug_assert!(slab_nnz > 0);
70 let nnz = triplets.len();
71 let mut end = (start + slab_nnz).min(nnz);
72 while end < nnz && major(&triplets[end]) == major(&triplets[end - 1]) {
73 end += 1;
74 }
75 let band_end = if end == nnz {
76 n_major as u64
77 } else {
78 major(&triplets[end])
79 };
80 (end, band_end)
81}
82
83pub trait SparseIo: Sync + Send {
84 type IndexIter: IntoIterator<Item = usize> + FromIterator<usize>;
85
86 #[cfg(feature = "ndarray")]
91 fn read_columns_ndarray(&self, columns: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
95 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
96 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
97 }
98
99 #[cfg(feature = "tensor")]
100 fn read_columns_tensor(&self, columns: Self::IndexIter) -> anyhow::Result<Tensor> {
104 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
105 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
106 }
107
108 fn read_columns_dmatrix(&self, columns: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
112 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
113 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
114 }
115
116 fn read_columns_csr(&self, columns: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
120 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
121 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
122 }
123
124 fn read_columns_csc(&self, columns: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
128 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
129 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
130 }
131
132 fn csc_column_arrays(&self) -> Option<(&[u64], &[u64], &[f32])> {
138 None
139 }
140
141 #[cfg(feature = "ndarray")]
142 fn read_rows_ndarray(&self, rows: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
146 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
147 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
148 }
149
150 #[cfg(feature = "tensor")]
151 fn read_rows_tensor(&self, rows: Self::IndexIter) -> anyhow::Result<Tensor> {
155 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
156 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
157 }
158
159 fn read_rows_dmatrix(&self, rows: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
163 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
164 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
165 }
166
167 fn read_rows_csr(&self, rows: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
171 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
172 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
173 }
174
175 fn read_rows_csc(&self, rows: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
179 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
180 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
181 }
182
183 fn import_mtx_file(&mut self, mtx_file: &str, index_by_row: bool) -> anyhow::Result<()> {
193 let (mut mtx_triplets, mtx_shape) = read_mtx_triplets(mtx_file)?;
194 info!("read mtx file: {}", mtx_file);
195 if mtx_triplets.is_empty() {
196 return Err(anyhow::anyhow!("No data in mtx file"));
197 }
198 self.record_mtx_shape(Some(mtx_shape))?;
199 info!("recording the column index");
200 self.record_triplets_by_col(&mut mtx_triplets)?;
201 if index_by_row {
202 info!("recording the row index");
203 self.record_triplets_by_row(&mut mtx_triplets)?;
204 }
205 Ok(())
206 }
207
208 fn import_dmatrix_by_row(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
215 let (nrow, ncol) = matrix.shape();
216 let mut mtx_triplets = dmatrix_to_triplets(matrix);
217 let mtx_shape = (nrow, ncol, mtx_triplets.len());
218 self.record_mtx_shape(Some(mtx_shape))?;
219 self.record_triplets_by_row(&mut mtx_triplets)
220 }
221
222 fn import_dmatrix_by_col(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
225 let (nrow, ncol) = matrix.shape();
226 let mut mtx_triplets = dmatrix_to_triplets(matrix);
227 let mtx_shape = (nrow, ncol, mtx_triplets.len());
228 self.record_mtx_shape(Some(mtx_shape))?;
229 self.record_triplets_by_col(&mut mtx_triplets)
230 }
231
232 #[cfg(feature = "ndarray")]
237 fn import_ndarray_by_row(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
240 let nrow = array.shape()[0];
241 let ncol = array.shape()[1];
242
243 let mut mtx_triplets = ndarray_to_triplets(array);
245
246 let nnz = mtx_triplets.len();
247 let mtx_shape = (nrow, ncol, nnz);
248 self.record_mtx_shape(Some(mtx_shape))?;
249
250 self.record_triplets_by_row(&mut mtx_triplets)
253 }
254
255 #[cfg(feature = "ndarray")]
256 fn import_ndarray_by_col(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
259 let nrow = array.shape()[0];
260 let ncol = array.shape()[1];
261
262 let mut mtx_triplets = ndarray_to_triplets(array);
264
265 let nnz = mtx_triplets.len();
266 let mtx_shape = (nrow, ncol, nnz);
267 self.record_mtx_shape(Some(mtx_shape))?;
268
269 self.record_triplets_by_col(&mut mtx_triplets)
272 }
273
274 #[allow(clippy::type_complexity)]
282 fn read_triplets_by_rows(
283 &self,
284 rows: Self::IndexIter,
285 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
286
287 #[allow(clippy::type_complexity)]
291 fn read_triplets_by_columns(
292 &self,
293 columns: Self::IndexIter,
294 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
295
296 #[allow(clippy::type_complexity)]
300 fn read_triplets_by_single_column(
301 &self,
302 col: usize,
303 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
304
305 fn to_mtx_file(&self, mtx_file: &str) -> anyhow::Result<()>;
308
309 fn num_rows(&self) -> Option<usize>;
311
312 fn num_columns(&self) -> Option<usize>;
314
315 fn num_non_zeros(&self) -> Option<usize>;
317
318 fn reopen_backend(&mut self) -> anyhow::Result<()>;
323
324 #[doc(hidden)]
328 fn note_streamed_nnz(&mut self, n: u64);
336
337 #[doc(hidden)]
339 fn streamed_nnz(&self) -> u64;
340
341 #[doc(hidden)]
343 fn reset_streamed_nnz(&mut self);
344
345 fn column_indptr(&self) -> &[u64];
350
351 fn column_nnz(&self, col: usize) -> Option<u64> {
357 let indptr = self.column_indptr();
358 let hi = *indptr.get(col + 1)?;
359 let lo = *indptr.get(col)?;
360 hi.checked_sub(lo)
361 }
362
363 fn register_row_names_file(&mut self, row_name_file: &str);
366
367 fn register_column_names_file(&mut self, column_name_file: &str);
370
371 fn register_row_names_vec(&mut self, rows: &[Box<str>]);
374
375 fn register_column_names_vec(&mut self, columns: &[Box<str>]);
378
379 fn register_names_file(
385 &mut self,
386 key: &str,
387 name_file: &str,
388 name_columns: Range<usize>,
389 name_sep: &str,
390 ) -> anyhow::Result<()>;
391
392 fn register_names_vec(&mut self, key: &str, names: &[Box<str>]) -> anyhow::Result<()>;
396
397 fn row_names(&self) -> anyhow::Result<Vec<Box<str>>>;
398
399 fn column_names(&self) -> anyhow::Result<Vec<Box<str>>>;
400
401 fn retrieve_registered_names(&self, key: &str) -> anyhow::Result<Vec<Box<str>>>;
404
405 fn subset_columns_rows(
413 &mut self,
414 columns: Option<&Vec<usize>>,
415 rows: Option<&Vec<usize>>,
416 ) -> anyhow::Result<()> {
417 let ncol_data = self
418 .num_columns()
419 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
420 let nrow_data = self
421 .num_rows()
422 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
423
424 let distinct = |sel: &[usize], what: &str| -> anyhow::Result<()> {
433 anyhow::ensure!(!sel.is_empty(), "subset: empty {what} selection");
434 let mut seen = sel.to_vec();
435 seen.sort_unstable();
436 seen.dedup();
437 anyhow::ensure!(
438 seen.len() == sel.len(),
439 "subset: the {what} selection repeats an index ({} of {} are distinct)",
440 seen.len(),
441 sel.len()
442 );
443 Ok(())
444 };
445 if let Some(cols) = columns {
446 distinct(cols, "column")?;
447 }
448 if let Some(rs) = rows {
449 distinct(rs, "row")?;
450 }
451
452 let (old2new_cols, new_col_names) =
457 take_subset_indices_names_if_needed(columns, Some(ncol_data), self.column_names()?);
458 let (old2new_rows, new_row_names) =
459 take_subset_indices_names_if_needed(rows, Some(nrow_data), self.row_names()?);
460 let (new_ncol, new_nrow) = (new_col_names.len(), new_row_names.len());
461 anyhow::ensure!(new_ncol > 0, "subset: no column survived the selection");
462 anyhow::ensure!(new_nrow > 0, "subset: no row survived the selection");
463
464 let mut cols_new_order: Vec<(u64, u64)> =
466 old2new_cols.iter().map(|(&o, &n)| (n, o)).collect();
467 cols_new_order.sort_unstable();
468
469 let mut row_map: Vec<Option<u64>> = vec![None; nrow_data];
474 for (&old, &new) in &old2new_rows {
475 row_map[old as usize] = Some(new);
476 }
477 let monotone_rows = row_map.iter().flatten().is_sorted_by(|a, b| a < b);
478
479 let full_rows = rows.is_none();
487 let per_col_nnz: Vec<u64> = if full_rows {
488 cols_new_order
489 .iter()
490 .map(|&(_, old)| {
491 self.column_nnz(old as usize)
492 .ok_or_else(|| anyhow::anyhow!("subset: no indptr for column {old}"))
493 })
494 .collect::<anyhow::Result<_>>()?
495 } else {
496 let mut counts = vec![0u64; cols_new_order.len()];
501 let coarse = legume_numeric::matrix::utils::generate_minibatch_intervals(
502 cols_new_order.len(),
503 0,
504 Some(8192),
505 );
506 for (lb, ub) in coarse {
507 let old_cols: Vec<usize> = cols_new_order[lb..ub]
508 .iter()
509 .map(|&(_, o)| o as usize)
510 .collect();
511 let (_, _, triplets) =
512 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
513 for (i, c_local, _) in triplets {
514 if row_map[i as usize].is_some() {
515 counts[lb + c_local as usize] += 1;
516 }
517 }
518 }
519 counts
520 };
521 let new_nnz: u64 = per_col_nnz.iter().sum();
522
523 let final_path = self.get_backend_file_name().to_string();
541 anyhow::ensure!(
542 !final_path.ends_with(".zip"),
543 "subset: {final_path} is a zip archive; convert it to a directory \
544 backend first (data-beans convert)"
545 );
546 let temp_path = format!("{final_path}.subset_tmp");
547 if std::path::Path::new(&temp_path).exists() {
548 crate::sparse_io::remove_backend_path(&temp_path)?;
549 }
550
551 {
552 let backend_kind = self.backend_type();
553 let mut out = crate::sparse_io::create_sparse_streaming_empty(
554 Some(&temp_path),
555 Some(&backend_kind),
556 )?;
557 out.begin_streaming_csc((new_nrow, new_ncol, new_nnz as usize))?;
558
559 let blocks = legume_numeric::matrix::utils::byte_budget_intervals(
562 &per_col_nnz,
563 crate::sparse_io::SLAB_BUDGET_BYTES,
564 crate::sparse_io::TRIPLET_BYTES,
565 );
566
567 let mut nnz_offset = 0u64;
568 for (lb, ub) in blocks {
569 let old_cols: Vec<usize> = cols_new_order[lb..ub]
573 .iter()
574 .map(|&(_, o)| o as usize)
575 .collect();
576 let (_, _, triplets) =
577 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
578
579 let n_block = ub - lb;
580 let mut per_col: Vec<Vec<(u64, f32)>> = vec![Vec::new(); n_block];
581 for (i, c_local, x) in triplets {
582 if let Some(new_row) = row_map[i as usize] {
583 per_col[c_local as usize].push((new_row, x));
584 }
585 }
586 let mut local_colptr = Vec::with_capacity(n_block);
587 let mut row_indices = Vec::new();
588 let mut values = Vec::new();
589 for entries in &mut per_col {
590 if !monotone_rows {
591 entries.sort_unstable_by_key(|&(r, _)| r);
594 }
595 local_colptr.push(row_indices.len() as u64);
596 for &(r, x) in entries.iter() {
597 row_indices.push(r);
598 values.push(x);
599 }
600 }
601 out.append_csc_slab(lb as u64, nnz_offset, &local_colptr, &row_indices, &values)?;
602 nnz_offset += values.len() as u64;
603 }
604
605 out.finalize_streaming_csc()?;
606 out.build_csr_from_csc_streaming()?;
607 out.register_row_names_vec(&new_row_names);
608 out.register_column_names_vec(&new_col_names);
609 }
610
611 self.remove_backend_file()?;
616 std::fs::rename(&temp_path, &final_path)?;
617 self.reopen_backend()?;
618 self.clean_preloaded_columns();
619 self.clean_preloaded_rows();
620 info!("registered new data to {}", self.get_backend_file_name());
621 Ok(())
622 }
623
624 fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()> {
627 let new_col_names = self.column_names()?.clone();
628 let name2new = build_name2index_map(row_names_order);
629
630 let block_size = 100;
631
632 let old2new: HashMap<u64, u64> = self
633 .row_names()?
634 .into_par_iter()
635 .enumerate()
636 .filter_map(|(idx_old, name)| {
637 name2new
638 .get(&name)
639 .map(|&idx_new| (idx_old as u64, idx_new as u64))
640 })
641 .collect();
642
643 if let Some(ncol) = self.num_columns() {
644 let arc_triplets = Arc::new(Mutex::new(vec![]));
649
650 let nblock = ncol.div_ceil(block_size);
651
652 info!("remapping triplets ...");
653
654 (0..nblock)
655 .into_par_iter()
656 .progress_with(styled_progress_bar(nblock as u64, "blocks"))
657 .map(|b| {
658 let lb = (b * block_size) as u64;
659 let ub = ((b + 1) * block_size).min(ncol) as u64;
660 (lb, ub)
661 })
662 .for_each(|(lb, ub)| {
663 let (_, _, _triplets_b) = self
664 .read_triplets_by_columns(((lb as usize)..(ub as usize)).collect())
665 .unwrap();
666
667 let _triplets_b = _triplets_b.into_iter().filter_map(|(i, j_loc, x)| {
668 let j_glob = j_loc + lb;
669 old2new.get(&i).map(|&i_new| (i_new, j_glob, x))
670 });
671
672 {
673 let mut triplets = arc_triplets.lock().unwrap();
674 triplets.extend(_triplets_b);
675 }
676 });
677
678 self.remove_backend_file()?;
682
683 self.initialize_backend()?;
687
688 {
690 let mut row_col_val_triplets =
691 arc_triplets.lock().expect("failed to lock triplets");
692
693 let nnz = row_col_val_triplets.len();
694 debug_assert!(row_col_val_triplets.len() <= nnz); let new_nrow = row_names_order.len();
696 let mtx_shape = (new_nrow, ncol, nnz);
697
698 info!("sorting triplets ...");
699
700 self.record_mtx_shape(Some(mtx_shape))?;
701 self.record_triplets_by_col(&mut row_col_val_triplets)?;
702 self.record_triplets_by_row(&mut row_col_val_triplets)?;
703 }
704 self.read_column_indptr()?;
705 self.read_row_indptr()?;
706
707 self.register_row_names_vec(row_names_order);
708 self.register_column_names_vec(&new_col_names);
709 info!("registered new data to {}", self.get_backend_file_name());
710 }
711
712 self.clean_preloaded_columns();
713 self.clean_preloaded_rows();
714 Ok(())
715 }
716 fn remove_backend_file(&self) -> anyhow::Result<()>;
720
721 fn initialize_backend(&mut self) -> anyhow::Result<()>;
723
724 fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()>;
725
726 fn record_triplets_by_row(
729 &mut self,
730 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
731 ) -> anyhow::Result<()> {
732 let nrow = self.num_rows().expect("should have `nrow`");
733 let ncol = self.num_columns().expect("should have `ncol`");
734 let nnz = row_col_val_triplets.len();
735
736 if nnz == 0 {
737 let csr_rowptr = vec![0u64; nrow + 1];
738 return self.record_csr_dataset_backend(&[], &[], &csr_rowptr);
739 }
740
741 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (row, col));
745
746 self.begin_streaming_csr((nrow, ncol, nnz))?;
747
748 let mut local_rowptr: Vec<u64> = Vec::new();
749 let mut cols: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
750 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
751
752 let mut start = 0_usize;
753 let mut row_offset = 0_u64;
754 while (row_offset as usize) < nrow {
755 let (end, band_end_row) =
756 slab_end(row_col_val_triplets, start, SLAB_NNZ, nrow, |t| t.0);
757
758 local_rowptr.clear();
759 cols.clear();
760 vals.clear();
761 let mut i = start;
762 for row in row_offset..band_end_row {
763 local_rowptr.push((i - start) as u64);
764 while i < end && row_col_val_triplets[i].0 == row {
765 cols.push(row_col_val_triplets[i].1);
766 vals.push(row_col_val_triplets[i].2);
767 i += 1;
768 }
769 }
770 debug_assert_eq!(i, end, "every entry of the band belongs to one of its rows");
771
772 self.append_csr_slab(row_offset, start as u64, &local_rowptr, &cols, &vals)?;
773 start = end;
774 row_offset = band_end_row;
775 }
776
777 self.finalize_streaming_csr()
778 }
779
780 fn record_triplets_by_col(
787 &mut self,
788 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
789 ) -> anyhow::Result<()> {
790 let nrow = self.num_rows().expect("should have `nrow`");
791 let ncol = self.num_columns().expect("should have `ncol`");
792 let nnz = row_col_val_triplets.len();
793
794 if nnz == 0 {
795 let csc_colptr = vec![0u64; ncol + 1];
796 return self.record_csc_dataset_backend(&[], &[], &csc_colptr);
797 }
798
799 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (col, row));
801
802 self.begin_streaming_csc((nrow, ncol, nnz))?;
803
804 let mut local_colptr: Vec<u64> = Vec::new();
805 let mut rows: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
806 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
807
808 let mut start = 0_usize;
809 let mut col_offset = 0_u64;
810 while (col_offset as usize) < ncol {
811 let (end, band_end_col) =
812 slab_end(row_col_val_triplets, start, SLAB_NNZ, ncol, |t| t.1);
813
814 local_colptr.clear();
815 rows.clear();
816 vals.clear();
817 let mut i = start;
818 for col in col_offset..band_end_col {
819 local_colptr.push((i - start) as u64);
820 while i < end && row_col_val_triplets[i].1 == col {
821 rows.push(row_col_val_triplets[i].0);
822 vals.push(row_col_val_triplets[i].2);
823 i += 1;
824 }
825 }
826 debug_assert_eq!(
827 i, end,
828 "every entry of the band belongs to one of its columns"
829 );
830
831 self.append_csc_slab(col_offset, start as u64, &local_colptr, &rows, &vals)?;
832 start = end;
833 col_offset = band_end_col;
834 }
835
836 self.finalize_streaming_csc()
837 }
838
839 fn record_csr_dataset_backend(
848 &mut self,
849 csr_cols: &[u64],
850 csr_vals: &[f32],
851 csr_rowptr: &[u64],
852 ) -> anyhow::Result<()>;
853
854 fn record_csc_dataset_backend(
864 &mut self,
865 csc_rows: &[u64],
866 csc_vals: &[f32],
867 csc_colptr: &[u64],
868 ) -> anyhow::Result<()>;
869
870 fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()>;
873
874 fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()>;
877
878 fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()>;
881
882 fn begin_streaming_csc(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
887 self.reset_streamed_nnz();
891 let (_, ncol, nnz) = shape;
892 self.record_mtx_shape(Some(shape))?;
893 self.cs_create(CsKey::CscData, nnz)?;
894 self.cs_create(CsKey::CscIndices, nnz)?;
895 self.cs_create(CsKey::CscIndptr, ncol + 1)?;
896 Ok(())
897 }
898
899 fn append_csc_slab(
908 &mut self,
909 col_offset: u64,
910 nnz_offset: u64,
911 local_colptr: &[u64],
912 row_indices: &[u64],
913 values: &[f32],
914 ) -> anyhow::Result<()> {
915 anyhow::ensure!(
921 row_indices.len() == values.len(),
922 "append_csc_slab: {} row indices vs {} values",
923 row_indices.len(),
924 values.len()
925 );
926 anyhow::ensure!(
927 local_colptr.first().copied() == Some(0) || local_colptr.is_empty(),
928 "append_csc_slab: local_colptr must start at 0"
929 );
930 anyhow::ensure!(
931 local_colptr.windows(2).all(|w| w[0] <= w[1]),
932 "append_csc_slab: local_colptr must be monotone non-decreasing"
933 );
934 if let Some(&last) = local_colptr.last() {
935 anyhow::ensure!(
936 last <= values.len() as u64,
937 "append_csc_slab: colptr claims {last} entries, slab holds {}",
938 values.len()
939 );
940 }
941 if let Some(nrow) = self.num_rows() {
942 if let Some(&bad) = row_indices.iter().find(|&&r| r >= nrow as u64) {
943 anyhow::bail!("append_csc_slab: row index {bad} outside the {nrow}-row matrix");
944 }
945 }
946 for (c, &start) in local_colptr.iter().enumerate() {
949 let end = local_colptr
950 .get(c + 1)
951 .copied()
952 .unwrap_or(values.len() as u64) as usize;
953 anyhow::ensure!(
954 row_indices[start as usize..end]
955 .windows(2)
956 .all(|w| w[0] < w[1]),
957 "append_csc_slab: rows within column {} of this band must be \
958 strictly ascending — repeated rows usually mean duplicate \
959 (row, col) coordinates in the source (an MTX with repeated \
960 entries, or a union remap folding rows together)",
961 col_offset as usize + c
962 );
963 }
964
965 let shifted: Vec<u64> = local_colptr.iter().map(|&p| p + nnz_offset).collect();
966 self.cs_write_u64(CsKey::CscIndptr, col_offset, &shifted)?;
967 self.cs_write_u64(CsKey::CscIndices, nnz_offset, row_indices)?;
968 self.cs_write_f32(CsKey::CscData, nnz_offset, values)?;
969 self.note_streamed_nnz(values.len() as u64);
970 Ok(())
971 }
972
973 fn finalize_streaming_csc(&mut self) -> anyhow::Result<()> {
976 let ncol = self
977 .num_columns()
978 .ok_or_else(|| anyhow::anyhow!("ncol not set before finalize_streaming_csc"))?;
979 let nnz = self
980 .num_non_zeros()
981 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csc"))?;
982 self.cs_write_u64(CsKey::CscIndptr, ncol as u64, &[nnz as u64])?;
983 self.read_column_indptr()?;
984
985 let indptr = self.column_indptr();
994 anyhow::ensure!(
995 indptr.len() == ncol + 1,
996 "finalize_streaming_csc: indptr has {} entries, expected {}",
997 indptr.len(),
998 ncol + 1
999 );
1000 anyhow::ensure!(
1001 indptr.first().copied() == Some(0),
1002 "finalize_streaming_csc: indptr[0] = {:?}, expected 0 — the first \
1003 slab was never appended",
1004 indptr.first()
1005 );
1006 if let Some(w) = indptr.windows(2).position(|w| w[0] > w[1]) {
1007 anyhow::bail!(
1008 "finalize_streaming_csc: indptr decreases at column {w} — slabs \
1009 were appended with a gap or overlap in their nnz offsets"
1010 );
1011 }
1012 let appended = self.streamed_nnz();
1018 anyhow::ensure!(
1019 appended == nnz as u64,
1020 "finalize_streaming_csc: {appended} entries appended but {nnz} \
1021 declared — the difference reads back as fill values wearing real \
1022 entries' positions"
1023 );
1024 Ok(())
1025 }
1026
1027 fn begin_streaming_csr(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
1030 self.reset_streamed_nnz();
1031 let (nrow, _, nnz) = shape;
1032 self.record_mtx_shape(Some(shape))?;
1033 self.cs_create(CsKey::CsrData, nnz)?;
1034 self.cs_create(CsKey::CsrIndices, nnz)?;
1035 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1036 Ok(())
1037 }
1038
1039 fn append_csr_slab(
1049 &mut self,
1050 row_offset: u64,
1051 nnz_offset: u64,
1052 local_rowptr: &[u64],
1053 col_indices: &[u64],
1054 values: &[f32],
1055 ) -> anyhow::Result<()> {
1056 anyhow::ensure!(
1057 col_indices.len() == values.len(),
1058 "append_csr_slab: {} column indices vs {} values",
1059 col_indices.len(),
1060 values.len()
1061 );
1062 anyhow::ensure!(
1063 local_rowptr.first().copied() == Some(0) || local_rowptr.is_empty(),
1064 "append_csr_slab: local_rowptr must start at 0"
1065 );
1066 anyhow::ensure!(
1067 local_rowptr.windows(2).all(|w| w[0] <= w[1]),
1068 "append_csr_slab: local_rowptr must be monotone non-decreasing"
1069 );
1070 if let Some(&last) = local_rowptr.last() {
1071 anyhow::ensure!(
1072 last <= values.len() as u64,
1073 "append_csr_slab: rowptr claims {last} entries, slab holds {}",
1074 values.len()
1075 );
1076 }
1077 if let Some(ncol) = self.num_columns() {
1078 if let Some(&bad) = col_indices.iter().find(|&&c| c >= ncol as u64) {
1079 anyhow::bail!(
1080 "append_csr_slab: column index {bad} outside the {ncol}-column matrix"
1081 );
1082 }
1083 }
1084 for (r, &start) in local_rowptr.iter().enumerate() {
1085 let end = local_rowptr
1086 .get(r + 1)
1087 .copied()
1088 .unwrap_or(values.len() as u64) as usize;
1089 anyhow::ensure!(
1090 col_indices[start as usize..end]
1091 .windows(2)
1092 .all(|w| w[0] < w[1]),
1093 "append_csr_slab: columns within row {} of this band must be \
1094 strictly ascending — repeated columns usually mean duplicate \
1095 (row, col) coordinates in the source",
1096 row_offset as usize + r
1097 );
1098 }
1099
1100 let shifted: Vec<u64> = local_rowptr.iter().map(|&p| p + nnz_offset).collect();
1101 self.cs_write_u64(CsKey::CsrIndptr, row_offset, &shifted)?;
1102 self.cs_write_u64(CsKey::CsrIndices, nnz_offset, col_indices)?;
1103 self.cs_write_f32(CsKey::CsrData, nnz_offset, values)?;
1104 self.note_streamed_nnz(values.len() as u64);
1105 Ok(())
1106 }
1107
1108 fn finalize_streaming_csr(&mut self) -> anyhow::Result<()> {
1112 let nrow = self
1113 .num_rows()
1114 .ok_or_else(|| anyhow::anyhow!("nrow not set before finalize_streaming_csr"))?;
1115 let nnz = self
1116 .num_non_zeros()
1117 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csr"))?;
1118 self.cs_write_u64(CsKey::CsrIndptr, nrow as u64, &[nnz as u64])?;
1119 self.read_row_indptr()?;
1120
1121 let appended = self.streamed_nnz();
1122 anyhow::ensure!(
1123 appended == nnz as u64,
1124 "finalize_streaming_csr: {appended} entries appended but {nnz} \
1125 declared — the slabs did not cover the matrix"
1126 );
1127 Ok(())
1128 }
1129
1130 fn build_csr_from_csc_streaming(&mut self) -> anyhow::Result<()> {
1134 let nrow = self
1135 .num_rows()
1136 .ok_or_else(|| anyhow::anyhow!("nrow not set before build_csr_from_csc_streaming"))?;
1137 let ncol = self
1138 .num_columns()
1139 .ok_or_else(|| anyhow::anyhow!("ncol not set before build_csr_from_csc_streaming"))?;
1140 let nnz = self
1141 .num_non_zeros()
1142 .ok_or_else(|| anyhow::anyhow!("nnz not set before build_csr_from_csc_streaming"))?;
1143
1144 if nnz == 0 {
1145 self.cs_create(CsKey::CsrData, 0)?;
1146 self.cs_create(CsKey::CsrIndices, 0)?;
1147 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1148 let zeros = vec![0u64; nrow + 1];
1149 self.cs_write_u64(CsKey::CsrIndptr, 0, &zeros)?;
1150 self.read_row_indptr()?;
1151 return Ok(());
1152 }
1153
1154 const COL_BLOCK: usize = 1024;
1155 let n_col_blocks = ncol.div_ceil(COL_BLOCK);
1156 let bar1 = styled_progress_bar(n_col_blocks as u64, "transpose count");
1157 let mut row_counts = vec![0u64; nrow];
1158 let mut col_lo = 0usize;
1159 while col_lo < ncol {
1160 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1161 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1162 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1163 for (row_i, _, _) in &triplets {
1164 row_counts[*row_i as usize] += 1;
1165 }
1166 col_lo = col_hi;
1167 bar1.inc(1);
1168 }
1169 bar1.finish_and_clear();
1170
1171 let mut rowptr = vec![0u64; nrow + 1];
1172 let mut acc = 0u64;
1173 for i in 0..nrow {
1174 rowptr[i] = acc;
1175 acc += row_counts[i];
1176 }
1177 rowptr[nrow] = acc;
1178 debug_assert_eq!(acc, nnz as u64);
1179
1180 self.cs_create(CsKey::CsrData, nnz)?;
1181 self.cs_create(CsKey::CsrIndices, nnz)?;
1182 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1183 self.cs_write_u64(CsKey::CsrIndptr, 0, &rowptr)?;
1184
1185 const TRANSPOSE_BAND_BYTES: usize = 256 * 1024 * 1024;
1189 let avg_density = nnz.div_ceil(nrow.max(1));
1190 let band_rows = (TRANSPOSE_BAND_BYTES / (12 * avg_density.max(1)))
1191 .max(1)
1192 .min(nrow);
1193 let n_bands = nrow.div_ceil(band_rows);
1194
1195 let bar2 = styled_progress_bar(n_bands as u64, "transpose scatter");
1196 let mut band_lo = 0usize;
1197 while band_lo < nrow {
1198 let band_hi = (band_lo + band_rows).min(nrow);
1199 let band_nnz_start = rowptr[band_lo];
1200 let band_nnz_end = rowptr[band_hi];
1201 let band_nnz = (band_nnz_end - band_nnz_start) as usize;
1202
1203 if band_nnz == 0 {
1204 band_lo = band_hi;
1205 bar2.inc(1);
1206 continue;
1207 }
1208
1209 let mut out_indices = vec![0u64; band_nnz];
1210 let mut out_values = vec![0f32; band_nnz];
1211 let mut cursor = vec![0u64; band_hi - band_lo];
1212
1213 let mut col_lo = 0usize;
1214 while col_lo < ncol {
1215 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1216 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1217 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1218 for &(row_i, col_j_local, x) in &triplets {
1219 let row_i_us = row_i as usize;
1220 if row_i_us >= band_lo && row_i_us < band_hi {
1221 let band_idx = row_i_us - band_lo;
1222 let col_j_global = col_j_local + col_lo as u64;
1227 let offset_in_band =
1228 (rowptr[band_lo + band_idx] - band_nnz_start) + cursor[band_idx];
1229 out_indices[offset_in_band as usize] = col_j_global;
1230 out_values[offset_in_band as usize] = x;
1231 cursor[band_idx] += 1;
1232 }
1233 }
1234 col_lo = col_hi;
1235 }
1236
1237 self.cs_write_u64(CsKey::CsrIndices, band_nnz_start, &out_indices)?;
1238 self.cs_write_f32(CsKey::CsrData, band_nnz_start, &out_values)?;
1239
1240 band_lo = band_hi;
1241 bar2.inc(1);
1242 }
1243 bar2.finish_and_clear();
1244
1245 self.read_row_indptr()?;
1246 Ok(())
1247 }
1248
1249 fn read_row_indptr(&mut self) -> anyhow::Result<()>;
1251
1252 fn read_column_indptr(&mut self) -> anyhow::Result<()>;
1254
1255 fn preload_columns(&mut self) -> anyhow::Result<()>;
1257
1258 fn clean_preloaded_columns(&mut self);
1260
1261 fn preload_rows(&mut self) -> anyhow::Result<()>;
1263
1264 fn clean_preloaded_rows(&mut self);
1266
1267 fn get_backend_file_name(&self) -> &str;
1269
1270 fn backend_type(&self) -> SparseIoBackend;
1272}