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 name = std::path::Path::new(&final_path)
607 .file_name()
608 .map(|n| n.to_string_lossy().into_owned())
609 .unwrap_or_default();
610 let temp_path = crate::zarr_io::in_stage(&final_path, &format!("{name}.subset_tmp"));
611 crate::zarr_io::make_stage(&temp_path)?;
612 if std::path::Path::new(&temp_path).exists() {
613 crate::sparse_io::remove_backend_path(&temp_path)?;
614 }
615
616 {
617 let backend_kind = self.backend_type();
618 let mut out = crate::sparse_io::create_sparse_streaming_empty(
619 Some(&temp_path),
620 Some(&backend_kind),
621 )?;
622 out.begin_streaming_csc((new_nrow, new_ncol, new_nnz as usize))?;
623
624 let blocks = legume_numeric::matrix::utils::byte_budget_intervals(
627 &per_col_nnz,
628 crate::sparse_io::SLAB_BUDGET_BYTES,
629 crate::sparse_io::TRIPLET_BYTES,
630 );
631
632 let mut nnz_offset = 0u64;
633 for (lb, ub) in blocks {
634 let old_cols: Vec<usize> = cols_new_order[lb..ub]
638 .iter()
639 .map(|&(_, o)| o as usize)
640 .collect();
641 let (_, _, triplets) =
642 self.read_triplets_by_columns(old_cols.into_iter().collect())?;
643
644 let n_block = ub - lb;
645 let mut per_col: Vec<Vec<(u64, f32)>> = vec![Vec::new(); n_block];
646 for (i, c_local, x) in triplets {
647 if let Some(new_row) = row_map[i as usize] {
648 per_col[c_local as usize].push((new_row, x));
649 }
650 }
651 let mut local_colptr = Vec::with_capacity(n_block);
652 let mut row_indices = Vec::new();
653 let mut values = Vec::new();
654 for entries in &mut per_col {
655 if !monotone_rows {
656 entries.sort_unstable_by_key(|&(r, _)| r);
659 }
660 local_colptr.push(row_indices.len() as u64);
661 for &(r, x) in entries.iter() {
662 row_indices.push(r);
663 values.push(x);
664 }
665 }
666 out.append_csc_slab(lb as u64, nnz_offset, &local_colptr, &row_indices, &values)?;
667 nnz_offset += values.len() as u64;
668 }
669
670 out.finalize_streaming_csc()?;
671 out.build_csr_from_csc_streaming()?;
672 out.register_row_names_vec(&new_row_names);
673 out.register_column_names_vec(&new_col_names);
674 out.set_metadata(&self.metadata())?;
675 }
676
677 self.remove_backend_file()?;
682 std::fs::rename(&temp_path, &final_path)?;
683 crate::zarr_io::leave_stage(&temp_path);
684 self.reopen_backend()?;
685 self.clean_preloaded_columns();
686 self.clean_preloaded_rows();
687 info!("registered new data to {}", self.get_backend_file_name());
688 Ok(())
689 }
690
691 fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()> {
694 let meta = self.metadata();
696 let new_col_names = self.column_names()?.clone();
697 let name2new = build_name2index_map(row_names_order);
698
699 let block_size = 100;
700
701 let old2new: HashMap<u64, u64> = self
702 .row_names()?
703 .into_par_iter()
704 .enumerate()
705 .filter_map(|(idx_old, name)| {
706 name2new
707 .get(&name)
708 .map(|&idx_new| (idx_old as u64, idx_new as u64))
709 })
710 .collect();
711
712 if let Some(ncol) = self.num_columns() {
713 let arc_triplets = Arc::new(Mutex::new(vec![]));
718
719 let nblock = ncol.div_ceil(block_size);
720
721 info!("remapping triplets ...");
722
723 (0..nblock)
724 .into_par_iter()
725 .progress_with(styled_progress_bar(nblock as u64, "blocks"))
726 .map(|b| {
727 let lb = (b * block_size) as u64;
728 let ub = ((b + 1) * block_size).min(ncol) as u64;
729 (lb, ub)
730 })
731 .for_each(|(lb, ub)| {
732 let (_, _, _triplets_b) = self
733 .read_triplets_by_columns(((lb as usize)..(ub as usize)).collect())
734 .unwrap();
735
736 let _triplets_b = _triplets_b.into_iter().filter_map(|(i, j_loc, x)| {
737 let j_glob = j_loc + lb;
738 old2new.get(&i).map(|&i_new| (i_new, j_glob, x))
739 });
740
741 {
742 let mut triplets = arc_triplets.lock().unwrap();
743 triplets.extend(_triplets_b);
744 }
745 });
746
747 self.remove_backend_file()?;
751
752 self.initialize_backend()?;
756
757 {
759 let mut row_col_val_triplets =
760 arc_triplets.lock().expect("failed to lock triplets");
761
762 let nnz = row_col_val_triplets.len();
763 debug_assert!(row_col_val_triplets.len() <= nnz); let new_nrow = row_names_order.len();
765 let mtx_shape = (new_nrow, ncol, nnz);
766
767 info!("sorting triplets ...");
768
769 self.record_mtx_shape(Some(mtx_shape))?;
770 self.record_triplets_by_col(&mut row_col_val_triplets)?;
771 self.record_triplets_by_row(&mut row_col_val_triplets)?;
772 }
773 self.read_column_indptr()?;
774 self.read_row_indptr()?;
775
776 self.register_row_names_vec(row_names_order);
777 self.register_column_names_vec(&new_col_names);
778 self.set_metadata(&meta)?;
779 info!("registered new data to {}", self.get_backend_file_name());
780 }
781
782 self.clean_preloaded_columns();
783 self.clean_preloaded_rows();
784 Ok(())
785 }
786 fn remove_backend_file(&self) -> anyhow::Result<()>;
790
791 fn initialize_backend(&mut self) -> anyhow::Result<()>;
793
794 fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()>;
795
796 fn record_triplets_by_row(
799 &mut self,
800 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
801 ) -> anyhow::Result<()> {
802 let nrow = self.num_rows().expect("should have `nrow`");
803 let ncol = self.num_columns().expect("should have `ncol`");
804 let nnz = row_col_val_triplets.len();
805
806 if nnz == 0 {
807 let csr_rowptr = vec![0u64; nrow + 1];
808 return self.record_csr_dataset_backend(&[], &[], &csr_rowptr);
809 }
810
811 check_triplet_bounds("record_triplets_by_row", row_col_val_triplets, nrow, ncol)?;
812
813 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (row, col));
817
818 self.begin_streaming_csr((nrow, ncol, nnz))?;
819
820 let mut local_rowptr: Vec<u64> = Vec::new();
821 let mut cols: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
822 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
823
824 let mut start = 0_usize;
825 let mut row_offset = 0_u64;
826 while (row_offset as usize) < nrow {
827 let (end, band_end_row) =
828 slab_end(row_col_val_triplets, start, SLAB_NNZ, nrow, |t| t.0);
829
830 local_rowptr.clear();
831 cols.clear();
832 vals.clear();
833 let mut i = start;
834 for row in row_offset..band_end_row {
835 local_rowptr.push((i - start) as u64);
836 while i < end && row_col_val_triplets[i].0 == row {
837 cols.push(row_col_val_triplets[i].1);
838 vals.push(row_col_val_triplets[i].2);
839 i += 1;
840 }
841 }
842 if i != end {
843 anyhow::bail!(
844 "record_triplets_by_row: band walk left triplet {:?} at {i}",
845 row_col_val_triplets[i]
846 );
847 }
848
849 self.append_csr_slab(row_offset, start as u64, &local_rowptr, &cols, &vals)?;
850 start = end;
851 row_offset = band_end_row;
852 }
853
854 self.finalize_streaming_csr()
855 }
856
857 fn record_triplets_by_col(
864 &mut self,
865 row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
866 ) -> anyhow::Result<()> {
867 let nrow = self.num_rows().expect("should have `nrow`");
868 let ncol = self.num_columns().expect("should have `ncol`");
869 let nnz = row_col_val_triplets.len();
870
871 if nnz == 0 {
872 let csc_colptr = vec![0u64; ncol + 1];
873 return self.record_csc_dataset_backend(&[], &[], &csc_colptr);
874 }
875
876 check_triplet_bounds("record_triplets_by_col", row_col_val_triplets, nrow, ncol)?;
877
878 row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (col, row));
880
881 self.begin_streaming_csc((nrow, ncol, nnz))?;
882
883 let mut local_colptr: Vec<u64> = Vec::new();
884 let mut rows: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
885 let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
886
887 let mut start = 0_usize;
888 let mut col_offset = 0_u64;
889 while (col_offset as usize) < ncol {
890 let (end, band_end_col) =
891 slab_end(row_col_val_triplets, start, SLAB_NNZ, ncol, |t| t.1);
892
893 local_colptr.clear();
894 rows.clear();
895 vals.clear();
896 let mut i = start;
897 for col in col_offset..band_end_col {
898 local_colptr.push((i - start) as u64);
899 while i < end && row_col_val_triplets[i].1 == col {
900 rows.push(row_col_val_triplets[i].0);
901 vals.push(row_col_val_triplets[i].2);
902 i += 1;
903 }
904 }
905 if i != end {
906 anyhow::bail!(
907 "record_triplets_by_col: band walk left triplet {:?} at {i}",
908 row_col_val_triplets[i]
909 );
910 }
911
912 self.append_csc_slab(col_offset, start as u64, &local_colptr, &rows, &vals)?;
913 start = end;
914 col_offset = band_end_col;
915 }
916
917 self.finalize_streaming_csc()
918 }
919
920 fn record_csr_dataset_backend(
929 &mut self,
930 csr_cols: &[u64],
931 csr_vals: &[f32],
932 csr_rowptr: &[u64],
933 ) -> anyhow::Result<()>;
934
935 fn record_csc_dataset_backend(
945 &mut self,
946 csc_rows: &[u64],
947 csc_vals: &[f32],
948 csc_colptr: &[u64],
949 ) -> anyhow::Result<()>;
950
951 fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()>;
954
955 fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()>;
958
959 fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()>;
962
963 fn begin_streaming_csc(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
968 self.reset_streamed_nnz();
972 let (_, ncol, nnz) = shape;
973 self.record_mtx_shape(Some(shape))?;
974 self.cs_create(CsKey::CscData, nnz)?;
975 self.cs_create(CsKey::CscIndices, nnz)?;
976 self.cs_create(CsKey::CscIndptr, ncol + 1)?;
977 Ok(())
978 }
979
980 fn append_csc_slab(
989 &mut self,
990 col_offset: u64,
991 nnz_offset: u64,
992 local_colptr: &[u64],
993 row_indices: &[u64],
994 values: &[f32],
995 ) -> anyhow::Result<()> {
996 anyhow::ensure!(
1002 row_indices.len() == values.len(),
1003 "append_csc_slab: {} row indices vs {} values",
1004 row_indices.len(),
1005 values.len()
1006 );
1007 anyhow::ensure!(
1008 local_colptr.first().copied() == Some(0) || local_colptr.is_empty(),
1009 "append_csc_slab: local_colptr must start at 0"
1010 );
1011 anyhow::ensure!(
1012 local_colptr.windows(2).all(|w| w[0] <= w[1]),
1013 "append_csc_slab: local_colptr must be monotone non-decreasing"
1014 );
1015 if let Some(&last) = local_colptr.last() {
1016 anyhow::ensure!(
1017 last <= values.len() as u64,
1018 "append_csc_slab: colptr claims {last} entries, slab holds {}",
1019 values.len()
1020 );
1021 }
1022 if let Some(nrow) = self.num_rows() {
1023 if let Some(&bad) = row_indices.iter().find(|&&r| r >= nrow as u64) {
1024 anyhow::bail!("append_csc_slab: row index {bad} outside the {nrow}-row matrix");
1025 }
1026 }
1027 for (c, &start) in local_colptr.iter().enumerate() {
1030 let end = local_colptr
1031 .get(c + 1)
1032 .copied()
1033 .unwrap_or(values.len() as u64) as usize;
1034 anyhow::ensure!(
1035 row_indices[start as usize..end]
1036 .windows(2)
1037 .all(|w| w[0] < w[1]),
1038 "append_csc_slab: rows within column {} of this band must be \
1039 strictly ascending — repeated rows usually mean duplicate \
1040 (row, col) coordinates in the source (an MTX with repeated \
1041 entries, or a union remap folding rows together)",
1042 col_offset as usize + c
1043 );
1044 }
1045
1046 let shifted: Vec<u64> = local_colptr.iter().map(|&p| p + nnz_offset).collect();
1047 self.cs_write_u64(CsKey::CscIndptr, col_offset, &shifted)?;
1048 self.cs_write_u64(CsKey::CscIndices, nnz_offset, row_indices)?;
1049 self.cs_write_f32(CsKey::CscData, nnz_offset, values)?;
1050 self.note_streamed_nnz(values.len() as u64);
1051 Ok(())
1052 }
1053
1054 fn finalize_streaming_csc(&mut self) -> anyhow::Result<()> {
1057 let ncol = self
1058 .num_columns()
1059 .ok_or_else(|| anyhow::anyhow!("ncol not set before finalize_streaming_csc"))?;
1060 let nnz = self
1061 .num_non_zeros()
1062 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csc"))?;
1063 self.cs_write_u64(CsKey::CscIndptr, ncol as u64, &[nnz as u64])?;
1064 self.read_column_indptr()?;
1065
1066 index_audit::check_indptr("finalize_streaming_csc", self.column_indptr(), ncol, nnz)
1075 .context(
1076 "the slabs were appended with a gap or overlap in their nnz offsets, \
1077 or the first slab was never appended",
1078 )?;
1079 let appended = self.streamed_nnz();
1085 anyhow::ensure!(
1086 appended == nnz as u64,
1087 "finalize_streaming_csc: {appended} entries appended but {nnz} \
1088 declared — the difference reads back as fill values wearing real \
1089 entries' positions"
1090 );
1091 Ok(())
1092 }
1093
1094 fn begin_streaming_csr(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
1097 self.reset_streamed_nnz();
1098 let (nrow, _, nnz) = shape;
1099 self.record_mtx_shape(Some(shape))?;
1100 self.cs_create(CsKey::CsrData, nnz)?;
1101 self.cs_create(CsKey::CsrIndices, nnz)?;
1102 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1103 Ok(())
1104 }
1105
1106 fn append_csr_slab(
1116 &mut self,
1117 row_offset: u64,
1118 nnz_offset: u64,
1119 local_rowptr: &[u64],
1120 col_indices: &[u64],
1121 values: &[f32],
1122 ) -> anyhow::Result<()> {
1123 anyhow::ensure!(
1124 col_indices.len() == values.len(),
1125 "append_csr_slab: {} column indices vs {} values",
1126 col_indices.len(),
1127 values.len()
1128 );
1129 anyhow::ensure!(
1130 local_rowptr.first().copied() == Some(0) || local_rowptr.is_empty(),
1131 "append_csr_slab: local_rowptr must start at 0"
1132 );
1133 anyhow::ensure!(
1134 local_rowptr.windows(2).all(|w| w[0] <= w[1]),
1135 "append_csr_slab: local_rowptr must be monotone non-decreasing"
1136 );
1137 if let Some(&last) = local_rowptr.last() {
1138 anyhow::ensure!(
1139 last <= values.len() as u64,
1140 "append_csr_slab: rowptr claims {last} entries, slab holds {}",
1141 values.len()
1142 );
1143 }
1144 if let Some(ncol) = self.num_columns() {
1145 if let Some(&bad) = col_indices.iter().find(|&&c| c >= ncol as u64) {
1146 anyhow::bail!(
1147 "append_csr_slab: column index {bad} outside the {ncol}-column matrix"
1148 );
1149 }
1150 }
1151 for (r, &start) in local_rowptr.iter().enumerate() {
1152 let end = local_rowptr
1153 .get(r + 1)
1154 .copied()
1155 .unwrap_or(values.len() as u64) as usize;
1156 anyhow::ensure!(
1157 col_indices[start as usize..end]
1158 .windows(2)
1159 .all(|w| w[0] < w[1]),
1160 "append_csr_slab: columns within row {} of this band must be \
1161 strictly ascending — repeated columns usually mean duplicate \
1162 (row, col) coordinates in the source",
1163 row_offset as usize + r
1164 );
1165 }
1166
1167 let shifted: Vec<u64> = local_rowptr.iter().map(|&p| p + nnz_offset).collect();
1168 self.cs_write_u64(CsKey::CsrIndptr, row_offset, &shifted)?;
1169 self.cs_write_u64(CsKey::CsrIndices, nnz_offset, col_indices)?;
1170 self.cs_write_f32(CsKey::CsrData, nnz_offset, values)?;
1171 self.note_streamed_nnz(values.len() as u64);
1172 Ok(())
1173 }
1174
1175 fn finalize_streaming_csr(&mut self) -> anyhow::Result<()> {
1179 let nrow = self
1180 .num_rows()
1181 .ok_or_else(|| anyhow::anyhow!("nrow not set before finalize_streaming_csr"))?;
1182 let nnz = self
1183 .num_non_zeros()
1184 .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csr"))?;
1185 self.cs_write_u64(CsKey::CsrIndptr, nrow as u64, &[nnz as u64])?;
1186 self.read_row_indptr()?;
1187
1188 index_audit::check_indptr("finalize_streaming_csr", self.row_indptr(), nrow, nnz).context(
1190 "the slabs were appended with a gap or overlap in their nnz offsets, \
1191 or the first slab was never appended",
1192 )?;
1193
1194 let appended = self.streamed_nnz();
1195 anyhow::ensure!(
1196 appended == nnz as u64,
1197 "finalize_streaming_csr: {appended} entries appended but {nnz} \
1198 declared — the slabs did not cover the matrix"
1199 );
1200 Ok(())
1201 }
1202
1203 fn build_csr_from_csc_streaming(&mut self) -> anyhow::Result<()> {
1207 let nrow = self
1208 .num_rows()
1209 .ok_or_else(|| anyhow::anyhow!("nrow not set before build_csr_from_csc_streaming"))?;
1210 let ncol = self
1211 .num_columns()
1212 .ok_or_else(|| anyhow::anyhow!("ncol not set before build_csr_from_csc_streaming"))?;
1213 let nnz = self
1214 .num_non_zeros()
1215 .ok_or_else(|| anyhow::anyhow!("nnz not set before build_csr_from_csc_streaming"))?;
1216
1217 if nnz == 0 {
1218 self.cs_create(CsKey::CsrData, 0)?;
1219 self.cs_create(CsKey::CsrIndices, 0)?;
1220 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1221 let zeros = vec![0u64; nrow + 1];
1222 self.cs_write_u64(CsKey::CsrIndptr, 0, &zeros)?;
1223 self.read_row_indptr()?;
1224 return Ok(());
1225 }
1226
1227 const COL_BLOCK: usize = 1024;
1228 let n_col_blocks = ncol.div_ceil(COL_BLOCK);
1229 let bar1 = styled_progress_bar(n_col_blocks as u64, "transpose count");
1230 let mut row_counts = vec![0u64; nrow];
1231 let mut col_lo = 0usize;
1232 while col_lo < ncol {
1233 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1234 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1235 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1236 for t in &triplets {
1237 let slot = row_counts.get_mut(t.0 as usize).ok_or_else(|| {
1238 anyhow::anyhow!(
1239 "build_csr_from_csc_streaming: row index {} (0x{:016x}) outside \
1240 the {nrow}-row matrix",
1241 t.0,
1242 t.0
1243 )
1244 })?;
1245 *slot += 1;
1246 }
1247 col_lo = col_hi;
1248 bar1.inc(1);
1249 }
1250 bar1.finish_and_clear();
1251
1252 let mut rowptr = vec![0u64; nrow + 1];
1253 let mut acc = 0u64;
1254 for i in 0..nrow {
1255 rowptr[i] = acc;
1256 acc += row_counts[i];
1257 }
1258 rowptr[nrow] = acc;
1259 anyhow::ensure!(
1260 acc == nnz as u64,
1261 "build_csr_from_csc_streaming: counted {acc} entries in the CSC arrays, \
1262 but {nnz} are declared"
1263 );
1264
1265 self.cs_create(CsKey::CsrData, nnz)?;
1266 self.cs_create(CsKey::CsrIndices, nnz)?;
1267 self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1268 self.cs_write_u64(CsKey::CsrIndptr, 0, &rowptr)?;
1269
1270 const TRANSPOSE_BAND_BYTES: usize = 256 * 1024 * 1024;
1274 let avg_density = nnz.div_ceil(nrow.max(1));
1275 let band_rows = (TRANSPOSE_BAND_BYTES / (12 * avg_density.max(1)))
1276 .max(1)
1277 .min(nrow);
1278 let n_bands = nrow.div_ceil(band_rows);
1279
1280 let bar2 = styled_progress_bar(n_bands as u64, "transpose scatter");
1281 let mut band_lo = 0usize;
1282 while band_lo < nrow {
1283 let band_hi = (band_lo + band_rows).min(nrow);
1284 let band_nnz_start = rowptr[band_lo];
1285 let band_nnz_end = rowptr[band_hi];
1286 let band_nnz = (band_nnz_end - band_nnz_start) as usize;
1287
1288 if band_nnz == 0 {
1289 band_lo = band_hi;
1290 bar2.inc(1);
1291 continue;
1292 }
1293
1294 let mut out_indices = vec![0u64; band_nnz];
1295 let mut out_values = vec![0f32; band_nnz];
1296 let mut cursor = vec![0u64; band_hi - band_lo];
1297
1298 let mut col_lo = 0usize;
1299 while col_lo < ncol {
1300 let col_hi = (col_lo + COL_BLOCK).min(ncol);
1301 let cols: Self::IndexIter = (col_lo..col_hi).collect();
1302 let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1303 for &(row_i, col_j_local, x) in &triplets {
1304 let row_i_us = row_i as usize;
1305 if row_i_us >= band_lo && row_i_us < band_hi {
1306 let band_idx = row_i_us - band_lo;
1307 let col_j_global = col_j_local + col_lo as u64;
1312 let offset_in_band =
1313 (rowptr[band_lo + band_idx] - band_nnz_start) + cursor[band_idx];
1314 out_indices[offset_in_band as usize] = col_j_global;
1315 out_values[offset_in_band as usize] = x;
1316 cursor[band_idx] += 1;
1317 }
1318 }
1319 col_lo = col_hi;
1320 }
1321
1322 self.cs_write_u64(CsKey::CsrIndices, band_nnz_start, &out_indices)?;
1323 self.cs_write_f32(CsKey::CsrData, band_nnz_start, &out_values)?;
1324
1325 band_lo = band_hi;
1326 bar2.inc(1);
1327 }
1328 bar2.finish_and_clear();
1329
1330 self.read_row_indptr()?;
1331 Ok(())
1332 }
1333
1334 fn read_row_indptr(&mut self) -> anyhow::Result<()>;
1336
1337 fn read_column_indptr(&mut self) -> anyhow::Result<()>;
1339
1340 fn preload_columns(&mut self) -> anyhow::Result<()>;
1342
1343 fn clean_preloaded_columns(&mut self);
1345
1346 fn preload_rows(&mut self) -> anyhow::Result<()>;
1348
1349 fn clean_preloaded_rows(&mut self);
1351
1352 fn get_backend_file_name(&self) -> &str;
1354
1355 fn backend_type(&self) -> SparseIoBackend;
1357}