Skip to main content

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}