diff --git a/UserDocMD/usage/select.md b/UserDocMD/usage/select.md index 75bbce08..981da1a2 100644 --- a/UserDocMD/usage/select.md +++ b/UserDocMD/usage/select.md @@ -1,9 +1,9 @@ # select -Project and/or aggregate the genome columns of an index into a new (or in-place) index. Where [`filter`](filter.md) selects rows (kmers), `select` operates on columns (genomes): grouping several genomes into one aggregated column, reordering columns, or dropping some. +Project and/or aggregate the genome columns of an index into a new index. Where [`filter`](filter.md) selects rows (kmers), `select` operates on columns (genomes): grouping several genomes into one aggregated column, reordering columns, or dropping some. ```bash -obikmer select SOURCE (--output OUTPUT | --in-place) [OPTIONS] +obikmer select SOURCE --output OUTPUT [OPTIONS] ``` ## Arguments @@ -16,8 +16,7 @@ obikmer select SOURCE (--output OUTPUT | --in-place) [OPTIONS] | Option | Default | Description | |---|---|---| -| `--output` | — | Output index directory (mutually exclusive with `--in-place`) | -| `--in-place` | off | Rewrite the source index in place (mutually exclusive with `--output`) | +| `-o, --output` | — | Output index directory (required) | | `-f, --force` | off | Overwrite an existing output directory | | `--group NAME:PRED` | none | Define a named group of genomes by predicate (repeatable; mutually exclusive with `--aggregate-by`) | | `--group-op NAME:OP` | none | Aggregation operator for a named group | @@ -25,6 +24,8 @@ obikmer select SOURCE (--output OUTPUT | --in-place) [OPTIONS] | `--aggregate-op OP` | none | Aggregation operator applied to every auto-generated group | | `--select COL,...` | all columns | Output columns, in order (group names or genome labels) | | `--presence-threshold` | `0` | Minimum count for a genome to be considered a carrier (logical operators only) | +| `--dense` | off | Pack the output's presence matrices in the dense format instead of the default sparse one | +| `--force-copy` | off | Copy each layer's unchanged kmer-identity files (mphf/unitigs/evidence/fingerprint) instead of hard-linking them | ## Aggregation operators @@ -32,4 +33,16 @@ obikmer select SOURCE (--output OUTPUT | --in-place) [OPTIONS] A `select` never changes the underlying kmer set — only the per-genome data (counts or presence) is rewritten, so an unaggregated pass-through column (a plain genome label in `--select`) is a cheap copy. -At least one of `--output`/`--in-place` is required, and at least one output column must be defined; every name listed in `--select` must resolve to either a defined group or an existing genome label. See [Genome predicates and taxonomy paths](predicates.md) for the predicate syntax used by `--group`. +At least one output column must be defined; every name listed in `--select` must resolve to either a defined group or an existing genome label. See [Genome predicates and taxonomy paths](predicates.md) for the predicate syntax used by `--group`. + +## Disk usage + +`select` always writes to a new output directory — there is no in-place mode. Each layer's kmer-identity files (MPHF, unitigs, evidence, fingerprint) never change under a column projection/aggregation, so they are hard-linked into the output rather than copied: no extra disk is used for them, even on a very large index. Linking falls back to a real copy automatically if it fails (e.g. `SOURCE`/`OUTPUT` on different filesystems). Use `--force-copy` to always copy instead — needed when the output must be able to survive independently of the source on disk (a hard link shares the same underlying data, so overwriting one path outside `select` itself would affect the other). + +To replace an index with a selected version of itself, select to a temporary directory and swap it in: + +```bash +obikmer select INDEX --output INDEX.tmp --group ... --group-op ... --select ... +rm -rf INDEX +mv INDEX.tmp INDEX +``` diff --git a/benchmark/filter_one_count.sh b/benchmark/filter_one_count.sh index 115ed3ce..51a7a293 100755 --- a/benchmark/filter_one_count.sh +++ b/benchmark/filter_one_count.sh @@ -1,7 +1,10 @@ #!/usr/bin/env bash # Usage: filter_one_count.sh SPECIES # Filters global_index_count to keep only kmers specific to SPECIES, -# then selects the SPECIES column in-place. +# then selects the SPECIES column (to a temp dir, swapped over the +# original — obikmer select has no --in-place mode, but its output hard- +# links each layer's unchanged kmer-identity files rather than copying +# them, so this costs no extra disk for those). # Outputs: # specific_index_count/SPECIES/index.done (written by obikmer select) # stats/specific_kmer_count/SPECIES.stats (one CSV data row, no header) @@ -37,8 +40,11 @@ trap 'rm -f "${LOG_FILTER}" "${LOG_SELECT}"' EXIT cat "${LOG_FILTER}" >&2 +SELECT_TMP="${OUTPUT}.select_tmp" +rm -rf "${SELECT_TMP}" + "${BINARY}" select \ - --in-place \ + --output "${SELECT_TMP}" \ --group "${SPECIES}:species=${SPECIES}" \ --group-op "${SPECIES}:any" \ --select "${SPECIES}" \ @@ -47,6 +53,9 @@ cat "${LOG_FILTER}" >&2 cat "${LOG_SELECT}" >&2 +rm -rf "${OUTPUT}" +mv "${SELECT_TMP}" "${OUTPUT}" + python3 - "${SPECIES}" "${LOG_FILTER}" "${LOG_SELECT}" <<'PYEOF' >"${STATS_FILE}" import sys, re diff --git a/benchmark/filter_one_presence.sh b/benchmark/filter_one_presence.sh index 12099ce9..5e61a659 100755 --- a/benchmark/filter_one_presence.sh +++ b/benchmark/filter_one_presence.sh @@ -1,7 +1,10 @@ #!/usr/bin/env bash # Usage: filter_one_presence.sh SPECIES # Filters global_index_presence to keep only kmers specific to SPECIES, -# then selects the SPECIES column in-place. +# then selects the SPECIES column (to a temp dir, swapped over the +# original — obikmer select has no --in-place mode, but its output hard- +# links each layer's unchanged kmer-identity files rather than copying +# them, so this costs no extra disk for those). # Outputs: # specific_index_presence/SPECIES/index.done (written by obikmer select) # stats/specific_kmer_presence/SPECIES.stats (one CSV data row, no header) @@ -37,8 +40,11 @@ trap 'rm -f "${LOG_FILTER}" "${LOG_SELECT}"' EXIT cat "${LOG_FILTER}" >&2 +SELECT_TMP="${OUTPUT}.select_tmp" +rm -rf "${SELECT_TMP}" + "${BINARY}" select \ - --in-place \ + --output "${SELECT_TMP}" \ --group "${SPECIES}:species=${SPECIES}" \ --group-op "${SPECIES}:any" \ --select "${SPECIES}" \ @@ -47,6 +53,9 @@ cat "${LOG_FILTER}" >&2 cat "${LOG_SELECT}" >&2 +rm -rf "${OUTPUT}" +mv "${SELECT_TMP}" "${OUTPUT}" + python3 - "${SPECIES}" "${LOG_FILTER}" "${LOG_SELECT}" <<'PYEOF' >"${STATS_FILE}" import sys, re diff --git a/src/obicompactvec/src/bitmatrix/group_ops.rs b/src/obicompactvec/src/bitmatrix/group_ops.rs index ac449cc2..75805855 100644 --- a/src/obicompactvec/src/bitmatrix/group_ops.rs +++ b/src/obicompactvec/src/bitmatrix/group_ops.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::io; use crate::colgroup::{chunked_presence_count, ColGroup, MatrixGroupOps}; @@ -6,6 +7,110 @@ use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder}; use super::persistent::PersistentBitMatrix; +// ── batch_presence_counts ─────────────────────────────────────────────────── +// +// One shared pass computing every group's presence count at once, instead +// of `select`'s old per-spec loop over `MatrixGroupOps::partial_group_*` +// (one independent `col_view`-driven pass per output column). Two reasons +// this exists, not just one: +// +// - `Sparse` has no on-disk column representation at all (`col_view` is a +// deliberate `panic!` there — see `PersistentBitMatrix::col_view`'s own +// docs) — a real column, however cheap, could only ever be synthesised by +// scanning every row, which is exactly the row-major work this function +// already does directly, natively, via +// `PersistentSparseBitMatrix::for_each_genome_in_row` (touches only the +// columns actually present per row — 1 for a singleton row, k̄ for a +// multi-genome row). +// - Even for `Columnar`/`Packed`, doing this once for *every* group +// together is strictly less work than the old per-group loop whenever two +// groups share a column (`--aggregate-by`'s auto-groups are disjoint, but +// nothing stops `--group` from overlapping) — a column referenced by `g` +// groups is now read from disk once, not `g` times. +// +// For a bit matrix, presence count at threshold 1 is the *only* primitive +// needed: `sum` (bit matrix: sum of 0/1 = presence count), `any` +// (count ≥ 1), `max` (= any), `min`/`all` (count == group size) are all +// cheap derivations from this one vector — see +// `obikselect::select_layer::agg_result_from_count`, the caller this +// exists for. +pub fn batch_presence_counts( + mat: &PersistentBitMatrix, + groups: &[ColGroup], +) -> io::Result> { + let n = mat.n(); + + // Column -> every group index that references it — built once, so a + // shared column is attributed to every one of its groups from a single + // read, whichever storage format this turns out to be. + let mut col_to_groups: HashMap> = HashMap::new(); + for (gi, g) in groups.iter().enumerate() { + for &c in &g.indices { + col_to_groups.entry(c).or_default().push(gi); + } + } + + let mut builders: Vec = groups + .iter() + .map(|_| TempCompactIntVecBuilder::new(n)) + .collect::>()?; + + let bump = |builders: &mut [TempCompactIntVecBuilder], gi: usize, slot: usize| { + let b = &mut builders[gi]; + let v = b.get(slot); + b.set(slot, v + 1); + }; + + match mat { + PersistentBitMatrix::Sparse(m) => { + // Native row-major decode — the whole reason this exists. + for slot in 0..n { + m.for_each_genome_in_row(slot, |col| { + if let Some(gs) = col_to_groups.get(&col) { + for &gi in gs { + bump(&mut builders, gi, slot); + } + } + }); + } + } + PersistentBitMatrix::Implicit { n_rows, .. } => { + // Single column (index 0), present at every slot. + if let Some(gs) = col_to_groups.get(&0) { + for &gi in gs { + for slot in 0..*n_rows { + bump(&mut builders, gi, slot); + } + } + } + } + PersistentBitMatrix::Columnar(_) | PersistentBitMatrix::Packed(_) => { + // One `col_view` per *distinct referenced column*, not per + // group — the common case (a column belongs to exactly one + // group) reuses the existing fast bulk primitive directly; + // only a column shared by several groups falls back to a + // per-bit scan. + for (&col, gs) in &col_to_groups { + let view = mat.col_view(col); + match gs.as_slice() { + [gi] => builders[*gi].inc_present_fast(view), + _ => { + for (slot, present) in view.iter().enumerate() { + if present { + for &gi in gs { + bump(&mut builders, gi, slot); + } + } + } + } + } + } + } + } + + builders.into_iter().map(|b| b.freeze()).collect() +} + // ── MatrixGroupOps ──────────────────────────────────────────────────────────── impl MatrixGroupOps for PersistentBitMatrix { diff --git a/src/obicompactvec/src/bitmatrix/mod.rs b/src/obicompactvec/src/bitmatrix/mod.rs index 4aefa62b..1a69806d 100644 --- a/src/obicompactvec/src/bitmatrix/mod.rs +++ b/src/obicompactvec/src/bitmatrix/mod.rs @@ -20,6 +20,7 @@ mod persistent; mod sparse; pub use builder::PersistentBitMatrixBuilder; +pub use group_ops::batch_presence_counts; pub use packed::pack_bit_matrix; pub use persistent::PersistentBitMatrix; pub use sparse::{PersistentSparseBitMatrix, PersistentSparseBitMatrixBuilder, pack_sparse_bit_matrix}; diff --git a/src/obicompactvec/src/lib.rs b/src/obicompactvec/src/lib.rs index a965ff39..3fbe47cd 100644 --- a/src/obicompactvec/src/lib.rs +++ b/src/obicompactvec/src/lib.rs @@ -21,7 +21,7 @@ mod views; pub use bitmatrix::{ PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentSparseBitMatrix, - PersistentSparseBitMatrixBuilder, pack_bit_matrix, pack_sparse_bit_matrix, + PersistentSparseBitMatrixBuilder, batch_presence_counts, pack_bit_matrix, pack_sparse_bit_matrix, }; pub use bitvec::{BitIter, PersistentBitVec, PersistentBitVecBuilder}; pub use builder::PersistentCompactIntVecBuilder; diff --git a/src/obicompactvec/src/tests/colgroup.rs b/src/obicompactvec/src/tests/colgroup.rs index b11ba39f..4d88a7f2 100644 --- a/src/obicompactvec/src/tests/colgroup.rs +++ b/src/obicompactvec/src/tests/colgroup.rs @@ -3,6 +3,7 @@ use tempfile::tempdir; use crate::{ ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentIntMatrix, PersistentCompactIntMatrixBuilder, + PersistentSparseBitMatrixBuilder, batch_presence_counts, }; use crate::{PersistentBitVecBuilder, PersistentCompactIntVec, PersistentCompactIntVecBuilder}; @@ -234,3 +235,66 @@ fn int_presence_count_additive_across_split() { assert_eq!(pb.get(1), 2); assert_eq!(pb.get(2), 1); } + +// ── batch_presence_counts: Columnar vs Sparse, disjoint and overlapping groups ── + +/// `batch_presence_counts` must agree with the naive "count group members +/// present at this row" definition, for both a `Columnar` (col_view-driven) +/// and a `Sparse` (native row-major) source — and in particular must still +/// be correct when two groups share a column (`g1`/`g3` share column 1, +/// `g2`/`g3` share column 2), the case that forces the per-bit fallback +/// instead of the single-target `inc_present_fast` fast path. +#[test] +fn batch_presence_counts_matches_naive_for_columnar_and_sparse() { + // Column-major: cols[c][row] + let cols: [&[bool]; 4] = [ + &[true, false, true, false, true], // col0 + &[true, true, false, false, true], // col1 + &[false, true, true, true, false], // col2 + &[false, false, false, true, true], // col3 + ]; + let (_dcol, columnar) = make_bit_matrix(&cols); + + // Same data, row-major, for the sparse builder. + let n_rows = 5; + let n_cols = 4; + let dir = tempdir().unwrap(); + let sparse_dir = dir.path().join("sparse"); + let mut b = PersistentSparseBitMatrixBuilder::new(n_rows, n_cols, &sparse_dir).unwrap(); + let mut genomes = Vec::new(); + for row in 0..n_rows { + genomes.clear(); + genomes.extend((0..n_cols).filter(|&c| cols[c][row]).map(|c| c as u32)); + b.push_row(&genomes); + } + let sparse = PersistentBitMatrix::Sparse(b.finish().unwrap()); + + let groups = [ + ColGroup::new("g1", vec![0, 1]), + ColGroup::new("g2", vec![2, 3]), + ColGroup::new("g3", vec![1, 2]), // overlaps g1 (col1) and g2 (col2) + ]; + + let expected: [[u32; 5]; 3] = [ + [2, 1, 1, 0, 2], // g1: col0+col1 + [0, 1, 1, 2, 1], // g2: col2+col3 + [1, 2, 1, 1, 1], // g3: col1+col2 + ]; + + for (label, mat) in [("columnar", &columnar), ("sparse", &sparse)] { + let counts = batch_presence_counts(mat, &groups).unwrap(); + assert_eq!(counts.len(), groups.len(), "{label}: wrong number of results"); + for (gi, exp) in expected.iter().enumerate() { + for (row, &e) in exp.iter().enumerate() { + assert_eq!(counts[gi].get(row), e, "{label}: group {gi} row {row}"); + } + } + } +} + +#[test] +fn batch_presence_counts_empty_groups_is_empty() { + let (_d, m) = make_bit_matrix(&[&[true, false], &[false, true]]); + let counts = batch_presence_counts(&m, &[]).unwrap(); + assert!(counts.is_empty()); +} diff --git a/src/obikselect/src/select_layer.rs b/src/obikselect/src/select_layer.rs index 9fda0cea..d2630ce0 100644 --- a/src/obikselect/src/select_layer.rs +++ b/src/obikselect/src/select_layer.rs @@ -12,7 +12,7 @@ use std::path::Path; use obicompactvec::{ ColGroup, MatrixBuilder, MatrixGroupOps, PersistentBitMatrix, PersistentIntMatrix, TempBitVec, - TempCompactIntVec, + TempBitVecBuilder, TempCompactIntVec, TempCompactIntVecBuilder, batch_presence_counts, }; use obikindex::layer::{KmerLayer, LayerContent}; use obikindex::{KmerIndex, OKIError, OKIResult}; @@ -93,6 +93,44 @@ fn compute_group( }) } +/// A bit-matrix source's every `AggOp` is a cheap derivation of one shared +/// presence count (see [`obicompactvec::batch_presence_counts`]'s own +/// docs for why: `sum` = the count itself, `any`/`max` = `count ≥ 1`, +/// `all`/`min` = `count == group size`, `none` = `count == 0`) — no further +/// matrix access needed once `count` is in hand. +fn agg_result_from_count(op: AggOp, group_len: usize, count: TempCompactIntVec) -> io::Result { + let n = count.len(); + let group_len = group_len as u32; + Ok(match op { + AggOp::Sum => AggResult::Int(count), + AggOp::Any => { + let mut b = TempBitVecBuilder::new(n)?; + b.or_where(count.view(), |v| v >= 1); + AggResult::Bit(b.freeze()?) + } + AggOp::All => { + let mut b = TempBitVecBuilder::new(n)?; + b.or_where(count.view(), |v| v == group_len); + AggResult::Bit(b.freeze()?) + } + AggOp::None => { + let mut b = TempBitVecBuilder::new(n)?; + b.or_where(count.view(), |v| v == 0); + AggResult::Bit(b.freeze()?) + } + AggOp::Min => { + let mut b = TempCompactIntVecBuilder::new(n)?; + b.inc_predicate_fast(count.view(), |v| v == group_len); + AggResult::Int(b.freeze()?) + } + AggOp::Max => { + let mut b = TempCompactIntVecBuilder::new(n)?; + b.inc_predicate_fast(count.view(), |v| v >= 1); + AggResult::Int(b.freeze()?) + } + }) +} + // ── AggResult → MatrixBuilder ───────────────────────────────────────────────── /// Add one already-aggregated column to `mb` — the only piece `MatrixBuilder` @@ -161,15 +199,6 @@ pub(crate) fn select_partition( .to_path_buf(); copy_layer_files(src_layer.dir(), &dst_layer_dir, force_copy).map_err(OKIError::Io)?; - let group_mat: Box = match src_layer.content() { - LayerContent::Count => { - Box::new(PersistentIntMatrix::open(src_layer.dir()).map_err(OKIError::Io)?) - } - LayerContent::Presence => { - Box::new(PersistentBitMatrix::open(src_layer.dir()).map_err(OKIError::Io)?) - } - }; - let data_subdir = if output_presence { "presence" } else { @@ -180,10 +209,37 @@ pub(crate) fn select_partition( let mut builder = MatrixBuilder::new(output_presence, n, &data_dir).map_err(OKIError::Io)?; - for spec in specs { - let r = compute_group(group_mat.as_ref(), spec, threshold).map_err(OKIError::Io)?; - add_result(&mut builder, r).map_err(OKIError::Io)?; + + match src_layer.content() { + // One shared row-major (Sparse) / deduplicated column-major + // (Columnar/Packed/Implicit) pass covers every spec at once — + // see `batch_presence_counts`'s own docs. Every `AggOp` for a + // bit-matrix source is a cheap derivation of that one count. + LayerContent::Presence => { + let mat = PersistentBitMatrix::open(src_layer.dir()).map_err(OKIError::Io)?; + let groups: Vec = specs + .iter() + .map(|s| ColGroup::new(s.label.clone(), s.indices.clone())) + .collect(); + let counts = batch_presence_counts(&mat, &groups).map_err(OKIError::Io)?; + for (spec, count) in specs.iter().zip(counts) { + let r = agg_result_from_count(spec.op, spec.indices.len(), count) + .map_err(OKIError::Io)?; + add_result(&mut builder, r).map_err(OKIError::Io)?; + } + } + // Count-matrix `sum`/`min`/`max` are genuine per-value + // reductions, not derivable from a single presence count — + // unchanged, one `col_view`-driven pass per spec. + LayerContent::Count => { + let mat = PersistentIntMatrix::open(src_layer.dir()).map_err(OKIError::Io)?; + for spec in specs { + let r = compute_group(&mat, spec, threshold).map_err(OKIError::Io)?; + add_result(&mut builder, r).map_err(OKIError::Io)?; + } + } } + builder.close().map_err(OKIError::Io)?; } diff --git a/src/obisys/src/numa/runner.rs b/src/obisys/src/numa/runner.rs index ffe2db7e..29c583b8 100644 --- a/src/obisys/src/numa/runner.rs +++ b/src/obisys/src/numa/runner.rs @@ -1,3 +1,5 @@ +use std::any::Any; +use std::panic::{self, AssertUnwindSafe}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -164,6 +166,7 @@ impl PartitionRunner { let f = &f; let mut first_err: Option = None; + let mut first_panic: Option> = None; std::thread::scope(|s| { // ── Timer thread ────────────────────────────────────────────────── @@ -208,12 +211,27 @@ impl PartitionRunner { for i in &prx { debug!(?tid, partition = i, "PartitionRunner worker: picked partition"); let t = Instant::now(); - let r = match &pool { + // Caught, not left to unwind straight through this + // spawned thread: a panicking `f(i)` would + // otherwise never reach `etx.send(...)` below, so + // the controller's `completed < n_total` loop + // waits forever for an event this partition can + // no longer produce (see `run`'s own docs on the + // termination protocol) — silently hanging + // instead of surfacing the panic. Caught here and + // re-raised on the caller's thread once `run` + // returns, so the original message/backtrace + // still surfaces, just from the right place. + let outcome = panic::catch_unwind(AssertUnwindSafe(|| match &pool { Some(p) => p.install(|| f(i)), None => f(i), - }; + })); debug!(?tid, partition = i, "PartitionRunner worker: partition done"); - etx.send(WorkerEvent::Completed(i, r, t.elapsed())).ok(); + let event = match outcome { + Ok(r) => WorkerEvent::Completed(i, r, t.elapsed()), + Err(payload) => WorkerEvent::Panicked(i, payload), + }; + etx.send(event).ok(); } debug!(?tid, "PartitionRunner worker: no more partitions, exiting"); }); @@ -260,6 +278,21 @@ impl PartitionRunner { n_total, ); } + WorkerEvent::Panicked(_i, payload) => { + // Counts toward `completed` like any other outcome — + // this partition will never produce a `Completed` + // event, so not counting it here is exactly what + // used to hang the controller forever. The payload + // is re-raised on the caller's thread once `run` + // returns (see below), not here: unwinding out of + // this `recv` loop would leak the still-running + // worker/timer threads this `thread::scope` owns. + if first_panic.is_none() { + first_panic = Some(payload); + } + completed += 1; + reset_tx.send(()).ok(); + } WorkerEvent::TimerTick => { maybe_activate( &mut activation, @@ -279,6 +312,15 @@ impl PartitionRunner { drop(reset_tx); }); + // A panic takes priority over a plain `Err`: it means `f` itself hit + // a bug (an unhandled case, an assertion) rather than a normal, + // typed failure — worth surfacing with its original message/ + // backtrace via unwinding, not silently downgraded to whatever `Err` + // another, unrelated partition happened to return first. + if let Some(payload) = first_panic { + panic::resume_unwind(payload); + } + match first_err { Some(e) => Err(e), None => Ok(()), @@ -290,6 +332,11 @@ impl PartitionRunner { enum WorkerEvent { Completed(usize, Result, Duration), + /// `f(i)` panicked instead of returning — see `run`'s own docs on why + /// this is caught at all (never letting the panic unwind straight + /// through the spawned worker thread) rather than a bug fix that could + /// be skipped. + Panicked(usize, Box), TimerTick, }