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};
7pub use ndarray::prelude::*;
8
9pub const MAX_ROW_NAME_IDX: usize = 3;
10pub const MAX_COLUMN_NAME_IDX: usize = 10;
11pub const COLUMN_SEP: &str = "@";
12pub const ROW_SEP: &str = "_";
13
14use super::helpers::*;
15
16use crate::sparse_data_visitors::styled_progress_bar;
17use clap::ValueEnum;
18use indicatif::ParallelProgressIterator;
19use legume_numeric::matrix::mtx_io::*;
20use legume_numeric::matrix::traits::*;
21use log::info;
22use rayon::prelude::*;
23use rustc_hash::FxHashMap as HashMap;
24use std::ops::Range;
25use std::sync::{Arc, Mutex};
26
27#[cfg(test)]
28mod tests;
29
30#[derive(ValueEnum, Clone, Debug, PartialEq)]
31#[clap(rename_all = "lowercase")]
32pub enum SparseIoBackend {
33 Zarr,
34 HDF5,
35}
36
37#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41pub enum CsKey {
42 CscData,
43 CscIndices,
44 CscIndptr,
45 CsrData,
46 CsrIndices,
47 CsrIndptr,
48}
49
50const SLAB_NNZ: usize = 1 << 20;
54
55fn slab_end(
62 triplets: &[(u64, u64, f32)],
63 start: usize,
64 slab_nnz: usize,
65 n_major: usize,
66 major: impl Fn(&(u64, u64, f32)) -> u64,
67) -> (usize, u64) {
68 debug_assert!(slab_nnz > 0);
69 let nnz = triplets.len();
70 let mut end = (start + slab_nnz).min(nnz);
71 while end < nnz && major(&triplets[end]) == major(&triplets[end - 1]) {
72 end += 1;
73 }
74 let band_end = if end == nnz {
75 n_major as u64
76 } else {
77 major(&triplets[end])
78 };
79 (end, band_end)
80}
81
82pub trait SparseIo: Sync + Send {
83 type IndexIter: IntoIterator<Item = usize> + FromIterator<usize>;
84
85 fn read_columns_ndarray(&self, columns: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
93 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
94 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
95 }
96
97 #[cfg(feature = "tensor")]
98 fn read_columns_tensor(&self, columns: Self::IndexIter) -> anyhow::Result<Tensor> {
102 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
103 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
104 }
105
106 fn read_columns_dmatrix(&self, columns: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
110 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
111 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
112 }
113
114 fn read_columns_csr(&self, columns: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
118 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
119 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
120 }
121
122 fn read_columns_csc(&self, columns: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
126 let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
127 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
128 }
129
130 fn csc_column_arrays(&self) -> Option<(&[u64], &[u64], &[f32])> {
136 None
137 }
138
139 fn read_rows_ndarray(&self, rows: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
143 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
144 Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
145 }
146
147 #[cfg(feature = "tensor")]
148 fn read_rows_tensor(&self, rows: Self::IndexIter) -> anyhow::Result<Tensor> {
152 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
153 Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
154 }
155
156 fn read_rows_dmatrix(&self, rows: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
160 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
161 DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
162 }
163
164 fn read_rows_csr(&self, rows: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
168 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
169 CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
170 }
171
172 fn read_rows_csc(&self, rows: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
176 let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
177 CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
178 }
179
180 fn import_mtx_file(&mut self, mtx_file: &str, index_by_row: bool) -> anyhow::Result<()> {
190 let (mut mtx_triplets, mtx_shape) = read_mtx_triplets(mtx_file)?;
191 info!("read mtx file: {}", mtx_file);
192 if mtx_triplets.is_empty() {
193 return Err(anyhow::anyhow!("No data in mtx file"));
194 }
195 self.record_mtx_shape(Some(mtx_shape))?;
196 info!("recording the column index");
197 self.record_triplets_by_col(&mut mtx_triplets)?;
198 if index_by_row {
199 info!("recording the row index");
200 self.record_triplets_by_row(&mut mtx_triplets)?;
201 }
202 Ok(())
203 }
204
205 fn import_dmatrix_by_row(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
212 let (nrow, ncol) = matrix.shape();
213 let mut mtx_triplets = dmatrix_to_triplets(matrix);
214 let mtx_shape = (nrow, ncol, mtx_triplets.len());
215 self.record_mtx_shape(Some(mtx_shape))?;
216 self.record_triplets_by_row(&mut mtx_triplets)
217 }
218
219 fn import_dmatrix_by_col(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
222 let (nrow, ncol) = matrix.shape();
223 let mut mtx_triplets = dmatrix_to_triplets(matrix);
224 let mtx_shape = (nrow, ncol, mtx_triplets.len());
225 self.record_mtx_shape(Some(mtx_shape))?;
226 self.record_triplets_by_col(&mut mtx_triplets)
227 }
228
229 fn import_ndarray_by_row(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
236 let nrow = array.shape()[0];
237 let ncol = array.shape()[1];
238
239 let mut mtx_triplets = ndarray_to_triplets(array);
241
242 let nnz = mtx_triplets.len();
243 let mtx_shape = (nrow, ncol, nnz);
244 self.record_mtx_shape(Some(mtx_shape))?;
245
246 self.record_triplets_by_row(&mut mtx_triplets)
249 }
250
251 fn import_ndarray_by_col(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
254 let nrow = array.shape()[0];
255 let ncol = array.shape()[1];
256
257 let mut mtx_triplets = ndarray_to_triplets(array);
259
260 let nnz = mtx_triplets.len();
261 let mtx_shape = (nrow, ncol, nnz);
262 self.record_mtx_shape(Some(mtx_shape))?;
263
264 self.record_triplets_by_col(&mut mtx_triplets)
267 }
268
269 #[allow(clippy::type_complexity)]
277 fn read_triplets_by_rows(
278 &self,
279 rows: Self::IndexIter,
280 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
281
282 #[allow(clippy::type_complexity)]
286 fn read_triplets_by_columns(
287 &self,
288 columns: Self::IndexIter,
289 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
290
291 #[allow(clippy::type_complexity)]
295 fn read_triplets_by_single_column(
296 &self,
297 col: usize,
298 ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
299
300 fn to_mtx_file(&self, mtx_file: &str) -> anyhow::Result<()>;
303
304 fn num_rows(&self) -> Option<usize>;
306
307 fn num_columns(&self) -> Option<usize>;
309
310 fn num_non_zeros(&self) -> Option<usize>;
312
313 fn reopen_backend(&mut self) -> anyhow::Result<()>;
318
319 #[doc(hidden)]
323 fn note_streamed_nnz(&mut self, n: u64);
331
332 #[doc(hidden)]
334 fn streamed_nnz(&self) -> u64;
335
336 #[doc(hidden)]
338 fn reset_streamed_nnz(&mut self);
339
340 fn column_indptr(&self) -> &[u64];
345
346 fn column_nnz(&self, col: usize) -> Option<u64> {
352 let indptr = self.column_indptr();
353 let hi = *indptr.get(col + 1)?;
354 let lo = *indptr.get(col)?;
355 hi.checked_sub(lo)
356 }
357
358 fn register_row_names_file(&mut self, row_name_file: &str);
361
362 fn register_column_names_file(&mut self, column_name_file: &str);
365
366 fn register_row_names_vec(&mut self, rows: &[Box<str>]);
369
370 fn register_column_names_vec(&mut self, columns: &[Box<str>]);
373
374 fn register_names_file(
380 &mut self,
381 key: &str,
382 name_file: &str,
383 name_columns: Range<usize>,
384 name_sep: &str,
385 ) -> anyhow::Result<()>;
386
387 fn register_names_vec(&mut self, key: &str, names: &[Box<str>]) -> anyhow::Result<()>;
391
392 fn row_names(&self) -> anyhow::Result<Vec<Box<str>>>;
393
394 fn column_names(&self) -> anyhow::Result<Vec<Box<str>>>;
395
396 fn retrieve_registered_names(&self, key: &str) -> anyhow::Result<Vec<Box<str>>>;
399
400 fn subset_columns_rows(
408 &mut self,
409 columns: Option<&Vec<usize>>,
410 rows: Option<&Vec<usize>>,
411 ) -> anyhow::Result<()> {
412 let ncol_data = self
413 .num_columns()
414 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
415 let nrow_data = self
416 .num_rows()
417 .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
418
419 let distinct = |sel: &[usize], what: &str| -> anyhow::Result<()> {
428 anyhow::ensure!(!sel.is_empty(), "subset: empty {what} selection");
429 let mut seen = sel.to_vec();
430 seen.sort_unstable();
431 seen.dedup();
432 anyhow::ensure!(
433 seen.len() == sel.len(),
434 "subset: the {what} selection repeats an index ({} of {} are distinct)",
435 seen.len(),
436 sel.len()
437 );
438 Ok(())
439 };
440 if let Some(cols) = columns {
441 distinct(cols, "column")?;
442 }
443 if let Some(rs) = rows {
444 distinct(rs, "row")?;
445 }
446
447 let (old2new_cols, new_col_names) =
452 take_subset_indices_names_if_needed(columns, Some(ncol_data), self.column_names()?);
453 let (old2new_rows, new_row_names) =
454 take_subset_indices_names_if_needed(rows, Some(nrow_data), self.row_names()?);
455 let (new_ncol, new_nrow) = (new_col_names.len(), new_row_names.len());
456 anyhow::ensure!(new_ncol > 0, "subset: no column survived the selection");
457 anyhow::ensure!(new_nrow > 0, "subset: no row survived the selection");
458
459 let mut cols_new_order: Vec<(u64, u64)> =
461 old2new_cols.iter().map(|(&o, &n)| (n, o)).collect();
462 cols_new_order.sort_unstable();
463
464 let mut row_map: Vec<Option<u64>> = vec![None; nrow_data];
469 for (&old, &new) in &old2new_rows {
470 row_map[old as usize] = Some(new);
471 }
472 let monotone_rows = row_map.iter().flatten().is_sorted_by(|a, b| a < b);
473
474 let full_rows = rows.is_none();
482 let per_col_nnz: Vec<u64> = if full_rows {
483 cols_new_order
484 .iter()
485 .map(|&(_, old)| {
486 self.column_nnz(old as usize)
487 .ok_or_else(|| anyhow::anyhow!("subset: no indptr for column {old}"))
488 })
489 .collect::<anyhow::Result<_>>()?
490 } else {
491 let mut counts = vec![0u64; cols_new_order.len()];
496 let coarse = legume_numeric::matrix::utils::generate_minibatch_intervals(
497 cols_new_order.len(),
498 0,
499 Some(8192),
500 );
501 for (lb, ub) in coarse {
502 let old_cols: Vec<usize> = cols_new_order[lb..ub]
503 .iter()
504 .map(|&(_, o)| o as usize)
505 .collect();
506 let (_, _, triplets) =
507 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
508 for (i, c_local, _) in triplets {
509 if row_map[i as usize].is_some() {
510 counts[lb + c_local as usize] += 1;
511 }
512 }
513 }
514 counts
515 };
516 let new_nnz: u64 = per_col_nnz.iter().sum();
517
518 let final_path = self.get_backend_file_name().to_string();
536 anyhow::ensure!(
537 !final_path.ends_with(".zip"),
538 "subset: {final_path} is a zip archive; convert it to a directory \
539 backend first (data-beans convert)"
540 );
541 let temp_path = format!("{final_path}.subset_tmp");
542 if std::path::Path::new(&temp_path).exists() {
543 crate::sparse_io::remove_backend_path(&temp_path)?;
544 }
545
546 {
547 let backend_kind = self.backend_type();
548 let mut out = crate::sparse_io::create_sparse_streaming_empty(
549 Some(&temp_path),
550 Some(&backend_kind),
551 )?;
552 out.begin_streaming_csc((new_nrow, new_ncol, new_nnz as usize))?;
553
554 let blocks = legume_numeric::matrix::utils::byte_budget_intervals(
557 &per_col_nnz,
558 crate::sparse_io::SLAB_BUDGET_BYTES,
559 crate::sparse_io::TRIPLET_BYTES,
560 );
561
562 let mut nnz_offset = 0u64;
563 for (lb, ub) in blocks {
564 let old_cols: Vec<usize> = cols_new_order[lb..ub]
568 .iter()
569 .map(|&(_, o)| o as usize)
570 .collect();
571 let (_, _, triplets) =
572 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
573
574 let n_block = ub - lb;
575 let mut per_col: Vec<Vec<(u64, f32)>> = vec![Vec::new(); n_block];
576 for (i, c_local, x) in triplets {
577 if let Some(new_row) = row_map[i as usize] {
578 per_col[c_local as usize].push((new_row, x));
579 }
580 }
581 let mut local_colptr = Vec::with_capacity(n_block);
582 let mut row_indices = Vec::new();
583 let mut values = Vec::new();
584 for entries in &mut per_col {
585 if !monotone_rows {
586 entries.sort_unstable_by_key(|&(r, _)| r);
589 }
590 local_colptr.push(row_indices.len() as u64);
591 for &(r, x) in entries.iter() {
592 row_indices.push(r);
593 values.push(x);
594 }
595 }
596 out.append_csc_slab(lb as u64, nnz_offset, &local_colptr, &row_indices, &values)?;
597 nnz_offset += values.len() as u64;
598 }
599
600 out.finalize_streaming_csc()?;
601 out.build_csr_from_csc_streaming()?;
602 out.register_row_names_vec(&new_row_names);
603 out.register_column_names_vec(&new_col_names);
604 }
605
606 self.remove_backend_file()?;
611 std::fs::rename(&temp_path, &final_path)?;
612 self.reopen_backend()?;
613 self.clean_preloaded_columns();
614 self.clean_preloaded_rows();
615 info!("registered new data to {}", self.get_backend_file_name());
616 Ok(())
617 }
618
619 fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()> {
622 let new_col_names = self.column_names()?.clone();
623 let name2new = build_name2index_map(row_names_order);
624
625 let block_size = 100;
626
627 let old2new: HashMap<u64, u64> = self
628 .row_names()?
629 .into_par_iter()
630 .enumerate()
631 .filter_map(|(idx_old, name)| {
632 name2new
633 .get(&name)
634 .map(|&idx_new| (idx_old as u64, idx_new as u64))
635 })
636 .collect();
637
638 if let Some(ncol) = self.num_columns() {
639 let arc_triplets = Arc::new(Mutex::new(vec![]));
644
645 let nblock = ncol.div_ceil(block_size);
646
647 info!("remapping triplets ...");
648
649 (0..nblock)
650 .into_par_iter()
651 .progress_with(styled_progress_bar(nblock as u64, "blocks"))
652 .map(|b| {
653 let lb = (b * block_size) as u64;
654 let ub = ((b + 1) * block_size).min(ncol) as u64;
655 (lb, ub)
656 })
657 .for_each(|(lb, ub)| {
658 let (_, _, _triplets_b) = self
659 .read_triplets_by_columns(((lb as usize)..(ub as usize)).collect())
660 .unwrap();
661
662 let _triplets_b = _triplets_b.into_iter().filter_map(|(i, j_loc, x)| {
663 let j_glob = j_loc + lb;
664 old2new.get(&i).map(|&i_new| (i_new, j_glob, x))
665 });
666
667 {
668 let mut triplets = arc_triplets.lock().unwrap();
669 triplets.extend(_triplets_b);
670 }
671 });
672
673 self.remove_backend_file()?;
677
678 self.initialize_backend()?;
682
683 {
685 let mut row_col_val_triplets =
686 arc_triplets.lock().expect("failed to lock triplets");
687
688 let nnz = row_col_val_triplets.len();
689 debug_assert!(row_col_val_triplets.len() <= nnz); let new_nrow = row_names_order.len();
691 let mtx_shape = (new_nrow, ncol, nnz);
692
693 info!("sorting triplets ...");
694
695 self.record_mtx_shape(Some(mtx_shape))?;
696 self.record_triplets_by_col(&mut row_col_val_triplets)?;
697 self.record_triplets_by_row(&mut row_col_val_triplets)?;
698 }
699 self.read_column_indptr()?;
700 self.read_row_indptr()?;
701
702 self.register_row_names_vec(row_names_order);
703 self.register_column_names_vec(&new_col_names);
704 info!("registered new data to {}", self.get_backend_file_name());
705 }
706
707 self.clean_preloaded_columns();
708 self.clean_preloaded_rows();
709 Ok(())
710 }
711 fn remove_backend_file(&self) -> anyhow::Result<()>;
715
716 fn initialize_backend(&mut self) -> anyhow::Result<()>;
718
719 fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()>;
720
721 fn record_triplets_by_row(
724 &mut self,
725 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
726 ) -> anyhow::Result<()> {
727 let nrow = self.num_rows().expect("should have `nrow`");
728 let ncol = self.num_columns().expect("should have `ncol`");
729 let nnz = row_col_val_triplets.len();
730
731 if nnz == 0 {
732 let csr_rowptr = vec![0u64; nrow + 1];
733 return self.record_csr_dataset_backend(&[], &[], &csr_rowptr);
734 }
735
736 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (row, col));
740
741 self.begin_streaming_csr((nrow, ncol, nnz))?;
742
743 let mut local_rowptr: Vec<u64> = Vec::new();
744 let mut cols: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
745 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
746
747 let mut start = 0_usize;
748 let mut row_offset = 0_u64;
749 while (row_offset as usize) < nrow {
750 let (end, band_end_row) =
751 slab_end(row_col_val_triplets, start, SLAB_NNZ, nrow, |t| t.0);
752
753 local_rowptr.clear();
754 cols.clear();
755 vals.clear();
756 let mut i = start;
757 for row in row_offset..band_end_row {
758 local_rowptr.push((i - start) as u64);
759 while i < end && row_col_val_triplets[i].0 == row {
760 cols.push(row_col_val_triplets[i].1);
761 vals.push(row_col_val_triplets[i].2);
762 i += 1;
763 }
764 }
765 debug_assert_eq!(i, end, "every entry of the band belongs to one of its rows");
766
767 self.append_csr_slab(row_offset, start as u64, &local_rowptr, &cols, &vals)?;
768 start = end;
769 row_offset = band_end_row;
770 }
771
772 self.finalize_streaming_csr()
773 }
774
775 fn record_triplets_by_col(
782 &mut self,
783 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
784 ) -> anyhow::Result<()> {
785 let nrow = self.num_rows().expect("should have `nrow`");
786 let ncol = self.num_columns().expect("should have `ncol`");
787 let nnz = row_col_val_triplets.len();
788
789 if nnz == 0 {
790 let csc_colptr = vec![0u64; ncol + 1];
791 return self.record_csc_dataset_backend(&[], &[], &csc_colptr);
792 }
793
794 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (col, row));
796
797 self.begin_streaming_csc((nrow, ncol, nnz))?;
798
799 let mut local_colptr: Vec<u64> = Vec::new();
800 let mut rows: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
801 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
802
803 let mut start = 0_usize;
804 let mut col_offset = 0_u64;
805 while (col_offset as usize) < ncol {
806 let (end, band_end_col) =
807 slab_end(row_col_val_triplets, start, SLAB_NNZ, ncol, |t| t.1);
808
809 local_colptr.clear();
810 rows.clear();
811 vals.clear();
812 let mut i = start;
813 for col in col_offset..band_end_col {
814 local_colptr.push((i - start) as u64);
815 while i < end && row_col_val_triplets[i].1 == col {
816 rows.push(row_col_val_triplets[i].0);
817 vals.push(row_col_val_triplets[i].2);
818 i += 1;
819 }
820 }
821 debug_assert_eq!(
822 i, end,
823 "every entry of the band belongs to one of its columns"
824 );
825
826 self.append_csc_slab(col_offset, start as u64, &local_colptr, &rows, &vals)?;
827 start = end;
828 col_offset = band_end_col;
829 }
830
831 self.finalize_streaming_csc()
832 }
833
834 fn record_csr_dataset_backend(
843 &mut self,
844 csr_cols: &[u64],
845 csr_vals: &[f32],
846 csr_rowptr: &[u64],
847 ) -> anyhow::Result<()>;
848
849 fn record_csc_dataset_backend(
859 &mut self,
860 csc_rows: &[u64],
861 csc_vals: &[f32],
862 csc_colptr: &[u64],
863 ) -> anyhow::Result<()>;
864
865 fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()>;
868
869 fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()>;
872
873 fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()>;
876
877 fn begin_streaming_csc(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
882 self.reset_streamed_nnz();
886 let (_, ncol, nnz) = shape;
887 self.record_mtx_shape(Some(shape))?;
888 self.cs_create(CsKey::CscData, nnz)?;
889 self.cs_create(CsKey::CscIndices, nnz)?;
890 self.cs_create(CsKey::CscIndptr, ncol + 1)?;
891 Ok(())
892 }
893
894 fn append_csc_slab(
903 &mut self,
904 col_offset: u64,
905 nnz_offset: u64,
906 local_colptr: &[u64],
907 row_indices: &[u64],
908 values: &[f32],
909 ) -> anyhow::Result<()> {
910 anyhow::ensure!(
916 row_indices.len() == values.len(),
917 "append_csc_slab: {} row indices vs {} values",
918 row_indices.len(),
919 values.len()
920 );
921 anyhow::ensure!(
922 local_colptr.first().copied() == Some(0) || local_colptr.is_empty(),
923 "append_csc_slab: local_colptr must start at 0"
924 );
925 anyhow::ensure!(
926 local_colptr.windows(2).all(|w| w[0] <= w[1]),
927 "append_csc_slab: local_colptr must be monotone non-decreasing"
928 );
929 if let Some(&last) = local_colptr.last() {
930 anyhow::ensure!(
931 last <= values.len() as u64,
932 "append_csc_slab: colptr claims {last} entries, slab holds {}",
933 values.len()
934 );
935 }
936 if let Some(nrow) = self.num_rows() {
937 if let Some(&bad) = row_indices.iter().find(|&&r| r >= nrow as u64) {
938 anyhow::bail!("append_csc_slab: row index {bad} outside the {nrow}-row matrix");
939 }
940 }
941 for (c, &start) in local_colptr.iter().enumerate() {
944 let end = local_colptr
945 .get(c + 1)
946 .copied()
947 .unwrap_or(values.len() as u64) as usize;
948 anyhow::ensure!(
949 row_indices[start as usize..end]
950 .windows(2)
951 .all(|w| w[0] < w[1]),
952 "append_csc_slab: rows within column {} of this band must be \
953 strictly ascending — repeated rows usually mean duplicate \
954 (row, col) coordinates in the source (an MTX with repeated \
955 entries, or a union remap folding rows together)",
956 col_offset as usize + c
957 );
958 }
959
960 let shifted: Vec<u64> = local_colptr.iter().map(|&p| p + nnz_offset).collect();
961 self.cs_write_u64(CsKey::CscIndptr, col_offset, &shifted)?;
962 self.cs_write_u64(CsKey::CscIndices, nnz_offset, row_indices)?;
963 self.cs_write_f32(CsKey::CscData, nnz_offset, values)?;
964 self.note_streamed_nnz(values.len() as u64);
965 Ok(())
966 }
967
968 fn finalize_streaming_csc(&mut self) -> anyhow::Result<()> {
971 let ncol = self
972 .num_columns()
973 .ok_or_else(|| anyhow::anyhow!("ncol not set before finalize_streaming_csc"))?;
974 let nnz = self
975 .num_non_zeros()
976 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csc"))?;
977 self.cs_write_u64(CsKey::CscIndptr, ncol as u64, &[nnz as u64])?;
978 self.read_column_indptr()?;
979
980 let indptr = self.column_indptr();
989 anyhow::ensure!(
990 indptr.len() == ncol + 1,
991 "finalize_streaming_csc: indptr has {} entries, expected {}",
992 indptr.len(),
993 ncol + 1
994 );
995 anyhow::ensure!(
996 indptr.first().copied() == Some(0),
997 "finalize_streaming_csc: indptr[0] = {:?}, expected 0 — the first \
998 slab was never appended",
999 indptr.first()
1000 );
1001 if let Some(w) = indptr.windows(2).position(|w| w[0] > w[1]) {
1002 anyhow::bail!(
1003 "finalize_streaming_csc: indptr decreases at column {w} — slabs \
1004 were appended with a gap or overlap in their nnz offsets"
1005 );
1006 }
1007 let appended = self.streamed_nnz();
1013 anyhow::ensure!(
1014 appended == nnz as u64,
1015 "finalize_streaming_csc: {appended} entries appended but {nnz} \
1016 declared — the difference reads back as fill values wearing real \
1017 entries' positions"
1018 );
1019 Ok(())
1020 }
1021
1022 fn begin_streaming_csr(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
1025 self.reset_streamed_nnz();
1026 let (nrow, _, nnz) = shape;
1027 self.record_mtx_shape(Some(shape))?;
1028 self.cs_create(CsKey::CsrData, nnz)?;
1029 self.cs_create(CsKey::CsrIndices, nnz)?;
1030 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1031 Ok(())
1032 }
1033
1034 fn append_csr_slab(
1044 &mut self,
1045 row_offset: u64,
1046 nnz_offset: u64,
1047 local_rowptr: &[u64],
1048 col_indices: &[u64],
1049 values: &[f32],
1050 ) -> anyhow::Result<()> {
1051 anyhow::ensure!(
1052 col_indices.len() == values.len(),
1053 "append_csr_slab: {} column indices vs {} values",
1054 col_indices.len(),
1055 values.len()
1056 );
1057 anyhow::ensure!(
1058 local_rowptr.first().copied() == Some(0) || local_rowptr.is_empty(),
1059 "append_csr_slab: local_rowptr must start at 0"
1060 );
1061 anyhow::ensure!(
1062 local_rowptr.windows(2).all(|w| w[0] <= w[1]),
1063 "append_csr_slab: local_rowptr must be monotone non-decreasing"
1064 );
1065 if let Some(&last) = local_rowptr.last() {
1066 anyhow::ensure!(
1067 last <= values.len() as u64,
1068 "append_csr_slab: rowptr claims {last} entries, slab holds {}",
1069 values.len()
1070 );
1071 }
1072 if let Some(ncol) = self.num_columns() {
1073 if let Some(&bad) = col_indices.iter().find(|&&c| c >= ncol as u64) {
1074 anyhow::bail!(
1075 "append_csr_slab: column index {bad} outside the {ncol}-column matrix"
1076 );
1077 }
1078 }
1079 for (r, &start) in local_rowptr.iter().enumerate() {
1080 let end = local_rowptr
1081 .get(r + 1)
1082 .copied()
1083 .unwrap_or(values.len() as u64) as usize;
1084 anyhow::ensure!(
1085 col_indices[start as usize..end]
1086 .windows(2)
1087 .all(|w| w[0] < w[1]),
1088 "append_csr_slab: columns within row {} of this band must be \
1089 strictly ascending — repeated columns usually mean duplicate \
1090 (row, col) coordinates in the source",
1091 row_offset as usize + r
1092 );
1093 }
1094
1095 let shifted: Vec<u64> = local_rowptr.iter().map(|&p| p + nnz_offset).collect();
1096 self.cs_write_u64(CsKey::CsrIndptr, row_offset, &shifted)?;
1097 self.cs_write_u64(CsKey::CsrIndices, nnz_offset, col_indices)?;
1098 self.cs_write_f32(CsKey::CsrData, nnz_offset, values)?;
1099 self.note_streamed_nnz(values.len() as u64);
1100 Ok(())
1101 }
1102
1103 fn finalize_streaming_csr(&mut self) -> anyhow::Result<()> {
1107 let nrow = self
1108 .num_rows()
1109 .ok_or_else(|| anyhow::anyhow!("nrow not set before finalize_streaming_csr"))?;
1110 let nnz = self
1111 .num_non_zeros()
1112 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csr"))?;
1113 self.cs_write_u64(CsKey::CsrIndptr, nrow as u64, &[nnz as u64])?;
1114 self.read_row_indptr()?;
1115
1116 let appended = self.streamed_nnz();
1117 anyhow::ensure!(
1118 appended == nnz as u64,
1119 "finalize_streaming_csr: {appended} entries appended but {nnz} \
1120 declared — the slabs did not cover the matrix"
1121 );
1122 Ok(())
1123 }
1124
1125 fn build_csr_from_csc_streaming(&mut self) -> anyhow::Result<()> {
1129 let nrow = self
1130 .num_rows()
1131 .ok_or_else(|| anyhow::anyhow!("nrow not set before build_csr_from_csc_streaming"))?;
1132 let ncol = self
1133 .num_columns()
1134 .ok_or_else(|| anyhow::anyhow!("ncol not set before build_csr_from_csc_streaming"))?;
1135 let nnz = self
1136 .num_non_zeros()
1137 .ok_or_else(|| anyhow::anyhow!("nnz not set before build_csr_from_csc_streaming"))?;
1138
1139 if nnz == 0 {
1140 self.cs_create(CsKey::CsrData, 0)?;
1141 self.cs_create(CsKey::CsrIndices, 0)?;
1142 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1143 let zeros = vec![0u64; nrow + 1];
1144 self.cs_write_u64(CsKey::CsrIndptr, 0, &zeros)?;
1145 self.read_row_indptr()?;
1146 return Ok(());
1147 }
1148
1149 const COL_BLOCK: usize = 1024;
1150 let n_col_blocks = ncol.div_ceil(COL_BLOCK);
1151 let bar1 = styled_progress_bar(n_col_blocks as u64, "transpose count");
1152 let mut row_counts = vec![0u64; nrow];
1153 let mut col_lo = 0usize;
1154 while col_lo < ncol {
1155 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1156 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1157 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1158 for (row_i, _, _) in &triplets {
1159 row_counts[*row_i as usize] += 1;
1160 }
1161 col_lo = col_hi;
1162 bar1.inc(1);
1163 }
1164 bar1.finish_and_clear();
1165
1166 let mut rowptr = vec![0u64; nrow + 1];
1167 let mut acc = 0u64;
1168 for i in 0..nrow {
1169 rowptr[i] = acc;
1170 acc += row_counts[i];
1171 }
1172 rowptr[nrow] = acc;
1173 debug_assert_eq!(acc, nnz as u64);
1174
1175 self.cs_create(CsKey::CsrData, nnz)?;
1176 self.cs_create(CsKey::CsrIndices, nnz)?;
1177 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1178 self.cs_write_u64(CsKey::CsrIndptr, 0, &rowptr)?;
1179
1180 const TRANSPOSE_BAND_BYTES: usize = 256 * 1024 * 1024;
1184 let avg_density = nnz.div_ceil(nrow.max(1));
1185 let band_rows = (TRANSPOSE_BAND_BYTES / (12 * avg_density.max(1)))
1186 .max(1)
1187 .min(nrow);
1188 let n_bands = nrow.div_ceil(band_rows);
1189
1190 let bar2 = styled_progress_bar(n_bands as u64, "transpose scatter");
1191 let mut band_lo = 0usize;
1192 while band_lo < nrow {
1193 let band_hi = (band_lo + band_rows).min(nrow);
1194 let band_nnz_start = rowptr[band_lo];
1195 let band_nnz_end = rowptr[band_hi];
1196 let band_nnz = (band_nnz_end - band_nnz_start) as usize;
1197
1198 if band_nnz == 0 {
1199 band_lo = band_hi;
1200 bar2.inc(1);
1201 continue;
1202 }
1203
1204 let mut out_indices = vec![0u64; band_nnz];
1205 let mut out_values = vec![0f32; band_nnz];
1206 let mut cursor = vec![0u64; band_hi - band_lo];
1207
1208 let mut col_lo = 0usize;
1209 while col_lo < ncol {
1210 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1211 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1212 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1213 for &(row_i, col_j_local, x) in &triplets {
1214 let row_i_us = row_i as usize;
1215 if row_i_us >= band_lo && row_i_us < band_hi {
1216 let band_idx = row_i_us - band_lo;
1217 let col_j_global = col_j_local + col_lo as u64;
1222 let offset_in_band =
1223 (rowptr[band_lo + band_idx] - band_nnz_start) + cursor[band_idx];
1224 out_indices[offset_in_band as usize] = col_j_global;
1225 out_values[offset_in_band as usize] = x;
1226 cursor[band_idx] += 1;
1227 }
1228 }
1229 col_lo = col_hi;
1230 }
1231
1232 self.cs_write_u64(CsKey::CsrIndices, band_nnz_start, &out_indices)?;
1233 self.cs_write_f32(CsKey::CsrData, band_nnz_start, &out_values)?;
1234
1235 band_lo = band_hi;
1236 bar2.inc(1);
1237 }
1238 bar2.finish_and_clear();
1239
1240 self.read_row_indptr()?;
1241 Ok(())
1242 }
1243
1244 fn read_row_indptr(&mut self) -> anyhow::Result<()>;
1246
1247 fn read_column_indptr(&mut self) -> anyhow::Result<()>;
1249
1250 fn preload_columns(&mut self) -> anyhow::Result<()>;
1252
1253 fn clean_preloaded_columns(&mut self);
1255
1256 fn preload_rows(&mut self) -> anyhow::Result<()>;
1258
1259 fn clean_preloaded_rows(&mut self);
1261
1262 fn get_backend_file_name(&self) -> &str;
1264
1265 fn backend_type(&self) -> SparseIoBackend;
1267}