1#![allow(dead_code, unused_imports)]
2
3pub use legume_numeric::candle::candle_core::Tensor;
4pub use nalgebra::DMatrix;
5pub use nalgebra_sparse::{csc::CscMatrix, csr::CsrMatrix};
6pub use ndarray::prelude::*;
7
8pub const MAX_ROW_NAME_IDX: usize = 3;
9pub const MAX_COLUMN_NAME_IDX: usize = 10;
10pub const COLUMN_SEP: &str = "@";
11pub const ROW_SEP: &str = "_";
12
13use super::helpers::*;
14
15use crate::sparse_data_visitors::styled_progress_bar;
16use clap::ValueEnum;
17use indicatif::ParallelProgressIterator;
18use legume_numeric::matrix::mtx_io::*;
19use legume_numeric::matrix::traits::*;
20use log::info;
21use rayon::prelude::*;
22use rustc_hash::FxHashMap as HashMap;
23use std::ops::Range;
24use std::sync::{Arc, Mutex};
25
26#[cfg(test)]
27mod tests;
28
29#[derive(ValueEnum, Clone, Debug, PartialEq)]
30#[clap(rename_all = "lowercase")]
31pub enum SparseIoBackend {
32 Zarr,
33 HDF5,
34}
35
36#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub enum CsKey {
41 CscData,
42 CscIndices,
43 CscIndptr,
44 CsrData,
45 CsrIndices,
46 CsrIndptr,
47}
48
49const SLAB_NNZ: usize = 1 << 20;
53
54fn slab_end(
61 triplets: &[(u64, u64, f32)],
62 start: usize,
63 slab_nnz: usize,
64 n_major: usize,
65 major: impl Fn(&(u64, u64, f32)) -> u64,
66) -> (usize, u64) {
67 debug_assert!(slab_nnz > 0);
68 let nnz = triplets.len();
69 let mut end = (start + slab_nnz).min(nnz);
70 while end < nnz && major(&triplets[end]) == major(&triplets[end - 1]) {
71 end += 1;
72 }
73 let band_end = if end == nnz {
74 n_major as u64
75 } else {
76 major(&triplets[end])
77 };
78 (end, band_end)
79}
80
81pub trait SparseIo: Sync + Send {
82 type IndexIter: IntoIterator<Item = usize> + FromIterator<usize>;
83
84 fn read_columns_ndarray(&self, columns: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
92 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
93 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
94 }
95
96 fn read_columns_tensor(&self, columns: Self::IndexIter) -> anyhow::Result<Tensor> {
100 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
101 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
102 }
103
104 fn read_columns_dmatrix(&self, columns: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
108 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
109 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
110 }
111
112 fn read_columns_csr(&self, columns: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
116 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
117 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
118 }
119
120 fn read_columns_csc(&self, columns: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
124 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
125 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
126 }
127
128 fn csc_column_arrays(&self) -> Option<(&[u64], &[u64], &[f32])> {
134 None
135 }
136
137 fn read_rows_ndarray(&self, rows: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
141 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
142 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
143 }
144
145 fn read_rows_tensor(&self, rows: Self::IndexIter) -> anyhow::Result<Tensor> {
149 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
150 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
151 }
152
153 fn read_rows_dmatrix(&self, rows: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
157 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
158 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
159 }
160
161 fn read_rows_csr(&self, rows: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
165 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
166 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
167 }
168
169 fn read_rows_csc(&self, rows: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
173 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
174 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
175 }
176
177 fn import_mtx_file(&mut self, mtx_file: &str, index_by_row: bool) -> anyhow::Result<()> {
187 let (mut mtx_triplets, mtx_shape) = read_mtx_triplets(mtx_file)?;
188 info!("read mtx file: {}", mtx_file);
189 if mtx_triplets.is_empty() {
190 return Err(anyhow::anyhow!("No data in mtx file"));
191 }
192 self.record_mtx_shape(Some(mtx_shape))?;
193 info!("recording the column index");
194 self.record_triplets_by_col(&mut mtx_triplets)?;
195 if index_by_row {
196 info!("recording the row index");
197 self.record_triplets_by_row(&mut mtx_triplets)?;
198 }
199 Ok(())
200 }
201
202 fn import_dmatrix_by_row(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
209 let (nrow, ncol) = matrix.shape();
210 let mut mtx_triplets = dmatrix_to_triplets(matrix);
211 let mtx_shape = (nrow, ncol, mtx_triplets.len());
212 self.record_mtx_shape(Some(mtx_shape))?;
213 self.record_triplets_by_row(&mut mtx_triplets)
214 }
215
216 fn import_dmatrix_by_col(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
219 let (nrow, ncol) = matrix.shape();
220 let mut mtx_triplets = dmatrix_to_triplets(matrix);
221 let mtx_shape = (nrow, ncol, mtx_triplets.len());
222 self.record_mtx_shape(Some(mtx_shape))?;
223 self.record_triplets_by_col(&mut mtx_triplets)
224 }
225
226 fn import_ndarray_by_row(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
233 let nrow = array.shape()[0];
234 let ncol = array.shape()[1];
235
236 let mut mtx_triplets = ndarray_to_triplets(array);
238
239 let nnz = mtx_triplets.len();
240 let mtx_shape = (nrow, ncol, nnz);
241 self.record_mtx_shape(Some(mtx_shape))?;
242
243 self.record_triplets_by_row(&mut mtx_triplets)
246 }
247
248 fn import_ndarray_by_col(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
251 let nrow = array.shape()[0];
252 let ncol = array.shape()[1];
253
254 let mut mtx_triplets = ndarray_to_triplets(array);
256
257 let nnz = mtx_triplets.len();
258 let mtx_shape = (nrow, ncol, nnz);
259 self.record_mtx_shape(Some(mtx_shape))?;
260
261 self.record_triplets_by_col(&mut mtx_triplets)
264 }
265
266 #[allow(clippy::type_complexity)]
274 fn read_triplets_by_rows(
275 &self,
276 rows: Self::IndexIter,
277 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
278
279 #[allow(clippy::type_complexity)]
283 fn read_triplets_by_columns(
284 &self,
285 columns: Self::IndexIter,
286 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
287
288 #[allow(clippy::type_complexity)]
292 fn read_triplets_by_single_column(
293 &self,
294 col: usize,
295 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
296
297 fn to_mtx_file(&self, mtx_file: &str) -> anyhow::Result<()>;
300
301 fn num_rows(&self) -> Option<usize>;
303
304 fn num_columns(&self) -> Option<usize>;
306
307 fn num_non_zeros(&self) -> Option<usize>;
309
310 fn reopen_backend(&mut self) -> anyhow::Result<()>;
315
316 #[doc(hidden)]
320 fn note_streamed_nnz(&mut self, n: u64);
328
329 #[doc(hidden)]
331 fn streamed_nnz(&self) -> u64;
332
333 #[doc(hidden)]
335 fn reset_streamed_nnz(&mut self);
336
337 fn column_indptr(&self) -> &[u64];
342
343 fn column_nnz(&self, col: usize) -> Option<u64> {
349 let indptr = self.column_indptr();
350 let hi = *indptr.get(col + 1)?;
351 let lo = *indptr.get(col)?;
352 hi.checked_sub(lo)
353 }
354
355 fn register_row_names_file(&mut self, row_name_file: &str);
358
359 fn register_column_names_file(&mut self, column_name_file: &str);
362
363 fn register_row_names_vec(&mut self, rows: &[Box<str>]);
366
367 fn register_column_names_vec(&mut self, columns: &[Box<str>]);
370
371 fn register_names_file(
377 &mut self,
378 key: &str,
379 name_file: &str,
380 name_columns: Range<usize>,
381 name_sep: &str,
382 ) -> anyhow::Result<()>;
383
384 fn register_names_vec(&mut self, key: &str, names: &[Box<str>]) -> anyhow::Result<()>;
388
389 fn row_names(&self) -> anyhow::Result<Vec<Box<str>>>;
390
391 fn column_names(&self) -> anyhow::Result<Vec<Box<str>>>;
392
393 fn retrieve_registered_names(&self, key: &str) -> anyhow::Result<Vec<Box<str>>>;
396
397 fn subset_columns_rows(
405 &mut self,
406 columns: Option<&Vec<usize>>,
407 rows: Option<&Vec<usize>>,
408 ) -> anyhow::Result<()> {
409 let ncol_data = self
410 .num_columns()
411 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
412 let nrow_data = self
413 .num_rows()
414 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
415
416 let distinct = |sel: &[usize], what: &str| -> anyhow::Result<()> {
425 anyhow::ensure!(!sel.is_empty(), "subset: empty {what} selection");
426 let mut seen = sel.to_vec();
427 seen.sort_unstable();
428 seen.dedup();
429 anyhow::ensure!(
430 seen.len() == sel.len(),
431 "subset: the {what} selection repeats an index ({} of {} are distinct)",
432 seen.len(),
433 sel.len()
434 );
435 Ok(())
436 };
437 if let Some(cols) = columns {
438 distinct(cols, "column")?;
439 }
440 if let Some(rs) = rows {
441 distinct(rs, "row")?;
442 }
443
444 let (old2new_cols, new_col_names) =
449 take_subset_indices_names_if_needed(columns, Some(ncol_data), self.column_names()?);
450 let (old2new_rows, new_row_names) =
451 take_subset_indices_names_if_needed(rows, Some(nrow_data), self.row_names()?);
452 let (new_ncol, new_nrow) = (new_col_names.len(), new_row_names.len());
453 anyhow::ensure!(new_ncol > 0, "subset: no column survived the selection");
454 anyhow::ensure!(new_nrow > 0, "subset: no row survived the selection");
455
456 let mut cols_new_order: Vec<(u64, u64)> =
458 old2new_cols.iter().map(|(&o, &n)| (n, o)).collect();
459 cols_new_order.sort_unstable();
460
461 let mut row_map: Vec<Option<u64>> = vec![None; nrow_data];
466 for (&old, &new) in &old2new_rows {
467 row_map[old as usize] = Some(new);
468 }
469 let monotone_rows = row_map.iter().flatten().is_sorted_by(|a, b| a < b);
470
471 let full_rows = rows.is_none();
479 let per_col_nnz: Vec<u64> = if full_rows {
480 cols_new_order
481 .iter()
482 .map(|&(_, old)| {
483 self.column_nnz(old as usize)
484 .ok_or_else(|| anyhow::anyhow!("subset: no indptr for column {old}"))
485 })
486 .collect::<anyhow::Result<_>>()?
487 } else {
488 let mut counts = vec![0u64; cols_new_order.len()];
493 let coarse = legume_numeric::matrix::utils::generate_minibatch_intervals(
494 cols_new_order.len(),
495 0,
496 Some(8192),
497 );
498 for (lb, ub) in coarse {
499 let old_cols: Vec<usize> = cols_new_order[lb..ub]
500 .iter()
501 .map(|&(_, o)| o as usize)
502 .collect();
503 let (_, _, triplets) =
504 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
505 for (i, c_local, _) in triplets {
506 if row_map[i as usize].is_some() {
507 counts[lb + c_local as usize] += 1;
508 }
509 }
510 }
511 counts
512 };
513 let new_nnz: u64 = per_col_nnz.iter().sum();
514
515 let final_path = self.get_backend_file_name().to_string();
533 anyhow::ensure!(
534 !final_path.ends_with(".zip"),
535 "subset: {final_path} is a zip archive; convert it to a directory \
536 backend first (data-beans convert)"
537 );
538 let temp_path = format!("{final_path}.subset_tmp");
539 if std::path::Path::new(&temp_path).exists() {
540 crate::sparse_io::remove_backend_path(&temp_path)?;
541 }
542
543 {
544 let backend_kind = self.backend_type();
545 let mut out = crate::sparse_io::create_sparse_streaming_empty(
546 Some(&temp_path),
547 Some(&backend_kind),
548 )?;
549 out.begin_streaming_csc((new_nrow, new_ncol, new_nnz as usize))?;
550
551 let blocks = legume_numeric::matrix::utils::byte_budget_intervals(
554 &per_col_nnz,
555 crate::sparse_io::SLAB_BUDGET_BYTES,
556 crate::sparse_io::TRIPLET_BYTES,
557 );
558
559 let mut nnz_offset = 0u64;
560 for (lb, ub) in blocks {
561 let old_cols: Vec<usize> = cols_new_order[lb..ub]
565 .iter()
566 .map(|&(_, o)| o as usize)
567 .collect();
568 let (_, _, triplets) =
569 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
570
571 let n_block = ub - lb;
572 let mut per_col: Vec<Vec<(u64, f32)>> = vec![Vec::new(); n_block];
573 for (i, c_local, x) in triplets {
574 if let Some(new_row) = row_map[i as usize] {
575 per_col[c_local as usize].push((new_row, x));
576 }
577 }
578 let mut local_colptr = Vec::with_capacity(n_block);
579 let mut row_indices = Vec::new();
580 let mut values = Vec::new();
581 for entries in &mut per_col {
582 if !monotone_rows {
583 entries.sort_unstable_by_key(|&(r, _)| r);
586 }
587 local_colptr.push(row_indices.len() as u64);
588 for &(r, x) in entries.iter() {
589 row_indices.push(r);
590 values.push(x);
591 }
592 }
593 out.append_csc_slab(lb as u64, nnz_offset, &local_colptr, &row_indices, &values)?;
594 nnz_offset += values.len() as u64;
595 }
596
597 out.finalize_streaming_csc()?;
598 out.build_csr_from_csc_streaming()?;
599 out.register_row_names_vec(&new_row_names);
600 out.register_column_names_vec(&new_col_names);
601 }
602
603 self.remove_backend_file()?;
608 std::fs::rename(&temp_path, &final_path)?;
609 self.reopen_backend()?;
610 self.clean_preloaded_columns();
611 self.clean_preloaded_rows();
612 info!("registered new data to {}", self.get_backend_file_name());
613 Ok(())
614 }
615
616 fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()> {
619 let new_col_names = self.column_names()?.clone();
620 let name2new = build_name2index_map(row_names_order);
621
622 let block_size = 100;
623
624 let old2new: HashMap<u64, u64> = self
625 .row_names()?
626 .into_par_iter()
627 .enumerate()
628 .filter_map(|(idx_old, name)| {
629 name2new
630 .get(&name)
631 .map(|&idx_new| (idx_old as u64, idx_new as u64))
632 })
633 .collect();
634
635 if let Some(ncol) = self.num_columns() {
636 let arc_triplets = Arc::new(Mutex::new(vec![]));
641
642 let nblock = ncol.div_ceil(block_size);
643
644 info!("remapping triplets ...");
645
646 (0..nblock)
647 .into_par_iter()
648 .progress_with(styled_progress_bar(nblock as u64, "blocks"))
649 .map(|b| {
650 let lb = (b * block_size) as u64;
651 let ub = ((b + 1) * block_size).min(ncol) as u64;
652 (lb, ub)
653 })
654 .for_each(|(lb, ub)| {
655 let (_, _, _triplets_b) = self
656 .read_triplets_by_columns(((lb as usize)..(ub as usize)).collect())
657 .unwrap();
658
659 let _triplets_b = _triplets_b.into_iter().filter_map(|(i, j_loc, x)| {
660 let j_glob = j_loc + lb;
661 old2new.get(&i).map(|&i_new| (i_new, j_glob, x))
662 });
663
664 {
665 let mut triplets = arc_triplets.lock().unwrap();
666 triplets.extend(_triplets_b);
667 }
668 });
669
670 self.remove_backend_file()?;
674
675 self.initialize_backend()?;
679
680 {
682 let mut row_col_val_triplets =
683 arc_triplets.lock().expect("failed to lock triplets");
684
685 let nnz = row_col_val_triplets.len();
686 debug_assert!(row_col_val_triplets.len() <= nnz); let new_nrow = row_names_order.len();
688 let mtx_shape = (new_nrow, ncol, nnz);
689
690 info!("sorting triplets ...");
691
692 self.record_mtx_shape(Some(mtx_shape))?;
693 self.record_triplets_by_col(&mut row_col_val_triplets)?;
694 self.record_triplets_by_row(&mut row_col_val_triplets)?;
695 }
696 self.read_column_indptr()?;
697 self.read_row_indptr()?;
698
699 self.register_row_names_vec(row_names_order);
700 self.register_column_names_vec(&new_col_names);
701 info!("registered new data to {}", self.get_backend_file_name());
702 }
703
704 self.clean_preloaded_columns();
705 self.clean_preloaded_rows();
706 Ok(())
707 }
708 fn remove_backend_file(&self) -> anyhow::Result<()>;
712
713 fn initialize_backend(&mut self) -> anyhow::Result<()>;
715
716 fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()>;
717
718 fn record_triplets_by_row(
721 &mut self,
722 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
723 ) -> anyhow::Result<()> {
724 let nrow = self.num_rows().expect("should have `nrow`");
725 let ncol = self.num_columns().expect("should have `ncol`");
726 let nnz = row_col_val_triplets.len();
727
728 if nnz == 0 {
729 let csr_rowptr = vec![0u64; nrow + 1];
730 return self.record_csr_dataset_backend(&[], &[], &csr_rowptr);
731 }
732
733 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (row, col));
737
738 self.begin_streaming_csr((nrow, ncol, nnz))?;
739
740 let mut local_rowptr: Vec<u64> = Vec::new();
741 let mut cols: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
742 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
743
744 let mut start = 0_usize;
745 let mut row_offset = 0_u64;
746 while (row_offset as usize) < nrow {
747 let (end, band_end_row) =
748 slab_end(row_col_val_triplets, start, SLAB_NNZ, nrow, |t| t.0);
749
750 local_rowptr.clear();
751 cols.clear();
752 vals.clear();
753 let mut i = start;
754 for row in row_offset..band_end_row {
755 local_rowptr.push((i - start) as u64);
756 while i < end && row_col_val_triplets[i].0 == row {
757 cols.push(row_col_val_triplets[i].1);
758 vals.push(row_col_val_triplets[i].2);
759 i += 1;
760 }
761 }
762 debug_assert_eq!(i, end, "every entry of the band belongs to one of its rows");
763
764 self.append_csr_slab(row_offset, start as u64, &local_rowptr, &cols, &vals)?;
765 start = end;
766 row_offset = band_end_row;
767 }
768
769 self.finalize_streaming_csr()
770 }
771
772 fn record_triplets_by_col(
779 &mut self,
780 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
781 ) -> anyhow::Result<()> {
782 let nrow = self.num_rows().expect("should have `nrow`");
783 let ncol = self.num_columns().expect("should have `ncol`");
784 let nnz = row_col_val_triplets.len();
785
786 if nnz == 0 {
787 let csc_colptr = vec![0u64; ncol + 1];
788 return self.record_csc_dataset_backend(&[], &[], &csc_colptr);
789 }
790
791 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (col, row));
793
794 self.begin_streaming_csc((nrow, ncol, nnz))?;
795
796 let mut local_colptr: Vec<u64> = Vec::new();
797 let mut rows: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
798 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
799
800 let mut start = 0_usize;
801 let mut col_offset = 0_u64;
802 while (col_offset as usize) < ncol {
803 let (end, band_end_col) =
804 slab_end(row_col_val_triplets, start, SLAB_NNZ, ncol, |t| t.1);
805
806 local_colptr.clear();
807 rows.clear();
808 vals.clear();
809 let mut i = start;
810 for col in col_offset..band_end_col {
811 local_colptr.push((i - start) as u64);
812 while i < end && row_col_val_triplets[i].1 == col {
813 rows.push(row_col_val_triplets[i].0);
814 vals.push(row_col_val_triplets[i].2);
815 i += 1;
816 }
817 }
818 debug_assert_eq!(
819 i, end,
820 "every entry of the band belongs to one of its columns"
821 );
822
823 self.append_csc_slab(col_offset, start as u64, &local_colptr, &rows, &vals)?;
824 start = end;
825 col_offset = band_end_col;
826 }
827
828 self.finalize_streaming_csc()
829 }
830
831 fn record_csr_dataset_backend(
840 &mut self,
841 csr_cols: &[u64],
842 csr_vals: &[f32],
843 csr_rowptr: &[u64],
844 ) -> anyhow::Result<()>;
845
846 fn record_csc_dataset_backend(
856 &mut self,
857 csc_rows: &[u64],
858 csc_vals: &[f32],
859 csc_colptr: &[u64],
860 ) -> anyhow::Result<()>;
861
862 fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()>;
865
866 fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()>;
869
870 fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()>;
873
874 fn begin_streaming_csc(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
879 self.reset_streamed_nnz();
883 let (_, ncol, nnz) = shape;
884 self.record_mtx_shape(Some(shape))?;
885 self.cs_create(CsKey::CscData, nnz)?;
886 self.cs_create(CsKey::CscIndices, nnz)?;
887 self.cs_create(CsKey::CscIndptr, ncol + 1)?;
888 Ok(())
889 }
890
891 fn append_csc_slab(
900 &mut self,
901 col_offset: u64,
902 nnz_offset: u64,
903 local_colptr: &[u64],
904 row_indices: &[u64],
905 values: &[f32],
906 ) -> anyhow::Result<()> {
907 anyhow::ensure!(
913 row_indices.len() == values.len(),
914 "append_csc_slab: {} row indices vs {} values",
915 row_indices.len(),
916 values.len()
917 );
918 anyhow::ensure!(
919 local_colptr.first().copied() == Some(0) || local_colptr.is_empty(),
920 "append_csc_slab: local_colptr must start at 0"
921 );
922 anyhow::ensure!(
923 local_colptr.windows(2).all(|w| w[0] <= w[1]),
924 "append_csc_slab: local_colptr must be monotone non-decreasing"
925 );
926 if let Some(&last) = local_colptr.last() {
927 anyhow::ensure!(
928 last <= values.len() as u64,
929 "append_csc_slab: colptr claims {last} entries, slab holds {}",
930 values.len()
931 );
932 }
933 if let Some(nrow) = self.num_rows() {
934 if let Some(&bad) = row_indices.iter().find(|&&r| r >= nrow as u64) {
935 anyhow::bail!("append_csc_slab: row index {bad} outside the {nrow}-row matrix");
936 }
937 }
938 for (c, &start) in local_colptr.iter().enumerate() {
941 let end = local_colptr
942 .get(c + 1)
943 .copied()
944 .unwrap_or(values.len() as u64) as usize;
945 anyhow::ensure!(
946 row_indices[start as usize..end]
947 .windows(2)
948 .all(|w| w[0] < w[1]),
949 "append_csc_slab: rows within column {} of this band must be \
950 strictly ascending — repeated rows usually mean duplicate \
951 (row, col) coordinates in the source (an MTX with repeated \
952 entries, or a union remap folding rows together)",
953 col_offset as usize + c
954 );
955 }
956
957 let shifted: Vec<u64> = local_colptr.iter().map(|&p| p + nnz_offset).collect();
958 self.cs_write_u64(CsKey::CscIndptr, col_offset, &shifted)?;
959 self.cs_write_u64(CsKey::CscIndices, nnz_offset, row_indices)?;
960 self.cs_write_f32(CsKey::CscData, nnz_offset, values)?;
961 self.note_streamed_nnz(values.len() as u64);
962 Ok(())
963 }
964
965 fn finalize_streaming_csc(&mut self) -> anyhow::Result<()> {
968 let ncol = self
969 .num_columns()
970 .ok_or_else(|| anyhow::anyhow!("ncol not set before finalize_streaming_csc"))?;
971 let nnz = self
972 .num_non_zeros()
973 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csc"))?;
974 self.cs_write_u64(CsKey::CscIndptr, ncol as u64, &[nnz as u64])?;
975 self.read_column_indptr()?;
976
977 let indptr = self.column_indptr();
986 anyhow::ensure!(
987 indptr.len() == ncol + 1,
988 "finalize_streaming_csc: indptr has {} entries, expected {}",
989 indptr.len(),
990 ncol + 1
991 );
992 anyhow::ensure!(
993 indptr.first().copied() == Some(0),
994 "finalize_streaming_csc: indptr[0] = {:?}, expected 0 — the first \
995 slab was never appended",
996 indptr.first()
997 );
998 if let Some(w) = indptr.windows(2).position(|w| w[0] > w[1]) {
999 anyhow::bail!(
1000 "finalize_streaming_csc: indptr decreases at column {w} — slabs \
1001 were appended with a gap or overlap in their nnz offsets"
1002 );
1003 }
1004 let appended = self.streamed_nnz();
1010 anyhow::ensure!(
1011 appended == nnz as u64,
1012 "finalize_streaming_csc: {appended} entries appended but {nnz} \
1013 declared — the difference reads back as fill values wearing real \
1014 entries' positions"
1015 );
1016 Ok(())
1017 }
1018
1019 fn begin_streaming_csr(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
1022 self.reset_streamed_nnz();
1023 let (nrow, _, nnz) = shape;
1024 self.record_mtx_shape(Some(shape))?;
1025 self.cs_create(CsKey::CsrData, nnz)?;
1026 self.cs_create(CsKey::CsrIndices, nnz)?;
1027 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1028 Ok(())
1029 }
1030
1031 fn append_csr_slab(
1041 &mut self,
1042 row_offset: u64,
1043 nnz_offset: u64,
1044 local_rowptr: &[u64],
1045 col_indices: &[u64],
1046 values: &[f32],
1047 ) -> anyhow::Result<()> {
1048 anyhow::ensure!(
1049 col_indices.len() == values.len(),
1050 "append_csr_slab: {} column indices vs {} values",
1051 col_indices.len(),
1052 values.len()
1053 );
1054 anyhow::ensure!(
1055 local_rowptr.first().copied() == Some(0) || local_rowptr.is_empty(),
1056 "append_csr_slab: local_rowptr must start at 0"
1057 );
1058 anyhow::ensure!(
1059 local_rowptr.windows(2).all(|w| w[0] <= w[1]),
1060 "append_csr_slab: local_rowptr must be monotone non-decreasing"
1061 );
1062 if let Some(&last) = local_rowptr.last() {
1063 anyhow::ensure!(
1064 last <= values.len() as u64,
1065 "append_csr_slab: rowptr claims {last} entries, slab holds {}",
1066 values.len()
1067 );
1068 }
1069 if let Some(ncol) = self.num_columns() {
1070 if let Some(&bad) = col_indices.iter().find(|&&c| c >= ncol as u64) {
1071 anyhow::bail!(
1072 "append_csr_slab: column index {bad} outside the {ncol}-column matrix"
1073 );
1074 }
1075 }
1076 for (r, &start) in local_rowptr.iter().enumerate() {
1077 let end = local_rowptr
1078 .get(r + 1)
1079 .copied()
1080 .unwrap_or(values.len() as u64) as usize;
1081 anyhow::ensure!(
1082 col_indices[start as usize..end]
1083 .windows(2)
1084 .all(|w| w[0] < w[1]),
1085 "append_csr_slab: columns within row {} of this band must be \
1086 strictly ascending — repeated columns usually mean duplicate \
1087 (row, col) coordinates in the source",
1088 row_offset as usize + r
1089 );
1090 }
1091
1092 let shifted: Vec<u64> = local_rowptr.iter().map(|&p| p + nnz_offset).collect();
1093 self.cs_write_u64(CsKey::CsrIndptr, row_offset, &shifted)?;
1094 self.cs_write_u64(CsKey::CsrIndices, nnz_offset, col_indices)?;
1095 self.cs_write_f32(CsKey::CsrData, nnz_offset, values)?;
1096 self.note_streamed_nnz(values.len() as u64);
1097 Ok(())
1098 }
1099
1100 fn finalize_streaming_csr(&mut self) -> anyhow::Result<()> {
1104 let nrow = self
1105 .num_rows()
1106 .ok_or_else(|| anyhow::anyhow!("nrow not set before finalize_streaming_csr"))?;
1107 let nnz = self
1108 .num_non_zeros()
1109 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csr"))?;
1110 self.cs_write_u64(CsKey::CsrIndptr, nrow as u64, &[nnz as u64])?;
1111 self.read_row_indptr()?;
1112
1113 let appended = self.streamed_nnz();
1114 anyhow::ensure!(
1115 appended == nnz as u64,
1116 "finalize_streaming_csr: {appended} entries appended but {nnz} \
1117 declared — the slabs did not cover the matrix"
1118 );
1119 Ok(())
1120 }
1121
1122 fn build_csr_from_csc_streaming(&mut self) -> anyhow::Result<()> {
1126 let nrow = self
1127 .num_rows()
1128 .ok_or_else(|| anyhow::anyhow!("nrow not set before build_csr_from_csc_streaming"))?;
1129 let ncol = self
1130 .num_columns()
1131 .ok_or_else(|| anyhow::anyhow!("ncol not set before build_csr_from_csc_streaming"))?;
1132 let nnz = self
1133 .num_non_zeros()
1134 .ok_or_else(|| anyhow::anyhow!("nnz not set before build_csr_from_csc_streaming"))?;
1135
1136 if nnz == 0 {
1137 self.cs_create(CsKey::CsrData, 0)?;
1138 self.cs_create(CsKey::CsrIndices, 0)?;
1139 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1140 let zeros = vec![0u64; nrow + 1];
1141 self.cs_write_u64(CsKey::CsrIndptr, 0, &zeros)?;
1142 self.read_row_indptr()?;
1143 return Ok(());
1144 }
1145
1146 const COL_BLOCK: usize = 1024;
1147 let n_col_blocks = ncol.div_ceil(COL_BLOCK);
1148 let bar1 = styled_progress_bar(n_col_blocks as u64, "transpose count");
1149 let mut row_counts = vec![0u64; nrow];
1150 let mut col_lo = 0usize;
1151 while col_lo < ncol {
1152 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1153 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1154 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1155 for (row_i, _, _) in &triplets {
1156 row_counts[*row_i as usize] += 1;
1157 }
1158 col_lo = col_hi;
1159 bar1.inc(1);
1160 }
1161 bar1.finish_and_clear();
1162
1163 let mut rowptr = vec![0u64; nrow + 1];
1164 let mut acc = 0u64;
1165 for i in 0..nrow {
1166 rowptr[i] = acc;
1167 acc += row_counts[i];
1168 }
1169 rowptr[nrow] = acc;
1170 debug_assert_eq!(acc, nnz as u64);
1171
1172 self.cs_create(CsKey::CsrData, nnz)?;
1173 self.cs_create(CsKey::CsrIndices, nnz)?;
1174 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1175 self.cs_write_u64(CsKey::CsrIndptr, 0, &rowptr)?;
1176
1177 const TRANSPOSE_BAND_BYTES: usize = 256 * 1024 * 1024;
1181 let avg_density = nnz.div_ceil(nrow.max(1));
1182 let band_rows = (TRANSPOSE_BAND_BYTES / (12 * avg_density.max(1)))
1183 .max(1)
1184 .min(nrow);
1185 let n_bands = nrow.div_ceil(band_rows);
1186
1187 let bar2 = styled_progress_bar(n_bands as u64, "transpose scatter");
1188 let mut band_lo = 0usize;
1189 while band_lo < nrow {
1190 let band_hi = (band_lo + band_rows).min(nrow);
1191 let band_nnz_start = rowptr[band_lo];
1192 let band_nnz_end = rowptr[band_hi];
1193 let band_nnz = (band_nnz_end - band_nnz_start) as usize;
1194
1195 if band_nnz == 0 {
1196 band_lo = band_hi;
1197 bar2.inc(1);
1198 continue;
1199 }
1200
1201 let mut out_indices = vec![0u64; band_nnz];
1202 let mut out_values = vec![0f32; band_nnz];
1203 let mut cursor = vec![0u64; band_hi - band_lo];
1204
1205 let mut col_lo = 0usize;
1206 while col_lo < ncol {
1207 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1208 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1209 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1210 for &(row_i, col_j_local, x) in &triplets {
1211 let row_i_us = row_i as usize;
1212 if row_i_us >= band_lo && row_i_us < band_hi {
1213 let band_idx = row_i_us - band_lo;
1214 let col_j_global = col_j_local + col_lo as u64;
1219 let offset_in_band =
1220 (rowptr[band_lo + band_idx] - band_nnz_start) + cursor[band_idx];
1221 out_indices[offset_in_band as usize] = col_j_global;
1222 out_values[offset_in_band as usize] = x;
1223 cursor[band_idx] += 1;
1224 }
1225 }
1226 col_lo = col_hi;
1227 }
1228
1229 self.cs_write_u64(CsKey::CsrIndices, band_nnz_start, &out_indices)?;
1230 self.cs_write_f32(CsKey::CsrData, band_nnz_start, &out_values)?;
1231
1232 band_lo = band_hi;
1233 bar2.inc(1);
1234 }
1235 bar2.finish_and_clear();
1236
1237 self.read_row_indptr()?;
1238 Ok(())
1239 }
1240
1241 fn read_row_indptr(&mut self) -> anyhow::Result<()>;
1243
1244 fn read_column_indptr(&mut self) -> anyhow::Result<()>;
1246
1247 fn preload_columns(&mut self) -> anyhow::Result<()>;
1249
1250 fn clean_preloaded_columns(&mut self);
1252
1253 fn preload_rows(&mut self) -> anyhow::Result<()>;
1255
1256 fn clean_preloaded_rows(&mut self);
1258
1259 fn get_backend_file_name(&self) -> &str;
1261
1262 fn backend_type(&self) -> SparseIoBackend;
1264}