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}