data_beans/sparse_io_vector/mod.rs
1#![allow(dead_code)]
2
3use crate::sparse_io::*;
4
5use indicatif::{ParallelProgressIterator, ProgressIterator};
6use legume_numeric::matrix::knn_match::ColumnDict;
7use legume_numeric::matrix::knn_match::MakeVecPoint;
8use legume_numeric::matrix::traits::*;
9use legume_numeric::matrix::utils::*;
10use log::info;
11use rayon::prelude::*;
12use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
13use std::ops::Index;
14use std::sync::Arc;
15
16// `impl SparseIoVec` is split across sibling modules by concern. Inherent
17// impls compose across modules, so every method stays callable at the same
18// path; siblings see the struct's private fields as descendant modules and
19// inherit this module's imports via `use super::*`.
20mod batch;
21mod groups;
22mod matched;
23mod push;
24mod read;
25pub mod scales;
26
27/// Where a single global cell's nonzeros come from in one of the
28/// underlying [`SparseIo`] backends. Under
29/// [`ColumnAlignment::Disjoint`] each global column has exactly one
30/// `BackendLocation`. Under [`ColumnAlignment::Union`] a global column
31/// can carry one entry per backend that observed the cell.
32#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub struct BackendLocation {
34 /// Index into [`SparseIoVec::data_vec`].
35 pub backend: u32,
36 /// Local column index within that backend.
37 pub local_col: u32,
38}
39
40type SparseData = dyn SparseIo<IndexIter = Vec<usize>>;
41
42/// Optional canonicalizer applied to every backend row name during
43/// `push`. Two raw names that produce the same canonical form are
44/// treated as the same row. The displayed `row_names_by_global` stores
45/// the canonical form, so downstream readers see a single consistent
46/// label for each row across all backends. Default = `None` preserves
47/// exact-match behavior.
48pub type RowNameCanonicalizer = Arc<dyn Fn(&str) -> Box<str> + Send + Sync>;
49
50/// Optional canonicalizer applied to every backend *column* (barcode)
51/// name during `push` under [`ColumnAlignment::Union`]. Mirror of
52/// [`RowNameCanonicalizer`]. Two raw barcodes whose canonical forms
53/// match are treated as the same biological cell — backends share a
54/// single global column id, and `read_columns_*` merges their nonzeros
55/// into one output column. Default = `None` preserves exact-match
56/// behavior.
57pub type ColumnNameCanonicalizer = Arc<dyn Fn(&str) -> Box<str> + Send + Sync>;
58
59/// The lazily-built group- and batch-membership caches for a [`SparseIoVec`],
60/// grouped into one field so the three places that reset them (`new`,
61/// `mask_columns`, and clone-for-collapse) each collapse to a single
62/// assignment.
63///
64/// **Cloning drops every cache** — `Clone` returns [`Default`]. The caches are
65/// derived state, re-registered from scratch by `assign_groups` /
66/// `register_batch_membership` / the collapse, so a cloned `SparseIoVec` starts
67/// from the correct pre-collapse state. Isolating the non-`Clone` HNSW indices
68/// in `batch_knn_lookup` here is also what lets the rest of `SparseIoVec`
69/// `#[derive(Clone)]`. Do not rely on `.clone()` preserving these.
70#[derive(Default)]
71struct DerivedCaches {
72 // group caches (built by `assign_groups`)
73 // per column, `usize::MAX` for a column left out of every group
74 col_to_group: Option<Vec<usize>>,
75 group_to_cols: Option<Vec<Vec<usize>>>,
76 group_keys: Option<Vec<Box<str>>>,
77 // batch caches (built by `register_batches` / `register_batch_membership`)
78 batch_knn_lookup: Option<Vec<ColumnDict<usize>>>,
79 col_to_batch: Option<Vec<usize>>,
80 batch_to_cols: Option<Vec<Vec<usize>>>,
81 batch_idx_to_name: Option<Vec<Box<str>>>,
82 between_batch_proximity: Option<Vec<Vec<usize>>>,
83}
84
85impl Clone for DerivedCaches {
86 /// Drops all caches (returns [`Default`]); see the type docs. They rebuild
87 /// lazily, so a cloned `SparseIoVec` is the correct fresh pre-collapse view.
88 fn clone(&self) -> Self {
89 Self::default()
90 }
91}
92
93#[derive(Clone)]
94pub struct SparseIoVec {
95 data_vec: Vec<Arc<SparseData>>,
96 /// Per global column → one or more [`BackendLocation`] entries.
97 /// Under [`ColumnAlignment::Disjoint`] each inner `Vec` always
98 /// has length 1. Under [`ColumnAlignment::Union`] a global
99 /// column can be observed by multiple backends — one entry per
100 /// backend that contributes triplets for that cell.
101 col_to_data: Vec<Vec<BackendLocation>>,
102 data_to_cols: HashMap<usize, Vec<usize>>,
103 offset: usize,
104 // Row-name alignment across backends. Each backend may have a
105 // different subset/ordering of rows; we keep only the intersection
106 // and remap local indices through `data_local_to_global_row` and
107 // `global_to_compact_row` at read time.
108 row_canonicalizer: Option<RowNameCanonicalizer>,
109 row_name_position: HashMap<Box<str>, usize>,
110 row_names_by_global: Vec<Box<str>>,
111 data_local_to_global_row: Vec<Vec<usize>>,
112 /// Per-backend inverse of `data_local_to_global_row[didx]`: every
113 /// `(global, local)` row pair, sorted, so the local rows of one global
114 /// row are a contiguous run found by binary search -- more than one where
115 /// the canonicalizer folded rows together. Built once at `push` time,
116 /// never mutated thereafter, so row reads don't have to rebuild it per
117 /// call.
118 data_global_to_local_row: Vec<Vec<(usize, usize)>>,
119 /// `true` for datasets where `local_to_global` is non-injective —
120 /// i.e. the canonicalizer collapsed two or more local rows to the
121 /// same global row. Read paths use this to switch from a fast
122 /// pass-through emit to a `(row, col) → sum` merge so the resulting
123 /// COO doesn't carry duplicate entries (which break CSC builders).
124 data_has_intra_row_merges: Vec<bool>,
125 row_count_by_global: Vec<usize>,
126 global_to_compact_row: Vec<Option<usize>>,
127 /// Inverse of `global_to_compact_row` restricted to compact rows.
128 /// `compact_to_global_row[c]` is the raw-global index for compact `c`.
129 /// Kept in sync with `global_to_compact_row` whenever it changes.
130 compact_to_global_row: Vec<usize>,
131 column_names_with_data_tag: Vec<Box<str>>,
132 /// `Union` mode only: canonical barcode → global column id. Populated
133 /// at push time so subsequent pushes can match cells across backends.
134 /// Unused (empty) under `Disjoint`.
135 col_name_position: HashMap<Box<str>, u32>,
136 /// Lazily-built group/batch membership caches; dropped on clone. See
137 /// [`DerivedCaches`].
138 derived: DerivedCaches,
139 cached_num_rows: usize,
140 cached_num_columns: usize,
141 /// How row names align across pushed backends. Default
142 /// [`RowAlignment::Union`] keeps every row from any backend, with
143 /// cells from a backend that doesn't contain row `g` implicitly
144 /// observing zero at that position — used for multi-modal data
145 /// where features (e.g. peaks vs genes) are disjoint.
146 /// [`RowAlignment::Intersect`] reverts to the historical
147 /// "common rows only" semantics.
148 row_alignment: RowAlignment,
149 /// How column (cell) names align across pushed backends. Default
150 /// [`ColumnAlignment::Disjoint`] preserves the historical
151 /// concatenate-cells semantics, with `@<basename>` suffixing when
152 /// multiple files are loaded. [`ColumnAlignment::Union`] matches
153 /// cells by canonical barcode across backends and lets one global
154 /// column carry triplets from multiple backends — the foundation
155 /// for patchy multi-modal (multiome) integration.
156 column_alignment: ColumnAlignment,
157 /// Optional canonicalizer for barcode matching under `Union` mode.
158 /// See [`ColumnNameCanonicalizer`].
159 column_canonicalizer: Option<ColumnNameCanonicalizer>,
160 /// Optional per-backend feature-name (row) suffix, indexed by push
161 /// order (`didx`). When set, every row of backend `b` is renamed
162 /// `{canon(row)}/{suffix[b]}` *after* canonicalization — so the
163 /// canonical gene/locus rule still applies to the bare name, and the
164 /// suffix only namespaces it onto a modality-specific row. Two
165 /// backends sharing a raw name (e.g. spliced vs unspliced `TSPAN6`)
166 /// thus stay separate, while the same name + same suffix (same
167 /// modality across donors) still merges. `None` = no suffixing.
168 per_backend_row_suffix: Option<Vec<Box<str>>>,
169 /// Per global column, the factor every read divides it by
170 /// ([`Self::set_column_scales`]). `None` reads the counts as stored.
171 column_scales: Option<Vec<f32>>,
172 /// Per global column, how many observations it stands for
173 /// ([`Self::register_column_multiplicity`]); `None` means one apiece.
174 /// Declared by the caller and derivable from nothing else, so unlike
175 /// [`DerivedCaches`] it survives cloning and
176 /// [`Self::clear_column_membership`], and [`Self::mask_columns`] keeps the
177 /// surviving columns' weights.
178 col_multiplicity: Option<Vec<f32>>,
179}
180
181/// Strategy for aligning row names across multiple pushed backends.
182#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
183pub enum RowAlignment {
184 /// Keep every row from any backend; non-observing backends emit zero
185 /// at that position. The default — strictly more permissive than
186 /// intersection: it reduces to the same result when all files share
187 /// their row set, and enables multi-modal load (e.g. peaks ∪ genes)
188 /// when they don't.
189 #[default]
190 Union,
191 /// Keep only rows present in every backend. Opt-in for callers that
192 /// want strict single-modality semantics.
193 Intersect,
194}
195
196/// Strategy for aligning column (cell) names across multiple pushed
197/// backends.
198#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
199pub enum ColumnAlignment {
200 /// Today's behavior. Each push appends `ncol_data` brand-new global
201 /// cells; raw barcodes get `@<basename>` suffixed for disambiguation
202 /// when more than one backend is pushed. Two backends naming the
203 /// same biological cell stay disjoint.
204 #[default]
205 Disjoint,
206 /// Match each pushed column's canonical barcode against existing
207 /// global cells. New barcodes extend the cell pool; matched barcodes
208 /// share one global column across backends (the column carries
209 /// triplets from every backend that observes the cell). No
210 /// `@<basename>` suffix. Foundation for patchy multi-modal
211 /// (multiome) integration.
212 Union,
213}
214
215pub struct TripletsMatched {
216 pub shape: (usize, usize),
217 pub triplets: Vec<(u64, u64, f32)>,
218 pub source_columns: Vec<usize>,
219 pub matched_columns: Vec<usize>,
220 pub distances: Vec<f32>,
221}
222
223/// The backend pushed `idx`-th, as stored: reads through it bypass the row
224/// alignment and the column scales ([`SparseIoVec::set_column_scales`]).
225impl Index<usize> for SparseIoVec {
226 type Output = Arc<SparseData>;
227 fn index(&self, idx: usize) -> &Self::Output {
228 &self.data_vec[idx]
229 }
230}
231
232impl Default for SparseIoVec {
233 /// an empty sparse io vector for horizontal data integration
234 fn default() -> Self {
235 Self::new()
236 }
237}
238
239impl SparseIoVec {
240 /// an empty sparse io vector for horizontal data integration
241 pub fn new() -> Self {
242 Self {
243 data_vec: vec![],
244 col_to_data: vec![],
245 data_to_cols: HashMap::default(),
246 offset: 0,
247 row_canonicalizer: None,
248 row_name_position: HashMap::default(),
249 row_names_by_global: vec![],
250 data_local_to_global_row: vec![],
251 data_global_to_local_row: vec![],
252 data_has_intra_row_merges: vec![],
253 row_count_by_global: vec![],
254 global_to_compact_row: vec![],
255 compact_to_global_row: vec![],
256 column_names_with_data_tag: vec![],
257 col_name_position: HashMap::default(),
258 derived: DerivedCaches::default(),
259 cached_num_rows: 0,
260 cached_num_columns: 0,
261 row_alignment: RowAlignment::default(),
262 column_alignment: ColumnAlignment::default(),
263 column_canonicalizer: None,
264 per_backend_row_suffix: None,
265 column_scales: None,
266 col_multiplicity: None,
267 }
268 }
269
270 /// Switch row-name alignment between intersection (default) and union.
271 /// Must be called BEFORE any push — the row-mapping recompute uses
272 /// the current value of `row_alignment`. Errors if any backend has
273 /// already been added.
274 pub fn with_row_alignment(mut self, mode: RowAlignment) -> anyhow::Result<Self> {
275 anyhow::ensure!(
276 self.data_vec.is_empty(),
277 "row alignment must be set before any push"
278 );
279 self.row_alignment = mode;
280 Ok(self)
281 }
282
283 /// Install a row-name canonicalizer for fuzzy cross-backend row
284 /// alignment. Must be called BEFORE the first [`push`](Self::push) —
285 /// returns an error if any backend has already been added (the
286 /// existing `row_names_by_global` would otherwise mix raw and
287 /// canonicalized forms).
288 ///
289 /// Typical use: pass a `GeneIndexResolver`-style canonicalizer so
290 /// `ENSG00000000003_TSPAN6` (file A) and `TSPAN6` (file B) collapse
291 /// to a single row, instead of being silently dropped from the
292 /// shared-row intersection.
293 pub fn with_row_canonicalizer(
294 mut self,
295 canon: impl Fn(&str) -> Box<str> + Send + Sync + 'static,
296 ) -> anyhow::Result<Self> {
297 anyhow::ensure!(
298 self.data_vec.is_empty(),
299 "row canonicalizer must be set before any push"
300 );
301 self.row_canonicalizer = Some(Arc::new(canon));
302 Ok(self)
303 }
304
305 /// Install a per-backend feature-name suffix (one entry per backend,
306 /// in push order). Each backend `b`'s rows are renamed
307 /// `{canon(row)}/{suffix[b]}`, so files sharing raw feature names stay
308 /// on separate rows unless they also share the suffix (same modality).
309 /// Must be called BEFORE the first [`push`](Self::push). The vec length
310 /// must match the number of backends that will be pushed; `push` errors
311 /// if `didx` is out of range.
312 pub fn with_per_backend_row_suffix(mut self, suffix: Vec<Box<str>>) -> anyhow::Result<Self> {
313 anyhow::ensure!(
314 self.data_vec.is_empty(),
315 "per-backend row suffix must be set before any push"
316 );
317 self.per_backend_row_suffix = Some(suffix);
318 Ok(self)
319 }
320
321 /// Switch column (cell) alignment between disjoint concatenation
322 /// (default) and barcode-keyed union. Must be called BEFORE any
323 /// push — the push branches off the current value. Errors if any
324 /// backend has already been added.
325 pub fn with_column_alignment(mut self, mode: ColumnAlignment) -> anyhow::Result<Self> {
326 anyhow::ensure!(
327 self.data_vec.is_empty(),
328 "column alignment must be set before any push"
329 );
330 self.column_alignment = mode;
331 Ok(self)
332 }
333
334 /// Install a column-name canonicalizer for fuzzy cross-backend
335 /// barcode matching under [`ColumnAlignment::Union`]. Must be
336 /// called BEFORE the first [`push`](Self::push) — returns an error
337 /// if any backend has already been added. Has no effect under
338 /// [`ColumnAlignment::Disjoint`] (barcodes are never compared
339 /// across backends in that mode).
340 pub fn with_column_canonicalizer(
341 mut self,
342 canon: impl Fn(&str) -> Box<str> + Send + Sync + 'static,
343 ) -> anyhow::Result<Self> {
344 anyhow::ensure!(
345 self.data_vec.is_empty(),
346 "column canonicalizer must be set before any push"
347 );
348 self.column_canonicalizer = Some(Arc::new(canon));
349 Ok(self)
350 }
351
352 /// number of data sets
353 pub fn len(&self) -> usize {
354 self.data_vec.len()
355 }
356
357 /// check if the vector is empty
358 pub fn is_empty(&self) -> bool {
359 self.data_vec.is_empty()
360 }
361
362 pub fn num_rows(&self) -> usize {
363 self.cached_num_rows
364 }
365
366 /// Number of canonical rows observed by **at least** `k` of the
367 /// pushed backends. Useful for detecting multi-modal-shaped inputs
368 /// (`num_rows_in_at_least(n_backends)` is the strict intersection
369 /// size; comparing it against per-backend row counts reveals how
370 /// disjoint the feature axes are).
371 pub fn num_rows_in_at_least(&self, k: usize) -> usize {
372 self.row_count_by_global.iter().filter(|&&c| c >= k).count()
373 }
374
375 /// Current column-alignment mode. Mirrors [`Self::row_alignment`].
376 pub fn column_alignment(&self) -> ColumnAlignment {
377 self.column_alignment
378 }
379
380 /// Per-backend row coverage on the exposed (compact) row axis:
381 /// `coverage[d][r]` is true when backend `d` measures row `r`.
382 ///
383 /// Under [`RowAlignment::Union`] a backend with a smaller panel simply
384 /// has no entry at the rows it lacks — reads return zero there, which is
385 /// indistinguishable from "measured, and absent". This is the map that
386 /// lets a consumer tell the two apart: unmeasured is *no evidence*, not
387 /// evidence of zero.
388 ///
389 /// `None` when every backend covers every row (single backend, identical
390 /// panels, or intersect alignment) — the common case, so callers can skip
391 /// observability handling entirely on `None`.
392 #[must_use]
393 pub fn row_coverage_by_backend(&self) -> Option<Vec<Vec<bool>>> {
394 let n = self.cached_num_rows;
395 let mut coverage = vec![vec![false; n]; self.data_vec.len()];
396 let mut any_gap = false;
397 for (d, locals) in self.data_local_to_global_row.iter().enumerate() {
398 for &g in locals {
399 if let Some(r) = self.global_to_compact_row[g] {
400 coverage[d][r] = true;
401 }
402 }
403 any_gap |= coverage[d].iter().any(|&c| !c);
404 }
405 any_gap.then_some(coverage)
406 }
407
408 /// The single backend a global column comes from, or `None` when the
409 /// column merges several backends (column-union alignment). Observability
410 /// accounting needs one source per column; a merged column has a *set* of
411 /// panels, which callers must handle (or refuse) explicitly.
412 #[must_use]
413 pub fn column_source(&self, col: usize) -> Option<usize> {
414 match self.col_to_data.get(col)?.as_slice() {
415 [one] => Some(one.backend as usize),
416 _ => None,
417 }
418 }
419
420 /// Every backend location of global column `col`: one entry under
421 /// `Disjoint`, one per observing backend under `Union`. Empty when out
422 /// of range.
423 #[must_use]
424 pub fn column_locations(&self, col: usize) -> &[BackendLocation] {
425 self.col_to_data.get(col).map_or(&[], Vec::as_slice)
426 }
427
428 pub fn num_non_zeros(&self) -> anyhow::Result<usize> {
429 let mut ret = 0;
430 for dat in self.data_vec.iter() {
431 let nnz = dat
432 .num_non_zeros()
433 .ok_or(anyhow::anyhow!("can't figure out the number of non-zeros"))?;
434 ret += nnz;
435 }
436 Ok(ret)
437 }
438
439 /// total number of columns across all data files
440 pub fn num_columns(&self) -> usize {
441 self.cached_num_columns
442 }
443
444 /// Structural clone for an independent projection / collapse pass.
445 ///
446 /// Named entry point for the collapse copy; it is exactly `self.clone()`,
447 /// kept as a method so call sites read as intent. Cloning copies the data
448 /// handles (matrices are `Arc`-shared, so only the index `Vec`s/`HashMap`s
449 /// are duplicated — cheap), the row/column alignment state and the
450 /// per-column scales and multiplicities, while
451 /// **dropping** every derived batch / group / HNSW cache — see
452 /// [`DerivedCaches`], whose drop-on-clone behavior is what isolates the
453 /// non-`Clone` `batch_knn_lookup` HNSW indices and lets `SparseIoVec`
454 /// `#[derive(Clone)]` at all. Those caches are re-registered from scratch
455 /// by `register_batch_membership` / the collapse, so a fresh clone is the
456 /// correct pre-collapse state.
457 ///
458 /// Prefer this name over a bare `.clone()` at collapse sites; note that any
459 /// `.clone()` is likewise lossy on the derived caches.
460 ///
461 /// Intended use: `let mut spliced = vec.clone_for_collapse();
462 /// spliced.mask_rows(&spliced_keep)?;` — gives a spliced-only view that
463 /// drives RP + collapse + refinement without disturbing the full backend
464 /// (still needed at all rows for per-modality aggregation).
465 pub fn clone_for_collapse(&self) -> Self {
466 // `Clone` drops the derived group/batch caches (see `DerivedCaches`) —
467 // exactly the fresh pre-collapse state this method documents.
468 self.clone()
469 }
470
471 /// Exclude rows (genes) from the working set. `keep[compact_row]`
472 /// is `true` for rows to keep, `false` for rows to exclude.
473 /// The compact row indices are renumbered after filtering.
474 /// This affects all downstream operations (projection, collapse,
475 /// training, inference).
476 pub fn mask_rows(&mut self, keep: &[bool]) -> anyhow::Result<()> {
477 let n_compact = self.cached_num_rows;
478 if keep.len() != n_compact {
479 return Err(anyhow::anyhow!(
480 "mask_rows: keep.len()={} != num_rows={}",
481 keep.len(),
482 n_compact
483 ));
484 }
485 // Build old_compact → new_compact mapping
486 let mut old_to_new: Vec<Option<usize>> = vec![None; n_compact];
487 let mut next = 0usize;
488 for (old, &k) in keep.iter().enumerate() {
489 if k {
490 old_to_new[old] = Some(next);
491 next += 1;
492 }
493 }
494 // Update global_to_compact_row
495 for entry in self.global_to_compact_row.iter_mut() {
496 *entry = entry.and_then(|old_compact| old_to_new[old_compact]);
497 }
498 self.cached_num_rows = next;
499 self.compact_to_global_row.clear();
500 self.compact_to_global_row.resize(next, 0);
501 for (g, &c_opt) in self.global_to_compact_row.iter().enumerate() {
502 if let Some(c) = c_opt {
503 self.compact_to_global_row[c] = g;
504 }
505 }
506 log::info!(
507 "mask_rows: {} → {} rows ({} excluded)",
508 n_compact,
509 next,
510 n_compact - next
511 );
512 Ok(())
513 }
514
515 /// Exclude columns (cells) from the working set. `keep[global_col]`
516 /// is `true` for cells to keep, `false` to exclude. Global column
517 /// indices are renumbered after filtering. This affects all
518 /// downstream operations (projection, collapse, training, inference).
519 ///
520 /// Cell-axis mirror of [`Self::mask_rows`]. MUST be called before
521 /// batch/group registration: `register_batch_membership` and group
522 /// assignment index by global column id and would be corrupted by a
523 /// renumber, so we `debug_assert!` they are unset and defensively
524 /// clear them. Dropped cells leave their backend-local columns in
525 /// place but unmapped (`usize::MAX` in `data_to_cols`), which the
526 /// row-wise read path (`rows_triplets`) skips.
527 pub fn mask_columns(&mut self, keep: &[bool]) -> anyhow::Result<()> {
528 let n = self.cached_num_columns;
529 if keep.len() != n {
530 return Err(anyhow::anyhow!(
531 "mask_columns: keep.len()={} != num_columns={}",
532 keep.len(),
533 n
534 ));
535 }
536 debug_assert!(
537 self.derived.col_to_batch.is_none() && self.derived.col_to_group.is_none(),
538 "mask_columns must be called before batch/group registration"
539 );
540
541 // old_global → Some(new_global) | None (dropped)
542 let mut old_to_new: Vec<Option<usize>> = vec![None; n];
543 let mut next = 0usize;
544 for (old, &k) in keep.iter().enumerate() {
545 if k {
546 old_to_new[old] = Some(next);
547 next += 1;
548 }
549 }
550
551 // Compact col_to_data + column_names_with_data_tag in one pass.
552 let mut new_col_to_data: Vec<Vec<BackendLocation>> = Vec::with_capacity(next);
553 let mut new_names: Vec<Box<str>> = Vec::with_capacity(next);
554 for (old, &k) in keep.iter().enumerate() {
555 if k {
556 new_col_to_data.push(std::mem::take(&mut self.col_to_data[old]));
557 new_names.push(self.column_names_with_data_tag[old].clone());
558 }
559 }
560 self.col_to_data = new_col_to_data;
561 self.column_names_with_data_tag = new_names;
562 for per_column in [&mut self.column_scales, &mut self.col_multiplicity] {
563 if let Some(values) = per_column.as_mut() {
564 *values = values
565 .iter()
566 .zip(keep)
567 .filter(|(_, &k)| k)
568 .map(|(&v, _)| v)
569 .collect();
570 }
571 }
572
573 // Remap data_to_cols in place. It is a positional
574 // local-col → global-col map per backend (consumed by
575 // `rows_triplets` / `take_backend_columns`); dropped cells map to
576 // `usize::MAX`, which both consumers skip. Already-`usize::MAX` entries
577 // (from a prior `mask_columns`) stay dropped — guard them so a second
578 // mask (e.g. up-front cell QC followed by a refine-pass mask) is
579 // re-entrant rather than indexing `old_to_new` out of bounds.
580 for cols in self.data_to_cols.values_mut() {
581 for c in cols.iter_mut() {
582 *c = if *c == usize::MAX {
583 usize::MAX
584 } else {
585 old_to_new[*c].unwrap_or(usize::MAX)
586 };
587 }
588 }
589
590 // Rebuild the Union-mode canonical-barcode → global lookup
591 // (empty / no-op under Disjoint).
592 if !self.col_name_position.is_empty() {
593 let mut pos: HashMap<Box<str>, u32> =
594 HashMap::with_capacity_and_hasher(next, Default::default());
595 for (g, name) in self.column_names_with_data_tag.iter().enumerate() {
596 let canon: Box<str> = match self.column_canonicalizer.as_ref() {
597 Some(c) => c(name),
598 None => name.clone(),
599 };
600 let g_u32: u32 = g
601 .try_into()
602 .map_err(|_| anyhow::anyhow!("global col overflows u32"))?;
603 pos.insert(canon, g_u32);
604 }
605 self.col_name_position = pos;
606 }
607
608 self.offset = next;
609 self.cached_num_columns = next;
610
611 // Any cell-indexed membership is now stale (asserted unset above).
612 self.derived = DerivedCaches::default();
613
614 log::info!(
615 "mask_columns: {} → {} cells ({} excluded)",
616 n,
617 next,
618 n - next
619 );
620 Ok(())
621 }
622
623 /// Drop the cell-indexed group/batch membership caches so the columns can
624 /// be re-masked and the membership re-registered from scratch. Needed
625 /// before [`Self::mask_columns`] when a collapse already registered groups
626 /// on this backend (it asserts the caches are unset); the next
627 /// `assign_groups` / `register_batch_membership` / collapse rebuilds them.
628 /// This is the fourth reset site alluded to in [`DerivedCaches`]. Column
629 /// multiplicities and scales are not membership and stay registered.
630 pub fn clear_column_membership(&mut self) {
631 self.derived = DerivedCaches::default();
632 }
633
634 pub fn row_names(&self) -> anyhow::Result<Vec<Box<str>>> {
635 let ntot = self.num_rows();
636 let mut ret = vec![Box::from(""); ntot];
637 for (raw_global, name) in self.row_names_by_global.iter().enumerate() {
638 if let Some(compact) = self
639 .global_to_compact_row
640 .get(raw_global)
641 .copied()
642 .flatten()
643 {
644 ret[compact] = name.clone();
645 }
646 }
647 Ok(ret)
648 }
649}