Remove in-place flag, optimize presence counting, and handle NUMA panics

The `obikmer select` CLI no longer supports in-place index rewriting; the `--output` flag is now required, with benchmarks updated to use temporary directories for atomic replacement. Added `--dense` and `--force-copy` flags. Introduced `batch_presence_counts` to compute presence counts across multiple column groups in a single pass, eliminating redundant I/O. Refactored the aggregation pipeline to branch on layer content, applying the optimized batched counting for `Presence` layers. Enhanced the NUMA runner to catch worker panics, track them, and re-raise after thread join to prevent indefinite blocking.
This commit is contained in:
Eric Coissac
2026-08-28 23:16:37 +02:00
parent 7363370065
commit 93fe838f97
9 changed files with 330 additions and 26 deletions
+18 -5
View File
@@ -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
```
+11 -2
View File
@@ -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
+11 -2
View File
@@ -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
@@ -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<Vec<TempCompactIntVec>> {
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<usize, Vec<usize>> = 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<TempCompactIntVecBuilder> = groups
.iter()
.map(|_| TempCompactIntVecBuilder::new(n))
.collect::<io::Result<_>>()?;
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 {
+1
View File
@@ -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};
+1 -1
View File
@@ -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;
+64
View File
@@ -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());
}
+69 -13
View File
@@ -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<AggResult> {
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<dyn MatrixGroupOps> = 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<ColGroup> = 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)?;
}
+50 -3
View File
@@ -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<E> = None;
let mut first_panic: Option<Box<dyn Any + Send + 'static>> = 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<R, E> {
Completed(usize, Result<R, E>, 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<dyn Any + Send + 'static>),
TimerTick,
}