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