From 1910ed2f21f71a4dc880d55984b6dbbe49489370 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Mon, 21 Sep 2026 19:11:23 +0200 Subject: [PATCH 01/35] feat(run-experiment): drive grouped runs, sharing the adapted bands across runs `run-experiment` refused `groups.window_groups > 1`, so a grouped search of several files meant one `mumdia run` per file and assembling the pooled rescore, the match between runs and the cross-run quantities by hand. It now routes each run's middle through `run_groups`: the run's competed table and chromatograms come back under the names the pooled stages already read, so the experiment-wide rescore, the per-source `run_psm_q`, MBR, per-run quant and the cross-run report are unchanged. `experiment.rt_library_scope` carries over. A grouped run adapts its bands rather than one library table, so the first run's `groups/` directory is what the rest reuse: each band takes that run's table for its own band, keeping the ids and row order that `fragment_offset` and the per-band seed views key on, and still fits its own per-run LOESS on top. Under `per_run` every run adapts its own bands as before. A run whose shared bands do not match its plan fails with the missing path named rather than searching a mismatched band. `run_groups` takes its manifest as an `Option`, because `run-experiment` keeps one experiment-level manifest rather than one per run. Measured on the fixture: a two-file grouped experiment plans three bands per run, pools each run to 288 competed rows and rescores 576 PSMs experiment-wide to 151 peptides at 1%, matching the ungrouped fixture experiment. Co-Authored-By: Claude Fable 5.1 --- rust/mumdia/crates/mumdia/src/stages/run.rs | 3 +- .../mumdia/src/stages/run_experiment.rs | 43 +++- .../crates/mumdia/src/stages/run_groups.rs | 220 ++++++++++++------ 3 files changed, 180 insertions(+), 86 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/run.rs b/rust/mumdia/crates/mumdia/src/stages/run.rs index bd5673b5..bd190eb7 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run.rs @@ -334,7 +334,8 @@ pub fn run(p: RunParams) -> Result<()> { lib_precursors: &lib_p, lib_fragments: &lib_f, out_dir: p.out_dir, - man: &mut man, + man: Some(&mut man), + shared_bands: None, mh_heads, library_input: p.lib_precursors.is_some(), })?; diff --git a/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs b/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs index 48f2d0da..1e7f5aed 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs @@ -77,14 +77,6 @@ fn preflight(p: &RunExperimentParams) -> Result<()> { if cfg.rt_im_train.finetune_deeplc && cfg.predict_frag.deeplc_python.is_none() { anyhow::bail!("rt_im_train.finetune_deeplc requires predict_frag.deeplc_python"); } - // The grouped search lives in the single-run orchestrator for now; refusing here beats - // reading the key and searching every run against the whole library anyway. - if cfg.groups.window_groups > 1 { - anyhow::bail!( - "groups.window_groups = {} is not supported by run-experiment yet: run each file with `mumdia run --mzml ` (which searches it group by group) and pool the competed tables with `mumdia rescore --competed a b c`", - cfg.groups.window_groups - ); - } // Explicit count: hard requirement. Automatic default: degrade with a warning, so a // native Python-free experiment stays runnable (`run.rs` says the same). if cfg.rt_im_train.multihead_calibration.is_some_and(|n| n > 0) @@ -203,6 +195,35 @@ fn process_run( top_peaks_ms1: 0, config_hash: &convert_hash, })?; + let has_deeplc = cfg.predict_frag.deeplc_python.is_some(); + let mh_heads = cfg + .rt_im_train + .multihead_heads(has_deeplc, cfg.deeplc_rt_source(library_input, has_deeplc)); + // Grouped: seed, calibration, extract, features and compete happen one isolation-window + // group at a time, and the pooled competed table and chromatograms come back under the + // names the pooled stages below read, exactly as in the single-run orchestrator. The + // third return value is then this run's `groups/` directory rather than an adapted + // library, which is what `experiment.rt_library_scope = first_run_only` hands to the + // runs that follow (`shared_rt_lib` here). + if cfg.groups.window_groups > 1 { + let pooled = crate::stages::run_groups::run(crate::stages::run_groups::GroupRun { + cfg, + config_hash: ch, + converted: &co, + lib_precursors: lib_p_base, + lib_fragments: lib_f, + out_dir: out, + man: None, + shared_bands: shared_rt_lib, + mh_heads, + library_input, + })?; + return Ok(( + pooled.competed, + pooled.chromatograms, + Some(format!("{out}/groups")), + )); + } let seed = d("seed_psms.parquet"); search_seed::run(search_seed::SearchSeedParams { fragment_offset: None, @@ -220,10 +241,6 @@ fn process_run( // drift is then absorbed by `rt_im_train`'s per-run calibration below, which is fitted // separately for every run regardless. let mut produced_rt_lib: Option = None; - let has_deeplc = cfg.predict_frag.deeplc_python.is_some(); - let mh_heads = cfg - .rt_im_train - .multihead_heads(has_deeplc, cfg.deeplc_rt_source(library_input, has_deeplc)); let lib_p = if let Some(shared) = shared_rt_lib { // A previous run already adapted the library and `experiment.rt_library_scope` // says to reuse it. Which mechanism produced it does not matter here: the @@ -664,6 +681,8 @@ pub fn run(p: RunExperimentParams) -> Result<()> { .rt_im_train .multihead_heads(has_deeplc_for_scope, deeplc_rt_source_for_scope) > 0; + // A grouped run adapts its bands rather than one library table, and shares them the + // same way: the first run's `groups/` directory is what the rest reuse. let share_ft = adapts_rt_library && matches!( cfg.experiment.rt_library_scope, diff --git a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs index a1f3b97b..4fac773f 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs @@ -36,7 +36,14 @@ pub struct GroupRun<'a> { pub lib_precursors: &'a str, pub lib_fragments: &'a str, pub out_dir: &'a str, - pub man: &'a mut Manifest, + /// Where to record the band and pooled artifacts. `None` under `run-experiment`, + /// which keeps one experiment-level manifest instead of one per run. + pub man: Option<&'a mut Manifest>, + /// A previous run's `groups/` directory, whose per-band precursor tables already carry + /// adapted retention times (`experiment.rt_library_scope = first_run_only`). Set, no + /// band re-predicts: each takes that run's table for its band and fits its own per-run + /// LOESS on top, which is what the ungrouped path does with a shared library. + pub shared_bands: Option<&'a str>, /// Multi-head calibration heads resolved by the caller (0 = off). pub mh_heads: usize, pub library_input: bool, @@ -60,7 +67,25 @@ fn windows_of(path: &str) -> Result> { Ok(lo.into_iter().zip(hi).collect()) } -pub fn run(g: GroupRun) -> Result { +/// The file name a band's adapted precursor table takes, so a later run can find it. +fn band_lib_name(rt_model: &str) -> &'static str { + if rt_model.starts_with("multihead") { + "lib_precursors_multihead.parquet" + } else if rt_model == "finetuned" { + "lib_precursors_ft.parquet" + } else { + "lib_precursors_deeplc.parquet" + } +} + +/// Record an artifact when this run keeps a manifest of its own. +fn record_opt(man: Option<&mut Manifest>, rec: mumdia_core::manifest::ArtifactRecord) { + if let Some(m) = man { + m.record(rec); + } +} + +pub fn run(mut g: GroupRun) -> Result { let t0 = Instant::now(); let cfg = g.cfg; let ch = g.config_hash; @@ -151,14 +176,17 @@ pub fn run(g: GroupRun) -> Result { config_hash: ch, fragment_offset: Some(first as u32), })?; - g.man.record(record_artifact( - &format!("{}[g{:02}]", artifact::SEED_PSMS.0, b.index), - artifact::SEED_PSMS, - &seed, - rows, - "search-seed", - ch, - )?); + record_opt( + g.man.as_deref_mut(), + record_artifact( + &format!("{}[g{:02}]", artifact::SEED_PSMS.0, b.index), + artifact::SEED_PSMS, + &seed, + rows, + "search-seed", + ch, + )?, + ); bands.push(Band { index: b.index, offset: first as u32, @@ -193,14 +221,17 @@ pub fn run(g: GroupRun) -> Result { masscals: &masscals, out: &pooled_seed, })?; - g.man.record(record_artifact( - artifact::SEED_PSMS.0, - artifact::SEED_PSMS, - &pooled_seed, - n, - "seed-pool", - ch, - )?); + record_opt( + g.man.as_deref_mut(), + record_artifact( + artifact::SEED_PSMS.0, + artifact::SEED_PSMS, + &pooled_seed, + n, + "seed-pool", + ch, + )?, + ); let global = cfg.groups.calibration == GroupCalibration::Global; // --- RT model per band, against the pooled or the band's own anchors @@ -221,7 +252,30 @@ pub fn run(g: GroupRun) -> Result { "library".to_string() }; let mut repredicted: Vec<(String, u32)> = Vec::new(); - for b in &mut bands { + if let Some(shared) = g.shared_bands { + // Take a previous run's adapted bands wholesale: same ids, same row order, which is + // what `fragment_offset` and the per-band seed views key on. Only the retention + // times inside differ from the raw library, and this run still fits its own LOESS. + let name = band_lib_name(&rt_model); + for b in &mut bands { + let from = format!("{shared}/g{:02}/{name}", b.index); + if !std::path::Path::new(&from).exists() { + bail!( + "groups: {from} is missing, so the shared bands do not match this run: \ + reusing them needs the same group plan and the same retention-time model" + ); + } + repredicted.push((from.clone(), b.offset)); + b.prec = from; + } + info!( + groups = bands.len(), + model = %rt_model, + source = %shared, + "groups: reusing a previous run's adapted bands" + ); + } + for b in bands.iter_mut().filter(|_| g.shared_bands.is_none()) { let anchors = if global { &pooled_seed } else { &b.seed }; let out = if g.mh_heads > 0 { let out = gd(b.index, "lib_precursors_multihead.parquet"); @@ -271,18 +325,21 @@ pub fn run(g: GroupRun) -> Result { }; if let Some(out) = out { let rows = mumdia_io::table::nrows(&out)?; - g.man.record(record_artifact( - &format!( - "{}[g{:02}]", - artifact::FRAGMENT_LIBRARY_PRECURSORS.0, - b.index - ), - artifact::FRAGMENT_LIBRARY_PRECURSORS, - &out, - rows, - &rt_model, - ch, - )?); + record_opt( + g.man.as_deref_mut(), + record_artifact( + &format!( + "{}[g{:02}]", + artifact::FRAGMENT_LIBRARY_PRECURSORS.0, + b.index + ), + artifact::FRAGMENT_LIBRARY_PRECURSORS, + &out, + rows, + &rt_model, + ch, + )?, + ); repredicted.push((out.clone(), b.offset)); b.prec = out; } @@ -321,14 +378,17 @@ pub fn run(g: GroupRun) -> Result { config_hash: ch, anchor_irt_from_seed: from_seed, })?; - g.man.record(record_artifact( - &format!("{}[g{:02}]", artifact::RUN_WINDOWS.0, b.index), - artifact::RUN_WINDOWS, - &windows, - rows, - "rt-im-train", - ch, - )?); + record_opt( + g.man.as_deref_mut(), + record_artifact( + &format!("{}[g{:02}]", artifact::RUN_WINDOWS.0, b.index), + artifact::RUN_WINDOWS, + &windows, + rows, + "rt-im-train", + ch, + )?, + ); let mass_cal = if global { format!("{pooled_seed}.masscal.json") } else { @@ -351,22 +411,28 @@ pub fn run(g: GroupRun) -> Result { config_hash: ch, fragment_offset: Some(b.offset), })?; - g.man.record(record_artifact( - &format!("{}[g{:02}]", artifact::PSMS_EXTRACTED.0, b.index), - artifact::PSMS_EXTRACTED, - &psms, - npsm, - "extract", - ch, - )?); - g.man.record(record_artifact( - &format!("{}[g{:02}]", artifact::CHROMATOGRAMS.0, b.index), - artifact::CHROMATOGRAMS, - &chrom, - nchr, - "extract", - ch, - )?); + record_opt( + g.man.as_deref_mut(), + record_artifact( + &format!("{}[g{:02}]", artifact::PSMS_EXTRACTED.0, b.index), + artifact::PSMS_EXTRACTED, + &psms, + npsm, + "extract", + ch, + )?, + ); + record_opt( + g.man.as_deref_mut(), + record_artifact( + &format!("{}[g{:02}]", artifact::CHROMATOGRAMS.0, b.index), + artifact::CHROMATOGRAMS, + &chrom, + nchr, + "extract", + ch, + )?, + ); let feats = gd(b.index, "features.parquet"); let pin = gd(b.index, "run.pin"); info!(stage = %"features", group = b.index, "run: stage start"); @@ -381,14 +447,17 @@ pub fn run(g: GroupRun) -> Result { cfg: &cfg.features, config_hash: ch, })?; - g.man.record(record_artifact( - &format!("{}[g{:02}]", artifact::FEATURES.0, b.index), - artifact::FEATURES, - &feats, - nf, - "features", - ch, - )?); + record_opt( + g.man.as_deref_mut(), + record_artifact( + &format!("{}[g{:02}]", artifact::FEATURES.0, b.index), + artifact::FEATURES, + &feats, + nf, + "features", + ch, + )?, + ); let competed = gd(b.index, "psms_competed.parquet"); info!(stage = %"compete", group = b.index, "run: stage start"); let nc = compete::run(compete::CompeteParams { @@ -397,14 +466,17 @@ pub fn run(g: GroupRun) -> Result { cfg: &cfg.compete, config_hash: ch, })?; - g.man.record(record_artifact( - &format!("{}[g{:02}]", artifact::PSMS_COMPETED.0, b.index), - artifact::PSMS_COMPETED, - &competed, - nc, - "compete", - ch, - )?); + record_opt( + g.man.as_deref_mut(), + record_artifact( + &format!("{}[g{:02}]", artifact::PSMS_COMPETED.0, b.index), + artifact::PSMS_COMPETED, + &competed, + nc, + "compete", + ch, + )?, + ); arts.push(pool::BandArtifacts { offset: b.offset, psms, @@ -461,8 +533,10 @@ pub fn run(g: GroupRun) -> Result { stats.competed, ), ] { - g.man - .record(record_artifact(name, schema, path, rows, "pool", ch)?); + record_opt( + g.man.as_deref_mut(), + record_artifact(name, schema, path, rows, "pool", ch)?, + ); ArtifactReport { logical_name: name.to_string(), schema_name: name.to_string(), From b8292fdb3ca940b979ef84ea39f3e2b5b8b9e673 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Mon, 21 Sep 2026 21:14:07 +0200 Subject: [PATCH 02/35] feat(groups): search several bands at once inside a run (groups.parallel) A grouped run searched its bands one after another, so a large machine sat idle: measured on a seven-file experiment with four runs in flight, basilisk was 97.6% idle with one process using 12 of 256 cores, because four bands (one per run) were all that was ever in flight, each using about three. Both per-band phases now run `groups.parallel` bands at a time: the slice-and-seed phase and the windows-extract-features-compete phase. Chunked rather than a free-running pool, because the chunk is exactly how many extraction working sets are resident, which is what banding exists to bound. Each band's manifest records are collected and merged in band order, so the manifest is what the sequential run wrote, and the artifacts are pooled in band order regardless of which band finishes first. The default stays 1, which is the previous behaviour exactly. Raise it after checking one band's peak RSS: on the 203M-precursor library at 63 bands the largest band took 39 GB and the median far less, so the low-m/z bands, which are both the largest and the slowest, set the ceiling. Byte-identical output: on the fixture, three bands searched three at a time reproduce the sequential grouped run's `psms_extracted`, `chromatograms`, `features`, `psms_competed`, `psms_scored` and `peptides.tsv` exactly. Co-Authored-By: Claude Fable 5.1 --- configs/config-schema.json | 15 +- docs/24_config_reference.md | 7 +- rust/mumdia/crates/mumdia-core/src/config.rs | 9 + .../crates/mumdia/src/stages/run_groups.rs | 241 ++++++++++-------- 4 files changed, 166 insertions(+), 106 deletions(-) diff --git a/configs/config-schema.json b/configs/config-schema.json index 2c15fef3..29621cff 100644 --- a/configs/config-schema.json +++ b/configs/config-schema.json @@ -33,7 +33,7 @@ "default_text": "0", "help": "", "gates": [], - "source_line": 1921 + "source_line": 1929 }, { "path": "convert.thermo_raw_parser", @@ -2446,6 +2446,19 @@ "global", "per_group" ] + }, + { + "path": "groups.parallel", + "name": "parallel", + "section": "groups", + "rust_type": "usize", + "kind": "integer", + "optional": false, + "default": 1, + "default_text": "1", + "help": "Bands searched at the same time inside one run. `1` (the default) is one band at a time, which is what bounds the memory: each band in flight holds its own extraction working set, so the peak is this many bands' worth. Raise it to fill a large machine, after checking one band's peak RSS: on a 203M-precursor library at 63 bands the largest band took 39 GB and the median far less. Results do not depend on it; bands are independent and their artifacts are pooled in band order either way.", + "gates": [], + "source_line": 1912 } ], "list_sections": [ diff --git a/docs/24_config_reference.md b/docs/24_config_reference.md index a192130f..af4b85c7 100644 --- a/docs/24_config_reference.md +++ b/docs/24_config_reference.md @@ -68,12 +68,12 @@ undocumented on purpose; those fields are counted under "Coverage". | [`quant`](#quant) | `QuantConfig` | 17 | [docs/12_quant_lfq_align_mbr_report_audit.md](12_quant_lfq_align_mbr_report_audit.md) | | [`mbr`](#mbr) | `MbrConfig` | 9 | [docs/12_quant_lfq_align_mbr_report_audit.md](12_quant_lfq_align_mbr_report_audit.md) | | [`experiment`](#experiment) | `ExperimentConfig` | 2 | [docs/01_overview_and_dataflow.md](01_overview_and_dataflow.md) | -| [`groups`](#groups) | `GroupsConfig` | 2 | | +| [`groups`](#groups) | `GroupsConfig` | 3 | | | [`peptidoforms.fixed_mods[] / peptidoforms.variable_mods[]`](#peptidoformsfixed_mods--peptidoformsvariable_mods) | `ResidueMod` | 2 | [docs/05_digest_peptidoforms.md](05_digest_peptidoforms.md) | ## (top level) -`Config` (rust/mumdia/crates/mumdia-core/src/config.rs:1918). stage document: [docs/02_config_and_data_model.md](02_config_and_data_model.md). +`Config` (rust/mumdia/crates/mumdia-core/src/config.rs:1926). stage document: [docs/02_config_and_data_model.md](02_config_and_data_model.md). | Field | Type | Default | Gated | Description | |---|---|---|---|---| @@ -396,6 +396,7 @@ Searching a run one isolation-window group at a time. A group of isolation windo |---|---|---|---|---| | `window_groups` | `usize` | `1` | | Number of window groups. `1` (the default) is the ordinary single-library search. Groups are contiguous bands of isolation windows balanced by the number of library precursors they select, read from the precursor table's row-group statistics. | | `calibration` | `GroupCalibration` | `global` | | Anchors for the RT calibration of each group; see `GroupCalibration`. | +| `parallel` | `usize` | `1` | | Bands searched at the same time inside one run. `1` (the default) is one band at a time, which is what bounds the memory: each band in flight holds its own extraction working set, so the peak is this many bands' worth. Raise it to fill a large machine, after checking one band's peak RSS: on a 203M-precursor library at 63 bands the largest band took 39 GB and the median far less. Results do not depend on it; bands are independent and their artifacts are pooled in band order either way. | ## peptidoforms.fixed_mods[] / peptidoforms.variable_mods[] @@ -852,6 +853,6 @@ Every field whose struct has an `impl Default` resolved from the source. ## Coverage -19 structs and 192 fields emitted from `rust/mumdia/crates/mumdia-core/src/config.rs`, plus 25 enumerations, 1 named profile(s), 66 environment variables read and 19 set. +19 structs and 193 fields emitted from `rust/mumdia/crates/mumdia-core/src/config.rs`, plus 25 enumerations, 1 named profile(s), 66 environment variables read and 19 set. 20 field(s) carry a gating marker in their doc comment. 48 field(s) carry no doc comment at all, so their description is empty above. 0 default(s) could not be resolved and 2 have none by design. diff --git a/rust/mumdia/crates/mumdia-core/src/config.rs b/rust/mumdia/crates/mumdia-core/src/config.rs index cca80bd2..b2efb5f9 100644 --- a/rust/mumdia/crates/mumdia-core/src/config.rs +++ b/rust/mumdia/crates/mumdia-core/src/config.rs @@ -1903,12 +1903,20 @@ pub struct GroupsConfig { pub window_groups: usize, /// Anchors for the RT calibration of each group; see [`GroupCalibration`]. pub calibration: GroupCalibration, + /// Bands searched at the same time inside one run. `1` (the default) is one band at a + /// time, which is what bounds the memory: each band in flight holds its own extraction + /// working set, so the peak is this many bands' worth. Raise it to fill a large + /// machine, after checking one band's peak RSS: on a 203M-precursor library at 63 + /// bands the largest band took 39 GB and the median far less. Results do not depend on + /// it; bands are independent and their artifacts are pooled in band order either way. + pub parallel: usize, } impl Default for GroupsConfig { fn default() -> Self { Self { window_groups: 1, calibration: GroupCalibration::Global, + parallel: 1, } } } @@ -2344,6 +2352,7 @@ impl Config { ("experiment.parallel_runs", self.experiment.parallel_runs), ("rescore.seeds", self.rescore.seeds), ("groups.window_groups", self.groups.window_groups), + ("groups.parallel", self.groups.parallel), ] { if value == 0 { return Err(Invalid(format!("{name} must be >= 1"))); diff --git a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs index 4fac773f..7e30ecbc 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs @@ -21,6 +21,7 @@ use mumdia_core::schema::artifact; use mumdia_io::record_artifact; use mumdia_io::report::ArtifactReport; use mumdia_io::table::TableFile; +use rayon::prelude::*; use serde_json::json; use tracing::info; @@ -151,50 +152,72 @@ pub fn run(mut g: GroupRun) -> Result { seed_view: String, } let mut bands: Vec = Vec::new(); - for b in &plan.bands { - let (first, n) = Library::precursor_row_span(g.lib_precursors, b.mz_lo, b.mz_hi)?; - if n == 0 { - info!( - group = b.index, - "groups: band selects no precursor; skipped" - ); - continue; - } - std::fs::create_dir_all(gd(b.index, ""))?; - let prec = gd(b.index, "lib_precursors.parquet"); - info!(stage = %"band-slice", group = b.index, rows = n, "run: stage start"); - groups::write_band_slice(g.lib_precursors, first, n, &prec)?; - let seed = gd(b.index, "seed_psms.parquet"); - info!(stage = %"search-seed", group = b.index, "run: stage start"); - let rows = search_seed::run(search_seed::SearchSeedParams { - ms2: &g.converted.ms2, - library_precursors: &prec, - library_fragments: g.lib_fragments, - out: &seed, - cfg: &cfg.search_seed, - bucket_size: cfg.extract.bucket_size, - config_hash: ch, - fragment_offset: Some(first as u32), - })?; - record_opt( - g.man.as_deref_mut(), - record_artifact( + let par = cfg.groups.parallel.max(1); + // Bands are independent, so `groups.parallel` of them are sliced and seeded at once. + // Chunked rather than a free-running pool: the chunk bounds how many extraction working + // sets are resident, which is the whole point of banding. Results do not depend on it, + // and the records below are merged in band order. + let slice_one = + |b: &groups::Band| -> Result)>> { + let (first, n) = Library::precursor_row_span(g.lib_precursors, b.mz_lo, b.mz_hi)?; + if n == 0 { + info!( + group = b.index, + "groups: band selects no precursor; skipped" + ); + return Ok(None); + } + std::fs::create_dir_all(gd(b.index, ""))?; + let prec = gd(b.index, "lib_precursors.parquet"); + info!(stage = %"band-slice", group = b.index, rows = n, "run: stage start"); + groups::write_band_slice(g.lib_precursors, first, n, &prec)?; + let seed = gd(b.index, "seed_psms.parquet"); + info!(stage = %"search-seed", group = b.index, "run: stage start"); + let rows = search_seed::run(search_seed::SearchSeedParams { + ms2: &g.converted.ms2, + library_precursors: &prec, + library_fragments: g.lib_fragments, + out: &seed, + cfg: &cfg.search_seed, + bucket_size: cfg.extract.bucket_size, + config_hash: ch, + fragment_offset: Some(first as u32), + })?; + let rec = vec![record_artifact( &format!("{}[g{:02}]", artifact::SEED_PSMS.0, b.index), artifact::SEED_PSMS, &seed, rows, "search-seed", ch, - )?, - ); - bands.push(Band { - index: b.index, - offset: first as u32, - n, - prec, - seed, - seed_view: gd(b.index, "seed_psms_pooled.parquet"), - }); + )?]; + Ok(Some(( + Band { + index: b.index, + offset: first as u32, + n, + prec, + seed, + seed_view: gd(b.index, "seed_psms_pooled.parquet"), + }, + rec, + ))) + }; + for chunk in plan.bands.chunks(par) { + let done: Vec)>> = if par == 1 { + chunk.iter().map(slice_one).collect::>>()? + } else { + chunk + .par_iter() + .map(slice_one) + .collect::>>()? + }; + for (band, recs) in done.into_iter().flatten() { + for r in recs { + record_opt(g.man.as_deref_mut(), r); + } + bands.push(band); + } } if bands.is_empty() { bail!("groups: no band selects any precursor of the library"); @@ -359,10 +382,17 @@ pub fn run(mut g: GroupRun) -> Result { // --- windows, extract, features, compete per band let mut arts: Vec = Vec::new(); let mut cals: Vec<(usize, String)> = Vec::new(); - for b in &bands { + // `groups.parallel` bands at a time. Each band in flight holds its own extraction + // working set, so this chunk is what the stage's memory scales with; the artifacts are + // pooled in band order regardless of which band finishes first. + let band_one = |b: &Band| -> Result<( + pool::BandArtifacts, + (usize, String), + Vec, + )> { + let mut recs: Vec = Vec::new(); let windows = gd(b.index, "run_windows.parquet"); let cal = gd(b.index, "cal.json"); - cals.push((b.index, cal.clone())); let (seed_for_windows, from_seed) = match (&anchors_for_windows, global) { (Some(refreshed), _) => (refreshed.clone(), true), (None, true) => (pooled_seed.clone(), true), @@ -378,17 +408,14 @@ pub fn run(mut g: GroupRun) -> Result { config_hash: ch, anchor_irt_from_seed: from_seed, })?; - record_opt( - g.man.as_deref_mut(), - record_artifact( - &format!("{}[g{:02}]", artifact::RUN_WINDOWS.0, b.index), - artifact::RUN_WINDOWS, - &windows, - rows, - "rt-im-train", - ch, - )?, - ); + recs.push(record_artifact( + &format!("{}[g{:02}]", artifact::RUN_WINDOWS.0, b.index), + artifact::RUN_WINDOWS, + &windows, + rows, + "rt-im-train", + ch, + )?); let mass_cal = if global { format!("{pooled_seed}.masscal.json") } else { @@ -411,28 +438,22 @@ pub fn run(mut g: GroupRun) -> Result { config_hash: ch, fragment_offset: Some(b.offset), })?; - record_opt( - g.man.as_deref_mut(), - record_artifact( - &format!("{}[g{:02}]", artifact::PSMS_EXTRACTED.0, b.index), - artifact::PSMS_EXTRACTED, - &psms, - npsm, - "extract", - ch, - )?, - ); - record_opt( - g.man.as_deref_mut(), - record_artifact( - &format!("{}[g{:02}]", artifact::CHROMATOGRAMS.0, b.index), - artifact::CHROMATOGRAMS, - &chrom, - nchr, - "extract", - ch, - )?, - ); + recs.push(record_artifact( + &format!("{}[g{:02}]", artifact::PSMS_EXTRACTED.0, b.index), + artifact::PSMS_EXTRACTED, + &psms, + npsm, + "extract", + ch, + )?); + recs.push(record_artifact( + &format!("{}[g{:02}]", artifact::CHROMATOGRAMS.0, b.index), + artifact::CHROMATOGRAMS, + &chrom, + nchr, + "extract", + ch, + )?); let feats = gd(b.index, "features.parquet"); let pin = gd(b.index, "run.pin"); info!(stage = %"features", group = b.index, "run: stage start"); @@ -447,17 +468,14 @@ pub fn run(mut g: GroupRun) -> Result { cfg: &cfg.features, config_hash: ch, })?; - record_opt( - g.man.as_deref_mut(), - record_artifact( - &format!("{}[g{:02}]", artifact::FEATURES.0, b.index), - artifact::FEATURES, - &feats, - nf, - "features", - ch, - )?, - ); + recs.push(record_artifact( + &format!("{}[g{:02}]", artifact::FEATURES.0, b.index), + artifact::FEATURES, + &feats, + nf, + "features", + ch, + )?); let competed = gd(b.index, "psms_competed.parquet"); info!(stage = %"compete", group = b.index, "run: stage start"); let nc = compete::run(compete::CompeteParams { @@ -466,24 +484,43 @@ pub fn run(mut g: GroupRun) -> Result { cfg: &cfg.compete, config_hash: ch, })?; - record_opt( - g.man.as_deref_mut(), - record_artifact( - &format!("{}[g{:02}]", artifact::PSMS_COMPETED.0, b.index), - artifact::PSMS_COMPETED, - &competed, - nc, - "compete", - ch, - )?, - ); - arts.push(pool::BandArtifacts { - offset: b.offset, - psms, - chromatograms: chrom, - features: feats, - competed, - }); + recs.push(record_artifact( + &format!("{}[g{:02}]", artifact::PSMS_COMPETED.0, b.index), + artifact::PSMS_COMPETED, + &competed, + nc, + "compete", + ch, + )?); + Ok(( + pool::BandArtifacts { + offset: b.offset, + psms, + chromatograms: chrom, + features: feats, + competed, + }, + (b.index, gd(b.index, "cal.json")), + recs, + )) + }; + for chunk in bands.chunks(par) { + let done: Vec<( + pool::BandArtifacts, + (usize, String), + Vec, + )> = if par == 1 { + chunk.iter().map(band_one).collect::>>()? + } else { + chunk.par_iter().map(band_one).collect::>>()? + }; + for (art, cal, recs) in done { + for r in recs { + record_opt(g.man.as_deref_mut(), r); + } + cals.push(cal); + arts.push(art); + } } // The run-level cal.json, as an ungrouped run writes it. From e8a2d951ec62c13e6f3e1f305a2f9a70a82400dc Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 10:22:33 +0200 Subject: [PATCH 03/35] fix(groups): cut the band plan at the crossings, not by a resetting accumulator `plan` closed a band once its running total reached the target share and then reset the accumulator. Every band therefore overshoots a little, the overshoot compounds along the m/z axis, and the windows left at the end fall into the last band. Measured on the 8-12-mer immunopeptidomics library, whose first 14 isolation windows select nothing because the library starts at m/z 326: 63 bands asked for came back as 36, ranging from 3.23M to 37.89M precursors against a 3.2M target, one of them holding 28 of the 114 windows. That band alone drove a 233 GB extract peak with only six bands in flight, which is what made the experiment unschedulable. Cuts are now taken where the running total crosses each k/n share of the whole, so they cannot compound: each is measured against the same running total rather than a reset one. A regression test covers the shape that broke it, a leading run of empty windows followed by equal ones, and asserts every band asked for comes back, within 3x of each other, with every window placed exactly once. Co-Authored-By: Claude Fable 5.1 --- rust/mumdia/crates/mumdia/src/groups.rs | 85 ++++++++++++++++++++----- 1 file changed, 69 insertions(+), 16 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/groups.rs b/rust/mumdia/crates/mumdia/src/groups.rs index 4cd5ba64..4b1b86f2 100644 --- a/rust/mumdia/crates/mumdia/src/groups.rs +++ b/rust/mumdia/crates/mumdia/src/groups.rs @@ -114,24 +114,41 @@ pub fn plan(windows: &[(f64, f64)], stats: &[RowGroupStats], n: usize) -> Result let n = n.min(windows.len()).max(1); let target = per_window.iter().sum::() / n as f64; - // Greedy cut: close a group once it holds at least the target and enough windows remain - // for the groups still to make. - let mut groups: Vec> = Vec::new(); - let mut cur: Vec = Vec::new(); - let mut acc = 0.0; - for (i, w) in per_window.iter().enumerate() { - cur.push(i); - acc += w; - let remaining_windows = windows.len() - i - 1; - let remaining_groups = n - groups.len() - 1; - if groups.len() + 1 < n && acc >= target && remaining_windows >= remaining_groups { - groups.push(std::mem::take(&mut cur)); - acc = 0.0; + // Cut where the running total crosses each k/n share of the whole. A greedy "close the + // group once it reaches the target" walk looks equivalent and is not: it closes a group + // on the window that first reaches the target, so every group overshoots a little, the + // overshoot compounds, and the windows left at the end are swallowed by the last group. + // Measured on the 8-12-mer library, whose first 14 windows select nothing because the + // library starts at m/z 326: 63 groups asked for came back as 36, one of them holding + // 28 windows and 37.9M precursors against a 3.2M target. Crossing points cannot + // compound, because each is taken against the running total rather than a reset + // accumulator. + let cum: Vec = per_window + .iter() + .scan(0.0, |acc, w| { + *acc += w; + Some(*acc) + }) + .collect(); + let mut cuts: Vec = vec![0]; + for k in 1..n { + let want = target * k as f64; + // First window whose running total reaches this share, and never a cut that would + // leave fewer windows than groups still to place. + let at = cum.partition_point(|&c| c < want) + 1; + let at = at + .min(windows.len() - (n - k)) + .max(*cuts.last().expect("seeded with 0") + 1); + if at > *cuts.last().expect("seeded with 0") && at < windows.len() { + cuts.push(at); } } - if !cur.is_empty() { - groups.push(cur); - } + cuts.push(windows.len()); + let groups: Vec> = cuts + .windows(2) + .filter(|c| c[1] > c[0]) + .map(|c| (c[0]..c[1]).collect()) + .collect(); // Bands from the groups, then merge any band that selects nothing into a neighbour. let mut bands: Vec = groups @@ -434,6 +451,42 @@ mod tests { assert!(plan(&windows, &no_stats, 2).is_err()); } + /// The cut must not compound: a leading run of windows that select nothing, and then + /// windows of equal weight, has to come back as the number of bands asked for, evenly + /// filled. The greedy walk this replaced returned 36 bands for 63 here, one of them + /// holding a third of the library. + #[test] + fn empty_leading_windows_do_not_swallow_the_plan() { + // 14 windows below the library's first precursor, then 100 windows over it. + let mut windows: Vec<(f64, f64)> = (0..14) + .map(|i| (100.0 + i as f64, 101.0 + i as f64)) + .collect(); + windows.extend((0..100).map(|i| (400.0 + i as f64, 401.0 + i as f64))); + let stats: Vec = (0..100) + .map(|i| RowGroupStats { + rows: 1_000_000, + min: Some(400.0 + i as f64), + max: Some(401.0 + i as f64), + }) + .collect(); + let p = plan(&windows, &stats, 25).expect("plans"); + assert_eq!(p.bands.len(), 25, "every band asked for"); + let est: Vec = p.bands.iter().map(|b| b.est_precursors).collect(); + let (lo, hi) = ( + est.iter().cloned().fold(f64::MAX, f64::min), + est.iter().cloned().fold(0.0, f64::max), + ); + assert!( + hi <= 3.0 * lo.max(1.0), + "bands within 3x of each other, got {lo} to {hi}" + ); + assert_eq!( + p.bands.iter().map(|b| b.windows.len()).sum::(), + windows.len(), + "every window placed exactly once" + ); + } + #[test] fn run_level_cal_keeps_the_ungrouped_keys_and_summarises_the_bands() { let dir = std::env::temp_dir().join(format!("mumdia_groups_cal_{}", std::process::id())); From 509f953af40b0a0c8be38864a0991695a15d5f03 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 11:59:07 +0200 Subject: [PATCH 04/35] perf(groups): write less, three ways A grouped run of the 8-12-mer library writes about 380 GB per file and the disk does 188 MB/s, so on a seven-file experiment the wall clock is partly the disk. Three changes, each measurable on its own: 1. A band's stages emit library-wide candidate ids. `extract` loads one band as a library whose rows are numbered 0..n and used to write those local numbers, so every downstream table only meant something next to the band's offset and pooling had to rewrite every id. Outputs now carry `local + Library::global_offset`, so pooling is a concatenation and a band's table can be read on its own. This is also what makes 2 and 3 simple. 2. The pooled `features.parquet` is not written. Nothing reads it: the competed table carries the feature columns, and the run-level copy was 55 GB per run. The per-band views of the pooled seed go with it, because features can now read the pooled seed directly, its ids matching the band's tables. 3. `MUMDIA_PARQUET_COMPRESSION=zstd` writes every artifact with zstd instead of snappy (`uncompressed` is also accepted). Both are read transparently whatever wrote them. Snappy stays the default, because it is what released artifacts use and what the sidecars' pyarrow reads without configuration, and because the setting changes every artifact's bytes: two runs compared by content hash must agree on it. On the CI fixture zstd is 3-11% smaller, which is too small a table to predict the float-heavy chromatograms of a real run. Unchanged output: the CI smoke passes with its TSV hashes intact, and a grouped fixture run reproduces the previous grouped run's `peptides.tsv` exactly, with and without zstd. Co-Authored-By: Claude Fable 5.1 --- rust/mumdia/Cargo.toml | 2 +- rust/mumdia/crates/mumdia-io/src/table.rs | 24 ++++- .../crates/mumdia/src/stages/extract.rs | 26 ++++- rust/mumdia/crates/mumdia/src/stages/pool.rs | 98 ++++++++----------- .../crates/mumdia/src/stages/run_groups.rs | 38 +++---- .../crates/mumdia/src/stages/seed_pool.rs | 44 +-------- 6 files changed, 98 insertions(+), 134 deletions(-) diff --git a/rust/mumdia/Cargo.toml b/rust/mumdia/Cargo.toml index 09712632..13f718c6 100644 --- a/rust/mumdia/Cargo.toml +++ b/rust/mumdia/Cargo.toml @@ -34,7 +34,7 @@ rayon = "1" clap = { version = "4", features = ["derive"] } arrow = { version = "59", features = ["prettyprint"] } # Pure-Rust build: drop C zlib-ng (needs cmake); SNAPPY codec is pure Rust. -parquet = { version = "59", default-features = false, features = ["arrow", "snap"] } +parquet = { version = "59", default-features = false, features = ["arrow", "snap", "zstd"] } # Pure-Rust zlib backend (miniz_oxide) instead of zlib-ng-compat (needs cmake). mzdata = { version = "0.66", default-features = false, features = ["mzml", "miniz_oxide"] } # Per-thread-arena allocator. The extraction accumulation is allocation-heavy and only diff --git a/rust/mumdia/crates/mumdia-io/src/table.rs b/rust/mumdia/crates/mumdia-io/src/table.rs index 11467c59..4a5e0ae7 100644 --- a/rust/mumdia/crates/mumdia-io/src/table.rs +++ b/rust/mumdia/crates/mumdia-io/src/table.rs @@ -21,6 +21,24 @@ use parquet::basic::{Compression, LogicalType}; use parquet::file::properties::WriterProperties; use parquet::file::statistics::Statistics; +/// The codec every artifact is written with. Snappy by default, which is what released +/// artifacts use and what the sidecars' pyarrow reads without configuration; +/// `MUMDIA_PARQUET_COMPRESSION=zstd` writes zstd instead, which is much smaller on the +/// float-heavy chromatogram and feature tables and therefore that much less to write on a +/// run whose wall clock is disk-bound. Both are read transparently, whatever wrote them. +/// It changes every artifact's bytes, so two runs compared by content hash must agree on it. +fn codec() -> Compression { + match std::env::var("MUMDIA_PARQUET_COMPRESSION") + .unwrap_or_default() + .to_ascii_lowercase() + .as_str() + { + "zstd" => Compression::ZSTD(Default::default()), + "uncompressed" | "none" => Compression::UNCOMPRESSED, + _ => Compression::SNAPPY, + } +} + /// One named, typed column for writing. pub enum Col { I64(String, Vec), @@ -203,7 +221,7 @@ fn cols_to_batch(path: &str, cols: Vec) -> Result<(Arc, RecordBatch } fn snappy_props(row_group_rows: Option) -> WriterProperties { - let mut b = WriterProperties::builder().set_compression(Compression::SNAPPY); + let mut b = WriterProperties::builder().set_compression(codec()); if let Some(n) = row_group_rows { b = b.set_max_row_group_row_count(Some(n.max(1))); } @@ -490,9 +508,7 @@ pub fn write_batches(path: &str, schema: Arc, batches: &[RecordBatch]) - let target = AtomicPath::new(path)?; let file = std::fs::File::create(target.tmp()) .with_context(|| format!("creating {}", target.tmp().display()))?; - let props = WriterProperties::builder() - .set_compression(Compression::SNAPPY) - .build(); + let props = WriterProperties::builder().set_compression(codec()).build(); let mut writer = ArrowWriter::try_new(file, schema, Some(props))?; let mut n = 0u64; for b in batches { diff --git a/rust/mumdia/crates/mumdia/src/stages/extract.rs b/rust/mumdia/crates/mumdia/src/stages/extract.rs index 75d4d1f2..05b2f8c7 100644 --- a/rust/mumdia/crates/mumdia/src/stages/extract.rs +++ b/rust/mumdia/crates/mumdia/src/stages/extract.rs @@ -142,7 +142,14 @@ struct ChromChunk { } impl ChromChunk { - fn cols(self) -> Vec { + /// `offset` is the library row of this band's local id 0 (`Library::global_offset`), + /// so the table carries library-wide ids even when the stage searched one band. + fn cols(mut self, offset: u32) -> Vec { + if offset != 0 { + for c in &mut self.cid { + *c += offset; + } + } vec![ Col::U32("candidate_id".into(), self.cid), Col::Str("frag_name".into(), self.name), @@ -2630,6 +2637,9 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { let (mut pk_apex, mut pk_start, mut pk_end): (Vec, Vec, Vec) = (Vec::new(), Vec::new(), Vec::new()); let (mut pk_ev, mut pk_area): (Vec, Vec) = (Vec::new(), Vec::new()); + // The chromatogram writer runs on its own thread and cannot borrow the library, so the + // band's offset travels with the chunks. + let chrom_offset = lib.global_offset; let n_chrom = std::thread::scope(|sc| -> Result { let (tx, rx) = std::sync::mpsc::sync_channel::>(2); let writer = sc.spawn(move || -> Result { @@ -2722,7 +2732,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { chrom_bytes_max_chunk = chrom_bytes_max_chunk.max(chunk_bytes); // Hand the chunk's chromatogram rows to the writer thread. A send error means // the writer failed; its error surfaces at the join below. - if tx.send(ch.cols()).is_err() { + if tx.send(ch.cols(chrom_offset)).is_err() { return false; } } @@ -2778,7 +2788,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { emit_batch(&mut cand_hits); } // A final empty chunk fixes the schema when no candidate was accepted at all. - let _ = tx.send(ChromChunk::default().cols()); + let _ = tx.send(ChromChunk::default().cols(0)); drop(tx); let n = writer .join() @@ -2865,7 +2875,10 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { } let mut psms_cols = vec![ - Col::U32("candidate_id".into(), cid_c), + Col::U32( + "candidate_id".into(), + cid_c.iter().map(|c| c + lib.global_offset).collect(), + ), Col::I32("peak_rank".into(), peakrank_c), Col::F64("apex_rt".into(), apexrt_c), Col::OptF64("apex_im".into(), apexim_c), @@ -2919,7 +2932,10 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { let n_peaks = write_table( &pk_path, vec![ - Col::U32("candidate_id".into(), pk_cid), + Col::U32( + "candidate_id".into(), + pk_cid.iter().map(|c| c + lib.global_offset).collect(), + ), Col::I32("peak_rank".into(), pk_rank), Col::F64("apex_rt".into(), pk_apex), Col::F64("start_rt".into(), pk_start), diff --git a/rust/mumdia/crates/mumdia/src/stages/pool.rs b/rust/mumdia/crates/mumdia/src/stages/pool.rs index a30ea8e8..89d982d3 100644 --- a/rust/mumdia/crates/mumdia/src/stages/pool.rs +++ b/rust/mumdia/crates/mumdia/src/stages/pool.rs @@ -18,23 +18,20 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::Instant; -use anyhow::{anyhow, bail, Context, Result}; +use anyhow::{anyhow, bail, Result}; use arrow::array::{Array, BooleanArray, UInt32Array}; use arrow::compute::filter_record_batch; -use arrow::record_batch::RecordBatch; use mumdia_io::table::{BatchWriter, TableFile}; use tracing::info; const BATCH_ROWS: usize = 1 << 16; const ROW_GROUP_ROWS: usize = 1 << 17; -/// One group's artifacts and where its band starts in the library. +/// One group's artifacts. The ids inside are already library-wide. #[derive(Clone, Debug)] pub struct BandArtifacts { - pub offset: u32, pub psms: String, pub chromatograms: String, - pub features: String, pub competed: String, } @@ -42,7 +39,6 @@ pub struct PoolParams<'a> { pub bands: &'a [BandArtifacts], pub out_psms: &'a str, pub out_chromatograms: &'a str, - pub out_features: &'a str, pub out_competed: &'a str, } @@ -51,7 +47,6 @@ pub struct PoolParams<'a> { pub struct PoolStats { pub psms: u64, pub chromatograms: u64, - pub features: u64, pub competed: u64, /// Candidates that appeared in two bands and were kept from one. pub duplicates: u64, @@ -81,10 +76,7 @@ fn overlap_losers(bands: &[BandArtifacts]) -> Result<(Vec>, u64)> { *e = *s; } } - for (local, s) in per_band { - let global = local - .checked_add(band.offset) - .ok_or_else(|| anyhow!("candidate id overflow pooling {}", band.competed))?; + for (global, s) in per_band { match best.get(&global) { None => { best.insert(global, (b, s)); @@ -92,10 +84,10 @@ fn overlap_losers(bands: &[BandArtifacts]) -> Result<(Vec>, u64)> { Some(&(prev_b, prev_s)) => { duplicates += 1; if s > prev_s { - losers[prev_b].insert(global - bands[prev_b].offset); + losers[prev_b].insert(global); best.insert(global, (b, s)); } else { - losers[b].insert(local); + losers[b].insert(global); } } } @@ -105,11 +97,11 @@ fn overlap_losers(bands: &[BandArtifacts]) -> Result<(Vec>, u64)> { } /// Append every band's table into `out`, ids offset, losers dropped. Returns the row count. -fn pool_table(paths: impl Iterator)>, out: &str) -> Result { +fn pool_table(paths: impl Iterator)>, out: &str) -> Result { let mut writer: Option = None; let mut schema: Option> = None; let mut rows = 0u64; - for (path, offset, drop) in paths { + for (path, drop) in paths { let t = TableFile::open(&path)?; let reader = t.batches(None, BATCH_ROWS)?; let this_schema = reader.schema(); @@ -137,29 +129,23 @@ fn pool_table(paths: impl Iterator)>, out: &st let w = writer.as_mut().expect("opened above"); for b in reader { let b = b?; - let cid = b - .column(cid_ix) - .as_any() - .downcast_ref::() - .ok_or_else(|| anyhow!("{path}: candidate_id is not u32"))?; - if cid.null_count() > 0 { - bail!("{path}: candidate_id has nulls"); - } - let mut global = Vec::with_capacity(cid.len()); - for &c in cid.values().iter() { - global.push( - c.checked_add(offset) - .ok_or_else(|| anyhow!("candidate id overflow pooling {path}"))?, - ); - } - let mut cols = b.columns().to_vec(); - cols[cid_ix] = Arc::new(UInt32Array::from(global)); - let mut batch = RecordBatch::try_new(this_schema.clone(), cols) - .with_context(|| format!("rebuilding a batch of {path}"))?; - if !drop.is_empty() { - let keep: Vec = cid.values().iter().map(|c| !drop.contains(c)).collect(); - batch = filter_record_batch(&batch, &BooleanArray::from(keep))?; - } + let keep: Option> = { + let cid = b + .column(cid_ix) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow!("{path}: candidate_id is not u32"))?; + if cid.null_count() > 0 { + bail!("{path}: candidate_id has nulls"); + } + (!drop.is_empty()).then(|| cid.values().iter().map(|c| !drop.contains(c)).collect()) + }; + // The band stages already wrote library-wide ids (`Library::global_offset` is + // added on the way out of extract), so pooling appends rather than rewrites. + let batch = match keep { + Some(k) => filter_record_batch(&b, &BooleanArray::from(k))?, + None => b, + }; if batch.num_rows() > 0 { w.write(&batch)?; rows += batch.num_rows() as u64; @@ -185,13 +171,12 @@ pub fn run(p: PoolParams) -> Result { p.bands .iter() .zip(losers.iter()) - .map(move |(b, l)| (pick(b).clone(), b.offset, l.clone())) + .map(move |(b, l)| (pick(b).clone(), l.clone())) .collect::>() }; let stats = PoolStats { psms: pool_table(with(|b| &b.psms).into_iter(), p.out_psms)?, chromatograms: pool_table(with(|b| &b.chromatograms).into_iter(), p.out_chromatograms)?, - features: pool_table(with(|b| &b.features).into_iter(), p.out_features)?, competed: pool_table(with(|b| &b.competed).into_iter(), p.out_competed)?, duplicates, }; @@ -199,7 +184,6 @@ pub fn run(p: PoolParams) -> Result { groups = p.bands.len(), psms = stats.psms, chromatograms = stats.chromatograms, - features = stats.features, competed = stats.competed, duplicates, elapsed_ms = t0.elapsed().as_millis() as u64, @@ -216,7 +200,7 @@ mod tests { fn band( dir: &std::path::Path, tag: &str, - offset: u32, + marker: f64, cids: &[u32], scores: &[f64], ) -> BandArtifacts { @@ -230,9 +214,7 @@ mod tests { Col::U32("candidate_id".into(), cids.to_vec()), Col::F64( "x".into(), - cids.iter() - .map(|c| *c as f64 + offset as f64 * 0.001) - .collect(), + cids.iter().map(|c| *c as f64 + marker).collect(), ), ]; if extra { @@ -242,10 +224,8 @@ mod tests { p }; BandArtifacts { - offset, psms: mk("psms", false), chromatograms: mk("chrom", false), - features: mk("feat", false), competed: mk("comp", true), } } @@ -254,35 +234,41 @@ mod tests { fn pooled_tables_carry_global_ids_and_one_row_per_overlap_candidate() { let dir = std::env::temp_dir().join(format!("mumdia_pool_{}", std::process::id())); std::fs::create_dir_all(&dir).unwrap(); - // Band 0: global 0..4; band 1: offset 3, local 0..3 = global 3..6. Globals 3 and 4 - // are in both; band 0 wins 3 (higher score), band 1 wins 4. - let b0 = band(&dir, "b0", 0, &[0, 1, 2, 3, 4], &[9.0, 9.0, 9.0, 7.0, 2.0]); - let b1 = band(&dir, "b1", 3, &[0, 1, 2, 3], &[5.0, 6.0, 9.0, 9.0]); + // The bands already carry library-wide ids: band 0 holds 0..4, band 1 holds 3..6, + // so 3 and 4 are in both. Band 0 wins 3 (higher score), band 1 wins 4. The `x` + // marker says which band a row came from. + let b0 = band( + &dir, + "b0", + 0.0, + &[0, 1, 2, 3, 4], + &[9.0, 9.0, 9.0, 7.0, 2.0], + ); + let b1 = band(&dir, "b1", 0.003, &[3, 4, 5, 6], &[5.0, 6.0, 9.0, 9.0]); let out = |n: &str| { dir.join(format!("out_{n}.parquet")) .to_str() .unwrap() .to_string() }; - let (op, oc, of, ok) = (out("psms"), out("chrom"), out("feat"), out("comp")); + let (op, oc, ok) = (out("psms"), out("chrom"), out("comp")); let stats = run(PoolParams { bands: &[b0, b1], out_psms: &op, out_chromatograms: &oc, - out_features: &of, out_competed: &ok, }) .unwrap(); assert_eq!(stats.duplicates, 2); assert_eq!(stats.competed, 7); - for path in [&op, &oc, &of, &ok] { + for path in [&op, &oc, &ok] { let t = TableFile::open(path).unwrap(); let cid = t.u32("candidate_id").unwrap(); assert_eq!(cid, vec![0, 1, 2, 3, 4, 5, 6], "{path}"); let x = t.f64("x").unwrap(); - // Global 3 came from band 0 (x = 3.000), global 4 from band 1 (x = 1.003). + // Id 3 came from band 0 (x = 3.000), id 4 from band 1 (x = 4.003). assert!( - (x[3] - 3.0).abs() < 1e-9 && (x[4] - 1.003).abs() < 1e-9, + (x[3] - 3.0).abs() < 1e-9 && (x[4] - 4.003).abs() < 1e-9, "{path}: {x:?}" ); } diff --git a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs index a1f3b97b..91d47388 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs @@ -122,8 +122,6 @@ pub fn run(g: GroupRun) -> Result { n: usize, prec: String, seed: String, - /// The band's rows of the pooled seed (local ids, pooled q). - seed_view: String, } let mut bands: Vec = Vec::new(); for b in &plan.bands { @@ -165,7 +163,6 @@ pub fn run(g: GroupRun) -> Result { n, prec, seed, - seed_view: gd(b.index, "seed_psms_pooled.parquet"), }); } if bands.is_empty() { @@ -180,7 +177,6 @@ pub fn run(g: GroupRun) -> Result { path: b.seed.clone(), offset: b.offset, rows: b.n as u32, - view: b.seed_view.clone(), }) .collect(); let masscals: Vec = bands @@ -373,9 +369,11 @@ pub fn run(g: GroupRun) -> Result { let nf = features::run(features::FeaturesParams { psms: &psms, chromatograms: &chrom, - // Corroboration and the confident elution boundary key on the seed by local - // id; under global calibration "confident" means the pooled q. - seed: Some(if global { &b.seed_view } else { &b.seed }), + // Corroboration and the confident elution boundary key on the seed by + // candidate id, and the band's tables carry library-wide ids, so the pooled + // seed is the one that matches. Its q is the pooled one either way, which is + // what "confident" has to mean once the bands are scored together. + seed: Some(&pooled_seed), out: &feats, out_pin: &pin, cfg: &cfg.features, @@ -406,10 +404,8 @@ pub fn run(g: GroupRun) -> Result { ch, )?); arts.push(pool::BandArtifacts { - offset: b.offset, psms, chromatograms: chrom, - features: feats, competed, }); } @@ -422,7 +418,9 @@ pub fn run(g: GroupRun) -> Result { seed: pooled_seed, psms: d("psms_extracted.parquet"), chromatograms: d("chromatograms.parquet"), - features: d("features.parquet"), + // Not pooled: nothing reads a run-level features table (compete's output carries + // the feature columns), and on a real run it is 55 GB of writes per run. + features: String::new(), competed: d("psms_competed.parquet"), rt_model, }; @@ -431,7 +429,6 @@ pub fn run(g: GroupRun) -> Result { bands: &arts, out_psms: &out.psms, out_chromatograms: &out.chromatograms, - out_features: &out.features, out_competed: &out.competed, }) .context("pooling the window groups")?; @@ -448,12 +445,6 @@ pub fn run(g: GroupRun) -> Result { &out.chromatograms, stats.chromatograms, ), - ( - artifact::FEATURES.0, - artifact::FEATURES, - &out.features, - stats.features, - ), ( artifact::PSMS_COMPETED.0, artifact::PSMS_COMPETED, @@ -486,15 +477,12 @@ pub fn run(g: GroupRun) -> Result { } .write_for(path)?; } - // The feature and competed schema companions (`.schema.json`) describe the - // classifier's columns; every band wrote the same one, so the first band's is the pool's. - for (band, pooled) in [ - (&arts[0].features, &out.features), - (&arts[0].competed, &out.competed), - ] { - let src = format!("{band}.schema.json"); + // The competed table's schema companion (`
.schema.json`) names the classifier's + // columns; every band wrote the same one, so the first band's is the pool's. + { + let src = format!("{}.schema.json", arts[0].competed); if std::path::Path::new(&src).exists() { - std::fs::copy(&src, format!("{pooled}.schema.json")) + std::fs::copy(&src, format!("{}.schema.json", out.competed)) .with_context(|| format!("copying {src} beside the pooled table"))?; } } diff --git a/rust/mumdia/crates/mumdia/src/stages/seed_pool.rs b/rust/mumdia/crates/mumdia/src/stages/seed_pool.rs index e7e5a3e8..bcdd675f 100644 --- a/rust/mumdia/crates/mumdia/src/stages/seed_pool.rs +++ b/rust/mumdia/crates/mumdia/src/stages/seed_pool.rs @@ -32,10 +32,6 @@ pub struct BandSeed { pub offset: u32, /// Rows of the band; the band's candidates are `offset..offset + rows`. pub rows: u32, - /// Where to write the band's view of the pooled seed: its own rows with band-local - /// ids and the pooled `spectrum_q`, for the band's stages that key on the seed by - /// candidate id (features' corroboration and elution boundary). - pub view: String, } pub struct SeedPoolParams<'a> { @@ -169,25 +165,7 @@ pub fn run(p: SeedPoolParams) -> Result { let pairs: Vec<(f64, bool)> = rows.iter().map(|r| (r.score, r.label == "decoy")).collect(); let q = target_decoy_q(&pairs); let n_out = rows.len(); - let n = write_rows(p.out, &rows, q.clone())?; - // Each band's rows again, with local ids and the pooled q. - for band in p.seeds { - let (view_rows, view_q): (Vec, Vec) = rows - .iter() - .zip(&q) - .filter(|(r, _)| r.cid >= band.offset && r.cid - band.offset < band.rows) - .map(|(r, q)| { - ( - Row { - cid: r.cid - band.offset, - ..r.clone() - }, - *q, - ) - }) - .unzip(); - write_rows(&band.view, &view_rows, view_q)?; - } + let n = write_rows(p.out, &rows, q)?; // Mass calibration: the bands' scalar offsets and learned tolerances combined by // calibrant count. The optional m/z grids are not combined (they would need the @@ -357,25 +335,17 @@ mod tests { .unwrap(); } let out = dir.join("pooled.parquet").to_str().unwrap().to_string(); - let view = |tag: &str| { - dir.join(format!("{tag}_view.parquet")) - .to_str() - .unwrap() - .to_string() - }; let n = run(SeedPoolParams { seeds: &[ BandSeed { path: a.clone(), offset: 0, rows: 4, - view: view("a"), }, BandSeed { path: b.clone(), offset: 3, rows: 3, - view: view("b"), }, ], masscals: &[format!("{a}.masscal.json"), format!("{b}.masscal.json")], @@ -408,16 +378,6 @@ mod tests { assert!((cal["frag_tol_ppm"].as_f64().unwrap() - 9.0).abs() < 1e-9); assert_eq!(cal["n_dev"].as_u64().unwrap(), 40); - // Band views: the band's own rows, local ids, the pooled q (band b's local 0 is - // the shared candidate, which the pool holds with band a's score). - let va = TableFile::open(&view("a")).unwrap(); - assert_eq!(va.u32("candidate_id").unwrap(), vec![0, 1, 2, 3]); - assert_eq!(va.f64("spectrum_q").unwrap(), q[0..4].to_vec()); - let vb = TableFile::open(&view("b")).unwrap(); - assert_eq!(vb.u32("candidate_id").unwrap(), vec![0, 1, 2]); - assert_eq!(vb.f64("spectrum_q").unwrap(), q[3..6].to_vec()); - assert_eq!(vb.f64("score").unwrap()[0], 15.0); - // Refresh: band b's re-predicted table says local 1 (global 4) now has iRT 50. let lib_b = dir.join("lib_b.parquet").to_str().unwrap().to_string(); write_table( @@ -453,13 +413,11 @@ mod tests { ) .unwrap(); let out = dir.join("pooled.parquet").to_str().unwrap().to_string(); - let view = dir.join("a_view.parquet").to_str().unwrap().to_string(); run(SeedPoolParams { seeds: &[BandSeed { path: a.clone(), offset: 0, rows: 2, - view, }], masscals: &[format!("{a}.masscal.json")], out: &out, From 906b24b38f22e439f1e10a248e25450a37cf3586 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 12:02:02 +0200 Subject: [PATCH 05/35] build: lock the zstd parquet feature Co-Authored-By: Claude Fable 5.1 --- rust/mumdia/Cargo.lock | 66 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 65 insertions(+), 1 deletion(-) diff --git a/rust/mumdia/Cargo.lock b/rust/mumdia/Cargo.lock index 93086f2d..4f16d6e7 100644 --- a/rust/mumdia/Cargo.lock +++ b/rust/mumdia/Cargo.lock @@ -427,6 +427,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -739,10 +741,21 @@ checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", "libc", - "r-efi", + "r-efi 5.3.0", "wasip2", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + [[package]] name = "gimli" version = "0.32.3" @@ -842,6 +855,16 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + [[package]] name = "js-sys" version = "0.3.103" @@ -1257,6 +1280,7 @@ dependencies = [ "seq-macro", "snap", "twox-hash", + "zstd", ] [[package]] @@ -1265,6 +1289,12 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + [[package]] name = "proc-macro2" version = "1.0.106" @@ -1299,6 +1329,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rayon" version = "1.12.0" @@ -1853,3 +1889,31 @@ name = "zmij" version = "1.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] From aecaf124e717168154702fc5454548ad760b9c34 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 14:57:31 +0200 Subject: [PATCH 06/35] perf(extract): size the probing fan-out per band, and do not sort hits that are sorted Two costs that only appear when several bands of a run are searched at once. The candidate-range fan-out took `rayon::current_num_threads()`, which is the whole pool, and every band in flight computed it independently: 24 bands each fanning out to twice the pool gave thousands of live narrowed-bin caches of 1-2 MB, allocated in lockstep, for pure scratch. `ExtractParams::sibling_bands` now says how many bands share the pool, and the fan-out is this band's share of it. The grouped orchestrator passes `groups.parallel`; every other caller passes 1, which is the previous behaviour exactly. Probing walks a window's scans in ascending index, so a candidate's hits are already rt-ascending in the ordinary case, and the sort then allocated scratch as large as the vector itself. Equal-rt hits collapse to a per-(rt, fragment) maximum afterwards whichever order they arrive in, so skipping a sort that would not move anything is exact. Those vectors are in the 64 KB to 1 MB class that dominates the process's memory mappings (measured on a live run: 92,030 mappings of 64-256 KB and 36,368 of 256 KB-1 MB out of 129,393 at 180 GB resident), so halving their churn matters beyond the copy itself. Output unchanged: the CI smoke passes with its hashes intact, and a grouped fixture run reproduces the previous run's `psms_extracted`, `chromatograms`, `psms_competed` and `peptides.tsv` byte for byte. Co-Authored-By: Claude Fable 5.1 --- rust/mumdia/crates/mumdia/src/main.rs | 1 + .../crates/mumdia/src/stages/extract.rs | 22 ++++++++++++++++--- rust/mumdia/crates/mumdia/src/stages/run.rs | 1 + .../mumdia/src/stages/run_experiment.rs | 1 + .../crates/mumdia/src/stages/run_groups.rs | 1 + rust/mumdia/crates/mumdia/tests/pipeline.rs | 1 + 6 files changed, 24 insertions(+), 3 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/main.rs b/rust/mumdia/crates/mumdia/src/main.rs index 6837b127..e0650be3 100644 --- a/rust/mumdia/crates/mumdia/src/main.rs +++ b/rust/mumdia/crates/mumdia/src/main.rs @@ -1226,6 +1226,7 @@ fn real_main() -> Result<()> { let ch = mumdia_io::hash::blake3_str(&cfg.canonical_json()); stages::extract::run(stages::extract::ExtractParams { fragment_offset, + sibling_bands: 1, ms2: &ms2, library_precursors: &lib_precursors, library_fragments: &lib_fragments, diff --git a/rust/mumdia/crates/mumdia/src/stages/extract.rs b/rust/mumdia/crates/mumdia/src/stages/extract.rs index 05b2f8c7..73f93e98 100644 --- a/rust/mumdia/crates/mumdia/src/stages/extract.rs +++ b/rust/mumdia/crates/mumdia/src/stages/extract.rs @@ -107,6 +107,10 @@ pub struct ExtractParams<'a> { /// are still read, and scans of other windows find no candidate. `None` is the ordinary /// whole-library extract. pub fragment_offset: Option, + /// How many bands of the same run are being searched beside this one + /// (`groups.parallel`). The probing fan-out is this band's share of the thread pool, + /// not the whole pool, because every band in flight computes it independently. + pub sibling_bands: usize, } /// One observed hit: scan RT, candidate-local fragment index, observed intensity @@ -753,6 +757,7 @@ fn window_groups(idx: &FragIndex, scans: &[Ms2Scan]) -> Vec { #[allow(clippy::too_many_arguments)] fn accumulate_groups( idx: &FragIndex, + sibling_bands: usize, groups: &[WinGroup], scans: &[Ms2Scan], rt_lo: &[f64], @@ -780,7 +785,11 @@ fn accumulate_groups( // Splitting the candidate axis rather than the scan axis is also where the work is: a // peak's cost is dominated by walking its posting list, which the narrowed sub-range // divides, while the peak loop itself is repeated per task. - let threads = rayon::current_num_threads().max(1); + // `current_num_threads()` is the whole pool, and under `groups.parallel` every band in + // flight computes this independently: 24 bands each fanning out to twice the pool gave + // thousands of live narrowed-bin caches (1-2 MB each) allocated in lockstep. Divide by + // the bands beside this one so the fan-out describes this band's share. + let threads = (rayon::current_num_threads() / sibling_bands.max(1)).max(1); let tasks_per_window = (threads * 2).div_ceil(groups.len().max(1)).max(1); let probe = |_gi: usize, g: &WinGroup| -> Vec<(u32, Vec)> { let ids = &g.scans; @@ -2017,7 +2026,13 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { } // Group hits into scan groups by RT (dedupe same fragment in a scan by max). - hits.sort_by(|a, b| a.rt.total_cmp(&b.rt)); + // Probing walks a window's scans in ascending index, so the hits of a candidate are + // already rt-ascending in the ordinary case; the sort then allocates scratch as large + // as the vector itself for nothing. Equal-rt hits collapse to `max` per (rt, frag) + // below either way, so skipping a sort that would not move anything is exact. + if !hits.windows(2).all(|w| w[0].rt <= w[1].rt) { + hits.sort_by(|a, b| a.rt.total_cmp(&b.rt)); + } // scan groups: Vec<(rt, BTreeMap)>. A BTreeMap keeps the // per-scan fragment order fixed so the f32 apex sum is deterministic. let mut groups: Vec<(f64, BTreeMap)> = Vec::new(); @@ -2752,6 +2767,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { accumulate_groups( fidx.as_ref() .expect("streamed path implies a fragment index"), + p.sibling_bands, &groups[gi..upto], &scans, &rt_lo, @@ -3237,7 +3253,7 @@ mod accumulate_tests { let mut acc: HashMap> = HashMap::new(); pool.install(|| { accumulate_groups( - &idx, &groups, &sc, &rt_lo, &rt_hi, &mass_off, &cfg, None, &mut acc, + &idx, 1, &groups, &sc, &rt_lo, &rt_hi, &mass_off, &cfg, None, &mut acc, ) }); acc diff --git a/rust/mumdia/crates/mumdia/src/stages/run.rs b/rust/mumdia/crates/mumdia/src/stages/run.rs index bd190eb7..771164ab 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run.rs @@ -543,6 +543,7 @@ pub fn run(p: RunParams) -> Result<()> { info!(stage = %"extract", "run: stage start"); let (npsm, nchr) = extract::run(extract::ExtractParams { fragment_offset: None, + sibling_bands: 1, ms2: &co.ms2, library_precursors: &lib_p, library_fragments: &lib_f, diff --git a/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs b/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs index 1e7f5aed..a5b5f32a 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs @@ -316,6 +316,7 @@ fn process_run( let chrom = d("chromatograms.parquet"); extract::run(extract::ExtractParams { fragment_offset: None, + sibling_bands: 1, ms2: &co.ms2, library_precursors: &lib_p, library_fragments: lib_f, diff --git a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs index 4038b254..92b0aff8 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs @@ -433,6 +433,7 @@ pub fn run(mut g: GroupRun) -> Result { cfg: &cfg.extract, config_hash: ch, fragment_offset: Some(b.offset), + sibling_bands: par, })?; recs.push(record_artifact( &format!("{}[g{:02}]", artifact::PSMS_EXTRACTED.0, b.index), diff --git a/rust/mumdia/crates/mumdia/tests/pipeline.rs b/rust/mumdia/crates/mumdia/tests/pipeline.rs index 921d3918..5daebd2e 100644 --- a/rust/mumdia/crates/mumdia/tests/pipeline.rs +++ b/rust/mumdia/crates/mumdia/tests/pipeline.rs @@ -170,6 +170,7 @@ fn run_extract(prec: &str, frag: &str, ms2: &str, win: &str, tag: &str) -> (Stri let chrom = tmp(&format!("chrom_{tag}.parquet")); stages::extract::run(stages::extract::ExtractParams { fragment_offset: None, + sibling_bands: 1, ms2, library_precursors: prec, library_fragments: frag, From 9cb65be742352573f9c6a7fc1d51c4659613a280 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:17:30 +0200 Subject: [PATCH 07/35] perf(index,seed): read the label as a bit, stop holding dead library columns The engine dies at ~290 GB resident with ~1.7 TB free because it is out of per-process memory MAPPINGS (1,048,576), not out of memory. Five changes in the library loader and the seed, all of which reduce live blocks or live bytes, plus one bug that has been printing a warning about a problem that does not exist. 1. `label` is read as one boolean per row instead of one `String` per row. The column is two-valued and its only consumer is `Candidate::is_decoy`, so a `Vec` was 203M heap blocks and about 6.5 GB on the full 203M-precursor library to produce one bit each. `read_is_decoy` streams the column in 64k-row batches; the validity rule and the message stay in `fdr::validate_labels`, which now sees only the offending value, the one string ever allocated. 2. `candidate_id` is verified the same way and never held. The ids are a precondition, never data -- every use is the row position -- and the column was 4 B per precursor (812 MB) live across the whole fragment load and index build. 3. `Library::prec_mz` is the MOVED `precursor_mz` column rather than a second copy of it: same values, same order, row c is candidate c. 8 B per candidate, 1.6 GB and one large block off the load's peak. (The other two copies of this value, `Candidate::precursor_mz` and `FragIndex::prec_mz`, are read from extract.rs and cloned in fragindex.rs, so they are outside this change's file set and stay.) 4. `Library::release_fragment_payload` frees `frag_int`, `frag_name_id` and the name dictionary, and search-seed calls it once its fragindex exists. The seed reads neither predicted intensity nor fragment name: the hyperscore is count plus observed intensity, and the mass recalibration needs only `frag_mz`, which is kept and is now reached through `cand_frag_mz`. 6 B per library fragment, held for the whole search, roughly 7 GB and two large mappings on a 1.2G-fragment library. `cand_frags` panics after a release rather than handing back a short slice; the bucketed path keeps the payload. 5. The seed's output columns are reserved once at the known row count, the `label` column is built FROM the `is_decoy` boolean at write time instead of the boolean being rebuilt by re-scanning the column, and the per-scan `touched()` copy (up to one candidate window wide, per scan) is a borrow. BUG, unrelated to memory: the seed's window-width survey mapped every window group of the run, including the ones outside the loaded band. An out-of-band group resolves to an empty candidate range and contributed a width of 1, so with 63 bands the median was 1 on every band -- the scratch was sized to a single slot and grown on first use -- and the `widest > 8 * median` guard fired on all 63, telling the user to check for a zero-width or missing isolation window in a run that has none. `window_width_stats` now excludes groups with `hi <= lo`, the same test the worker already uses to skip a group, so both the sizing and the warning describe this search. Output equality: the boolean is exactly `label[c] == "decoy"` and an unknown label is still refused; the id check is the same comparison against the same file row; `prec_mz` is the same array; the released columns are unread on that path; the seed's `label` column is the same text and `is_dec` the same booleans; the scratch size never affects a result (SeedScratch grows on demand and new slots are unstamped), only the allocation. Tests added for each claim, including before/after equality of `cand_frag_mz` across a release and the old/new identity of the label round trip. Not done, blocked by the file set: a seed-only `FragIndex` build omitting `post_int`/`post_frag` (fragindex.rs), and moving `Candidate`'s two owned Strings to a flat buffer plus a dictionary id (extract.rs, fragindex.rs). Co-Authored-By: Claude Fable 5.1 --- rust/mumdia/crates/mumdia/src/index.rs | 316 ++++++++++++++++-- .../crates/mumdia/src/stages/search_seed.rs | 180 +++++++--- 2 files changed, 438 insertions(+), 58 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/index.rs b/rust/mumdia/crates/mumdia/src/index.rs index 5c3d42a8..91296210 100644 --- a/rust/mumdia/crates/mumdia/src/index.rs +++ b/rust/mumdia/crates/mumdia/src/index.rs @@ -23,6 +23,10 @@ use rayon::prelude::*; /// Fragment rows per decoded batch while streaming the fragment table (a few MB). const FRAG_BATCH_ROWS: usize = 1 << 16; +/// Precursor rows per decoded batch for the two columns that are validated in a streaming +/// pass and never kept (`label`, `candidate_id`). +const PREC_BATCH_ROWS: usize = 1 << 16; + #[derive(Clone, Debug)] pub struct Candidate { pub candidate_id: u32, @@ -104,6 +108,69 @@ fn require_finite_f32(v: &[f32], column: &str, path: &str) -> Result<()> { Ok(()) } +/// Read the `label` column as one boolean per row, streaming, without materialising a +/// `String` per row. +/// +/// The column is two-valued and its only consumer is [`Candidate::is_decoy`], so reading it +/// as a `Vec` cost one heap block per precursor -- 203M blocks and about 6.5 GB on +/// the full library -- to produce one bit each. The validity rule is unchanged and stays in +/// one place: an unexpected value is handed to [`crate::fdr::validate_labels`], which is +/// also the only value ever turned into a `String`. +fn read_is_decoy(pt: &TableFile, path: &str) -> Result> { + let mut out: Vec = Vec::with_capacity(pt.nrows); + for b in pt.batches(Some(&["label"]), PREC_BATCH_ROWS)? { + let b = b?; + let a = b + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("column 'label' is not utf8"))?; + // `value()` ignores the validity bitmap, so a NULL would read as "" and become a + // target. Same contract as the fragment columns below. + require_no_nulls(a, "label", path, out.len())?; + for k in 0..a.len() { + match a.value(k) { + "decoy" => out.push(true), + "target" => out.push(false), + other => crate::fdr::validate_labels(&[other.to_string()])?, + } + } + } + Ok(out) +} + +/// Check, streaming, that `candidate_id` is the contiguous row-aligned range +/// `offset..offset + ncand`. +/// +/// The ids are a precondition, never data: `index.rs` uses the row position everywhere and +/// the column exists only to be verified. Holding it cost 4 bytes per precursor (812 MB on +/// the full library) for the whole of the fragment load and index build. +fn check_candidate_ids(pt: &TableFile, path: &str, offset: usize, ncand: usize) -> Result<()> { + let mut row = 0usize; + for b in pt.batches(Some(&["candidate_id"]), PREC_BATCH_ROWS)? { + let b = b?; + let a = b + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow::anyhow!("column 'candidate_id' is not u32"))?; + require_no_nulls(a, "candidate_id", path, row)?; + for &candidate_id in a.values().iter().take(ncand.saturating_sub(row)) { + if candidate_id as usize != row + offset { + anyhow::bail!( + "library precursor row {} has candidate_id {} but candidate_id must \ + be the contiguous range 0..n in row order; reindex the library \ + (e.g. via the decoy-builder scripts)", + row + offset, + candidate_id + ); + } + row += 1; + } + } + Ok(()) +} + impl Library { /// Load a library and build the bucketed inverted index. /// @@ -301,16 +368,16 @@ impl Library { None => TableFile::open(precursors)?, Some((first, n)) => TableFile::open_rows(precursors, first, n)?, }; - let cid = pt.u32("candidate_id")?; let pfid = pt.u32("peptidoform_id")?; let baseid = pt.u32("base_peptide_id")?; let mut pform = pt.str("peptidoform")?; let charge = pt.i32("charge")?; let pmz = pt.f64("precursor_mz")?; let irt = pt.f32("predicted_irt")?; - let label = pt.str("label")?; let mut protein = pt.str("protein")?; - crate::fdr::validate_labels(&label)?; + // `label` is validated and reduced to a bit in one streaming pass rather than being + // held as a `Vec`; `candidate_id` is checked the same way further down. + let is_decoy = read_is_decoy(&pt, precursors)?; // A Parquet NULL decodes to NaN (mumdia-io `Table::f64`/`f32`), and NaN is // accepted rather than rejected by every downstream guard that should catch it: // the ascending-m/z check below is `<`, the extract RT-window guards are @@ -380,24 +447,14 @@ impl Library { } let ncand = pt.nrows; - drop(pt); // Precondition: candidate_id is the contiguous, row-aligned range 0..ncand // (the library + decoy builders guarantee this). An external library that // violates it would misgroup fragments or panic on the index below, so // check explicitly and fail with a clear error instead. // For a range load the same invariant holds against the file row: local id c is // file row c + offset, so the slice's ids must be exactly offset..offset + ncand. - for (c, &candidate_id) in cid.iter().enumerate().take(ncand) { - if candidate_id as usize != c + offset { - anyhow::bail!( - "library precursor row {} has candidate_id {} but candidate_id must \ - be the contiguous range 0..n in row order; reindex the library \ - (e.g. via the decoy-builder scripts)", - c + offset, - candidate_id - ); - } - } + check_candidate_ids(&pt, precursors, offset, ncand)?; + drop(pt); // Fragment table: two streaming passes over the four columns the library needs (the // artifact also carries `ion_type`, `ordinal`, `frag_charge` and `cardinality`, which @@ -589,8 +646,12 @@ impl Library { } drop(name_lookup); + // `prec_mz` IS the decoded `precursor_mz` column: row c is candidate c, in the same + // order, so the second array was a copy of the first. Move it instead of pushing a + // duplicate, which removes 8 bytes per candidate (1.6 GB on the full library) and + // one large heap block from the load's peak. + let prec_mz = pmz; let mut cands = Vec::with_capacity(ncand); - let mut prec_mz = Vec::with_capacity(ncand); for c in 0..ncand { let start = frag_offsets[c] as usize; let n = frag_offsets[c + 1] as usize - start; @@ -602,14 +663,13 @@ impl Library { // Move the strings out of the column Vecs instead of cloning them. peptidoform: std::mem::take(&mut pform[c]), charge: charge[c], - precursor_mz: pmz[c], + precursor_mz: prec_mz[c], predicted_irt: irt[c], - is_decoy: label[c] == "decoy", + is_decoy: is_decoy[c], protein: std::mem::take(&mut protein[c]), frag_start: start, n_frag: n, }); - prec_mz.push(pmz[c]); } drop(frag_offsets); @@ -733,6 +793,11 @@ impl Library { /// Per-candidate fragment m/z, predicted intensity, and INTERNED name ids (resolve /// with [`Library::frag_name_str`]). pub fn cand_frags(&self, cid: u32) -> (&[f32], &[f32], &[u16]) { + assert!( + !self.fragment_payload_released(), + "the library's fragment payload was released (Library::release_fragment_payload); \ + only cand_frag_mz is available after that" + ); let c = &self.cands[cid as usize]; let s = c.frag_start; let e = s + c.n_frag; @@ -743,6 +808,38 @@ impl Library { ) } + /// Per-candidate fragment m/z alone. The only fragment column that survives + /// [`Library::release_fragment_payload`], so this is what a stage that has already built + /// its index must use. + pub fn cand_frag_mz(&self, cid: u32) -> &[f32] { + let c = &self.cands[cid as usize]; + &self.frag_mz[c.frag_start..c.frag_start + c.n_frag] + } + + /// Free the fragment columns that only an index build reads: `frag_int`, + /// `frag_name_id` and the name dictionary. `frag_mz` is kept, because the seed's mass + /// recalibration still matches every library fragment against the scan. + /// + /// Predicted intensity and fragment name are copied into the index (or, for the + /// bucketed arrays, into `idx_int`) at build time and never read from the library + /// afterwards by a stage that has an index, so holding them is 6 bytes per fragment of + /// dead weight for the whole search -- about 7 GB on a 1.2G-fragment library, and two + /// of the large mappings the process is rationed on. Call it only after the index the + /// stage will actually probe exists; [`Library::cand_frags`] panics afterwards rather + /// than returning a short or stale slice. + pub fn release_fragment_payload(&mut self) { + self.frag_int = Vec::new(); + self.frag_name_id = Vec::new(); + self.frag_name_dict = Vec::new(); + } + + /// Whether [`Library::release_fragment_payload`] has run. Derived rather than stored: + /// the three fragment arrays are parallel, so a `frag_int` shorter than `frag_mz` can + /// only mean the payload was dropped. + pub fn fragment_payload_released(&self) -> bool { + self.frag_int.len() != self.frag_mz.len() + } + /// Local fragment index (0..n_frag) of the candidate whose stored m/z is /// closest to `frag_mz_f32` (returned by `page_search`). pub fn local_frag_index(&self, cid: u32, frag_mz_f32: f32) -> usize { @@ -1284,6 +1381,187 @@ mod tests { (p, f) } + /// A two-candidate library whose `candidate_id` and `label` columns are the caller's, + /// for the two streaming-validated columns below. + fn library_ids_labels( + dir: &std::path::Path, + ids: [u32; 2], + labels: [&str; 2], + ) -> (String, String) { + std::fs::create_dir_all(dir).unwrap(); + let p = dir.join("prec.parquet").to_str().unwrap().to_string(); + let f = dir.join("frag.parquet").to_str().unwrap().to_string(); + write_table( + &p, + vec![ + Col::U32("candidate_id".into(), ids.to_vec()), + Col::U32("peptidoform_id".into(), vec![0, 1]), + Col::U32("base_peptide_id".into(), vec![0, 1]), + Col::Str( + "peptidoform".into(), + vec!["PEPTIDEK".into(), "SAMPLER".into()], + ), + Col::I32("charge".into(), vec![2, 2]), + Col::F64("precursor_mz".into(), vec![400.0, 500.0]), + Col::F32("predicted_irt".into(), vec![10.0, 20.0]), + Col::Str( + "label".into(), + vec![labels[0].to_string(), labels[1].to_string()], + ), + Col::Str("protein".into(), vec!["P1".into(), "P2".into()]), + ], + ) + .unwrap(); + write_table( + &f, + vec![ + Col::U32("candidate_id".into(), vec![0, 1]), + Col::F64("mz".into(), vec![200.1, 250.5]), + Col::F32("predicted_intensity".into(), vec![1.0, 0.9]), + Col::Str("name".into(), vec!["b2".into(), "y3".into()]), + ], + ) + .unwrap(); + (p, f) + } + + /// The `label` column is read as a boolean per row instead of a `String` per row (203M + /// heap blocks and ~6.5 GB on the full library, to set one bit each). The bit must be + /// exactly what `label[c] == "decoy"` produced, and an unexpected value must still be + /// refused by `fdr::validate_labels`'s rule and message. + #[test] + fn the_label_column_becomes_the_boolean_it_always_meant() { + let dir = unique_dir("label_bool"); + let (p, f) = library_ids_labels(&dir, [0, 1], ["target", "decoy"]); + let lib = Library::load_with(&p, &f, 8, false).unwrap(); + assert_eq!( + lib.cands.iter().map(|c| c.is_decoy).collect::>(), + vec![false, true] + ); + // Reversed, so a constant-false or index-shifted read cannot pass both cases. + let (p2, f2) = library_ids_labels(&unique_dir("label_bool2"), [0, 1], ["decoy", "target"]); + let lib2 = Library::load_with(&p2, &f2, 8, false).unwrap(); + assert_eq!( + lib2.cands.iter().map(|c| c.is_decoy).collect::>(), + vec![true, false] + ); + std::fs::remove_dir_all(&dir).ok(); + std::fs::remove_dir_all(unique_dir("label_bool2")).ok(); + } + + #[test] + fn an_unknown_label_is_still_refused_by_name() { + let dir = unique_dir("label_bad"); + let (p, f) = library_ids_labels(&dir, [0, 1], ["target", "REVERSE"]); + let err = match Library::load_with(&p, &f, 8, false) { + Ok(_) => panic!("an unknown label must be rejected at load"), + Err(e) => e.to_string(), + }; + assert!(err.contains("unknown PSM label"), "{err}"); + assert!(err.contains("REVERSE"), "{err}"); + std::fs::remove_dir_all(&dir).ok(); + } + + /// `candidate_id` is verified in a streaming pass and never held (4 B per precursor, + /// 812 MB on the full library, live for the whole fragment load). The check itself is + /// unchanged: row-aligned or refused, naming the row. + #[test] + fn a_non_contiguous_candidate_id_is_rejected_with_the_row() { + let dir = unique_dir("cid_gap"); + let (p, f) = library_ids_labels(&dir, [0, 2], ["target", "decoy"]); + let err = match Library::load_with(&p, &f, 8, false) { + Ok(_) => panic!("a non-contiguous candidate_id must be rejected at load"), + Err(e) => e.to_string(), + }; + assert!(err.contains("contiguous range"), "{err}"); + assert!( + err.contains("row 1"), + "should name the offending row: {err}" + ); + std::fs::remove_dir_all(&dir).ok(); + } + + /// A range load checks the ids against the FILE row, not the local one, and the + /// streaming check must keep that: local id c is file row c + offset. + #[test] + fn the_streamed_id_check_uses_the_file_row_for_a_range_load() { + let dir = std::env::temp_dir().join(format!("mumdia_index_cidspan_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let (p, f) = build_six_lib(&dir, "cidspan", &[0, 1, 2, 3, 4, 5]); + // Rows 1..4 carry file ids 1, 2, 3; the load accepts them against offset 1. + let part = Library::load_range_with(&p, &f, 440.0, 530.0, 8, false).unwrap(); + assert_eq!(part.n_candidates(), 3); + assert_eq!( + part.cands + .iter() + .map(|c| c.candidate_id) + .collect::>(), + vec![0, 1, 2] + ); + } + + /// `prec_mz` is now the moved `precursor_mz` column rather than a second copy of it. + /// It must still be exactly the per-candidate value, since `candidate_range` and the + /// fragment index both index it by candidate id. + #[test] + fn prec_mz_is_the_per_candidate_precursor_mz() { + let dir = std::env::temp_dir().join(format!("mumdia_index_precmz_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let (p, f) = build_six_lib(&dir, "precmz", &[0, 1, 2, 3, 4, 5]); + let lib = Library::load_with(&p, &f, 8, false).unwrap(); + assert_eq!(lib.prec_mz.len(), lib.n_candidates()); + for c in 0..lib.n_candidates() { + assert_eq!(lib.prec_mz[c], lib.cands[c].precursor_mz); + } + assert_eq!(lib.prec_mz, vec![400.0, 450.0, 500.0, 520.0, 600.0, 650.0]); + // The same holds for a band, whose `prec_mz` is the band's slice of the column. + let part = Library::load_range_with(&p, &f, 440.0, 530.0, 8, false).unwrap(); + assert_eq!(part.prec_mz, vec![450.0, 500.0, 520.0]); + for c in 0..part.n_candidates() { + assert_eq!(part.prec_mz[c], part.cands[c].precursor_mz); + } + } + + /// Releasing the fragment payload frees `frag_int`/`frag_name_id` and keeps `frag_mz` + /// byte for byte, because the seed's mass recalibration still walks every library + /// fragment m/z after its index exists. + #[test] + fn releasing_the_fragment_payload_keeps_every_fragment_mz() { + let dir = std::env::temp_dir().join(format!("mumdia_index_release_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let (p, f) = build_six_lib(&dir, "release", &[0, 1, 2, 3, 4, 5]); + let mut lib = Library::load_with(&p, &f, 8, false).unwrap(); + assert!(!lib.fragment_payload_released()); + let before: Vec> = (0..lib.n_candidates()) + .map(|c| lib.cand_frag_mz(c as u32).to_vec()) + .collect(); + // Before the release the two accessors agree, which is what makes the swap safe. + for c in 0..lib.n_candidates() { + assert_eq!(lib.cand_frags(c as u32).0, lib.cand_frag_mz(c as u32)); + } + lib.release_fragment_payload(); + assert!(lib.fragment_payload_released()); + assert!(lib.frag_int.is_empty() && lib.frag_name_id.is_empty()); + let after: Vec> = (0..lib.n_candidates()) + .map(|c| lib.cand_frag_mz(c as u32).to_vec()) + .collect(); + assert_eq!(before, after); + } + + /// The released payload must fail loudly rather than hand back a short or stale slice: + /// `cand_frags` is what extract reads predicted intensities through. + #[test] + #[should_panic(expected = "fragment payload was released")] + fn cand_frags_after_a_release_panics_instead_of_lying() { + let dir = + std::env::temp_dir().join(format!("mumdia_index_relpanic_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let (p, f) = build_six_lib(&dir, "relpanic", &[0, 1, 2, 3, 4, 5]); + let mut lib = Library::load_with(&p, &f, 8, false).unwrap(); + lib.release_fragment_payload(); + let _ = lib.cand_frags(0); + } + #[test] fn an_empty_protein_loads_as_the_unassigned_group() { // DIA-NN leaves the protein empty for peptides it did not map (the iRT-kit diff --git a/rust/mumdia/crates/mumdia/src/stages/search_seed.rs b/rust/mumdia/crates/mumdia/src/stages/search_seed.rs index 7edc0c24..925c9b71 100644 --- a/rust/mumdia/crates/mumdia/src/stages/search_seed.rs +++ b/rust/mumdia/crates/mumdia/src/stages/search_seed.rs @@ -62,7 +62,7 @@ pub fn run(p: SearchSeedParams) -> Result { )?; // See extract: the bucketed index is dead weight on the fragindex path. let build_bucketed = !matches!(p.cfg.matcher, MatcherKind::Fragindex); - let lib = match p.fragment_offset { + let mut lib = match p.fragment_offset { None => Library::load_with( p.library_precursors, p.library_fragments, @@ -87,6 +87,15 @@ pub fn run(p: SearchSeedParams) -> Result { // fragindex backend, built once at the seed's fragment tolerance when selected. let fidx = matches!(p.cfg.matcher, MatcherKind::Fragindex) .then(|| FragIndex::build(&lib, p.cfg.fragment_tol_ppm)); + if fidx.is_some() { + // The index owns its own copy of every posting, and the seed reads neither the + // predicted intensity nor the fragment name from either side: the hyperscore is + // count + observed intensity, and the mass recalibration below needs only + // `frag_mz`. So the library's `frag_int` and `frag_name_id` are dead from here on + // -- 6 bytes per library fragment, held for the whole search. The bucketed path + // keeps them, because `page_search` serves `idx_int` out of arrays built from them. + lib.release_fragment_payload(); + } // Best-per-candidate PSM. The fragindex path parallelizes across isolation-window // groups (each scan belongs to exactly one window, so groups are independent) and @@ -148,17 +157,23 @@ pub fn run(p: SearchSeedParams) -> Result { .collect(); let q = target_decoy_q(&sd); - let (mut cid_c, mut pform_c, mut charge_c, mut mz_c) = - (Vec::new(), Vec::new(), Vec::new(), Vec::new()); - let (mut base_c, mut prot_c, mut label_c) = (Vec::new(), Vec::new(), Vec::new()); - let (mut score_c, mut q_c, mut rt_c, mut matched_c, mut scan_c, mut irt_c) = ( - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - ); + // One allocation per column instead of a growth sequence per column: the row count is + // known, and these are the tables that make the seed's peak. + let n_rows = rows.len(); + let mut cid_c = Vec::with_capacity(n_rows); + let mut pform_c = Vec::with_capacity(n_rows); + let mut charge_c = Vec::with_capacity(n_rows); + let mut mz_c = Vec::with_capacity(n_rows); + let mut base_c = Vec::with_capacity(n_rows); + let mut prot_c = Vec::with_capacity(n_rows); + // The label is the boolean, not a string that a later pass has to parse back into one. + let mut is_dec: Vec = Vec::with_capacity(n_rows); + let mut score_c = Vec::with_capacity(n_rows); + let mut q_c = Vec::with_capacity(n_rows); + let mut rt_c = Vec::with_capacity(n_rows); + let mut matched_c = Vec::with_capacity(n_rows); + let mut scan_c = Vec::with_capacity(n_rows); + let mut irt_c = Vec::with_capacity(n_rows); for (i, (cid, b)) in rows.iter().enumerate() { let c = &lib.cands[*cid as usize]; cid_c.push(*cid); @@ -167,7 +182,7 @@ pub fn run(p: SearchSeedParams) -> Result { mz_c.push(c.precursor_mz); base_c.push(c.base_peptide_id); prot_c.push(c.protein.clone()); - label_c.push(if c.is_decoy { "decoy" } else { "target" }.to_string()); + is_dec.push(c.is_decoy); score_c.push(b.score); q_c.push(q[i]); rt_c.push(b.rt); @@ -176,7 +191,6 @@ pub fn run(p: SearchSeedParams) -> Result { irt_c.push(c.predicted_irt); } - let is_dec: Vec = label_c.iter().map(|l| l == "decoy").collect(); let n_at_1pct = count_targets_at_q(&q_c, &is_dec, p.cfg.fdr_seed); // Per-run fragment mass recalibration + learned tolerance @@ -197,7 +211,9 @@ pub fn run(p: SearchSeedParams) -> Result { continue; } if let Some(scan) = scan_by_index.get(&b.scan_index) { - let (mzs, _, _) = lib.cand_frags(*cid); + // m/z only: the predicted intensity and the fragment name are not part of the + // mass calibration, and on the fragindex path they no longer exist. + let mzs = lib.cand_frag_mz(*cid); // The calibrant collection window has to be at least as wide as the search // tolerance, or the deviation percentile that SETS the learned tolerance is // truncated by the collection window itself. A literal 50 ppm silently did @@ -328,7 +344,11 @@ pub fn run(p: SearchSeedParams) -> Result { Col::F64("precursor_mz".into(), mz_c), Col::U32("base_peptide_id".into(), base_c), Col::Str("protein".into(), prot_c), - Col::Str("label".into(), label_c), + // Materialised here and nowhere else: the artifact's column is text, but the + // engine never needed the strings for anything but the write itself. (A `Col` + // variant over a two-valued dictionary would remove these too; that is a + // mumdia-io change.) + Col::Str("label".into(), label_column(&is_dec)), Col::F64("score".into(), score_c), Col::F64("spectrum_q".into(), q_c), Col::F64("observed_rt".into(), rt_c), @@ -391,6 +411,55 @@ fn select_peaks(scan: &Ms2Scan, top_n: usize) -> Vec { } } +/// The seed artifact's `label` column, built from the boolean the library already carries. +/// +/// It used to be the other way round: the column was assembled as `Vec` while the +/// rows were walked and the `is_decoy` vector was then rebuilt by comparing each of those +/// strings to `"decoy"`, which is a full extra pass over the rows to recover a bit that +/// `Candidate::is_decoy` had all along. +fn label_column(is_decoy: &[bool]) -> Vec { + is_decoy + .iter() + .map(|&d| if d { "decoy" } else { "target" }.to_string()) + .collect() +} + +/// Median and maximum candidate-window width over the window groups THIS search actually +/// serves, from each group's `(lo, hi)` candidate range. +/// +/// Size each worker's scratch to the MEDIAN candidate window, not the widest one and +/// certainly not the whole library. The scratch is indexed window-relative, and it grows on +/// demand, so an underestimate costs a few reallocations while an overestimate costs +/// 16 B x width per rayon worker up front. +/// +/// The max is the wrong statistic because one window can be the whole library: +/// `convert.rs` maps BOTH a zero-width reported isolation window and a missing precursor to +/// the full range (0, 1e6), which `candidate_range` resolves to every candidate. So a single +/// malformed or all-ion scan in an otherwise 50-window run sized every worker's scratch to +/// the library -- 877 MB per worker on the profiled 54.8M-candidate library, about 28 GB of +/// commit charge on 32 cores, for arrays a worker inside one narrow window never touches. +/// +/// Groups that select nothing (`hi <= lo`) are EXCLUDED, which is the same test the worker +/// uses to skip a group. Under an isolation-window-group search the library is one m/z band +/// while the run's windows are all of them, so with 63 bands 62 of every 63 groups select +/// nothing and contribute a width of 1: the median was then 1 on every band, the scratch was +/// sized to a single slot and grown on first use, and `widest > 8 * median` fired on every +/// band, printing a warning about a zero-width or missing isolation window once per band for +/// a run that has no such scan. Restricted to the in-band groups, both the sizing and the +/// warning describe this search. +/// +/// Widths are collected in group order and then sorted, so the value is deterministic. +fn window_width_stats(ranges: &[(u32, u32)]) -> (usize, Option) { + let mut widths: Vec = ranges + .iter() + .filter(|(lo, hi)| hi > lo) + .map(|(lo, hi)| hi.saturating_sub(*lo) as usize + 1) + .collect(); + widths.sort_unstable(); + let median = widths.get(widths.len() / 2).copied().unwrap_or(1).max(1); + (median, widths.last().copied()) +} + /// fragindex seed over isolation-window groups, in parallel. Each scan belongs to /// exactly one isolation window, so grouping scans by window gives independent /// parallel units (the candidate axis overlaps between adjacent windows, handled by @@ -417,37 +486,19 @@ fn seed_fragindex_windows( .push(si); } let group_vec: Vec> = groups.into_values().collect(); - // Size each worker's scratch to the MEDIAN candidate window, not the widest one and - // certainly not the whole library. The scratch is indexed window-relative, and it - // grows on demand, so an underestimate costs a few reallocations while an - // overestimate costs 16 B x width per rayon worker up front. - // - // The max is the wrong statistic because one window can be the whole library: - // `convert.rs` maps BOTH a zero-width reported isolation window and a missing - // precursor to the full range (0, 1e6), which `candidate_range` resolves to every - // candidate. So a single malformed or all-ion scan in an otherwise 50-window run - // sized every worker's scratch to the library -- 877 MB per worker on the profiled - // 54.8M-candidate library, about 28 GB of commit charge on 32 cores, for arrays a - // worker inside one narrow window never touches. The only mitigation was `--threads`, - // which the user had to know to reach for. - // - // The median is robust to that outlier and needs no configuration. Widths are - // collected in group order, so the value is deterministic. - let mut widths: Vec = group_vec + let ranges: Vec<(u32, u32)> = group_vec .iter() .filter_map(|ids| ids.first()) .map(|&si| { let w = &scans[si].window; - let (lo, hi) = idx.candidate_range(w.lower_mz, w.upper_mz); - hi.saturating_sub(lo) as usize + 1 + idx.candidate_range(w.lower_mz, w.upper_mz) }) .collect(); - widths.sort_unstable(); - let median_window_width = widths.get(widths.len() / 2).copied().unwrap_or(1).max(1); + let (median_window_width, widest) = window_width_stats(&ranges); // Say so when the two disagree by a lot: a window spanning most of the library is // either a genuine all-ion acquisition or a malformed scan, and both are worth // knowing about before wondering why extraction is slow. - if let Some(&widest) = widths.last() { + if let Some(widest) = widest { if widest > 8 * median_window_width { warn!( median_window_candidates = median_window_width, @@ -483,7 +534,10 @@ fn seed_fragindex_windows( .map(|&pi| (scan.peaks[pi].mz, scan.peaks[pi].intensity)) .collect(); scratch.accumulate(idx, &peaks, lo, hi); - let touched: Vec = scratch.touched().to_vec(); + // Borrowed, not copied: `touched` can be as long as the candidate + // window, so the copy was one allocation of up to a window's width per + // scan. Same slice, same order, so the scored list is unchanged. + let touched = scratch.touched(); let mut scored: Vec<(u32, f64, u32)> = touched .iter() .filter(|&&cid| scratch.count(cid) as usize >= cfg.min_matched_peaks) @@ -544,6 +598,54 @@ fn hyperscore(matched: u32, sum_obs: f64) -> f64 { ln_factorial(matched) + (1.0 + sum_obs).ln() } +#[cfg(test)] +mod survey_tests { + use super::{label_column, window_width_stats}; + + #[test] + fn window_width_stats_ignores_the_groups_this_band_does_not_serve() { + // A grouped search loads ONE m/z band of the library and sees the run's whole + // window list, so every group outside the band resolves to an empty candidate + // range. With 63 bands that is 62 empty groups against one real one; counting the + // empty groups as width 1 put the median at 1 (scratch sized to a single slot, + // resized on first use) and made `widest > 8 * median` true on every band, so the + // zero-width-isolation-window warning fired 63 times per run for a run with no such + // scan. + let mut ranges = vec![(0u32, 0u32); 62]; + ranges.push((1000, 6000)); + let (median, widest) = window_width_stats(&ranges); + assert_eq!((median, widest), (5001, Some(5001))); + // ... which is what the guard needs in order not to fire. + assert!(widest.unwrap() <= 8 * median); + } + + #[test] + fn window_width_stats_is_the_median_of_the_served_groups_and_still_flags_an_outlier() { + // Widths 11, 11, 11, and one group covering the library: the median is still 11 and + // the outlier is still reported. Order-independent, and the empty groups mixed in + // change neither answer. + let ranges = [(0, 10), (100, 110), (0, 0), (50, 60), (0, 5_000_000)]; + let (median, widest) = window_width_stats(&ranges); + assert_eq!((median, widest), (11, Some(5_000_001))); + assert!(widest.unwrap() > 8 * median); + // No group selects anything: nothing to size from, nothing to warn about. + assert_eq!(window_width_stats(&[(7, 7), (0, 0)]), (1, None)); + assert_eq!(window_width_stats(&[]), (1, None)); + } + + #[test] + fn label_column_is_the_boolean_the_old_code_parsed_back_out_of_it() { + // The artifact text is unchanged, and the boolean the stage uses is now its source + // rather than its product: `is_decoy` -> column -> `l == "decoy"` is the identity. + let is_dec = [false, true, true, false]; + let col = label_column(&is_dec); + assert_eq!(col, vec!["target", "decoy", "decoy", "target"]); + let round_trip: Vec = col.iter().map(|l| l == "decoy").collect(); + assert_eq!(round_trip, is_dec); + assert!(label_column(&[]).is_empty()); + } +} + #[cfg(test)] mod peak_selection_tests { use super::select_peaks; From 8f5386c0a335cb549fdfd9593121d6bef9ed1983 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:21:33 +0200 Subject: [PATCH 08/35] perf(compete): stop allocating one heap block per PSM, and cap the competed row group Four terms, all of them live heap BLOCKS rather than bytes, which is what the engine actually dies on: a grouped search holds several bands in one process and hits the kernel's per-process mapping limit (1,048,576) at ~290 GB resident with 1.7 TB free (measured on a live run: 129,393 mappings at 180 GB, 92,030 of them 64-256 KB and 36,368 of them 256 KB-1 MB). The grouping was a `HashMap>` with one `Vec` per competition group. Under the shipped `group_by = peptidoform_charge` nearly every group is a singleton, so that is about one small heap block per PSM, plus a 48-byte-per-entry table at up to 87.5% load. It is now a `Vec<(key, row)>`, sorted once: 32 bytes per PSM in ONE allocation, so 100 MB and one mapping for the six-run Astral pool's 3.13 M PSMs. The resolver walks contiguous runs of equal keys instead of a sorted key list. It visits the same groups in the same order with members in the same ascending row order, and the winner rule was already order-independent with an index tiebreak, so the kept set and the removal pairs are unchanged; a test pins that against the previous HashMap resolver over all five modes on a population with ties, a NaN prelim and a signed-zero pair. `copy_kept_rows` built a take index and called `take` on all ~398 columns of every batch even when the batch had lost nothing, which under the default grouping is the normal case: the widest artifact of the run was copied through `take` to reproduce itself. An f64 column of a 16,384-row batch is exactly 128 KB, i.e. the 64-256 KB band above, 398 of them per batch and 191 batches for that pool. When as many rows survive as the batch holds they ARE the batch, in order, so the column is now passed through as an Arc clone. The synthesised all-zero `peak_rank` has no source array and is still built. The peptidoform column was read as a `Vec` for the whole table and, unless the competition audit is on, only used to number dense ids; the label column was read as `Vec` to be compared against two literals. Both are now streamed: the label as one byte per row, the peptidoform into first-appearance ids whose map holds one String per DISTINCT peptidoform rather than one per row. `base_peptide_id` is read only for the groupings that key on it and `candidate_id` only for the audit sidecar, which is the only consumer of either as a value; both columns are still required and type-checked by the pass-through copy. At 3.13 M rows that is ~6.2 M fewer String blocks and ~175 MB less. The competed table is written in 131,072-row groups, the cap the rescore handoff already uses and the one every other wide writer in the engine sets. This file took parquet's default 1,048,576, which at ~387 f64 feature columns is ~3.2 GB decoded per group: that is what the writer buffers before each flush, and rescore is the reader that decodes a whole group before slicing batches out of it. 131,072 rows is ~405 MB. Output equality: the competed parquet, its `.schema.json`, its report's row/kept/removed stats and the audit sidecar are byte-identical to the previous implementation over seven configurations (all three groupings, all four modes, audit on and off) on a 500-row crafted features table, and value- and order-identical on a 300,000-row one, where the cap binds and splits one row group into three. Only row-group boundaries move, so the competed table's `content_hash` changes for any run above 131,072 kept rows; nothing downstream reads that hash as an equality assertion. The one behavioural drift is error text, not values: a features table missing or mistyping `candidate_id` or `base_peptide_id` now fails in the pass-through copy's column check rather than in the typed getter. Co-Authored-By: Claude Fable 5.1 --- .../crates/mumdia/src/stages/compete.rs | 571 ++++++++++++++++-- 1 file changed, 506 insertions(+), 65 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/compete.rs b/rust/mumdia/crates/mumdia/src/stages/compete.rs index 3441ccf4..4762d88e 100644 --- a/rust/mumdia/crates/mumdia/src/stages/compete.rs +++ b/rust/mumdia/crates/mumdia/src/stages/compete.rs @@ -17,12 +17,19 @@ use mumdia_core::config::{CompeteConfig, CompeteGroupBy, CompetitionMode}; use mumdia_core::rejection::RejectionReason; use mumdia_core::schema::artifact; use mumdia_io::report::ArtifactReport; -use mumdia_io::table::{write_table, BatchWriter, Col, TableFile}; +use mumdia_io::table::{require_no_nulls, write_table, BatchWriter, Col, TableFile}; use serde_json::json; use tracing::{info, warn}; use crate::stages::features::FeatureSchema; +/// Competition group key: `(base-or-peptidoform id, label code, bucket, peak rank)`. +/// Fixed size, so a whole run's keys are one flat buffer rather than a String per PSM. +type GroupKey = (u32, u8, i64, i32); + +/// Rows per streamed batch of a single key column. +const KEY_BATCH_ROWS: usize = 1 << 16; + pub struct CompeteParams<'a> { pub features: &'a str, pub out: &'a str, @@ -43,9 +50,12 @@ pub fn run(p: CompeteParams) -> Result { // batch while the surviving rows are copied through to the output. let t = TableFile::open(p.features)?; let n = t.nrows; - let cid = t.u32("candidate_id")?; - let label = t.str("label")?; - let base = t.u32("base_peptide_id")?; + let audit = p.cfg.emit_competition_audit; + // The label is a one-byte code in the key and nowhere else, so it is read as codes. + // A `Vec` here was one small heap block per PSM for a value only ever compared + // against two literals; the Strings are read back only for the opt-in audit sidecar, + // which is the one consumer that prints them. + let label_code = label_codes(&t)?; let prelim = t.f64("prelim_score")?; // Top-K peak rank (#7). Part of the competition key so peaks of one candidate // compete only within their own rank (a lower-scoring peak of a candidate must @@ -59,11 +69,14 @@ pub fn run(p: CompeteParams) -> Result { let feat_names = &schema.feature_columns; let by_pform_charge = matches!(p.cfg.group_by, CompeteGroupBy::PeptidoformCharge); - // Only read what the grouping (and the optional audit) needs. - let pform: Vec = if by_pform_charge || p.cfg.emit_competition_audit { - t.str("peptidoform")? - } else { + // Only read what the grouping (and the optional audit) needs. `base_peptide_id` is + // dead weight under the default peptidoform-charge grouping, and `candidate_id` is + // read by the audit sidecar alone; both columns are still required (and type-checked) + // by the pass-through copy below, which writes every bookkeeping column. + let base: Vec = if by_pform_charge { Vec::new() + } else { + t.u32("base_peptide_id")? }; let apex_rt: Vec = if matches!(p.cfg.group_by, CompeteGroupBy::Apex) { t.f64("apex_rt")? @@ -81,15 +94,11 @@ pub fn run(p: CompeteParams) -> Result { }; // Dense peptidoform id by first appearance (deterministic) so the fixed-size // tuple key can separate modforms without allocating a String per PSM. Built - // only for the peptidoform-charge grouping; empty otherwise. + // only for the peptidoform-charge grouping; empty otherwise. Streamed, so the + // peptidoform column is never materialised: the map holds one String per DISTINCT + // peptidoform, not one per row. let pform_id: Vec = if by_pform_charge { - let mut ids = Vec::with_capacity(n); - let mut seen: HashMap<&str, u32> = HashMap::new(); - for peptidoform in pform.iter().take(n) { - let next = seen.len() as u32; - ids.push(*seen.entry(peptidoform.as_str()).or_insert(next)); - } - ids + dense_peptidoform_ids(&t)? } else { Vec::new() }; @@ -102,29 +111,33 @@ pub fn run(p: CompeteParams) -> Result { // Key is a fixed-size tuple (base-or-pform id, label_code, bucket) instead of a // freshly-allocated String per PSM. Precursor grouping uses a constant bucket // (0) so its equivalence classes are unchanged. - let mut groups: HashMap<(u32, u8, i64, i32), Vec> = HashMap::new(); + // + // The groups are a SORTED `(key, row)` array, not a `HashMap>`. Under + // the shipped `group_by = peptidoform_charge` nearly every group is a singleton, so + // the map was about one small heap block per PSM and it was this stage's dominant + // memory term; a grouped search that holds several bands at once dies on the kernel's + // per-process mapping limit long before it runs out of bytes. The resolver walks + // contiguous runs of equal keys instead, which visits the same groups in the same + // sorted-key order with members in the same ascending row order. + let mut entries: Vec<(GroupKey, usize)> = Vec::with_capacity(n); for i in 0..n { - let label_code = match label[i].as_str() { - "target" => 0u8, - "decoy" => 1u8, - _ => 2u8, - }; let pk = peak_rank[i]; let key = match p.cfg.group_by { - CompeteGroupBy::BasePeptide => (base[i], label_code, 0i64, pk), + CompeteGroupBy::BasePeptide => (base[i], label_code[i], 0i64, pk), CompeteGroupBy::Apex => { let bucket = (apex_rt[i] / p.cfg.apex_rt_tolerance_s).round() as i64; - (base[i], label_code, bucket, pk) + (base[i], label_code[i], bucket, pk) } CompeteGroupBy::PeptidoformCharge => { // pform_id separates modforms; charge in the bucket separates // charges -> one group per peptidoform+charge (precursor-level). let c = charge.as_ref().unwrap()[i].round() as i64; - (pform_id[i], label_code, c, pk) + (pform_id[i], label_code[i], c, pk) } }; - groups.entry(key).or_default().push(i); + entries.push((key, i)); } + entries.sort_unstable(); // Per-candidate unique-fragment evidence for the `unique_evidence` mode. Prefers // an explicit `unique_fragment_count` feature; otherwise approximates it as @@ -172,32 +185,44 @@ pub fn run(p: CompeteParams) -> Result { // Resolve each group under the configured competition mode. let (keep, removed) = resolve_competition( - &groups, + &entries, &prelim, p.cfg.mode, p.cfg.margin, p.cfg.unique_evidence_min_fragments, unique_ev.as_deref(), ); - drop(groups); + drop(entries); - let rows = copy_kept_rows(&t, p.out, feat_names, synth_peak_rank, &keep)?; + let rows = copy_kept_rows( + &t, + p.out, + feat_names, + synth_peak_rank, + &keep, + COMPETED_ROW_GROUP_ROWS, + )?; // Feature schema companion: unchanged feature list, so rescore validates the same schema. mumdia_io::json::write_json(&format!("{}.schema.json", p.out), &schema)?; // Competition audit sidecar (opt-in): one row per removed PSM with its winner. Lets a - // post-hoc analysis see what competition removed without re-running the stage. - if p.cfg.emit_competition_audit { + // post-hoc analysis see what competition removed without re-running the stage. This is + // the only consumer of the candidate_id / label / peptidoform columns as VALUES, so + // they are read here rather than held for the whole stage. + if audit { + let cid = t.u32("candidate_id")?; + let label = t.str("label")?; + let pform = t.str("peptidoform")?; let reason_of = |i: usize| { - if label[i] == "decoy" { + if label_code[i] == 1 { RejectionReason::OutcompetedByDecoy } else { RejectionReason::OutcompetedByTarget } }; - let audit = format!("{}.compete_audit.parquet", p.out); + let audit_path = format!("{}.compete_audit.parquet", p.out); write_table( - &audit, + &audit_path, vec![ Col::U32( "candidate_id".into(), @@ -266,6 +291,60 @@ pub fn run(p: CompeteParams) -> Result { Ok(rows) } +/// The label column as one code per row: `target` -> 0, `decoy` -> 1, anything else -> 2, +/// which is exactly the resolution the competition key applies. Streamed, so the peak is +/// one byte per row plus one decoded batch. Nulls are refused, as `TableFile::str` refuses +/// them for a required column. +fn label_codes(t: &TableFile) -> Result> { + let mut out: Vec = Vec::with_capacity(t.nrows); + t.for_each_batch(Some(&["label"]), KEY_BATCH_ROWS, |b| { + let a = b + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow!("compete: column 'label' is not utf8"))?; + require_no_nulls(a, "label", t.path(), out.len())?; + for k in 0..a.len() { + out.push(match a.value(k) { + "target" => 0u8, + "decoy" => 1u8, + _ => 2u8, + }); + } + Ok(()) + })?; + Ok(out) +} + +/// Dense peptidoform ids by first appearance, streamed. Same values as numbering a +/// materialised `Vec` of the column, without the String per row: the map holds one +/// key per DISTINCT peptidoform, and row order (hence the numbering) is the file's. +fn dense_peptidoform_ids(t: &TableFile) -> Result> { + let mut ids: Vec = Vec::with_capacity(t.nrows); + let mut seen: HashMap = HashMap::new(); + t.for_each_batch(Some(&["peptidoform"]), KEY_BATCH_ROWS, |b| { + let a = b + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow!("compete: column 'peptidoform' is not utf8"))?; + require_no_nulls(a, "peptidoform", t.path(), ids.len())?; + for k in 0..a.len() { + let s = a.value(k); + match seen.get(s) { + Some(&id) => ids.push(id), + None => { + let next = seen.len() as u32; + seen.insert(s.to_string(), next); + ids.push(next); + } + } + } + Ok(()) + })?; + Ok(ids) +} + /// The bookkeeping columns every competed table starts with, in order, with the types the /// features stage writes them in. One place, so the pass-through below cannot drift from /// the typed schema this stage used to re-declare column by column. @@ -286,6 +365,16 @@ const META_COLUMNS: [(&str, DataType); 11] = [ /// Input rows per streamed batch of the pass-through copy (~50 MB at ~400 f64 columns). const COPY_BATCH_ROWS: usize = 1 << 14; +/// Row-group cap of the competed table, the same cap the rescore handoff writes with. +/// +/// Rescore reads this file back batch by batch, and a parquet reader decodes a WHOLE row +/// group before it slices batches out of it, so on a wide table the row-group size is the +/// reader's working set; it is also what the writer buffers before each flush. At parquet's +/// default 1,048,576 rows and ~387 f64 feature columns that is ~3.2 GB decoded per group. +/// 131,072 rows is ~400 MB. Row-group boundaries are the only thing this moves; values and +/// row order are unchanged. +const COMPETED_ROW_GROUP_ROWS: usize = 131_072; + /// Copy the surviving rows (`keep`, sorted ascending) of the features table into `out`, /// one input batch at a time, in exactly the column set and order the previous typed /// rewrite produced: the 11 bookkeeping columns, then the schema's feature columns, all @@ -297,6 +386,7 @@ fn copy_kept_rows( feat_names: &[String], synth_peak_rank: bool, keep: &[usize], + row_group_rows: usize, ) -> Result { let mut fields: Vec = Vec::with_capacity(META_COLUMNS.len() + feat_names.len()); // Source column per output field; None = synthesised zeros (a pre-v2 features table @@ -340,7 +430,7 @@ fn copy_kept_rows( .map(|s| in_schema.index_of(s).expect("validated above")) }) .collect(); - let mut w = BatchWriter::new(out, out_schema.clone())?; + let mut w = BatchWriter::with_row_group_rows(out, out_schema.clone(), row_group_rows)?; let (mut row0, mut kp) = (0usize, 0usize); for b in reader { let b = b?; @@ -350,17 +440,30 @@ fn copy_kept_rows( kp += 1; } if kp > start { - let idx = UInt32Array::from( - keep[start..kp] - .iter() - .map(|&r| (r - row0) as u32) - .collect::>(), - ); + // `keep` is sorted and unique and every index in `keep[start..kp]` lies in this + // batch, so when as many rows survive as the batch holds they ARE the batch, in + // order, and `take` would be an identity copy. Under the shipped grouping that + // is the normal case, and this batch is the widest artifact of the run: all + // ~398 columns were being rebuilt to reproduce themselves. Pass the column + // through (an Arc clone) instead. The synthesised peak_rank column has no + // source array to pass through and is still built. + let n_kept = kp - start; + let idx = if n_kept == b.num_rows() { + None + } else { + Some(UInt32Array::from( + keep[start..kp] + .iter() + .map(|&r| (r - row0) as u32) + .collect::>(), + )) + }; let mut arrays: Vec = Vec::with_capacity(src_idx.len()); for (si, f) in src_idx.iter().zip(out_schema.fields()) { - arrays.push(match si { - Some(i) => densify(take(b.column(*i).as_ref(), &idx, None)?, f)?, - None => Arc::new(Int32Array::from(vec![0i32; idx.len()])), + arrays.push(match (si, &idx) { + (Some(i), None) => densify(b.column(*i).clone(), f)?, + (Some(i), Some(ix)) => densify(take(b.column(*i).as_ref(), ix, None)?, f)?, + (None, _) => Arc::new(Int32Array::from(vec![0i32; n_kept])), }); } w.write(&RecordBatch::try_new(out_schema.clone(), arrays)?)?; @@ -465,42 +568,51 @@ fn col_f64(t: &TableFile, name: &str) -> Option> { /// Resolve within-group competition per [`CompetitionMode`]. Returns the /// sorted-unique kept row indices and the `(loser, winner)` removal pairs. -/// Deterministic: groups are visited in sorted key order; the winner is the -/// highest `prelim` (ties broken by smallest index). +/// +/// `entries` is `(key, row)` sorted ascending, so a group is a contiguous run of equal +/// keys and its members are in ascending row order. That is the same visiting order the +/// previous `HashMap>` produced (keys sorted, members pushed in row order), +/// so the winner, the kept set and the removal pairs are unchanged. Deterministic: the +/// winner is the highest `prelim` (ties broken by smallest index). fn resolve_competition( - groups: &HashMap<(u32, u8, i64, i32), Vec>, + entries: &[(GroupKey, usize)], prelim: &[f64], mode: CompetitionMode, margin: f64, unique_min: usize, unique_ev: Option<&[f64]>, ) -> (Vec, Vec<(usize, usize)>) { - let mut group_keys: Vec<&(u32, u8, i64, i32)> = groups.keys().collect(); - group_keys.sort_unstable(); let mut keep: Vec = Vec::new(); let mut removed: Vec<(usize, usize)> = Vec::new(); - for gk in group_keys { - let members = &groups[gk]; - let win = *members + let mut g = 0usize; + while g < entries.len() { + let mut h = g + 1; + while h < entries.len() && entries[h].0 == entries[g].0 { + h += 1; + } + let members = &entries[g..h]; + g = h; + let win = members .iter() + .map(|&(_, i)| i) // `total_cmp`, not `partial_cmp(..).unwrap_or(Equal)`: this picks the // single row that survives competition, and treating every NaN // prelim_score as equal to every other score made that choice depend on // iteration order. `total_cmp` is a genuine total order, so the winner // is well defined even then, and the `.then(a.cmp(&b))` index tiebreak // keeps it deterministic. - .min_by(|&&a, &&b| prelim[b].total_cmp(&prelim[a]).then(a.cmp(&b))) - .unwrap(); + .min_by(|&a, &b| prelim[b].total_cmp(&prelim[a]).then(a.cmp(&b))) + .expect("a run of equal keys is never empty"); match mode { CompetitionMode::None | CompetitionMode::FeaturesOnly => { - keep.extend(members.iter().copied()); + keep.extend(members.iter().map(|&(_, i)| i)); } CompetitionMode::WinnerTakeAll => { keep.push(win); removed.extend( members .iter() - .copied() + .map(|&(_, i)| i) .filter(|&m| m != win) .map(|m| (m, win)), ); @@ -508,7 +620,7 @@ fn resolve_competition( CompetitionMode::UniqueEvidence => { keep.push(win); let thr = unique_min as f64; - for &m in members { + for &(_, m) in members { if m == win { continue; } @@ -521,7 +633,7 @@ fn resolve_competition( } CompetitionMode::MarginGated => { keep.push(win); - for &m in members { + for &(_, m) in members { if m == win { continue; } @@ -542,6 +654,7 @@ fn resolve_competition( #[cfg(test)] mod tests { use super::*; + use mumdia_io::table::Col as IoCol; #[test] fn writing_the_output_over_the_input_is_refused() { @@ -581,10 +694,14 @@ mod tests { let _ = std::fs::remove_dir_all(&dir); } - fn one_group(members: Vec) -> HashMap<(u32, u8, i64, i32), Vec> { - let mut g = HashMap::new(); - g.insert((0u32, 0u8, 0i64, 0i32), members); - g + /// One group of `members`, as the sorted `(key, row)` array the resolver takes. + fn one_group(members: Vec) -> Vec<(GroupKey, usize)> { + let mut e: Vec<(GroupKey, usize)> = members + .into_iter() + .map(|i| ((0u32, 0u8, 0i64, 0i32), i)) + .collect(); + e.sort_unstable(); + e } #[test] @@ -656,9 +773,13 @@ mod tests { #[test] fn winner_take_all_is_deterministic_across_groups() { - let mut g = HashMap::new(); - g.insert((0u32, 0u8, 0i64, 0i32), vec![0, 1]); - g.insert((1u32, 1u8, 0i64, 0i32), vec![2, 3]); + let mut g: Vec<(GroupKey, usize)> = vec![ + ((0u32, 0u8, 0i64, 0i32), 0), + ((0u32, 0u8, 0i64, 0i32), 1), + ((1u32, 1u8, 0i64, 0i32), 2), + ((1u32, 1u8, 0i64, 0i32), 3), + ]; + g.sort_unstable(); let prelim = [0.2, 0.8, 0.9, 0.3]; let (keep, _) = resolve_competition(&g, &prelim, CompetitionMode::WinnerTakeAll, 0.0, 2, None); @@ -673,4 +794,324 @@ mod tests { resolve_competition(&g, &prelim, CompetitionMode::WinnerTakeAll, 0.0, 2, None); assert_eq!(keep, vec![0]); } + + /// The grouping this stage used until 2026-09-22: one `Vec` per key in a `HashMap`, + /// keys visited in sorted order, members in ascending row order. Kept here as the + /// reference the sorted-run resolver has to reproduce exactly. + fn resolve_via_hashmap( + entries: &[(GroupKey, usize)], + prelim: &[f64], + mode: CompetitionMode, + margin: f64, + unique_min: usize, + unique_ev: Option<&[f64]>, + ) -> (Vec, Vec<(usize, usize)>) { + let mut by_row: Vec<(GroupKey, usize)> = entries.to_vec(); + by_row.sort_unstable_by_key(|&(_, i)| i); + let mut groups: HashMap> = HashMap::new(); + for (k, i) in by_row { + groups.entry(k).or_default().push(i); + } + let mut group_keys: Vec<&GroupKey> = groups.keys().collect(); + group_keys.sort_unstable(); + let mut keep: Vec = Vec::new(); + let mut removed: Vec<(usize, usize)> = Vec::new(); + for gk in group_keys { + let members = &groups[gk]; + let win = *members + .iter() + .min_by(|&&a, &&b| prelim[b].total_cmp(&prelim[a]).then(a.cmp(&b))) + .unwrap(); + match mode { + CompetitionMode::None | CompetitionMode::FeaturesOnly => { + keep.extend(members.iter().copied()); + } + CompetitionMode::WinnerTakeAll => { + keep.push(win); + removed.extend( + members + .iter() + .copied() + .filter(|&m| m != win) + .map(|m| (m, win)), + ); + } + CompetitionMode::UniqueEvidence => { + keep.push(win); + let thr = unique_min as f64; + for &m in members { + if m == win { + continue; + } + if unique_ev.map(|u| u[m] >= thr).unwrap_or(false) { + keep.push(m); + } else { + removed.push((m, win)); + } + } + } + CompetitionMode::MarginGated => { + keep.push(win); + for &m in members { + if m == win { + continue; + } + if prelim[win] - prelim[m] >= margin { + removed.push((m, win)); + } else { + keep.push(m); + } + } + } + } + } + keep.sort_unstable(); + keep.dedup(); + (keep, removed) + } + + #[test] + fn sorted_run_grouping_reproduces_the_hashmap_grouping() { + // A deterministic population with singletons, multi-member groups, ties, a NaN + // prelim and a signed-zero pair, over all five modes. + let n = 600usize; + let mut state = 0x2545_F491_4F6C_DD1Du64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + let mut entries: Vec<(GroupKey, usize)> = Vec::with_capacity(n); + let mut prelim: Vec = Vec::with_capacity(n); + let mut ev: Vec = Vec::with_capacity(n); + for i in 0..n { + let r = next(); + let key = ( + (r % 97) as u32, + ((r >> 8) % 3) as u8, + ((r >> 16) % 5) as i64, + ((r >> 24) % 3) as i32, + ); + entries.push((key, i)); + // Coarse quantisation so ties are common. + prelim.push(match i { + 7 => f64::NAN, + 11 => -0.0, + 12 => 0.0, + _ => ((r >> 32) % 11) as f64 / 10.0, + }); + ev.push(((r >> 40) % 5) as f64); + } + let reference = entries.clone(); + entries.sort_unstable(); + for mode in [ + CompetitionMode::None, + CompetitionMode::FeaturesOnly, + CompetitionMode::WinnerTakeAll, + CompetitionMode::MarginGated, + CompetitionMode::UniqueEvidence, + ] { + let got = resolve_competition(&entries, &prelim, mode, 0.2, 2, Some(&ev)); + let want = resolve_via_hashmap(&reference, &prelim, mode, 0.2, 2, Some(&ev)); + assert_eq!(got.0, want.0, "kept rows differ for {mode:?}"); + assert_eq!(got.1, want.1, "removal pairs differ for {mode:?}"); + } + } + + /// A minimal features table: the 11 bookkeeping columns plus two feature columns. + fn write_features(path: &str, n: usize) { + let f = |g: fn(usize) -> f64| (0..n).map(g).collect::>(); + mumdia_io::table::write_table( + path, + vec![ + IoCol::U32("candidate_id".into(), (0..n as u32).collect()), + IoCol::I32("peak_rank".into(), vec![0; n]), + IoCol::Str( + "label".into(), + (0..n) + .map(|i| if i % 3 == 0 { "decoy" } else { "target" }.to_string()) + .collect(), + ), + IoCol::U32( + "base_peptide_id".into(), + (0..n as u32).map(|i| i / 2).collect(), + ), + IoCol::Str( + "peptidoform".into(), + (0..n).map(|i| format!("PEP{}", i / 2)).collect(), + ), + IoCol::Str("protein".into(), (0..n).map(|i| format!("P{i}")).collect()), + IoCol::F64("apex_rt".into(), f(|i| i as f64)), + IoCol::F64("elution_lo".into(), f(|i| i as f64 - 1.0)), + IoCol::F64("elution_hi".into(), f(|i| i as f64 + 1.0)), + IoCol::F64("precursor_mz".into(), f(|i| 400.0 + i as f64)), + IoCol::F64("prelim_score".into(), f(|i| (i % 7) as f64 / 7.0)), + IoCol::F64("charge".into(), f(|i| 2.0 + (i % 2) as f64)), + IoCol::F64("n_matched_fragments".into(), f(|i| (i % 5) as f64)), + ], + ) + .unwrap(); + } + + fn tmp_dir(tag: &str) -> std::path::PathBuf { + let d = std::env::temp_dir().join(format!( + "mumdia_compete_{tag}_{}_{:?}", + std::process::id(), + std::thread::current().id() + )); + let _ = std::fs::remove_dir_all(&d); + std::fs::create_dir_all(&d).unwrap(); + d + } + + #[test] + fn passing_a_whole_batch_through_equals_taking_every_row_of_it() { + // The pass-through branch must be indistinguishable from the `take` it replaces. + // Row 9 is dropped in the second copy, so that batch goes through `take`; the rows + // both copies share have to agree column for column. + let dir = tmp_dir("passthrough"); + let src_path = dir.join("features.parquet"); + let src = src_path.to_str().unwrap(); + write_features(src, 10); + let t = TableFile::open(src).unwrap(); + let feats = vec!["charge".to_string(), "n_matched_fragments".to_string()]; + + let all_path = dir.join("all.parquet"); + let all = all_path.to_str().unwrap(); + let keep_all: Vec = (0..10).collect(); + assert_eq!( + copy_kept_rows(&t, all, &feats, false, &keep_all, 1 << 20).unwrap(), + 10 + ); + + let sub_path = dir.join("sub.parquet"); + let sub = sub_path.to_str().unwrap(); + let keep_sub: Vec = (0..9).collect(); + assert_eq!( + copy_kept_rows(&t, sub, &feats, false, &keep_sub, 1 << 20).unwrap(), + 9 + ); + + let a = TableFile::open(all).unwrap(); + let b = TableFile::open(sub).unwrap(); + assert_eq!(a.column_names(), b.column_names()); + for name in ["candidate_id", "base_peptide_id"] { + assert_eq!(a.u32(name).unwrap()[..9], b.u32(name).unwrap()[..]); + } + assert_eq!( + a.i32("peak_rank").unwrap()[..9], + b.i32("peak_rank").unwrap()[..] + ); + for name in ["label", "peptidoform", "protein"] { + assert_eq!(a.str(name).unwrap()[..9], b.str(name).unwrap()[..]); + } + for name in [ + "apex_rt", + "elution_lo", + "elution_hi", + "precursor_mz", + "prelim_score", + "charge", + "n_matched_fragments", + ] { + assert_eq!(a.f64(name).unwrap()[..9], b.f64(name).unwrap()[..]); + } + // And the pass-through copy really is the input, values and order. + assert_eq!( + a.u32("candidate_id").unwrap(), + (0..10u32).collect::>() + ); + assert_eq!( + a.f64("precursor_mz").unwrap(), + t.f64("precursor_mz").unwrap() + ); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn the_competed_table_is_written_in_capped_row_groups() { + // Rescore reads this file back a row group at a time, so the cap is its working + // set. It is the cap the rescore handoff writes with. + assert_eq!(COMPETED_ROW_GROUP_ROWS, 131_072); + let dir = tmp_dir("rowgroups"); + let src_path = dir.join("features.parquet"); + let src = src_path.to_str().unwrap(); + write_features(src, 10); + let t = TableFile::open(src).unwrap(); + let out_path = dir.join("competed.parquet"); + let out = out_path.to_str().unwrap(); + let keep: Vec = (0..10).collect(); + copy_kept_rows(&t, out, &["charge".to_string()], false, &keep, 4).unwrap(); + let file = std::fs::File::open(out).unwrap(); + let builder = + parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + let meta = builder.metadata(); + let sizes: Vec = (0..meta.num_row_groups()) + .map(|i| meta.row_group(i).num_rows()) + .collect(); + assert_eq!(sizes, vec![4, 4, 2]); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn label_codes_match_the_string_column() { + let dir = tmp_dir("labels"); + let p_path = dir.join("t.parquet"); + let p = p_path.to_str().unwrap(); + mumdia_io::table::write_table( + p, + vec![IoCol::Str( + "label".into(), + vec![ + "target".to_string(), + "decoy".to_string(), + "target".to_string(), + "entrapment".to_string(), + ], + )], + ) + .unwrap(); + let t = TableFile::open(p).unwrap(); + // The pre-2026-09-22 coding, from the materialised column. + let want: Vec = t + .str("label") + .unwrap() + .iter() + .map(|s| match s.as_str() { + "target" => 0u8, + "decoy" => 1u8, + _ => 2u8, + }) + .collect(); + assert_eq!(label_codes(&t).unwrap(), want); + assert_eq!(want, vec![0, 1, 0, 2]); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn streamed_peptidoform_ids_match_first_appearance_numbering() { + let dir = tmp_dir("pformids"); + let p_path = dir.join("t.parquet"); + let p = p_path.to_str().unwrap(); + let pforms: Vec = ["B", "A", "B", "C", "A", "C", "D"] + .iter() + .map(|s| s.to_string()) + .collect(); + mumdia_io::table::write_table(p, vec![IoCol::Str("peptidoform".into(), pforms.clone())]) + .unwrap(); + let t = TableFile::open(p).unwrap(); + // The pre-2026-09-22 numbering, from the materialised column. + let mut seen: HashMap<&str, u32> = HashMap::new(); + let want: Vec = pforms + .iter() + .map(|s| { + let next = seen.len() as u32; + *seen.entry(s.as_str()).or_insert(next) + }) + .collect(); + assert_eq!(dense_peptidoform_ids(&t).unwrap(), want); + assert_eq!(want, vec![0, 1, 0, 2, 1, 2, 3]); + let _ = std::fs::remove_dir_all(&dir); + } } From 9d6388efe2b457d4bcee3b36ebefd018ccd058c7 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:22:20 +0200 Subject: [PATCH 09/35] perf(groups): stop pooling a table nothing reads, hash each pooled file once Three changes to the grouped middle, none of which changes a value. 1. `psms_extracted.parquet` is pooled only when something reads it. Its one consumer is the candidate audit (`extract.emit_candidate_audit`, off by default and, per CLAUDE.md, reading a per-candidate audit table extract does not yet write). Pooling it read every band's extracted rows and wrote them again into the run's second-widest artifact, for a file nothing opened. The band tables are written either way, so the audit path is unchanged and the run says in the log when it skips the pool. 2. Each pooled artifact is hashed once. `record_artifact` hashed the file for the manifest record and `ArtifactReport` hashed it again for the report beside it; on a real run those are tens of GB each, and blake3_file reads all of it. `record_artifact_with_hash` takes the hash the caller already has. 3. Under `experiment.rt_library_scope = first_run_only` the runs after the first reuse the first run's band precursor slices. A slice is a deterministic function of (library, row span) and a shared-band run searches the same library under the same plan, so rewriting it produced the same bytes; only the seed reads it, and the adapted table replaces it everywhere else. On a 203M-precursor library that is the whole precursor table not written, per run after the first. The reuse is guarded on the file existing with the row count the plan expects, and falls back to writing the slice. Two more, outside the grouped path: - `split_by_source` caps its row groups at 131,072 rows, as the rescore handoff does. The scored table is ~390 float columns wide, so parquet's default 1,048,576-row group is about 3 GB decoded per group for every reader of the per-run tables. Row-group boundaries move; rows, order and values do not. - `load_ms1` copies each peak list once instead of twice. Both lists were copied whole and then copied again truncated to the shorter of the two, which on MS1 scans of tens of thousands of peaks is two allocations per scan per column that are freed immediately. New unit test covers the truncation and the RT ordering, neither of which had one. Measured on the running seven-file experiment (63 bands, 8 in flight): extract 218 CPU-min per run, features 156, search-seed 99, rt-im-train 1.2. The RT calibration is not worth hoisting; the writes are. Co-Authored-By: Claude Opus 5 (1M context) --- docs/33_window_groups.md | 42 +++++++-- rust/mumdia/crates/mumdia-io/src/lib.rs | 31 +++++- rust/mumdia/crates/mumdia/src/spectra.rs | 64 ++++++++++--- rust/mumdia/crates/mumdia/src/stages/pool.rs | 57 ++++++++++- .../mumdia/src/stages/run_experiment.rs | 8 +- .../crates/mumdia/src/stages/run_groups.rs | 94 ++++++++++++++----- 6 files changed, 242 insertions(+), 54 deletions(-) diff --git a/docs/33_window_groups.md b/docs/33_window_groups.md index f8e87f9c..d9e28ba4 100644 --- a/docs/33_window_groups.md +++ b/docs/33_window_groups.md @@ -84,11 +84,19 @@ The stages then run unchanged on that table, as if it were the library: precursor file, which is why the band is a file rather than an in-memory slice. - `rt-im-train` writes the band's `run_windows.parquet` and `cal.json`. - `extract`, `features` and `compete` write the band's `psms_extracted`, `chromatograms`, - `features` and `psms_competed` with local ids. + `features` and `psms_competed`. The id columns are library-wide on the way out + (local id plus `Library::global_offset`), so a band's table means the same thing as the + run's and pooling is a concatenation. Nothing outside the band is resident during any of these. The converted spectra are read by each stage as they always are. +Under `experiment.rt_library_scope = first_run_only` the runs after the first reuse the first +run's adapted bands, and then they reuse its band slices too: a slice is a deterministic +function of the library and the row span, so rewriting it would produce the same bytes. Only +the seed reads it, and the adapted table replaces it everywhere else. On a 203M-precursor +library that is the whole precursor table not written, per run after the first. + ## 4. The seed pool and the two calibration modes Each band's seed has two problems as a calibration anchor set. Its `spectrum_q` was @@ -140,14 +148,27 @@ pays for. On the CI fixture it cannot fit at all (no band has a confident anchor ## 5. The artifact pool -`pool` (`stages/pool.rs`) rewrites the four band tables with library-wide ids and appends -them, batch by batch, into the standard `psms_extracted.parquet`, `chromatograms.parquet`, -`features.parquet` and `psms_competed.parquet`. Where two bands hold the same candidate -(window overlap across a cut), the competed row with the higher `prelim_score` wins and the -loser's rows are dropped from all four tables, so the pooled tables hold each candidate -once and rescore's competition is unchanged. One batch is resident at a time, and the writer's -row groups are capped, so the pool costs one read and one write of the group artifacts and -no more memory than any single stage. The feature and competed schema companions +`pool` (`stages/pool.rs`) appends the band tables, batch by batch, into the standard +`chromatograms.parquet` and `psms_competed.parquet`. The bands already wrote library-wide +ids, so this is a concatenation rather than a rewrite. Where two bands hold the same +candidate (window overlap across a cut), the competed row with the higher `prelim_score` +wins and the loser's rows are dropped from every pooled table, so the pooled tables hold +each candidate once and rescore's competition is unchanged. One batch is resident at a time, +and the writer's row groups are capped, so the pool costs one read and one write of the +group artifacts and no more memory than any single stage. + +Two of the four tables an ungrouped run writes are not pooled, because nothing reads them: + +- `features.parquet`: the competed table carries the feature columns, and no stage opens a + run-level feature table. On a real run it was 55 GB of writes per run. +- `psms_extracted.parquet`: its one consumer is the candidate audit, which is off by + default (`extract.emit_candidate_audit`). With the audit on it is pooled as before; with + it off the band tables stay where they are and the run says so in the log. This is the + run's second-widest artifact, so pooling it cost a full read and a full write for a file + nothing opened. + +Each pooled table is hashed once, for the manifest record and the report beside it +together. At experiment scale those tables are tens of GB, and hashing reads all of it. The feature and competed schema companions (`
.schema.json`, the classifier's column list) are copied from the first band; every band wrote the same one. Each pooled table gets a `.report.json` whose stage is `pool` and whose stats record the number of groups and the overlap duplicates removed. @@ -172,7 +193,8 @@ out/ seed_psms.parquet pooled seed, library-wide ids (+ .masscal.json) seed_psms_calibrated.parquet pooled seed with refreshed iRT (global, after re-prediction) cal.json run-level RT calibration record (section 7) - {psms_extracted,chromatograms,features,psms_competed}.parquet pooled (+ .report.json) + {chromatograms,psms_competed}.parquet pooled (+ .report.json) + psms_extracted.parquet pooled only under extract.emit_candidate_audit psms_scored.parquet, quant, peptides.tsv, proteins.tsv, manifest.json as always ``` diff --git a/rust/mumdia/crates/mumdia-io/src/lib.rs b/rust/mumdia/crates/mumdia-io/src/lib.rs index a13aa635..94bb960f 100644 --- a/rust/mumdia/crates/mumdia-io/src/lib.rs +++ b/rust/mumdia/crates/mumdia-io/src/lib.rs @@ -86,17 +86,42 @@ pub fn record_artifact( stage: &str, config_hash: &str, ) -> Result { - Ok(ArtifactRecord { + let content_hash = hash::blake3_file(path)?; + Ok(record_artifact_with_hash( + logical_name, + schema, + path, + rows, + stage, + config_hash, + content_hash, + )) +} + +/// The same record from a hash the caller already has. A stage that writes both a manifest +/// record and an [`report::ArtifactReport`] for one file would otherwise read it twice, and +/// on a grouped run those files are tens of GB each. +#[allow(clippy::too_many_arguments)] +pub fn record_artifact_with_hash( + logical_name: &str, + schema: (&str, u32), + path: &str, + rows: u64, + stage: &str, + config_hash: &str, + content_hash: String, +) -> ArtifactRecord { + ArtifactRecord { logical_name: logical_name.to_string(), path: path.to_string(), format: "parquet".to_string(), schema_name: schema.0.to_string(), schema_version: schema.1, rows, - content_hash: hash::blake3_file(path)?, + content_hash, producing_stage: stage.to_string(), config_hash: config_hash.to_string(), - }) + } } /// Human-readable schema + head sample + row count for any Parquet artifact diff --git a/rust/mumdia/crates/mumdia/src/spectra.rs b/rust/mumdia/crates/mumdia/src/spectra.rs index df61d102..749f6677 100644 --- a/rust/mumdia/crates/mumdia/src/spectra.rs +++ b/rust/mumdia/crates/mumdia/src/spectra.rs @@ -180,29 +180,29 @@ pub fn load_ms1(path: &str) -> Result> { let mza = list_col(b, "mz")?; let ina = list_col(b, "intensity")?; for k in 0..mza.len() { - let mz: Vec = if mza.is_null(k) { - Vec::new() - } else { - let mv = mza.value(k); - inner_f32(&mv, "mz")?.values().to_vec() - }; - let intensity: Vec = if ina.is_null(k) { - Vec::new() - } else { - let iv = ina.value(k); - inner_f32(&iv, "intensity")?.values().to_vec() - }; // Truncate to the shorter list, as `load_ms2` does at :68. The two list // columns are decoded independently and either can be null, so a spectra // artifact whose m/z and intensity lists disagree in length would otherwise // be carried into `sum_near` and the MS1 isotope features, where the loop // bound comes from one array and the body indexes the other. - let n = mz.len().min(intensity.len()); + // + // The length is taken from the decoded arrays before either is copied: each + // list was previously copied whole and then copied again truncated, which on a + // run with MS1 scans of tens of thousands of peaks is two allocations per scan + // per column that are freed immediately. + let mv = (!mza.is_null(k)).then(|| mza.value(k)); + let iv = (!ina.is_null(k)).then(|| ina.value(k)); + let mf = mv.as_ref().map(|v| inner_f32(v, "mz")).transpose()?; + let iff = iv.as_ref().map(|v| inner_f32(v, "intensity")).transpose()?; + let n = match (&mf, &iff) { + (Some(m), Some(x)) => m.len().min(x.len()), + _ => 0, + }; out.push(Ms1Scan { scan_index: scan_index[i], rt_seconds: rt[i], - mz: mz[..n].to_vec(), - intensity: intensity[..n].to_vec(), + mz: mf.map(|m| m.values()[..n].to_vec()).unwrap_or_default(), + intensity: iff.map(|x| x.values()[..n].to_vec()).unwrap_or_default(), }); i += 1; } @@ -227,3 +227,37 @@ pub fn load_ms1(path: &str) -> Result> { ); Ok(out) } + +#[cfg(test)] +mod tests { + use super::*; + use mumdia_io::table::{write_table, Col}; + + /// The MS1 loader truncates a scan to the shorter of its two peak lists and keeps the + /// scans in retention-time order. Both are contracts the isotope features depend on: + /// they take the loop bound from one array and index the other. + #[test] + fn ms1_scans_are_rt_sorted_and_truncated_to_the_shorter_list() { + let dir = std::env::temp_dir().join(format!("mumdia_ms1_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("ms1.parquet").to_str().unwrap().to_string(); + write_table( + &path, + vec![ + Col::U32("scan_index".into(), vec![7, 3]), + Col::F64("rt_seconds".into(), vec![20.0, 10.0]), + // The second scan carries one more intensity than it has m/z values. + Col::ListF32("mz".into(), vec![vec![100.0, 200.0], vec![300.0]]), + Col::ListF32("intensity".into(), vec![vec![1.0, 2.0], vec![3.0, 4.0]]), + ], + ) + .unwrap(); + let scans = load_ms1(&path).unwrap(); + assert_eq!(scans.len(), 2); + assert_eq!(scans[0].scan_index, 3, "sorted by retention time"); + assert_eq!(scans[0].mz, vec![300.0]); + assert_eq!(scans[0].intensity, vec![3.0], "truncated to the m/z count"); + assert_eq!(scans[1].mz, vec![100.0, 200.0]); + assert_eq!(scans[1].intensity, vec![1.0, 2.0]); + } +} diff --git a/rust/mumdia/crates/mumdia/src/stages/pool.rs b/rust/mumdia/crates/mumdia/src/stages/pool.rs index 89d982d3..330c6b20 100644 --- a/rust/mumdia/crates/mumdia/src/stages/pool.rs +++ b/rust/mumdia/crates/mumdia/src/stages/pool.rs @@ -13,6 +13,11 @@ //! Streaming: one batch is resident at a time per table, and the writer's row groups are //! capped, so the pooling costs one read and one write of the group artifacts and no more //! memory than any single stage. +//! +//! `psms_extracted` is pooled only when something reads it (the candidate audit). The +//! extracted rows are already on disk per band; copying them into a run-level table that no +//! stage opens was a full read and a full write of the run's widest non-chromatogram +//! artifact. use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -37,7 +42,11 @@ pub struct BandArtifacts { pub struct PoolParams<'a> { pub bands: &'a [BandArtifacts], - pub out_psms: &'a str, + /// Where to pool the extracted rows, or `None` to leave them per band. Nothing in the + /// pipeline reads this table -- features ran per band, and rescore reads the competed + /// one -- so the caller passes `None` unless the candidate audit is on, which is the + /// one consumer. The per-band tables are written either way. + pub out_psms: Option<&'a str>, pub out_chromatograms: &'a str, pub out_competed: &'a str, } @@ -175,7 +184,10 @@ pub fn run(p: PoolParams) -> Result { .collect::>() }; let stats = PoolStats { - psms: pool_table(with(|b| &b.psms).into_iter(), p.out_psms)?, + psms: match p.out_psms { + Some(out) => pool_table(with(|b| &b.psms).into_iter(), out)?, + None => 0, + }, chromatograms: pool_table(with(|b| &b.chromatograms).into_iter(), p.out_chromatograms)?, competed: pool_table(with(|b| &b.competed).into_iter(), p.out_competed)?, duplicates, @@ -183,6 +195,7 @@ pub fn run(p: PoolParams) -> Result { info!( groups = p.bands.len(), psms = stats.psms, + psms_pooled = p.out_psms.is_some(), chromatograms = stats.chromatograms, competed = stats.competed, duplicates, @@ -254,7 +267,7 @@ mod tests { let (op, oc, ok) = (out("psms"), out("chrom"), out("comp")); let stats = run(PoolParams { bands: &[b0, b1], - out_psms: &op, + out_psms: Some(op.as_str()), out_chromatograms: &oc, out_competed: &ok, }) @@ -273,4 +286,42 @@ mod tests { ); } } + + #[test] + fn without_a_psms_path_the_other_tables_are_pooled_unchanged() { + let dir = std::env::temp_dir().join(format!("mumdia_pool_nopsms_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let b0 = band( + &dir, + "n0", + 0.0, + &[0, 1, 2, 3, 4], + &[9.0, 9.0, 9.0, 7.0, 2.0], + ); + let b1 = band(&dir, "n1", 0.003, &[3, 4, 5, 6], &[5.0, 6.0, 9.0, 9.0]); + let out = |n: &str| { + dir.join(format!("nopsms_{n}.parquet")) + .to_str() + .unwrap() + .to_string() + }; + let (op, oc, ok) = (out("psms"), out("chrom"), out("comp")); + let stats = run(PoolParams { + bands: &[b0, b1], + out_psms: None, + out_chromatograms: &oc, + out_competed: &ok, + }) + .unwrap(); + // The skipped table is not written, and the pooled tables that are read downstream + // are exactly what they are when it is. + assert_eq!(stats.psms, 0); + assert!(!std::path::Path::new(&op).exists()); + assert_eq!(stats.competed, 7); + assert_eq!(stats.duplicates, 2); + for path in [&oc, &ok] { + let t = TableFile::open(path).unwrap(); + assert_eq!(t.u32("candidate_id").unwrap(), vec![0, 1, 2, 3, 4, 5, 6]); + } + } } diff --git a/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs b/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs index a5b5f32a..8e596b75 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_experiment.rs @@ -361,9 +361,15 @@ fn split_by_source(scored: &str, out_paths: &[String]) -> Result<()> { .schema .index_of("source") .map_err(|_| anyhow::anyhow!("scored table has no `source` column for split"))?; + // Row groups capped as the rescore handoff caps them. The scored table is ~390 float + // columns wide, so parquet's default 1,048,576-row group is about 3 GB decoded, and + // every reader of these per-run tables (quant, report) then pays that in one allocation. + // Only the row-group boundaries change; the rows, their order and their values do not. let mut writers: Vec = out_paths .iter() - .map(|out| mumdia_io::table::BatchWriter::new(out, t.schema.clone())) + .map(|out| { + mumdia_io::table::BatchWriter::with_row_group_rows(out, t.schema.clone(), 1 << 17) + }) .collect::>()?; let mut written = 0usize; t.for_each_batch(None, 1 << 14, |b| { diff --git a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs index 92b0aff8..bf9b011b 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs @@ -166,9 +166,35 @@ pub fn run(mut g: GroupRun) -> Result { return Ok(None); } std::fs::create_dir_all(gd(b.index, ""))?; - let prec = gd(b.index, "lib_precursors.parquet"); - info!(stage = %"band-slice", group = b.index, rows = n, "run: stage start"); - groups::write_band_slice(g.lib_precursors, first, n, &prec)?; + let mut prec = gd(b.index, "lib_precursors.parquet"); + // A band's slice is a deterministic function of the library and the row span, + // and a shared-band run searches the same library under the same plan, so the + // first run's slice is the same bytes. Take it rather than write the whole + // precursor table again: only the seed reads it, and the adapted table replaces + // it for everything downstream. + let shared_slice = g + .shared_bands + .map(|s| format!("{s}/g{:02}/lib_precursors.parquet", b.index)) + .filter(|p| { + std::path::Path::new(p).exists() + && mumdia_io::table::nrows(p).is_ok_and(|r| r == n as u64) + }); + match shared_slice { + Some(p) => { + info!( + stage = %"band-slice", + group = b.index, + rows = n, + reused = %p, + "run: stage skipped" + ); + prec = p; + } + None => { + info!(stage = %"band-slice", group = b.index, rows = n, "run: stage start"); + groups::write_band_slice(g.lib_precursors, first, n, &prec)?; + } + } let seed = gd(b.index, "seed_psms.parquet"); info!(stage = %"search-seed", group = b.index, "run: stage start"); let rows = search_seed::run(search_seed::SearchSeedParams { @@ -534,37 +560,61 @@ pub fn run(mut g: GroupRun) -> Result { competed: d("psms_competed.parquet"), rt_model, }; + // The pooled extracted table has exactly one reader, the candidate audit, and that is + // off by default. Pooling it anyway read and rewrote every band's extracted rows for a + // file nothing opens; the per-band tables stay where they are either way. + let pool_psms = cfg.extract.emit_candidate_audit; + if !pool_psms { + info!( + groups = arts.len(), + "groups: psms_extracted stays per band (extract.emit_candidate_audit is off, and nothing else reads the pooled table)" + ); + } info!(stage = %"pool", groups = arts.len(), "run: stage start"); let stats = pool::run(pool::PoolParams { bands: &arts, - out_psms: &out.psms, + out_psms: pool_psms.then_some(out.psms.as_str()), out_chromatograms: &out.chromatograms, out_competed: &out.competed, }) .context("pooling the window groups")?; - for (name, schema, path, rows) in [ - ( + let pooled_artifacts = pool_psms + .then_some(( artifact::PSMS_EXTRACTED.0, artifact::PSMS_EXTRACTED, &out.psms, stats.psms, - ), - ( - artifact::CHROMATOGRAMS.0, - artifact::CHROMATOGRAMS, - &out.chromatograms, - stats.chromatograms, - ), - ( - artifact::PSMS_COMPETED.0, - artifact::PSMS_COMPETED, - &out.competed, - stats.competed, - ), - ] { + )) + .into_iter() + .chain([ + ( + artifact::CHROMATOGRAMS.0, + artifact::CHROMATOGRAMS, + &out.chromatograms, + stats.chromatograms, + ), + ( + artifact::PSMS_COMPETED.0, + artifact::PSMS_COMPETED, + &out.competed, + stats.competed, + ), + ]); + for (name, schema, path, rows) in pooled_artifacts { + // One hash for both the manifest record and the report beside the file: these are + // the run's largest artifacts, and hashing reads all of it. + let content_hash = mumdia_io::hash::blake3_file(path)?; record_opt( g.man.as_deref_mut(), - record_artifact(name, schema, path, rows, "pool", ch)?, + mumdia_io::record_artifact_with_hash( + name, + schema, + path, + rows, + "pool", + ch, + content_hash.clone(), + ), ); ArtifactReport { logical_name: name.to_string(), @@ -572,7 +622,7 @@ pub fn run(mut g: GroupRun) -> Result { schema_version: schema.1, stage: "pool".to_string(), rows, - content_hash: mumdia_io::hash::blake3_file(path)?, + content_hash, params: json!({ "window_groups": arts.len(), "calibration": if global { "global" } else { "per_group" }, From 47857bc66c4ce2b675660ed4dde97cba821757bc Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:22:38 +0200 Subject: [PATCH 10/35] perf(rescore): flatten the per-fold training matrix, and stop copying columns that already exist Five allocation-count reductions in the rescore stage. Output is unchanged: the crafted features -> compete -> rescore pipeline writes a scored parquet with the same blake3 (e535d04b2821a47d4fc8f9bc0a7ba57461b701b607496ffef0700a0352cd135a) before and after, and the two kernels that were rewritten are checked against transcriptions of their previous selves at bit level. 1. `percolator_lite` built its standardised training slice as a `Vec>`, one heap block per training row, inside the parallel fold map, so every fold's copy is live at once. That is the exact layout `FeatureMatrix`'s own comment records was removed from the matrix for this reason, surviving one level down: at 11.6M PSMs x 3 folds it is ~23M live blocks, now 3 flat buffers sliced by `chunks_exact`. The held-out fold is scored through the standardizer without materialising a row at all. The f32 narrowing is kept at both sites, because the score consumed the standardised value AFTER it was rounded to f32. 2. The competed-input merge cloned every metadata string (label, peptidoform, protein) into a second vector while the reader's stayed alive: twice the string bytes and 3 extra live allocations per PSM at peak. `merge_col` moves the first input's column and reserves the footer row count, so a single-input rescore copies nothing and a multi-input one reallocates once. 3. `target_decoy_q` sorted an index permutation with a comparator that dereferenced a separate key column (two random loads per comparison, two more per row in the tie walk) and allocated four n-length vectors. It now sorts one 16-byte `(key, row, is_decoy)` record in parallel and merges the monotonization into the block walk, so `fdr_at` is gone: 32 -> 24 bytes per row, 4 -> 2 allocations, and no indirection. The tie-block handling is unchanged, and the permutation is the same one the stable sort produced (`key` desc, row asc). 4. The non-finite validation was a serial `n_rows x n_features` scan over hundreds of GB. Both scans are parallel `find_first` now, and their answers are combined so the row reported is still the first offending row, features before scalars within a row, exactly as the interleaved serial loop reported it. 5. `psms_scored` has three duplicate columns by construction (`protein_group` is `protein`; `global_q_value` and `experiment_psm_q` are the pooled `q_value`). Through `write_table` each had to arrive as its own `Vec`, so the stage carried a second full copy of the protein strings and two more of the q column. They share one refcounted Arrow array now. A test asserts the parquet is byte-for-byte what `write_table` wrote from separate equal vectors. Risk to output equality: (1) and (3) are arithmetic-identical by construction and covered by reference-implementation tests; (5) is byte-tested; (2) and (4) move or parallelise without touching values. The one behaviour that is not identical is which error a malformed input reports when it has both a non-finite feature and a non-finite prelim_score in the SAME row -- it is still the feature, and still the first offending row, so this is preserved too. Co-Authored-By: Claude Fable 5.1 --- rust/mumdia/crates/mumdia/src/fdr.rs | 172 +++++++-- rust/mumdia/crates/mumdia/src/rescoring.rs | 298 ++++++++++++++- .../crates/mumdia/src/stages/rescore.rs | 351 ++++++++++++++---- 3 files changed, 714 insertions(+), 107 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/fdr.rs b/rust/mumdia/crates/mumdia/src/fdr.rs index 00c6be77..a14168a4 100644 --- a/rust/mumdia/crates/mumdia/src/fdr.rs +++ b/rust/mumdia/crates/mumdia/src/fdr.rs @@ -2,6 +2,8 @@ //! estimator `q = (n_decoys + 1) / max(1, n_targets)`, monotonized, with tied //! scores collapsed to a single block q. Shared by search-seed and rescore. +use rayon::prelude::*; + /// Rank key for the q-value kernels: finite scores pass through, non-finite ones /// (NaN, +/-inf) become the worst possible key. /// @@ -27,45 +29,82 @@ fn rank_key(x: f64) -> f64 { /// Compute per-record q-values from (score, is_decoy). Higher score is better. /// Returns q aligned to the input order. +/// +/// Layout, not statistics: the walk below is the same tied-block estimator it has always +/// been, but it ranks ONE sortable record per row -- `(key, row, is_decoy)`, 16 bytes -- +/// instead of sorting an index permutation with a comparator that dereferences a separate +/// key column. The old comparator paid two random loads per comparison (`key[a]`, +/// `key[b]`) and the tied-block walk paid two more per row (`key[order[end]]`, +/// `scores[order[end]].1`); here the sort compares the key in hand and the walk is a +/// sequential scan. It also allocates two vectors rather than four (`key`, `order`, +/// `fdr_at`, `q`): the `fdr_at` column is gone because the monotonization now runs in the +/// same backward pass that computes each block's FDR, from the per-block counts and the +/// totals, so 32 bytes per row becomes 24. At the experiment scale this kernel runs at +/// (pooled PSM q over ~11.6M rows, plus `folds x num_iter` calls inside +/// `rescoring::percolator_lite`) that is the difference between a random-access sort and +/// a cache-resident one. +/// +/// Output is unchanged, including tie behaviour. The previous stable `sort_by` over an +/// ascending index vector ordered ties by row index; sorting `(key desc, row asc)` is the +/// same permutation, and no two records compare equal, so the unstable parallel sort is +/// deterministic. pub fn target_decoy_q(scores: &[(f64, bool)]) -> Vec { let n = scores.len(); if n == 0 { return Vec::new(); } - let key: Vec = scores.iter().map(|&(s, _)| rank_key(s)).collect(); - let mut order: Vec = (0..n).collect(); - order.sort_by(|&a, &b| key[b].total_cmp(&key[a])); - let (mut td, mut tt) = (0usize, 0usize); - let mut fdr_at = vec![1.0f64; n]; + // A u32 row index halves the sort record. 2^32 PSMs is not a scale this engine can + // reach (the feature matrix alone would be 6.6 TB), but an assert is still better + // than a silent truncation that would scatter q onto the wrong rows. + assert!( + n <= u32::MAX as usize, + "target_decoy_q: {n} rows exceeds the u32 row index" + ); + let mut ranked: Vec<(f64, u32, bool)> = scores + .iter() + .enumerate() + .map(|(i, &(s, d))| (rank_key(s), i as u32, d)) + .collect(); + ranked.par_sort_unstable_by(|a, b| b.0.total_cmp(&a.0).then_with(|| a.1.cmp(&b.1))); + let total_d = scores.iter().filter(|&&(_, d)| d).count(); + let total_t = n - total_d; // Walk in score order, processing tied-score blocks together so every PSM in // a block gets the same FDR (its within-tie order is arbitrary and must not // change the q). Numerator uses `n_decoys + 1` (the standard conservative // target-decoy estimate); the bare `n_decoys / n_targets` is optimistic in // the low-count regime. - let mut rank = 0usize; - while rank < n { - let s = key[order[rank]]; - let mut end = rank; - while end < n && key[order[end]] == s { - if scores[order[end]].1 { - td += 1; + // + // Walked from the WORST-scoring end so the monotonization (q non-increasing with + // score) happens in the same pass: the counts at or above a block are the totals + // minus what is strictly below it, which is exactly what the forward walk + // accumulated, and `qmin` over the blocks already visited is exactly what the + // separate backward pass over `fdr_at` used to compute. + let mut q = vec![1.0f64; n]; + let (mut below_d, mut below_t) = (0usize, 0usize); + let mut qmin = 1.0f64; + let mut end = n; + while end > 0 { + let s = ranked[end - 1].0; + let mut start = end; + let (mut block_d, mut block_t) = (0usize, 0usize); + while start > 0 && ranked[start - 1].0 == s { + start -= 1; + if ranked[start].2 { + block_d += 1; } else { - tt += 1; + block_t += 1; } - end += 1; } + let td = total_d - below_d; + let tt = total_t - below_t; let f = (td as f64 + 1.0) / (tt.max(1) as f64); - for value in fdr_at.iter_mut().take(end).skip(rank) { - *value = f; + qmin = qmin.min(f); + for r in &ranked[start..end] { + q[r.1 as usize] = qmin; } - rank = end; - } - // Monotonize from worst-scoring to best so q is non-increasing with score. - let mut q = vec![1.0f64; n]; - let mut qmin = 1.0f64; - for rank in (0..n).rev() { - qmin = qmin.min(fdr_at[rank]); - q[order[rank]] = qmin; + below_d += block_d; + below_t += block_t; + end = start; } q } @@ -164,6 +203,91 @@ pub fn ln_factorial(n: u32) -> f64 { mod tests { use super::*; + /// The previous `target_decoy_q`, transcribed unchanged: an index permutation sorted + /// with an indirect comparator, a forward tied-block walk into `fdr_at`, and a + /// separate backward monotonization. Kept as the reference the rewritten kernel is + /// checked against, because the claim being made is equality of output, not merely + /// that the new code is self-consistent. + fn target_decoy_q_reference(scores: &[(f64, bool)]) -> Vec { + let n = scores.len(); + if n == 0 { + return Vec::new(); + } + let key: Vec = scores.iter().map(|&(s, _)| rank_key(s)).collect(); + let mut order: Vec = (0..n).collect(); + order.sort_by(|&a, &b| key[b].total_cmp(&key[a])); + let (mut td, mut tt) = (0usize, 0usize); + let mut fdr_at = vec![1.0f64; n]; + let mut rank = 0usize; + while rank < n { + let s = key[order[rank]]; + let mut end = rank; + while end < n && key[order[end]] == s { + if scores[order[end]].1 { + td += 1; + } else { + tt += 1; + } + end += 1; + } + let f = (td as f64 + 1.0) / (tt.max(1) as f64); + for value in fdr_at.iter_mut().take(end).skip(rank) { + *value = f; + } + rank = end; + } + let mut q = vec![1.0f64; n]; + let mut qmin = 1.0f64; + for rank in (0..n).rev() { + qmin = qmin.min(fdr_at[rank]); + q[order[rank]] = qmin; + } + q + } + + #[test] + fn q_is_bit_identical_to_the_previous_kernel() { + // A deterministic pseudo-random population with every case the rewrite touches: + // heavy score ties (the tied-block walk), signed zeros (`total_cmp` separates + // them, `==` groups them), non-finite scores (mapped to the worst key), decoys + // above and below targets, and single-class stretches. + let mut state = 0x2545_F491_4F6C_DD1Du64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + for n in [1usize, 2, 3, 7, 64, 513, 4096] { + let mut scores: Vec<(f64, bool)> = Vec::with_capacity(n); + for _ in 0..n { + let r = next(); + // Coarse quantisation so ties are common rather than incidental. + let s = match r % 11 { + 0 => 0.0, + 1 => -0.0, + 2 => f64::NAN, + 3 => f64::INFINITY, + 4 => f64::NEG_INFINITY, + _ => ((r >> 8) % 37) as f64 * 0.25 - 4.0, + }; + scores.push((s, r % 3 == 0)); + } + let got = target_decoy_q(&scores); + let want = target_decoy_q_reference(&scores); + assert_eq!( + got.iter().map(|v| v.to_bits()).collect::>(), + want.iter().map(|v| v.to_bits()).collect::>(), + "n = {n}" + ); + } + // All-target and all-decoy populations, where the `max(1, .)` guard bites. + for flag in [false, true] { + let scores: Vec<(f64, bool)> = (0..32).map(|i| ((i % 4) as f64, flag)).collect(); + assert_eq!(target_decoy_q(&scores), target_decoy_q_reference(&scores)); + } + } + #[test] fn perfect_separation_q_is_conservative_plus_one() { // all targets score above all decoys; with the (n_decoys+1)/n_targets diff --git a/rust/mumdia/crates/mumdia/src/rescoring.rs b/rust/mumdia/crates/mumdia/src/rescoring.rs index 8827cb47..bc7a428a 100644 --- a/rust/mumdia/crates/mumdia/src/rescoring.rs +++ b/rust/mumdia/crates/mumdia/src/rescoring.rs @@ -46,6 +46,36 @@ impl FeatureMatrix { self.values.push(v as f32); } + /// Append a whole row that a caller has already narrowed to f32. + #[inline] + pub fn push_row(&mut self, row: &[f32]) { + self.values.extend_from_slice(row); + } + + /// `(row, feature)` of the first non-finite value in flat row order, or `None`. + /// + /// Parallel because the serial form was `n_rows * n_features` `is_finite` checks on one + /// thread over a matrix that is hundreds of gigabytes at experiment scale, for a + /// predicate that is embarrassingly parallel. `find_first` is order-preserving: it + /// returns the earliest matching row regardless of which worker found it, so the error + /// a malformed table produces is the same row and the same message as before. + pub fn find_non_finite(&self) -> Option<(usize, usize)> { + if self.n_features == 0 { + return None; + } + self.values + .par_chunks_exact(self.n_features) + .enumerate() + .find_first(|(_, row)| row.iter().any(|v| !v.is_finite())) + .map(|(row, values)| { + let col = values + .iter() + .position(|v| !v.is_finite()) + .expect("the row matched the predicate"); + (row, col) + }) + } + /// Fail loudly rather than silently mis-striding if a caller pushed a partial row. pub fn finish(self) -> anyhow::Result { if self.n_features > 0 && !self.values.len().is_multiple_of(self.n_features) { @@ -115,14 +145,27 @@ fn fit_standardizer(x: &FeatureMatrix, idx: &[usize]) -> (Vec, Vec) { (mean, std) } -/// Standardise one row. The subtraction and division are f64; the result is kept f32 -/// because `xtr` below holds one standardised copy of the training slice per fold, and -/// those copies are as large as the matrix itself. +/// Standardise one row into `out`. The subtraction and division are f64; the result is +/// kept f32 because `xtr` below holds one standardised copy of the training slice per +/// fold, and those copies are as large as the matrix itself. #[inline] -fn std_row(row: &[f32], mean: &[f64], std: &[f64]) -> Vec { - (0..row.len()) - .map(|j| ((row[j] as f64 - mean[j]) / std[j]) as f32) - .collect() +fn std_row_into(row: &[f32], mean: &[f64], std: &[f64], out: &mut Vec) { + out.extend((0..row.len()).map(|j| ((row[j] as f64 - mean[j]) / std[j]) as f32)); +} + +/// Score a row through the standardizer without materialising it. +/// +/// The f32 narrowing is kept: `score_row(w, &std_row(..))` scored the standardised value +/// AFTER it had been rounded to f32, so computing the product from the f64 quotient +/// instead would change the last bits of every test-fold score. +#[inline] +fn score_std_row(w: &[f64], row: &[f32], mean: &[f64], std: &[f64]) -> f64 { + let mut z = w[0]; + for j in 0..row.len() { + let v = ((row[j] as f64 - mean[j]) / std[j]) as f32; + z += w[j + 1] * v as f64; + } + z } /// Logistic regression by full-batch gradient descent with L2. Weight[0] = bias. @@ -199,11 +242,20 @@ pub fn percolator_lite(inp: RescoreInput) -> Vec { return Vec::new(); } let (mean, std) = fit_standardizer(inp.features, &train_idx); - // standardized train matrix (owned, so we can take &[f64] slices) - let xtr: Vec> = train_idx - .iter() - .map(|&i| std_row(inp.features.row(i), &mean, &std)) - .collect(); + // Standardized train matrix: ONE flat allocation of `train_rows * d`, row-major, + // sliced with `chunks_exact(d)` below. It used to be a `Vec>`, which is + // one heap block per training row -- the exact layout the comment on + // `FeatureMatrix` records was removed from the matrix itself, surviving one + // level down and multiplied by `folds`, since every fold is live at once inside + // this parallel map. At experiment scale (11.6M PSMs, 3 folds) that is ~23M + // live blocks here; now it is 3. Same values, same order, so the fit is + // bit-identical. + let d = inp.features.n_features(); + let mut xtr: Vec = Vec::with_capacity(train_idx.len().saturating_mul(d)); + for &i in &train_idx { + std_row_into(inp.features.row(i), &mean, &std, &mut xtr); + } + let xrow = |k: usize| -> &[f32] { &xtr[k * d..(k + 1) * d] }; let mut train_scores: Vec = train_idx.iter().map(|&i| inp.init_score[i]).collect(); let mut w = vec![0.0; inp.features.n_features() + 1]; let mut sd: Vec<(f64, bool)> = Vec::with_capacity(train_idx.len()); @@ -222,10 +274,10 @@ pub fn percolator_lite(inp: RescoreInput) -> Vec { let mut n_pos = 0; for (k, &i) in train_idx.iter().enumerate() { if inp.is_decoy[i] { - rows.push(&xtr[k]); + rows.push(xrow(k)); ys.push(0.0); } else if q[k] <= inp.train_fdr { - rows.push(&xtr[k]); + rows.push(xrow(k)); ys.push(1.0); n_pos += 1; } @@ -240,23 +292,23 @@ pub fn percolator_lite(inp: RescoreInput) -> Vec { for (rank, &k) in order.iter().enumerate() { let i = train_idx[k]; if inp.is_decoy[i] { - rows.push(&xtr[k]); + rows.push(xrow(k)); ys.push(0.0); } else if rank < take { - rows.push(&xtr[k]); + rows.push(xrow(k)); ys.push(1.0); } } } w = logreg_fit(&rows, &ys, 1e-3, 200, 0.5); train_scores = (0..train_idx.len()) - .map(|k| score_row(&w, &xtr[k])) + .map(|k| score_row(&w, xrow(k))) .collect(); } // score the held-out test fold with this fold's scaler + weights test_idx .iter() - .map(|&i| (i, score_row(&w, &std_row(inp.features.row(i), &mean, &std)))) + .map(|&i| (i, score_std_row(&w, inp.features.row(i), &mean, &std))) .collect() }) .collect(); @@ -274,6 +326,216 @@ pub fn percolator_lite(inp: RescoreInput) -> Vec { mod tests { use super::*; + /// The previous per-row standardiser: one `Vec` (one heap block) per row. + fn std_row_ref(row: &[f32], mean: &[f64], std: &[f64]) -> Vec { + (0..row.len()) + .map(|j| ((row[j] as f64 - mean[j]) / std[j]) as f32) + .collect() + } + + /// `percolator_lite` as it was before the flattening, transcribed: `xtr` as a + /// `Vec>`, and the test fold scored through a materialised `std_row`. The + /// claim the rewrite makes is bit-equality of the scores, so the old layout is kept + /// here as the thing to compare against rather than asserting properties of the new + /// one. + fn percolator_lite_reference(inp: RescoreInput) -> Vec { + let n = inp.features.rows(); + if n == 0 { + return Vec::new(); + } + let folds = inp.folds.max(1); + let fold_of: Vec = inp.fold_key.iter().map(|c| (*c as usize) % folds).collect(); + let per_fold: Vec> = (0..folds) + .into_par_iter() + .map(|test_fold| { + let train_idx: Vec = (0..n).filter(|&i| fold_of[i] != test_fold).collect(); + let test_idx: Vec = (0..n).filter(|&i| fold_of[i] == test_fold).collect(); + if train_idx.is_empty() || test_idx.is_empty() { + return Vec::new(); + } + let (mean, std) = fit_standardizer(inp.features, &train_idx); + let xtr: Vec> = train_idx + .iter() + .map(|&i| std_row_ref(inp.features.row(i), &mean, &std)) + .collect(); + let mut train_scores: Vec = + train_idx.iter().map(|&i| inp.init_score[i]).collect(); + let mut w = vec![0.0; inp.features.n_features() + 1]; + let mut sd: Vec<(f64, bool)> = Vec::with_capacity(train_idx.len()); + for _ in 0..inp.num_iter.max(1) { + sd.clear(); + sd.extend( + train_idx + .iter() + .enumerate() + .map(|(k, &i)| (train_scores[k], inp.is_decoy[i])), + ); + let q = target_decoy_q(&sd); + let mut rows: Vec<&[f32]> = Vec::new(); + let mut ys: Vec = Vec::new(); + let mut n_pos = 0; + for (k, &i) in train_idx.iter().enumerate() { + if inp.is_decoy[i] { + rows.push(&xtr[k]); + ys.push(0.0); + } else if q[k] <= inp.train_fdr { + rows.push(&xtr[k]); + ys.push(1.0); + n_pos += 1; + } + } + if n_pos < 10 { + let mut order: Vec = (0..train_idx.len()).collect(); + order.sort_by(|&a, &b| train_scores[b].total_cmp(&train_scores[a])); + let take = (train_idx.len() / 2).max(1); + rows.clear(); + ys.clear(); + for (rank, &k) in order.iter().enumerate() { + let i = train_idx[k]; + if inp.is_decoy[i] { + rows.push(&xtr[k]); + ys.push(0.0); + } else if rank < take { + rows.push(&xtr[k]); + ys.push(1.0); + } + } + } + w = logreg_fit(&rows, &ys, 1e-3, 200, 0.5); + train_scores = (0..train_idx.len()) + .map(|k| score_row(&w, &xtr[k])) + .collect(); + } + test_idx + .iter() + .map(|&i| { + ( + i, + score_row(&w, &std_row_ref(inp.features.row(i), &mean, &std)), + ) + }) + .collect() + }) + .collect(); + let mut final_score = inp.init_score.to_vec(); + for fold_scores in per_fold { + for (i, s) in fold_scores { + final_score[i] = s; + } + } + final_score + } + + /// A synthetic population wide enough that the flattening matters and mixed enough + /// that both the confident-target branch and the top-half fallback are exercised. + fn crafted_population(n: usize, d: usize) -> (FeatureMatrix, Vec, Vec, Vec) { + let mut features = FeatureMatrix::with_capacity(n, d); + let (mut is_decoy, mut key, mut init) = (Vec::new(), Vec::new(), Vec::new()); + let mut state = 0x9E37_79B9_7F4A_7C15u64; + for i in 0..n { + let decoy = i % 3 == 0; + for j in 0..d { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + let noise = (state >> 40) as f64 / 16_777_216.0; + // Column 0 separates; the rest are noise of varying scale, including one + // constant column (std guard) and one with a large offset. + let v = match j { + 0 => (if decoy { 0.0 } else { 2.5 }) + noise, + 1 => 7.0, + 2 => 1e6 * noise, + _ => noise - 0.5, + }; + features.push(v); + } + is_decoy.push(decoy); + key.push((i / 2) as u32); + init.push(if decoy { -0.5 } else { 0.5 } + (i % 7) as f64 * 0.01); + } + (features.finish().unwrap(), is_decoy, key, init) + } + + #[test] + fn flat_training_matrix_scores_bit_identically_to_the_per_row_layout() { + let (features, is_decoy, key, init) = crafted_population(600, 9); + let mk = || RescoreInput { + features: &features, + is_decoy: &is_decoy, + fold_key: &key, + init_score: &init, + folds: 3, + num_iter: 4, + train_fdr: 0.05, + }; + let got = percolator_lite(mk()); + let want = percolator_lite_reference(mk()); + assert_eq!( + got.iter().map(|v| v.to_bits()).collect::>(), + want.iter().map(|v| v.to_bits()).collect::>() + ); + // And the fallback branch (too few confident targets to train on) as well. + let (features, is_decoy, key, init) = crafted_population(40, 5); + let mk = || RescoreInput { + features: &features, + is_decoy: &is_decoy, + fold_key: &key, + init_score: &init, + folds: 2, + num_iter: 3, + train_fdr: 1e-9, + }; + assert_eq!( + percolator_lite(mk()) + .iter() + .map(|v| v.to_bits()) + .collect::>(), + percolator_lite_reference(mk()) + .iter() + .map(|v| v.to_bits()) + .collect::>() + ); + } + + #[test] + fn non_finite_scan_reports_the_first_row_in_flat_order() { + let mut m = FeatureMatrix::with_capacity(4, 3); + for v in [ + 1.0, + 2.0, + 3.0, // row 0: clean + 4.0, + 5.0, + 6.0, // row 1: clean + 7.0, + f64::NAN, + f64::INFINITY, // row 2: feature 1 first + f64::NEG_INFINITY, + 0.0, + 0.0, // row 3: also bad, must not win + ] { + m.push(v); + } + let m = m.finish().unwrap(); + assert_eq!(m.find_non_finite(), Some((2, 1))); + // Same answer as the serial scan it replaces, on a clean matrix too. + let serial = m + .iter_rows() + .enumerate() + .find_map(|(r, row)| row.iter().position(|v| !v.is_finite()).map(|c| (r, c))); + assert_eq!(m.find_non_finite(), serial); + let mut clean = FeatureMatrix::with_capacity(2, 2); + for v in [1.0, 2.0, 3.0, 4.0] { + clean.push(v); + } + assert_eq!(clean.finish().unwrap().find_non_finite(), None); + // An f64 that overflows f32 is non-finite once narrowed, and must be caught: the + // scan reads the stored f32, not the value the reader decoded. + let mut overflow = FeatureMatrix::with_capacity(1, 1); + overflow.push(1e300); + assert_eq!(overflow.finish().unwrap().find_non_finite(), Some((0, 0))); + } + #[test] fn separates_targets_from_decoys() { // targets have high feature[0], decoys low, plus noise diff --git a/rust/mumdia/crates/mumdia/src/stages/rescore.rs b/rust/mumdia/crates/mumdia/src/stages/rescore.rs index bee7ed51..750b7ea9 100644 --- a/rust/mumdia/crates/mumdia/src/stages/rescore.rs +++ b/rust/mumdia/crates/mumdia/src/stages/rescore.rs @@ -12,6 +12,7 @@ use mumdia_core::config::{FeaturePreset, RescoreConfig, RescorerKind}; use mumdia_core::schema::artifact; use mumdia_io::report::ArtifactReport; use mumdia_io::table::{write_table, Col, TableFile}; +use rayon::prelude::*; use serde_json::json; use tracing::{info, warn}; @@ -88,6 +89,23 @@ fn human_bytes(bytes: f64) -> String { } } +/// Append one competed input's column to the concatenated column. +/// +/// The first input is MOVED in, not copied, and the destination is then reserved to the +/// total row count known from the parquet footers, so a multi-input concatenation +/// reallocates once and every later input appends into spare capacity. For the common +/// single-input rescore nothing is copied at all. +fn merge_col(dst: &mut Vec, mut src: Vec, total_rows: usize) { + if dst.is_empty() { + *dst = src; + if dst.len() < total_rows { + dst.reserve(total_rows - dst.len()); + } + return; + } + dst.append(&mut src); +} + /// Refuse a rescore whose population cannot support target-decoy FDR. /// /// A one-sided population is not a hard rescore failure in any technical sense: @@ -261,21 +279,30 @@ pub fn run(p: RescoreParams) -> Result { } } } - for i in 0..t.nrows { - cid.push(c[i]); - peak_rank.push(pkr[i]); - label.push(l[i].clone()); - base.push(b[i]); - pform.push(pf[i].clone()); - protein.push(pr[i].clone()); - charge.push(z[i] as i32); - prelim.push(pl[i]); - mz.push(pm[i]); - apex_rt.push(ar[i]); - elution_lo.push(elo[i]); - elution_hi.push(ehi[i]); - source.push(src as u32); - } + // Take the reader's columns rather than copying them row by row into a second set. + // The row loop this replaces cloned every metadata STRING -- label, peptidoform, + // protein -- into a fresh allocation while the source vector stayed alive until the + // end of the iteration, so at peak both copies existed: twice the string bytes and, + // more importantly, twice the live allocation count, three per PSM. `merge_col` + // moves the whole vector for the first (usually only) input and appends the rest, + // which copies `String` headers but never the heap blocks they own. + merge_col(&mut cid, c, total_rows); + merge_col(&mut peak_rank, pkr, total_rows); + merge_col(&mut label, l, total_rows); + merge_col(&mut base, b, total_rows); + merge_col(&mut pform, pf, total_rows); + merge_col(&mut protein, pr, total_rows); + merge_col( + &mut charge, + z.into_iter().map(|v| v as i32).collect(), + total_rows, + ); + merge_col(&mut prelim, pl, total_rows); + merge_col(&mut mz, pm, total_rows); + merge_col(&mut apex_rt, ar, total_rows); + merge_col(&mut elution_lo, elo, total_rows); + merge_col(&mut elution_hi, ehi, total_rows); + merge_col(&mut source, vec![src as u32; t.nrows], total_rows); } let feats = matrix.finish()?; crate::memlog::report( @@ -297,23 +324,30 @@ pub fn run(p: RescoreParams) -> Result { let n_decoys = is_decoy.iter().filter(|&&v| v).count(); let n_targets = n - n_decoys; require_both_labels(n_targets, n_decoys)?; - for (row, values) in feats.iter_rows().enumerate() { - if let Some((feature, value)) = values - .iter() - .enumerate() - .find(|(_, value)| !value.is_finite()) - { - anyhow::bail!( - "rescore input contains non-finite feature '{}' at flat row {row}: {value}", - feat_names[feature] - ); - } - if !prelim[row].is_finite() || !mz[row].is_finite() { + // The validation this replaces was one serial pass of `n_rows x n_features` + // `is_finite` checks -- hundreds of gigabytes on one thread at experiment scale -- + // for a predicate with no carried state. Both scans are parallel now and both are + // order-preserving (`find_first`), and the two answers are combined so the row + // reported is still the FIRST offending row, features before scalars within a row, + // exactly as the interleaved serial loop reported it. + let bad_feature = feats.find_non_finite(); + let bad_scalar = (0..n) + .into_par_iter() + .find_first(|&row| !prelim[row].is_finite() || !mz[row].is_finite()); + if let Some((row, feature)) = bad_feature { + if bad_scalar.is_none_or(|scalar_row| row <= scalar_row) { anyhow::bail!( - "rescore input contains non-finite prelim_score/precursor_mz at flat row {row}" + "rescore input contains non-finite feature '{}' at flat row {row}: {}", + feat_names[feature], + feats.row(row)[feature] ); } } + if let Some(row) = bad_scalar { + anyhow::bail!( + "rescore input contains non-finite prelim_score/precursor_mz at flat row {row}" + ); + } } info!( psms = n, @@ -599,10 +633,10 @@ pub fn run(p: RescoreParams) -> Result { // Multi-context q-values (docs/11_compete_rescore_fdr.md). The pooled per-PSM q // is `experiment_psm_q`; `run_psm_q` re-runs TDA within each source (run) so a // per-run quant/report gets a real per-run FDR rather than the pooled value; - // `precursor_q` groups on peptidoform+charge. `global_q` is kept as a - // byte-identical alias of the pooled q for backward-compat. - let global_q = psm_q.clone(); - let experiment_psm_q = psm_q.clone(); + // `precursor_q` groups on peptidoform+charge. `global_q_value` and + // `experiment_psm_q` are kept as byte-identical aliases of the pooled q for + // backward-compat, and are written from the SAME Arrow array below rather than from + // two more copies of the column (see `write_scored_table`). // Per-run PSM q: TDA within each source separately, scattered back by row index. // Single-run (source all-zero) => equals `q_value`. Sorted (BTree) source // iteration keeps it deterministic; no floats are summed. @@ -703,40 +737,28 @@ pub fn run(p: RescoreParams) -> Result { seen.len() }; - let rows = write_table( + let rows = write_scored_table( p.out, - vec![ - Col::U32("candidate_id".into(), cid), - Col::Str("peptidoform".into(), pform), - Col::I32("charge".into(), charge), - Col::Str("label".into(), label), - // `protein` feeds two columns; clone once here and move the original below. - Col::Str("protein".into(), protein.clone()), - Col::U32("base_peptide_id".into(), base), - Col::F64("apex_rt".into(), apex_rt), - Col::F64("elution_lo".into(), elution_lo), - Col::F64("elution_hi".into(), elution_hi), - Col::F64("score".into(), scores), - Col::F64("q_value".into(), psm_q), - Col::F64("peptide_q_value".into(), peptide_q), - Col::Str("protein_group".into(), protein), - Col::F64("pg_q_value".into(), pg_q), - Col::F64("global_q_value".into(), global_q), - Col::F64("prelim_score".into(), prelim), - // Run identity for experiment-wide rescore (index into --competed); - // all-zero for a single-run rescore. Lets quant map scores per file. - Col::U32("source".into(), source), - // Multi-context q columns (docs/11_compete_rescore_fdr.md). - // run_psm_q = per-run PSM FDR; experiment_psm_q = pooled PSM FDR - // (== q_value/global_q_value); precursor_q = per (peptidoform+charge) - // FDR. - Col::F64("run_psm_q".into(), run_psm_q), - Col::F64("experiment_psm_q".into(), experiment_psm_q), - Col::F64("precursor_q".into(), precursor_q), - // Which chromatographic peak the rescorer selected for this candidate - // (#7). 0 = the up-front apex; > 0 = a promoted alternate peak won. - Col::I32("selected_peak_rank".into(), peak_rank), - ], + ScoredColumns { + cid, + pform, + charge, + label, + protein, + base, + apex_rt, + elution_lo, + elution_hi, + scores, + psm_q, + peptide_q, + pg_q, + prelim, + source, + run_psm_q, + precursor_q, + peak_rank, + }, )?; let elapsed = t0.elapsed().as_millis(); @@ -807,6 +829,112 @@ pub fn run(p: RescoreParams) -> Result { Ok(rows) } +/// The `psms_scored` columns, in schema order. +struct ScoredColumns { + cid: Vec, + pform: Vec, + charge: Vec, + label: Vec, + protein: Vec, + base: Vec, + apex_rt: Vec, + elution_lo: Vec, + elution_hi: Vec, + scores: Vec, + psm_q: Vec, + peptide_q: Vec, + pg_q: Vec, + prelim: Vec, + source: Vec, + run_psm_q: Vec, + precursor_q: Vec, + peak_rank: Vec, +} + +/// Write `psms_scored`, sharing the Arrow array behind every column that repeats. +/// +/// Three of the 21 columns are duplicates by construction: `protein_group` is `protein`, +/// and `global_q_value` and `experiment_psm_q` are both the pooled `q_value` (the comments +/// at their definitions say so). Through `write_table` each duplicate had to arrive as its +/// own `Vec`, so the stage carried a second full copy of the protein strings -- one heap +/// block per PSM -- and two more copies of the pooled q column, all three alive at the +/// same moment as the originals. An `ArrayRef` is refcounted, so passing the same array in +/// several slots of the batch writes the same bytes from one buffer. +/// +/// The parquet is byte-identical to the previous `write_table` call: the same schema +/// (names, types, all non-nullable, same order), one record batch, and `BatchWriter` opens +/// the identical SNAPPY writer on the identical `AtomicPath` temp-then-rename. The +/// duplicate columns are encoded independently, exactly as they were when they were +/// separate vectors. +fn write_scored_table(path: &str, c: ScoredColumns) -> Result { + use arrow::array::{ArrayRef, Int32Array, StringArray, UInt32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use std::sync::Arc; + + let f = |name: &str, t: DataType| Field::new(name, t, false); + let schema = Arc::new(Schema::new(vec![ + f("candidate_id", DataType::UInt32), + f("peptidoform", DataType::Utf8), + f("charge", DataType::Int32), + f("label", DataType::Utf8), + f("protein", DataType::Utf8), + f("base_peptide_id", DataType::UInt32), + f("apex_rt", DataType::Float64), + f("elution_lo", DataType::Float64), + f("elution_hi", DataType::Float64), + f("score", DataType::Float64), + f("q_value", DataType::Float64), + f("peptide_q_value", DataType::Float64), + f("protein_group", DataType::Utf8), + f("pg_q_value", DataType::Float64), + f("global_q_value", DataType::Float64), + f("prelim_score", DataType::Float64), + // Run identity for experiment-wide rescore (index into --competed); all-zero for + // a single-run rescore. Lets quant map scores per file. + f("source", DataType::UInt32), + // Multi-context q columns (docs/11_compete_rescore_fdr.md). run_psm_q = per-run + // PSM FDR; experiment_psm_q = pooled PSM FDR (== q_value/global_q_value); + // precursor_q = per (peptidoform+charge) FDR. + f("run_psm_q", DataType::Float64), + f("experiment_psm_q", DataType::Float64), + f("precursor_q", DataType::Float64), + // Which chromatographic peak the rescorer selected for this candidate (#7). + // 0 = the up-front apex; > 0 = a promoted alternate peak won. + f("selected_peak_rank", DataType::Int32), + ])); + let protein: ArrayRef = Arc::new(StringArray::from(c.protein)); + let q: ArrayRef = Arc::new(Float64Array::from(c.psm_q)); + let arrays: Vec = vec![ + Arc::new(UInt32Array::from(c.cid)), + Arc::new(StringArray::from(c.pform)), + Arc::new(Int32Array::from(c.charge)), + Arc::new(StringArray::from(c.label)), + protein.clone(), + Arc::new(UInt32Array::from(c.base)), + Arc::new(Float64Array::from(c.apex_rt)), + Arc::new(Float64Array::from(c.elution_lo)), + Arc::new(Float64Array::from(c.elution_hi)), + Arc::new(Float64Array::from(c.scores)), + q.clone(), + Arc::new(Float64Array::from(c.peptide_q)), + protein, + Arc::new(Float64Array::from(c.pg_q)), + q.clone(), + Arc::new(Float64Array::from(c.prelim)), + Arc::new(UInt32Array::from(c.source)), + Arc::new(Float64Array::from(c.run_psm_q)), + q, + Arc::new(Float64Array::from(c.precursor_q)), + Arc::new(Int32Array::from(c.peak_rank)), + ]; + let batch = RecordBatch::try_new(schema.clone(), arrays) + .with_context(|| format!("building the scored record batch for {path}"))?; + let mut w = mumdia_io::table::BatchWriter::new(path, schema)?; + w.write(&batch)?; + w.close() +} + /// Reject a concatenation whose feature companions differ in either identity or /// ordered feature columns. Both checks are intentional: the ID catches a /// declared contract mismatch, while the explicit column comparison protects @@ -1577,6 +1705,99 @@ b assert_eq!(got, vec!["b".to_string(), "c".to_string()]); } + #[test] + fn merge_col_moves_the_first_input_and_appends_the_rest() { + // One input: the destination IS the reader's vector, not a copy of it. Checked on + // the pointer, because "moved rather than copied" is the whole claim. + let src = vec![String::from("a"), String::from("b")]; + let ptr = src.as_ptr(); + let mut dst: Vec = Vec::new(); + merge_col(&mut dst, src, 2); + assert_eq!(dst.as_ptr(), ptr); + assert_eq!(dst, vec!["a".to_string(), "b".to_string()]); + // Several inputs: the concatenation is in input order, and the reserve from the + // footer row count means the appends do not realloc. + let mut dst: Vec = Vec::new(); + merge_col(&mut dst, vec![1, 2], 5); + let after_first = dst.as_ptr(); + merge_col(&mut dst, vec![3], 5); + merge_col(&mut dst, vec![4, 5], 5); + assert_eq!(dst, vec![1, 2, 3, 4, 5]); + assert_eq!( + dst.as_ptr(), + after_first, + "the reserve should have sufficed" + ); + // An empty first input must not lose the rows of the second. + let mut dst: Vec = Vec::new(); + merge_col(&mut dst, Vec::new(), 2); + merge_col(&mut dst, vec![7, 8], 2); + assert_eq!(dst, vec![7, 8]); + } + + #[test] + fn a_shared_column_array_writes_the_same_parquet_as_two_copies() { + use arrow::array::{ArrayRef, Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use std::sync::Arc; + + // `write_scored_table` hands the same Arrow array to `protein`/`protein_group` and + // to the three aliases of the pooled q. This asserts the file that produces is + // byte-for-byte the file `write_table` produced from separate, equal vectors -- + // the output-equality claim of that change, on the mechanism it rests on. + let dir = std::env::temp_dir().join("mumdia_scored_share_test"); + std::fs::create_dir_all(&dir).unwrap(); + let pid = std::process::id(); + let a = dir.join(format!("copies_{pid}.parquet")); + let b = dir.join(format!("shared_{pid}.parquet")); + let prot: Vec = (0..2000) + .map(|i| format!("sp|P{:05}|PROT_HUMAN", i % 700)) + .collect(); + let q: Vec = (0..2000).map(|i| (i % 97) as f64 / 97.0).collect(); + let rank: Vec = (0..2000).map(|i| i % 3).collect(); + write_table( + a.to_str().unwrap(), + vec![ + Col::Str("protein".into(), prot.clone()), + Col::F64("q_value".into(), q.clone()), + Col::Str("protein_group".into(), prot.clone()), + Col::F64("global_q_value".into(), q.clone()), + Col::I32("selected_peak_rank".into(), rank.clone()), + ], + ) + .unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("protein", DataType::Utf8, false), + Field::new("q_value", DataType::Float64, false), + Field::new("protein_group", DataType::Utf8, false), + Field::new("global_q_value", DataType::Float64, false), + Field::new("selected_peak_rank", DataType::Int32, false), + ])); + let protein: ArrayRef = Arc::new(StringArray::from(prot)); + let qq: ArrayRef = Arc::new(Float64Array::from(q)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + protein.clone(), + qq.clone(), + protein, + qq, + Arc::new(Int32Array::from(rank)), + ], + ) + .unwrap(); + let mut w = mumdia_io::table::BatchWriter::new(b.to_str().unwrap(), schema).unwrap(); + w.write(&batch).unwrap(); + assert_eq!(w.close().unwrap(), 2000); + assert_eq!( + std::fs::read(&a).unwrap(), + std::fs::read(&b).unwrap(), + "sharing one array across duplicate columns must not change the parquet" + ); + } + #[test] fn a_one_sided_population_is_refused_with_both_counts_named() { assert!(require_both_labels(10, 5).is_ok()); From 6a605c4c3db3c40974b8e9bebd4bb28b24e92ea0 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:26:07 +0200 Subject: [PATCH 11/35] perf(io): parse a parquet footer once per handle, and write a table in chunks Four changes to the Arrow/Parquet table layer, all of them about how many live heap blocks and how much repeated work one band of a grouped search costs. No artifact changes: the scalar write path is asserted byte-identical to the one it replaces, and every new reader is asserted equal to the reader it stands beside. 1. `TableFile` holds its parsed footer (`ArrowReaderMetadata`) and every reader it builds reuses it. `ParquetRecordBatchReaderBuilder::try_new` re-reads and re-parses the footer on every call, and it was being called once per `open`, once per `row_group_stats` and once per typed getter, so reading nine columns of one band cost ten parses. Measured on a 13,000-row-group, six-column, 1.3-million-row file (the shape of a library fragment table): 26.6 ms per parse against 118 ms to read a whole f64 column, so each getter paid about a fifth of its own cost again for a footer it already had. `row_group_stats` now touches no file at all. `TableFile::span(first_row, n_rows)` is the span read against already-parsed metadata: `open_rows` is now `open` + `span`, and a caller that cuts one library into many bands can open once and span per band instead of re-parsing per band per stage. Spans do not nest, so a span's rows are always file rows. `span_of_an_open_handle_matches_open_rows` asserts a span equals `open_rows` column for column and equals the whole table's slice, over an empty span, the whole file, inside one row group, across several and the last row. 2. `str_eq(name, value)` on both read paths: a string column as a boolean equality test. `str` returns one `String` per row, which for `label` on a 203-million-row precursor table is 203 million live allocations and 4.9 GB of `String` spine to carry one bit per row. No caller changed; the test asserts `str_eq` is exactly `str(...).iter().map(|s| s == value)`, error for error, including the NULL refusal. `str_flat` is the same idea for the high-cardinality columns where an equality test will not do (`peptidoform`, `protein`): one buffer plus offsets, the counterpart of the `list_f32_flat` that is already here. 3. `write_table` encodes in 65,536-row chunks instead of building one `RecordBatch` for the whole table. The full batch was a second complete Arrow copy of every column beside the source vectors, one very large heap block per column; mimalloc maps those individually, which is what the per-process mapping limit counts. The chunk size divides the writer's default 1,048,576-row row group and is a multiple of its 1,024-value write batch, so the row groups and the pages fall where they fell: `write_table_matches_one_batch_byte_for_byte_on_scalars` compares the new output with the old code path (`cols_to_batch` + `write_batches`) byte for byte over 196,615 rows, and again at 0, 1 and 1,000 rows. List columns hold a variable number of leaf values per row, so their pages are the encoder's business; there the test asserts the rows and the row-group count (measured byte-identical too, at 0-6 and 0-136 values per row, but not asserted). Validation moved into `validate_cols` so the whole column set is still checked once, before anything is written, with the messages it always had. Risk to output equality: the scalar write path is byte-asserted; the list write path is row-asserted and was measured byte-identical. The read changes cannot alter values -- the metadata is the same metadata, and `str_eq` / `str_flat` / `span` are new entry points with equality tests against the existing ones. One behavioural difference worth naming: a `TableFile` is now a snapshot of the footer as of `open`, where before each getter re-read it. Readers still open the file by path per call, so a file replaced under an open handle was already inconsistent between `row_group_stats` and a getter; a run does not rewrite an artifact it is reading. cargo fmt, cargo clippy -p mumdia-io --all-targets -D warnings, cargo test --workspace all pass. Co-Authored-By: Claude Opus 5 (1M context) --- rust/mumdia/crates/mumdia-io/src/table.rs | 684 ++++++++++++++++++++-- 1 file changed, 646 insertions(+), 38 deletions(-) diff --git a/rust/mumdia/crates/mumdia-io/src/table.rs b/rust/mumdia/crates/mumdia-io/src/table.rs index 11467c59..e440aead 100644 --- a/rust/mumdia/crates/mumdia-io/src/table.rs +++ b/rust/mumdia/crates/mumdia-io/src/table.rs @@ -15,9 +15,13 @@ use arrow::array::{ }; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; -use parquet::arrow::arrow_reader::{ParquetRecordBatchReaderBuilder, RowSelection, RowSelector}; +use parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection, + RowSelector, +}; use parquet::arrow::ArrowWriter; use parquet::basic::{Compression, LogicalType}; +use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::WriterProperties; use parquet::file::statistics::Statistics; @@ -163,17 +167,20 @@ impl Col { } } -/// Validate a set of typed columns and turn them into one Arrow schema + record batch. -/// Shared by [`write_table`] (one batch = the whole table) and [`TableWriter`] (one batch -/// per chunk), so both write paths declare a column identically. -fn cols_to_batch(path: &str, cols: Vec) -> Result<(Arc, RecordBatch)> { +/// Validate a set of typed columns: at least one column, unique names, equal lengths. +/// Returns the row count. +/// +/// Split out of [`cols_to_batch`] so [`write_table`] can check the WHOLE column set once, +/// before it starts consuming the columns chunk by chunk, and still report exactly the +/// errors it always reported. +fn validate_cols(path: &str, cols: &[Col]) -> Result { if cols.is_empty() { return Err(anyhow!("write_table: no columns for {path}")); } // Reject duplicate column names: Arrow allows them but readers resolve a // name to the first match, silently hiding the second column. let mut names = std::collections::HashSet::new(); - for c in &cols { + for c in cols { if !names.insert(c.name()) { return Err(anyhow!( "write_table: duplicate column '{}' for {path}", @@ -182,7 +189,7 @@ fn cols_to_batch(path: &str, cols: Vec) -> Result<(Arc, RecordBatch } } let nrows = cols[0].len(); - for c in &cols { + for c in cols { if c.len() != nrows { return Err(anyhow!( "write_table: column '{}' has {} rows, expected {}", @@ -192,6 +199,14 @@ fn cols_to_batch(path: &str, cols: Vec) -> Result<(Arc, RecordBatch )); } } + Ok(nrows) +} + +/// Validate a set of typed columns and turn them into one Arrow schema + record batch. +/// Shared by [`write_table`] (one batch per chunk) and [`TableWriter`] (one batch per +/// caller chunk), so both write paths declare a column identically. +fn cols_to_batch(path: &str, cols: Vec) -> Result<(Arc, RecordBatch)> { + validate_cols(path, &cols)?; let fields: Vec = cols.iter().map(|c| c.field()).collect(); let schema = Arc::new(Schema::new(fields)); // Consume the columns so each Vec is moved into its Arrow array rather than @@ -210,19 +225,87 @@ fn snappy_props(row_group_rows: Option) -> WriterProperties { b.build() } +/// Rows per internal chunk in [`write_table`]. +/// +/// A multiple of the parquet writer's 1,024-value write batch, so the encoder sees the +/// same sequence of mini-batches it would see from one big batch, and a divisor of the +/// writer's default 1,048,576-row row group, so the row-group boundaries are where they +/// were. 65,536 rows is ~0.5 MB for an f64 column and ~1.5 MB for a string column. +const WRITE_TABLE_CHUNK_ROWS: usize = 1 << 16; + +/// One column of [`write_table`] mid-flight: the source `Vec` turned into an iterator so +/// each chunk MOVES its rows out of it. A `String` or an inner `Vec` is handed to the +/// chunk and freed with the chunk's Arrow arrays, rather than every column being copied +/// into Arrow in full while the source is still resident. +macro_rules! col_chunks { + ($($v:ident => $t:ty),+ $(,)?) => { + enum ColChunks { $($v(String, std::vec::IntoIter<$t>)),+ } + + impl ColChunks { + fn of(c: Col) -> ColChunks { + match c { $(Col::$v(n, v) => ColChunks::$v(n, v.into_iter())),+ } + } + + /// The next `k` rows as a `Col` of the same variant and name. Fewer than `k` + /// only when the column is exhausted, which `write_table` never asks for. + fn take(&mut self, k: usize) -> Col { + match self { + $(ColChunks::$v(n, it) => Col::$v(n.clone(), it.by_ref().take(k).collect())),+ + } + } + } + }; +} + +col_chunks!( + I64 => i64, + I32 => i32, + U32 => u32, + F64 => f64, + F32 => f32, + Bool => bool, + Str => String, + OptF64 => Option, + OptF32 => Option, + OptI32 => Option, + OptStr => Option, + ListF32 => Vec, + ListF64 => Vec, + LargeListF32 => Vec, +); + /// Write columns to a Parquet file. Returns the row count. All columns must /// share the same length. +/// +/// The columns are validated as a set, then encoded in [`WRITE_TABLE_CHUNK_ROWS`] chunks +/// through [`TableWriter`]. Building ONE record batch for the whole table first, as this +/// used to, meant a second full Arrow copy of every column existed beside the source +/// vectors: on a wide artifact that is a second copy of the whole table, and each column +/// is one very large heap block, which on a grouped run is what the per-process mapping +/// limit counts. Chunking holds one chunk's Arrow arrays instead. +/// +/// Rows are written in order and row groups still fall where the parquet writer's default +/// 1,048,576-row maximum puts them (the writer accumulates across chunks; the chunk size +/// divides it), so the file is the same file: +/// `write_table_matches_one_batch_byte_for_byte_on_scalars` asserts exactly that, and +/// `..._row_for_row_on_lists` asserts the rows where the encoder's page boundaries are its +/// own business. pub fn write_table(path: &str, cols: Vec) -> Result { - let (schema, batch) = cols_to_batch(path, cols)?; - let nrows = batch.num_rows(); - let target = AtomicPath::new(path)?; - let file = std::fs::File::create(target.tmp()) - .with_context(|| format!("creating {}", target.tmp().display()))?; - let mut writer = ArrowWriter::try_new(file, schema, Some(snappy_props(None)))?; - writer.write(&batch)?; - writer.close()?; - target.publish()?; - Ok(nrows as u64) + let nrows = validate_cols(path, &cols)?; + let mut chunks: Vec = cols.into_iter().map(ColChunks::of).collect(); + let mut w = TableWriter::new(path); + let mut written = 0usize; + loop { + let k = (nrows - written).min(WRITE_TABLE_CHUNK_ROWS); + // A zero-row table still writes its one empty chunk, which fixes the schema, so + // an empty artifact keeps its columns. + w.write_cols(chunks.iter_mut().map(|c| c.take(k)).collect())?; + written += k; + if written >= nrows { + break; + } + } + w.close() } /// Incremental typed writer: the chunked counterpart of [`write_table`]. Feed `Vec` @@ -681,6 +764,18 @@ impl Table { Ok(out) } + /// A string column as `column == value`, row by row: exactly + /// `self.str(name)?.iter().map(|s| s == value).collect()`, without building the + /// `String`s. See [`TableFile::str_eq`] for what that costs on a two-valued column. + pub fn str_eq(&self, name: &str, value: &str) -> Result> { + let i = self.idx(name)?; + let mut out = Vec::with_capacity(self.nrows); + for b in &self.batches { + push_str_eq(&mut out, b.column(i), name, value)?; + } + Ok(out) + } + pub fn opt_f64(&self, name: &str) -> Result>> { let i = self.idx(name)?; let mut out = Vec::with_capacity(self.nrows); @@ -864,6 +959,44 @@ fn push_str(out: &mut Vec, col: &ArrayRef, name: &str) -> Result<()> { Ok(()) } +/// Flat layout: one `String` of concatenated values plus `offsets` (row `r` is +/// `data[offsets[r]..offsets[r + 1]]`). One allocation instead of one per row. Same null +/// policy as `push_str`: a NULL in a required column is an error. +fn push_str_flat( + offsets: &mut Vec, + data: &mut String, + col: &ArrayRef, + name: &str, +) -> Result<()> { + let a: &StringArray = downcast(col, name, "utf8")?; + if offsets.is_empty() { + offsets.push(0); + } + for k in 0..a.len() { + if a.is_null(k) { + // `offsets` holds one entry per row pushed plus the leading 0, so this is the + // same absolute row `push_str` would name. + return Err(reject_null(name, offsets.len() - 1)); + } + data.push_str(a.value(k)); + offsets.push(data.len()); + } + Ok(()) +} + +/// The equality test of `push_str`'s output against `value`, without the `String` per row. +/// Same null policy: a NULL in a required column is an error, not `false`. +fn push_str_eq(out: &mut Vec, col: &ArrayRef, name: &str, value: &str) -> Result<()> { + let a: &StringArray = downcast(col, name, "utf8")?; + for k in 0..a.len() { + if a.is_null(k) { + return Err(reject_null(name, out.len())); + } + out.push(a.value(k) == value); + } + Ok(()) +} + fn push_opt_f64(out: &mut Vec>, col: &ArrayRef, name: &str) -> Result<()> { let a: &Float64Array = downcast(col, name, "f64")?; for k in 0..a.len() { @@ -1055,6 +1188,19 @@ pub struct TableFile { pub schema: Arc, /// Row count from the parquet footer, or the span's row count for a partial open. pub nrows: usize, + /// The parsed parquet footer, held for the life of the handle. + /// + /// Every reader this handle builds reuses it instead of re-reading and re-parsing the + /// footer from disk, which is what [`ParquetRecordBatchReaderBuilder::try_new`] does + /// each time it is called. The footer is not small on the library tables -- a + /// 1.7-billion-row fragment table has ~13,000 row groups, and one + /// `ColumnChunkMetaData` with statistics per column per group -- and it was being + /// parsed once per open, once per `row_group_stats` and once per typed getter, so + /// reading nine columns of one band cost ten parses. Measured on a 13,000-row-group, + /// six-column, 1.3-million-row file: 26.6 ms per parse against 118 ms for the whole f64 + /// column, so every getter was paying about a fifth of its own cost again for a footer + /// it already had. Cloning it into a span ([`TableFile::span`]) is an `Arc` clone. + meta: ArrowReaderMetadata, /// `Some` when the table was opened on a row span ([`TableFile::open_rows`]): the row /// groups that cover the span and the selection that trims them to it. Every getter and /// batch reader applies it, so a partial table behaves as a smaller file. @@ -1092,17 +1238,20 @@ impl RowGroupStats { } impl TableFile { - /// Open `path` and read its footer. No column data is decoded. + /// Open `path` and read its footer. No column data is decoded. The parsed footer is + /// kept on the handle, so every getter, batch reader and [`TableFile::span`] taken + /// from it reads the file without parsing the footer again. pub fn open(path: &str) -> Result { let file = std::fs::File::open(path).with_context(|| format!("opening {path}"))?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file) + let meta = ArrowReaderMetadata::load(&file, ArrowReaderOptions::default()) .with_context(|| format!("reading parquet footer {path}"))?; - let nrows = builder.metadata().file_metadata().num_rows().max(0) as usize; - let schema = builder.schema().clone(); + let nrows = meta.metadata().file_metadata().num_rows().max(0) as usize; + let schema = meta.schema().clone(); Ok(TableFile { path: path.to_string(), schema, nrows, + meta, selection: None, }) } @@ -1116,10 +1265,26 @@ impl TableFile { /// are ordered by m/z with row-aligned ids, so an isolation-window group is a row span, /// and the group's search never holds the rest of the library. pub fn open_rows(path: &str, first_row: usize, n_rows: usize) -> Result { - let file = std::fs::File::open(path).with_context(|| format!("opening {path}"))?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file) - .with_context(|| format!("reading parquet footer {path}"))?; - let meta = builder.metadata(); + TableFile::open(path)?.span(first_row, n_rows) + } + + /// The row span `[first_row, first_row + n_rows)` of an ALREADY-OPEN handle, reusing + /// the footer this handle has already parsed. [`TableFile::open_rows`] is this plus the + /// open, and is the right call for a single span; this is the right call for a grouped + /// search, which cuts one library into many bands and would otherwise re-read and + /// re-parse a ~13,000-row-group footer for each band of each stage. + /// + /// Spans do not nest: take every span from the whole-file handle. + pub fn span(&self, first_row: usize, n_rows: usize) -> Result { + if self.selection.is_some() { + anyhow::bail!( + "TableFile::span: {} is already a row span; take spans from the whole-file \ + handle so their rows are file rows", + self.path + ); + } + let path = &self.path; + let meta: &ParquetMetaData = self.meta.metadata(); let total = meta.file_metadata().num_rows().max(0) as usize; if first_row.saturating_add(n_rows) > total { anyhow::bail!( @@ -1145,9 +1310,10 @@ impl TableFile { } let skip_after = covered.saturating_sub(skip_before + n_rows); Ok(TableFile { - path: path.to_string(), - schema: builder.schema().clone(), + path: self.path.clone(), + schema: self.schema.clone(), nrows: n_rows, + meta: self.meta.clone(), selection: Some(RowSpan { row_groups, skip_before, @@ -1158,15 +1324,15 @@ impl TableFile { } /// Row count and min/max statistics of a numeric column per row group, in file order, - /// from the footer alone. Integer and floating columns are reported as f64; unsigned - /// integers are read back through their logical type so a value above `i32::MAX` is not - /// returned as its two's-complement image. + /// from the footer this handle already holds -- no file access at all. Integer and + /// floating columns are reported as f64; unsigned integers are read back through their + /// logical type so a value above `i32::MAX` is not returned as its two's-complement + /// image. + /// + /// On a span handle these are the statistics of the WHOLE file, as they always were: + /// the caller plans a span from them. pub fn row_group_stats(&self, name: &str) -> Result> { - let file = - std::fs::File::open(&self.path).with_context(|| format!("opening {}", self.path))?; - let builder = ParquetRecordBatchReaderBuilder::try_new(file) - .with_context(|| format!("reading parquet footer {}", self.path))?; - let meta = builder.metadata(); + let meta: &ParquetMetaData = self.meta.metadata(); let mut out = Vec::with_capacity(meta.num_row_groups()); for i in 0..meta.num_row_groups() { let rg = meta.row_group(i); @@ -1245,9 +1411,11 @@ impl TableFile { pub fn batches(&self, columns: Option<&[&str]>, batch_size: usize) -> Result { let file = std::fs::File::open(&self.path).with_context(|| format!("opening {}", self.path))?; - let mut builder = ParquetRecordBatchReaderBuilder::try_new(file) - .with_context(|| format!("reading parquet {}", self.path))? - .with_batch_size(batch_size.max(1)); + // The footer this handle parsed at `open`, not a fresh parse: a typed getter is one + // call to this, and a stage reads a dozen columns. + let mut builder = + ParquetRecordBatchReaderBuilder::new_with_metadata(file, self.meta.clone()) + .with_batch_size(batch_size.max(1)); if let Some(span) = &self.selection { // The selection counts rows of the SELECTED row groups only, front to back, so // it is skip / take / skip over exactly the groups named here. @@ -1364,6 +1532,24 @@ impl TableFile { Ok(out) } + /// A string column as `column == value`, row by row: exactly + /// `self.str(name)?.iter().map(|s| s == value).collect()`, without building the + /// `String`s. + /// + /// For the two-valued columns -- `label` is "target" or "decoy" on every row of a + /// 203-million-row precursor table -- `str` returns one `String` per row: a 24-byte + /// spine plus its own heap block, so about 4.9 GB of spine and 203 million live + /// allocations to carry one bit per row. This returns the bit: one byte per row and one + /// allocation. The null policy is `str`'s, so a NULL is still refused rather than + /// quietly reading as `false`. + pub fn str_eq(&self, name: &str, value: &str) -> Result> { + let mut out = Vec::with_capacity(self.nrows); + for b in self.column(name, SCALAR_BATCH_ROWS)? { + push_str_eq(&mut out, b?.column(0), name, value)?; + } + Ok(out) + } + pub fn opt_f64(&self, name: &str) -> Result>> { let mut out = Vec::with_capacity(self.nrows); for b in self.column(name, SCALAR_BATCH_ROWS)? { @@ -1381,6 +1567,27 @@ impl TableFile { Ok(out) } + /// Read a string column into one buffer plus `nrows + 1` byte offsets (row `r` is + /// `data[offsets[r]..offsets[r + 1]]`, always a char boundary because whole values are + /// concatenated): one allocation for the column instead of one `String` per row. + /// + /// The counterpart of [`TableFile::list_f32_flat`] for the high-cardinality string + /// columns, where an equality test will not do. `peptidoform` and `protein` on a + /// 203-million-row precursor library are 203 million live allocations and 4.9 GB of + /// `String` spine through [`TableFile::str`] before a byte of text; here they are one + /// `Vec` and one `String`. Same null policy as `str`. + pub fn str_flat(&self, name: &str) -> Result<(Vec, String)> { + let mut offsets = Vec::with_capacity(self.nrows + 1); + let mut data = String::new(); + for b in self.column(name, SCALAR_BATCH_ROWS)? { + push_str_flat(&mut offsets, &mut data, b?.column(0), name)?; + } + if offsets.is_empty() { + offsets.push(0); + } + Ok((offsets, data)) + } + /// Read an f32 list column into one flat values buffer plus `nrows + 1` offsets /// (row `r` is `values[offsets[r]..offsets[r + 1]]`): one allocation for the whole /// column instead of one per row, which is what makes a chromatogram table with tens @@ -1538,6 +1745,186 @@ mod projection_tests { } } +#[cfg(test)] +mod write_chunking_tests { + use super::*; + + fn tmp(name: &str) -> String { + let dir = std::env::temp_dir().join(format!("mumdia_table_chunk_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + dir.join(name).to_str().unwrap().to_string() + } + + /// Scalars and a string column: one value per row per column, so the parquet encoder + /// sees the same 1,024-value mini-batches whether the rows arrive in one batch or in + /// chunks that are multiples of 1,024. + fn scalar_cols(n: usize) -> Vec{ + vec![ + Col::U32("id".into(), (0..n as u32).collect()), + Col::F64("mz".into(), (0..n).map(|i| i as f64 * 1.5 - 3.0).collect()), + Col::F32("irt".into(), (0..n).map(|i| (i % 977) as f32).collect()), + Col::Bool("flag".into(), (0..n).map(|i| i % 3 == 0).collect()), + Col::Str( + "label".into(), + (0..n) + .map(|i| if i % 2 == 0 { "target" } else { "decoy" }.to_string()) + .collect(), + ), + Col::OptF64( + "cal".into(), + (0..n) + .map(|i| (i % 5 != 2).then_some(i as f64 * 0.25)) + .collect(), + ), + ] + } + + fn list_cols(n: usize) -> Vec{ + vec![ + Col::U32("id".into(), (0..n as u32).collect()), + Col::ListF32( + "trace".into(), + (0..n) + .map(|i| (0..(i % 7)).map(|k| k as f32 + i as f32).collect()) + .collect(), + ), + Col::LargeListF32( + "rt".into(), + (0..n) + .map(|i| (0..(i % 5)).map(|k| k as f32 * 0.5).collect()) + .collect(), + ), + ] + } + + /// The pre-change write path, verbatim: validate, build ONE record batch for the whole + /// table, write it with the same Snappy properties. + fn write_one_batch(path: &str, cols: Vec) { + let (schema, batch) = cols_to_batch(path, cols).unwrap(); + write_batches(path, schema, &[batch]).unwrap(); + } + + /// `write_table` now encodes in 65,536-row chunks. On scalar columns that must produce + /// the same FILE, not merely the same values: the chunk size divides the writer's + /// 1,048,576-row row group, so the row groups fall where they fell, and it is a + /// multiple of the 1,024-value write batch, so the pages do too. + #[test] + fn write_table_matches_one_batch_byte_for_byte_on_scalars() { + let n = 3 * WRITE_TABLE_CHUNK_ROWS + 7; // several chunks and a short last one + let chunked = tmp("scalars_chunked.parquet"); + let once = tmp("scalars_once.parquet"); + assert_eq!(write_table(&chunked, scalar_cols(n)).unwrap(), n as u64); + write_one_batch(&once, scalar_cols(n)); + assert_eq!( + std::fs::read(&chunked).unwrap(), + std::fs::read(&once).unwrap(), + "chunked write_table must produce the same parquet file as the single-batch write" + ); + // And the row groups are the writer's default, not one per chunk. + let meta = TableFile::open(&chunked).unwrap(); + assert_eq!( + meta.row_group_stats("id").unwrap().len(), + 1, + "{n} rows is one 1,048,576-row row group, chunked or not" + ); + assert_eq!(meta.nrows, n); + std::fs::remove_file(&chunked).ok(); + std::fs::remove_file(&once).ok(); + } + + /// List columns hold a variable number of leaf values per row, so a row-chunk boundary + /// need not fall on a 1,024-level mini-batch boundary and the data pages could in + /// principle be cut elsewhere. The contract there is the rows: same order, same values, + /// same row groups. (Measured on parquet 59.3.0 the files came out byte-identical here + /// too, at 0-6 and at 0-136 values per row, but that is the encoder's business, not a + /// promise this layer makes.) + #[test] + fn write_table_matches_one_batch_row_for_row_on_lists() { + let n = WRITE_TABLE_CHUNK_ROWS + 1_234; + let chunked = tmp("lists_chunked.parquet"); + let once = tmp("lists_once.parquet"); + write_table(&chunked, list_cols(n)).unwrap(); + write_one_batch(&once, list_cols(n)); + let (a, b) = (Table::read(&chunked).unwrap(), Table::read(&once).unwrap()); + assert_eq!(a.schema, b.schema); + assert_eq!(a.nrows, n); + assert_eq!(a.u32("id").unwrap(), b.u32("id").unwrap()); + assert_eq!(a.list_f32("trace").unwrap(), b.list_f32("trace").unwrap()); + assert_eq!(a.list_f32("rt").unwrap(), b.list_f32("rt").unwrap()); + assert_eq!( + TableFile::open(&chunked) + .unwrap() + .row_group_stats("id") + .unwrap() + .len(), + TableFile::open(&once) + .unwrap() + .row_group_stats("id") + .unwrap() + .len() + ); + std::fs::remove_file(&chunked).ok(); + std::fs::remove_file(&once).ok(); + } + + /// A table shorter than one chunk, and an empty one: both must still be the file the + /// single-batch path wrote, including the schema of a zero-row artifact. + #[test] + fn write_table_matches_one_batch_when_short_or_empty() { + for n in [0usize, 1, 1_000] { + let chunked = tmp(&format!("short_chunked_{n}.parquet")); + let once = tmp(&format!("short_once_{n}.parquet")); + assert_eq!(write_table(&chunked, scalar_cols(n)).unwrap(), n as u64); + write_one_batch(&once, scalar_cols(n)); + assert_eq!( + std::fs::read(&chunked).unwrap(), + std::fs::read(&once).unwrap(), + "{n} rows" + ); + let t = Table::read(&chunked).unwrap(); + assert_eq!(t.nrows, n); + assert_eq!( + t.column_names().len(), + 6, + "a zero-row table keeps its columns" + ); + std::fs::remove_file(&chunked).ok(); + std::fs::remove_file(&once).ok(); + } + } + + /// The validation is of the whole column set, before any chunk is encoded: a length + /// mismatch past the first chunk must still be refused, with the same message, and + /// must not leave a partial file behind. + #[test] + fn write_table_validates_the_whole_column_set_before_writing() { + let p = tmp("rejected.parquet"); + let n = WRITE_TABLE_CHUNK_ROWS + 10; + let mut cols = scalar_cols(n); + cols.push(Col::I32("short".into(), vec![1, 2, 3])); + let err = write_table(&p, cols).unwrap_err().to_string(); + assert!(err.contains("'short'") && err.contains("expected"), "{err}"); + assert!( + !std::path::Path::new(&p).exists(), + "a refused write must not publish a file" + ); + let dup = write_table( + &p, + vec![ + Col::U32("id".into(), vec![1]), + Col::U32("id".into(), vec![2]), + ], + ) + .unwrap_err() + .to_string(); + assert!(dup.contains("duplicate column 'id'"), "{dup}"); + assert!(write_table(&p, vec![]) + .unwrap_err() + .to_string() + .contains("no columns")); + } +} + #[cfg(test)] mod batch_writer_tests { use super::*; @@ -1819,6 +2206,227 @@ mod streaming_tests { assert!(TableWriter::new(&tmp("never.parquet")).close().is_err()); std::fs::remove_file(&p).ok(); } + + /// A span taken from an open handle must be the table `open_rows` opens, column for + /// column, and must be the matching slice of the whole table. That equality is what + /// lets a grouped search take every band from one handle and parse the footer once + /// instead of once per band per stage. + #[test] + fn span_of_an_open_handle_matches_open_rows() { + let p = tmp("spans.parquet"); + let n = 10_000; + let mut w = TableWriter::new(&p).with_row_group_rows(1_000); + w.write_cols(mixed_cols(n)).unwrap(); + w.close().unwrap(); + + let whole = TableFile::open(&p).unwrap(); + let all_id = whole.u32("id").unwrap(); + let all_name = whole.str("name").unwrap(); + let all_trace = whole.list_f32("trace").unwrap(); + // Empty, whole-file, inside one row group, across several, and the last row. + for (first, len) in [(0, 0), (0, n), (999, 1), (1_500, 3_000), (n - 1, 1)] { + let a = whole.span(first, len).unwrap(); + let b = TableFile::open_rows(&p, first, len).unwrap(); + assert_eq!(a.nrows, len); + assert_eq!(b.nrows, len); + assert_eq!(a.u32("id").unwrap(), b.u32("id").unwrap(), "{first}+{len}"); + assert_eq!(a.f64("mz").unwrap(), b.f64("mz").unwrap()); + assert_eq!(a.str("name").unwrap(), b.str("name").unwrap()); + assert_eq!(a.list_f32("trace").unwrap(), b.list_f32("trace").unwrap()); + // ... and the same rows as the whole table's slice. + assert_eq!(a.u32("id").unwrap(), all_id[first..first + len]); + assert_eq!(a.str("name").unwrap(), all_name[first..first + len]); + assert_eq!(a.list_f32("trace").unwrap(), all_trace[first..first + len]); + } + // Out of range is refused with the same message as open_rows. + let msg = |r: Result| r.err().expect("must be refused").to_string(); + assert_eq!( + msg(whole.span(n - 1, 2)), + msg(TableFile::open_rows(&p, n - 1, 2)) + ); + // Spans do not nest: a span's rows would otherwise mean two different things. + let part = whole.span(10, 20).unwrap(); + let err = msg(part.span(0, 1)); + assert!(err.contains("already a row span"), "{err}"); + // Statistics are the whole file's on either handle, which is how a caller plans + // a span from them. + assert_eq!( + part.row_group_stats("id").unwrap(), + whole.row_group_stats("id").unwrap() + ); + assert_eq!(whole.row_group_stats("id").unwrap().len(), 10); + std::fs::remove_file(&p).ok(); + } + + /// The handle carries its parsed footer across getters, spans and batch readers. If + /// anything in that plumbing went wrong it would show up as a wrong row count or a + /// wrong column, so read every getter twice from one handle and once through a span. + #[test] + fn one_handle_serves_repeated_reads() { + let p = tmp("reuse.parquet"); + let n = 3_000; + write_table(&p, mixed_cols(n)).unwrap(); + let f = TableFile::open(&p).unwrap(); + for _ in 0..3 { + assert_eq!(f.u32("id").unwrap().len(), n); + assert_eq!(f.f64("mz").unwrap().len(), n); + assert_eq!(f.str("name").unwrap().len(), n); + assert_eq!(f.list_f32("trace").unwrap().len(), n); + let mut seen = 0; + f.for_each_batch(Some(&["id", "mz"]), 512, |b| { + seen += b.num_rows(); + Ok(()) + }) + .unwrap(); + assert_eq!(seen, n); + } + assert_eq!(f.span(100, 50).unwrap().u32("id").unwrap()[0], 100); + std::fs::remove_file(&p).ok(); + } + + /// `str_eq` must be `str` plus a comparison, error for error: it is offered as a + /// drop-in for a caller that only asks whether a two-valued column equals a value. + #[test] + fn str_eq_is_the_string_comparison_without_the_strings() { + let p = tmp("labels.parquet"); + let n: usize = 5_000; + write_table( + &p, + vec![ + Col::U32("id".into(), (0..n as u32).collect()), + Col::Str( + "label".into(), + (0..n) + .map(|i| { + match i % 3 { + 0 => "target", + 1 => "decoy", + _ => "target_x", + } + .to_string() + }) + .collect(), + ), + Col::OptStr( + "note".into(), + (0..n) + .map(|i| (i % 7 != 4).then(|| format!("n{i}"))) + .collect(), + ), + ], + ) + .unwrap(); + let t = Table::read(&p).unwrap(); + let f = TableFile::open(&p).unwrap(); + let want: Vec = t + .str("label") + .unwrap() + .iter() + .map(|s| s == "target") + .collect(); + assert_eq!(want.iter().filter(|b| **b).count(), n.div_ceil(3)); + assert_eq!(t.str_eq("label", "target").unwrap(), want); + assert_eq!(f.str_eq("label", "target").unwrap(), want); + // A prefix is not a match, and a value the column does not hold matches nothing. + assert!(f.str_eq("label", "targe").unwrap().iter().all(|b| !b)); + assert!(f.str_eq("label", "TARGET").unwrap().iter().all(|b| !b)); + // A span reads its own rows only: rows 10..16, where 10 % 3 == 1 is the decoy. + assert_eq!( + f.span(10, 6).unwrap().str_eq("label", "decoy").unwrap(), + vec![true, false, false, true, false, false] + ); + // A NULL is refused exactly as `str` refuses it, and so are a missing column and + // the wrong type -- same wording on both read paths. + for (a, b) in [ + (f.str_eq("note", "n0"), f.str("note")), + (f.str_eq("nope", "x"), f.str("nope")), + (f.str_eq("id", "x"), f.str("id")), + ] { + assert_eq!( + a.unwrap_err().to_string(), + b.unwrap_err().to_string(), + "str_eq must fail exactly as str does" + ); + } + assert_eq!( + t.str_eq("note", "n0").unwrap_err().to_string(), + f.str_eq("note", "n0").unwrap_err().to_string() + ); + std::fs::remove_file(&p).ok(); + } + + /// `str_flat` must rebuild exactly what `str` returns, row for row, including empty + /// values and multi-byte ones, and refuse a NULL the same way. + #[test] + fn str_flat_rebuilds_the_string_column() { + let p = tmp("flat_str.parquet"); + let n = 10_000; // several SCALAR_BATCH_ROWS-independent batches of rows + write_table(&p, mixed_cols(n)).unwrap(); + let f = TableFile::open(&p).unwrap(); + let rows = f.str("name").unwrap(); + let (off, data) = f.str_flat("name").unwrap(); + assert_eq!(off.len(), n + 1); + assert_eq!(off[0], 0); + assert_eq!(*off.last().unwrap(), data.len()); + for (r, want) in rows.iter().enumerate() { + assert_eq!(&data[off[r]..off[r + 1]], want, "row {r}"); + } + // A span reads its own rows only, and a NULL is refused as `str` refuses it. + let part = f.span(3, 4).unwrap(); + let (o, d) = part.str_flat("name").unwrap(); + assert_eq!(o.len(), 5); + assert_eq!(&d[o[0]..o[1]], &rows[3]); + assert_eq!( + f.str_flat("note").unwrap_err().to_string(), + f.str("note").unwrap_err().to_string() + ); + assert_eq!( + f.str_flat("id").unwrap_err().to_string(), + f.str("id").unwrap_err().to_string() + ); + // Multi-byte values: the offsets are byte offsets, and every one of them must land + // on a char boundary or the slicing panics. + let u = tmp("flat_unicode.parquet"); + write_table( + &u, + vec![Col::Str( + "s".into(), + vec![ + String::new(), + "\u{03bc}\u{00b2}".into(), + "plain".into(), + "\u{1f9ea}\u{00e9}".into(), + ], + )], + ) + .unwrap(); + let tu = TableFile::open(&u).unwrap(); + let want = tu.str("s").unwrap(); + let (o, d) = tu.str_flat("s").unwrap(); + for (r, w) in want.iter().enumerate() { + assert_eq!(&d[o[r]..o[r + 1]], w, "row {r}"); + } + // An empty table still returns the leading offset. + let e = tmp("flat_empty.parquet"); + write_table(&e, vec![Col::Str("s".into(), vec![])]).unwrap(); + let (o, d) = TableFile::open(&e).unwrap().str_flat("s").unwrap(); + assert_eq!((o, d.as_str()), (vec![0], "")); + std::fs::remove_file(&p).ok(); + std::fs::remove_file(&u).ok(); + std::fs::remove_file(&e).ok(); + } + + /// A grouped search hands band handles to worker threads. Holding the parsed footer on + /// the handle must not have taken that away. + #[test] + fn table_handles_stay_send_and_sync() { + fn assert_send_sync() {} + assert_send_sync::(); + assert_send_sync::
(); + // A reader is moved to the thread that drains it, not shared. + fn assert_send() {} + assert_send::(); + } } #[cfg(test)] From 461a42af62194843b9e390b5687d7d3c62ef0328 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:30:25 +0200 Subject: [PATCH 12/35] perf(quant): flat chromatogram store, one union map, one window computation quant was the last consumer of the chromatogram table still holding it as `Vec>` plus a `String` per row: three heap blocks per chromatogram row, on top of four to six more per candidate. That is the shape that kills a grouped search at ~290 GB resident with 1.7 TB free, because the process runs out of per-process memory MAPPINGS (1,048,576) rather than out of memory, and `--out-peak-bounds` bypasses the accepted-candidate filter so the whole table is held that way. What changed, in the order of the findings: 1. `ChromStore` gives the table the flat treatment `features` already gives it: one buffer for intensities, one for RT axes with the axis shared between the rows of a candidate that sample the same grid (every row under extract's window grid), interned fragment names, and flat per-row index arrays. The MS1 isotope pseudo-traces are dropped at load rather than stored and skipped later, after the rt/intensity length check, which stays a guard on the table. Live heap blocks per row go 3 -> 0. Derived from the layout at a 60-sample trace and 8 fragments per candidate: ~624 B/row -> ~314 B/row, before allocator per-block overhead, which the figure does not count. 2. `CandIndex` replaces `BTreeMap>` (a node and a `Vec` per candidate) with one grouped row array, and `win` / `applied_win` / the areas become flat arrays indexed by candidate position. The per-candidate `Vec` of areas, the `(area, predicted)` pairs and the `(name, area)` pairs are gone: rayon writes the areas into disjoint sub-slices of one buffer, so the positional pairing the serial fold used to assert is now structural. Blocks per candidate go 4-6 -> 0; what remains transient is bounded by the thread count, not the candidate count. 3. `peak_window` built the union of the RT axes by inserting every sample into TWO BTreeMaps keyed identically, one tree walk per sample each. It now either accumulates directly into the shared axis (the common case) or sorts the samples once by RT bit pattern. Both reduce each RT's f64 sum in the same order the tree did -- row order, then sample order -- and the sort is stable for that reason. 4. The fixed-window indices were computed twice per fragment row, once for the area and once to report the bounds; they are computed once. The nearest-sample search inside them linear-scanned the whole trace, which is O(trace) where the integration is O(window); it binary-searches when the store says every axis is non-decreasing and NaN-free, and falls back to the scan otherwise. 5. `select_fragment_areas` collected and SORTED a candidate's areas once per scored row mapping to it; memoised per candidate. 6. The cross-run combine cloned every feature vector into the protein matrix and again into each of the two sibling levels. `lfq_profile` now takes anything that borrows as `[Option]`, so all three borrow from one owner. The protein-group key is no longer allocated per row in either the rollup or the combine. Risk to output equality: the arithmetic and every reduction order are unchanged, and the equality claims are tested rather than argued. `peak_window_both` runs every existing `peak_window` test through BOTH union constructions and asserts they are bit-identical; `fixed_window_indices_both` runs the window tests through both nearest-sample searches; `nearest_index_binary_search_matches_the_scan` covers the tie-breaks the scan's strict `<` produced (duplicate RTs below, at and above the target, exact midpoints, targets off both ends); `a_repeated_rt_sample_merges_into_one_profile_point` pins the merge the BTreeMap did; `the_candidate_index_groups_exactly_as_the_map_it_replaces` compares the grouping against the map for a grouped, an interleaved and a single-row table; `borrowed_rows_give_the_same_profile_as_owned_rows` pins the LFQ core. The residual risk is the shared-axis shortcut in `peak_window`, which is taken only for an axis checked strictly increasing and non-negative when it is stored, and the binary search, which is taken only when every stored axis is non-decreasing and non-negative; either property failing sends the code back to the merge and the scan. Not run end to end on a real chromatogram artifact. Not done: streaming quant per candidate rather than materialising the run's chromatograms (finding 7). It needs each candidate's rows to be contiguous in the table, which this stage does not currently require, and it doubles the parquet read on the normal path, where the accepted-candidate filter already keeps only a few percent of the table. Reported rather than attempted. Co-Authored-By: Claude Fable 5.1 --- rust/mumdia/crates/mumdia/src/quant_lfq.rs | 31 +- rust/mumdia/crates/mumdia/src/stages/quant.rs | 1259 +++++++++++++---- 2 files changed, 1047 insertions(+), 243 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/quant_lfq.rs b/rust/mumdia/crates/mumdia/src/quant_lfq.rs index 170001e8..39fe802f 100644 --- a/rust/mumdia/crates/mumdia/src/quant_lfq.rs +++ b/rust/mumdia/crates/mumdia/src/quant_lfq.rs @@ -82,7 +82,13 @@ fn solve_fixed(l: &[Vec], c: &[f64]) -> Vec { /// feature-sharing graph is fit by least squares to the pairwise median log /// ratios, then anchored so the component preserves its measured total intensity; /// isolated samples fall back to their column sum. -pub fn lfq_profile(mat: &[Vec>], n_samples: usize) -> Vec { +/// +/// The rows are taken as anything that borrows as `[Option]`, so a caller that +/// already owns the per-feature vectors (the cross-run combine, which rolls the same +/// features up at protein, precursor and peptide level) passes references rather than +/// three clones of every feature vector. Read-only: the arithmetic and the iteration +/// order over rows are unchanged. +pub fn lfq_profile]>>(mat: &[R], n_samples: usize) -> Vec { let mut out = vec![0.0f64; n_samples]; if n_samples == 0 { return out; @@ -90,6 +96,7 @@ pub fn lfq_profile(mat: &[Vec>], n_samples: usize) -> Vec { let mut colsum = vec![0.0f64; n_samples]; let mut colcount = vec![0usize; n_samples]; for row in mat { + let row = row.as_ref(); for s in 0..n_samples { if let Some(v) = row.get(s).copied().flatten() { if v > 0.0 { @@ -109,6 +116,7 @@ pub fn lfq_profile(mat: &[Vec>], n_samples: usize) -> Vec { for b in (a + 1)..n_samples { let mut d = Vec::new(); for row in mat { + let row = row.as_ref(); if let (Some(va), Some(vb)) = (row.get(a).copied().flatten(), row.get(b).copied().flatten()) { @@ -189,14 +197,14 @@ pub fn lfq_profile(mat: &[Vec>], n_samples: usize) -> Vec { } /// MaxLFQ (Cox et al. 2014): peptide-level cross-sample profile. -pub fn maxlfq(peptides_by_sample: &[Vec>], n_samples: usize) -> Vec { +pub fn maxlfq]>>(peptides_by_sample: &[R], n_samples: usize) -> Vec { lfq_profile(peptides_by_sample, n_samples) } /// directLFQ (Ammar et al. 2023) core: the same ratio-alignment applied at the /// ion/fragment level rather than the peptide level (the caller passes an /// ion-by-sample matrix). -pub fn directlfq(ions_by_sample: &[Vec>], n_samples: usize) -> Vec { +pub fn directlfq]>>(ions_by_sample: &[R], n_samples: usize) -> Vec { lfq_profile(ions_by_sample, n_samples) } @@ -264,6 +272,23 @@ mod tests { ); } + #[test] + fn borrowed_rows_give_the_same_profile_as_owned_rows() { + // The cross-run combine passes `&Vec>` rows so the same feature + // vector is not cloned once per rollup level. The profile must be bit-identical + // to the owned form it replaced. + let owned = vec![ + vec![Some(10.0), Some(20.0), None], + vec![None, Some(10.0), Some(20.0)], + vec![Some(3.0), None, Some(12.0)], + vec![Some(5.0), Some(10.0), Some(20.0)], + ]; + let borrowed: Vec<&Vec>> = owned.iter().collect(); + assert_eq!(lfq_profile(&owned, 3), lfq_profile(&borrowed, 3)); + let slices: Vec<&[Option]> = owned.iter().map(|v| v.as_slice()).collect(); + assert_eq!(lfq_profile(&owned, 3), lfq_profile(&slices, 3)); + } + #[test] fn disconnected_sample_falls_back_to_sum() { // s2 shares no peptide with s0/s1 -> its own column sum. diff --git a/rust/mumdia/crates/mumdia/src/stages/quant.rs b/rust/mumdia/crates/mumdia/src/stages/quant.rs index c99823d3..00f45632 100644 --- a/rust/mumdia/crates/mumdia/src/stages/quant.rs +++ b/rust/mumdia/crates/mumdia/src/stages/quant.rs @@ -36,6 +36,301 @@ pub struct QuantParams<'a> { pub config_hash: &'a str, } +/// Interned fragment names. A chromatogram table has tens of millions of rows and only a +/// few dozen distinct fragment names (`y1`..`y30`, `b1`..), so a row carries a name id and +/// each string is stored once instead of the `String` per row this stage used to keep. +#[derive(Default)] +struct NameTab { + ids: HashMap, + names: Vec, +} + +impl NameTab { + fn intern(&mut self, s: &str) -> u32 { + if let Some(&i) = self.ids.get(s) { + return i; + } + let i = self.names.len() as u32; + self.names.push(s.to_string()); + self.ids.insert(s.to_string(), i); + i + } + + fn get(&self, id: u32) -> &str { + &self.names[id as usize] + } +} + +/// Marker for a row with no RT axis at all: extract emits an empty trace, not a +/// zero-filled one, for a predicted fragment that was never observed. +const NO_AXIS: u32 = u32::MAX; + +/// The fragment chromatogram rows quant keeps, in table order, in flat buffers. +/// +/// This is the same treatment `features` gives the same table ([`super::features`]'s +/// `ChromChunk`): one heap block per array instead of a `Vec` for the RT axis, a +/// `Vec` for the intensities and a `String` for the fragment name PER ROW. The +/// motivation is the number of live heap blocks rather than the bytes: a grouped search +/// died at ~290 GB resident with 1.7 TB free because the process had hit the kernel's +/// per-process memory-mapping limit (1,048,576) with one mapping per medium block, and +/// `--out-peak-bounds` bypasses the accepted-candidate filter and holds the whole table +/// this way. The RT axis is additionally shared between the rows of one candidate that +/// sample the same grid, which is every row of a candidate under extract's window-grid +/// mode, so the axis is stored once per candidate instead of once per fragment. +#[derive(Default)] +struct ChromStore { + names: NameTab, + /// Candidate id per row, in table order. + cid: Vec, + name_id: Vec, + pred: Vec, + /// Index into `axis_off`, or [`NO_AXIS`] for an empty trace. + axis_id: Vec, + /// Axis `a` is `axis_vals[axis_off[a]..axis_off[a + 1]]`. + axis_off: Vec, + axis_vals: Vec, + /// Axis `a` is strictly increasing and non-negative, so it carries no duplicate RT + /// and no NaN. [`peak_window`] needs that before it may treat the axis as the union + /// axis directly instead of merging the samples. + axis_strict: Vec, + /// Row `r`'s intensities are `int_vals[int_off[r]..int_off[r + 1]]`. + int_off: Vec, + int_vals: Vec, + /// First axis of the candidate being filled, so dedup only compares within it. + open_axis_lo: usize, + open_cid: Option, + /// Every stored axis is non-decreasing, non-negative and NaN-free, so the nearest + /// sample to an apex may be found by binary search rather than by scanning the whole + /// trace. False makes every search fall back to the scan, which is what this stage + /// always did. + rt_sorted: bool, +} + +impl ChromStore { + fn new() -> ChromStore { + ChromStore { + axis_off: vec![0], + int_off: vec![0], + rt_sorted: true, + ..Default::default() + } + } + + fn nrows(&self) -> usize { + self.cid.len() + } + + /// Store `rt` as an axis id, reusing an axis already stored for the open candidate + /// when the values are identical (the common case: one grid per candidate). + fn axis_for(&mut self, cid: u32, rt: &[f32]) -> u32 { + if self.open_cid != Some(cid) { + self.open_cid = Some(cid); + self.open_axis_lo = self.axis_off.len() - 1; + } + if rt.is_empty() { + return NO_AXIS; + } + for a in self.open_axis_lo..self.axis_off.len() - 1 { + let (lo, hi) = (self.axis_off[a], self.axis_off[a + 1]); + if hi - lo == rt.len() && self.axis_vals[lo..hi] == *rt { + return a as u32; + } + } + let nondecreasing = rt_is_sorted(rt); + self.rt_sorted &= nondecreasing; + self.axis_strict + .push(nondecreasing && rt.windows(2).all(|w| w[0] < w[1])); + self.axis_vals.extend_from_slice(rt); + self.axis_off.push(self.axis_vals.len()); + (self.axis_off.len() - 2) as u32 + } + + fn push(&mut self, cid: u32, name: &str, pred: f32, rt: &[f32], inten: &[f32]) { + let axis_id = self.axis_for(cid, rt); + let name_id = self.names.intern(name); + self.cid.push(cid); + self.name_id.push(name_id); + self.pred.push(pred); + self.axis_id.push(axis_id); + self.int_vals.extend_from_slice(inten); + self.int_off.push(self.int_vals.len()); + } + + fn axis(&self, id: u32) -> &[f32] { + if id == NO_AXIS { + return &[]; + } + let a = id as usize; + &self.axis_vals[self.axis_off[a]..self.axis_off[a + 1]] + } + + fn rt(&self, row: usize) -> &[f32] { + self.axis(self.axis_id[row]) + } + + fn inten(&self, row: usize) -> &[f32] { + &self.int_vals[self.int_off[row]..self.int_off[row + 1]] + } + + fn name(&self, row: usize) -> &str { + self.names.get(self.name_id[row]) + } + + /// Payload bytes by part, for [`crate::memlog::report`]. + fn mem_parts(&self) -> Vec<(&'static str, usize)> { + use crate::memlog::bytes_of; + vec![ + ("intensities", bytes_of(&self.int_vals)), + ("rt_axes", bytes_of(&self.axis_vals)), + ( + "row_index", + bytes_of(&self.cid) + + bytes_of(&self.name_id) + + bytes_of(&self.pred) + + bytes_of(&self.axis_id) + + bytes_of(&self.int_off), + ), + ( + "axis_index", + bytes_of(&self.axis_off) + bytes_of(&self.axis_strict), + ), + ] + } +} + +/// The fragment chromatogram rows grouped by candidate: candidate `ci` owns the rows +/// `rows[cand_off[ci]..cand_off[ci + 1]]` of [`ChromStore`], ascending by candidate id and +/// in table order within a candidate. That is exactly the iteration the +/// `BTreeMap>` this replaces produced, without a map node and a `Vec` per +/// candidate; the per-candidate results below (windows, areas) are flat arrays indexed by +/// the same `ci`. +struct CandIndex { + /// Distinct candidate ids, ascending. + cids: Vec, + cand_off: Vec, + rows: Vec, + /// `pred[slot]` for the row at `rows[slot]`, so the fragment-selection ranking reads + /// one contiguous slice per candidate instead of gathering pairs. + slot_pred: Vec, +} + +impl CandIndex { + fn build(store: &ChromStore) -> CandIndex { + // The table is written grouped by candidate, so a one-entry run cache resolves + // almost every row; an ungrouped table falls back to the binary search and is + // grouped just the same. + let mut distinct: std::collections::HashSet = std::collections::HashSet::new(); + let mut last: Option = None; + for &c in &store.cid { + if last != Some(c) { + distinct.insert(c); + last = Some(c); + } + } + let mut cids: Vec = distinct.into_iter().collect(); + cids.sort_unstable(); + let mut cand_off = vec![0usize; cids.len() + 1]; + let mut cache: Option<(u32, usize)> = None; + let lookup = |c: u32, cache: &mut Option<(u32, usize)>| -> usize { + if let Some((lc, li)) = *cache { + if lc == c { + return li; + } + } + let i = cids + .binary_search(&c) + .expect("candidate id was collected above"); + *cache = Some((c, i)); + i + }; + for &c in &store.cid { + cand_off[lookup(c, &mut cache) + 1] += 1; + } + for k in 1..=cids.len() { + cand_off[k] += cand_off[k - 1]; + } + let mut cursor = cand_off[..cids.len()].to_vec(); + let mut rows = vec![0usize; store.nrows()]; + let mut slot_pred = vec![0.0f32; store.nrows()]; + cache = None; + for (row, &c) in store.cid.iter().enumerate() { + let ci = lookup(c, &mut cache); + rows[cursor[ci]] = row; + slot_pred[cursor[ci]] = store.pred[row]; + cursor[ci] += 1; + } + CandIndex { + cids, + cand_off, + rows, + slot_pred, + } + } + + fn len(&self) -> usize { + self.cids.len() + } + + fn slots(&self, ci: usize) -> std::ops::Range { + self.cand_off[ci]..self.cand_off[ci + 1] + } + + fn rows_of(&self, ci: usize) -> &[usize] { + &self.rows[self.slots(ci)] + } + + fn find(&self, cid: u32) -> Option { + self.cids.binary_search(&cid).ok() + } +} + +/// Index of the sample nearest to `target`, reproducing the first-minimum tie-break of a +/// forward linear scan (`d < best`). `sorted` asserts the trace is non-decreasing and +/// NaN-free ([`ChromStore::rt_sorted`], or the union axis built by [`peak_window`]), and +/// only then is the answer found by binary search; otherwise the scan runs as before. +/// +/// The scan was the dominant per-fragment cost of a fixed-window integration: it is +/// `O(trace)` where the integration itself is `O(window)`, and on a window-grid +/// chromatogram the trace is a few thousand samples against a window of ten. +fn nearest_index(rt: &[f32], target: f64, sorted: bool) -> usize { + if !sorted { + let mut k = 0usize; + let mut best = f64::INFINITY; + for (i, &r) in rt.iter().enumerate() { + let d = (r as f64 - target).abs(); + if d < best { + best = d; + k = i; + } + } + return k; + } + // First sample at or after `target`. A non-finite target makes every predicate false + // and lands on 0, which is where the scan's untouched `k` sits too. + let j = rt.partition_point(|&r| (r as f64) < target); + // Backing up over an equal-valued run keeps the FIRST of several identical samples, + // which is the one the scan's strict `<` kept. + let back = |mut m: usize| { + while m > 0 && rt[m - 1] == rt[m] { + m -= 1; + } + m + }; + if j == 0 { + return 0; + } + if j >= rt.len() { + return back(rt.len() - 1); + } + // Sorted, so the distance is V-shaped and the minimum is at `j - 1` or `j`; an exact + // tie is the equidistant case, where the scan kept the earlier index. + if (target - rt[j - 1] as f64).abs() <= (rt[j] as f64 - target).abs() { + back(j - 1) + } else { + j + } +} + /// Trapezoidal integral of an intensity trace over RT (seconds). A single point /// yields its raw intensity. fn trapezoid(rt: &[f32], inten: &[f32]) -> f64 { @@ -128,20 +423,27 @@ fn flank_baseline(inten: &[f32], lo: usize, hi: usize, flank: usize, quantile: f /// otherwise takes `half` scans on each side. `None` for an empty trace. /// /// Shared by the integration below and by the applied-window contract in [`run`], -/// so the bounds reported for a quantity are the bounds it was integrated over. -fn fixed_window_indices(rt: &[f32], apex: f64, half: usize, half_s: f64) -> Option<(usize, usize)> { +/// so the bounds reported for a quantity are the bounds it was integrated over. [`run`] +/// computes this ONCE per fragment row and uses it for both, where it used to compute the +/// same indices a second time to report the bounds. +/// +/// `sorted` says the trace is non-decreasing and NaN-free (see [`nearest_index`]). +fn fixed_window_indices( + rt: &[f32], + apex: f64, + half: usize, + half_s: f64, + sorted: bool, +) -> Option<(usize, usize)> { if rt.is_empty() { return None; } - let mut k = 0usize; - let mut best = f64::INFINITY; - for (i, &r) in rt.iter().enumerate() { - let d = (r as f64 - apex).abs(); - if d < best { - best = d; - k = i; - } - } + let k = nearest_index(rt, apex, sorted); + // The scan this replaced left `best` at infinity when every distance was NaN (a + // non-finite apex, which makes the guard below reject the window); reproduce that + // rather than letting a NaN slip past the comparison as false. + let d = (rt[k] as f64 - apex).abs(); + let best = if d.is_nan() { f64::INFINITY } else { d }; let (mut lo, mut hi) = if half_s > 0.0 { // Distance guard: the nearest sample must itself lie inside the window. Without // it the nearest sample was always included however far away it was, so a weak @@ -181,9 +483,21 @@ fn fixed_window_indices(rt: &[f32], apex: f64, half: usize, half_s: f64) -> Opti Some((lo, hi)) } +/// True when `rt` is non-decreasing, non-negative and NaN-free, the property +/// [`nearest_index`] needs before it may binary-search. [`run`] takes it once per store +/// ([`ChromStore::rt_sorted`]) rather than per call. +fn rt_is_sorted(rt: &[f32]) -> bool { + rt.is_empty() || (rt[0] >= 0.0 && rt.windows(2).all(|w| w[0] <= w[1])) +} + /// Fixed-window integration over the samples chosen by [`fixed_window_indices`], /// with optional apex-outward envelope and optional flank-baseline subtraction /// (`baseline = Some((flank, quantile))`). Empty trace integrates to 0. +/// +/// Test-facing composition of [`fixed_window_indices`] and [`trapezoid_fixed_at`]; [`run`] +/// calls the two separately so the window indices are computed once and serve both the +/// area and the reported bounds. +#[cfg(test)] fn trapezoid_fixed_opts( rt: &[f32], inten: &[f32], @@ -193,9 +507,23 @@ fn trapezoid_fixed_opts( envelope: bool, baseline: Option<(usize, f64)>, ) -> f64 { - let Some((lo, hi)) = fixed_window_indices(rt, apex, half, half_s) else { + let Some((lo, hi)) = fixed_window_indices(rt, apex, half, half_s, rt_is_sorted(rt)) else { return 0.0; }; + trapezoid_fixed_at(rt, inten, lo, hi, envelope, baseline) +} + +/// Fixed-window integration over the sample range `[lo, hi)` already chosen by +/// [`fixed_window_indices`], with optional apex-outward envelope and optional +/// flank-baseline subtraction (`baseline = Some((flank, quantile))`). +fn trapezoid_fixed_at( + rt: &[f32], + inten: &[f32], + lo: usize, + hi: usize, + envelope: bool, + baseline: Option<(usize, f64)>, +) -> f64 { let mut w: Vec = inten[lo..hi].to_vec(); if let Some((flank, quantile)) = baseline { let b = flank_baseline(inten, lo, hi, flank, quantile); @@ -212,23 +540,26 @@ fn trapezoid_fixed_opts( /// Top-N sum with the fragment ranking chosen by `selection`. `observed_area` /// delegates to [`summarize_fragment_areas`] (legacy, byte-identical); `predicted` /// ranks the positive finite areas by library intensity and sums the top N. +/// +/// The areas and their library intensities arrive as two parallel slices of one +/// candidate's fragment rows rather than a `Vec<(f64, f32)>`, because [`run`] holds them +/// as flat per-candidate slices of one buffer each. `None` is a candidate with no +/// chromatogram row at all. fn select_fragment_areas( - areas: Option<&[(f64, f32)]>, + areas: Option<(&[f64], &[f32])>, top_n: usize, selection: FragmentSelection, ) -> (Option, usize, &'static str) { match selection { - FragmentSelection::ObservedArea => { - let plain: Option> = areas.map(|a| a.iter().map(|x| x.0).collect()); - summarize_fragment_areas(plain.as_deref(), top_n) - } + FragmentSelection::ObservedArea => summarize_fragment_areas(areas.map(|a| a.0), top_n), FragmentSelection::Predicted => { - let Some(areas) = areas else { + let Some((areas, preds)) = areas else { return (None, 0, "no_fragment_traces"); }; let mut positive: Vec<(f64, f32)> = areas .iter() .copied() + .zip(preds.iter().copied()) // The ranking key must be finite too. `total_cmp` orders NaN ABOVE every // real value, so in the descending sort below a NaN predicted intensity // reached the front and was preferentially selected into the top N. The @@ -256,7 +587,7 @@ fn select_fragment_areas( /// Elution-peak RT window `[lo, hi]` for one candidate, from the summed XIC across /// all its fragment chromatograms. Fragments are aligned on the union of their RT -/// samples via a BTreeMap keyed by the f32 RT bit pattern: for the non-negative +/// samples ordered by the f32 RT bit pattern: for the non-negative /// RTs here the bit order matches the value order, so both the union axis and the /// f64 summation order are fixed (determinism, /// docs/14_build_test_deploy_gotchas.md). When a finite identification apex is @@ -266,54 +597,105 @@ fn select_fragment_areas( /// artifacts and missing/non-finite hints retain the legacy co-elution apex /// detector. Returns an unbounded window when there are fewer than two distinct /// RT samples (nothing to bound). +/// +/// The union used to be built by inserting every sample into TWO `BTreeMap`s keyed +/// identically, one tree walk per sample each. It is now either a direct accumulation +/// (every row of the candidate shares one strictly increasing axis, which is what +/// extract's window grid writes, so the union IS that axis) or one stable sort of the +/// samples by RT bits. Both reduce each RT's f64 sum in the same order the tree did -- +/// row order, then sample order -- so the profile, the apex and the bounds are +/// bit-identical. fn peak_window( rows: &[usize], - ch_rt: &[Vec], - ch_int: &[Vec], + store: &ChromStore, frac: f64, grace: usize, apex_hint: Option, ) -> (f64, f64, f64) { - let mut prof_map: BTreeMap = BTreeMap::new(); - // Per-scan count of co-eluting (nonzero) fragments, aligned to prof_map keys. - let mut cnt_map: BTreeMap = BTreeMap::new(); + // Does every non-empty row of this candidate sample the same strictly increasing + // axis? An empty trace contributes no sample and so cannot widen the union. + let mut shared: Option = None; + let mut uniform = true; for &i in rows { - let rts = &ch_rt[i]; - let ins = &ch_int[i]; - for k in 0..rts.len() { - *prof_map.entry(rts[k].to_bits()).or_insert(0.0) += ins[k] as f64; - if ins[k] > 0.0 { - *cnt_map.entry(rts[k].to_bits()).or_insert(0) += 1; + let a = store.axis_id[i]; + if a == NO_AXIS { + continue; + } + match shared { + None => shared = Some(a), + Some(s) if s == a => {} + _ => { + uniform = false; + break; } } } - if prof_map.len() < 2 { + let shared = shared.filter(|&a| uniform && store.axis_strict[a as usize]); + + let mut owned_axis: Vec = Vec::new(); + let (prof, cnt) = if let Some(a) = shared { + let n = store.axis(a).len(); + let mut prof = vec![0.0f64; n]; + let mut cnt = vec![0u32; n]; + for &i in rows { + for (k, &v) in store.inten(i).iter().enumerate() { + prof[k] += v as f64; + if v > 0.0 { + cnt[k] += 1; + } + } + } + (prof, cnt) + } else { + let total: usize = rows.iter().map(|&i| store.inten(i).len()).sum(); + let mut samples: Vec<(u32, f32)> = Vec::with_capacity(total); + for &i in rows { + let rts = store.rt(i); + let ins = store.inten(i); + for k in 0..rts.len() { + samples.push((rts[k].to_bits(), ins[k])); + } + } + // Stable, so within one RT the contributions keep their (row, sample) order -- + // the order the map accumulated them in. + samples.sort_by_key(|s| s.0); + let mut prof: Vec = Vec::new(); + let mut cnt: Vec = Vec::new(); + let mut k = 0usize; + while k < samples.len() { + let bits = samples[k].0; + let mut sum = 0.0f64; + let mut c = 0u32; + while k < samples.len() && samples[k].0 == bits { + sum += samples[k].1 as f64; + if samples[k].1 > 0.0 { + c += 1; + } + k += 1; + } + owned_axis.push(f32::from_bits(bits)); + prof.push(sum); + cnt.push(c); + } + (prof, cnt) + }; + let axis: &[f32] = match shared { + Some(a) => store.axis(a), + None => &owned_axis, + }; + if axis.len() < 2 { // Nothing to bound; apex is the lone RT if present, else NaN. - let apex = prof_map - .keys() - .next() - .map_or(f64::NAN, |b| f32::from_bits(*b) as f64); + let apex = axis.first().map_or(f64::NAN, |&r| r as f64); return (f64::NEG_INFINITY, f64::INFINITY, apex); } - let axis: Vec = prof_map.keys().map(|b| f32::from_bits(*b) as f64).collect(); - let prof: Vec = prof_map.values().cloned().collect(); - let cnt: Vec = prof_map - .keys() - .map(|b| *cnt_map.get(b).unwrap_or(&0)) - .collect(); + // Both constructions leave the axis ascending: the shared one is checked strictly + // increasing when it is stored, and the merged one is deduped and ordered by bit + // pattern, which is value order as long as no negative or NaN RT is present. + let axis_sorted = shared.is_some() || axis.last().is_some_and(|&r| r >= 0.0); let ai = if let Some(hint) = apex_hint.filter(|v| v.is_finite()) { // The identified apex need not exactly equal a chromatogram sample (for // example after serialization/calibration), so anchor to the nearest RT. - let mut nearest = 0usize; - let mut distance = f64::INFINITY; - for (i, &rt) in axis.iter().enumerate() { - let d = (rt - hint).abs(); - if d < distance { - nearest = i; - distance = d; - } - } - nearest + nearest_index(axis, hint, axis_sorted) } else { // Legacy robust apex: among scans whose co-eluting-fragment count is // within 1 of the maximum ("-1 for robustness"), take the one with the @@ -356,7 +738,7 @@ fn peak_window( } lo = lo.saturating_sub(1); } - (axis[lo], axis[hi], axis[ai]) + (axis[lo] as f64, axis[hi] as f64, axis[ai] as f64) } fn finite_option(value: f64) -> Option { @@ -419,7 +801,17 @@ fn add_protein_base_quantity( base_peptide_id: u32, quantity: Option, ) { - let bases = groups.entry(protein_group.to_string()).or_default(); + // `entry` needs an owned key, so this allocated a `String` for every scored row even + // though the group almost always exists already. Look it up first and allocate only + // for a group that is new. The group is still created for an unquantifiable row: a + // protein whose every peptide is unquantifiable is REPORTED, with status + // `no_quantifiable_peptide`, not omitted. + if !groups.contains_key(protein_group) { + groups.insert(protein_group.to_string(), BTreeMap::new()); + } + let bases = groups + .get_mut(protein_group) + .expect("the group was just inserted"); if let Some(quantity) = quantity.filter(|v| v.is_finite() && *v > 0.0) { bases .entry(base_peptide_id) @@ -606,14 +998,15 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { ); } let ch = TableFile::open(p.chromatograms)?; - let mut ch_cid: Vec = Vec::new(); - let mut ch_name: Vec = Vec::new(); - let mut ch_rt: Vec> = Vec::new(); - let mut ch_int: Vec> = Vec::new(); - // Zero, not NaN, for the absent column: the value is only a ranking key, and - // `total_cmp` orders NaN above every real intensity, which would silently invert the - // `predicted` ranking rather than fail. - let mut ch_pred: Vec = Vec::new(); + // Flat, grouped-by-candidate store (see [`ChromStore`]). The MS1 isotope XIC + // pseudo-traces (frag_name "ms1_*") are precursor channels, not fragment ions, and + // neither the peak-window detection nor the top-N sum ever reads one, so they are not + // stored at all rather than loaded and skipped later. + // + // Zero, not NaN, for an absent `predicted_intensity`: the value is only a ranking key, + // and `total_cmp` orders NaN above every real intensity, which would silently invert + // the `predicted` ranking rather than fail. + let mut store = ChromStore::new(); { let cols: Vec<&str> = if has_pred { vec![ @@ -643,6 +1036,11 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { } else { None }; + // Trace values are appended into scratch buffers and then copied into the store, + // because the RT axis is deduplicated against the candidate's earlier rows before + // it is stored. + let mut rt_buf: Vec = Vec::new(); + let mut int_buf: Vec = Vec::new(); for b in reader { let b = b?; let a_cid = b @@ -671,8 +1069,16 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { if !keep_all && !wanted.contains(&c) { continue; } - let rt = a_rt.row(k, "rt")?; - let it = a_int.row(k, "intensity")?; + let nm = if a_name.is_null(k) { + "" + } else { + a_name.value(k) + }; + rt_buf.clear(); + int_buf.clear(); + a_rt.append_row(k, &mut rt_buf, "rt")?; + a_int.append_row(k, &mut int_buf, "intensity")?; + let (rt, it) = (&rt_buf, &int_buf); // `rt` and `intensity` are two independent list columns, and every // integration below slices `intensity` with indices computed from the LENGTH // OF `rt`. Extract writes them from paired vectors so they always match, but @@ -691,69 +1097,74 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { p.chromatograms ); } - ch_cid.push(c); - ch_name.push(if a_name.is_null(k) { - String::new() - } else { - a_name.value(k).to_string() - }); - ch_rt.push(rt); - ch_int.push(it); - ch_pred.push(match a_pred { + // The MS1 isotope XIC pseudo-traces are precursor channels, not fragment + // ions, and no phase below reads one, so they are dropped rather than + // stored and skipped later. Dropped AFTER the length check, which is a + // guard on the table rather than on what this stage happens to consume. + if nm.starts_with("ms1_") { + continue; + } + let pred = match a_pred { Some(a) => a.value(k), None => 0.0, - }); + }; + store.push(c, nm, pred, rt, it); } } } drop(wanted); - // Group the b/y fragment chromatogram rows by candidate. The MS1 isotope XIC - // pseudo-traces (frag_name "ms1_*") are precursor channels, not fragment ions, - // and are excluded from both the peak-window detection and the top-N sum. The - // BTreeMap keeps candidate iteration order deterministic. - let mut cand_rows: BTreeMap> = BTreeMap::new(); - for i in 0..ch_cid.len() { - if ch_name[i].starts_with("ms1_") { - continue; - } - cand_rows.entry(ch_cid[i]).or_default().push(i); + let store = store; + // Group the fragment chromatogram rows by candidate, ascending by candidate id and in + // table order within a candidate: the order the `BTreeMap>` this + // replaces iterated in, so every per-candidate reduction below runs over the same rows + // in the same order. + let index = CandIndex::build(&store); + { + let mut parts = store.mem_parts(); + parts.push(( + "cand_index", + crate::memlog::bytes_of(&index.rows) + + crate::memlog::bytes_of(&index.slot_pred) + + crate::memlog::bytes_of(&index.cids) + + crate::memlog::bytes_of(&index.cand_off), + )); + crate::memlog::report("quant chromatogram store", &parts); } - let mut areas: HashMap> = HashMap::new(); - // Store the fragment name by reference (borrowed from `ch_name`, which outlives - // this map) to avoid a per-row String clone; it is materialized once at export. - let mut frag_areas: HashMap> = HashMap::new(); + info!( + chromatogram_rows = store.nrows(), + candidates = index.len(), + "quant: chromatograms loaded" + ); // Optional peak-window diagnostic: (candidate_id, lo_rt, hi_rt) for finite windows. let emit_bounds = p.out_peak_bounds.is_some() && p.cfg.bound_peak; let (mut pb_cid, mut pb_lo, mut pb_hi) = (Vec::new(), Vec::new(), Vec::new()); // Phase 1: per-candidate summed-XIC window (lo_rt, hi_rt, apex_rt), anchored at // the identification apex when available and otherwise using the legacy robust - // co-elution detector. Kept keyed by candidate for the consensus estimate. - let mut win: BTreeMap = BTreeMap::new(); - if p.cfg.bound_peak { - // Each candidate's window depends only on its own chromatogram rows, so this is - // embarrassingly parallel. Results are collected into a Vec and only then folded - // into the BTreeMap, so the map is built from a deterministically ordered sequence - // and every float inside `peak_window` is still reduced per candidate in the same - // order as before -- bit-identical, not merely equivalent. - let computed: Vec<(u32, (f64, f64, f64))> = cand_rows - .par_iter() - .map(|(&c, rows)| { - ( - c, - peak_window( - rows, - &ch_rt, - &ch_int, - p.cfg.peak_fraction, - p.cfg.peak_grace, - apex_by_cid.get(&c).copied(), - ), + // co-elution detector. Indexed by candidate position in `index` for the consensus + // estimate and for phase 2, which is the same ascending-candidate_id order the + // `BTreeMap` this replaces iterated in, without its node per candidate. + // + // Each candidate's window depends only on its own chromatogram rows, so this is + // embarrassingly parallel; a range is an indexed parallel iterator, so the collected + // Vec is in candidate order and every float inside `peak_window` is still reduced per + // candidate in the same order as before. + let win: Vec<(f64, f64, f64)> = if p.cfg.bound_peak { + (0..index.len()) + .into_par_iter() + .map(|ci| { + peak_window( + index.rows_of(ci), + &store, + p.cfg.peak_fraction, + p.cfg.peak_grace, + apex_by_cid.get(&index.cids[ci]).copied(), ) }) - .collect(); - win.extend(computed); - } + .collect() + } else { + Vec::new() + }; // Consensus mode: peak width is a near-constant instrument/gradient property, so // take the median left/right half-width over CONFIDENT peptides (q <= reliable_q) @@ -773,11 +1184,13 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { } } let (mut left, mut right) = (Vec::new(), Vec::new()); - for (c, &(lo, hi, apex)) in &win { + for (ci, &(lo, hi, apex)) in win.iter().enumerate() { if lo.is_finite() && hi.is_finite() && apex.is_finite() - && q_by_cid.get(c).is_some_and(|&q| q <= p.cfg.reliable_q) + && q_by_cid + .get(&index.cids[ci]) + .is_some_and(|&q| q <= p.cfg.reliable_q) { left.push(apex - lo); right.push(hi - apex); @@ -806,18 +1219,38 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { // Phase 2: integrate each fragment over the chosen window and retain the // actually applied apex/bounds for the peptide-quant contract. - let mut applied_win: BTreeMap = BTreeMap::new(); - // `(candidate_id, (lo_rt, hi_rt, integration_apex), one area per chromatogram row)`. - type Integrated = (u32, (f64, f64, f64), Vec); + // // Integrate each candidate's fragment traces. Parallel across candidates for the same // reason the peak-window phase above is: a candidate reads only its own chromatogram // rows and every float reduction happens inside one candidate's `trapezoid*` call. - // Rayon's indexed `collect` keeps candidate order, and `cand_rows` is a `BTreeMap`, so - // the maps below are filled in exactly the order the serial loop filled them. - let integrated: Vec = cand_rows - .par_iter() - .map(|(&c, rows)| { - let want_fixed = p.cfg.fixed_scan_halfwidth > 0 || p.cfg.fixed_window_s > 0.0; + // The areas go into ONE flat buffer laid out exactly like `index.rows`, so candidate + // `ci`'s areas are `area_by_slot[index.slots(ci)]` and nothing is allocated per + // candidate. Rayon writes into disjoint sub-slices of it, carved once below, which is + // also why the positional pairing the serial fold used to assert is gone: each + // candidate writes the slice that is its own by construction. + let want_fixed = p.cfg.fixed_scan_halfwidth > 0 || p.cfg.fixed_window_s > 0.0; + let baseline = if p.cfg.baseline_subtract { + Some((p.cfg.baseline_flank_scans, p.cfg.baseline_quantile)) + } else { + None + }; + let mut area_by_slot: Vec = vec![0.0; index.rows.len()]; + let mut slices: Vec<&mut [f64]> = Vec::with_capacity(index.len()); + { + let mut rest: &mut [f64] = &mut area_by_slot; + for ci in 0..index.len() { + let n = index.cand_off[ci + 1] - index.cand_off[ci]; + let (head, tail) = rest.split_at_mut(n); + slices.push(head); + rest = tail; + } + } + let applied_win: Vec<(f64, f64, f64)> = slices + .par_iter_mut() + .enumerate() + .map(|(ci, out)| { + let c = index.cids[ci]; + let rows = index.rows_of(ci); let (lo_rt, hi_rt, integration_apex) = if !p.cfg.bound_peak { // A fixed window needs only an apex to centre on, not the descent walk, // so take the identification apex directly. Previously this branch @@ -839,7 +1272,7 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { }; (f64::NEG_INFINITY, f64::INFINITY, apex) } else { - let (lo, hi, apex) = win[&c]; + let (lo, hi, apex) = win[ci]; match consensus { Some((ml, mr)) if apex.is_finite() => (apex - ml, apex + mr, apex), _ => (lo, hi, apex), @@ -848,87 +1281,63 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { // A fixed window replaces the walked bounds entirely; it needs a finite apex // to centre on, so an unknown apex falls back to the configured window. let fixed = want_fixed && integration_apex.is_finite(); - let a: Vec = rows - .iter() - .map(|&i| { - if fixed { - trapezoid_fixed_opts( - &ch_rt[i], - &ch_int[i], - integration_apex, - p.cfg.fixed_scan_halfwidth, - p.cfg.fixed_window_s, - p.cfg.interference_envelope, - if p.cfg.baseline_subtract { - Some((p.cfg.baseline_flank_scans, p.cfg.baseline_quantile)) - } else { - None - }, - ) - } else if p.cfg.bound_peak { - trapezoid_window( - &ch_rt[i], - &ch_int[i], - lo_rt, - hi_rt, - p.cfg.interference_envelope, - ) - } else { - trapezoid(&ch_rt[i], &ch_int[i]) - } - }) - .collect(); // Applied-window contract: under a fixed window the walked bounds are NOT the // integration range, so report the RT extent actually covered (union over this // candidate's traces, whose sample grids may differ). Otherwise // `integration_lo_rt`/`integration_hi_rt` and the peak-bounds diagnostic would - // describe a window that produced no part of `quantity`. - let (lo_rt, hi_rt) = if fixed { - let mut flo = f64::INFINITY; - let mut fhi = f64::NEG_INFINITY; - for &i in rows.iter() { - if let Some((lo, hi)) = fixed_window_indices( - &ch_rt[i], + // describe a window that produced no part of `quantity`. The window indices are + // computed ONCE per row and serve both the area and this union; the union used + // to recompute the identical indices in a second pass over the same rows. + let mut flo = f64::INFINITY; + let mut fhi = f64::NEG_INFINITY; + for (slot, &i) in rows.iter().enumerate() { + let rt = store.rt(i); + let it = store.inten(i); + out[slot] = if fixed { + match fixed_window_indices( + rt, integration_apex, p.cfg.fixed_scan_halfwidth, p.cfg.fixed_window_s, + store.rt_sorted, ) { - flo = flo.min(ch_rt[i][lo] as f64); - fhi = fhi.max(ch_rt[i][hi - 1] as f64); + Some((lo, hi)) => { + flo = flo.min(rt[lo] as f64); + fhi = fhi.max(rt[hi - 1] as f64); + trapezoid_fixed_at( + rt, + it, + lo, + hi, + p.cfg.interference_envelope, + baseline, + ) + } + None => 0.0, } - } - if flo.is_finite() && fhi.is_finite() { - (flo, fhi) + } else if p.cfg.bound_peak { + trapezoid_window(rt, it, lo_rt, hi_rt, p.cfg.interference_envelope) } else { - (lo_rt, hi_rt) - } + trapezoid(rt, it) + }; + } + let (lo_rt, hi_rt) = if fixed && flo.is_finite() && fhi.is_finite() { + (flo, fhi) } else { (lo_rt, hi_rt) }; - (c, (lo_rt, hi_rt, integration_apex), a) + (lo_rt, hi_rt, integration_apex) }) .collect(); - for ((c, w, computed), (key, rows)) in integrated.into_iter().zip(cand_rows.iter()) { - // The pairing is positional. It holds because both sides derive from the same - // `BTreeMap` and rayon's `collect` into a `Vec` preserves order, but the comment - // that justified it named the wrong reason (`BTreeMap::par_iter()` is not an - // `IndexedParallelIterator`; the order survives incidentally). If it ever - // diverged, one candidate's areas would attach to another's fragment rows with no - // error at all, so assert the invariant instead of arguing it. - debug_assert_eq!(c, *key, "quant: integrated order diverged from cand_rows"); - let (lo_rt, hi_rt, _) = w; - applied_win.insert(c, w); - if emit_bounds && lo_rt.is_finite() && hi_rt.is_finite() { - pb_cid.push(c); - pb_lo.push(lo_rt); - pb_hi.push(hi_rt); - } - for (&i, a) in rows.iter().zip(computed) { - areas.entry(c).or_default().push((a, ch_pred[i])); - frag_areas - .entry(c) - .or_default() - .push((ch_name[i].as_str(), a)); + drop(slices); + let area_by_slot = area_by_slot; + if emit_bounds { + for (ci, &(lo_rt, hi_rt, _)) in applied_win.iter().enumerate() { + if lo_rt.is_finite() && hi_rt.is_finite() { + pb_cid.push(index.cids[ci]); + pb_lo.push(lo_rt); + pb_hi.push(hi_rt); + } } } @@ -974,20 +1383,37 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { ); let mut per_group: ProteinBaseQuant = BTreeMap::new(); let mut n_quantified_peptides = 0u64; + // The top-N selection depends on the candidate alone, so it is computed once per + // candidate rather than once per scored row: it collects and SORTS the candidate's + // areas, and every extra row mapping to the same candidate repeated that. + let mut selected: HashMap = HashMap::new(); + let mut selections: Vec<(Option, usize, &'static str)> = Vec::new(); for i in 0..ps.nrows { if !passes_quant_filter(&label[i], pep_q[i], p.cfg.q_threshold, is_transferred[i]) { continue; } - let (quantity, used, status) = select_fragment_areas( - areas.get(&cid[i]).map(Vec::as_slice), - p.cfg.top_n_fragments, - p.cfg.fragment_selection, - ); + let ci = index.find(cid[i]); + let (quantity, used, status) = match ci { + None => select_fragment_areas(None, p.cfg.top_n_fragments, p.cfg.fragment_selection), + Some(ci) => { + let slot = *selected.entry(ci).or_insert_with(|| { + let r = index.slots(ci); + selections.push(select_fragment_areas( + Some((&area_by_slot[r.clone()], &index.slot_pred[r])), + p.cfg.top_n_fragments, + p.cfg.fragment_selection, + )); + selections.len() - 1 + }); + selections[slot] + } + }; if quantity.is_some() { n_quantified_peptides += 1; } - let (integration_lo, integration_hi, integration_apex) = match applied_win.get(&cid[i]) { - Some(&(lo, hi, apex)) => (finite_option(lo), finite_option(hi), finite_option(apex)), + let (integration_lo, integration_hi, integration_apex) = match ci.map(|ci| applied_win[ci]) + { + Some((lo, hi, apex)) => (finite_option(lo), finite_option(hi), finite_option(apex)), None => (None, None, None), }; q_cid.push(cid[i]); @@ -1061,17 +1487,21 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { if !passes_quant_filter(&label[i], pep_q[i], p.cfg.q_threshold, is_transferred[i]) { continue; } - if let Some(fa) = frag_areas.get(&cid[i]) { - for (nm, a) in fa { - if !a.is_finite() || *a <= 0.0 { + // The fragment name comes from the interned table by row rather than from a + // per-candidate `Vec<(&str, f64)>` built alongside the areas; the rows are the + // same rows in the same order. + if let Some(ci) = index.find(cid[i]) { + for slot in index.slots(ci) { + let a = area_by_slot[slot]; + if !a.is_finite() || a <= 0.0 { continue; } f_cid.push(cid[i]); f_pf.push(pform[i].clone()); f_z.push(charge[i]); f_pg.push(pg[i].clone()); - f_name.push(nm.to_string()); - f_area.push(*a); + f_name.push(store.name(index.rows[slot]).to_string()); + f_area.push(a); } } } @@ -1213,9 +1643,15 @@ pub fn run_lfq_combine( Some(fnm) => format!("{}|{}|{}", pform[i], z[i], fnm[i]), None => format!("{}|{}", pform[i], z[i]), }; + // Same reason as `add_protein_base_quantity`: `entry` would clone the protein + // group name on every row, and a run's quant table repeats each group many + // times over. + if !data.contains_key(&pgc[i]) { + data.insert(pgc[i].clone(), BTreeMap::new()); + } let slot = &mut data - .entry(pgc[i].clone()) - .or_default() + .get_mut(&pgc[i]) + .expect("the group was just inserted") .entry(key) .or_insert_with(|| vec![None; n])[ri]; *slot = Some(slot.map_or(quantity, |previous| previous.max(quantity))); @@ -1237,8 +1673,13 @@ pub fn run_lfq_combine( } } let (mut c_pg, mut c_run, mut c_q, mut c_nf) = (Vec::new(), Vec::new(), Vec::new(), Vec::new()); + // The rollup levels below read the SAME feature vectors, so each level borrows them + // from `data` rather than cloning them: with the protein level and the two sibling + // levels this held three or four copies of every feature vector at once. + let mut mat: Vec<&Vec>> = Vec::new(); for (pgname, feats) in &data { - let mat: Vec>> = feats.values().cloned().collect(); + mat.clear(); + mat.extend(feats.values()); let prof = crate::quant_lfq::lfq_profile(&mat, n); for (r, &v) in prof.iter().enumerate() { c_pg.push(pgname.clone()); @@ -1263,25 +1704,24 @@ pub fn run_lfq_combine( // base sequence. Both roll their member features up with the same LFQ engine. // Purely additive analysis granularity; strictly post-FDR, no identification // or FDR change. - let mut prec: BTreeMap<(String, i32), Vec>>> = BTreeMap::new(); - let mut pep: BTreeMap>>> = BTreeMap::new(); + let mut prec: BTreeMap<(String, i32), Vec<&Vec>>> = BTreeMap::new(); + let mut pep: BTreeMap>>> = BTreeMap::new(); for feats in data.values() { for (key, vec) in feats { // key = "peptidoform|charge" (maxlfq) or "peptidoform|charge|fragment" // (directlfq); peptidoform strings never contain '|'. let mut it = key.splitn(3, '|'); - let pform = it.next().unwrap_or("").to_string(); + let pform = it.next().unwrap_or(""); let charge: i32 = it.next().and_then(|s| s.parse().ok()).unwrap_or(0); - prec.entry((pform.clone(), charge)) + prec.entry((pform.to_string(), charge)) .or_default() - .push(vec.clone()); - pep.entry(base_sequence(&pform)) - .or_default() - .push(vec.clone()); + .push(vec); + pep.entry(base_sequence(pform)).or_default().push(vec); } } - // (group key, charge, feature-by-run matrix) for one sibling-matrix level. - type LevelGroup = (String, i32, Vec>>); + // (group key, charge, feature-by-run matrix) for one sibling-matrix level. The matrix + // borrows its rows from `data`. + type LevelGroup<'a> = (String, i32, Vec<&'a Vec>>); let write_level = |path: String, groups: Vec| -> Result<()> { let (mut g_key, mut g_z, mut g_run, mut g_q, mut g_nf) = (Vec::new(), Vec::new(), Vec::new(), Vec::new(), Vec::new()); @@ -1441,6 +1881,77 @@ mod tests { use super::*; use mumdia_io::table::Table; + /// A store built from plain per-row traces, all rows belonging to one candidate, so + /// the axes deduplicate exactly as extract's window grid makes them. + fn store_of(rt: &[Vec], inten: &[Vec]) -> ChromStore { + let mut s = ChromStore::new(); + for (i, (r, it)) in rt.iter().zip(inten).enumerate() { + s.push(0, &format!("y{i}"), 0.0, r, it); + } + s + } + + /// The same rows under a DISTINCT candidate id each, so no two rows can share an axis + /// id and [`peak_window`] must merge the samples instead of reading one shared axis. + /// The row indices are passed to `peak_window` explicitly, so the candidate ids only + /// steer the store's axis dedup. + fn store_of_unshared(rt: &[Vec], inten: &[Vec]) -> ChromStore { + let mut s = ChromStore::new(); + for (i, (r, it)) in rt.iter().zip(inten).enumerate() { + s.push(i as u32, &format!("y{i}"), 0.0, r, it); + } + s + } + + /// [`peak_window`] through BOTH union constructions, asserting they are bit-identical. + /// The shared-axis path is an accumulation into one profile array; the merged path is + /// the sort over every sample that replaced the two BTreeMaps. Every `peak_window` + /// test below goes through this, so each is also an equality test between the paths. + fn peak_window_both( + rows: &[usize], + rt: &[Vec], + inten: &[Vec], + frac: f64, + grace: usize, + hint: Option, + ) -> (f64, f64, f64) { + let shared = peak_window(rows, &store_of(rt, inten), frac, grace, hint); + let merged = peak_window(rows, &store_of_unshared(rt, inten), frac, grace, hint); + let bits = |w: (f64, f64, f64)| (w.0.to_bits(), w.1.to_bits(), w.2.to_bits()); + assert_eq!( + bits(shared), + bits(merged), + "shared-axis and merged-sample unions disagree: {shared:?} vs {merged:?}" + ); + shared + } + + /// [`fixed_window_indices`] through both nearest-sample searches, asserting they + /// agree: `false` forces the linear first-minimum scan the binary search replaced. + fn fixed_window_indices_both( + rt: &[f32], + apex: f64, + half: usize, + half_s: f64, + ) -> Option<(usize, usize)> { + let scanned = fixed_window_indices(rt, apex, half, half_s, false); + let searched = fixed_window_indices(rt, apex, half, half_s, rt_is_sorted(rt)); + assert_eq!( + scanned, searched, + "scan and binary search disagree on {rt:?} at apex {apex}" + ); + scanned + } + + /// One candidate's fragment areas as the two parallel slices `select_fragment_areas` + /// reads (it used to take a `Vec<(f64, f32)>`). + fn split_areas(pairs: &[(f64, f32)]) -> (Vec, Vec) { + ( + pairs.iter().map(|p| p.0).collect(), + pairs.iter().map(|p| p.1).collect(), + ) + } + fn quant_test_path(name: &str) -> String { use std::sync::atomic::{AtomicU64, Ordering}; static COUNTER: AtomicU64 = AtomicU64::new(0); @@ -1519,7 +2030,7 @@ mod tests { ]; let ch_rt = vec![grid.clone(), grid.clone()]; let ch_int = vec![real, interf]; - let (lo, hi, _) = peak_window(&[0, 1], &ch_rt, &ch_int, 1.0 / 6.0, 1, None); + let (lo, hi, _) = peak_window_both(&[0, 1], &ch_rt, &ch_int, 1.0 / 6.0, 1, None); // Apex rt=6 (sum=10); 1/6 threshold ~1.67. Left: idx4(2)>=thr, idx3/idx2=0 // -> 2 consecutive misses stop at rt=4. Right: symmetric stop at rt=8. assert_eq!(lo, 4.0); @@ -1537,8 +2048,8 @@ mod tests { let prof = vec![0.0f32, 0.0, 1.0, 5.0, 10.0, 5.0, 1.0, 5.0, 0.0]; let ch_rt = vec![grid.clone()]; let ch_int = vec![prof]; - let (_, hi1, _) = peak_window(&[0], &ch_rt, &ch_int, 1.0 / 3.0, 1, None); - let (_, hi0, _) = peak_window(&[0], &ch_rt, &ch_int, 1.0 / 3.0, 0, None); + let (_, hi1, _) = peak_window_both(&[0], &ch_rt, &ch_int, 1.0 / 3.0, 1, None); + let (_, hi0, _) = peak_window_both(&[0], &ch_rt, &ch_int, 1.0 / 3.0, 0, None); // grace=1 bridges the idx6 dip (1.0 < 3.33) and includes idx7 (5.0) -> rt 7. assert_eq!(hi1, 7.0); // grace=0 stops at the first sub-threshold scan -> last above-threshold rt 5. @@ -1625,7 +2136,7 @@ mod tests { real.clone(), interf, ]; - let (_, _, apex) = peak_window(&[0, 1, 2, 3, 4], &ch_rt, &ch_int, 1.0 / 6.0, 1, None); + let (_, _, apex) = peak_window_both(&[0, 1, 2, 3, 4], &ch_rt, &ch_int, 1.0 / 6.0, 1, None); assert_eq!( apex, 5.0, "apex must be the 4-fragment co-elution scan, not the lone interferent spike" @@ -1642,16 +2153,16 @@ mod tests { let ch_rt = vec![grid.clone(), grid.clone(), grid.clone()]; let ch_int = vec![trace.clone(), trace.clone(), trace]; - let (_, _, legacy_apex) = peak_window(&[0, 1, 2], &ch_rt, &ch_int, 1.0 / 6.0, 1, None); + let (_, _, legacy_apex) = peak_window_both(&[0, 1, 2], &ch_rt, &ch_int, 1.0 / 6.0, 1, None); assert_eq!(legacy_apex, 1.0); let (lo, hi, anchored_apex) = - peak_window(&[0, 1, 2], &ch_rt, &ch_int, 1.0 / 6.0, 1, Some(5.1)); + peak_window_both(&[0, 1, 2], &ch_rt, &ch_int, 1.0 / 6.0, 1, Some(5.1)); assert_eq!(anchored_apex, 5.0); assert!(lo > 1.0 && hi >= 5.0, "anchored window was [{lo}, {hi}]"); let (_, _, nonfinite_fallback) = - peak_window(&[0, 1, 2], &ch_rt, &ch_int, 1.0 / 6.0, 1, Some(f64::NAN)); + peak_window_both(&[0, 1, 2], &ch_rt, &ch_int, 1.0 / 6.0, 1, Some(f64::NAN)); assert_eq!(nonfinite_fallback, legacy_apex); } @@ -1828,7 +2339,7 @@ mod tests { let spike = vec![0.0f32, 0.0, 10.0, 0.0, 0.0]; let ch_rt = vec![grid.clone()]; let ch_int = vec![spike]; - let (lo, hi, apex) = peak_window(&[0], &ch_rt, &ch_int, 1.0 / 6.0, 1, None); + let (lo, hi, apex) = peak_window_both(&[0], &ch_rt, &ch_int, 1.0 / 6.0, 1, None); assert_eq!(apex, 4.0, "apex should be the summed-XIC max rt"); assert!(hi > lo, "window must have nonzero width: lo={lo} hi={hi}"); let a = trapezoid_window(&ch_rt[0], &ch_int[0], lo, hi, false); @@ -1860,7 +2371,7 @@ mod tests { trapezoid_fixed_opts(&rt, &it, 5.0, 0, 0.0, false, None), 7.5 ); - assert_eq!(fixed_window_indices(&rt, 5.0, 0, 0.0), Some((5, 7))); + assert_eq!(fixed_window_indices_both(&rt, 5.0, 0, 0.0), Some((5, 7))); // A window wider than the trace integrates the whole trace (area 20). assert!((trapezoid_fixed_opts(&rt, &it, 5.0, 0, 100.0, false, None) - 20.0).abs() < 1e-9); // An apex off the sampled range still integrates around the nearest sample. @@ -1873,9 +2384,9 @@ mod tests { trapezoid_fixed_opts(&[], &[], 5.0, 3, 0.0, false, None), 0.0 ); - assert_eq!(fixed_window_indices(&[], 5.0, 3, 0.0), None); - assert_eq!(fixed_window_indices(&rt, 5.0, 1, 0.0), Some((4, 7))); - assert_eq!(fixed_window_indices(&rt, 5.0, 0, 1.0), Some((4, 7))); + assert_eq!(fixed_window_indices_both(&[], 5.0, 3, 0.0), None); + assert_eq!(fixed_window_indices_both(&rt, 5.0, 1, 0.0), Some((4, 7))); + assert_eq!(fixed_window_indices_both(&rt, 5.0, 0, 1.0), Some((4, 7))); // Distance guard on the seconds form: an apex far outside the sampled range has // NO sample inside +/- half_s, so the window is empty and the area is 0. Before @@ -1883,15 +2394,18 @@ mod tests { // fragment sampled 45 s from the apex contributed its off-peak intensity as this // candidate's area -- while `trapezoid_window`, given the same RT bounds, // correctly returns 0. - assert_eq!(fixed_window_indices(&rt, -50.0, 0, 5.0), None); + assert_eq!(fixed_window_indices_both(&rt, -50.0, 0, 5.0), None); assert_eq!( trapezoid_fixed_opts(&rt, &it, -50.0, 0, 5.0, false, None), 0.0 ); // Just inside the guard, the sample is kept (and widened to two). - assert_eq!(fixed_window_indices(&rt, -4.0, 0, 5.0), Some((0, 2))); + assert_eq!(fixed_window_indices_both(&rt, -4.0, 0, 5.0), Some((0, 2))); // A single-sample trace cannot be widened, and must not panic. - assert_eq!(fixed_window_indices(&rt[0..1], 0.0, 0, 1.0), Some((0, 1))); + assert_eq!( + fixed_window_indices_both(&rt[0..1], 0.0, 0, 1.0), + Some((0, 1)) + ); } #[test] @@ -1918,22 +2432,25 @@ mod tests { fn select_fragment_areas_ranks_by_predicted_intensity() { // Fragment 0 has the largest observed area but the smallest library intensity: // the interference case `fragment_selection = predicted` exists to avoid. - let areas = [(100.0f64, 0.1f32), (40.0, 1.0), (30.0, 0.8)]; + let pairs = [(100.0f64, 0.1f32), (40.0, 1.0), (30.0, 0.8)]; + let (a, pr) = split_areas(&pairs); + let areas = Some((a.as_slice(), pr.as_slice())); assert_eq!( - select_fragment_areas(Some(&areas), 2, FragmentSelection::ObservedArea), + select_fragment_areas(areas, 2, FragmentSelection::ObservedArea), (Some(140.0), 2, "quantified") ); assert_eq!( - select_fragment_areas(Some(&areas), 2, FragmentSelection::Predicted), + select_fragment_areas(areas, 2, FragmentSelection::Predicted), (Some(70.0), 2, "quantified") ); // `observed_area` must stay byte-identical to the legacy summariser. - let plain: Vec = areas.iter().map(|a| a.0).collect(); assert_eq!( - select_fragment_areas(Some(&areas), 2, FragmentSelection::ObservedArea), - summarize_fragment_areas(Some(&plain), 2) + select_fragment_areas(areas, 2, FragmentSelection::ObservedArea), + summarize_fragment_areas(Some(&a), 2) ); // Both rankings report the same statuses on the degenerate inputs. + let (zero_a, zero_p) = split_areas(&[(0.0, 1.0)]); + let (nan_a, nan_p) = split_areas(&[(f64::NAN, 1.0), (10.0, 0.5)]); for sel in [ FragmentSelection::ObservedArea, FragmentSelection::Predicted, @@ -1943,16 +2460,16 @@ mod tests { (None, 0, "no_fragment_traces") ); assert_eq!( - select_fragment_areas(Some(&[(0.0, 1.0)]), 3, sel), + select_fragment_areas(Some((&zero_a, &zero_p)), 3, sel), (None, 0, "no_positive_fragment_area") ); assert_eq!( - select_fragment_areas(Some(&areas), 0, sel), + select_fragment_areas(areas, 0, sel), (None, 0, "no_fragments_selected") ); // Non-finite areas are dropped, not summed into a NaN quantity. assert_eq!( - select_fragment_areas(Some(&[(f64::NAN, 1.0), (10.0, 0.5)]), 3, sel), + select_fragment_areas(Some((&nan_a, &nan_p)), 3, sel), (Some(10.0), 1, "quantified") ); } @@ -2191,14 +2708,276 @@ mod tests { // `total_cmp` orders NaN above every real value, so in the descending sort a NaN // predicted intensity reached the front of the ranking and was preferentially // summed into the top N. - let areas: [(f64, f32); 3] = [(10.0, f32::NAN), (20.0, 0.9), (30.0, 0.8)]; + let pairs: [(f64, f32); 3] = [(10.0, f32::NAN), (20.0, 0.9), (30.0, 0.8)]; + let (a, pr) = split_areas(&pairs); let (q, used, status) = - select_fragment_areas(Some(&areas), 1, FragmentSelection::Predicted); + select_fragment_areas(Some((&a, &pr)), 1, FragmentSelection::Predicted); assert_eq!(status, "quantified"); assert_eq!(used, 1); // The 0.9-intensity fragment wins, not the NaN one. assert_eq!(q, Some(20.0)); } + + #[test] + fn nearest_index_binary_search_matches_the_scan() { + // The search replaces a forward linear scan whose strict `<` kept the FIRST + // minimum, so ties, runs of identical RTs and targets off either end all have to + // land on the same index. + let traces: [&[f32]; 7] = [ + &[], + &[5.0], + &[1.0, 1.0, 5.0], + &[1.0, 3.0, 3.0, 10.0], + &[0.0, 1.0, 2.0, 3.0, 4.0], + &[2.0, 2.0, 2.0], + &[0.0, 0.0, 1.0, 1.0, 1.0, 9.0], + ]; + for rt in traces { + for &t in &[-10.0f64, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 4.5, 9.0, 100.0] { + assert_eq!( + nearest_index(rt, t, false), + nearest_index(rt, t, true), + "rt={rt:?} target={t}" + ); + } + } + } + + #[test] + fn the_chromatogram_store_returns_the_rows_it_was_given() { + let grid: Vec = (0..5).map(|k| k as f32).collect(); + let offset: Vec = (0..5).map(|k| k as f32 + 0.5).collect(); + let mut s = ChromStore::new(); + s.push(7, "y1", 0.25, &grid, &[1.0, 2.0, 3.0, 2.0, 1.0]); + s.push(7, "y2", 0.75, &grid, &[0.0, 1.0, 2.0, 1.0, 0.0]); + s.push(7, "b3", 0.5, &offset, &[9.0; 5]); + s.push(7, "b4", 0.1, &[], &[]); + s.push(8, "y1", 0.6, &grid, &[4.0; 5]); + assert_eq!(s.nrows(), 5); + assert_eq!(s.rt(0), grid.as_slice()); + assert_eq!(s.inten(1), &[0.0, 1.0, 2.0, 1.0, 0.0]); + assert_eq!(s.name(2), "b3"); + assert_eq!(s.pred[4], 0.6); + assert_eq!(s.rt(3), &[] as &[f32]); + assert_eq!(s.inten(3), &[] as &[f32]); + // One stored axis per distinct grid within a candidate, and the interned name + // table holds one string per distinct name. + assert_eq!(s.axis_id[0], s.axis_id[1]); + assert_ne!(s.axis_id[0], s.axis_id[2]); + assert_eq!(s.axis_id[3], NO_AXIS); + assert_eq!(s.names.names.len(), 4); + // A new candidate starts a new axis even for identical values: that costs a copy, + // never a value. + assert_ne!(s.axis_id[0], s.axis_id[4]); + assert_eq!(s.rt(4), grid.as_slice()); + assert!(s.rt_sorted); + assert!(s.axis_strict[s.axis_id[0] as usize]); + } + + #[test] + fn an_unsorted_or_duplicated_axis_refuses_the_fast_paths() { + let mut descending = ChromStore::new(); + descending.push(1, "y1", 0.0, &[3.0, 1.0, 2.0], &[1.0, 2.0, 3.0]); + assert!( + !descending.rt_sorted, + "a descending step must send every nearest-sample search back to the scan" + ); + assert!(!descending.axis_strict[0]); + + let mut duplicated = ChromStore::new(); + duplicated.push(1, "y1", 0.0, &[1.0, 1.0, 2.0], &[1.0, 2.0, 3.0]); + assert!( + duplicated.rt_sorted, + "non-decreasing is all the nearest-sample search needs" + ); + assert!( + !duplicated.axis_strict[0], + "a repeated RT must merge into one profile point, so the axis cannot be \ + used as the union axis" + ); + + let mut negative = ChromStore::new(); + negative.push(1, "y1", 0.0, &[-1.0, 2.0], &[1.0, 2.0]); + assert!( + !negative.rt_sorted, + "bit order is value order only above zero" + ); + } + + #[test] + fn a_repeated_rt_sample_merges_into_one_profile_point() { + // The union used to be a BTreeMap keyed on the RT bit pattern, so two samples at + // the same RT became ONE profile point carrying their sum. Expressed as two rows + // the same two samples must give exactly the same window. + let duplicated = peak_window( + &[0], + &store_of( + &[vec![0.0f32, 1.0, 1.0, 2.0, 3.0]], + &[vec![0.0f32, 4.0, 6.0, 0.0, 0.0]], + ), + 1.0 / 6.0, + 1, + None, + ); + let split = peak_window( + &[0, 1], + &store_of( + &[vec![0.0f32, 1.0, 2.0, 3.0], vec![1.0f32]], + &[vec![0.0f32, 4.0, 0.0, 0.0], vec![6.0f32]], + ), + 1.0 / 6.0, + 1, + None, + ); + assert_eq!(duplicated, split); + assert_eq!(duplicated.2, 1.0, "the merged point is the apex"); + } + + #[test] + fn the_candidate_index_groups_exactly_as_the_map_it_replaces() { + for cids in [ + vec![3u32, 3, 3, 1, 1, 9], // grouped by candidate, as extract writes it + vec![3u32, 1, 3, 9, 1, 3], // interleaved: the run cache cannot shortcut this + vec![5u32], + ] { + let mut s = ChromStore::new(); + for (r, &c) in cids.iter().enumerate() { + s.push(c, &format!("y{r}"), r as f32, &[r as f32], &[1.0]); + } + let index = CandIndex::build(&s); + let mut want: BTreeMap> = BTreeMap::new(); + for (r, &c) in cids.iter().enumerate() { + want.entry(c).or_default().push(r); + } + let got: BTreeMap> = (0..index.len()) + .map(|ci| (index.cids[ci], index.rows_of(ci).to_vec())) + .collect(); + assert_eq!(got, want, "grouping of {cids:?}"); + for ci in 0..index.len() { + assert_eq!(index.find(index.cids[ci]), Some(ci)); + for slot in index.slots(ci) { + assert_eq!(index.slot_pred[slot], s.pred[index.rows[slot]]); + } + } + assert_eq!(index.find(u32::MAX), None); + } + } + + #[test] + fn several_scored_rows_of_one_candidate_agree_and_ms1_traces_stay_out() { + // The top-N selection is memoised per candidate, so two scored rows on candidate 1 + // must carry exactly what the per-row computation produced. The `ms1_*` traces are + // precursor channels, not fragment ions: here they are ten times brighter than any + // fragment, so including one would move both the window and the sum. Candidate 3 + // has no scored row at all and reaches only the peak-bounds diagnostic, which is + // the export that keeps every candidate's chromatogram rows. + let scored = quant_test_path("memo_scored.parquet"); + let chrom = quant_test_path("memo_chrom.parquet"); + let peptide = quant_test_path("memo_peptide.parquet"); + let protein = quant_test_path("memo_protein.parquet"); + let fragment = quant_test_path("memo_fragment.parquet"); + let bounds = quant_test_path("memo_bounds.parquet"); + + write_table( + &scored, + vec![ + Col::U32("candidate_id".into(), vec![1, 1, 2]), + Col::U32("base_peptide_id".into(), vec![10, 11, 12]), + Col::Str( + "peptidoform".into(), + vec!["PEPA".into(), "PEPB".into(), "PEPC".into()], + ), + Col::I32("charge".into(), vec![2, 3, 2]), + Col::Str("label".into(), vec!["target".into(); 3]), + Col::Str( + "protein_group".into(), + vec!["PG".into(), "PG".into(), "PG2".into()], + ), + Col::F64("peptide_q_value".into(), vec![0.0; 3]), + Col::F64("apex_rt".into(), vec![5.0; 3]), + ], + ) + .unwrap(); + + let grid: Vec = (0..10).map(|rt| rt as f32).collect(); + let c1a = vec![0.0f32, 50.0, 0.0, 0.0, 5.0, 10.0, 5.0, 0.0, 0.0, 0.0]; + let c1b = vec![0.0f32, 25.0, 0.0, 0.0, 2.5, 5.0, 2.5, 0.0, 0.0, 0.0]; + let c2 = vec![0.0f32, 0.0, 0.0, 0.0, 10.0, 20.0, 10.0, 0.0, 0.0, 0.0]; + let ms1 = vec![100.0f32; 10]; + write_table( + &chrom, + vec![ + Col::U32("candidate_id".into(), vec![1, 1, 1, 2, 3]), + Col::Str( + "frag_name".into(), + vec![ + "b2".into(), + "ms1_mono".into(), + "y3".into(), + "b4".into(), + "y5".into(), + ], + ), + Col::ListF32( + "rt".into(), + vec![ + grid.clone(), + grid.clone(), + grid.clone(), + grid.clone(), + grid.clone(), + ], + ), + Col::ListF32("intensity".into(), vec![c1a, ms1, c1b, c2.clone(), c2]), + ], + ) + .unwrap(); + + let cfg = QuantConfig::default(); + let rows = run(QuantParams { + psms_scored: &scored, + chromatograms: &chrom, + out_peptide: &peptide, + out_protein: &protein, + out_fragment: Some(&fragment), + out_peak_bounds: Some(&bounds), + cfg: &cfg, + config_hash: "test", + }) + .unwrap(); + assert_eq!(rows, (3, 2)); + + let pq = Table::read(&peptide).unwrap(); + let quantities = pq.opt_f64("quantity").unwrap(); + // b2 + y3 over the walked window [4, 6]: (5+10)/2 + (10+5)/2 = 15, and half that. + assert_eq!(quantities[0], Some(22.5)); + assert_eq!( + quantities[0], quantities[1], + "the two rows of candidate 1 must quantify identically" + ); + assert_eq!(quantities[2], Some(30.0)); + assert_eq!(pq.i32("n_fragments_used").unwrap(), vec![2, 2, 1]); + assert_eq!( + pq.opt_f64("integration_lo_rt").unwrap(), + vec![Some(4.0); 3], + "the ms1 channel must not widen the window" + ); + + let fq = Table::read(&fragment).unwrap(); + let names = fq.str("fragment_name").unwrap(); + assert_eq!(names.len(), 5, "two fragments x two rows, plus candidate 2"); + assert!( + !names.iter().any(|n| n.starts_with("ms1_")), + "ms1 pseudo-traces are not fragment quantities: {names:?}" + ); + + let bq = Table::read(&bounds).unwrap(); + assert_eq!( + bq.u32("candidate_id").unwrap(), + vec![1, 2, 3], + "the diagnostic keeps candidates no scored row selected" + ); + } } /// Distinct values of the scored table's `source` column, or `None` when the column is From c85052203c4e9ecfbb78ccf31e42c4d8fc8143b7 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:33:59 +0200 Subject: [PATCH 13/35] perf(rescore): stream the competed features into the sidecar handoff instead of materialising them Under `rescore.strict` with a PIN sidecar the engine's own `FeatureMatrix` is never read. The handoff file is written FROM it, it is released immediately after, and strict turns a sidecar failure into an error rather than a fall back to `native_tda`, so the native path that would read it is unreachable. It was still materialised in full first -- one allocation of rows x features x 4 bytes, about 250 GB at experiment scale (203M-precursor library, ~11.6M PSMs x 387 features and up) -- purely so that it could be copied to disk column by column. The competed feature columns now go straight into the handoff. What stays resident is one staging block of HANDOFF_BATCH_ROWS rows, about 390 MB at 387 features, plus its transposed columns while a block is encoded. That is the largest single allocation in the stage removed outright, and with it the largest contributor to the mapping count the OOM at ~290 GB resident was traced to. How it is arranged: - the load is two passes. Pass 1 reads the per-PSM metadata columns of every input; pass 2 reads the feature values, either into the matrix (`load_feature_matrix`) or through `for_each_feature_row` into the handoff. Splitting them is what makes the feature pass skippable, and it keeps the order of the population checks: labels and the target/decoy population are still validated before anything reads a feature value. It costs one extra parquet footer read per input. - `HandoffWriter` writes either encoding a row at a time, so the same writer serves the matrix path and the streamed path. The parquet blocks are the same 250k-row blocks, the row groups the same 131,072 rows, the values narrowed through the same `as f32` (including `f64::NAN as f32` for a null cell), so the file the worker reads is unchanged. A test writes both ways over two inputs with a batch small enough to force several blocks and compares the bytes. - the non-finite feature validation moves into the streaming pass, inline with the narrowing while the row is still in cache. `bail_non_finite` then picks between it and the scalar scan on the row index, so the error is still the first offending row, feature before scalar, with the same message. - `sidecar_paths` names the per-invocation files once; `run_pin_sidecar` takes the already-written handoff when there is one. A guard refuses a handoff whose row count disagrees with the metadata: the two passes cannot disagree, but if they ever did the handoff would be row-misaligned with every label. Equality: the streamed and matrix-built handoffs are byte-identical (tested, both encodings); the scored artifact of the crafted features -> compete -> rescore pipeline keeps the same blake3 (e535d04b2821a47d4fc8f9bc0a7ba57461b701b607496ffef0700a0352cd135a) as before this branch. `rescore.max_feature_matrix_gib` is deliberately left applying to the streamed path too: the worker still builds a matrix of exactly that size, so the ceiling still describes a real allocation. Residual risk: the streamed path is covered at unit level (row values equal to the matrix rows, handoff bytes equal, and `run` driven end to end with an unspawnable interpreter so the handoff is written and then strict fails), but no test in this repository runs a real Python sidecar, so the worker's side of the contract is unexercised here. It reads the same bytes it read before. Co-Authored-By: Claude Fable 5.1 --- .../crates/mumdia/src/stages/rescore.rs | 939 ++++++++++++++---- 1 file changed, 732 insertions(+), 207 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/rescore.rs b/rust/mumdia/crates/mumdia/src/stages/rescore.rs index 750b7ea9..31c982d8 100644 --- a/rust/mumdia/crates/mumdia/src/stages/rescore.rs +++ b/rust/mumdia/crates/mumdia/src/stages/rescore.rs @@ -219,7 +219,12 @@ pub fn run(p: RescoreParams) -> Result { feat_names.len(), ); } - let mut matrix = FeatureMatrix::with_capacity(total_rows, feat_names.len()); + // Pass 1: the per-PSM metadata columns of every input. Reading them in their own pass, + // separately from the ~390 feature columns, is what lets the feature pass below be + // skipped entirely, and it preserves the order of the checks: the labels and the + // target/decoy population are validated before anything reads a feature value, as they + // were when both were read together. The cost is one extra parquet footer read per + // input. for (src, path) in p.competed.iter().enumerate() { let actual_schema = FeatureSchema::read(path)?; validate_feature_schema(&expected_schema, &actual_schema, path)?; @@ -236,49 +241,6 @@ pub fn run(p: RescoreParams) -> Result { let elo = t.f64("elution_lo")?; let ehi = t.f64("elution_hi")?; let pkr = t.i32("peak_rank").unwrap_or_else(|_| vec![0; t.nrows]); - // Feature values: ONE streaming pass over just the feature columns, appending each - // row's values contiguously into the flat matrix. The previous path materialised all - // ~390 columns as owned `Vec`s first (`fcols`) and then built a `Vec` per PSM from - // them, so the whole matrix existed twice, plus the Arrow batches of the full table. - // Values, and the null policy (a null f64 reads as NaN), are unchanged. - { - let names: Vec<&str> = feat_names.iter().map(String::as_str).collect(); - let reader = t.batches(Some(&names), FEATURE_BATCH_ROWS)?; - let sch = reader.schema(); - let order: Vec = feat_names - .iter() - .map(|n| { - sch.index_of(n) - .map_err(|_| anyhow!("competed table {path} has no feature column '{n}'")) - }) - .collect::>()?; - for b in reader { - let b = b?; - let cols: Vec<&Float64Array> = order - .iter() - .map(|&i| { - b.column(i) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - anyhow!("feature column '{}' is not f64", sch.field(i).name()) - }) - }) - .collect::>()?; - let any_null = cols.iter().any(|c| c.null_count() > 0); - for k in 0..b.num_rows() { - if any_null { - for c in &cols { - matrix.push(if c.is_null(k) { f64::NAN } else { c.value(k) }); - } - } else { - for c in &cols { - matrix.push(c.value(k)); - } - } - } - } - } // Take the reader's columns rather than copying them row by row into a second set. // The row loop this replaces cloned every metadata STRING -- label, peptidoform, // protein -- into a fresh allocation while the source vector stayed alive until the @@ -304,61 +266,105 @@ pub fn run(p: RescoreParams) -> Result { merge_col(&mut elution_hi, ehi, total_rows); merge_col(&mut source, vec![src as u32; t.nrows], total_rows); } - let feats = matrix.finish()?; - crate::memlog::report( - "rescore feature matrix", - &[ - ("feats", feats.bytes()), - ( - "metadata_columns", - meta_bytes(&cid, &label, &pform, &protein), - ), - ], - ); crate::fdr::validate_labels(&label)?; let is_decoy: Vec = label.iter().map(|l| l == "decoy").collect(); let (mut is_entrapment, mut is_real_target) = classify_entrapment(p.cfg, &protein, &is_decoy); let mut is_decoy = is_decoy; let mut n = cid.len(); + // First row with a non-finite prelim_score/precursor_mz, if any. Found here but NOT + // reported here: the serial validation this replaces checked a row's features before + // its scalars and stopped at the first offending ROW, so the two answers are combined + // once the feature scan has run (`bail_non_finite`). + let mut bad_scalar = None; if n > 0 { let n_decoys = is_decoy.iter().filter(|&&v| v).count(); let n_targets = n - n_decoys; require_both_labels(n_targets, n_decoys)?; - // The validation this replaces was one serial pass of `n_rows x n_features` - // `is_finite` checks -- hundreds of gigabytes on one thread at experiment scale -- - // for a predicate with no carried state. Both scans are parallel now and both are - // order-preserving (`find_first`), and the two answers are combined so the row - // reported is still the FIRST offending row, features before scalars within a row, - // exactly as the interleaved serial loop reported it. - let bad_feature = feats.find_non_finite(); - let bad_scalar = (0..n) + bad_scalar = (0..n) .into_par_iter() .find_first(|&row| !prelim[row].is_finite() || !mz[row].is_finite()); - if let Some((row, feature)) = bad_feature { - if bad_scalar.is_none_or(|scalar_row| row <= scalar_row) { - anyhow::bail!( - "rescore input contains non-finite feature '{}' at flat row {row}: {}", - feat_names[feature], - feats.row(row)[feature] - ); + } + + // Pass 2: the feature values, either into the engine's own matrix or straight through + // to the sidecar's handoff file. + // + // Under `rescore.strict` with a PIN sidecar, nothing in this process ever reads that + // matrix. The handoff file is written FROM it and it is released immediately after + // (that release is itself measured: docs/28), and strict turns a sidecar failure into + // an error rather than a fall back to `native_tda`, so the native path that would read + // it is unreachable. It was still materialised in full first -- one allocation of + // `rows x features x 4` bytes, ~250 GB at experiment scale -- purely so that it could + // be copied to disk. Streaming the competed feature columns straight into the handoff + // removes it: what stays resident is one staging block of `HANDOFF_BATCH_ROWS` rows + // (about 390 MB at 387 features), and the non-finite validation moves into the same + // pass, inline with the narrowing, while the row is still in cache. + let sidecar_script = match p.cfg.classifier { + RescorerKind::Mokapot => Some("mokapot_worker.py"), + RescorerKind::NnTorch => Some("nn_rescore_worker.py"), + _ => None, + }; + // `python.is_some()` belongs in this condition: without an interpreter the sidecar + // fails before it reads anything, and a streamed handoff would be work thrown away. + let stream_to_handoff = + n > 0 && p.cfg.strict && p.cfg.python.is_some() && sidecar_script.is_some(); + let mut feats_slot: Option = None; + let mut prewritten: Option = None; + if stream_to_handoff { + let paths = sidecar_paths(&p, sidecar_script.expect("checked in the condition")); + let mut w = HandoffWriter::new(&paths, &feat_names, &label, &pform, &protein, &mz)?; + let mut bad_feature = None; + for_each_feature_row(p.competed, &feat_names, |row, values| { + if bad_feature.is_none() { + if let Some(col) = values.iter().position(|v| !v.is_finite()) { + bad_feature = Some((row, col, values[col])); + } } - } - if let Some(row) = bad_scalar { + w.push_row(row, values) + })?; + let rows = w.finish()?; + // The metadata and the feature values now come from two separate passes over the + // same files. Nothing should be able to make them disagree, and if anything does, + // the handoff is row-misaligned with every metadata column: refuse rather than + // train on it. + if rows as usize != n { anyhow::bail!( - "rescore input contains non-finite prelim_score/precursor_mz at flat row {row}" + "rescore: the streamed handoff wrote {rows} rows but the competed metadata \ + has {n}; the two passes disagree about the input" ); } + info!( + path = %paths.handoff, + rows, + features = feat_names.len(), + "rescore: streamed the competed features into the sidecar handoff \ + (no engine-side feature matrix)" + ); + bail_non_finite(bad_feature, bad_scalar, &feat_names)?; + prewritten = Some(paths); + } else { + let feats = load_feature_matrix(p.competed, &feat_names, total_rows)?; + let bad_feature = feats + .find_non_finite() + .map(|(row, col)| (row, col, feats.row(row)[col])); + bail_non_finite(bad_feature, bad_scalar, &feat_names)?; + feats_slot = Some(feats); } + crate::memlog::report( + "rescore feature matrix", + &[ + ("feats", feats_slot.as_ref().map_or(0, |f| f.bytes())), + ( + "metadata_columns", + meta_bytes(&cid, &label, &pform, &protein), + ), + ], + ); info!( psms = n, features = feat_names.len(), folds = p.cfg.folds, "rescore: loaded competed PSMs" ); - // From here on the matrix is optional. A sidecar run under `rescore.strict` releases it - // as soon as the worker has its own copy on disk (`run_pin_sidecar`), because no native - // fallback can follow; every native path reads it through `kept_matrix`. - let mut feats_slot = Some(feats); // Track the path actually taken so the report reflects reality rather than a // hardcoded label, and pick the null the q-values are computed against. @@ -383,6 +389,7 @@ pub fn run(p: RescoreParams) -> Result { &mz, &mut feats_slot, &base, + prewritten.take(), ) { Ok(s) => { info!("rescore: using Mokapot scores"); @@ -413,6 +420,7 @@ pub fn run(p: RescoreParams) -> Result { &mz, &mut feats_slot, &base, + prewritten.take(), ) { Ok(s) => { info!("rescore: using PyTorch NN sidecar scores"); @@ -1288,60 +1296,192 @@ impl Drop for ChildGuard { /// /// Batched because `feats` is already resident (one flat row-major matrix, 8 bytes per /// value); materialising 387 full columns as well would add ~12.8 GB for nothing. -/// The feature matrix, which is released only under `rescore.strict` and only after a -/// PIN sidecar has its own copy; strict has no native fallback, so no reader is left. +/// The feature matrix, for the paths that read it. +/// +/// It is absent in exactly two situations, both of them `rescore.strict` with a PIN +/// sidecar: it was streamed past (never built) or released once the worker had its own copy +/// on disk. Strict has no native fallback, so neither situation leaves a reader; every +/// caller of this function is on a non-strict or native path, where the matrix was built. fn kept_matrix(slot: &Option) -> &FeatureMatrix { slot.as_ref().expect( - "the feature matrix is released only under rescore.strict, after the sidecar handoff", + "the feature matrix is absent only under rescore.strict with a sidecar classifier, \ + which has no native fallback", ) } -fn write_features_parquet( - path: &str, - feat_names: &[String], - label: &[String], - pform: &[String], - protein: &[String], - mz: &[f64], - feats: &FeatureMatrix, -) -> Result { - use arrow::array::{ArrayRef, Float32Array, Float64Array, Int32Array, StringArray}; - use arrow::datatypes::{DataType, Field, Schema}; - use arrow::record_batch::RecordBatch; - use std::sync::Arc; +/// ~250k rows x 387 f32 is about 390 MB per batch, which keeps the encoder's working set +/// modest. +const HANDOFF_BATCH_ROWS: usize = 250_000; +/// Row groups are capped well below the batch: the worker reads this file back with +/// `ParquetFile.iter_batches`, which decodes a whole row group before it slices batches out +/// of it, so at parquet's default 1,048,576-row groups the load transient was 1.6 GB of +/// Arrow buffers per group on top of the matrix. 131,072 rows x 387 f32 is 200 MB. +const HANDOFF_ROW_GROUP_ROWS: usize = 131_072; + +/// The parquet encoding's state. Boxed in `HandoffSink` because the arrow writer it owns +/// dwarfs a `BufWriter`. +struct ParquetHandoff { + schema: std::sync::Arc, + writer: mumdia_io::table::BatchWriter, + /// Row-major staging buffer for the current block, transposed to columns on flush. + stage: Vec, + block_start: usize, +} - let mut fields: Vec = vec![ - Field::new("SpecId", DataType::Utf8, false), - Field::new("Label", DataType::Int32, false), - Field::new("ScanNr", DataType::Int32, false), - Field::new("ExpMass", DataType::Float64, false), - Field::new("CalcMass", DataType::Float64, false), - ]; - for n in feat_names { - fields.push(Field::new(n, DataType::Float32, false)); +enum HandoffSink { + Parquet(Box), + Pin(std::io::BufWriter), +} + +/// The sidecar handoff file, written one row at a time. +/// +/// Same logical contract in both encodings - `SpecId`, `Label`, `ScanNr`, `ExpMass`, +/// `CalcMass`, the feature columns, `Peptide`, `Proteins` - so the worker reads either +/// format from the same column names. Parquet applies to `nn_torch` only; `mokapot_worker` +/// goes through `mokapot.read_pin()`, which requires the tab-separated form. +/// +/// Features are `f32` in the parquet: the TSV writes `{:.6}` and the worker casts to +/// float32 anyway, so f32 matches what is actually used, halves the file, and does not +/// silently increase precision relative to the validated TSV reference. +/// +/// Row-at-a-time, rather than a function over a resident matrix, because the source is +/// either that matrix or the competed parquet itself (`for_each_feature_row`) and neither +/// caller should hold a second copy. The parquet blocks are the same +/// `HANDOFF_BATCH_ROWS`-row blocks written before, with the same values in the same order, +/// so the file is unchanged. +struct HandoffWriter<'a> { + sink: HandoffSink, + /// Rows per parquet batch. `HANDOFF_BATCH_ROWS` outside tests. + block_rows: usize, + /// Rows accepted so far, so `finish` reports what was written rather than what the + /// caller's metadata columns happen to be long. + rows: u64, + feat_names: &'a [String], + label: &'a [String], + pform: &'a [String], + protein: &'a [String], + mz: &'a [f64], +} + +impl<'a> HandoffWriter<'a> { + fn new( + paths: &SidecarPaths, + feat_names: &'a [String], + label: &'a [String], + pform: &'a [String], + protein: &'a [String], + mz: &'a [f64], + ) -> Result> { + use arrow::datatypes::{DataType, Field, Schema}; + use std::io::Write as _; + use std::sync::Arc; + + let sink = if paths.use_pq { + let mut fields: Vec = vec![ + Field::new("SpecId", DataType::Utf8, false), + Field::new("Label", DataType::Int32, false), + Field::new("ScanNr", DataType::Int32, false), + Field::new("ExpMass", DataType::Float64, false), + Field::new("CalcMass", DataType::Float64, false), + ]; + for n in feat_names { + fields.push(Field::new(n, DataType::Float32, false)); + } + fields.push(Field::new("Peptide", DataType::Utf8, false)); + fields.push(Field::new("Proteins", DataType::Utf8, false)); + let schema = Arc::new(Schema::new(fields)); + let writer = mumdia_io::table::BatchWriter::with_row_group_rows( + &paths.handoff, + schema.clone(), + HANDOFF_ROW_GROUP_ROWS, + )?; + HandoffSink::Parquet(Box::new(ParquetHandoff { + schema, + writer, + stage: Vec::with_capacity(HANDOFF_BATCH_ROWS.saturating_mul(feat_names.len())), + block_start: 0, + })) + } else { + // Streamed through a BufWriter. It was previously accumulated in ONE + // un-reserved String: at ~1M rows x 387 features that is a >5 GB allocation + // (plus realloc churn) held entirely in RAM before the first byte reaches disk. + let mut w = + std::io::BufWriter::with_capacity(1 << 20, std::fs::File::create(&paths.handoff)?); + w.write_all(b"SpecId\tLabel\tScanNr\tExpMass\tCalcMass\t")?; + w.write_all(feat_names.join("\t").as_bytes())?; + w.write_all(b"\tPeptide\tProteins\n")?; + HandoffSink::Pin(w) + }; + Ok(HandoffWriter { + sink, + block_rows: HANDOFF_BATCH_ROWS, + rows: 0, + feat_names, + label, + pform, + protein, + mz, + }) } - fields.push(Field::new("Peptide", DataType::Utf8, false)); - fields.push(Field::new("Proteins", DataType::Utf8, false)); - let schema = Arc::new(Schema::new(fields)); - - let n = label.len(); - let nf = feat_names.len(); - // ~250k rows x 387 f32 is about 390 MB per batch, which keeps the encoder's working set - // modest. Row groups are capped well below that: the worker reads this file back with - // `ParquetFile.iter_batches`, which decodes a whole row group before it slices batches - // out of it, so at parquet's default 1,048,576-row groups the load transient was 1.6 GB - // of Arrow buffers per group on top of the matrix. 131,072 rows x 387 f32 is 200 MB. - const BATCH: usize = 250_000; - const ROW_GROUP_ROWS: usize = 131_072; - let mut w = - mumdia_io::table::BatchWriter::with_row_group_rows(path, schema.clone(), ROW_GROUP_ROWS)?; - let mut start = 0usize; - while start < n { - let end = (start + BATCH).min(n); - let k = end - start; + + /// Shrink the parquet batch so a test can exercise several blocks. Production always + /// uses `HANDOFF_BATCH_ROWS`. + #[cfg(test)] + fn with_block_rows(mut self, rows: usize) -> Self { + self.block_rows = rows.max(1); + self + } + + /// Append flat row `i`. Callers push rows in ascending order from 0. + /// + /// SpecId / ScanNr key on the row index, NOT candidate_id: candidate_id is the library + /// index and repeats across runs, so an experiment-wide table would collide. + fn push_row(&mut self, i: usize, values: &[f32]) -> Result<()> { + use std::io::Write as _; + let nf = self.feat_names.len(); + self.rows += 1; + match &mut self.sink { + HandoffSink::Parquet(pq) => { + pq.stage.extend_from_slice(&values[..nf]); + if pq.stage.len() >= self.block_rows.saturating_mul(nf.max(1)) { + self.flush_block()?; + } + Ok(()) + } + HandoffSink::Pin(w) => { + let lab = if self.label[i] == "decoy" { -1 } else { 1 }; + write!( + w, + "psm_{}\t{}\t{}\t{:.5}\t{:.5}\t", + i, lab, i, self.mz[i], self.mz[i] + )?; + for v in values.iter().take(nf) { + write!(w, "{:.6}\t", v)?; + } + writeln!(w, "-.{}.-\t{}", self.pform[i], self.protein[i])?; + Ok(()) + } + } + } + + /// Transpose the staged rows into one column per field and write them as one batch. + fn flush_block(&mut self) -> Result<()> { + use arrow::array::{ArrayRef, Float32Array, Int32Array, StringArray}; + use arrow::record_batch::RecordBatch; + use std::sync::Arc; + + let nf = self.feat_names.len(); + let (label, pform, protein, mz) = (self.label, self.pform, self.protein, self.mz); + let HandoffSink::Parquet(pq) = &mut self.sink else { + return Ok(()); + }; + if pq.stage.is_empty() { + return Ok(()); + } + let start = pq.block_start; + let k = pq.stage.len() / nf.max(1); + let end = start + k; let mut arrays: Vec = Vec::with_capacity(nf + 7); - // SpecId / ScanNr key on the row index, NOT candidate_id: candidate_id is the library - // index and repeats across runs, so an experiment-wide table would collide. arrays.push(Arc::new(StringArray::from( (start..end).map(|i| format!("psm_{i}")).collect::>(), ))); @@ -1356,57 +1496,156 @@ fn write_features_parquet( let mzv: Vec = mz[start..end].to_vec(); arrays.push(Arc::new(Float64Array::from(mzv.clone()))); arrays.push(Arc::new(Float64Array::from(mzv))); - // Transpose this row block into one column per feature. `feats` is row-major, so a - // column read strides it; doing that per batch keeps the working set to one block. for fi in 0..nf { - let col: Vec = (start..end).map(|i| feats.row(i)[fi]).collect(); + let col: Vec = (0..k).map(|r| pq.stage[r * nf + fi]).collect(); arrays.push(Arc::new(Float32Array::from(col))); } - let _ = k; arrays.push(Arc::new(StringArray::from( (start..end) .map(|i| format!("-.{}.-", pform[i])) .collect::>(), ))); arrays.push(Arc::new(StringArray::from(protein[start..end].to_vec()))); - w.write(&RecordBatch::try_new(schema.clone(), arrays)?)?; - start = end; + pq.writer + .write(&RecordBatch::try_new(pq.schema.clone(), arrays)?)?; + pq.stage.clear(); + pq.block_start = end; + Ok(()) + } + + /// Finish the file and return the rows written. + fn finish(mut self) -> Result { + use std::io::Write as _; + self.flush_block()?; + match self.sink { + HandoffSink::Parquet(pq) => pq.writer.close(), + HandoffSink::Pin(mut w) => { + w.flush()?; + Ok(self.rows) + } + } } - w.close() } -/// Run a PIN-contract Python rescorer sidecar (`mokapot_worker.py` or -/// `nn_rescore_worker.py`) over a PIN written from the competed set; return scores -/// aligned to the input candidate order (the file contract in -/// docs/13_sidecars.md). Both sidecars share this exact contract: PIN in, -/// `candidate_id`+`score` parquet out. -#[allow(clippy::too_many_arguments)] -fn run_pin_sidecar( - p: &RescoreParams, - script_name: &str, +/// Stream the selected feature columns of every competed input, one row at a time, in flat +/// row order, narrowed to f32 exactly as `FeatureMatrix::push` narrows them. +/// +/// One pass over just the feature columns. The null policy is unchanged: a null f64 reads +/// as NaN, which the validation then rejects. +fn for_each_feature_row( + competed: &[String], feat_names: &[String], - cid: &[u32], - label: &[String], - pform: &[String], - protein: &[String], - mz: &[f64], - feats: &mut Option, - base: &[u32], -) -> Result> { - use std::io::Write as _; - let python = p.cfg.python.as_deref().ok_or_else(|| { - anyhow::anyhow!("classifier sidecar {script_name} requires rescore.python") - })?; - let matrix = feats.as_ref().ok_or_else(|| { - anyhow::anyhow!("rescore: the feature matrix was released before the {script_name} handoff") + mut f: impl FnMut(usize, &[f32]) -> Result<()>, +) -> Result<()> { + let names: Vec<&str> = feat_names.iter().map(String::as_str).collect(); + let mut row: Vec = vec![0.0; feat_names.len()]; + let mut flat = 0usize; + for path in competed { + let t = TableFile::open(path)?; + let reader = t.batches(Some(&names), FEATURE_BATCH_ROWS)?; + let sch = reader.schema(); + let order: Vec = feat_names + .iter() + .map(|n| { + sch.index_of(n) + .map_err(|_| anyhow!("competed table {path} has no feature column '{n}'")) + }) + .collect::>()?; + for b in reader { + let b = b?; + let cols: Vec<&Float64Array> = order + .iter() + .map(|&i| { + b.column(i) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + anyhow!("feature column '{}' is not f64", sch.field(i).name()) + }) + }) + .collect::>()?; + let any_null = cols.iter().any(|c| c.null_count() > 0); + for k in 0..b.num_rows() { + if any_null { + for (slot, c) in row.iter_mut().zip(&cols) { + // `f64::NAN as f32`, not `f32::NAN`: the same expression the + // matrix path narrowed through, so a null cell keeps its bits. + *slot = if c.is_null(k) { + f64::NAN as f32 + } else { + c.value(k) as f32 + }; + } + } else { + for (slot, c) in row.iter_mut().zip(&cols) { + *slot = c.value(k) as f32; + } + } + f(flat, &row)?; + flat += 1; + } + } + } + Ok(()) +} + +/// The engine's own copy of the feature values: one contiguous f32 buffer. +fn load_feature_matrix( + competed: &[String], + feat_names: &[String], + total_rows: usize, +) -> Result { + let mut m = FeatureMatrix::with_capacity(total_rows, feat_names.len()); + for_each_feature_row(competed, feat_names, |_, values| { + m.push_row(values); + Ok(()) })?; + m.finish() +} + +/// Report the first non-finite input value, whichever kind of column it is in. +/// +/// The validation these two answers come from used to be one serial row loop that checked a +/// row's features before its scalars and stopped at the first offending ROW. Both scans run +/// separately now (one of them inside the handoff stream), so the choice between them is +/// made here, on the row index, to keep the message identical. +fn bail_non_finite( + bad_feature: Option<(usize, usize, f32)>, + bad_scalar: Option, + feat_names: &[String], +) -> Result<()> { + if let Some((row, feature, value)) = bad_feature { + if bad_scalar.is_none_or(|scalar_row| row <= scalar_row) { + anyhow::bail!( + "rescore input contains non-finite feature '{}' at flat row {row}: {value}", + feat_names[feature] + ); + } + } + if let Some(row) = bad_scalar { + anyhow::bail!( + "rescore input contains non-finite prelim_score/precursor_mz at flat row {row}" + ); + } + Ok(()) +} + +/// The per-invocation sidecar file paths, and which handoff encoding they name. +struct SidecarPaths { + /// The feature handoff: a parquet for `nn_torch`, the tab-separated PIN otherwise. + handoff: String, + out: String, + foldkeys: String, + use_pq: bool, +} + +/// Per-invocation sidecar filenames. Fixed names (`rescore.pin`, +/// `rescore_sidecar_out.parquet`) made two concurrent rescores clobber each other and let a +/// killed run's orphaned Python worker hold `*.feat.mm` open forever, so every later +/// rescore failed with `OSError: [Errno 22]` on a path it did not own. Keying on the output +/// artifact plus the PID makes collisions impossible. +fn sidecar_paths(p: &RescoreParams, script_name: &str) -> SidecarPaths { std::fs::create_dir_all(p.work_dir).ok(); - // Per-invocation sidecar filenames. Fixed names (`rescore.pin`, - // `rescore_sidecar_out.parquet`) made two concurrent rescores clobber each other and - // let a killed run's orphaned Python worker hold `*.feat.mm` open forever, so every - // later rescore failed with `OSError: [Errno 22]` on a path it did not own. Keying on - // the output artifact plus the PID makes collisions impossible, and the guard below - // reaps the child so a killed parent cannot orphan it. let tag = format!( "{}_{}", std::path::Path::new(p.out) @@ -1427,48 +1666,77 @@ fn run_pin_sidecar( a PIN through mokapot.read_pin); writing the tab-separated PIN instead" ); } - let pin = if use_pq { + let handoff = if use_pq { format!("{}/rescore_{tag}.features.parquet", p.work_dir) } else { format!("{}/rescore_{tag}.pin", p.work_dir) }; - let outp = format!("{}/rescore_{tag}_out.parquet", p.work_dir); - - // Stream the PIN through a BufWriter. It was previously accumulated in ONE - // un-reserved String: at ~1M rows x 387 features that is a >5 GB allocation (plus - // realloc churn) held entirely in RAM before the first byte reaches disk. - if use_pq { - let rows = write_features_parquet(&pin, feat_names, label, pform, protein, mz, matrix)?; - tracing::info!( - path = %pin, - rows, - features = feat_names.len(), - "rescore: wrote the sidecar feature table as parquet" - ); - } else { - let mut w = std::io::BufWriter::with_capacity(1 << 20, std::fs::File::create(&pin)?); - w.write_all(b"SpecId\tLabel\tScanNr\tExpMass\tCalcMass\t")?; - w.write_all(feat_names.join("\t").as_bytes())?; - w.write_all(b"\tPeptide\tProteins\n")?; - // Key the PIN on the unique row index i (SpecId=psm_i, ScanNr=i), NOT - // candidate_id: candidate_id is the library index and repeats across runs, so - // an experiment-wide (multi-file) PIN would collide on ScanNr and mokapot's - // per-spectrum competition would collapse the runs. The row index is unique - // across the whole concatenation. Single-run behaviour is unchanged (the - // mapping is bijective and mokapot does not use SpecId/ScanNr as features). - for i in 0..cid.len() { - let lab = if label[i] == "decoy" { -1 } else { 1 }; - write!(w, "psm_{}\t{}\t{}\t{:.5}\t{:.5}\t", i, lab, i, mz[i], mz[i])?; - // one flat matrix row, already in `feat_names` order - let row = matrix.row(i); - for v in row.iter().take(feat_names.len()) { - write!(w, "{:.6}\t", v)?; + SidecarPaths { + handoff, + out: format!("{}/rescore_{tag}_out.parquet", p.work_dir), + foldkeys: format!("{}/rescore_{tag}.foldkeys.parquet", p.work_dir), + use_pq, + } +} + +/// Run a PIN-contract Python rescorer sidecar (`mokapot_worker.py` or +/// `nn_rescore_worker.py`) over a PIN written from the competed set; return scores +/// aligned to the input candidate order (the file contract in +/// docs/13_sidecars.md). Both sidecars share this exact contract: PIN in, +/// `candidate_id`+`score` parquet out. +/// +/// `prewritten` is `Some` when the handoff was already streamed straight from the competed +/// parquet, which is what happens under `rescore.strict` (see `run`); then there is no +/// engine-side matrix and `feats` is `None`. Otherwise the handoff is written here, from +/// the matrix, and released afterwards if strict allows it. +#[allow(clippy::too_many_arguments)] +fn run_pin_sidecar( + p: &RescoreParams, + script_name: &str, + feat_names: &[String], + cid: &[u32], + label: &[String], + pform: &[String], + protein: &[String], + mz: &[f64], + feats: &mut Option, + base: &[u32], + prewritten: Option, +) -> Result> { + let python = p.cfg.python.as_deref().ok_or_else(|| { + anyhow::anyhow!("classifier sidecar {script_name} requires rescore.python") + })?; + let paths = match prewritten { + Some(paths) => paths, + None => { + let paths = sidecar_paths(p, script_name); + let matrix = feats.as_ref().ok_or_else(|| { + anyhow::anyhow!( + "rescore: the feature matrix was released before the {script_name} handoff" + ) + })?; + // Key the handoff on the unique row index i (SpecId=psm_i, ScanNr=i), NOT + // candidate_id: candidate_id is the library index and repeats across runs, so + // an experiment-wide (multi-file) PIN would collide on ScanNr and mokapot's + // per-spectrum competition would collapse the runs. The row index is unique + // across the whole concatenation. Single-run behaviour is unchanged (the + // mapping is bijective and mokapot does not use SpecId/ScanNr as features). + let mut w = HandoffWriter::new(&paths, feat_names, label, pform, protein, mz)?; + for i in 0..cid.len() { + w.push_row(i, matrix.row(i))?; } - writeln!(w, "-.{}.-\t{}", pform[i], protein[i])?; + let rows = w.finish()?; + tracing::info!( + path = %paths.handoff, + rows, + features = feat_names.len(), + parquet = paths.use_pq, + "rescore: wrote the sidecar feature table" + ); + paths } - w.flush()?; - drop(w); - } + }; + let (pin, outp) = (&paths.handoff, &paths.out); // Cross-validation fold keys, row-aligned to the PIN. // @@ -1484,21 +1752,24 @@ fn run_pin_sidecar( // By environment variable rather than argv: the sidecar contract is positional and // shared with `mokapot_worker.py`, and the NN hyperparameters already travel this way. // A worker that does not read it simply keeps its previous behaviour. - let foldkeys = format!("{}/rescore_{tag}.foldkeys.parquet", p.work_dir); - write_table(&foldkeys, vec![Col::U32("fold_key".into(), base.to_vec())])?; + let foldkeys = &paths.foldkeys; + write_table(foldkeys, vec![Col::U32("fold_key".into(), base.to_vec())])?; // Everything the worker reads is on disk now. Under `strict` a sidecar failure is an // error rather than a fall back to native_tda, so nothing downstream reads the engine's // copy of the matrix again: release it before the child starts instead of holding it // idle beside the worker's own for the whole training run. Measured on the six-run // Astral pool (3.13M x 387): the engine sat at 5.3 GB of a 17.9 GB process-tree peak. + // Under strict the matrix is usually never built at all (the handoff was streamed), + // and then there is nothing to release. if p.cfg.strict { - let bytes = matrix.bytes(); - *feats = None; - tracing::info!( - released = %human_bytes(bytes as f64), - "rescore: released the engine's feature matrix for the sidecar run" - ); + if let Some(bytes) = feats.as_ref().map(|m| m.bytes()) { + *feats = None; + tracing::info!( + released = %human_bytes(bytes as f64), + "rescore: released the engine's feature matrix for the sidecar run" + ); + } } let script = crate::sidecar::resolve_script(p.script_dir, script_name); @@ -1508,8 +1779,8 @@ fn run_pin_sidecar( // stale lock with no hint about which PID held it. let child = std::process::Command::new(python) .arg(&script) - .arg(&pin) - .arg(&outp) + .arg(pin) + .arg(outp) .env("PYTHONUTF8", "1") // Pass the configured NN hyperparameters so the worker uses them instead // of its own defaults, and so the folds/num_iter/train_fdr recorded in the @@ -1543,7 +1814,7 @@ fn run_pin_sidecar( .env("MUMDIA_NN_WARM_EPOCHS", p.cfg.train_warm_epochs.to_string()) .env("MUMDIA_NN_MARGIN_FRAC", p.cfg.train_margin_frac.to_string()) .env("MUMDIA_NN_SEEDS", p.cfg.seeds.max(1).to_string()) - .env("MUMDIA_NN_FOLD_KEYS", &foldkeys) + .env("MUMDIA_NN_FOLD_KEYS", foldkeys) .spawn() .map_err(|e| { // A bare `.spawn()?` reported only "No such file or directory (os error 2)" with @@ -1561,7 +1832,7 @@ fn run_pin_sidecar( // flat row index. Exact, unique, finite coverage is part of the classifier // contract: silently assigning a worst score to missing rows changes the // trained population and can invalidate sensitivity/FDR comparisons. - let t = TableFile::open(&outp)?; + let t = TableFile::open(outp)?; let orow = t.u32("candidate_id")?; let osc = t.f64("score")?; align_sidecar_scores(&orow, &osc, cid.len(), script_name) @@ -1735,6 +2006,260 @@ b assert_eq!(dst, vec![7, 8]); } + /// A competed-shaped parquet: the named feature columns as nullable f64, values keyed + /// on the row so the concatenation order is checkable, and (when `nulls`) one null per + /// five rows, which the reader must present as NaN. + fn crafted_competed(path: &str, rows: usize, tag: f64, feat_names: &[String], nulls: bool) { + let cols: Vec= feat_names + .iter() + .enumerate() + .map(|(j, name)| { + let v: Vec> = (0..rows) + .map(|i| { + if nulls && j == 1 && i % 5 == 0 { + None + } else { + Some(tag + i as f64 + j as f64 * 0.125) + } + }) + .collect(); + Col::OptF64(name.clone(), v) + }) + .collect(); + write_table(path, cols).unwrap(); + } + + fn scratch(name: &str) -> String { + let dir = std::env::temp_dir().join("mumdia_rescore_stream_test"); + std::fs::create_dir_all(&dir).unwrap(); + dir.join(format!("{}_{name}", std::process::id())) + .to_str() + .unwrap() + .to_string() + } + + #[test] + fn streamed_feature_rows_are_the_matrix_rows() { + // The streamed handoff never builds a `FeatureMatrix`, so the equality that has to + // hold is between what `for_each_feature_row` yields and what the matrix would have + // held: same values, same flat order across several inputs, same null policy. + let names: Vec = ["f0", "f1", "f2"].iter().map(|s| s.to_string()).collect(); + let a = scratch("cmp_a.parquet"); + let b = scratch("cmp_b.parquet"); + crafted_competed(&a, 7, 100.0, &names, false); + crafted_competed(&b, 5, 900.0, &names, true); + let competed = vec![a, b]; + + let m = load_feature_matrix(&competed, &names, 12).unwrap(); + assert_eq!(m.rows(), 12); + let mut seen: Vec<(usize, Vec)> = Vec::new(); + for_each_feature_row(&competed, &names, |i, v| { + seen.push((i, v.to_vec())); + Ok(()) + }) + .unwrap(); + assert_eq!(seen.len(), 12); + for (i, values) in &seen { + assert_eq!(*i, seen[*i].0, "flat row index must be dense and ascending"); + assert_eq!( + values.iter().map(|v| v.to_bits()).collect::>(), + m.row(*i).iter().map(|v| v.to_bits()).collect::>(), + "row {i}" + ); + } + // Row 7 is the first row of the second input, and its f1 is the null: NaN in both + // paths, and the validation that runs inside the stream finds the same row the + // matrix scan does. + assert!(seen[7].1[1].is_nan()); + assert_eq!(m.find_non_finite(), Some((7, 1))); + } + + #[test] + fn the_streamed_handoff_is_byte_identical_to_the_handoff_from_the_matrix() { + // The claim of the streaming change: the sidecar receives exactly the file it + // received when the engine materialised the whole feature matrix first. Both + // encodings, several inputs, and a batch small enough to force several parquet + // blocks with a block boundary inside the second input. + let names: Vec = ["f0", "f1", "f2"].iter().map(|s| s.to_string()).collect(); + let a = scratch("hs_a.parquet"); + let b = scratch("hs_b.parquet"); + crafted_competed(&a, 7, 100.0, &names, false); + crafted_competed(&b, 5, 900.0, &names, true); + let competed = vec![a, b]; + let n = 12; + let label: Vec = (0..n) + .map(|i| if i % 3 == 0 { "decoy" } else { "target" }.to_string()) + .collect(); + let pform: Vec = (0..n).map(|i| format!("PEPTIDEK/{i}")).collect(); + let protein: Vec = (0..n).map(|i| format!("sp|P{i:05}|X")).collect(); + let mz: Vec = (0..n).map(|i| 400.0 + i as f64 * 1.5).collect(); + let m = load_feature_matrix(&competed, &names, n).unwrap(); + + for use_pq in [true, false] { + let from_matrix = SidecarPaths { + handoff: scratch(if use_pq { "m.parquet" } else { "m.pin" }), + out: String::new(), + foldkeys: String::new(), + use_pq, + }; + let from_stream = SidecarPaths { + handoff: scratch(if use_pq { "s.parquet" } else { "s.pin" }), + out: String::new(), + foldkeys: String::new(), + use_pq, + }; + let mut w = HandoffWriter::new(&from_matrix, &names, &label, &pform, &protein, &mz) + .unwrap() + .with_block_rows(5); + for i in 0..n { + w.push_row(i, m.row(i)).unwrap(); + } + assert_eq!(w.finish().unwrap(), n as u64); + + let mut w = HandoffWriter::new(&from_stream, &names, &label, &pform, &protein, &mz) + .unwrap() + .with_block_rows(5); + for_each_feature_row(&competed, &names, |i, v| w.push_row(i, v)).unwrap(); + assert_eq!(w.finish().unwrap(), n as u64); + + assert_eq!( + std::fs::read(&from_matrix.handoff).unwrap(), + std::fs::read(&from_stream.handoff).unwrap(), + "handoff differs (parquet = {use_pq})" + ); + } + // And the PIN really is the PIN contract: header, then one row per PSM. + let pin = scratch("m.pin"); + let text = std::fs::read_to_string(&pin).unwrap(); + assert!(text.starts_with( + "SpecId\tLabel\tScanNr\tExpMass\tCalcMass\tf0\tf1\tf2\tPeptide\tProteins\n" + )); + assert_eq!(text.lines().count(), n + 1); + assert!(text.lines().nth(1).unwrap().starts_with("psm_0\t-1\t0\t")); + } + + /// A minimal competed table: the metadata columns `rescore::run` reads plus two + /// feature columns. No `.schema.json` companion, so the feature list is reconstructed + /// from the parquet's own columns (`NON_FEATURE_COLUMNS`). + fn crafted_competed_table(path: &str, rows: usize) { + let cols = vec![ + Col::U32("candidate_id".into(), (0..rows as u32).collect()), + Col::Str( + "label".into(), + (0..rows) + .map(|i| if i % 2 == 0 { "target" } else { "decoy" }.to_string()) + .collect(), + ), + Col::U32( + "base_peptide_id".into(), + (0..rows as u32).map(|i| i / 2).collect(), + ), + Col::Str( + "peptidoform".into(), + (0..rows).map(|i| format!("PEPTIDEK{i}")).collect(), + ), + Col::Str( + "protein".into(), + (0..rows).map(|i| format!("sp|P{i:04}|X")).collect(), + ), + Col::F64("charge".into(), (0..rows).map(|_| 2.0).collect()), + Col::F64( + "prelim_score".into(), + (0..rows) + .map(|i| if i % 2 == 0 { 9.0 } else { 1.0 } + i as f64 * 0.01) + .collect(), + ), + Col::F64( + "precursor_mz".into(), + (0..rows).map(|i| 400.0 + i as f64).collect(), + ), + Col::F64( + "apex_rt".into(), + (0..rows).map(|i| 10.0 + i as f64).collect(), + ), + Col::F64( + "elution_lo".into(), + (0..rows).map(|i| 9.0 + i as f64).collect(), + ), + Col::F64( + "elution_hi".into(), + (0..rows).map(|i| 11.0 + i as f64).collect(), + ), + Col::I32("peak_rank".into(), vec![0; rows]), + Col::F64( + "feat_a".into(), + (0..rows) + .map(|i| if i % 2 == 0 { 3.0 } else { 0.5 }) + .collect(), + ), + Col::F64( + "feat_b".into(), + (0..rows).map(|i| (i % 7) as f64 * 0.25).collect(), + ), + ]; + write_table(path, cols).unwrap(); + } + + #[test] + fn strict_sidecar_streams_the_handoff_without_building_a_matrix() { + // End-to-end through `run`: under strict with a PIN sidecar the features go from + // the competed parquet straight into the handoff, and no `FeatureMatrix` is + // allocated. Driven with an interpreter that cannot be spawned, so the run fails at + // the child rather than needing a Python environment -- but only AFTER the handoff + // has been written from the streamed features, which is what is asserted. + let competed = scratch("run_competed.parquet"); + crafted_competed_table(&competed, 24); + let work = scratch("run_work"); + let out = scratch("run_scored.parquet"); + let cfg = RescoreConfig { + classifier: RescorerKind::NnTorch, + strict: true, + python: Some("mumdia-no-such-interpreter-for-this-test".to_string()), + ..Default::default() + }; + let err = run(RescoreParams { + competed: &[competed], + out: &out, + work_dir: &work, + script_dir: "scripts", + cfg: &cfg, + config_hash: "test", + }) + .unwrap_err() + .to_string(); + assert!( + err.contains("NnTorch sidecar failed") && err.contains("strict"), + "{err}" + ); + // The handoff is a parquet (the default `handoff`), named per invocation, and holds + // every row exactly once in flat order with the PIN column contract. + let tag = format!( + "{}_{}", + std::path::Path::new(&out) + .file_stem() + .unwrap() + .to_str() + .unwrap(), + std::process::id() + ); + let handoff = format!("{work}/rescore_{tag}.features.parquet"); + let t = mumdia_io::table::Table::read(&handoff).unwrap(); + assert_eq!(t.nrows, 24); + assert_eq!(t.i32("ScanNr").unwrap(), (0..24).collect::>()); + assert_eq!( + t.i32("Label").unwrap(), + (0..24) + .map(|i| if i % 2 == 0 { 1 } else { -1 }) + .collect::>() + ); + assert_eq!( + t.f32("feat_a").unwrap(), + (0..24) + .map(|i| if i % 2 == 0 { 3.0f32 } else { 0.5 }) + .collect::>() + ); + } + #[test] fn a_shared_column_array_writes_the_same_parquet_as_two_copies() { use arrow::array::{ArrayRef, Int32Array, StringArray}; From 1ac1b443af750d983957533948a01ed86f8e6e50 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:36:28 +0200 Subject: [PATCH 14/35] perf(features): flat per-PSM buffers, dense peptidoform ids, a writer thread, and a pruned confident-bounds pass Six changes to the features stage, all output-equal. The stage's problem is the NUMBER of live medium heap blocks, not bytes: mimalloc maps blocks of that size individually and a grouped search dies at ~290 GB resident against the kernel's 1,048,576 per-process mapping limit, with 92,030 mappings of 64-256 KB and 36,368 of 256 KB-1 MB measured on a live process at 180 GB. - The per-PSM extended values were one `Vec` each inside a `Vec`: 337 f64 (2,696 B) per block, 70k-175k live from the collect() to the drop at the end of every chunk. They are one flat rows x n_ext buffer now, filled by par_chunks_mut. `FragFeatures` (no heap) stays a Vec. - The cross-charge reduction was a `HashMap<&str, HashSet>` plus a `HashMap<&str, f64>`: one HashSet block per distinct peptidoform, about 1M live for the whole run. Peptidoforms get a dense id (compete.rs's `pform_id` pattern) and the reductions become a `Vec` charge bitmask and a `Vec` of sums. The bitmask is exact for charges in 0..32 and falls back to the set-per-peptidoform build otherwise; the sums accumulate in the same row order, so the f64 addition order is unchanged. - The serial tail hashed a feature NAME per value per row, ~390 lookups per row in the Extended set, to resolve a permutation fixed before the loop. `ColIx` resolves it once. Its field names ARE the feature names (macro-generated), and a test asserts the fields of each set account for every active column exactly once, so a typo cannot silently drop a column. - The value matrix's columns are owned, so the writer takes them by move instead of `column(c).to_vec()`. That copy is why the matrix existed twice at the write and why `memlog` understated the peak by exactly 2x; the report now names the in-flight peak as well. - The parquet encode moved to a writer thread mirroring the loader thread, order preserved. Both channels are rendezvous (depth 0): `sync_channel(1)` held THREE chunks, not the two its comment claimed. - `confident_global_bounds` streamed the entire chromatogram table (up to 88 GB) to learn two scalars from a few thousand candidates, on by default. It now reads only the row groups whose `candidate_id` statistics can contain a confident candidate; a candidate's rows are contiguous, so one that spans two groups is inside both ranges and neither is skipped. Chunk boundaries inside a span stay on the same ABSOLUTE row grid, so a candidate straddling one is split where it was split before and the percentiles are identical. On the test fixture (100 row groups, every fifth candidate an anchor) it reads 40% of the rows. - Hoisted allocations, all bit-identical: the per-scan `obs_k` in the pCos loop and the per-fragment `others` in the contrast loop (features.rs); the summed-positive gather in nonzero.rs, which was two Vecs per fragment PAIR for a per-fragment quantity, plus scratch reuse for the third; the k x k `corr` (coelution.rs) and m x m Gram matrix (interference.rs) are one flat buffer each instead of a Vec per row, and power_top's per-iteration `nv` is reused. The trace alignment was built twice per PSM, once in `fragment_features` and once in `build_evidence`, each fragment through its own `HashMap` keyed on the RT bit pattern. It is built once per PSM and handed to both (Evidence takes ownership), and it has a fast path: the chunk store shares ONE axis across a candidate's rows, so when every row with a trace samples that same strictly ascending slice the union axis IS that slice and every map is the identity. The two paths' bounding behaviour (`fragment_features` honours `bound_features`, `build_evidence` does not) is untouched. Two hazards the threading uncovered, both fixed rather than inherited: - the loader's receiver lived in the enclosing frame, so an error return from the scope closure left the loader blocked in `send` on a channel nobody would receive from again while the scope waited to join it: a hang instead of an error. Depth 0 makes that more reachable, so both channels are created inside the scope and their receivers drop with it. - with the write on its own thread, an error would drop the sender, end the writer's loop normally and PUBLISH a truncated features table over the previous good one. A `None` commit marker is sent only after the last chunk; without it the writer drops unpublished and `AtomicPath` removes its temp file. Output equality: eight configurations (Extended/Rich/Minimal x bound_from_confident on/off x chunk sizes 64, 4096 and 2^20, with and without a seed table) over a 97-candidate fixture produce byte-identical feature values and PIN bytes before and after, compared as raw f64 bit patterns per column. New tests assert the new code against the old: the shared-axis alignment against the union-and-map build (and that it refuses an unsorted or repeated axis), the charge bitmask against the set build, the pruned confident-bounds pass against the whole-table scan at four chunk sizes, ColIx against the name lookup, and the whole chunked Extended pass for chunk invariance, which the existing `features_chunking_is_value_preserving` never reached because it runs the default Minimal set. Residual risk: the confident-bounds pass already depended on the chunk size before this change (a candidate straddling a chunk boundary is bounded from each part separately), which the alignment preserves but does not fix. cargo fmt, cargo clippy -p mumdia --all-targets -D warnings and cargo test --workspace all pass. Co-Authored-By: Claude Fable 5.1 --- .../crates/mumdia/src/stages/features.rs | 1495 ++++++++++++++--- .../mumdia/src/stages/features/coelution.rs | 24 +- .../src/stages/features/interference.rs | 37 +- .../mumdia/src/stages/features/nonzero.rs | 30 +- 4 files changed, 1298 insertions(+), 288 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/features.rs b/rust/mumdia/crates/mumdia/src/stages/features.rs index a8548033..a8bcb900 100644 --- a/rust/mumdia/crates/mumdia/src/stages/features.rs +++ b/rust/mumdia/crates/mumdia/src/stages/features.rs @@ -357,6 +357,102 @@ struct ChromRow<'a> { inten: &'a [f32], } +/// One PSM's fragment traces on a single union RT axis: `axis_full` ascending and +/// deduplicated, `traces_full[i]` the intensities of row `i` sampled on it (zero where +/// that row has no point). +/// +/// Both feature paths need exactly this, and both used to rebuild it independently, each +/// row through its own `HashMap` keyed on the RT bit pattern -- two maps per +/// fragment per PSM. It is built once per PSM now and handed to both. +struct TraceAlign { + axis_full: Vec, + traces_full: Vec>, +} + +/// Build the union alignment of a PSM's rows, by the shared-axis fast path when it +/// applies and by the union-and-map build otherwise. +fn align_traces(rows: &[ChromRow]) -> TraceAlign { + align_shared_axis(rows).unwrap_or_else(|| align_union(rows)) +} + +/// Fast path: every row with a trace samples the SAME stored axis, and that axis is +/// strictly ascending. This is the normal case -- [`ChromChunk::axis_for`] stores one +/// axis per candidate because extract samples every fragment of a candidate on the same +/// window grid -- and then the union axis IS that slice and every per-fragment map is +/// the identity, so the sort, the dedup and the maps are all pure overhead. +/// +/// Returns `None` unless the precondition holds, and the two builds then agree bit for +/// bit: sorting and deduplicating a strictly ascending axis returns it unchanged, and +/// looking a row's own RT value up in its own map returns that row's intensity at the +/// same position. A row with no trace at all gets zeros either way (its map is empty). +fn align_shared_axis(rows: &[ChromRow]) -> Option { + let mut shared: Option<&[f32]> = None; + for r in rows { + if r.rt.is_empty() { + continue; + } + // A row whose value count does not match its axis is aligned by the map build, + // which truncates or zero-fills; that is rare enough not to be worth mirroring. + if r.rt.len() != r.inten.len() { + return None; + } + match shared { + None => shared = Some(r.rt), + Some(a) => { + if a.len() != r.rt.len() || !std::ptr::eq(a.as_ptr(), r.rt.as_ptr()) { + return None; + } + } + } + } + let axis = shared?; + // Strictly ascending rules out the two cases the sort-and-dedup build would change: + // an out-of-order axis, and repeated RT values (where the map keeps the last). + if !axis.windows(2).all(|w| w[0] < w[1]) { + return None; + } + let traces_full = rows + .iter() + .map(|r| { + if r.rt.is_empty() { + vec![0.0; axis.len()] + } else { + r.inten.iter().map(|&v| v as f64).collect() + } + }) + .collect(); + Some(TraceAlign { + axis_full: axis.to_vec(), + traces_full, + }) +} + +/// General path: the sorted, deduplicated union of every row's RT values, with each row +/// resampled onto it through a bit-pattern-keyed map. +fn align_union(rows: &[ChromRow]) -> TraceAlign { + let mut axis_full: Vec = rows.iter().flat_map(|r| r.rt.iter().cloned()).collect(); + axis_full.sort_by(|a, b| a.total_cmp(b)); + axis_full.dedup(); + let traces_full: Vec> = rows + .iter() + .map(|r| { + let map: HashMap = + r.rt.iter() + .zip(r.inten.iter()) + .map(|(&t, &v)| (t.to_bits(), v)) + .collect(); + axis_full + .iter() + .map(|t| *map.get(&t.to_bits()).unwrap_or(&0.0) as f64) + .collect() + }) + .collect(); + TraceAlign { + axis_full, + traces_full, + } +} + /// Per-PSM evidence handed to the extended feature families. All arrays are /// f64. Fragment-indexed arrays share one order; time-series share `axis` /// (elution-peak-bounded) or `axis_full` (whole extracted window). Built once @@ -442,6 +538,7 @@ fn parse_ion(name: &str) -> (bool, u32, u32) { /// same elution peak the legacy features use. fn build_evidence( rows: &[ChromRow], + al: TraceAlign, ms1_rows: &[ChromRow], apex_rt: f64, frac: f64, @@ -478,23 +575,12 @@ fn build_evidence( mass_err_ppm.push(ppm_diff(r.frag_obs_mz, r.frag_mz)); } - let mut axis_full: Vec = rows.iter().flat_map(|r| r.rt.iter().cloned()).collect(); - axis_full.sort_by(|a, b| a.total_cmp(b)); - axis_full.dedup(); - let traces_full: Vec> = rows - .iter() - .map(|r| { - let map: HashMap = - r.rt.iter() - .zip(r.inten.iter()) - .map(|(&t, &v)| (t.to_bits(), v)) - .collect(); - axis_full - .iter() - .map(|t| *map.get(&t.to_bits()).unwrap_or(&0.0) as f64) - .collect() - }) - .collect(); + // The alignment is built once per PSM by the caller and read by `fragment_features` + // first; [`Evidence`] then owns it, so it is moved in rather than rebuilt here. + let TraceAlign { + axis_full, + traces_full, + } = al; let (lo_i, hi_i) = if axis_full.len() >= 3 { let ai = axis_full @@ -1053,46 +1139,148 @@ fn plan_chunks( Ok(chunks) } -/// Feature values of one chunk, column-major in one flat buffer: `vals[c * rows + r]`. -/// Replaces the `HashMap<&str, Vec>` that held one `Vec` per feature for the whole -/// run (7.5 GiB at 2.6M rows x 387 features). A column is contiguous, so handing it to -/// the writer is one copy and no transpose. +/// Feature values of one chunk, one owned column per feature. Replaces the +/// `HashMap<&str, Vec>` that held one `Vec` per feature for the WHOLE RUN (7.5 GiB +/// at 2.6M rows x 387 features); a chunk's columns are a few hundred blocks, not a few +/// hundred thousand. +/// +/// The columns are owned rather than slices of one flat buffer so that +/// [`ValueMatrix::take_column`] can hand each one to the parquet writer by move. The +/// flat layout it replaced had to copy every column (`column(c).to_vec()`), which meant +/// the whole matrix existed twice at the moment of the write -- the memory report said +/// `largest_chunk` and meant half of the real peak. struct ValueMatrix { - idx: HashMap, rows: usize, - vals: Vec, + cols: Vec>, } impl ValueMatrix { - fn new(cols: &[String], rows: usize) -> ValueMatrix { + fn new(n_cols: usize, rows: usize) -> ValueMatrix { ValueMatrix { - idx: cols - .iter() - .enumerate() - .map(|(i, c)| (c.clone(), i)) - .collect(), rows, - vals: vec![0.0; cols.len() * rows], + cols: (0..n_cols).map(|_| vec![0.0; rows]).collect(), } } - /// Set feature `name` for chunk-local row `r`. A name outside the active set is - /// dropped, exactly as the old map was filtered by `cols_active` at write time. - fn set(&mut self, name: &str, r: usize, v: f64) { - if let Some(&c) = self.idx.get(name) { - self.vals[c * self.rows + r] = v; + /// Set the column `c` (as resolved once by [`ColIx`]) for chunk-local row `r`. A + /// feature outside the active set resolves to `None` and is dropped, exactly as the + /// old map was filtered by `cols_active` at write time. + fn set(&mut self, c: Option, r: usize, v: f64) { + if let Some(c) = c { + self.cols[c][r] = v; } } fn column(&self, c: usize) -> &[f64] { - &self.vals[c * self.rows..(c + 1) * self.rows] + &self.cols[c] + } + + /// Move column `c` out, leaving an empty Vec behind. The matrix is dropped straight + /// after, so the emptied columns are never read again. + fn take_column(&mut self, c: usize) -> Vec { + std::mem::take(&mut self.cols[c]) } fn payload_bytes(&self) -> usize { - crate::memlog::bytes_of(&self.vals) + self.cols.len() * self.rows * std::mem::size_of::() } } +/// Column index of every feature the serial assembly writes, resolved ONCE from the +/// active feature list. +/// +/// The assembly used to name each feature as a string and hash it per row: ~390 lookups +/// for every PSM in the Extended set, all resolving to the same fixed permutation. The +/// permutation is known before the loop, so it is taken once here; `None` is a feature +/// the configured set does not carry, which the assembly then drops exactly as the +/// name-keyed setter did. +/// +/// The field name IS the feature name, so the list below cannot drift from the names it +/// resolves, and `colix_covers_every_active_column` checks that the fields of one set +/// account for every column of that set exactly once. +macro_rules! col_ix { + ($($field:ident),* $(,)?) => { + struct ColIx { + $($field: Option,)* + /// One entry per extended-battery feature, in registry order. + ext: Vec>, + } + + impl ColIx { + fn new(cols_active: &[String], ext_names: &[&'static str]) -> ColIx { + let idx: std::collections::HashMap<&str, usize> = cols_active + .iter() + .enumerate() + .map(|(i, c)| (c.as_str(), i)) + .collect(); + ColIx { + $($field: idx.get(stringify!($field)).copied(),)* + ext: ext_names.iter().map(|n| idx.get(n).copied()).collect(), + } + } + + /// Every non-extended column this resolved, for the coverage test. + #[cfg(test)] + fn named(&self) -> Vec> { + vec![$(self.$field,)*] + } + } + }; +} + +col_ix!( + rt_error_abs, + rt_error_rel, + n_matched_fragments, + coelution_run, + log_apex_intensity, + frag_corr, + frag_cosine, + spectral_angle, + coelution_mean, + coelution_best, + n_coelution_above, + charge, + peptide_length, + n_proteins, + library_norm_manhattan, + library_rmsd, + xcorr_coelution, + xcorr_shape, + sum_b_intensity, + sum_y_intensity, + diff_by_intensity, + n_b_ions, + n_y_ions, + weighted_mass_error, + mean_mass_error, + isotope_corr, + ms1_isom1_ratio, + log_mono_ms1, + has_ms1, + log_sn, + n_observations, + base_width_rt, + seed_score, + seed_identified, + matched_fraction, + profile_cos, + ref_corr, + best_ref_corr, + low_frag_coel, + evidence, + contrast_min, + resid_corr, + coel_clean, + shadow_frac, + peak_contested_frac, + peak_contested_count_frac, + peak_apportioned_frac, + n_charge_states, + charge_multi_flag, + cross_charge_intensity_log, +); + /// Percolator-style PIN written row by row as the chunks are computed. Nothing in the /// pipeline reads it (rescore builds its own), so it is gated by `features.emit_pin`; /// the byte output is unchanged from the single-shot writer it replaced. @@ -1148,12 +1336,101 @@ impl PinWriter { } } +/// Number of distinct charge states of each row's peptidoform, indexed by the dense +/// peptidoform id of that row. +/// +/// Prefers the bitmask build and falls back to the set-per-peptidoform build, which is +/// what this used to do for every run: one `HashSet` heap block per distinct +/// peptidoform, about a million live for the whole run on a real library. +fn charge_states_per_row(pf_of_row: &[u32], n_pf: usize, charge: &[i32]) -> Vec { + charge_states_by_mask(pf_of_row, n_pf, charge) + .unwrap_or_else(|| charge_states_by_set(pf_of_row, n_pf, charge)) +} + +/// Bitmask build: exact for charges in `0..32`, which covers every real precursor charge. +/// `None` when any charge falls outside that range, so the caller takes the set build and +/// the count is provably the same either way. +fn charge_states_by_mask(pf_of_row: &[u32], n_pf: usize, charge: &[i32]) -> Option> { + if !charge.iter().all(|&z| (0..32).contains(&z)) { + return None; + } + let mut mask: Vec = vec![0; n_pf]; + for (&pf, &z) in pf_of_row.iter().zip(charge) { + mask[pf as usize] |= 1u32 << z; + } + Some( + pf_of_row + .iter() + .map(|&pf| mask[pf as usize].count_ones() as f64) + .collect(), + ) +} + +/// Reference build: one set of charges per peptidoform. +fn charge_states_by_set(pf_of_row: &[u32], n_pf: usize, charge: &[i32]) -> Vec { + let mut sets: Vec> = vec![Default::default(); n_pf]; + for (&pf, &z) in pf_of_row.iter().zip(charge) { + sets[pf as usize].insert(z); + } + pf_of_row + .iter() + .map(|&pf| sets[pf as usize].len() as f64) + .collect() +} + +/// Contiguous row spans of the chromatogram table that CAN hold a confident candidate, +/// from the parquet row-group statistics on `candidate_id` alone. `confident` must be +/// sorted ascending. +/// +/// A group whose `[min, max]` contains no confident id contains no confident ROW either, +/// because min/max bound every value in the group, so skipping it changes nothing the +/// caller would have kept. The converse is only conservative: a kept group may hold no +/// confident row after all, and the row filter drops those as it always did. +/// +/// A candidate's rows are contiguous (`plan_chunks` verifies exactly that), so a confident +/// candidate that spans two row groups puts its id inside BOTH groups' ranges and both are +/// kept; a confident candidate is therefore never split across a span boundary. +/// +/// Returns `None` when the statistics cannot be used -- a writer that recorded none, or a +/// file with no row groups -- which means "read the whole table", the previous behaviour. +fn confident_row_spans( + stats: &[mumdia_io::table::RowGroupStats], + confident: &[u32], +) -> Option> { + if stats.is_empty() || stats.iter().any(|s| s.min.is_none() || s.max.is_none()) { + return None; + } + let mut spans: Vec<(usize, usize)> = Vec::new(); + let mut start = 0usize; + for s in stats { + let (lo, hi) = (s.min?, s.max?); + // First confident id at or above the group's minimum; the group is wanted when + // that id is also at or below its maximum. u32 is exact in f64. + let i = confident.partition_point(|&c| (c as f64) < lo); + if confident.get(i).is_some_and(|&c| (c as f64) <= hi) { + match spans.last_mut() { + Some(last) if last.0 + last.1 == start => last.1 += s.rows, + _ => spans.push((start, s.rows)), + } + } + start += s.rows; + } + Some(spans) +} + /// Global elution half-widths from the confident set, computed in one streaming pass that /// holds a single candidate's rows at a time. Returns None when fewer than 20 confident /// anchors have a resolvable peak (the caller then keeps per-candidate detection). +/// +/// `spans` are the row spans to read, from [`confident_row_spans`]. The whole table used +/// to be decoded here -- up to 88 GB of traces to learn two scalars from a few thousand +/// candidates -- with every row of a non-confident candidate decoded and thrown away. +/// Chunk boundaries inside a span stay on the same ABSOLUTE row grid the whole-table pass +/// used, so a candidate that straddles one is split exactly where it was split before and +/// the half-width samples, and therefore the percentiles, are unchanged. fn confident_global_bounds( - ch: &TableFile, chrom_path: &str, + spans: &[(usize, usize)], confident_rows: &HashMap>, apex_rt: &[f64], cfg: &FeaturesConfig, @@ -1162,34 +1439,38 @@ fn confident_global_bounds( let mut lefts: Vec = Vec::new(); let mut rights: Vec = Vec::new(); let mut names = NameTab::default(); - let mut stream = ChromStream::open(ch, chrom_path)?; let keep: std::collections::HashSet = confident_rows.keys().copied().collect(); // Chunk reading already groups rows by candidate, so this reuses it and keeps only // the confident candidates' rows long enough to bound their peak. - let mut read = 0usize; - while read < ch.nrows { - let chunk = stream.read_chunk_filtered(chunk_rows, &mut names, Some(&keep))?; - read += chunk_rows; - for (ci, &c) in chunk.cids.iter().enumerate() { - let Some(psm_rows) = confident_rows.get(&c) else { - continue; - }; - let rows = chunk.rows(&chunk.frag, ci, &names); - if rows.is_empty() { - continue; - } - for &i in psm_rows { - if let Some((lo, hi)) = elution_peak_rt_bounds( - &rows, - apex_rt[i], - cfg.bound_peak_fraction, - cfg.bound_peak_grace, - ) { - let l = apex_rt[i] - lo as f64; - let r = hi as f64 - apex_rt[i]; - if l >= 0.0 && r >= 0.0 { - lefts.push(l); - rights.push(r); + for &(first, n_rows) in spans { + let span = TableFile::open_rows(chrom_path, first, n_rows)?; + let mut stream = ChromStream::open(&span, chrom_path)?; + let (mut abs, end) = (first, first + n_rows); + while abs < end { + let want = (chunk_rows - abs % chunk_rows).min(end - abs); + let chunk = stream.read_chunk_filtered(want, &mut names, Some(&keep))?; + abs += want; + for (ci, &c) in chunk.cids.iter().enumerate() { + let Some(psm_rows) = confident_rows.get(&c) else { + continue; + }; + let rows = chunk.rows(&chunk.frag, ci, &names); + if rows.is_empty() { + continue; + } + for &i in psm_rows { + if let Some((lo, hi)) = elution_peak_rt_bounds( + &rows, + apex_rt[i], + cfg.bound_peak_fraction, + cfg.bound_peak_grace, + ) { + let l = apex_rt[i] - lo as f64; + let r = hi as f64 - apex_rt[i]; + if l >= 0.0 && r >= 0.0 { + lefts.push(l); + rights.push(r); + } } } } @@ -1352,9 +1633,26 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result confident_rows.entry(c).or_default().push(i); } } + // Read only the row groups whose candidate_id range can hold a confident + // candidate; the rest contain none by construction. + let mut sorted: Vec = confident_rows.keys().copied().collect(); + sorted.sort_unstable(); + let spans = ch + .row_group_stats("candidate_id") + .ok() + .and_then(|s| confident_row_spans(&s, &sorted)) + .unwrap_or_else(|| vec![(0, ch.nrows)]); + let span_rows: usize = spans.iter().map(|s| s.1).sum(); + info!( + spans = spans.len(), + rows = span_rows, + of_rows = ch.nrows, + confident = sorted.len(), + "features: confident-bounds pass reads this much of the chromatogram table" + ); confident_global_bounds( - &ch, p.chromatograms, + &spans, &confident_rows, &apex_rt, p.cfg, @@ -1375,18 +1673,26 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result // more than a shift decoy, and this evidence axis is invisible to the per-PSM // Evidence families since each charge is a separate candidate. It is a whole-run // reduction over PSM columns only, so it is computed before the chunk loop. - let mut pf_charges: HashMap<&str, std::collections::HashSet> = HashMap::new(); - let mut pf_int: HashMap<&str, f64> = HashMap::new(); + // + // The grouping used to be a `HashMap<&str, HashSet>` plus a `HashMap<&str, f64>`: + // ONE HashSet heap block per distinct peptidoform, about a million of them, live for + // the whole run. Peptidoforms get a dense id here instead -- the same thing compete.rs + // does with `pform_id` -- and the two reductions become a `Vec` charge bitmask and + // a `Vec` of sums, three allocations in total. The sums are accumulated in the same + // row order as before, so the f64 addition order, and therefore the value, is unchanged. + let mut pf_ids: HashMap<&str, u32> = HashMap::with_capacity(n); + let mut pf_of_row: Vec = Vec::with_capacity(n); + for p in pform.iter().take(n) { + let next = pf_ids.len() as u32; + pf_of_row.push(*pf_ids.entry(p.as_str()).or_insert(next)); + } + let n_pf = pf_ids.len(); + drop(pf_ids); + let mut pf_int: Vec = vec![0.0; n_pf]; for i in 0..n { - pf_charges - .entry(pform[i].as_str()) - .or_default() - .insert(charge[i]); - *pf_int.entry(pform[i].as_str()).or_insert(0.0) += apex_int[i] as f64; - } - let f_n_charge: Vec = (0..n) - .map(|i| pf_charges[pform[i].as_str()].len() as f64) - .collect(); + pf_int[pf_of_row[i] as usize] += apex_int[i] as f64; + } + let f_n_charge = charge_states_per_row(&pf_of_row, n_pf, &charge); let f_charge_multi: Vec = f_n_charge .iter() .map(|&c| if c >= 2.0 { 1.0 } else { 0.0 }) @@ -1394,9 +1700,9 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result // ln(1 + summed apex intensity of the OTHER charge states of this peptidoform): // how much independent charge-state evidence reinforces this PSM (unbounded). let f_cross_charge_int: Vec = (0..n) - .map(|i| (1.0 + (pf_int[pform[i].as_str()] - apex_int[i] as f64).max(0.0)).ln()) + .map(|i| (1.0 + (pf_int[pf_of_row[i] as usize] - apex_int[i] as f64).max(0.0)).ln()) .collect(); - drop((pf_charges, pf_int)); + drop((pf_of_row, pf_int)); let extended = matches!(p.cfg.set, FeatureSet::Extended); let ext_names = extended_name_refs(); @@ -1404,15 +1710,23 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result // The two expensive per-PSM computations (`fragment_features` and, when the // extended set is active, `build_evidence` + `extended_values`) are pure // functions of that PSM's own inputs, so they are computed in parallel over the - // chunk and indexed by row. The serial assembly below reads `per[r]` and is - // otherwise unchanged, so the feature values are identical to the whole-run + // chunk and indexed by row. The serial assembly below reads them back by row and + // is otherwise unchanged, so the feature values are identical to the whole-run // version this replaced. - struct PerPsm { - ff: FragFeatures, - ext: Vec, - } - - let mut writer = TableWriter::new(p.out).with_row_group_rows(FEATURE_ROW_GROUP_ROWS); + // + // The extended values used to come back as one `Vec` PER PSM inside a + // `Vec`: 337 f64 (2,696 B) in its own heap block, 70k-175k of them live + // from the `collect()` to the `drop` at the end of the chunk. That is the single + // largest producer of medium-sized live mappings in the stage, and mimalloc maps + // blocks of that size individually, so it counts against the kernel's per-process + // mapping limit (1,048,576) rather than against bytes. They are one flat + // `rows x n_ext` buffer now, filled by `par_chunks_mut` and addressed by offset; + // the values, their order and their row assignment are unchanged. + let n_ext = if extended { ext_names.len() } else { 0 }; + // The feature -> column permutation, taken once instead of once per value per row. + let ix = ColIx::new(&cols_active, &ext_names); + + let writer = TableWriter::new(p.out).with_row_group_rows(FEATURE_ROW_GROUP_ROWS); let mut pin = if p.cfg.emit_pin { Some(PinWriter::create(p.out_pin, &cols_active)?) } else { @@ -1436,10 +1750,55 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result // snapshot of the fragment-name table as it stood when the chunk was read, which is // exactly what the serial code saw at that point. let chunk_rows: Vec = chunks.iter().map(|c| c.chrom_rows).collect(); - let (tx, rx) = std::sync::mpsc::sync_channel::>(1); let ch_ref = &ch; let chrom_path = p.chromatograms; - std::thread::scope(|sc| -> Result<()> { + let rows = std::thread::scope(|sc| -> Result { + // A RENDEZVOUS channel, not a one-deep buffer. `sync_channel(1)` let the loader + // finish a chunk, park it in the buffer and start a third, so three chunks of + // traces were resident where the comment claimed two. At depth 0 the loader's + // `send` blocks until this loop takes the chunk, which still overlaps decode with + // compute fully (the loader is free again the instant the chunk is handed over) + // and holds exactly two: the one being computed and the one being decoded. + // + // Both channels are created HERE rather than outside the scope so that both + // receivers drop when this closure returns. An error return used to leave `rx` + // alive in the enclosing frame while the scope waited to join the loader, and the + // loader was blocked in `send` on a channel that would never be received from + // again: the stage hung instead of reporting the error. Depth 0 makes that more + // reachable, not less, so it is fixed rather than papered over. + let (tx, rx) = std::sync::mpsc::sync_channel::>(0); + // The parquet encode is single-threaded; on its own thread it overlaps with the + // next chunk's computation instead of stalling it. Rendezvous again, so the + // encoder never queues a backlog of value matrices. + // + // `None` is the commit marker: the writer publishes the artifact only after the + // chunk loop has sent every chunk. Without it, an error anywhere in this closure + // would drop `wtx`, end the writer's loop normally and publish a TRUNCATED table + // over the previous good one, where the old in-line write simply returned before + // `close` and let `AtomicPath` remove its temp file. + let (wtx, wrx) = std::sync::mpsc::sync_channel::>>(0); + let wh = sc.spawn(move || -> Result { + let mut writer = writer; + let mut commit = false; + for msg in wrx { + match msg { + Some(cols) => writer.write_cols(cols)?, + None => { + commit = true; + break; + } + } + } + if !commit { + // Dropping the writer unpublished takes its temp file with it. + drop(writer); + return Err(anyhow!( + "features: the chunk loop stopped before the last chunk; \ + the features table was not written" + )); + } + writer.close() + }); sc.spawn(move || { let mut stream = match ChromStream::open(ch_ref, chrom_path) { Ok(s) => s, @@ -1478,77 +1837,95 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result n_frag_rows += store.frag.nrows(); n_ms1_rows += store.ms1.nrows(); - let per: Vec = (lo..hi) - .into_par_iter() - .map(|i| { - let ci = store.index.get(&cid[i]).copied(); - let rows = ci - .map(|c| store.rows(&store.frag, c, &names)) + // One PSM's work, writing its extended values into `ext` (already zeroed and + // exactly `n_ext` wide) rather than returning a fresh Vec. + let per_psm = |i: usize, ext: &mut [f64]| -> FragFeatures { + let ci = store.index.get(&cid[i]).copied(); + let rows = ci + .map(|c| store.rows(&store.frag, c, &names)) + .unwrap_or_default(); + if rows.is_empty() { + // No chromatogram rows: default features and all-zero extended + // values, as before. + return FragFeatures::default(); + } + // One alignment per PSM, read by `fragment_features` and then moved + // into the Evidence: it used to be rebuilt inside each of them. + let al = align_traces(&rows); + let ff = fragment_features( + &rows, + &al, + apex_rt[i], + p.cfg.coelution_corr_threshold, + p.cfg.bound_features, + p.cfg.bound_peak_fraction, + p.cfg.bound_peak_grace, + global_bounds, + ); + if !ext.is_empty() { + let ms1_rows = ci + .map(|c| store.rows(&store.ms1, c, &names)) .unwrap_or_default(); - let ff = if rows.is_empty() { - FragFeatures::default() - } else { - fragment_features( - &rows, - apex_rt[i], - p.cfg.coelution_corr_threshold, - p.cfg.bound_features, - p.cfg.bound_peak_fraction, - p.cfg.bound_peak_grace, - global_bounds, - ) - }; - let ext = if extended { - if rows.is_empty() { - vec![0.0; ext_names.len()] - } else { - let ms1_rows = ci - .map(|c| store.rows(&store.ms1, c, &names)) - .unwrap_or_default(); - let mut ev = build_evidence( - &rows, - &ms1_rows, - apex_rt[i], - p.cfg.bound_peak_fraction, - p.cfg.bound_peak_grace, - global_bounds, - ); - ev.rt_pred_cal = rt_cal[i]; - ev.rt_err = calibrated_rt_error(apex_rt[i], rt_cal[i]); - ev.gradient = gradient; - ev.precursor_mz = mz[i]; - ev.charge = charge[i]; - ev.seq_len = peptide_length(&pform[i]); - ev.n_matched = n_matched[i]; - ev.n_predicted = n_pred[i]; - ev.seed_score = *seed_score_map.get(&cid[i]).unwrap_or(&0.0); - ev.seed_identified = *seed_id_map.get(&cid[i]).unwrap_or(&0.0); - ev.apex_intensity = apex_int[i] as f64; - ev.ms1_mono = ms1_mono[i]; - ev.ms1_iso1 = ms1_i1[i]; - ev.ms1_iso2 = ms1_i2[i]; - ev.ms1_isom1 = ms1_m1[i]; - ev.ms1_precursor_features = p.cfg.ms1_precursor_features; - ev.deconv_explained = deconv_expl[i] as f64; - ev.deconv_active = deconv_act[i] as f64; - ev.deconv_share = deconv_shr[i] as f64; - ev.deconv_max_collin = deconv_col[i] as f64; - ev.deconv_shadow = deconv_sha[i] as f64; - extended_values(&ev) - } - } else { - Vec::new() - }; - PerPsm { ff, ext } - }) - .collect(); + let mut ev = build_evidence( + &rows, + al, + &ms1_rows, + apex_rt[i], + p.cfg.bound_peak_fraction, + p.cfg.bound_peak_grace, + global_bounds, + ); + ev.rt_pred_cal = rt_cal[i]; + ev.rt_err = calibrated_rt_error(apex_rt[i], rt_cal[i]); + ev.gradient = gradient; + ev.precursor_mz = mz[i]; + ev.charge = charge[i]; + ev.seq_len = peptide_length(&pform[i]); + ev.n_matched = n_matched[i]; + ev.n_predicted = n_pred[i]; + ev.seed_score = *seed_score_map.get(&cid[i]).unwrap_or(&0.0); + ev.seed_identified = *seed_id_map.get(&cid[i]).unwrap_or(&0.0); + ev.apex_intensity = apex_int[i] as f64; + ev.ms1_mono = ms1_mono[i]; + ev.ms1_iso1 = ms1_i1[i]; + ev.ms1_iso2 = ms1_i2[i]; + ev.ms1_isom1 = ms1_m1[i]; + ev.ms1_precursor_features = p.cfg.ms1_precursor_features; + ev.deconv_explained = deconv_expl[i] as f64; + ev.deconv_active = deconv_act[i] as f64; + ev.deconv_share = deconv_shr[i] as f64; + ev.deconv_max_collin = deconv_col[i] as f64; + ev.deconv_shadow = deconv_sha[i] as f64; + // Truncating zip, exactly as the name/value zip in the assembly + // did: a family that returned fewer values leaves the tail at 0.0. + for (dst, v) in ext.iter_mut().zip(extended_values(&ev)) { + *dst = v; + } + } + ff + }; - let mut m = ValueMatrix::new(&cols_active, rows_in_chunk); + let mut frag_feats: Vec = vec![FragFeatures::default(); rows_in_chunk]; + let mut ext_vals: Vec = vec![0.0; rows_in_chunk * n_ext]; + if n_ext > 0 { + frag_feats + .par_iter_mut() + .zip(ext_vals.par_chunks_mut(n_ext)) + .enumerate() + .for_each(|(r, (ff, ext))| *ff = per_psm(lo + r, ext)); + } else { + frag_feats + .par_iter_mut() + .enumerate() + .for_each(|(r, ff)| *ff = per_psm(lo + r, &mut [])); + } + + let mut m = ValueMatrix::new(cols_active.len(), rows_in_chunk); let mut prelim = vec![0.0f64; rows_in_chunk]; let mut elu_lo = vec![0.0f64; rows_in_chunk]; let mut elu_hi = vec![0.0f64; rows_in_chunk]; for (r, i) in (lo..hi).enumerate() { - let ff = &per[r].ff; + let ff = &frag_feats[r]; elu_lo[r] = ff.elution_lo; elu_hi[r] = ff.elution_hi; let rt_err = calibrated_rt_error(apex_rt[i], rt_cal[i]); @@ -1557,89 +1934,91 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result let (iso_corr, isom1_ratio, log_mono, has_ms1) = isotope_features(ms1_m1[i], ms1_mono[i], ms1_i1[i], ms1_i2[i], neutral); - m.set("rt_error_abs", r, rt_err); - m.set("rt_error_rel", r, rt_err / gradient); - m.set("n_matched_fragments", r, n_matched[i] as f64); - m.set("coelution_run", r, corun[i] as f64); - m.set("log_apex_intensity", r, (1.0 + apex_int[i] as f64).ln()); - m.set("frag_corr", r, ff.frag_corr); - m.set("frag_cosine", r, ff.frag_cosine); - m.set("spectral_angle", r, ff.spectral_angle); - m.set("coelution_mean", r, ff.coelution_mean); - m.set("coelution_best", r, ff.coelution_best); - m.set("n_coelution_above", r, ff.n_coelution_above); - m.set("charge", r, charge[i] as f64); - m.set("peptide_length", r, peptide_length(&pform[i]) as f64); + m.set(ix.rt_error_abs, r, rt_err); + m.set(ix.rt_error_rel, r, rt_err / gradient); + m.set(ix.n_matched_fragments, r, n_matched[i] as f64); + m.set(ix.coelution_run, r, corun[i] as f64); + m.set(ix.log_apex_intensity, r, (1.0 + apex_int[i] as f64).ln()); + m.set(ix.frag_corr, r, ff.frag_corr); + m.set(ix.frag_cosine, r, ff.frag_cosine); + m.set(ix.spectral_angle, r, ff.spectral_angle); + m.set(ix.coelution_mean, r, ff.coelution_mean); + m.set(ix.coelution_best, r, ff.coelution_best); + m.set(ix.n_coelution_above, r, ff.n_coelution_above); + m.set(ix.charge, r, charge[i] as f64); + m.set(ix.peptide_length, r, peptide_length(&pform[i]) as f64); m.set( - "n_proteins", + ix.n_proteins, r, (protein[i].matches(';').count() + 1) as f64, ); - m.set("library_norm_manhattan", r, ff.norm_manhattan); - m.set("library_rmsd", r, ff.rmsd); - m.set("xcorr_coelution", r, ff.xcorr_coelution); - m.set("xcorr_shape", r, ff.xcorr_shape); - m.set("sum_b_intensity", r, ff.sum_b); - m.set("sum_y_intensity", r, ff.sum_y); - m.set("diff_by_intensity", r, ff.sum_b - ff.sum_y); - m.set("n_b_ions", r, ff.n_b); - m.set("n_y_ions", r, ff.n_y); - m.set("weighted_mass_error", r, ff.weighted_mass_error); - m.set("mean_mass_error", r, ff.mean_mass_error); - m.set("isotope_corr", r, iso_corr); - m.set("ms1_isom1_ratio", r, isom1_ratio); - m.set("log_mono_ms1", r, log_mono); - m.set("has_ms1", r, has_ms1); - m.set("log_sn", r, ff.log_sn); - m.set("n_observations", r, ff.n_observations); - m.set("base_width_rt", r, ff.base_width_rt); + m.set(ix.library_norm_manhattan, r, ff.norm_manhattan); + m.set(ix.library_rmsd, r, ff.rmsd); + m.set(ix.xcorr_coelution, r, ff.xcorr_coelution); + m.set(ix.xcorr_shape, r, ff.xcorr_shape); + m.set(ix.sum_b_intensity, r, ff.sum_b); + m.set(ix.sum_y_intensity, r, ff.sum_y); + m.set(ix.diff_by_intensity, r, ff.sum_b - ff.sum_y); + m.set(ix.n_b_ions, r, ff.n_b); + m.set(ix.n_y_ions, r, ff.n_y); + m.set(ix.weighted_mass_error, r, ff.weighted_mass_error); + m.set(ix.mean_mass_error, r, ff.mean_mass_error); + m.set(ix.isotope_corr, r, iso_corr); + m.set(ix.ms1_isom1_ratio, r, isom1_ratio); + m.set(ix.log_mono_ms1, r, log_mono); + m.set(ix.has_ms1, r, has_ms1); + m.set(ix.log_sn, r, ff.log_sn); + m.set(ix.n_observations, r, ff.n_observations); + m.set(ix.base_width_rt, r, ff.base_width_rt); m.set( - "seed_score", + ix.seed_score, r, *seed_score_map.get(&cid[i]).unwrap_or(&0.0), ); m.set( - "seed_identified", + ix.seed_identified, r, *seed_id_map.get(&cid[i]).unwrap_or(&0.0), ); m.set( - "matched_fraction", + ix.matched_fraction, r, n_matched[i] as f64 / (n_pred[i].max(1) as f64), ); - m.set("profile_cos", r, ff.profile_cos); - m.set("ref_corr", r, ff.ref_corr); - m.set("best_ref_corr", r, ff.best_ref_corr); - m.set("low_frag_coel", r, ff.low_frag_coel); - m.set("evidence", r, ff.evidence); - m.set("contrast_min", r, ff.contrast_min); - m.set("resid_corr", r, ff.resid_corr); - m.set("coel_clean", r, ff.coel_clean); - m.set("shadow_frac", r, ff.shadow_frac); - m.set("peak_contested_frac", r, contested[i]); - m.set("peak_contested_count_frac", r, contested_count[i]); - m.set("peak_apportioned_frac", r, apportioned[i]); + m.set(ix.profile_cos, r, ff.profile_cos); + m.set(ix.ref_corr, r, ff.ref_corr); + m.set(ix.best_ref_corr, r, ff.best_ref_corr); + m.set(ix.low_frag_coel, r, ff.low_frag_coel); + m.set(ix.evidence, r, ff.evidence); + m.set(ix.contrast_min, r, ff.contrast_min); + m.set(ix.resid_corr, r, ff.resid_corr); + m.set(ix.coel_clean, r, ff.coel_clean); + m.set(ix.shadow_frac, r, ff.shadow_frac); + m.set(ix.peak_contested_frac, r, contested[i]); + m.set(ix.peak_contested_count_frac, r, contested_count[i]); + m.set(ix.peak_apportioned_frac, r, apportioned[i]); // Extended battery (opt-in). Built once per PSM above and fanned out to the - // family modules; pushed here under the fixed registry-order names. - if extended { - for (k, v) in ext_names.iter().zip(&per[r].ext) { - m.set(k, r, *v); + // family modules; written here under the fixed registry-order columns. + if n_ext > 0 { + let row = &ext_vals[r * n_ext..(r + 1) * n_ext]; + for (c, v) in ix.ext.iter().zip(row) { + m.set(*c, r, *v); } } // Cross-charge corroboration features (whole-run reductions, indexed by row). - m.set("n_charge_states", r, f_n_charge[i]); - m.set("charge_multi_flag", r, f_charge_multi[i]); - m.set("cross_charge_intensity_log", r, f_cross_charge_int[i]); + m.set(ix.n_charge_states, r, f_n_charge[i]); + m.set(ix.charge_multi_flag, r, f_charge_multi[i]); + m.set(ix.cross_charge_intensity_log, r, f_cross_charge_int[i]); prelim[r] = n_matched[i] as f64 * (0.5 + ff.frag_corr.max(0.0)) + ff.coelution_mean.max(0.0) + (1.0 + apex_int[i] as f64).ln() * 0.1 - rt_err / gradient; } - drop(per); + drop(frag_feats); + drop(ext_vals); max_matrix_bytes = max_matrix_bytes.max(m.payload_bytes()); // PIN first: it reads the same values the columns below move into Arrow. @@ -1668,18 +2047,32 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result Col::F64("precursor_mz".into(), mz[lo..hi].to_vec()), Col::F64("prelim_score".into(), prelim), ]; + // Moved, not copied: the matrix is dropped on the next line, and copying + // every column is what made the matrix exist twice at this point. for (c, name) in cols_active.iter().enumerate() { - cols.push(Col::F64(name.clone(), m.column(c).to_vec())); + cols.push(Col::F64(name.clone(), m.take_column(c))); } drop(m); - writer.write_cols(cols)?; + // The parquet encode is single-threaded and ran here, between one chunk's + // computation and the next chunk's: the writer thread below takes it off the + // critical path, the way the loader thread took the decode off it. The + // channel is a rendezvous, so at most one chunk's columns wait behind the + // one being written. + if wtx.send(Some(cols)).is_err() { + // The writer failed; its error is the real one, so surface that. + break; + } } - Ok(()) + // Commit only after every chunk went out; a `break` above leaves this send + // failing, and the writer's own error is what `join` then returns. + let _ = wtx.send(None); + drop(wtx); + wh.join() + .map_err(|_| anyhow!("features: the parquet writer thread panicked"))? })?; if let Some(w) = pin { w.finish()?; } - let rows = writer.close()?; crate::memlog::report( "features chromatogram store", @@ -1697,9 +2090,17 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result chrom_rows = chrom_rows_total, "mem: features chromatogram store shape" ); + // Two matrices can be alive at once and the report used to name only one, so it + // understated the peak by exactly 2x. It did so for the wrong reason as well: the + // second copy was `column(c).to_vec()` into the writer's columns, pure waste. The + // columns are moved now, and the second matrix that remains is the next chunk's, + // built while the writer thread encodes this one -- overlap, not a copy. crate::memlog::report( "features value matrix", - &[("largest_chunk", max_matrix_bytes)], + &[ + ("largest_chunk", max_matrix_bytes), + ("in_flight_peak", 2 * max_matrix_bytes), + ], ); // Feature schema companion. @@ -1735,7 +2136,7 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result Ok(rows) } -#[derive(Default)] +#[derive(Clone, Default)] struct FragFeatures { frag_corr: f64, frag_cosine: f64, @@ -1895,25 +2296,13 @@ fn elution_peak_rt_bounds( frac: f64, grace: usize, ) -> Option<(f32, f32)> { - let mut axis: Vec = rows.iter().flat_map(|r| r.rt.iter().cloned()).collect(); - axis.sort_by(|a, b| a.total_cmp(b)); - axis.dedup(); + let TraceAlign { + axis_full: axis, + traces_full: traces, + } = align_traces(rows); if axis.len() < 3 { return None; } - let traces: Vec> = rows - .iter() - .map(|r| { - let map: HashMap = - r.rt.iter() - .zip(r.inten.iter()) - .map(|(&t, &v)| (t.to_bits(), v)) - .collect(); - axis.iter() - .map(|t| *map.get(&t.to_bits()).unwrap_or(&0.0) as f64) - .collect() - }) - .collect(); let mut ord: Vec = (0..rows.len()).collect(); ord.sort_by(|&a, &b| rows[b].pred_int.total_cmp(&rows[a].pred_int)); let k3: Vec = ord.into_iter().take(3).collect(); @@ -1935,8 +2324,12 @@ fn elution_peak_rt_bounds( Some((axis[lo], axis[hi])) } +// The peak-bounding knobs (`bound`, `frac`, `grace`, `global_bounds`) are passed through +// from the config as they always were; the alignment is the eighth. +#[allow(clippy::too_many_arguments)] fn fragment_features( rows: &[ChromRow], + al: &TraceAlign, apex_rt: f64, coel_thresh: f64, bound: bool, @@ -2008,26 +2401,12 @@ fn fragment_features( 0.0 }; - // Align traces on the union RT axis, then restrict to the elution PEAK so the - // trace-based features below are computed over the peak, not the whole extracted - // RT window (which spans +/- w_rt and would dilute co-elution/profile scores). - let mut axis_full: Vec = rows.iter().flat_map(|r| r.rt.iter().cloned()).collect(); - axis_full.sort_by(|a, b| a.total_cmp(b)); - axis_full.dedup(); - let traces_full: Vec> = rows - .iter() - .map(|r| { - let map: HashMap = - r.rt.iter() - .zip(r.inten.iter()) - .map(|(&t, &v)| (t.to_bits(), v)) - .collect(); - axis_full - .iter() - .map(|t| *map.get(&t.to_bits()).unwrap_or(&0.0) as f64) - .collect() - }) - .collect(); + // Traces aligned on the union RT axis (built once per PSM by the caller), then + // restricted to the elution PEAK so the trace-based features below are computed over + // the peak, not the whole extracted RT window (which spans +/- w_rt and would dilute + // co-elution/profile scores). + let axis_full = &al.axis_full; + let traces_full = &al.traces_full; let (lo_i, hi_i) = if bound && axis_full.len() >= 3 { // boundary on the smoothed summed top-3-predicted-fragment profile, around apex let ai = axis_full @@ -2041,7 +2420,7 @@ fn fragment_features( .map(|(i, _)| i) .unwrap_or(0); match global_bounds { - Some((l, r)) => global_bound_indices(&axis_full, apex_rt, ai, l, r), + Some((l, r)) => global_bound_indices(axis_full, apex_rt, ai, l, r), None => { let mut ord: Vec = (0..pred.len()).collect(); ord.sort_by(|&a, &b| pred[b].total_cmp(&pred[a])); @@ -2105,12 +2484,16 @@ fn fragment_features( // pCos: at each scan, cosine(observed fragment vector, predicted vector), // weighted by reference-profile^2 (concentrates on the elution peak). let (mut num, mut den) = (0.0, 0.0); + // One scratch vector for the whole scan loop; it used to be a fresh allocation + // per scan, tens of millions of them over a run. + let mut obs_k: Vec = Vec::with_capacity(traces.len()); for k in 0..np { let w = refp[k] * refp[k]; if w <= 0.0 { continue; } - let obs_k: Vec = traces.iter().map(|tr| tr[k]).collect(); + obs_k.clear(); + obs_k.extend(traces.iter().map(|tr| tr[k])); num += cosine(&obs_k, &pred) * w; den += w; } @@ -2142,8 +2525,12 @@ fn fragment_features( .map(|k| traces.iter().map(|tr| tr[k]).sum::()) .collect(); let mut contrasts = Vec::with_capacity(traces.len()); + // Scratch reused across fragments; the per-fragment allocation it replaces was + // one heap block per fragment per PSM. + let mut others: Vec = Vec::with_capacity(np); for tr in &traces { - let others: Vec = (0..np).map(|k| total[k] - tr[k]).collect(); + others.clear(); + others.extend((0..np).map(|k| total[k] - tr[k])); contrasts.push(pearson(tr, &others)); } f.contrast_min = contrasts.iter().cloned().fold(f64::MAX, f64::min); @@ -2359,4 +2746,608 @@ mod tests { assert_eq!(lag, 0); assert!(shape > 0.99); } + + // --- trace alignment (the fast path must reproduce the union-and-map build) --- + + fn row<'a>(name: &'a str, pred: f32, rt: &'a [f32], inten: &'a [f32]) -> ChromRow<'a> { + ChromRow { + frag_name: name, + frag_mz: 100.0, + frag_obs_mz: 100.0, + pred_int: pred, + rt, + inten, + } + } + + fn same(a: &TraceAlign, b: &TraceAlign) { + assert_eq!( + a.axis_full.iter().map(|x| x.to_bits()).collect::>(), + b.axis_full.iter().map(|x| x.to_bits()).collect::>(), + "axis differs" + ); + assert_eq!(a.traces_full.len(), b.traces_full.len(), "fragment count"); + for (i, (x, y)) in a.traces_full.iter().zip(&b.traces_full).enumerate() { + assert_eq!( + x.iter().map(|v| v.to_bits()).collect::>(), + y.iter().map(|v| v.to_bits()).collect::>(), + "fragment {i} trace differs" + ); + } + } + + #[test] + fn shared_axis_alignment_reproduces_the_union_build() { + // The normal case: every fragment of a candidate samples ONE stored axis, so the + // fast path applies and must be bit-identical to the map build it skips. + let grid = [10.0f32, 11.0, 12.0, 13.0]; + let a = [0.0f32, 4.0, 9.0, 2.0]; + let b = [1.0f32, 5.0, 8.0, 0.0]; + let rows = vec![row("y1", 1.0, &grid, &a), row("y2", 0.5, &grid, &b)]; + assert!( + align_shared_axis(&rows).is_some(), + "rows sharing one stored axis must take the fast path" + ); + same(&align_traces(&rows), &align_union(&rows)); + + // A predicted-but-unobserved fragment carries no axis at all; it is zero-filled + // by both builds and must not disable the fast path for the others. + let rows = vec![ + row("y1", 1.0, &grid, &a), + row("y2", 0.5, &[], &[]), + row("y3", 0.2, &grid, &b), + ]; + assert!(align_shared_axis(&rows).is_some()); + let al = align_traces(&rows); + same(&al, &align_union(&rows)); + assert_eq!( + al.traces_full[1], + vec![0.0; 4], + "empty trace is zero-filled" + ); + } + + #[test] + fn alignment_falls_back_when_the_axes_are_not_one_shared_grid() { + let g1 = [10.0f32, 11.0, 12.0]; + let g2 = [10.5f32, 11.0, 12.5]; + let v1 = [1.0f32, 2.0, 3.0]; + let v2 = [4.0f32, 5.0, 6.0]; + // Two different grids: the union has five points and each row is zero elsewhere. + let rows = vec![row("y1", 1.0, &g1, &v1), row("y2", 1.0, &g2, &v2)]; + assert!(align_shared_axis(&rows).is_none(), "distinct axes"); + let al = align_traces(&rows); + assert_eq!(al.axis_full, vec![10.0, 10.5, 11.0, 12.0, 12.5]); + assert_eq!(al.traces_full[0], vec![1.0, 0.0, 2.0, 3.0, 0.0]); + assert_eq!(al.traces_full[1], vec![0.0, 4.0, 5.0, 0.0, 6.0]); + + // Equal VALUES but stored separately are not the same slice: still correct, just + // through the general path. + let copy = g1; + let rows = vec![row("y1", 1.0, &g1, &v1), row("y2", 1.0, ©, &v2)]; + assert!(align_shared_axis(&rows).is_none()); + same(&align_traces(&rows), &align_union(&rows)); + + // A non-ascending axis would be reordered by the sort, and a repeated RT value + // would be collapsed by the dedup; both must refuse the fast path. + let unsorted = [12.0f32, 10.0, 11.0]; + let rows = vec![row("y1", 1.0, &unsorted, &v1)]; + assert!(align_shared_axis(&rows).is_none(), "unsorted axis"); + same(&align_traces(&rows), &align_union(&rows)); + let dup = [10.0f32, 10.0, 11.0]; + let rows = vec![row("y1", 1.0, &dup, &v1)]; + assert!(align_shared_axis(&rows).is_none(), "repeated RT"); + same(&align_traces(&rows), &align_union(&rows)); + } + + // --- cross-charge corroboration (dense ids must reproduce the set build) --- + + #[test] + fn charge_state_bitmask_matches_the_set_build() { + // Five rows over three peptidoforms, charges repeated within a peptidoform. + let pf = [0u32, 0, 1, 1, 2, 0]; + let z = [2i32, 3, 2, 2, 4, 2]; + let by_mask = charge_states_by_mask(&pf, 3, &z).expect("charges are in 0..32"); + assert_eq!(by_mask, charge_states_by_set(&pf, 3, &z)); + assert_eq!(by_mask, vec![2.0, 2.0, 1.0, 1.0, 1.0, 2.0]); + assert_eq!(charge_states_per_row(&pf, 3, &z), by_mask); + + // Charge 0 and charge 31 are inside the mask; 32 and a negative are not, and the + // fallback then carries the count. + let edge = [0i32, 31, 0, 31, 0, 0]; + assert_eq!( + charge_states_by_mask(&pf, 3, &edge).expect("0 and 31 fit"), + charge_states_by_set(&pf, 3, &edge) + ); + let out = [32i32, 3, 2, 2, 4, 2]; + assert!(charge_states_by_mask(&pf, 3, &out).is_none()); + assert_eq!( + charge_states_per_row(&pf, 3, &out), + charge_states_by_set(&pf, 3, &out) + ); + let neg = [-1i32, 3, 2, 2, 4, 2]; + assert!(charge_states_by_mask(&pf, 3, &neg).is_none()); + assert_eq!( + charge_states_per_row(&pf, 3, &neg), + charge_states_by_set(&pf, 3, &neg) + ); + } + + // --- the feature -> column permutation --- + + #[test] + fn colix_covers_every_active_column() { + for set in [FeatureSet::Minimal, FeatureSet::Rich, FeatureSet::Extended] { + let cols = active_features(set); + let ext = extended_name_refs(); + let ix = ColIx::new(&cols, &ext); + // Every column the assembly resolves must agree with the name lookup it + // replaced, and together they must account for each active column exactly + // once: a mistyped field would resolve to None and leave a column unwritten. + let mut seen: Vec = ix + .named() + .into_iter() + .chain(ix.ext.iter().copied()) + .flatten() + .collect(); + seen.sort_unstable(); + assert_eq!( + seen, + (0..cols.len()).collect::>(), + "{set:?}: the assembly does not write each active column exactly once" + ); + // The extended names resolve only under the Extended set. + if !matches!(set, FeatureSet::Extended) { + assert!(ix.ext.iter().all(|c| c.is_none()), "{set:?}"); + } + } + } + + #[test] + fn value_matrix_hands_its_columns_over_by_move() { + let mut m = ValueMatrix::new(3, 2); + m.set(Some(1), 0, 7.5); + m.set(Some(1), 1, -1.0); + m.set(None, 0, 99.0); // a feature outside the active set is dropped + assert_eq!(m.column(1), &[7.5, -1.0]); + assert_eq!(m.take_column(1), vec![7.5, -1.0]); + assert_eq!(m.column(0), &[0.0, 0.0]); + } + + // --- confident-bounds row-group pruning --- + + fn stats(groups: &[(usize, u32, u32)]) -> Vec { + groups + .iter() + .map(|&(rows, lo, hi)| mumdia_io::table::RowGroupStats { + rows, + min: Some(lo as f64), + max: Some(hi as f64), + }) + .collect() + } + + #[test] + fn confident_row_spans_keep_every_group_that_can_hold_an_anchor() { + // Four groups of 10 rows covering candidate ids 0-9, 10-19, 20-29, 30-39. + let s = stats(&[(10, 0, 9), (10, 10, 19), (10, 20, 29), (10, 30, 39)]); + // Anchors in groups 0 and 2 only. + assert_eq!( + confident_row_spans(&s, &[3, 25]).unwrap(), + vec![(0, 10), (20, 10)] + ); + // Adjacent kept groups merge into one span. + assert_eq!( + confident_row_spans(&s, &[3, 15]).unwrap(), + vec![(0, 20)], + "adjacent groups must merge" + ); + // An anchor in a gap between two groups' ranges keeps neither. + let gapped = stats(&[(10, 0, 9), (10, 20, 29)]); + assert!(confident_row_spans(&gapped, &[15]).unwrap().is_empty()); + // No anchors at all: nothing to read. + assert!(confident_row_spans(&s, &[]).unwrap().is_empty()); + // Every group holding an anchor is kept, and the kept spans cover them. + let all: Vec = (0..40).collect(); + assert_eq!(confident_row_spans(&s, &all).unwrap(), vec![(0, 40)]); + // Missing statistics mean "read everything", the behaviour before the pruning. + let blind = vec![mumdia_io::table::RowGroupStats { + rows: 10, + min: None, + max: None, + }]; + assert!(confident_row_spans(&blind, &[3]).is_none()); + assert!(confident_row_spans(&[], &[3]).is_none()); + } + + /// A chromatogram table with several row groups, its confident candidates, and the + /// PSM apex RTs, for the pruning-equality test below. + fn craft_chrom_for_bounds(path: &str) -> (HashMap>, Vec) { + let (mut cid, mut name): (Vec, Vec) = (Vec::new(), Vec::new()); + let (mut fmz, mut pint): (Vec, Vec) = (Vec::new(), Vec::new()); + let (mut rt, mut inten): (Vec>, Vec>) = (Vec::new(), Vec::new()); + let mut apex_rt: Vec = Vec::new(); + let mut confident: HashMap> = HashMap::new(); + for c in 0..200u32 { + let apex = 100.0 + c as f64 * 5.0; + apex_rt.push(apex); + // Every fifth candidate is a confident anchor, so most row groups hold none. + if c % 5 == 0 { + confident.insert(c, vec![c as usize]); + } + let grid: Vec = (0..9).map(|k| (apex - 4.0 + k as f64) as f32).collect(); + for f in 0..6u32 { + cid.push(c); + name.push(format!("y{}", f + 1)); + fmz.push(200.0 + f as f64 * 30.0); + pint.push(1.0 / (f + 1) as f32); + rt.push(grid.clone()); + inten.push( + (0..9) + .map(|k| { + let x = k as f32 - 4.0; + (100.0 - f as f32 * 10.0) / (1.0 + x * x) + }) + .collect(), + ); + } + } + // Small row groups (two candidates each), so the pruning has something to prune. + let mut w = TableWriter::new(path).with_row_group_rows(12); + w.write_cols(vec![ + Col::U32("candidate_id".into(), cid), + Col::Str("frag_name".into(), name), + Col::F64("frag_mz".into(), fmz), + Col::F32("predicted_intensity".into(), pint), + Col::ListF32("rt".into(), rt), + Col::ListF32("intensity".into(), inten), + ]) + .unwrap(); + w.close().unwrap(); + (confident, apex_rt) + } + + #[test] + fn pruned_confident_bounds_equal_the_whole_table_scan() { + let dir = std::env::temp_dir().join("mumdia_features_bounds"); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir + .join("chrom_bounds.parquet") + .to_string_lossy() + .to_string(); + let (confident, apex_rt) = craft_chrom_for_bounds(&path); + let ch = TableFile::open(&path).unwrap(); + let cfg = FeaturesConfig::default(); + + let mut sorted: Vec = confident.keys().copied().collect(); + sorted.sort_unstable(); + let rg = ch.row_group_stats("candidate_id").unwrap(); + let spans = confident_row_spans(&rg, &sorted).expect("the writer records statistics"); + let read: usize = spans.iter().map(|s| s.1).sum(); + assert!( + read < ch.nrows, + "the pruning must skip something: read {read} of {} rows", + ch.nrows + ); + assert!(read > 0); + + // Same half-widths at several chunk sizes, pruned against the whole-table scan. + // The chunk size matters: a candidate that straddles a chunk boundary is bounded + // from each part separately, and the pruned pass has to split it in exactly the + // same places for the percentiles to come out the same. + for chunk in [7usize, 16, 64, 1 << 20] { + let pruned = + confident_global_bounds(&path, &spans, &confident, &apex_rt, &cfg, chunk).unwrap(); + let whole = + confident_global_bounds(&path, &[(0, ch.nrows)], &confident, &apex_rt, &cfg, chunk) + .unwrap(); + let bits = |v: Option<(f64, f64)>| v.map(|(l, r)| (l.to_bits(), r.to_bits())); + assert_eq!( + bits(pruned), + bits(whole), + "chunk {chunk}: pruned half-widths differ from the whole-table scan" + ); + assert!(pruned.is_some(), "chunk {chunk}: 40 anchors is enough"); + } + } + + // --- the whole chunked pass, end to end --- + + /// Crafted psms_extracted + chromatograms for the Extended feature set: shared window + /// grids, a predicted-but-unobserved fragment, candidates with no chromatogram rows at + /// all, MS1 isotope XICs, b/y and multiply-charged fragment names, and several charge + /// states per peptidoform. + fn craft_extended_inputs(dir: &std::path::Path, n_cand: u32) -> (String, String) { + let psms = dir.join("psms_ext.parquet").to_string_lossy().to_string(); + let chrom = dir.join("chrom_ext.parquet").to_string_lossy().to_string(); + let (mut cid, mut base): (Vec, Vec) = (Vec::new(), Vec::new()); + let (mut apex_rt, mut cal, mut mz): (Vec, Vec, Vec) = + (Vec::new(), Vec::new(), Vec::new()); + let mut apex_int: Vec = Vec::new(); + let (mut nmatch, mut corun, mut z): (Vec, Vec, Vec) = + (Vec::new(), Vec::new(), Vec::new()); + let (mut label, mut pform, mut prot): (Vec, Vec, Vec) = + (Vec::new(), Vec::new(), Vec::new()); + let (mut ccid, mut cname): (Vec, Vec) = (Vec::new(), Vec::new()); + let (mut cfmz, mut cobsmz): (Vec, Vec) = (Vec::new(), Vec::new()); + let mut cpint: Vec = Vec::new(); + let (mut crt, mut cint): (Vec>, Vec>) = (Vec::new(), Vec::new()); + for c in 0..n_cand { + let apex = 100.0 + c as f64 * 7.0; + cid.push(c); + base.push(c / 3); + apex_rt.push(apex); + cal.push(apex - 1.5); + mz.push(500.0 + c as f64); + apex_int.push(1000.0 + c as f32); + nmatch.push(3); + corun.push(2); + z.push(2 + (c % 3) as i32); + label.push(if c % 3 == 0 { "decoy" } else { "target" }.to_string()); + // Three charge states share a peptidoform, so the cross-charge reduction has + // something to group. + pform.push(format!("PEPTIDEK{}", c / 3)); + prot.push(format!("P{c};Q{c}")); + if c % 7 == 6 { + continue; // a PSM row with no chromatogram rows + } + let npts = 5 + (c % 6) as usize; + let grid: Vec = (0..npts).map(|k| (apex - 4.0 + k as f64) as f32).collect(); + let nfrag = 2 + (c % 4); + for f in 0..nfrag { + ccid.push(c); + cname.push(if f % 2 == 0 { + format!("y{}", f + 1) + } else { + format!("b{}^2", f + 1) + }); + cfmz.push(200.0 + f as f64 * 30.0); + cobsmz.push(200.0 + f as f64 * 30.0 + 0.0001 * (c % 5) as f64); + cpint.push(1.0 / (f + 1) as f32); + if f == nfrag - 1 && c % 4 == 0 { + crt.push(Vec::new()); // predicted but never observed + cint.push(Vec::new()); + } else { + crt.push(grid.clone()); + cint.push( + (0..npts) + .map(|k| { + let x = k as f32 - (npts as f32) / 2.0; + ((-x * x / 3.0).exp() * (100.0 - f as f32 * 9.0)).max(0.0) + }) + .collect(), + ); + } + } + for (iso, nm) in ["ms1_mono", "ms1_iso1", "ms1_iso2"].iter().enumerate() { + ccid.push(c); + cname.push(nm.to_string()); + cfmz.push(500.0 + c as f64 + iso as f64 * 0.5); + cobsmz.push(500.0 + c as f64 + iso as f64 * 0.5); + cpint.push(0.0); + crt.push(grid.clone()); + cint.push((0..npts).map(|k| (k as f32 + 1.0) * 50.0).collect()); + } + } + let n = cid.len(); + mumdia_io::table::write_table( + &psms, + vec![ + Col::U32("candidate_id".into(), cid), + Col::F64("apex_rt".into(), apex_rt), + Col::F32("apex_intensity".into(), apex_int), + Col::I32("n_matched_fragments".into(), nmatch), + Col::I32("coelution_run".into(), corun), + Col::F64("rt_pred_cal".into(), cal), + Col::I32("charge".into(), z), + Col::Str("label".into(), label), + Col::U32("base_peptide_id".into(), base), + Col::Str("peptidoform".into(), pform), + Col::Str("protein".into(), prot), + Col::F64("precursor_mz".into(), mz), + Col::OptF64("ms1_mono".into(), vec![Some(900.0); n]), + Col::OptF64("ms1_iso1".into(), vec![Some(450.0); n]), + Col::OptF64("ms1_iso2".into(), vec![Some(120.0); n]), + Col::OptF64("ms1_isom1".into(), vec![Some(30.0); n]), + ], + ) + .unwrap(); + mumdia_io::table::write_table( + &chrom, + vec![ + Col::U32("candidate_id".into(), ccid), + Col::Str("frag_name".into(), cname), + Col::F64("frag_mz".into(), cfmz), + Col::F64("frag_obs_mz".into(), cobsmz), + Col::F32("predicted_intensity".into(), cpint), + Col::ListF32("rt".into(), crt), + Col::ListF32("intensity".into(), cint), + ], + ) + .unwrap(); + (psms, chrom) + } + + #[test] + fn extended_features_are_chunk_invariant() { + // The `features_chunking_is_value_preserving` integration test runs the DEFAULT + // (Minimal) set with no seed, so it never reaches the extended battery, the + // per-PSM extended buffer, the extended half of the column permutation, or the + // trace alignment's fast path. This one does, over the full chunked pass: + // loader thread, parallel per-PSM map, serial assembly, writer thread. + let dir = std::env::temp_dir().join("mumdia_features_extended"); + std::fs::create_dir_all(&dir).unwrap(); + let (psms, chrom) = craft_extended_inputs(&dir, 61); + let run = |tag: &str, chunk_rows: usize| -> String { + let out = dir + .join(format!("features_{tag}.parquet")) + .to_string_lossy() + .to_string(); + let pin = dir + .join(format!("features_{tag}.pin")) + .to_string_lossy() + .to_string(); + let mut cfg = FeaturesConfig { + set: FeatureSet::Extended, + emit_pin: true, + // The confident-bounds pass depends on the chunk size by construction + // (it bounds a candidate from each side of a chunk boundary separately), + // so it is off here; `pruned_confident_bounds_equal_the_whole_table_scan` + // covers that pass instead. + bound_from_confident: false, + ..Default::default() + }; + cfg.ms1_precursor_features = true; + run_with_chunk_rows( + FeaturesParams { + psms: &psms, + chromatograms: &chrom, + seed: None, + out: &out, + out_pin: &pin, + cfg: &cfg, + config_hash: "test", + }, + chunk_rows, + ) + .unwrap(); + out + }; + let one = run("one", 1 << 20); + let many = run("many", 1); + + let a = mumdia_io::table::Table::read(&one).unwrap(); + let b = mumdia_io::table::Table::read(&many).unwrap(); + assert_eq!(a.nrows, 61); + assert_eq!(a.nrows, b.nrows); + assert_eq!(a.column_names(), b.column_names()); + assert_eq!( + a.column_names().len(), + active_features(FeatureSet::Extended).len() + NON_FEATURE_COLUMNS.len() - 3, + "metadata columns plus every active feature" + ); + // Every f64 column bit for bit, so a feature that differs in the last ulp fails. + let mut checked = 0usize; + for name in a.column_names() { + if let (Ok(x), Ok(y)) = (a.f64(&name), b.f64(&name)) { + let xb: Vec = x.iter().map(|v| v.to_bits()).collect(); + let yb: Vec = y.iter().map(|v| v.to_bits()).collect(); + assert_eq!(xb, yb, "column '{name}' differs between chunk sizes"); + checked += 1; + } + } + assert!(checked > 300, "only {checked} f64 columns compared"); + // The cross-charge reduction is a whole-run quantity: three charge states per + // peptidoform, so it must read 3 for the peptidoforms that have all three. + let n_charge = a.f64("n_charge_states").unwrap(); + assert!( + n_charge.contains(&3.0), + "the charge-state count never reached 3" + ); + assert!(n_charge.iter().all(|&v| (1.0..=3.0).contains(&v))); + } + + #[test] + fn a_failed_chunk_pass_neither_hangs_nor_publishes_a_partial_table() { + // Both halves of this are about the stage's write and load having moved onto + // their own threads: an error must still come back (the loader must not be left + // blocked on a rendezvous send that nobody will receive), and the writer must not + // publish what it has over the previous good artifact. + let dir = std::env::temp_dir().join("mumdia_features_failure"); + std::fs::create_dir_all(&dir).unwrap(); + let (psms, chrom) = craft_extended_inputs(&dir, 13); + let out = dir + .join("features_fail.parquet") + .to_string_lossy() + .to_string(); + let pin = dir.join("features_fail.pin").to_string_lossy().to_string(); + let cfg = FeaturesConfig { + set: FeatureSet::Rich, + bound_from_confident: false, + ..Default::default() + }; + let go = |chrom_path: &str| { + run_with_chunk_rows( + FeaturesParams { + psms: &psms, + chromatograms: chrom_path, + seed: None, + out: &out, + out_pin: &pin, + cfg: &cfg, + config_hash: "test", + }, + 4, + ) + }; + go(&chrom).expect("the good run must succeed"); + let good = mumdia_io::hash::blake3_file(&out).unwrap(); + + // The same candidates, but the chromatogram table is missing a column the stream + // requires, so the loader thread fails after the plan has been made. + let broken = dir + .join("chrom_broken.parquet") + .to_string_lossy() + .to_string(); + let src = mumdia_io::table::Table::read(&chrom).unwrap(); + mumdia_io::table::write_table( + &broken, + vec![ + Col::U32("candidate_id".into(), src.u32("candidate_id").unwrap()), + Col::Str("frag_name".into(), src.str("frag_name").unwrap()), + Col::F64("frag_mz".into(), src.f64("frag_mz").unwrap()), + Col::ListF32("rt".into(), src.list_f32("rt").unwrap()), + Col::ListF32("intensity".into(), src.list_f32("intensity").unwrap()), + ], + ) + .unwrap(); + + // On a worker thread with a deadline, so a regression that reintroduces the hang + // fails this test instead of wedging the suite. + let (tx, rx) = std::sync::mpsc::channel(); + let (b, p2, o2, pin2) = (broken.clone(), psms.clone(), out.clone(), pin.clone()); + std::thread::spawn(move || { + let cfg = FeaturesConfig { + set: FeatureSet::Rich, + bound_from_confident: false, + ..Default::default() + }; + let r = run_with_chunk_rows( + FeaturesParams { + psms: &p2, + chromatograms: &b, + seed: None, + out: &o2, + out_pin: &pin2, + cfg: &cfg, + config_hash: "test", + }, + 4, + ); + let _ = tx.send(r.is_err()); + }); + let failed = rx + .recv_timeout(std::time::Duration::from_secs(60)) + .expect("the stage must return, not hang, when the loader fails"); + assert!( + failed, + "a chromatogram table without predicted_intensity must be an error" + ); + assert_eq!( + mumdia_io::hash::blake3_file(&out).unwrap(), + good, + "a failed run must leave the previous features table exactly as it was" + ); + // And no temp file left behind. + let leftovers: Vec<_> = std::fs::read_dir(&dir) + .unwrap() + .filter_map(|e| e.ok()) + .map(|e| e.file_name().to_string_lossy().to_string()) + .filter(|n| n.contains("features_fail.parquet.tmp-")) + .collect(); + assert!( + leftovers.is_empty(), + "temp files left behind: {leftovers:?}" + ); + } } diff --git a/rust/mumdia/crates/mumdia/src/stages/features/coelution.rs b/rust/mumdia/crates/mumdia/src/stages/features/coelution.rs index c7f6859d..d7381cbf 100644 --- a/rust/mumdia/crates/mumdia/src/stages/features/coelution.rs +++ b/rust/mumdia/crates/mumdia/src/stages/features/coelution.rs @@ -140,7 +140,8 @@ pub fn values(e: &Evidence) -> Vec { let f_gain = mean_all(&rc) - f_mean_full; // --- pairwise Pearson (symmetric matrix) + cross-correlation over pairs --- - let mut corr = vec![vec![0.0f64; k]; k]; + // Row-major flat k x k, one allocation instead of k. `corr[a * k + b]`. + let mut corr = vec![0.0f64; k * k]; let mut pair_p: Vec = Vec::new(); let mut pair_w: Vec = Vec::new(); // l_a * l_b weights aligned with pair_p let mut xvals: Vec = Vec::new(); @@ -153,11 +154,11 @@ pub fn values(e: &Evidence) -> Vec { let mut chg_cnt = 0.0f64; if has_traces { for a in 0..k { - corr[a][a] = 1.0; + corr[a * k + a] = 1.0; for b in (a + 1)..k { let p = pearson(&traces[a], &traces[b]); - corr[a][b] = p; - corr[b][a] = p; + corr[a * k + b] = p; + corr[b * k + a] = p; pair_p.push(p); pair_w.push(e.pred[a] * e.pred[b]); let (lag, xv) = best_xcorr(&traces[a], &traces[b], MAXLAG); @@ -207,8 +208,8 @@ pub fn values(e: &Evidence) -> Vec { let half = k / 2; let hi_idx: Vec = order.iter().take(half).cloned().collect(); let lo_idx: Vec = order.iter().skip(k - half).cloned().collect(); - let f_hi = subset_pair_mean(&corr, &hi_idx); - let f_lo = subset_pair_mean(&corr, &lo_idx); + let f_hi = subset_pair_mean(&corr, k, &hi_idx); + let f_lo = subset_pair_mean(&corr, k, &lo_idx); let f_contrast = f_hi - f_lo; // 19. coelution_corr_entropy: histogram of pairwise pearson in [-1,1] @@ -538,9 +539,10 @@ fn weighted_mean_topk(rc: &[f64], pred: &[f64], k: usize) -> f64 { } } -/// Mean pairwise Pearson (from a precomputed symmetric matrix) over a subset of -/// fragment indices. Returns 0.0 when fewer than two indices are given. -fn subset_pair_mean(corr: &[Vec], idx: &[usize]) -> f64 { +/// Mean pairwise Pearson over a subset of fragment indices, read from a precomputed +/// symmetric `k x k` matrix stored row-major in one buffer. Returns 0.0 when fewer than +/// two indices are given. +fn subset_pair_mean(corr: &[f64], k: usize, idx: &[usize]) -> f64 { let m = idx.len(); if m < 2 { return 0.0; @@ -550,8 +552,8 @@ fn subset_pair_mean(corr: &[Vec], idx: &[usize]) -> f64 { for a in 0..m { for b in (a + 1)..m { let (ia, ib) = (idx[a], idx[b]); - if ia < corr.len() && ib < corr.len() { - s += corr[ia][ib]; + if ia < k && ib < k { + s += corr[ia * k + ib]; n += 1.0; } } diff --git a/rust/mumdia/crates/mumdia/src/stages/features/interference.rs b/rust/mumdia/crates/mumdia/src/stages/features/interference.rs index b96fa87b..70bb7006 100644 --- a/rust/mumdia/crates/mumdia/src/stages/features/interference.rs +++ b/rust/mumdia/crates/mumdia/src/stages/features/interference.rs @@ -357,29 +357,30 @@ pub fn values(e: &Evidence) -> Vec { let m = matched.len(); let (mut explained_variance_ratio, mut second_component_fraction) = (0.0, 0.0); if m >= 1 { - // Gram matrix G = X X^T over matched fragment peak traces (m x m). - let mut g = vec![vec![0.0f64; m]; m]; + // Gram matrix G = X X^T over matched fragment peak traces (m x m), row-major in + // one buffer: `g[i * m + j]`. One allocation instead of m. + let mut g = vec![0.0f64; m * m]; for i in 0..m { for j in i..m { let d = dot(&e.traces[matched[i]], &e.traces[matched[j]]); - g[i][j] = d; - g[j][i] = d; + g[i * m + j] = d; + g[j * m + i] = d; } } let mut trace = 0.0; - for (i, row) in g.iter().enumerate().take(m) { - trace += row[i]; + for i in 0..m { + trace += g[i * m + i]; } if trace > 0.0 { - let (lam1, v1) = power_top(&g); + let (lam1, v1) = power_top(&g, m); explained_variance_ratio = (lam1 / trace).clamp(0.0, 1.0); // Deflate and extract the second eigenvalue. for i in 0..m { for j in 0..m { - g[i][j] -= lam1 * v1[i] * v1[j]; + g[i * m + j] -= lam1 * v1[i] * v1[j]; } } - let (lam2, _) = power_top(&g); + let (lam2, _) = power_top(&g, m); second_component_fraction = (lam2.max(0.0) / trace).clamp(0.0, 1.0); } } @@ -544,20 +545,22 @@ fn mean_loo(e: &Evidence, set: &[usize], tp: usize) -> f64 { } } -/// Top eigenvalue and (normalized) eigenvector of a symmetric matrix via power -/// iteration. Returns (0.0, zeros) for an empty or degenerate matrix. -fn power_top(g: &[Vec]) -> (f64, Vec) { - let m = g.len(); +/// Top eigenvalue and (normalized) eigenvector of a symmetric `m x m` matrix stored +/// row-major in one buffer, via power iteration. Returns (0.0, zeros) for an empty or +/// degenerate matrix. +fn power_top(g: &[f64], m: usize) -> (f64, Vec) { if m == 0 { return (0.0, Vec::new()); } let mut v = vec![1.0 / (m as f64).sqrt(); m]; + // Two buffers reused across the 100 iterations; `nv` used to be a fresh allocation + // on every one of them. + let mut nv = vec![0.0f64; m]; for _ in 0..100 { - let mut nv = vec![0.0f64; m]; for i in 0..m { let mut s = 0.0; for j in 0..m { - s += g[i][j] * v[j]; + s += g[i * m + j] * v[j]; } nv[i] = s; } @@ -568,14 +571,14 @@ fn power_top(g: &[Vec]) -> (f64, Vec) { for x in nv.iter_mut() { *x /= norm; } - v = nv; + std::mem::swap(&mut v, &mut nv); } // Rayleigh quotient: v is unit-norm, so lambda = v^T G v. let mut gv = vec![0.0f64; m]; for i in 0..m { let mut s = 0.0; for j in 0..m { - s += g[i][j] * v[j]; + s += g[i * m + j] * v[j]; } gv[i] = s; } diff --git a/rust/mumdia/crates/mumdia/src/stages/features/nonzero.rs b/rust/mumdia/crates/mumdia/src/stages/features/nonzero.rs index b551ac3a..58d0e727 100644 --- a/rust/mumdia/crates/mumdia/src/stages/features/nonzero.rs +++ b/rust/mumdia/crates/mumdia/src/stages/features/nonzero.rs @@ -85,20 +85,34 @@ pub fn values(e: &Evidence) -> Vec { // present (both-positive), and over scans where any signal is present. let mut both = Vec::new(); let mut summpos = Vec::new(); + // The summed-positive gather does not depend on the pair, so it is done ONCE per + // fragment rather than twice per pair: it was three of the five Vecs this loop + // allocated per pair, which is O(k^2) allocations per PSM for an O(k) quantity. + let sm_gathered: Vec> = if sm.len() >= 2 { + e.traces + .iter() + .map(|tr| sm.iter().map(|&t| tr[t]).collect()) + .collect() + } else { + Vec::new() + }; + // The both-positive index set IS pair-dependent, so only its buffers are reused. + let mut idx: Vec = Vec::with_capacity(np); + let mut xa: Vec = Vec::with_capacity(np); + let mut xb: Vec = Vec::with_capacity(np); for a in 0..e.traces.len() { for b in (a + 1)..e.traces.len() { - let idx: Vec = (0..np) - .filter(|&t| e.traces[a][t] > 0.0 && e.traces[b][t] > 0.0) - .collect(); + idx.clear(); + idx.extend((0..np).filter(|&t| e.traces[a][t] > 0.0 && e.traces[b][t] > 0.0)); if idx.len() >= 2 { - let xa: Vec = idx.iter().map(|&t| e.traces[a][t]).collect(); - let xb: Vec = idx.iter().map(|&t| e.traces[b][t]).collect(); + xa.clear(); + xa.extend(idx.iter().map(|&t| e.traces[a][t])); + xb.clear(); + xb.extend(idx.iter().map(|&t| e.traces[b][t])); both.push(pearson(&xa, &xb)); } if sm.len() >= 2 { - let xa: Vec = sm.iter().map(|&t| e.traces[a][t]).collect(); - let xb: Vec = sm.iter().map(|&t| e.traces[b][t]).collect(); - summpos.push(pearson(&xa, &xb)); + summpos.push(pearson(&sm_gathered[a], &sm_gathered[b])); } } } From 8437b86224388070d853cc8b0ec2157b4726aae0 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:41:08 +0200 Subject: [PATCH 15/35] perf(pool): splice the band row groups instead of decoding them Measured on the running seven-file experiment, at the moment it reached the pool: the stage was writing 3.3 MB/s and reading 3.9 MB/s at 110% CPU, on a disk that reads 221 MB/s (dd, direct). One run's chromatograms are 63 bands x 1.09 GB, so at that rate pooling one run would have taken about five hours, and the experiment pools seven of them. The rows were being decoded from parquet, filtered, and encoded again for no reason: the bands already write library-wide ids, so the pooled table is the bands' own rows in the bands' own order. `SpliceWriter` (mumdia-io) appends a source file's encoded column chunks into the output with the parquet writer's `append_column`, which is what that API exists for. The output takes the template's parquet schema AND its Arrow metadata, so the `LargeList` chromatogram columns read back as `LargeList` and not `List` -- that distinction lives only in the key-value metadata, and losing it would have changed the type of the run's largest artifact. A test asserts it. Only the row groups that can hold a candidate the overlap dedup drops are decoded: the statistics on `candidate_id` say which, the rest are copied. A dropped candidate comes from a window that two bands share, so those row groups are at the two ends of a band. The rewritten group is encoded to a temporary file beside the band and spliced in at its own position, so the pooled row order is unchanged whether a group was copied or rewritten. Also here: `groups.parallel` is clamped to one less than the thread count. A band in flight parks a rayon worker on its extraction's accumulation channel while the probing tasks run on the others, so as many bands as threads leaves nothing to feed them and the run deadlocks. Reproduced on the fixture with `parallel = 2, --threads 2`: the process sat at 0.1 s of CPU with no error and no output. The production runs are far from the bound (8 bands, 128 threads), which is why this was never hit. Tests: the three pool tests cover a clean splice, a band whose losers sit in one of ten row groups (mixed copy and rewrite, asserting the pooled ids are the bands' rows in order), and the skipped `psms_extracted`. Co-Authored-By: Claude Opus 5 (1M context) --- docs/33_window_groups.md | 23 +- rust/mumdia/crates/mumdia-core/src/config.rs | 5 + rust/mumdia/crates/mumdia-io/src/table.rs | 134 +++++++++- rust/mumdia/crates/mumdia/src/stages/pool.rs | 249 ++++++++++++++---- .../crates/mumdia/src/stages/run_groups.rs | 23 +- 5 files changed, 375 insertions(+), 59 deletions(-) diff --git a/docs/33_window_groups.md b/docs/33_window_groups.md index d9e28ba4..cba2acfb 100644 --- a/docs/33_window_groups.md +++ b/docs/33_window_groups.md @@ -168,7 +168,18 @@ Two of the four tables an ungrouped run writes are not pooled, because nothing r nothing opened. Each pooled table is hashed once, for the manifest record and the report beside it -together. At experiment scale those tables are tens of GB, and hashing reads all of it. The feature and competed schema companions +together. At experiment scale those tables are tens of GB, and hashing reads all of it. + +The pooling itself is a byte copy. A band's rows are already in the order the pooled table +wants and already encoded, so `pool` splices each band's parquet row groups into the output +without decoding them (`mumdia_io::table::SpliceWriter`, the column-chunk append the parquet +writer exposes for concatenation). Only the row groups that hold a candidate the overlap +dedup drops are decoded, filtered and re-encoded, and window overlap puts those at the two +ends of a band. Measured on the production run before the change: decoding and re-encoding +ran at 3.5 MB/s on one core, so one run's 68 GB of chromatograms would have taken about five +hours, on a disk that reads at 221 MB/s. The spliced output holds the same rows in the same +order with the same values; its row groups are the bands' own, so it is not byte-identical to +a re-encoded pool. The feature and competed schema companions (`
.schema.json`, the classifier's column list) are copied from the first band; every band wrote the same one. Each pooled table gets a `.report.json` whose stage is `pool` and whose stats record the number of groups and the overlap duplicates removed. @@ -342,6 +353,16 @@ Three things this says: On this data the useful range is therefore 48 to 64 bands: about 40 GB per band, which is what a 100 GB desktop can run two of at a time, or one with room to spare. +### `groups.parallel` and the thread count + +A band in flight occupies one rayon worker, which then blocks on its own extraction's +accumulation channel while the probing tasks run on the other workers. So `groups.parallel` +must stay below `--threads`: with as many bands as threads every worker parks and the run +makes no progress at all (reproduced on the fixture at `parallel = 2, --threads 2`: the +process sat at 0.1 s of CPU indefinitely, with no error and no output). The orchestrator +clamps the value to `threads - 1` and warns. The production runs are far from the bound (8 +bands on 128 threads), which is why this was not visible before. + ## 9. What is not there yet - Groups run one after another in one process. Running them as child processes in parallel diff --git a/rust/mumdia/crates/mumdia-core/src/config.rs b/rust/mumdia/crates/mumdia-core/src/config.rs index b2efb5f9..36f312d0 100644 --- a/rust/mumdia/crates/mumdia-core/src/config.rs +++ b/rust/mumdia/crates/mumdia-core/src/config.rs @@ -1909,6 +1909,11 @@ pub struct GroupsConfig { /// machine, after checking one band's peak RSS: on a 203M-precursor library at 63 /// bands the largest band took 39 GB and the median far less. Results do not depend on /// it; bands are independent and their artifacts are pooled in band order either way. + /// + /// It must stay below the thread count: a band in flight parks one worker on its + /// accumulation channel, so as many bands as there are threads leaves nothing to do the + /// probing and the run deadlocks. A larger value is clamped to `threads - 1` with a + /// warning rather than hanging. pub parallel: usize, } impl Default for GroupsConfig { diff --git a/rust/mumdia/crates/mumdia-io/src/table.rs b/rust/mumdia/crates/mumdia-io/src/table.rs index b48e4b3f..b16c6749 100644 --- a/rust/mumdia/crates/mumdia-io/src/table.rs +++ b/rust/mumdia/crates/mumdia-io/src/table.rs @@ -21,9 +21,11 @@ use parquet::arrow::arrow_reader::{ }; use parquet::arrow::ArrowWriter; use parquet::basic::{Compression, LogicalType}; -use parquet::file::metadata::ParquetMetaData; +use parquet::column::writer::ColumnCloseResult; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; use parquet::file::properties::WriterProperties; use parquet::file::statistics::Statistics; +use parquet::file::writer::SerializedFileWriter; /// The codec every artifact is written with. Snappy by default, which is what released /// artifacts use and what the sidecars' pyarrow reads without configuration; @@ -243,6 +245,136 @@ fn snappy_props(row_group_rows: Option) -> WriterProperties { b.build() } +/// A parquet file assembled from the row groups of other parquet files, copied as bytes. +/// +/// Pooling a grouped run's band artifacts is a concatenation: the rows are already in the +/// order the pooled table wants, already encoded and already compressed. Decoding and +/// re-encoding them is what the pool used to do, and on the chromatogram tables it ran at +/// 3.5 MB/s on one core against a disk that does 221 MB/s -- five hours for one run's 68 GB +/// (measured on the production seven-file experiment). Splicing the column chunks byte for +/// byte costs a copy, so the pool becomes disk-bound. +/// +/// The output carries the template's parquet schema and its Arrow metadata, so a spliced +/// table reads back as the same types as the tables it came from -- including the +/// `LargeList` columns, whose distinction from `List` lives only in that metadata. Every +/// source must have that same schema; splicing is a byte copy and cannot convert anything. +/// +/// The values and the row order are exactly those of the sources. The row-group boundaries +/// are the sources' own, so a spliced file is not byte-identical to a re-encoded one. +pub struct SpliceWriter { + writer: Option>, + /// The Arrow schema the sources must share, for the caller's own checks. + pub schema: Arc, + rows: u64, + target: Option, +} + +/// The parsed footer of a file to splice from, with its page index when it has one. +fn splice_meta(path: &str) -> Result<(std::fs::File, ParquetMetaData)> { + let file = std::fs::File::open(path).with_context(|| format!("opening {path}"))?; + let meta = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&file) + .with_context(|| format!("reading the parquet footer of {path}"))?; + Ok((file, meta)) +} + +impl SpliceWriter { + /// Create `out`, taking the schema and the Arrow metadata from `template`, which is + /// normally the first file whose row groups will be spliced in. + pub fn create(out: &str, template: &str) -> Result { + let (_, meta) = splice_meta(template)?; + let fm = meta.file_metadata(); + let props = WriterProperties::builder() + .set_key_value_metadata(fm.key_value_metadata().cloned()) + .build(); + let schema = + parquet::arrow::parquet_to_arrow_schema(fm.schema_descr(), fm.key_value_metadata()) + .with_context(|| format!("reading the arrow schema of {template}"))?; + let target = AtomicPath::new(out)?; + let file = std::fs::File::create(target.tmp()) + .with_context(|| format!("creating {}", target.tmp().display()))?; + let writer = + SerializedFileWriter::new(file, fm.schema_descr().root_schema_ptr(), Arc::new(props)) + .with_context(|| format!("opening {out} for splicing"))?; + Ok(SpliceWriter { + writer: Some(writer), + schema: Arc::new(schema), + rows: 0, + target: Some(target), + }) + } + + /// Splice row groups of `src`: those whose index `keep` accepts, or all of them when + /// `keep` accepts everything. Returns the rows appended. + pub fn append_row_groups(&mut self, src: &str, keep: impl Fn(usize) -> bool) -> Result { + let (file, meta) = splice_meta(src)?; + let w = self.writer.as_mut().expect("writer closed"); + if meta.file_metadata().schema_descr() != w.schema_descr() { + return Err(anyhow!( + "{src} has a different parquet schema from the table being spliced into; the band artifacts must come from the same configuration" + )); + } + let column_indexes = meta.column_index(); + let offset_indexes = meta.offset_index(); + let mut rows = 0u64; + for (i, rg) in meta.row_groups().iter().enumerate() { + if !keep(i) { + continue; + } + let rg_column = column_indexes.and_then(|ci| ci.get(i)); + let rg_offset = offset_indexes.and_then(|oi| oi.get(i)); + let mut out = w.next_row_group()?; + for (j, col) in rg.columns().iter().enumerate() { + out.append_column( + &file, + ColumnCloseResult { + bytes_written: col.compressed_size() as u64, + rows_written: rg.num_rows() as u64, + metadata: col.clone(), + // The engine writes no bloom filters; the page index is carried + // through when the source has one. + bloom_filter: None, + column_index: rg_column.and_then(|r| r.get(j)).cloned(), + offset_index: rg_offset.and_then(|r| r.get(j)).cloned(), + }, + )?; + } + out.close()?; + rows += rg.num_rows() as u64; + } + self.rows += rows; + Ok(rows) + } + + /// The row groups of `src`, as (first row, row count) in file order. + pub fn row_group_spans(src: &str) -> Result> { + let (_, meta) = splice_meta(src)?; + let mut out = Vec::with_capacity(meta.num_row_groups()); + let mut start = 0usize; + for rg in meta.row_groups() { + let n = rg.num_rows() as usize; + out.push((start, n)); + start += n; + } + Ok(out) + } + + pub fn rows(&self) -> u64 { + self.rows + } + + pub fn close(mut self) -> Result { + if let Some(w) = self.writer.take() { + w.close().context("closing the spliced parquet file")?; + } + if let Some(t) = self.target.take() { + t.publish()?; + } + Ok(self.rows) + } +} + /// Rows per internal chunk in [`write_table`]. /// /// A multiple of the parquet writer's 1,024-value write batch, so the encoder sees the diff --git a/rust/mumdia/crates/mumdia/src/stages/pool.rs b/rust/mumdia/crates/mumdia/src/stages/pool.rs index 330c6b20..e8b0f25d 100644 --- a/rust/mumdia/crates/mumdia/src/stages/pool.rs +++ b/rust/mumdia/crates/mumdia/src/stages/pool.rs @@ -20,13 +20,12 @@ //! artifact. use std::collections::{HashMap, HashSet}; -use std::sync::Arc; use std::time::Instant; use anyhow::{anyhow, bail, Result}; use arrow::array::{Array, BooleanArray, UInt32Array}; use arrow::compute::filter_record_batch; -use mumdia_io::table::{BatchWriter, TableFile}; +use mumdia_io::table::{BatchWriter, SpliceWriter, TableFile}; use tracing::info; const BATCH_ROWS: usize = 1 << 16; @@ -105,69 +104,116 @@ fn overlap_losers(bands: &[BandArtifacts]) -> Result<(Vec>, u64)> { Ok((losers, duplicates)) } -/// Append every band's table into `out`, ids offset, losers dropped. Returns the row count. +/// Append every band's table into `out`, losers dropped. Returns the row count. +/// +/// The bands already wrote library-wide ids (`Library::global_offset` is added on the way +/// out of extract), so pooling is a concatenation and almost every row group can be spliced +/// into the output as bytes, without being decoded. That matters at scale: decoding and +/// re-encoding one run's 68 GB of chromatograms ran at 3.5 MB/s on one core, five hours, +/// against a disk that does 221 MB/s. Only the row groups that actually hold a dropped +/// candidate are decoded, filtered and re-encoded, and window overlap puts those at the two +/// ends of a band. fn pool_table(paths: impl Iterator)>, out: &str) -> Result { - let mut writer: Option = None; - let mut schema: Option> = None; + let paths: Vec<(String, HashSet)> = paths.collect(); + let Some((template, _)) = paths.first() else { + bail!("pool: no group tables to pool into {out}"); + }; + let mut w = SpliceWriter::create(out, template)?; let mut rows = 0u64; - for (path, drop) in paths { - let t = TableFile::open(&path)?; - let reader = t.batches(None, BATCH_ROWS)?; - let this_schema = reader.schema(); - match &schema { - None => { - schema = Some(this_schema.clone()); - writer = Some(BatchWriter::with_row_group_rows( - out, - this_schema.clone(), - ROW_GROUP_ROWS, - )?); - } - Some(first) => { - if first.fields() != this_schema.fields() { - bail!( - "{path} has a different schema from the first pooled table; the group \ - artifacts must come from the same configuration" - ); - } - } + for (path, drop) in &paths { + if drop.is_empty() { + rows += w.append_row_groups(path, |_| true)?; + continue; + } + // Which row groups can hold a dropped id, from the statistics on candidate_id. + let t = TableFile::open(path)?; + let stats = t.row_group_stats("candidate_id")?; + let spans = SpliceWriter::row_group_spans(path)?; + if stats.len() != spans.len() { + bail!( + "{path}: {} row groups but {} statistics", + spans.len(), + stats.len() + ); } - let cid_ix = this_schema - .index_of("candidate_id") - .map_err(|_| anyhow!("{path} has no candidate_id column"))?; - let w = writer.as_mut().expect("opened above"); - for b in reader { - let b = b?; - let keep: Option> = { - let cid = b - .column(cid_ix) - .as_any() - .downcast_ref::() - .ok_or_else(|| anyhow!("{path}: candidate_id is not u32"))?; - if cid.null_count() > 0 { - bail!("{path}: candidate_id has nulls"); + let mut sorted: Vec = drop.iter().copied().collect(); + sorted.sort_unstable(); + let dirty: Vec = stats + .iter() + .map(|s| match (s.min, s.max) { + (Some(lo), Some(hi)) => { + let at = sorted.partition_point(|&d| f64::from(d) < lo); + at < sorted.len() && f64::from(sorted[at]) <= hi } - (!drop.is_empty()).then(|| cid.values().iter().map(|c| !drop.contains(c)).collect()) - }; - // The band stages already wrote library-wide ids (`Library::global_offset` is - // added on the way out of extract), so pooling appends rather than rewrites. - let batch = match keep { - Some(k) => filter_record_batch(&b, &BooleanArray::from(k))?, - None => b, - }; - if batch.num_rows() > 0 { - w.write(&batch)?; - rows += batch.num_rows() as u64; + // No statistics: assume it holds one and take the slow path. + _ => true, + }) + .collect(); + // File order is the pooled row order, so clean runs and dirty groups are appended as + // they come, not clean ones first. + let mut i = 0usize; + while i < spans.len() { + if dirty[i] { + rows += rewrite_row_group(&mut w, path, spans[i], drop)?; + i += 1; + continue; + } + let start = i; + while i < spans.len() && !dirty[i] { + i += 1; } + rows += w.append_row_groups(path, |k| k >= start && k < i)?; } } - match writer { - Some(w) => { - w.close()?; - Ok(rows) + let spliced = w.close()?; + if spliced != rows { + bail!("pool: counted {rows} rows into {out} but the file holds {spliced}"); + } + Ok(rows) +} + +/// Decode one row group, drop the losers, and splice the result back in. Used only for the +/// row groups whose candidate range holds a dropped id. +fn rewrite_row_group( + w: &mut SpliceWriter, + path: &str, + span: (usize, usize), + drop: &HashSet, +) -> Result { + let (first_row, n_rows) = span; + let t = TableFile::open_rows(path, first_row, n_rows)?; + let reader = t.batches(None, BATCH_ROWS)?; + let schema = reader.schema(); + let cid_ix = schema + .index_of("candidate_id") + .map_err(|_| anyhow!("{path} has no candidate_id column"))?; + let tmp = format!("{path}.pool-rg{first_row}.parquet"); + let mut bw = BatchWriter::with_row_group_rows(&tmp, schema.clone(), ROW_GROUP_ROWS)?; + let mut kept = 0u64; + for b in reader { + let b = b?; + let cid = b + .column(cid_ix) + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow!("{path}: candidate_id is not u32"))?; + if cid.null_count() > 0 { + bail!("{path}: candidate_id has nulls"); } - None => bail!("pool: no group tables to pool into {out}"), + let keep: Vec = cid.values().iter().map(|c| !drop.contains(c)).collect(); + let batch = filter_record_batch(&b, &BooleanArray::from(keep))?; + if batch.num_rows() > 0 { + bw.write(&batch)?; + kept += batch.num_rows() as u64; + } + } + bw.close()?; + let spliced = w.append_row_groups(&tmp, |_| true)?; + std::fs::remove_file(&tmp).ok(); + if spliced != kept { + bail!("{path}: rewrote {kept} rows but spliced {spliced}"); } + Ok(kept) } pub fn run(p: PoolParams) -> Result { @@ -229,6 +275,13 @@ mod tests { "x".into(), cids.iter().map(|c| *c as f64 + marker).collect(), ), + // A LargeList column, as the chromatogram table has: the distinction from + // `List` lives only in the file's Arrow metadata, which a spliced output has + // to carry over from the tables it was spliced from. + Col::LargeListF32( + "trace".into(), + cids.iter().map(|c| vec![*c as f32, *c as f32 + 0.5]).collect(), + ), ]; if extra { cols.push(Col::F64("prelim_score".into(), scores.to_vec())); @@ -278,6 +331,13 @@ mod tests { let t = TableFile::open(path).unwrap(); let cid = t.u32("candidate_id").unwrap(); assert_eq!(cid, vec![0, 1, 2, 3, 4, 5, 6], "{path}"); + assert!( + matches!( + t.schema.field_with_name("trace").unwrap().data_type(), + arrow::datatypes::DataType::LargeList(_) + ), + "{path}: the spliced table lost the LargeList type" + ); let x = t.f64("x").unwrap(); // Id 3 came from band 0 (x = 3.000), id 4 from band 1 (x = 4.003). assert!( @@ -287,6 +347,83 @@ mod tests { } } + /// A band whose losers sit in one row group: the groups around it are spliced as bytes + /// and that one is decoded, filtered and re-encoded. The pooled table must still be the + /// bands' rows, in band order, in row order, with the losers gone -- the splice must not + /// reorder anything, and the mixed path must not drop or duplicate a group. + #[test] + fn a_band_of_many_row_groups_keeps_its_order_when_only_one_holds_a_loser() { + use mumdia_io::table::BatchWriter; + let dir = std::env::temp_dir().join(format!("mumdia_pool_rg_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + // 500 rows in row groups of 50, so 10 groups; the losers are ids 220..230, which is + // group 4 alone. + let mk = |tag: &str, ids: std::ops::Range, score: f64| -> String { + let p = dir + .join(format!("{tag}.parquet")) + .to_str() + .unwrap() + .to_string(); + let cids: Vec = ids.collect(); + let n = cids.len(); + let schema = std::sync::Arc::new(arrow::datatypes::Schema::new(vec![ + arrow::datatypes::Field::new("candidate_id", arrow::datatypes::DataType::UInt32, false), + arrow::datatypes::Field::new("prelim_score", arrow::datatypes::DataType::Float64, false), + ])); + let mut w = BatchWriter::with_row_group_rows(&p, schema.clone(), 50).unwrap(); + for chunk in cids.chunks(25) { + let batch = arrow::record_batch::RecordBatch::try_new( + schema.clone(), + vec![ + std::sync::Arc::new(UInt32Array::from(chunk.to_vec())), + std::sync::Arc::new(arrow::array::Float64Array::from( + vec![score; chunk.len()], + )), + ], + ) + .unwrap(); + w.write(&batch).unwrap(); + } + w.close().unwrap(); + assert_eq!(TableFile::open(&p).unwrap().nrows, n); + p + }; + // Band 0 holds 0..500 and band 1 holds 220..230 with a higher score, so band 0 + // loses exactly those ten, which live in one of its ten row groups. + let b0 = BandArtifacts { + psms: mk("rg_b0_psms", 0..500, 1.0), + chromatograms: mk("rg_b0_chrom", 0..500, 1.0), + competed: mk("rg_b0_comp", 0..500, 1.0), + }; + let b1 = BandArtifacts { + psms: mk("rg_b1_psms", 220..230, 9.0), + chromatograms: mk("rg_b1_chrom", 220..230, 9.0), + competed: mk("rg_b1_comp", 220..230, 9.0), + }; + let out = |n: &str| { + dir.join(format!("rg_out_{n}.parquet")) + .to_str() + .unwrap() + .to_string() + }; + let (op, oc, ok) = (out("psms"), out("chrom"), out("comp")); + let stats = run(PoolParams { + bands: &[b0, b1], + out_psms: Some(op.as_str()), + out_chromatograms: &oc, + out_competed: &ok, + }) + .unwrap(); + assert_eq!(stats.duplicates, 10); + assert_eq!(stats.competed, 500); + let mut expect: Vec = (0..500).filter(|c| !(220..230).contains(c)).collect(); + expect.extend(220..230); + for path in [&op, &oc, &ok] { + let t = TableFile::open(path).unwrap(); + assert_eq!(t.u32("candidate_id").unwrap(), expect, "{path}"); + } + } + #[test] fn without_a_psms_path_the_other_tables_are_pooled_unchanged() { let dir = std::env::temp_dir().join(format!("mumdia_pool_nopsms_{}", std::process::id())); diff --git a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs index bf9b011b..1b5ff5fb 100644 --- a/rust/mumdia/crates/mumdia/src/stages/run_groups.rs +++ b/rust/mumdia/crates/mumdia/src/stages/run_groups.rs @@ -150,7 +150,28 @@ pub fn run(mut g: GroupRun) -> Result { seed: String, } let mut bands: Vec = Vec::new(); - let par = cfg.groups.parallel.max(1); + // Bands in flight are driven from the rayon pool, and each one's extraction blocks its + // own thread on the accumulation channel while the probing tasks run on the others. A + // band in flight therefore occupies a worker that cannot do the work it is waiting for, + // and `groups.parallel >= threads` deadlocks: every worker parks and no task is left to + // feed them. Reproduced on the fixture with `parallel = 2, --threads 2` (the process sat + // at 0.1 s of CPU indefinitely). Leave at least one worker free. + let threads = rayon::current_num_threads(); + let par = { + let want = cfg.groups.parallel.max(1); + let most = threads.saturating_sub(1).max(1); + if want > most { + tracing::warn!( + requested = want, + used = most, + threads, + "groups.parallel is at least the thread count, which would deadlock: every band in flight parks a worker on its accumulation channel. Using one fewer band than there are threads; raise --threads to run more at once" + ); + most + } else { + want + } + }; // Bands are independent, so `groups.parallel` of them are sliced and seeded at once. // Chunked rather than a free-running pool: the chunk bounds how many extraction working // sets are resident, which is the whole point of banding. Results do not depend on it, From 7f908a4f2a5356f6c063da9f1ebf4a5ce2392faf Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:44:53 +0200 Subject: [PATCH 16/35] fix(compete): refuse a null in a column the competed schema declares non-nullable The pass-through copy reads fewer columns through the typed getters than the typed path did -- `candidate_id` only for the audit, `base_peptide_id` only for the groupings that key on it -- and those getters are where nulls were refused. Without a check the copy would carry a null into the competed table and the failure would surface somewhere downstream, or nowhere. Raised by the adversarial review of the compete change. `Array::null_count` is a counter on the array, not a scan, so this is one comparison per column per batch. Co-Authored-By: Claude Opus 5 (1M context) --- rust/mumdia/crates/mumdia/src/stages/compete.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/rust/mumdia/crates/mumdia/src/stages/compete.rs b/rust/mumdia/crates/mumdia/src/stages/compete.rs index 4762d88e..b2ff65c7 100644 --- a/rust/mumdia/crates/mumdia/src/stages/compete.rs +++ b/rust/mumdia/crates/mumdia/src/stages/compete.rs @@ -460,6 +460,23 @@ fn copy_kept_rows( }; let mut arrays: Vec = Vec::with_capacity(src_idx.len()); for (si, f) in src_idx.iter().zip(out_schema.fields()) { + // A column the schema declares non-nullable must not carry nulls. The + // typed getters (`TableFile::u32`) refuse them, and reading fewer columns + // means fewer of them pass through a getter: `candidate_id` is now read + // only for the audit and `base_peptide_id` only for the groupings that key + // on it, so without this the pass-through would copy a null into the + // competed table and the failure would surface somewhere downstream, or + // not at all. `null_count` is a counter on the array, not a scan. + if let Some(i) = si { + let col = b.column(*i); + if !f.is_nullable() && col.null_count() > 0 { + return Err(anyhow!( + "column '{}' has {} null values in rows {row0}..{row1} of the features table, and the competed schema declares it non-nullable", + f.name(), + col.null_count() + )); + } + } arrays.push(match (si, &idx) { (Some(i), None) => densify(b.column(*i).clone(), f)?, (Some(i), Some(ix)) => densify(take(b.column(*i).as_ref(), ix, None)?, f)?, From 7b44cab315e8e14b77478d6743725c3c533323e8 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:50:35 +0200 Subject: [PATCH 17/35] test(io): pin what the chunked write really guarantees The adversarial review of the chunked `write_table` argued that its byte-equality claim does not hold for an entirely NULL optional column, because the definition levels are then RLE-run encoded and the writer takes its internal mini-batches in a different size, so a row-chunk boundary no longer lands on a mini-batch boundary. That is correct, and it is measured here rather than argued: 2,287,903 bytes against 2,287,919 at 196,615 rows, a 16-byte difference in page framing, with identical rows. It matters because every real run writes such a column: `apex_im` on `psms_extracted` and the three ion-mobility columns on `run_windows` have exactly one push site each and it pushes `None`. Those two artifacts therefore have different content hashes from the ones a pre-chunking build wrote. The rows, their order and their values are unchanged, which is what this layer promises; the byte change is of the same kind as the row-group caps elsewhere in this series. Two tests, neither of which existed: an all-null table compared row for row and row-group for row-group, and a table past the writer's 1,048,576-row row-group maximum, where the chunked and single-batch files come out byte-identical with the groups in the same places. The second is the property the 65,536-row chunk size was chosen for and nothing covered it: every other chunking test is inside one row group. Also: the all-null test gets its own temporary directory. The shared one is wiped with `remove_dir_all` by a test in another module while these run in parallel threads. Co-Authored-By: Claude Opus 5 (1M context) --- rust/mumdia/crates/mumdia-io/src/table.rs | 92 +++++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/rust/mumdia/crates/mumdia-io/src/table.rs b/rust/mumdia/crates/mumdia-io/src/table.rs index b16c6749..88673baf 100644 --- a/rust/mumdia/crates/mumdia-io/src/table.rs +++ b/rust/mumdia/crates/mumdia-io/src/table.rs @@ -1980,6 +1980,98 @@ mod write_chunking_tests { std::fs::remove_file(&once).ok(); } + /// An entirely NULL optional column, which is what every real run writes: `apex_im` on + /// `psms_extracted` and the three ion-mobility columns on `run_windows` have exactly one + /// push site each and it pushes `None`. + /// + /// Here the chunked write is NOT byte-identical to the single-batch write, and this test + /// pins that rather than hiding it. An all-null column's definition levels are RLE-run + /// encoded, which changes the size the writer takes its internal mini-batches in, so a + /// row-chunk boundary no longer coincides with a mini-batch boundary and one page header + /// lands elsewhere: measured on parquet 59.3.0, 2,287,903 bytes against 2,287,919 at + /// 196,615 rows, a 16-byte difference in page framing. The ROWS are identical, which is + /// the contract this layer makes. The consequence to know about is that + /// `psms_extracted.parquet` and `run_windows.parquet` have different content hashes from + /// the ones a pre-chunking build wrote, exactly as the row-group caps elsewhere do. + #[test] + fn an_entirely_null_column_keeps_its_rows_but_not_its_page_framing() { + let n = 3 * WRITE_TABLE_CHUNK_ROWS + 7; + let cols = |n: usize| -> Vec{ + vec![ + Col::U32("id".into(), (0..n as u32).collect()), + Col::F64("mz".into(), (0..n).map(|i| i as f64 * 1.5 - 3.0).collect()), + Col::OptF64("apex_im".into(), vec![None; n]), + Col::OptF32("im_lower".into(), vec![None; n]), + Col::OptI32("im_bin".into(), vec![None; n]), + Col::OptStr("note".into(), vec![None; n]), + ] + }; + // Its own directory: another test in this module wipes the shared one with + // `remove_dir_all` while these run in parallel threads. + let dir = std::env::temp_dir().join(format!("mumdia_allnull_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let p = |n: &str| dir.join(n).to_str().unwrap().to_string(); + let chunked = p("allnull_chunked.parquet"); + let once = p("allnull_once.parquet"); + assert_eq!(write_table(&chunked, cols(n)).unwrap(), n as u64); + write_one_batch(&once, cols(n)); + let (a, b) = (Table::read(&chunked).unwrap(), Table::read(&once).unwrap()); + assert_eq!(a.schema, b.schema); + assert_eq!(a.nrows, n); + assert_eq!(a.u32("id").unwrap(), b.u32("id").unwrap()); + assert_eq!(a.f64("mz").unwrap(), b.f64("mz").unwrap()); + assert_eq!(a.opt_f64("apex_im").unwrap(), b.opt_f64("apex_im").unwrap()); + assert!( + a.opt_f64("apex_im").unwrap().iter().all(|v| v.is_none()), + "the fixture's point is that the column is entirely null" + ); + let (ta, tb) = ( + TableFile::open(&chunked).unwrap(), + TableFile::open(&once).unwrap(), + ); + assert_eq!( + ta.row_group_stats("id").unwrap().len(), + tb.row_group_stats("id").unwrap().len(), + "the row groups fall in the same places whichever way the table was written" + ); + std::fs::remove_file(&chunked).ok(); + std::fs::remove_file(&once).ok(); + } + + /// Past the writer's 1,048,576-row row-group maximum the chunked write must still put + /// the group boundaries where the single-batch write puts them: the 65,536-row chunk + /// divides that maximum, so a row group closes on the same row either way. Nothing + /// tested that before -- the other chunking tests are all under one row group -- and it + /// is the property the whole chunk-size choice rests on. + #[test] + fn the_row_groups_fall_in_the_same_places_past_the_row_group_maximum() { + let n = 1_048_576 + 7; + let dir = std::env::temp_dir().join(format!("mumdia_bigrg_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let p = |x: &str| dir.join(x).to_str().unwrap().to_string(); + let cols = |n: usize| -> Vec{ + vec![ + Col::U32("id".into(), (0..n as u32).collect()), + Col::F64("mz".into(), (0..n).map(|i| i as f64 * 0.5).collect()), + ] + }; + let (chunked, once) = (p("big_chunked.parquet"), p("big_once.parquet")); + write_table(&chunked, cols(n)).unwrap(); + write_one_batch(&once, cols(n)); + let a = TableFile::open(&chunked).unwrap().row_group_stats("id").unwrap(); + let b = TableFile::open(&once).unwrap().row_group_stats("id").unwrap(); + let rows_a: Vec = a.iter().map(|g| g.rows).collect(); + let rows_b: Vec = b.iter().map(|g| g.rows).collect(); + assert_eq!(rows_a, rows_b, "row groups moved: {rows_a:?} against {rows_b:?}"); + assert_eq!(rows_a, vec![1_048_576, 7]); + assert_eq!( + std::fs::read(&chunked).unwrap(), + std::fs::read(&once).unwrap(), + "and the file is still byte-identical across a row-group boundary" + ); + let _ = std::fs::remove_dir_all(&dir); + } + /// List columns hold a variable number of leaf values per row, so a row-chunk boundary /// need not fall on a 1,024-level mini-batch boundary and the data pages could in /// principle be cut elsewhere. The contract there is the rows: same order, same values, From f9ef7645a6f14b4738acc3d63a0931b0a0a12ff3 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:51:43 +0200 Subject: [PATCH 18/35] perf(extract): hold the hit accumulator flat (CSR), not one Vec per candidate The grouped search dies at ~290 GB resident with "memory allocation of N bytes failed" while 1.7 TB of RAM is free. Measured on the live process, the limit it reaches is the per-process memory MAPPING table (vm.max_map_count = 1,048,576): at 180 GB it held 129,393 mappings, 92,030 of them 64-256 KB and 36,368 of them 256 KB-1 MB. Those sizes are one-heap-block-per-item structures, and the item here is a candidate: every candidate with evidence owned a `Vec` (24 B per hit) created on first collision and grown by doubling, in the task-local map inside `probe_range`, in the merge into the accumulator, and in the drain. A candidate with no `run_windows` row gets infinite RT bounds and so collects hits across the whole gradient, 10^4-10^5 of them, i.e. 240 KB-2.4 MB in one block, which mimalloc maps individually. Those hits now live in a CSR layout (`HitStore`: ascending `cids`, per-candidate `offs`, one flat `hits` buffer), so a band of 1.9-5.8M precursors holds three allocations per probing task instead of one per candidate with evidence. Live medium/large blocks drop from O(candidates with evidence) to O(tasks), which is 2x threads. Mechanics, all chosen so the bytes never double: - `probe_range` collects flat `(cid, hit)` pairs in probe order and groups them at the end of the task with a STABLE counting sort (`group_hits_by_candidate`). The permutation is applied in place by cycle following rather than scattered into a second buffer, and the `cid` array is reused as the slot array, so the extra is 4 B per hit plus one u32 per candidate of the sub-range, not another 24 B per hit. - A window's sub-range tasks cover disjoint ascending candidate spans, so their stores concatenate into one ascending run (`HitRun`) with no per-candidate merge at all. The consumer appends runs in window order; `gather_chunk` resolves the concatenation lazily at flush, one `CAND_CHUNK` of candidates at a time through a reused buffer. - `per_candidate` takes `&mut [Hit]` into that buffer instead of an owned `Vec`; it sorted a slice either way. - The eager (non-fragindex, two-pass, allowlist-free-of-fragindex) paths still build the whole-run `HashMap>` and are now moved into the CSR buffer a chunk at a time, freeing the map's vectors as they are consumed, so the flat copy never coexists with the whole map. Also hoists the per-peak query m/z out of the sub-range tasks: `peak.mz / mass_off.factor_at(peak.mz)` is a division, and a binary search into the masscal grid when one is present, and it was recomputed once per task per peak, up to 2x threads redundantly. It is now computed once per window into a flat buffer the window's tasks share, built only when the window is actually split. The bin computation (a `ln()` inside `probe_peak_win`) is the other half of that hoist and is NOT done here: it needs a `FragIndex` entry point taking a precomputed bin, which is outside this file. Risk to output equality: none intended, and the ordering is the whole argument. A candidate's hits must arrive in ascending scan order within a window and with windows concatenated in window order, because every float reduction downstream (rt sort, apex sum, wsum, chromatogram grids) is order-sensitive. The counting sort is stable, so within a task the order is arrival order; the sub-range spans are disjoint, so a candidate is produced by exactly one task of a window; and the gather walks runs in window order, appending each candidate's segments. PSM row order, chromatogram row order and the CAND_CHUNK partition (hence parquet row-group boundaries) are unchanged: the flush is the same candidate partition in the same sequence. Tests: `candidate_range_split_reproduces_the_unsplit_accumulation` now materialises the accumulator through the production gather and `slices_mut` and still asserts 2/8/16 threads agree hit for hit; `grouping_reproduces_per_candidate_push_order` asserts the counting sort reproduces the naive "push onto this candidate's Vec" accumulation it replaced, including untouched candidates inside the span; `gather_concatenates_a_candidate_across_runs_in_window_order`, `gather_stops_at_the_bound_and_keeps_the_rest` and `gather_chunking_partitions_the_accumulator` cover the flush. Co-Authored-By: Claude Opus 5 (1M context) --- .../crates/mumdia/src/stages/extract.rs | 660 +++++++++++++++--- 1 file changed, 578 insertions(+), 82 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/extract.rs b/rust/mumdia/crates/mumdia/src/stages/extract.rs index 75d4d1f2..e4c8a373 100644 --- a/rust/mumdia/crates/mumdia/src/stages/extract.rs +++ b/rust/mumdia/crates/mumdia/src/stages/extract.rs @@ -111,6 +111,7 @@ pub struct ExtractParams<'a> { /// One observed hit: scan RT, candidate-local fragment index, observed intensity /// and observed m/z (for mass-accuracy features). +#[derive(Clone, Copy)] #[cfg_attr(test, derive(Debug, PartialEq))] struct Hit { rt: f64, @@ -119,6 +120,287 @@ struct Hit { obs_mz: f64, } +/// Hits for many candidates in one flat buffer with per-candidate offsets (CSR). +/// +/// The accumulator used to be a `HashMap>`: one heap block per candidate +/// with evidence, created on first collision and grown by doubling. A candidate with no +/// `run_windows` row gets infinite RT bounds and therefore collects hits across the whole +/// gradient -- 10^4 to 10^5 hits, i.e. 240 KB to 2.4 MB in a single block -- and mimalloc +/// maps blocks that size individually. Measured on a live grouped search: 129,393 mappings +/// at 180 GB resident, 92,030 of them 64-256 KB and 36,368 of them 256 KB-1 MB, with the +/// process dying at ~290 GB on the per-process mapping limit (`vm.max_map_count` = +/// 1,048,576) while 1.7 TB of RAM was free. The CSR layout holds the same hits, for the +/// same candidates, in the same per-candidate order, in three allocations per store. +struct HitStore { + /// Candidate ids owning a slice of `hits`, strictly ascending. + cids: Vec, + /// `offs[i]..offs[i + 1]` is `cids[i]`'s slice of `hits`; length `cids.len() + 1` with + /// `offs[0] == 0`. `usize` rather than `u32`: one run has been measured at 1.6 billion + /// hits, past what a 32-bit offset can address. + offs: Vec, + hits: Vec, +} + +impl Default for HitStore { + fn default() -> Self { + HitStore { + cids: Vec::new(), + offs: vec![0], + hits: Vec::new(), + } + } +} + +impl HitStore { + fn len(&self) -> usize { + self.cids.len() + } + + fn is_empty(&self) -> bool { + self.cids.is_empty() + } + + fn clear(&mut self) { + self.cids.clear(); + self.offs.clear(); + self.offs.push(0); + self.hits.clear(); + } + + /// Candidate `i`'s hits. + fn slice(&self, i: usize) -> &[Hit] { + &self.hits[self.offs[i]..self.offs[i + 1]] + } + + /// Append `hits` to candidate `cid`, which must be >= the last id pushed. When it is + /// equal the hits extend that candidate's existing segment, which is how one + /// candidate's hits from several windows end up contiguous and in window order. + fn push_segment(&mut self, cid: u32, hits: &[Hit]) { + if hits.is_empty() { + return; + } + debug_assert!(self.cids.last().map(|&l| l <= cid).unwrap_or(true)); + if self.cids.last() != Some(&cid) { + self.cids.push(cid); + self.offs.push(0); + } + self.hits.extend_from_slice(hits); + *self.offs.last_mut().expect("offs is never empty") = self.hits.len(); + } + + /// One disjoint `&mut [Hit]` per candidate, ascending by id: what the per-candidate + /// pass consumes in place of the `Vec` it used to be handed. The sort it runs is + /// a slice sort either way, so the hits it sees and the order it sees them in are + /// exactly those of the owned vector. + fn slices_mut(&mut self) -> Vec<(u32, &mut [Hit])> { + let HitStore { cids, offs, hits } = self; + let mut out = Vec::with_capacity(cids.len()); + let mut rest: &mut [Hit] = hits.as_mut_slice(); + let mut base = 0usize; + for (i, &cid) in cids.iter().enumerate() { + let end = offs[i + 1]; + let (head, tail) = rest.split_at_mut(end - base); + out.push((cid, head)); + rest = tail; + base = end; + } + out + } +} + +/// Group `(cid, hit)` pairs by candidate with a STABLE counting sort applied in place, +/// and return the CSR spine (`cids`, `offs`) for the grouped `hits`. +/// +/// Stability is the whole contract: each candidate's hits must come out in exactly the +/// order the probe produced them (ascending scan order, and within a scan the order the +/// posting list emitted), because that is the order the per-candidate `Vec` held them +/// in and every float reduction downstream is order-sensitive. +/// +/// The permutation is applied by cycle following rather than scattering into a second +/// buffer: a second buffer would double the stage's largest structure. `cid` is consumed +/// as scratch (it becomes the destination slot of each hit). Every `cid` must lie in +/// `[lo, hi)`. +fn group_hits_by_candidate( + lo: u32, + hi: u32, + cid: &mut [u32], + hits: &mut [Hit], +) -> (Vec, Vec) { + debug_assert_eq!(cid.len(), hits.len()); + let n = hits.len(); + if n == 0 { + return (Vec::new(), vec![0]); + } + assert!( + n <= u32::MAX as usize, + "a single probing task produced {n} hits, past the u32 slot index" + ); + let span = (hi - lo) as usize; + // `cum[c]` is first candidate `lo + c`'s hit count, then (after the prefix sum) its + // first slot, and finally -- the scatter having advanced it once per hit -- one past + // its last, i.e. the cumulative count through `c`. One span-sized u32 array, 4 bytes + // per candidate of the sub-range, serves all three roles. + let mut cum: Vec = vec![0; span]; + for &c in cid.iter() { + debug_assert!(c >= lo && c < hi); + cum[(c - lo) as usize] += 1; + } + let mut running = 0u32; + for v in cum.iter_mut() { + let k = *v; + *v = running; + running += k; + } + debug_assert_eq!(running as usize, n); + // Walking in order is what makes the sort stable: `cid[i]` becomes the slot hit `i` + // belongs in, and equal candidates take consecutive slots in arrival order. + for c in cid.iter_mut() { + let b = (*c - lo) as usize; + *c = cum[b]; + cum[b] += 1; + } + for i in 0..n { + while cid[i] as usize != i { + let j = cid[i] as usize; + hits.swap(i, j); + cid.swap(i, j); + } + } + // `cum[c]` now holds the cumulative count through candidate `c`, so a candidate is + // occupied exactly where that count grew. + let mut cids: Vec = Vec::new(); + let mut offs: Vec = vec![0]; + let mut prev = 0u32; + for (c, &end) in cum.iter().enumerate() { + if end > prev { + cids.push(lo + c as u32); + offs.push(end as usize); + } + prev = end; + } + (cids, offs) +} + +/// One probed window's stores read as a single ascending candidate sequence. The +/// sub-range tasks of a window cover disjoint, ascending candidate spans, so their stores +/// concatenate without a merge. +struct HitRun { + parts: Vec, + pi: usize, + ci: usize, +} + +impl HitRun { + fn new(parts: Vec) -> HitRun { + let mut r = HitRun { + parts, + pi: 0, + ci: 0, + }; + r.settle(); + r + } + + fn settle(&mut self) { + while self.pi < self.parts.len() && self.ci >= self.parts[self.pi].len() { + self.pi += 1; + self.ci = 0; + } + } + + fn peek(&self) -> Option { + self.parts.get(self.pi).map(|p| p.cids[self.ci]) + } + + fn current(&self) -> &[Hit] { + self.parts[self.pi].slice(self.ci) + } + + fn advance(&mut self) { + self.ci += 1; + self.settle(); + } + + /// Hits not yet gathered out of this run. + fn hits_left(&self) -> usize { + match self.parts.get(self.pi) { + None => 0, + Some(p) => { + (p.hits.len() - p.offs[self.ci]) + + self.parts[self.pi + 1..] + .iter() + .map(|q| q.hits.len()) + .sum::() + } + } + } + + /// Free the stores already fully gathered. + fn release_consumed(&mut self) { + if self.pi > 0 { + self.parts.drain(..self.pi); + self.pi = 0; + } + } +} + +/// Gather up to `max_cands` candidates below `bound` out of `runs`, ascending by candidate +/// id, appending each candidate's hits run by run. Runs are held in window order, so a +/// candidate seen in several windows keeps exactly the hit sequence the serial path +/// produced: ascending scan order within a window, windows concatenated. +/// +/// `bound` is the first candidate a later window can still add hits to, and `u32::MAX` +/// when nothing is pending, which no real candidate id reaches (ids are dense `0..ncand`). +fn gather_chunk(runs: &mut [HitRun], bound: u32, max_cands: usize, out: &mut HitStore) { + out.clear(); + while out.len() < max_cands { + let mut min_cid: Option = None; + for r in runs.iter() { + if let Some(c) = r.peek() { + if min_cid.map(|m| c < m).unwrap_or(true) { + min_cid = Some(c); + } + } + } + let cid = match min_cid { + Some(c) if c < bound => c, + _ => break, + }; + for r in runs.iter_mut() { + if r.peek() == Some(cid) { + out.push_segment(cid, r.current()); + r.advance(); + } + } + } +} + +/// The streamed accumulator: one run per probed window, oldest first. +#[derive(Default)] +struct HitAcc { + runs: Vec, +} + +impl HitAcc { + fn n_hits(&self) -> usize { + self.runs.iter().map(|r| r.hits_left()).sum() + } + + /// Free what has been gathered and fold whatever is left into a single run, so the + /// number of live runs stays bounded however much the windows overlap. + fn compact(&mut self) { + for r in self.runs.iter_mut() { + r.release_consumed(); + } + self.runs.retain(|r| r.peek().is_some()); + if self.runs.len() > 1 { + let mut merged = HitStore::default(); + gather_chunk(&mut self.runs, u32::MAX, usize::MAX, &mut merged); + self.runs = vec![HitRun::new(vec![merged])]; + } + } +} + type ChromOutputRow = (u32, String, f64, f64, f32, Vec, Vec); /// Candidates per parallel chunk of the per-candidate pass: large enough to keep every @@ -753,15 +1035,15 @@ fn accumulate_groups( mass_off: &MassOffset, cfg: &ExtractConfig, restrict: Option<&std::collections::HashSet>, - acc: &mut HashMap>, + acc: &mut HitAcc, ) { - // Partials are merged into `acc` in window order as they complete, not collected first - // and merged after: with a collect, every window's hits existed twice at the merge (in - // the partials and, growing, in the merged map). The channel hands each finished - // partial to this thread; the reorder buffer holds the ones that finished ahead of an - // earlier window, so the merge order, and with it every candidate's hit sequence, is - // exactly the serial one. - let (tx, rx) = std::sync::mpsc::channel::<(usize, Vec<(u32, Vec)>)>(); + // Partials are appended to `acc` in window order as they complete, not collected first + // and merged after. The channel hands each finished partial to this thread; the + // reorder buffer holds the ones that finished ahead of an earlier window, so the run + // order, and with it every candidate's hit sequence, is exactly the serial one. No + // per-candidate merge happens here any more: a window's partial is a CSR store per + // sub-range task, and the concatenation is resolved lazily by `gather_chunk` at flush. + let (tx, rx) = std::sync::mpsc::channel::<(usize, Vec)>(); // A batch of `step` windows on `t` threads leaves `t - step` threads idle whenever the // batch is smaller than the pool, which is what a `windows_in_flight` set low for memory // does: measured at 4 in flight on 24 threads, 3.2-3.7 cores busy for the whole stage. @@ -775,21 +1057,54 @@ fn accumulate_groups( // divides, while the peak loop itself is repeated per task. let threads = rayon::current_num_threads().max(1); let tasks_per_window = (threads * 2).div_ceil(groups.len().max(1)).max(1); - let probe = |_gi: usize, g: &WinGroup| -> Vec<(u32, Vec)> { + let probe = |g: &WinGroup| -> Vec { let ids = &g.scans; - let probe_range = |lo: u32, hi: u32| -> Vec<(u32, Vec)> { - let mut local: HashMap> = HashMap::new(); + let n_cand = (g.hi_cid - g.lo_cid) as usize; + let tasks = tasks_per_window + .min(n_cand.div_ceil(MIN_CANDIDATES_PER_TASK)) + .max(1); + // Offset-corrected query m/z per peak, computed ONCE for the window instead of + // once per sub-range task. `MassOffset::factor_at` is a binary search into the + // m/z-dependent calibration grid when a masscal grid is present and a division + // either way, and both were repeated `tasks` times -- up to 2x threads -- for + // every peak of every scan of the window. Flat with one offset per scan, so it is + // two allocations rather than one per scan, and built only when the window is + // actually split (unsplit, there is nothing to share it with). + let qmz: Option<(Vec, Vec)> = (tasks > 1).then(|| { + let mut q: Vec = + Vec::with_capacity(ids.iter().map(|&si| scans[si].peaks.len()).sum()); + let mut off: Vec = Vec::with_capacity(ids.len() + 1); + off.push(0); + for &si in ids { + for peak in &scans[si].peaks { + q.push(peak.mz / mass_off.factor_at(peak.mz)); + } + off.push(q.len()); + } + (q, off) + }); + let probe_range = |lo: u32, hi: u32| -> HitStore { + // Flat `(cid, hit)` pairs in probe order, grouped by candidate at the end of + // the task with a stable counting sort. The task-local `HashMap>` this replaces was one of the two populations of medium heap + // blocks that exhausted the mapping table. + let mut flat_cid: Vec = Vec::new(); + let mut flat_hit: Vec = Vec::new(); let mut claimants: Vec<(u32, u16, f32)> = Vec::new(); // `(lo, hi)` is fixed for this whole sub-range and every scan of the window // reprobes the same bins, so cache each bin's narrowed posting range once // instead of binary-searching it per peak. let mut nw = idx.window_narrow(lo, hi); - for &si in ids { + for (k, &si) in ids.iter().enumerate() { let scan = &scans[si]; let rt = scan.rt_seconds; - for peak in &scan.peaks { + let qs: Option<&[f64]> = qmz.as_ref().map(|(q, off)| &q[off[k]..off[k + 1]]); + for (pk, peak) in scan.peaks.iter().enumerate() { let inten = peak.intensity; - let q_mz = peak.mz / mass_off.factor_at(peak.mz); + let q_mz = match qs { + Some(s) => s[pk], + None => peak.mz / mass_off.factor_at(peak.mz), + }; let obs_mz = peak.mz; claimants.clear(); idx.probe_peak_win(&mut nw, q_mz, |cid, _pmz, pint, frag| { @@ -821,7 +1136,8 @@ fn accumulate_groups( } } let (cid, frag, _) = claimants[best]; - local.entry(cid).or_default().push(Hit { + flat_cid.push(cid); + flat_hit.push(Hit { rt, frag, inten, @@ -836,7 +1152,8 @@ fn accumulate_groups( } else { inten / claimants.len() as f32 }; - local.entry(cid).or_default().push(Hit { + flat_cid.push(cid); + flat_hit.push(Hit { rt, frag, inten: share, @@ -846,7 +1163,8 @@ fn accumulate_groups( } _ => { for &(cid, frag, _) in &claimants { - local.entry(cid).or_default().push(Hit { + flat_cid.push(cid); + flat_hit.push(Hit { rt, frag, inten, @@ -857,14 +1175,16 @@ fn accumulate_groups( } } } - local.into_iter().collect() + let (cids, offs) = group_hits_by_candidate(lo, hi, &mut flat_cid, &mut flat_hit); + drop(flat_cid); + HitStore { + cids, + offs, + hits: flat_hit, + } }; - let n_cand = (g.hi_cid - g.lo_cid) as usize; - let tasks = tasks_per_window - .min(n_cand.div_ceil(MIN_CANDIDATES_PER_TASK)) - .max(1); if tasks <= 1 { - return probe_range(g.lo_cid, g.hi_cid); + return vec![probe_range(g.lo_cid, g.hi_cid)]; } let span = n_cand.div_ceil(tasks) as u32; let ranges: Vec<(u32, u32)> = (0..tasks as u32) @@ -876,12 +1196,11 @@ fn accumulate_groups( }) .filter(|&(lo, hi)| hi > lo) .collect(); + // Ascending, disjoint candidate spans, so the stores concatenate into one + // ascending run without a per-candidate merge. ranges .par_iter() .map(|&(lo, hi)| probe_range(lo, hi)) - .collect::>() - .into_iter() - .flatten() .collect() }; let n = groups.len(); @@ -895,26 +1214,21 @@ fn accumulate_groups( .par_iter() .enumerate() .for_each_with(tx, |tx, (gi, g)| { - let part = probe(gi, g); + let part = probe(g); // A closed receiver only happens if the consumer panicked; nothing to do. let _ = tx.send((gi, part)); }); }); let mut next = 0usize; - let mut pending: std::collections::BTreeMap)>> = + let mut pending: std::collections::BTreeMap> = std::collections::BTreeMap::new(); - let merge = |part: Vec<(u32, Vec)>, acc: &mut HashMap>| { - for (cid, hits) in part { - acc.entry(cid).or_default().extend(hits); - } - }; for _ in 0..n { let (gi, part) = rx.recv().expect("every window sends exactly one partial"); if gi == next { - merge(part, acc); + acc.runs.push(HitRun::new(part)); next += 1; while let Some(p) = pending.remove(&next) { - merge(p, acc); + acc.runs.push(HitRun::new(p)); next += 1; } } else { @@ -946,21 +1260,6 @@ fn groups_in_flight(cfg: &ExtractConfig) -> usize { .max(1) } -/// Take every candidate below `bound` out of the accumulator, ascending. Windows are -/// processed in ascending m/z and precursors are sorted by m/z, so once the next window's -/// candidate range starts at `bound`, no later window can add a hit below it: those -/// candidates are complete and can be scored and written. -fn drain_below(acc: &mut HashMap>, bound: u32) -> Vec<(u32, Vec)> { - let mut ids: Vec = acc.keys().copied().filter(|&c| c < bound).collect(); - ids.sort_unstable(); - ids.into_iter() - .map(|cid| { - let hits = acc.remove(&cid).expect("id came from the map"); - (cid, hits) - }) - .collect() -} - /// Parallel two-pass co-elution peak-claim. Each isolation-window group is /// processed independently (a candidate's precursor m/z places it in one window, /// and a peak's claimants come only from that window via `candidate_range`), so @@ -1918,23 +2217,20 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // Deterministic output order (a HashMap's iteration order is randomized, // and downstream floating-point sums in the rescorer are order-sensitive). + // Empty on the streamed path, where the driver gathers each flush's candidates out of + // the CSR accumulator in ascending order instead. let mut cand_ids: Vec = acc.keys().cloned().collect(); cand_ids.sort_unstable(); - // Empty on the streamed path; the driver builds one of these per flush instead. - - // Drain the accumulator into (cid, hits) in the deterministic sorted order, - // then process candidates in parallel. Each candidate's work depends only on - // its own hits plus read-only library/window/MS1 data, and every float - // reduction (apex sum, wsum, chromatogram grids) is self-contained, so - // collecting the results in cand_ids order and appending them serially below - // yields output (PSM rows and chromatogram rows) byte-identical to the serial - // loop. The optional Pearson gate now allocates its two scratch vectors per - // candidate (was a hoisted reused buffer) because buffers cannot be shared - // across parallel candidates. - let mut cand_hits: Vec<(u32, Vec)> = cand_ids - .iter() - .map(|&cid| (cid, acc.remove(&cid).unwrap())) - .collect(); + + // The eager accumulator is drained into CSR stores of `CAND_CHUNK` candidates at a + // time, in the deterministic sorted order, and candidates are then processed in + // parallel. Each candidate's work depends only on its own hits plus read-only + // library/window/MS1 data, and every float reduction (apex sum, wsum, chromatogram + // grids) is self-contained, so collecting the results in cand_ids order and appending + // them serially below yields output (PSM rows and chromatogram rows) byte-identical to + // the serial loop. The optional Pearson gate now allocates its two scratch vectors per + // candidate (was a hoisted reused buffer) because buffers cannot be shared across + // parallel candidates. struct CandOut { cid: u32, @@ -2000,7 +2296,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { HashMap::new() }; - let per_candidate = |cid: u32, mut hits: Vec| -> Vec { + let per_candidate = |cid: u32, hits: &mut [Hit]| -> Vec { // distinct matched fragments (tier b) let mut distinct: Vec = hits.iter().map(|h| h.frag).collect(); distinct.sort_unstable(); @@ -2014,7 +2310,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // scan groups: Vec<(rt, BTreeMap)>. A BTreeMap keeps the // per-scan fragment order fixed so the f32 apex sum is deterministic. let mut groups: Vec<(f64, BTreeMap)> = Vec::new(); - for h in &hits { + for h in hits.iter() { match groups.last_mut() { Some((rt, map)) if (*rt - h.rt).abs() < 1e-9 => { let e = map.entry(h.frag).or_insert(0.0); @@ -2349,7 +2645,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // Per-fragment intensity-weighted observed m/z (for mass accuracy). let mut wsum: HashMap = HashMap::new(); // frag -> (sum w*mz, sum w) - for h in &hits { + for h in hits.iter() { let e = wsum.entry(h.frag).or_insert((0.0, 0.0)); e.0 += h.obs_mz * h.inten as f64; e.1 += h.inten as f64; @@ -2644,11 +2940,12 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // window batch on the streamed path and once on the eager paths, so the code // that produces a row is the same either way. Returns false when the writer has // gone away; its error surfaces at the join below. - let mut emit_batch = |cand_hits: &mut Vec<(u32, Vec)>| -> bool { + let mut emit_batch = |store: &mut HitStore| -> bool { + let mut cand_hits = store.slices_mut(); for chunk in cand_hits.chunks_mut(CAND_CHUNK) { let outs: Vec> = chunk .par_iter_mut() - .map(|(cid, hits)| per_candidate(*cid, std::mem::take(hits))) + .map(|(cid, hits)| per_candidate(*cid, hits)) .collect(); let mut ch = ChromChunk::default(); for r in outs.into_iter().flatten() { @@ -2731,8 +3028,11 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // Streamed path: probe a batch of windows, then flush everything the next batch // cannot touch. `acc_stream` therefore holds the hits of the windows in flight - // rather than the hits of the run. - let mut acc_stream: HashMap> = HashMap::new(); + // rather than the hits of the run. `chunk` is the reusable gather buffer the + // candidates are flushed through; it grows to one chunk's hits and is then reused, + // so the flat accumulator is never copied whole. + let mut acc_stream = HitAcc::default(); + let mut chunk = HitStore::default(); let mut open_hits_max = 0usize; if let Some(groups) = stream_groups.as_ref() { let step = groups_in_flight(p.cfg); @@ -2752,12 +3052,25 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { &mut acc_stream, ); gi = upto; - open_hits_max = open_hits_max.max(acc_stream.values().map(|v| v.len()).sum()); - // Everything below the next window's first candidate is final. + open_hits_max = open_hits_max.max(acc_stream.n_hits()); + // Everything below the next window's first candidate is final: windows are + // processed in ascending m/z and precursors are sorted by m/z, so no later + // window can add a hit below `bound`. let bound = groups.get(gi).map(|g| g.lo_cid).unwrap_or(u32::MAX); - let mut flushed = drain_below(&mut acc_stream, bound); - n_materialized += flushed.len() as u64; - if !emit_batch(&mut flushed) { + let mut stopped = false; + loop { + gather_chunk(&mut acc_stream.runs, bound, CAND_CHUNK, &mut chunk); + if chunk.is_empty() { + break; + } + n_materialized += chunk.len() as u64; + if !emit_batch(&mut chunk) { + stopped = true; + break; + } + } + acc_stream.compact(); + if stopped { break; } } @@ -2775,7 +3088,23 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { "extract: candidates with evidence" ); } else { - emit_batch(&mut cand_hits); + // Eager paths: move each chunk of candidates out of the whole-run + // `HashMap>` into the CSR buffer in ascending id order, freeing + // the map's per-candidate vectors as they are consumed, so the flat copy never + // coexists with the whole map. + let mut i = 0usize; + while i < cand_ids.len() { + let end = (i + CAND_CHUNK).min(cand_ids.len()); + chunk.clear(); + for &cid in &cand_ids[i..end] { + let hits = acc.remove(&cid).expect("id came from the map"); + chunk.push_segment(cid, &hits); + } + i = end; + if !emit_batch(&mut chunk) { + break; + } + } } // A final empty chunk fixes the schema when no candidate was accepted at all. let _ = tx.send(ChromChunk::default().cols()); @@ -3213,12 +3542,12 @@ mod accumulate_tests { grid_ppm: Vec::new(), }; let cfg = ExtractConfig::default(); - let run = |threads: usize| -> HashMap> { + let run = |threads: usize| -> HitAcc { let pool = rayon::ThreadPoolBuilder::new() .num_threads(threads) .build() .expect("thread pool"); - let mut acc: HashMap> = HashMap::new(); + let mut acc = HitAcc::default(); pool.install(|| { accumulate_groups( &idx, &groups, &sc, &rt_lo, &rt_hi, &mass_off, &cfg, None, &mut acc, @@ -3227,14 +3556,14 @@ mod accumulate_tests { acc }; // Two threads: 2 * 2 / 4 windows = one sub-range per window, the unsplit path. - let unsplit = run(2); + let unsplit = materialize(run(2)); assert!(!unsplit.is_empty(), "the fixture must produce hits"); assert!( unsplit.values().any(|v| v.len() > 1), "candidates must collect several hits, or hit order proves nothing" ); for threads in [8, 16] { - let split = run(threads); + let split = materialize(run(threads)); assert_eq!( split.len(), unsplit.len(), @@ -3249,4 +3578,171 @@ mod accumulate_tests { } } } + + /// Flush the whole accumulator through the chunked gather the driver uses, and read the + /// CSR back as the `HashMap>` the accumulator used to be. This is the + /// assertion that the flat layout hands the per-candidate pass the same sequences: the + /// gather, the CSR spine and `slices_mut` all sit between `accumulate_groups` and + /// `per_candidate` in production. + fn materialize(mut acc: HitAcc) -> BTreeMap> { + let mut out: BTreeMap> = BTreeMap::new(); + let mut chunk = HitStore::default(); + let mut last: Option = None; + loop { + // A small chunk on purpose, so several gathers run and their boundaries are + // exercised against a fixture with far more candidates than that. + gather_chunk(&mut acc.runs, u32::MAX, 512, &mut chunk); + if chunk.is_empty() { + break; + } + assert_eq!( + chunk.offs.len(), + chunk.cids.len() + 1, + "the CSR spine must have one offset per candidate plus the terminator" + ); + assert_eq!( + *chunk.offs.last().expect("offs is never empty"), + chunk.hits.len(), + "the last offset must be the hit count" + ); + for (cid, hits) in chunk.slices_mut() { + assert!( + last.map(|l| l < cid).unwrap_or(true), + "candidates must be gathered strictly ascending" + ); + last = Some(cid); + assert!(!hits.is_empty(), "an empty segment must not be emitted"); + out.insert(cid, hits.to_vec()); + } + acc.compact(); + } + out + } + + /// The counting sort is the step that replaces "push onto this candidate's Vec", so it + /// must reproduce that exactly: same candidates, and each candidate's hits in arrival + /// order. Built here against the naive per-candidate accumulation it replaced. + #[test] + fn grouping_reproduces_per_candidate_push_order() { + let (lo, hi) = (10u32, 23u32); + let mut cid: Vec = Vec::new(); + let mut hits: Vec = Vec::new(); + let mut reference: BTreeMap> = BTreeMap::new(); + // Interleaved candidates, several repeats, and some candidates of the span never + // touched (11, 12 and 22 stay empty), which is what the occupancy scan must skip. + for i in 0..200u32 { + let c = lo + ((i * 7) % 11) + if i % 5 == 0 { 2 } else { 0 }; + let c = c.min(hi - 1); + let h = Hit { + rt: 100.0 + i as f64, + frag: (i % 6) as u16, + inten: 1.0 + i as f32, + obs_mz: 300.0 + i as f64 * 0.01, + }; + cid.push(c); + hits.push(h); + reference.entry(c).or_default().push(h); + } + let (cids, offs) = group_hits_by_candidate(lo, hi, &mut cid, &mut hits); + let store = HitStore { cids, offs, hits }; + assert_eq!( + store.len(), + reference.len(), + "every candidate with a hit must own a segment, and no other" + ); + for (i, (&rc, rhits)) in reference.iter().enumerate() { + assert_eq!(store.cids[i], rc, "segment {i} candidate"); + assert_eq!(store.slice(i), rhits.as_slice(), "candidate {rc} hit order"); + } + } + + /// A candidate seen in several windows keeps window order: the windows' runs are + /// appended to the accumulator in window order, and the gather concatenates a + /// candidate's segments run by run. + #[test] + fn gather_concatenates_a_candidate_across_runs_in_window_order() { + let h = |rt: f64| Hit { + rt, + frag: 0, + inten: 1.0, + obs_mz: 300.0, + }; + let mut a = HitStore::default(); + a.push_segment(4, &[h(1.0), h(2.0)]); + a.push_segment(9, &[h(3.0)]); + let mut b = HitStore::default(); + b.push_segment(4, &[h(4.0)]); + b.push_segment(7, &[h(5.0)]); + let mut acc = HitAcc { + runs: vec![HitRun::new(vec![a]), HitRun::new(vec![b])], + }; + let mut out = HitStore::default(); + gather_chunk(&mut acc.runs, u32::MAX, usize::MAX, &mut out); + assert_eq!(out.cids, vec![4, 7, 9]); + assert_eq!( + out.slice(0).iter().map(|x| x.rt).collect::>(), + vec![1.0, 2.0, 4.0], + "the earlier window's hits must come first" + ); + assert_eq!(out.slice(1).iter().map(|x| x.rt).collect::>(), [5.0]); + assert_eq!(out.slice(2).iter().map(|x| x.rt).collect::>(), [3.0]); + } + + /// `bound` is the only thing holding a candidate back, and what it holds back must + /// still be there, in order, for the next flush. + #[test] + fn gather_stops_at_the_bound_and_keeps_the_rest() { + let h = |rt: f64| Hit { + rt, + frag: 0, + inten: 1.0, + obs_mz: 300.0, + }; + let mut a = HitStore::default(); + for c in [2u32, 5, 8, 11] { + a.push_segment(c, &[h(c as f64)]); + } + let mut acc = HitAcc { + runs: vec![HitRun::new(vec![a])], + }; + let mut out = HitStore::default(); + gather_chunk(&mut acc.runs, 8, usize::MAX, &mut out); + assert_eq!(out.cids, vec![2, 5]); + assert_eq!(acc.n_hits(), 2, "8 and 11 stay open"); + acc.compact(); + gather_chunk(&mut acc.runs, u32::MAX, usize::MAX, &mut out); + assert_eq!(out.cids, vec![8, 11]); + assert_eq!(acc.n_hits(), 0); + } + + /// The chunk cap is a flush-size limit, not a filter: chunking must not lose or + /// reorder a candidate. + #[test] + fn gather_chunking_partitions_the_accumulator() { + let h = |rt: f64| Hit { + rt, + frag: 0, + inten: 1.0, + obs_mz: 300.0, + }; + let mut a = HitStore::default(); + for c in 0..50u32 { + a.push_segment(c, &[h(c as f64), h(c as f64 + 0.5)]); + } + let mut acc = HitAcc { + runs: vec![HitRun::new(vec![a])], + }; + let mut seen: Vec = Vec::new(); + let mut out = HitStore::default(); + loop { + gather_chunk(&mut acc.runs, u32::MAX, 7, &mut out); + if out.is_empty() { + break; + } + assert!(out.len() <= 7, "the chunk cap must bind"); + seen.extend_from_slice(&out.cids); + acc.compact(); + } + assert_eq!(seen, (0..50u32).collect::>()); + } } From 66928d770490742e96546c290ee6142d6531341f Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:54:50 +0200 Subject: [PATCH 19/35] fix(index): make the rejected-label path explicit, and test the batch boundary From the adversarial review of the library change. Reading the label column in 65,536-row batches means the loop must push exactly one value per row; the `other` arm called `fdr::validate_labels`, which refuses anything but target/decoy, and pushed nothing. That is correct only because that rule is absolute today: if it ever widened, the column would come out shorter than the table and every later row would read its neighbour's label. The arm now ends in an explicit error naming the row and the file. The batch boundary itself had no test: every other fixture in index.rs is six rows. The new one is 65,541, and covers the label values either side of the boundary and in the short last batch, the `ncand` cap, a wrong candidate_id in the second batch (which must be named by its absolute row, not its row within the batch), and a NULL label there. Co-Authored-By: Claude Opus 5 (1M context) --- rust/mumdia/crates/mumdia/src/index.rs | 92 +++++++++++++++++++++++++- 1 file changed, 91 insertions(+), 1 deletion(-) diff --git a/rust/mumdia/crates/mumdia/src/index.rs b/rust/mumdia/crates/mumdia/src/index.rs index 91296210..0742cc53 100644 --- a/rust/mumdia/crates/mumdia/src/index.rs +++ b/rust/mumdia/crates/mumdia/src/index.rs @@ -132,7 +132,18 @@ fn read_is_decoy(pt: &TableFile, path: &str) -> Result> { match a.value(k) { "decoy" => out.push(true), "target" => out.push(false), - other => crate::fdr::validate_labels(&[other.to_string()])?, + other => { + // `validate_labels` refuses anything but target/decoy, so this arm + // returns. Written as an explicit error rather than a fall-through: + // if that rule ever widens, not pushing here would leave the column + // shorter than the table and every later row would read its + // neighbour's label. + crate::fdr::validate_labels(&[other.to_string()])?; + anyhow::bail!( + "library precursor row {} of {path} has label '{other}', which is neither 'target' nor 'decoy'", + out.len() + ); + } } } } @@ -920,6 +931,85 @@ mod tests { use super::*; use mumdia_io::table::{write_table, Col, TableWriter}; + /// The label and candidate_id passes read in 65,536-row batches, and their row + /// accounting across a batch boundary -- the absolute row in an error, and the + /// `ncand` cap on the last batch -- is the only genuinely new logic in them. Nothing + /// covered it: every other fixture in this file is six rows. + #[test] + fn labels_and_ids_are_read_across_a_batch_boundary() { + let dir = std::env::temp_dir().join(format!("mumdia_idx_batch_{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let n = PREC_BATCH_ROWS + 5; + let path = |t: &str| dir.join(t).to_str().unwrap().to_string(); + + let good = path("good.parquet"); + write_table( + &good, + vec![ + Col::U32("candidate_id".into(), (0..n as u32).collect()), + Col::Str( + "label".into(), + (0..n) + .map(|i| if i % 2 == 0 { "target" } else { "decoy" }.to_string()) + .collect(), + ), + ], + ) + .unwrap(); + let t = TableFile::open(&good).unwrap(); + let decoy = read_is_decoy(&t, &good).unwrap(); + assert_eq!(decoy.len(), n, "one value per row across both batches"); + assert!(!decoy[0] && decoy[1]); + // The rows either side of the boundary, and the last row of the short batch. + assert_eq!(decoy[PREC_BATCH_ROWS - 1], (PREC_BATCH_ROWS - 1) % 2 == 1); + assert_eq!(decoy[PREC_BATCH_ROWS], PREC_BATCH_ROWS % 2 == 1); + assert_eq!(decoy[n - 1], (n - 1) % 2 == 1); + check_candidate_ids(&t, &good, 0, n).unwrap(); + // A band: the ids start at an offset and only `ncand` of them are checked. + check_candidate_ids(&t, &good, 0, PREC_BATCH_ROWS).unwrap(); + + // A wrong id in the SECOND batch must be named by its absolute row. + let bad = path("bad_id.parquet"); + let mut ids: Vec = (0..n as u32).collect(); + ids[PREC_BATCH_ROWS + 2] = 7; + write_table( + &bad, + vec![ + Col::U32("candidate_id".into(), ids), + Col::Str("label".into(), vec!["target".to_string(); n]), + ], + ) + .unwrap(); + let t = TableFile::open(&bad).unwrap(); + let err = check_candidate_ids(&t, &bad, 0, n).unwrap_err().to_string(); + assert!( + err.contains(&format!("row {}", PREC_BATCH_ROWS + 2)) && err.contains("candidate_id 7"), + "{err}" + ); + + // A NULL label in the second batch, likewise. + let nulls = path("null_label.parquet"); + let mut labels: Vec> = vec![Some("target".to_string()); n]; + labels[PREC_BATCH_ROWS + 1] = None; + write_table( + &nulls, + vec![ + Col::U32("candidate_id".into(), (0..n as u32).collect()), + Col::OptStr("label".into(), labels), + ], + ) + .unwrap(); + let t = TableFile::open(&nulls).unwrap(); + // The row is in the error's source and the path in its context, so the whole + // chain has to be formatted (`{:#}`), as mumdia-io's own test of this does. + let err = format!("{:#}", read_is_decoy(&t, &nulls).unwrap_err()); + assert!( + err.contains(&(PREC_BATCH_ROWS + 1).to_string()), + "the null must be named by its absolute row: {err}" + ); + let _ = std::fs::remove_dir_all(&dir); + } + /// Six candidates at m/z 400, 450, 500, 520, 600, 650 with two fragments each, written /// in row groups of two rows so a range crosses row-group boundaries on both tables. /// `frag_order` is the candidate order of the fragment table: `[0, 1, 2, 3, 4, 5]` is From 1d97500012e1203b8c4e7c8137f8c3d6bf6550d5 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 15:58:55 +0200 Subject: [PATCH 20/35] feat(pool): pool a grouped run's bands from the command line `run` pools at the end of a grouped search. Standalone the same stage is for the case where the search finished and the run did not: the band directories hold every artifact, and pooling them is now a byte copy of their parquet row groups, so a killed run costs a pool rather than a re-search. It also makes the pool testable at real scale without a search, which is how the splice was measured. `--psms` opts into pooling psms_extracted, which only the candidate audit reads. Co-Authored-By: Claude Opus 5 (1M context) --- docs/23_cli_reference.md | 3 - rust/mumdia/crates/mumdia/src/main.rs | 84 +++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 3 deletions(-) diff --git a/docs/23_cli_reference.md b/docs/23_cli_reference.md index 7b4f1671..2945d944 100644 --- a/docs/23_cli_reference.md +++ b/docs/23_cli_reference.md @@ -347,9 +347,6 @@ Options: --restrict-candidates Optional candidate allowlist (a prior run's psms.parquet): restrict extraction to these candidate_ids. For "gate first, then compete" - re-extract with a peak_claim strategy over only the gate-accepted survivors, keeping the two-pass profile map small - --fragment-offset - The library row that `--lib-precursors` row 0 came from, when that table is one isolation-window band of a larger library (`groups.window_groups` writes such bands, with ids rebased to `0..n`). The band's fragments are then read from the shared `--lib-fragments` table by that id range, selectively when the table is sorted by `candidate_id`, so a band needs no fragment table of its own - --config ``` diff --git a/rust/mumdia/crates/mumdia/src/main.rs b/rust/mumdia/crates/mumdia/src/main.rs index e0650be3..fbb5566b 100644 --- a/rust/mumdia/crates/mumdia/src/main.rs +++ b/rust/mumdia/crates/mumdia/src/main.rs @@ -255,6 +255,24 @@ enum Cmd { #[arg(long)] config: Option, }, + /// Pool a grouped run's band artifacts into the run-level tables. + /// + /// `run` does this itself at the end of a grouped search (`groups.window_groups`). + /// Standalone it is for the case where the search finished and the run did not: the + /// band directories hold everything, and pooling them is a byte copy of their parquet + /// row groups, so a killed run costs a pool rather than a re-search. + Pool { + /// The run's `groups/` directory, holding the `gNN/` band directories. + #[arg(long)] + groups_dir: String, + /// Where the pooled tables go. Default: the parent of `--groups-dir`, which is + /// where a run writes them. + #[arg(long)] + out_dir: Option, + /// Also pool `psms_extracted`, which only the candidate audit reads. + #[arg(long)] + psms: bool, + }, /// Keep the best candidate per competition group -> psms_competed.parquet. Compete { #[arg(long)] @@ -1294,6 +1312,72 @@ fn real_main() -> Result<()> { config_hash: &ch, })?; } + Cmd::Pool { + groups_dir, + out_dir, + psms, + } => { + let out_dir = out_dir.unwrap_or_else(|| { + std::path::Path::new(&groups_dir) + .parent() + .map(|p| p.to_string_lossy().to_string()) + .unwrap_or_else(|| ".".to_string()) + }); + // Band directories in band order: the pooled row order is the bands' order, + // and `gNN` sorts lexically into it. + let mut dirs: Vec = std::fs::read_dir(&groups_dir) + .with_context(|| format!("reading {groups_dir}"))? + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| { + p.is_dir() + && p.file_name() + .and_then(|n| n.to_str()) + .is_some_and(|n| n.starts_with('g')) + }) + .collect(); + dirs.sort(); + let bands: Vec = dirs + .iter() + .map(|d| { + let f = |n: &str| d.join(n).to_string_lossy().to_string(); + stages::pool::BandArtifacts { + psms: f("psms_extracted.parquet"), + chromatograms: f("chromatograms.parquet"), + competed: f("psms_competed.parquet"), + } + }) + .filter(|b| std::path::Path::new(&b.competed).exists()) + .collect(); + if bands.is_empty() { + anyhow::bail!( + "no band directory under {groups_dir} holds a psms_competed.parquet" + ); + } + let out = |n: &str| format!("{out_dir}/{n}"); + let (op, oc, ok) = ( + out("psms_extracted.parquet"), + out("chromatograms.parquet"), + out("psms_competed.parquet"), + ); + let stats = stages::pool::run(stages::pool::PoolParams { + bands: &bands, + out_psms: psms.then_some(op.as_str()), + out_chromatograms: &oc, + out_competed: &ok, + })?; + println!( + "pooled {} bands: {} competed rows, {} chromatogram rows, {} overlap duplicates removed{}", + bands.len(), + stats.competed, + stats.chromatograms, + stats.duplicates, + if psms { + format!(", {} extracted rows", stats.psms) + } else { + String::new() + } + ); + } Cmd::Audit { lib_precursors, psms_extracted, From a41b00ad56a31a089a5eefa8f3da02d273d9efca Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 16:00:38 +0200 Subject: [PATCH 21/35] perf(extract): flush on the candidate axis, not only on the window batch `window_groups` drops the windows that select no candidate of the band, so a band of a 63-band grouped plan overlaps only 1-3 windows: the whole band is one batch, the flush bound is `u32::MAX`, and nothing was ever flushed early. The accumulator held the entire band, which is 1.9-5.8M precursors' worth of hits, for the whole stage. `windows_in_flight` cannot help, because there is only one batch. The sub-ranges that already split each window's candidate range are now a GLOBAL grid shared by every window of the batch, and a sub-range is flushed as soon as every window that intersects it has reported. Per-window sub-ranges parallelise just as well but cannot be flushed: a candidate is final once every window of the batch has passed it, and that is a statement about a boundary the windows share. Tasks are one per (sub-range, window) pair that actually intersects, k-major, so the pool starts on what will be flushed first. The width is taken from the MEAN window span divided by `tasks_per_window` rather than from the global span divided by a task count, so each window still splits into about `tasks_per_window` pieces and the batch still runs about 2x threads of them. No extra probing, no extra `WindowNarrow` caches: these are the tasks that existed before, on a shared grid. On a one-window band at 32 threads the flush unit becomes 1/64 of the band; on a 16-window batch of a whole-library search it becomes roughly one window's span instead of sixteen. The end-of-batch drain stays, because a batch's windows can reach past the last window's range when the windows differ in width, and the sub-range grid stops at the widest window's end. Risk to output equality: PSM rows and values unchanged; chromatogram parquet row-group boundaries may move. Sub-ranges are flushed in ascending index and candidates ascending within, and the sub-range grid partitions the candidate axis, so the flushed sequence is the same ascending candidate sequence as before. A candidate's hits are still gathered run by run in window order, which is what every float reduction downstream depends on. What does change is where a flush ends: the ChromChunk partition now follows sub-range boundaries as well as `CAND_CHUNK`, so the chromatogram writer's row groups can fall differently. That is the same class of change `windows_in_flight` already documents as acceptable, and no row and no value moves with it. Also hoists the per-window query-m/z buffer out of the window probe and into the batch plan, since the per-window task count is now known there. Test: `candidate_range_split_reproduces_the_unsplit_accumulation` now runs the fixture twice, once with `bound = 0` (nothing flushes, the accumulation is compared as before across 2/8/16 threads) and once with `bound = u32::MAX`, and asserts the per-sub-range flush delivers exactly the same candidates, strictly ascending, with the same hits in the same order, and leaves nothing open. Co-Authored-By: Claude Opus 5 (1M context) --- .../crates/mumdia/src/stages/extract.rs | 352 ++++++++++++------ 1 file changed, 228 insertions(+), 124 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/extract.rs b/rust/mumdia/crates/mumdia/src/stages/extract.rs index e4c8a373..98abe675 100644 --- a/rust/mumdia/crates/mumdia/src/stages/extract.rs +++ b/rust/mumdia/crates/mumdia/src/stages/extract.rs @@ -1021,10 +1021,32 @@ fn window_groups(idx: &FragIndex, scans: &[Ms2Scan]) -> Vec { .collect() } -/// Probe one batch of isolation windows, in parallel across windows, and merge the hits -/// in window order. Merging in window order (not thread completion order) is what keeps -/// each candidate's hit sequence, and therefore every float reduction downstream, the -/// same as the serial path. +/// Probe one batch of isolation windows and flush what each candidate sub-range finalises, +/// in ascending candidate order. +/// +/// Two axes are cut. Across WINDOWS, because each scan belongs to exactly one of them. And +/// across a shared, GLOBAL grid of candidate sub-ranges, because a batch of `step` windows +/// on `t` threads leaves `t - step` threads idle whenever the batch is smaller than the +/// pool (measured at 4 in flight on 24 threads: 3.2-3.7 cores busy for the whole stage), +/// and because a candidate is final once every window of the batch has passed it, which is +/// a statement about a boundary all the windows share. Per-window sub-ranges would +/// parallelise just as well and could not be flushed. +/// +/// That matters most exactly where the stage hurts. `window_groups` drops the windows that +/// select no candidate of the band, so a band of a 63-band grouped plan overlaps only 1-3 +/// windows, the whole band is one batch, the flush bound is `u32::MAX` and nothing was ever +/// flushed early: the accumulator held the entire band. Cutting on the candidate axis makes +/// the flush unit a sub-range of it instead, at no extra probing, because the sub-ranges +/// are the tasks that already existed. +/// +/// Order is the contract. A candidate is flushed only once every window that can reach it +/// has reported, its hits are gathered run by run in window order, and sub-ranges are +/// flushed in ascending index, so the per-candidate hit sequence and the ascending PSM row +/// order are exactly what the serial path produced. Chromatogram parquet row-group +/// boundaries can move, because the flush is now cut on sub-ranges rather than on the +/// batch; no row and no value moves with them. +/// +/// Returns true when `flush` asked to stop (the chromatogram writer went away). #[allow(clippy::too_many_arguments)] fn accumulate_groups( idx: &FragIndex, @@ -1035,55 +1057,86 @@ fn accumulate_groups( mass_off: &MassOffset, cfg: &ExtractConfig, restrict: Option<&std::collections::HashSet>, + // First candidate a window of a LATER batch can still add hits to; nothing at or above + // it may be flushed here. + bound: u32, acc: &mut HitAcc, -) { - // Partials are appended to `acc` in window order as they complete, not collected first - // and merged after. The channel hands each finished partial to this thread; the - // reorder buffer holds the ones that finished ahead of an earlier window, so the run - // order, and with it every candidate's hit sequence, is exactly the serial one. No - // per-candidate merge happens here any more: a window's partial is a CSR store per - // sub-range task, and the concatenation is resolved lazily by `gather_chunk` at flush. - let (tx, rx) = std::sync::mpsc::channel::<(usize, Vec)>(); - // A batch of `step` windows on `t` threads leaves `t - step` threads idle whenever the - // batch is smaller than the pool, which is what a `windows_in_flight` set low for memory - // does: measured at 4 in flight on 24 threads, 3.2-3.7 cores busy for the whole stage. - // So each window is also split across its CANDIDATE range. A candidate belongs to - // exactly one sub-range, so the sub-results are disjoint: they hold the window's hits - // once between them (no duplication of the accumulator, which is the stage's memory), - // they concatenate instead of merging per candidate, and each candidate's hits are still - // appended in ascending scan order, so every float reduction downstream is unchanged. - // Splitting the candidate axis rather than the scan axis is also where the work is: a - // peak's cost is dominated by walking its posting list, which the narrowed sub-range - // divides, while the peak loop itself is repeated per task. + chunk: &mut HitStore, + flush: &mut (dyn FnMut(&mut HitStore) -> bool + Send), +) -> bool { + if groups.is_empty() { + return false; + } let threads = rayon::current_num_threads().max(1); - let tasks_per_window = (threads * 2).div_ceil(groups.len().max(1)).max(1); - let probe = |g: &WinGroup| -> Vec { - let ids = &g.scans; - let n_cand = (g.hi_cid - g.lo_cid) as usize; - let tasks = tasks_per_window - .min(n_cand.div_ceil(MIN_CANDIDATES_PER_TASK)) - .max(1); - // Offset-corrected query m/z per peak, computed ONCE for the window instead of - // once per sub-range task. `MassOffset::factor_at` is a binary search into the - // m/z-dependent calibration grid when a masscal grid is present and a division - // either way, and both were repeated `tasks` times -- up to 2x threads -- for - // every peak of every scan of the window. Flat with one offset per scan, so it is - // two allocations rather than one per scan, and built only when the window is - // actually split (unsplit, there is nothing to share it with). - let qmz: Option<(Vec, Vec)> = (tasks > 1).then(|| { - let mut q: Vec = - Vec::with_capacity(ids.iter().map(|&si| scans[si].peaks.len()).sum()); - let mut off: Vec = Vec::with_capacity(ids.len() + 1); - off.push(0); - for &si in ids { - for peak in &scans[si].peaks { - q.push(peak.mz / mass_off.factor_at(peak.mz)); + let tasks_per_window = (threads * 2).div_ceil(groups.len()).max(1); + // The sub-range WIDTH is taken from the mean window span divided by + // `tasks_per_window`, not from the global span divided by a task count, so each window + // still splits into about `tasks_per_window` pieces and the batch still runs about + // 2x threads of them -- the same parallelism as before, on a shared grid. + let g_lo = groups.iter().map(|g| g.lo_cid).min().expect("non-empty"); + let g_hi = groups.iter().map(|g| g.hi_cid).max().expect("non-empty"); + let mean_span = groups + .iter() + .map(|g| (g.hi_cid - g.lo_cid) as usize) + .sum::() + / groups.len(); + let width = mean_span + .div_ceil(tasks_per_window) + .max(MIN_CANDIDATES_PER_TASK) + .min(u32::MAX as usize) as u64; + let n_sub = ((g_hi - g_lo) as u64).div_ceil(width).max(1) as usize; + let sub_bounds = |k: usize| -> (u32, u32) { + let hi = g_hi as u64; + let s = (g_lo as u64 + k as u64 * width).min(hi); + ((s) as u32, (s + width).min(hi) as u32) + }; + // One task per (sub-range, window) pair that actually intersects, k-major so the pool + // starts on the sub-ranges that will be flushed first. `expected[k]` is how many + // windows must report before sub-range `k` is final. + let mut tasks: Vec<(usize, usize, u32, u32)> = Vec::new(); + let mut expected: Vec = vec![0; n_sub]; + let mut per_window: Vec = vec![0; groups.len()]; + for (k, exp) in expected.iter_mut().enumerate() { + let (s, e) = sub_bounds(k); + for (gi, g) in groups.iter().enumerate() { + let (lo, hi) = (g.lo_cid.max(s), g.hi_cid.min(e)); + if hi > lo { + tasks.push((k, gi, lo, hi)); + *exp += 1; + per_window[gi] += 1; + } + } + } + // Offset-corrected query m/z per peak, computed ONCE per window instead of once per + // sub-range task. `MassOffset::factor_at` is a binary search into the m/z-dependent + // calibration grid when a masscal grid is present, and a division either way, and both + // were repeated once per task -- up to 2x threads -- for every peak of every scan of + // the window. Flat with one offset per scan, so it is two allocations per window rather + // than one per scan, and built only where the window is actually split. + let qmz: Vec, Vec)>> = groups + .iter() + .enumerate() + .map(|(gi, g)| { + (per_window[gi] > 1).then(|| { + let mut q: Vec = + Vec::with_capacity(g.scans.iter().map(|&si| scans[si].peaks.len()).sum()); + let mut off: Vec = Vec::with_capacity(g.scans.len() + 1); + off.push(0); + for &si in &g.scans { + for peak in &scans[si].peaks { + q.push(peak.mz / mass_off.factor_at(peak.mz)); + } + off.push(q.len()); } - off.push(q.len()); - } - (q, off) - }); - let probe_range = |lo: u32, hi: u32| -> HitStore { + (q, off) + }) + }) + .collect(); + let (tx, rx) = std::sync::mpsc::channel::<(usize, usize, HitStore)>(); + { + let probe_range = |gi: usize, lo: u32, hi: u32| -> HitStore { + let ids = &groups[gi].scans; + let qmz = &qmz[gi]; // Flat `(cid, hit)` pairs in probe order, grouped by candidate at the end of // the task with a stable counting sort. The task-local `HashMap>` this replaces was one of the two populations of medium heap @@ -1183,70 +1236,73 @@ fn accumulate_groups( hits: flat_hit, } }; - if tasks <= 1 { - return vec![probe_range(g.lo_cid, g.hi_cid)]; - } - let span = n_cand.div_ceil(tasks) as u32; - let ranges: Vec<(u32, u32)> = (0..tasks as u32) - .map(|i| { - ( - g.lo_cid + i * span, - (g.lo_cid + (i + 1) * span).min(g.hi_cid), - ) - }) - .filter(|&(lo, hi)| hi > lo) - .collect(); - // Ascending, disjoint candidate spans, so the stores concatenate into one - // ascending run without a per-candidate merge. - ranges - .par_iter() - .map(|&(lo, hi)| probe_range(lo, hi)) - .collect() - }; - let n = groups.len(); - // `move`: the receiver is Send but not Sync, so it has to be owned by the scope closure - // (the consumer) rather than borrowed into it. - rayon::scope(move |sc| { - // The workers run under rayon; this thread is the single consumer. `spawn` so the - // scope does not block the consumer loop below until every window has finished. - sc.spawn(move |_| { - groups - .par_iter() - .enumerate() - .for_each_with(tx, |tx, (gi, g)| { - let part = probe(g); - // A closed receiver only happens if the consumer panicked; nothing to do. - let _ = tx.send((gi, part)); + let n = tasks.len(); + // `move`: the receiver is Send but not Sync, so it has to be owned by the scope + // closure (the consumer) rather than borrowed into it. + rayon::scope(move |sc| { + // The workers run under rayon; this thread is the single consumer. `spawn` so + // the scope does not block the consumer loop below until every task is done. + sc.spawn(move |_| { + tasks.par_iter().for_each_with(tx, |tx, &(k, gi, lo, hi)| { + // A closed receiver only happens if the consumer panicked. + let _ = tx.send((k, gi, probe_range(gi, lo, hi))); }); - }); - let mut next = 0usize; - let mut pending: std::collections::BTreeMap> = - std::collections::BTreeMap::new(); - for _ in 0..n { - let (gi, part) = rx.recv().expect("every window sends exactly one partial"); - if gi == next { - acc.runs.push(HitRun::new(part)); - next += 1; - while let Some(p) = pending.remove(&next) { - acc.runs.push(HitRun::new(p)); - next += 1; + }); + // Stores wait here, keyed by `(sub-range, window)`, until their whole + // sub-range has reported. `next_k` is the lowest sub-range not yet flushed, so + // the map's first keys are always that sub-range's. + let mut pending: std::collections::BTreeMap<(usize, usize), HitStore> = + std::collections::BTreeMap::new(); + let mut arrived: Vec = vec![0; n_sub]; + let mut next_k = 0usize; + let mut stopped = false; + for _ in 0..n { + let (k, gi, store) = rx.recv().expect("every task sends exactly one store"); + arrived[k] += 1; + pending.insert((k, gi), store); + while next_k < n_sub && arrived[next_k] == expected[next_k] { + let mut runs = std::mem::take(&mut acc.runs); + for _ in 0..expected[next_k] { + let key = *pending.keys().next().expect("the sub-range is complete"); + debug_assert_eq!(key.0, next_k); + let store = pending.remove(&key).expect("key came from the map"); + // One run per window, appended in window order, after whatever + // earlier windows already left in the accumulator. + runs.push(HitRun::new(vec![store])); + } + // Everything this sub-range owns is final, except what a later batch + // can still reach. + let sub_bound = sub_bounds(next_k).1.min(bound); + while !stopped { + gather_chunk(&mut runs, sub_bound, CAND_CHUNK, chunk); + if chunk.is_empty() { + break; + } + if !flush(chunk) { + stopped = true; + } + } + acc.runs = runs; + acc.compact(); + next_k += 1; } - } else { - pending.insert(gi, part); } - } - debug_assert!(pending.is_empty() && next == n); - }); + debug_assert!(pending.is_empty() && next_k == n_sub); + stopped + }) + } } /// Isolation windows probed before the driver flushes the candidates that are now final. -/// Memory scales with this: the accumulator holds the hits of the windows in flight, so a -/// batch of `n` windows of width `w` stepping `s` keeps about `n * s + w` Th of the -/// precursor axis open. It is no longer the unit of parallelism: each window is probed in -/// parallel over sub-ranges of its candidates (`accumulate_groups`), so a small batch keeps -/// every thread busy without holding more of the precursor axis open. Measured on the HYE benchmark at 32 threads before that change (docs/27 section -/// 3.10): 32 in flight 24.65 GiB / 5:00, 16 in flight 16.57 GiB / 5:04, 8 in flight -/// 12.31 GiB / 5:26, identical output throughout. +/// It is neither the unit of parallelism nor, since the sub-range flush, the unit of +/// memory: each window is probed in parallel over a shared grid of candidate sub-ranges, +/// and `accumulate_groups` flushes a sub-range as soon as every window of the batch has +/// passed it, so the accumulator holds a sub-range of the batch rather than the batch. +/// The batch still bounds how much of the precursor axis is open at once, because a +/// candidate is not final until the last window that can reach it has run. Measured on the +/// HYE benchmark at 32 threads before either change (docs/27 section 3.10): 32 in flight +/// 24.65 GiB / 5:00, 16 in flight 16.57 GiB / 5:04, 8 in flight 12.31 GiB / 5:26, +/// identical output throughout. const DEFAULT_MAX_WINDOWS_IN_FLIGHT: usize = 16; /// Fewest candidates a probing task takes: below this the per-task narrowed bin cache (one @@ -3035,11 +3091,23 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { let mut chunk = HitStore::default(); let mut open_hits_max = 0usize; if let Some(groups) = stream_groups.as_ref() { + // Counting the flushed candidates here rather than inside `emit_batch` keeps + // the eager path below able to borrow `emit_batch` on its own. + let mut flush = |c: &mut HitStore| -> bool { + n_materialized += c.len() as u64; + emit_batch(c) + }; let step = groups_in_flight(p.cfg); let mut gi = 0usize; while gi < groups.len() { let upto = (gi + step).min(groups.len()); - accumulate_groups( + // Everything below the next batch's first candidate is final: windows are + // processed in ascending m/z and precursors are sorted by m/z, so no later + // window can add a hit below `bound`. Within the batch, `accumulate_groups` + // flushes each candidate sub-range as soon as every window has passed it, + // so the accumulator holds a sub-range of the band rather than the band. + let bound = groups.get(upto).map(|g| g.lo_cid).unwrap_or(u32::MAX); + let mut stopped = accumulate_groups( fidx.as_ref() .expect("streamed path implies a fragment index"), &groups[gi..upto], @@ -3049,24 +3117,23 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { &mass_off, p.cfg, restrict.as_ref(), + bound, &mut acc_stream, + &mut chunk, + &mut flush, ); gi = upto; open_hits_max = open_hits_max.max(acc_stream.n_hits()); - // Everything below the next window's first candidate is final: windows are - // processed in ascending m/z and precursors are sorted by m/z, so no later - // window can add a hit below `bound`. - let bound = groups.get(gi).map(|g| g.lo_cid).unwrap_or(u32::MAX); - let mut stopped = false; - loop { + // Candidates past the batch's own sub-range grid (a window of this batch + // can reach past the last window's range when the windows differ in width) + // are final too, and the sub-range loop cannot have reached them. + while !stopped { gather_chunk(&mut acc_stream.runs, bound, CAND_CHUNK, &mut chunk); if chunk.is_empty() { break; } - n_materialized += chunk.len() as u64; - if !emit_batch(&mut chunk) { + if !flush(&mut chunk) { stopped = true; - break; } } acc_stream.compact(); @@ -3542,28 +3609,42 @@ mod accumulate_tests { grid_ppm: Vec::new(), }; let cfg = ExtractConfig::default(); - let run = |threads: usize| -> HitAcc { + // `bound = 0` flushes nothing, so the whole batch stays in the accumulator and the + // comparison is over the accumulation itself. `flush_all` below runs the same + // fixture with the flush live. + let run = |threads: usize, bound: u32| -> (HitAcc, Vec<(u32, Vec)>) { let pool = rayon::ThreadPoolBuilder::new() .num_threads(threads) .build() .expect("thread pool"); let mut acc = HitAcc::default(); + let mut chunk = HitStore::default(); + let mut flushed: Vec<(u32, Vec)> = Vec::new(); pool.install(|| { + let mut sink = |c: &mut HitStore| -> bool { + for (cid, hits) in c.slices_mut() { + flushed.push((cid, hits.to_vec())); + } + true + }; accumulate_groups( - &idx, &groups, &sc, &rt_lo, &rt_hi, &mass_off, &cfg, None, &mut acc, - ) + &idx, &groups, &sc, &rt_lo, &rt_hi, &mass_off, &cfg, None, bound, &mut acc, + &mut chunk, &mut sink, + ); }); - acc + (acc, flushed) }; // Two threads: 2 * 2 / 4 windows = one sub-range per window, the unsplit path. - let unsplit = materialize(run(2)); + let (acc, empty) = run(2, 0); + assert!(empty.is_empty(), "bound 0 must flush nothing"); + let unsplit = materialize(acc); assert!(!unsplit.is_empty(), "the fixture must produce hits"); assert!( unsplit.values().any(|v| v.len() > 1), "candidates must collect several hits, or hit order proves nothing" ); for threads in [8, 16] { - let split = materialize(run(threads)); + let split = materialize(run(threads, 0).0); assert_eq!( split.len(), unsplit.len(), @@ -3577,6 +3658,29 @@ mod accumulate_tests { ); } } + // Flushing per candidate sub-range instead of per batch must deliver exactly the + // same candidates, ascending, with the same hits in the same order: the sub-range + // grid is where a candidate becomes final, not where its evidence changes. + for threads in [2, 8, 16] { + let (acc, flushed) = run(threads, u32::MAX); + assert_eq!(acc.n_hits(), 0, "{threads} threads: nothing may stay open"); + assert!( + flushed.windows(2).all(|w| w[0].0 < w[1].0), + "{threads} threads: flushed candidates must be strictly ascending" + ); + assert_eq!( + flushed.len(), + unsplit.len(), + "{threads} threads: flushed candidate count" + ); + for (cid, hits) in &flushed { + assert_eq!( + unsplit.get(cid), + Some(hits), + "{threads} threads: flushed candidate {cid} hits differ" + ); + } + } } /// Flush the whole accumulator through the chunked gather the driver uses, and read the From 1c6f8c41f5f8091544b4810143b16c5aa9ddcc01 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 16:08:51 +0200 Subject: [PATCH 22/35] fix(rescore): validate the streamed handoff as it is written, not after An adversarial review of the feature-streaming change found five issues in `stages/rescore.rs`. Four are real and fixed here; one is a documentation claim, corrected in place rather than by changing code. Fixed: - The non-finite validation ran AFTER the whole handoff was written and published. It used to run before `run_pin_sidecar` was entered, so a malformed competed table cost nothing; deferring it made the same failure pay a full handoff write, hundreds of GB at experiment scale, for a run that was always going to abort. The scan now aborts the stream on the first offending row. The message is unchanged: `bad_scalar` is known before the stream starts, so `bail_non_finite` still makes the same choice between the two answers on the row index, and the stream stops at the scalar row once no earlier feature row can win the tie-break. A parquet handoff never reaches its final name (`AtomicPath`), but the PIN encoding writes it directly, so an aborted stream now removes it. - The row-count guard existed on the streamed path only. The matrix path gets the same guard, and both now go through one `refuse_row_disagreement` helper, as does a third check the review flagged as a silently dropped invariant: `merge_col` appends whatever the reader returned, while `source` is still sized from the footer, so every concatenated metadata column is now measured against the footer row count at the point `n` is taken from it. - An empty feature set is refused by name. The old code neither errored nor worked: `iter_rows` is `chunks_exact(n_features.max(1))` over an empty buffer, so the validation loop ran zero times, which also skipped the prelim_score/precursor_mz check for every row, and the run trained on no features. An explicitly empty `rescore.features` list was already an error, so this makes the two consistent. - `a_shared_column_array_writes_the_same_parquet_as_two_copies` ran at 2,000 rows, below `write_table`'s 65,536-row chunk, where the two writers cannot diverge. It now runs at 3 chunks plus a short last one. The row-group seam is pinned in mumdia-io at 1,048,583 rows. Rejected: - "`write_scored_table`'s memory claim inverts on the merge target": the byte output is unaffected and the mechanism is right, but the review is correct that on `perf/integrate` chunked `write_table` holds one chunk of Arrow copies while one batch re-packs each string column in full. The doc comment now says that instead of claiming an unqualified saving; the code is unchanged. - `assert!(n <= u32::MAX)` in `fdr::target_decoy_q` stays an assert. It is unreachable at any allocatable scale, the alternative it guards is silent truncation scattering q onto the wrong rows, and an `anyhow` error would change a signature used by search-seed and others, outside this file set. - `rescore.max_feature_matrix_gib` stays checked on the streamed path. The ceiling is `rows x features x 4`, which is exactly the matrix the Python worker builds from the handoff, so it still describes a real allocation; making it path-dependent would let a strict `nn_torch` pool that will OOM the worker start anyway. The default is 0.0 (off). - The commit-message residency figure for the staging block: corrected in the code comment (the staged f32 block plus what `flush_block` transposes out of it, roughly twice the stated 390 MB). Tests: the two abort tests fail against the pre-fix ordering (the handoff was written and published before the error), the tie-break tests pass before and after, which is the point. Not fixed, outside the file set: `cargo fmt --check` already fails on `stages/pool.rs` and `mumdia-io/src/table.rs` at this branch tip. Co-Authored-By: Claude Opus 5 (1M context) --- .../crates/mumdia/src/stages/rescore.rs | 372 ++++++++++++++++-- 1 file changed, 341 insertions(+), 31 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/rescore.rs b/rust/mumdia/crates/mumdia/src/stages/rescore.rs index 31c982d8..5776c757 100644 --- a/rust/mumdia/crates/mumdia/src/stages/rescore.rs +++ b/rust/mumdia/crates/mumdia/src/stages/rescore.rs @@ -106,6 +106,25 @@ fn merge_col(dst: &mut Vec, mut src: Vec, total_rows: usize) { dst.append(&mut src); } +/// Refuse a read of the competed inputs whose row count disagrees with the row count every +/// other read of them produced. +/// +/// The metadata columns, the streamed handoff and the feature matrix are now three separate +/// passes over the same files, where they used to be one loop that could not disagree with +/// itself. Any disagreement row-misaligns the features against `label`, `source` and every +/// other metadata column, and nothing downstream can detect that, so every pass is measured +/// against the row count the parquet footers declared. +fn refuse_row_disagreement(what: &str, got: usize, expected: usize) -> Result<()> { + if got != expected { + anyhow::bail!( + "rescore: {what} has {got} rows where {expected} were expected; the reads of the \ + competed inputs disagree about the input, and a disagreement row-misaligns the \ + feature rows against every metadata column" + ); + } + Ok(()) +} + /// Refuse a rescore whose population cannot support target-decoy FDR. /// /// A one-sided population is not a hard rescore failure in any technical sense: @@ -174,6 +193,25 @@ pub fn run(p: RescoreParams) -> Result { // with it. The list actually used is recorded in the artifact report below, which is // the source of truth for what the classifier saw. let feat_names = resolve_feature_subset(p.cfg, &expected_schema.feature_columns)?; + // An empty feature set is refused HERE rather than carried into the passes below. + // + // `resolve_feature_subset` already refuses an explicitly empty selection, and + // `FeatureSchema::read` already refuses to reconstruct an empty list from the parquet, + // so the one way left in is a `.schema.json` companion whose `feature_columns` is `[]`. + // What that used to do was an accident of the validation loop's shape, not a decision: + // it ran `for (row, values) in feats.iter_rows().enumerate()`, `iter_rows` is + // `chunks_exact(n_features.max(1))`, and with no feature columns the matrix buffer is + // empty, so the loop never ran a single iteration -- which silently skipped the + // prelim_score/precursor_mz finiteness check for EVERY row as well, and then trained a + // classifier on zero inputs. Neither is worth preserving; say so instead. + if feat_names.is_empty() { + anyhow::bail!( + "rescore: the feature schema of {} declares no feature columns, so the \ + classifier would have no input. Check the `.schema.json` companion next to \ + the competed table.", + p.competed[0] + ); + } if feat_names.len() != expected_schema.feature_columns.len() { info!( selected = feat_names.len(), @@ -266,6 +304,32 @@ pub fn run(p: RescoreParams) -> Result { merge_col(&mut elution_hi, ehi, total_rows); merge_col(&mut source, vec![src as u32; t.nrows], total_rows); } + // The row loop this pass replaces was `for i in 0..t.nrows { cid.push(c[i]); ... }`, + // which could not produce columns of different lengths: it indexed every column at the + // footer's row count and panicked if one was short. `merge_col` appends whatever the + // reader returned, and `source` is still built from `t.nrows`, so a reader column that + // disagreed with the footer would row-misalign `source` -- the column quant splits the + // scored table on -- against `cid`, `label` and the feature rows, with nothing + // downstream able to detect it (`n` is taken as `cid.len()` below). `TableFile` + // concatenates all batches and so returns exactly `nrows` today; this keeps the + // invariant checked at the point that relies on it rather than assumed. + for (name, len) in [ + ("candidate_id", cid.len()), + ("peak_rank", peak_rank.len()), + ("label", label.len()), + ("base_peptide_id", base.len()), + ("peptidoform", pform.len()), + ("protein", protein.len()), + ("charge", charge.len()), + ("prelim_score", prelim.len()), + ("precursor_mz", mz.len()), + ("apex_rt", apex_rt.len()), + ("elution_lo", elution_lo.len()), + ("elution_hi", elution_hi.len()), + ("source", source.len()), + ] { + refuse_row_disagreement(&format!("the competed column '{name}'"), len, total_rows)?; + } crate::fdr::validate_labels(&label)?; let is_decoy: Vec = label.iter().map(|l| l == "decoy").collect(); let (mut is_entrapment, mut is_real_target) = classify_entrapment(p.cfg, &protein, &is_decoy); @@ -296,8 +360,10 @@ pub fn run(p: RescoreParams) -> Result { // `rows x features x 4` bytes, ~250 GB at experiment scale -- purely so that it could // be copied to disk. Streaming the competed feature columns straight into the handoff // removes it: what stays resident is one staging block of `HANDOFF_BATCH_ROWS` rows - // (about 390 MB at 387 features), and the non-finite validation moves into the same - // pass, inline with the narrowing, while the row is still in cache. + // (387 MB of staged f32 at 387 features, plus what `flush_block` transposes out of it + // while a block is encoded, so call the transient roughly twice that), and the + // non-finite validation moves into the same pass, inline with the narrowing, while the + // row is still in cache. let sidecar_script = match p.cfg.classifier { RescorerKind::Mokapot => Some("mokapot_worker.py"), RescorerKind::NnTorch => Some("nn_rescore_worker.py"), @@ -311,27 +377,49 @@ pub fn run(p: RescoreParams) -> Result { let mut prewritten: Option = None; if stream_to_handoff { let paths = sidecar_paths(&p, sidecar_script.expect("checked in the condition")); - let mut w = HandoffWriter::new(&paths, &feat_names, &label, &pform, &protein, &mz)?; - let mut bad_feature = None; - for_each_feature_row(p.competed, &feat_names, |row, values| { - if bad_feature.is_none() { + // The validation aborts the stream on the first offending row rather than after the + // file is complete. A malformed competed table used to cost nothing: the serial + // scan ran before `run_pin_sidecar` was ever entered and nothing had been written. + // Deferring it until after `finish()` made the same failure cost a full handoff + // write, which at experiment scale is hundreds of GB of IO for a run that was + // always going to abort. The message is unchanged, because `bad_scalar` is already + // known here and `bail_non_finite` still makes the same choice between the two on + // the row index. + let streamed = { + let mut w = HandoffWriter::new(&paths, &feat_names, &label, &pform, &protein, &mz)?; + let scan = for_each_feature_row(p.competed, &feat_names, |row, values| { if let Some(col) = values.iter().position(|v| !v.is_finite()) { - bad_feature = Some((row, col, values[col])); + bail_non_finite(Some((row, col, values[col])), bad_scalar, &feat_names)?; + } + // Past the offending scalar row with no earlier bad feature, the choice + // `bail_non_finite` would make is already decided (a later feature row + // loses the tie-break), so there is nothing left to learn from the rest of + // the file. + if bad_scalar.is_some_and(|scalar_row| row >= scalar_row) { + bail_non_finite(None, bad_scalar, &feat_names)?; } + w.push_row(row, values) + }); + // `finish` consumes the writer, and the `Err` arm drops it, so the file is + // closed either way before the cleanup below. + scan.and_then(|()| w.finish()) + }; + let rows = match streamed { + Ok(rows) => rows, + Err(e) => { + // A parquet handoff is written through `AtomicPath` and never appears at + // its final name, but the PIN encoding writes the final path directly, so + // an aborted stream would leave a truncated PIN for the next reader to + // find. Take the rubble with us. + let _ = std::fs::remove_file(&paths.handoff); + return Err(e); } - w.push_row(row, values) - })?; - let rows = w.finish()?; - // The metadata and the feature values now come from two separate passes over the - // same files. Nothing should be able to make them disagree, and if anything does, - // the handoff is row-misaligned with every metadata column: refuse rather than - // train on it. - if rows as usize != n { - anyhow::bail!( - "rescore: the streamed handoff wrote {rows} rows but the competed metadata \ - has {n}; the two passes disagree about the input" - ); - } + }; + // The metadata and the feature values come from two separate passes over the same + // files. Nothing should be able to make them disagree, and if anything does, the + // handoff is row-misaligned with every metadata column: refuse rather than train + // on it. + refuse_row_disagreement("the streamed handoff", rows as usize, n)?; info!( path = %paths.handoff, rows, @@ -339,10 +427,12 @@ pub fn run(p: RescoreParams) -> Result { "rescore: streamed the competed features into the sidecar handoff \ (no engine-side feature matrix)" ); - bail_non_finite(bad_feature, bad_scalar, &feat_names)?; prewritten = Some(paths); } else { let feats = load_feature_matrix(p.competed, &feat_names, total_rows)?; + // Same guard as the streamed path's: the feature pass and the metadata pass are + // separate reads of the same files and only this comparison ties them together. + refuse_row_disagreement("the feature matrix", feats.rows(), n)?; let bad_feature = feats .find_non_finite() .map(|(row, col)| (row, col, feats.row(row)[col])); @@ -864,10 +954,19 @@ struct ScoredColumns { /// Three of the 21 columns are duplicates by construction: `protein_group` is `protein`, /// and `global_q_value` and `experiment_psm_q` are both the pooled `q_value` (the comments /// at their definitions say so). Through `write_table` each duplicate had to arrive as its -/// own `Vec`, so the stage carried a second full copy of the protein strings -- one heap -/// block per PSM -- and two more copies of the pooled q column, all three alive at the -/// same moment as the originals. An `ArrayRef` is refcounted, so passing the same array in -/// several slots of the batch writes the same bytes from one buffer. +/// own `Vec` -- `protein.clone()` and two `psm_q.clone()`s at the call site -- so the stage +/// carried a second full copy of the protein strings, one heap block per PSM, and two more +/// copies of the pooled q column, all three alive at the same moment as the originals. An +/// `ArrayRef` is refcounted, so passing the same array in several slots of the batch writes +/// the same bytes from one buffer. +/// +/// What this does NOT claim is a strict reduction in peak: `write_table` encodes in +/// `WRITE_TABLE_CHUNK_ROWS` chunks and moves each chunk's rows out of the source vectors, +/// so its own Arrow copy is one chunk wide, while one batch means the `StringArray` +/// conversion re-packs each of `peptidoform`, `label` and `protein` into one contiguous +/// offsets+values buffer at once. The three clones go away; a full-width Arrow copy of each +/// string column arrives. Which side is larger depends on the string widths of the run, so +/// the honest claim is the one the test below pins: the file is identical. /// /// The parquet is byte-identical to the previous `write_table` call: the same schema /// (names, types, all non-nullable, same order), one record batch, and `BatchWriter` opens @@ -2142,6 +2241,19 @@ b /// feature columns. No `.schema.json` companion, so the feature list is reconstructed /// from the parquet's own columns (`NON_FEATURE_COLUMNS`). fn crafted_competed_table(path: &str, rows: usize) { + crafted_competed_table_planting(path, rows, None, None); + } + + /// The same table with a NaN planted in the `feat_b` feature column and/or in the + /// `precursor_mz` metadata column, which is what a malformed competed input looks like + /// to the two validations (a null f64 cell reads back as NaN, so a null is the same + /// case). + fn crafted_competed_table_planting( + path: &str, + rows: usize, + nan_feature_row: Option, + nan_mz_row: Option, + ) { let cols = vec![ Col::U32("candidate_id".into(), (0..rows as u32).collect()), Col::Str( @@ -2171,7 +2283,15 @@ b ), Col::F64( "precursor_mz".into(), - (0..rows).map(|i| 400.0 + i as f64).collect(), + (0..rows) + .map(|i| { + if nan_mz_row == Some(i) { + f64::NAN + } else { + 400.0 + i as f64 + } + }) + .collect(), ), Col::F64( "apex_rt".into(), @@ -2194,7 +2314,15 @@ b ), Col::F64( "feat_b".into(), - (0..rows).map(|i| (i % 7) as f64 * 0.25).collect(), + (0..rows) + .map(|i| { + if nan_feature_row == Some(i) { + f64::NAN + } else { + (i % 7) as f64 * 0.25 + } + }) + .collect(), ), ]; write_table(path, cols).unwrap(); @@ -2260,6 +2388,175 @@ b ); } + /// Drive `run` down the streamed-handoff path (strict + a sidecar classifier + an + /// interpreter that cannot be spawned) and return the error, the handoff path the + /// invocation used, and its work directory. + fn run_streamed(competed: &str, name: &str) -> (String, String, String) { + let work = scratch(&format!("{name}_work")); + let out = scratch(&format!("{name}_scored.parquet")); + let cfg = RescoreConfig { + classifier: RescorerKind::NnTorch, + strict: true, + python: Some("mumdia-no-such-interpreter-for-this-test".to_string()), + ..Default::default() + }; + let err = run(RescoreParams { + competed: &[competed.to_string()], + out: &out, + work_dir: &work, + script_dir: "scripts", + cfg: &cfg, + config_hash: "test", + }) + .unwrap_err() + .to_string(); + let tag = format!( + "{}_{}", + std::path::Path::new(&out) + .file_stem() + .unwrap() + .to_str() + .unwrap(), + std::process::id() + ); + (err, format!("{work}/rescore_{tag}.features.parquet"), work) + } + + /// Nothing that looks like a handoff (published or half-written) is left in `work`. + fn no_handoff_rubble(work: &str) { + let left: Vec = std::fs::read_dir(work) + .map(|d| { + d.filter_map(|e| e.ok()) + .map(|e| e.file_name().to_string_lossy().to_string()) + .filter(|n| n.contains(".features.parquet") || n.ends_with(".pin")) + .collect() + }) + .unwrap_or_default(); + assert!(left.is_empty(), "handoff rubble left behind: {left:?}"); + } + + #[test] + fn a_non_finite_feature_aborts_the_stream_instead_of_writing_the_whole_handoff() { + // Before the features were streamed, the non-finite scan ran over the resident + // matrix and bailed before `run_pin_sidecar` was ever entered, so a malformed + // competed table cost nothing. Streaming moved the scan into the write and then + // reported it after `finish()`: the full handoff -- hundreds of GB at experiment + // scale -- was written and published for a run that was always going to abort. + // The row and column named must not change, and the file must not be there. + let competed = scratch("nanfeat_competed.parquet"); + crafted_competed_table_planting(&competed, 24, Some(3), None); + let (err, handoff, work) = run_streamed(&competed, "nanfeat"); + assert_eq!( + err, "rescore input contains non-finite feature 'feat_b' at flat row 3: NaN", + "the message must name the same row and column the serial scan named" + ); + assert!( + !std::path::Path::new(&handoff).exists(), + "the handoff must not survive a stream that aborted: {handoff}" + ); + no_handoff_rubble(&work); + } + + #[test] + fn a_non_finite_precursor_mz_aborts_the_stream_at_its_row() { + // The scalar scan runs before the stream, so its row is already known; the stream + // still has to reach that row before the message is decided (an earlier bad feature + // would win), and must stop there rather than write the remaining rows. + let competed = scratch("nanmz_competed.parquet"); + crafted_competed_table_planting(&competed, 24, None, Some(2)); + let (err, handoff, work) = run_streamed(&competed, "nanmz"); + assert_eq!( + err, + "rescore input contains non-finite prelim_score/precursor_mz at flat row 2" + ); + assert!(!std::path::Path::new(&handoff).exists(), "{handoff}"); + no_handoff_rubble(&work); + } + + #[test] + fn the_first_offending_row_still_decides_the_message_when_both_kinds_are_present() { + // The serial loop this replaces checked a row's features before its scalars and + // stopped at the first offending ROW. Aborting the stream early must not change + // which of the two answers is reported. + let a = scratch("both_scalar_first.parquet"); + crafted_competed_table_planting(&a, 24, Some(5), Some(2)); + let (err, _, _) = run_streamed(&a, "both_scalar_first"); + assert_eq!( + err, "rescore input contains non-finite prelim_score/precursor_mz at flat row 2", + "the scalar row is earlier, so it is the one reported" + ); + + let b = scratch("both_feature_first.parquet"); + crafted_competed_table_planting(&b, 24, Some(1), Some(4)); + let (err, _, _) = run_streamed(&b, "both_feature_first"); + assert_eq!( + err, "rescore input contains non-finite feature 'feat_b' at flat row 1: NaN", + "the feature row is earlier, so it is the one reported" + ); + } + + #[test] + fn the_earlier_row_wins_and_a_tie_goes_to_the_feature() { + // The tie-break `bail_non_finite` applies, pinned directly: within one row the old + // serial loop tested the features first, so an equal row index reports the feature. + let names: Vec = ["f0", "f1"].iter().map(|s| s.to_string()).collect(); + assert!(bail_non_finite(None, None, &names).is_ok()); + let feature_at = |row: usize| Some((row, 1, f32::NAN)); + for (feature, scalar, expect) in [ + (feature_at(3), None, "feature 'f1' at flat row 3"), + (None, Some(3), "prelim_score/precursor_mz at flat row 3"), + (feature_at(3), Some(7), "feature 'f1' at flat row 3"), + ( + feature_at(7), + Some(3), + "prelim_score/precursor_mz at flat row 3", + ), + (feature_at(3), Some(3), "feature 'f1' at flat row 3"), + ] { + let err = bail_non_finite(feature, scalar, &names) + .unwrap_err() + .to_string(); + assert!(err.contains(expect), "{err} should contain {expect}"); + } + } + + #[test] + fn a_feature_schema_with_no_feature_columns_is_refused() { + // Reachable only through a hand-written companion: an explicitly empty + // `rescore.features` list is already refused, and `FeatureSchema::read` refuses to + // reconstruct an empty list from the parquet. What the old validation did with it + // was an accident of `chunks_exact(n_features.max(1))` over an empty buffer: the + // loop ran zero times, which ALSO skipped the prelim_score/precursor_mz check for + // every row, and the run went on to train on no features at all. + let competed = scratch("nofeat_competed.parquet"); + crafted_competed_table(&competed, 8); + std::fs::write( + format!("{competed}.schema.json"), + r#"{"feature_columns":[],"schema_id":"empty-for-this-test"}"#, + ) + .unwrap(); + let (err, _, work) = run_streamed(&competed, "nofeat"); + std::fs::remove_file(format!("{competed}.schema.json")).ok(); + assert!( + err.contains("declares no feature columns"), + "an empty feature set must be refused by name: {err}" + ); + no_handoff_rubble(&work); + } + + #[test] + fn a_pass_that_disagrees_about_the_row_count_is_refused_with_both_counts() { + assert!(refuse_row_disagreement("the streamed handoff", 10, 10).is_ok()); + let err = refuse_row_disagreement("the feature matrix", 9, 10) + .unwrap_err() + .to_string(); + assert!( + err.contains("the feature matrix has 9 rows where 10 were expected"), + "{err}" + ); + assert!(err.contains("row-misaligns"), "{err}"); + } + #[test] fn a_shared_column_array_writes_the_same_parquet_as_two_copies() { use arrow::array::{ArrayRef, Int32Array, StringArray}; @@ -2271,16 +2568,29 @@ b // to the three aliases of the pooled q. This asserts the file that produces is // byte-for-byte the file `write_table` produced from separate, equal vectors -- // the output-equality claim of that change, on the mechanism it rests on. + // + // The row count is what makes the comparison capable of failing. `write_table` + // encodes in 65,536-row chunks and `write_scored_table` writes one batch, so below + // 65,536 rows both write exactly one chunk and no arrangement of the data could + // tell them apart -- the test asserted the new shape at a size where the two + // writers cannot diverge. Chunking is not free of that risk in general: mumdia-io's + // `an_entirely_null_column_keeps_its_rows_but_not_its_page_framing` shows a column + // type whose page framing DOES move at a chunk boundary. 3 chunks and a short last + // one is the shape that can catch it here, on the scored table's own duplicate-array + // arrangement and with a string column in it. The writer's 1,048,576-row row-group + // seam is the other place the two could part; mumdia-io pins that one directly, at + // 1,048,583 rows, in `the_row_groups_fall_in_the_same_places_past_the_row_group_maximum`. + let n = 3 * (1usize << 16) + 7; let dir = std::env::temp_dir().join("mumdia_scored_share_test"); std::fs::create_dir_all(&dir).unwrap(); let pid = std::process::id(); let a = dir.join(format!("copies_{pid}.parquet")); let b = dir.join(format!("shared_{pid}.parquet")); - let prot: Vec = (0..2000) + let prot: Vec = (0..n) .map(|i| format!("sp|P{:05}|PROT_HUMAN", i % 700)) .collect(); - let q: Vec = (0..2000).map(|i| (i % 97) as f64 / 97.0).collect(); - let rank: Vec = (0..2000).map(|i| i % 3).collect(); + let q: Vec = (0..n).map(|i| (i % 97) as f64 / 97.0).collect(); + let rank: Vec = (0..n).map(|i| (i % 3) as i32).collect(); write_table( a.to_str().unwrap(), vec![ @@ -2315,7 +2625,7 @@ b .unwrap(); let mut w = mumdia_io::table::BatchWriter::new(b.to_str().unwrap(), schema).unwrap(); w.write(&batch).unwrap(); - assert_eq!(w.close().unwrap(), 2000); + assert_eq!(w.close().unwrap(), n as u64); assert_eq!( std::fs::read(&a).unwrap(), std::fs::read(&b).unwrap(), From bfab6a12ff85a53597c5ea517bab0bb1afe9cc08 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:17:26 +0200 Subject: [PATCH 23/35] fix(quant,seed): the preconditions the two reviews found asserted rather than enforced Two adversarial reviews of the quant chromatogram store and the seed's window-width survey. Both found real defects; each fix is pinned by a test that fails when the fix is reverted, verified one fix at a time. quant.rs - `nearest_index`'s comment claimed any non-finite target "makes every predicate false and lands on 0". True for NaN and -inf, false for +inf: `r < inf` holds at every sample, so `partition_point` returned `rt.len()` and the LAST index came back where the scan returns 0. The precondition that actually holds is "finite", so it is now a guard returning the scan's answer, not a comment. Unreachable today (both call sites filter the apex hint on `is_finite`, and convert drops a spectrum whose scan start time is not finite), so no artifact moves. - `ChromStore::axis_for` deduplicated axes with slice `==`, and `-0.0 == 0.0` under f32's `PartialEq`, while `peak_window`'s merge keys the union on `to_bits`. A row whose axis differed from a stored one only in the sign of a zero was given the stored axis, folding two union points into one: different profile, apex, window and quantity. The comparison is now bitwise, so interning can only substitute an axis identical to the last bit. - `peak_window`'s `axis_sorted` guard asked whether the merged axis's last sample was `>= 0.0`, which `-0.0` satisfies although its bits sort after every positive f32. An axis ending in `-0.0` was declared ascending while descending at its last step, and `nearest_index` then binary-searched it, moving `ai` (the apex index that anchors `peak_bounds`). It now checks the order itself, one pass over an axis the sort already paid O(n log n) for. Neither -0.0 path is reachable from extract, whose grid is positive mzML scan times; both are reachable through `mumdia quant --chromatograms`, which is why the neighbouring rt/intensity length check exists. - `axis_for` could in principle mint the `NO_AXIS` (`u32::MAX`) sentinel, which every consumer reads as "no RT axis" and silently integrates to 0. Checked where the id is minted (one comparison per distinct axis); `ChromStore::push` and the loader now return `Result`. - `peak_window_both`'s claim to run every case through BOTH union constructions did not hold: `peak_window` takes the shared-axis fast path whenever the rows resolve to one strict axis id, which a single-row call always does whatever candidate id it was pushed under, so three of eight call sites compared the shared path against itself. The merged fixture now clears `axis_strict`, which withdraws the fast path without touching a stored value, and the helper asserts that routing. Reverting the fixture fails all six `peak_window_both` callers. search_seed.rs - ACCEPTED, the guard was weakened. Restricting the width survey to the in-band groups removed a 63x false positive and introduced a false negative: under `groups.window_groups` a full-range window resolves to the whole loaded band, and so does the band's own isolation window, so `widest > 8 * median` can never fire on a band. The malformed scan is now detected from the WINDOW instead: convert substitutes (0.0, 1e6) for a zero-width or missing isolation window, which reads the same on every band. A mixture of full-range and narrow windows warns; all full-range is a legitimate all-ion acquisition and stays silent. The width ratio is kept as the second, weaker diagnostic, with a comment saying what it can see under a grouped search. - ACCEPTED, the eager scratch sizing was a memory regression. "It grows on demand to the widest window it processes anyway" holds only for a leaf that processes a served group. `map_init` runs once per rayon LEAF, before the leaf knows which groups it drew, so in a banded search where 62 of 63 groups select nothing, nearly every leaf would allocate 16 B x band width and return empty: about 14 MB per worker on the profiled 54.8M-candidate library over 63 bands, roughly 445 MB across 32 workers, almost all never touched. `initial_scratch_width` now takes the served median only when a majority of groups is served; the leaf that does draw work still reaches the same size through `SeedScratch::ensure`, in one amortised realloc. Rejected - The reviews' remaining findings are all in `index.rs`, which is out of this change's scope and already addressed on the integration branch. - Not acted on, recorded instead: `peak_window`'s shared fast path indexes `prof` (sized from the AXIS) by the INTENSITY position, so its non-panicking behaviour rests on the loader's per-row length check. Making the loop `zip` would replace a panic with silent truncation, which is worse, so the invariant is asserted at `ChromStore::push` (debug only; the loader keeps the named, candidate-identifying error). - Neither review's output-equality reasoning is backed by an end-to-end artifact diff, and neither is this change. Every claim here is crafted-fixture evidence. The byte-identical contract still needs a before/after `blake3_file` of the quant tables and `seed_psms.parquet` on a real run, including one banded `groups.window_groups` run. cargo clippy -p mumdia --all-targets -D warnings clean; 283 lib + 5 bin + 2 + 3 integration tests pass. `cargo fmt --check` still reports pre-existing diffs in `pool.rs` and `mumdia-io/table.rs`, untouched here. Co-Authored-By: Claude Opus 5 (1M context) --- rust/mumdia/crates/mumdia/src/stages/quant.rs | 271 +++++++++++++++--- .../crates/mumdia/src/stages/search_seed.rs | 235 ++++++++++++--- 2 files changed, 439 insertions(+), 67 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/quant.rs b/rust/mumdia/crates/mumdia/src/stages/quant.rs index 00f45632..00f77f25 100644 --- a/rust/mumdia/crates/mumdia/src/stages/quant.rs +++ b/rust/mumdia/crates/mumdia/src/stages/quant.rs @@ -122,18 +122,33 @@ impl ChromStore { /// Store `rt` as an axis id, reusing an axis already stored for the open candidate /// when the values are identical (the common case: one grid per candidate). - fn axis_for(&mut self, cid: u32, rt: &[f32]) -> u32 { + fn axis_for(&mut self, cid: u32, rt: &[f32]) -> Result { if self.open_cid != Some(cid) { self.open_cid = Some(cid); self.open_axis_lo = self.axis_off.len() - 1; } if rt.is_empty() { - return NO_AXIS; + return Ok(NO_AXIS); } for a in self.open_axis_lo..self.axis_off.len() - 1 { let (lo, hi) = (self.axis_off[a], self.axis_off[a + 1]); - if hi - lo == rt.len() && self.axis_vals[lo..hi] == *rt { - return a as u32; + // BITWISE, not `==`. Interning makes every consumer read `store.rt(row)` + // instead of the row's own values, so it may only substitute an axis that is + // identical to the last bit. `-0.0 == 0.0` holds under f32's `PartialEq` + // while [`peak_window`]'s merge keys the union on `to_bits`, so an `==` match + // would fold two union points into one and move the profile, the apex, the + // walked window and the quantity. Extract's grid is positive mzML scan times, + // but `mumdia quant --chromatograms` takes a table written by anything, which + // is the same reason the rt/intensity length check below exists. + // (A NaN-carrying axis now dedups against an identical one, which `==` never + // did; the values stored are the same bits either way.) + if hi - lo == rt.len() + && self.axis_vals[lo..hi] + .iter() + .zip(rt) + .all(|(a, b)| a.to_bits() == b.to_bits()) + { + return Ok(a as u32); } } let nondecreasing = rt_is_sorted(rt); @@ -142,11 +157,34 @@ impl ChromStore { .push(nondecreasing && rt.windows(2).all(|w| w[0] < w[1])); self.axis_vals.extend_from_slice(rt); self.axis_off.push(self.axis_vals.len()); - (self.axis_off.len() - 2) as u32 - } - - fn push(&mut self, cid: u32, name: &str, pred: f32, rt: &[f32], inten: &[f32]) { - let axis_id = self.axis_for(cid, rt); + let id = (self.axis_off.len() - 2) as u32; + // [`NO_AXIS`] means "this row has no RT axis": [`ChromStore::axis`] hands back an + // empty slice for it, [`peak_window`]'s uniformity scan skips the row and + // [`fixed_window_indices`] returns `None`, so a real axis minted under that id + // would integrate to 0 without a word. One comparison per DISTINCT axis, at the + // point the id is minted, is cheaper than the argument that the count cannot get + // there -- and `--out-peak-bounds` is exactly the path that holds the whole table. + if id == NO_AXIS { + anyhow::bail!( + "chromatogram table has more than {NO_AXIS} distinct retention-time axes; \ + the axis id would collide with the empty-trace marker" + ); + } + Ok(id) + } + + fn push(&mut self, cid: u32, name: &str, pred: f32, rt: &[f32], inten: &[f32]) -> Result<()> { + // The shared-axis profile in [`peak_window`] is sized from the AXIS and indexed by + // the INTENSITY position, and `fixed_window_indices` slices the intensities with + // indices taken from the RT trace. `run` refuses a mismatched row while it can + // still name the candidate and the file; this catches a fixture or a future loader + // that does not. + debug_assert_eq!( + rt.len(), + inten.len(), + "chromatogram row for candidate_id {cid} has mismatched rt/intensity lengths" + ); + let axis_id = self.axis_for(cid, rt)?; let name_id = self.names.intern(name); self.cid.push(cid); self.name_id.push(name_id); @@ -154,6 +192,7 @@ impl ChromStore { self.axis_id.push(axis_id); self.int_vals.extend_from_slice(inten); self.int_off.push(self.int_vals.len()); + Ok(()) } fn axis(&self, id: u32) -> &[f32] { @@ -305,8 +344,20 @@ fn nearest_index(rt: &[f32], target: f64, sorted: bool) -> usize { } return k; } - // First sample at or after `target`. A non-finite target makes every predicate false - // and lands on 0, which is where the scan's untouched `k` sits too. + // The binary search below needs a FINITE target, and the earlier claim that any + // non-finite one "makes every predicate false and lands on 0" is wrong for `+inf`: + // `r < inf` holds at every sample, so `partition_point` returns `rt.len()` and the + // last index comes back, where the scan returns 0 (its `d` is `inf` everywhere and + // `inf < inf` is false, so `k` is never assigned). NaN and `-inf` do land on 0 by + // themselves; the precondition that actually holds is "finite", so it is stated once + // here instead of argued per case. Retention times reach this stage finite (convert + // drops a spectrum whose scan start time is not, `convert.rs`), and both call sites + // filter the apex hint on `is_finite`, so this is the guard for the third caller. + // Returning 0 is what the scan returns for every non-finite target. + if !target.is_finite() { + return 0; + } + // First sample at or after `target`. let j = rt.partition_point(|&r| (r as f64) < target); // Backing up over an equal-valued run keeps the FIRST of several identical samples, // which is the one the scan's strict `<` kept. @@ -688,10 +739,18 @@ fn peak_window( let apex = axis.first().map_or(f64::NAN, |&r| r as f64); return (f64::NEG_INFINITY, f64::INFINITY, apex); } - // Both constructions leave the axis ascending: the shared one is checked strictly - // increasing when it is stored, and the merged one is deduped and ordered by bit - // pattern, which is value order as long as no negative or NaN RT is present. - let axis_sorted = shared.is_some() || axis.last().is_some_and(|&r| r >= 0.0); + // [`nearest_index`] may binary-search only an ascending axis. The shared one was + // checked strictly increasing when it was stored ([`ChromStore::axis_strict`]). The + // merged one is ordered by BIT PATTERN, which is value order only while no RT is + // negative or NaN -- and the previous test, "the last sample is `>= 0.0`", passes for + // `-0.0`, whose bits sort after every positive f32. An axis ending in `-0.0` was + // therefore declared sorted while descending at its last step, and the search and the + // scan then disagreed on `ai`, which is the apex index that sizes the integration + // window. Check the order itself instead: one pass over an axis the sort above already + // paid O(n log n) for, and it needs no argument about the bit layout. Strict `<` is + // right because the merge already deduped by bits, so the only way two neighbours can + // compare equal is `0.0` against `-0.0`. + let axis_sorted = shared.is_some() || axis.windows(2).all(|w| w[0] < w[1]); let ai = if let Some(hint) = apex_hint.filter(|v| v.is_finite()) { // The identified apex need not exactly equal a chromatogram sample (for // example after serialization/calibration), so anchor to the nearest RT. @@ -1108,7 +1167,7 @@ pub fn run(p: QuantParams) -> Result<(u64, u64)> { Some(a) => a.value(k), None => 0.0, }; - store.push(c, nm, pred, rt, it); + store.push(c, nm, pred, rt, it)?; } } } @@ -1886,19 +1945,26 @@ mod tests { fn store_of(rt: &[Vec], inten: &[Vec]) -> ChromStore { let mut s = ChromStore::new(); for (i, (r, it)) in rt.iter().zip(inten).enumerate() { - s.push(0, &format!("y{i}"), 0.0, r, it); + s.push(0, &format!("y{i}"), 0.0, r, it).unwrap(); } s } - /// The same rows under a DISTINCT candidate id each, so no two rows can share an axis - /// id and [`peak_window`] must merge the samples instead of reading one shared axis. - /// The row indices are passed to `peak_window` explicitly, so the candidate ids only - /// steer the store's axis dedup. + /// The same rows, forced through [`peak_window`]'s MERGED union construction. + /// + /// Distinct candidate ids alone do not force it. They stop the axis dedup, so each row + /// gets its own axis id, but the uniformity scan then finds exactly ONE axis id + /// whenever `rows.len() == 1` and takes the shared-axis fast path anyway -- which made + /// every single-row `peak_window_both` case compare the shared path against itself. + /// Clearing `axis_strict` withdraws the fast path's permission without touching a + /// stored value, so the merge runs over exactly the same samples for any row count. fn store_of_unshared(rt: &[Vec], inten: &[Vec]) -> ChromStore { let mut s = ChromStore::new(); for (i, (r, it)) in rt.iter().zip(inten).enumerate() { - s.push(i as u32, &format!("y{i}"), 0.0, r, it); + s.push(i as u32, &format!("y{i}"), 0.0, r, it).unwrap(); + } + for strict in s.axis_strict.iter_mut() { + *strict = false; } s } @@ -1907,6 +1973,9 @@ mod tests { /// The shared-axis path is an accumulation into one profile array; the merged path is /// the sort over every sample that replaced the two BTreeMaps. Every `peak_window` /// test below goes through this, so each is also an equality test between the paths. + /// + /// The routing is asserted rather than assumed: see [`store_of_unshared`] for why + /// distinct candidate ids were not enough on their own. fn peak_window_both( rows: &[usize], rt: &[Vec], @@ -1915,8 +1984,13 @@ mod tests { grace: usize, hint: Option, ) -> (f64, f64, f64) { + let merged_store = store_of_unshared(rt, inten); + assert!( + merged_store.axis_strict.iter().all(|&s| !s), + "the merged fixture must refuse `peak_window`'s shared-axis fast path" + ); let shared = peak_window(rows, &store_of(rt, inten), frac, grace, hint); - let merged = peak_window(rows, &store_of_unshared(rt, inten), frac, grace, hint); + let merged = peak_window(rows, &merged_store, frac, grace, hint); let bits = |w: (f64, f64, f64)| (w.0.to_bits(), w.1.to_bits(), w.2.to_bits()); assert_eq!( bits(shared), @@ -2743,16 +2817,142 @@ mod tests { } } + #[test] + fn nearest_index_returns_the_scans_answer_for_a_non_finite_target() { + // The comment that used to stand in for a guard claimed any non-finite target + // "makes every predicate false and lands on 0". That holds for NaN and for -inf + // and is FALSE for +inf: `r < inf` is true at every sample, so `partition_point` + // returned `rt.len()` and the search handed back the LAST index where the scan + // hands back 0 (its `d` is `inf` everywhere and `inf < inf` never fires, so `k` is + // never assigned). Both call sites filter the apex hint on `is_finite` today, so + // no artifact moved; the precondition is now enforced rather than asserted. + let traces: [&[f32]; 5] = [ + &[], + &[5.0], + &[1.0, 2.0], + &[0.0, 1.0, 2.0, 3.0, 4.0], + &[2.0, 2.0, 2.0], + ]; + for rt in traces { + for t in [f64::INFINITY, f64::NEG_INFINITY, f64::NAN] { + assert_eq!( + nearest_index(rt, t, true), + nearest_index(rt, t, false), + "rt={rt:?} target={t}" + ); + assert_eq!( + nearest_index(rt, t, true), + 0, + "the scan leaves `k` at 0 for every non-finite target: rt={rt:?} target={t}" + ); + } + } + } + + #[test] + fn an_axis_differing_only_in_the_sign_of_a_zero_is_a_distinct_axis() { + // Interning made every consumer read `store.rt(row)` instead of the row's own + // values, and `-0.0 == 0.0` under f32's `PartialEq`, so a slice comparison handed + // the second row the FIRST row's axis. `peak_window`'s merge keys the union on + // `to_bits`, where `-0.0` and `0.0` are two points, so the substitution folded the + // union: profile [4.0, 6.0] over axis [0.0, 1.0] instead of [1.0, 6.0, 3.0] over + // [0.0, 1.0, -0.0]. Extract writes positive mzML scan times, but + // `mumdia quant --chromatograms` takes a table written by anything. + let mut s = ChromStore::new(); + s.push(0, "y1", 0.0, &[0.0, 1.0], &[1.0, 2.0]).unwrap(); + s.push(0, "y2", 0.0, &[-0.0, 1.0], &[3.0, 4.0]).unwrap(); + assert_ne!(s.axis_id[0], s.axis_id[1]); + assert_eq!( + s.rt(1)[0].to_bits(), + (-0.0f32).to_bits(), + "the row must read back the bits the table carried" + ); + // An axis that really is identical still interns to one copy, which is what the + // store exists for. + s.push(0, "y3", 0.0, &[0.0, 1.0], &[5.0, 6.0]).unwrap(); + assert_eq!(s.axis_id[0], s.axis_id[2]); + + // And the fold moved the apex, not just the profile: with the `-0.0` sample the + // brightest, the substitution summed it onto the `+0.0` point and put the apex on + // the wrong sample. + let mut moved = ChromStore::new(); + moved.push(0, "y1", 0.0, &[0.0, 1.0], &[1.0, 2.0]).unwrap(); + moved.push(0, "y2", 0.0, &[-0.0, 1.0], &[9.0, 0.0]).unwrap(); + let (_, _, apex) = peak_window(&[0, 1], &moved, 0.5, 0, None); + assert_eq!( + apex.to_bits(), + (-0.0f64).to_bits(), + "apex must land on the -0.0 sample (interning gave +0.0), got {apex}" + ); + } + + #[test] + fn a_merged_axis_ending_in_negative_zero_is_not_binary_searched() { + // The merged union is ordered by BIT PATTERN, and `-0.0`'s bits sort after every + // positive f32, so an axis that ends in `-0.0` descends at its last step. The old + // guard asked only whether the last sample was `>= 0.0`, which `-0.0` satisfies, + // so the axis was declared sorted and `nearest_index` binary-searched a descending + // trace. `ai` is the apex index that anchors `peak_bounds` and the returned + // window, so the disagreement is a different quantity. + let rt = [vec![1.0f32, 2.0], vec![-0.0f32, 3.0]]; + let inten = [vec![1.0f32, 1.0], vec![1.0f32, 1.0]]; + let store = store_of(&rt, &inten); + // Union axis, in bit order: [1.0, 2.0, 3.0, -0.0]. + assert_ne!(store.axis_id[0], store.axis_id[1]); + + // Below every sample the scan picks `-0.0` at index 3; the binary search returned + // index 0 because `partition_point` found nothing smaller than the target. + let (_, _, apex_low) = peak_window(&[0, 1], &store, 0.5, 0, Some(-1.0)); + assert_eq!( + apex_low.to_bits(), + (-0.0f64).to_bits(), + "expected the -0.0 sample, got {apex_low}" + ); + // Above every sample the scan picks 3.0 at index 2; the binary search ran off the + // end and returned index 3, which is the `-0.0`. + let (_, _, apex_high) = peak_window(&[0, 1], &store, 0.5, 0, Some(100.0)); + assert_eq!(apex_high, 3.0, "expected the 3.0 sample, got {apex_high}"); + } + + #[test] + fn the_merged_fixture_refuses_the_shared_axis_fast_path_even_for_one_row() { + // `peak_window` takes the shared-axis fast path whenever the rows resolve to ONE + // axis id that is marked strictly increasing, and a single-row call always does, + // whatever candidate id the row was pushed under. Distinct candidate ids therefore + // did NOT make the merged fixture merge, and the three single-row + // `peak_window_both` cases were comparing the shared path against itself. + let rt = [vec![0.0f32, 1.0, 2.0]]; + let inten = [vec![1.0f32, 5.0, 1.0]]; + let plain = store_of(&rt, &inten); + assert!( + plain.axis_strict[plain.axis_id[0] as usize], + "the fixture's axis is strictly increasing, so the fast path is available" + ); + let merged = store_of_unshared(&rt, &inten); + assert!( + merged.axis_strict.iter().all(|&s| !s), + "clearing `axis_strict` is what actually withdraws the fast path" + ); + // Which is the equality `peak_window_both` claims to be testing. + let bits = |w: (f64, f64, f64)| (w.0.to_bits(), w.1.to_bits(), w.2.to_bits()); + assert_eq!( + bits(peak_window(&[0], &plain, 0.5, 0, None)), + bits(peak_window(&[0], &merged, 0.5, 0, None)) + ); + } + #[test] fn the_chromatogram_store_returns_the_rows_it_was_given() { let grid: Vec = (0..5).map(|k| k as f32).collect(); let offset: Vec = (0..5).map(|k| k as f32 + 0.5).collect(); let mut s = ChromStore::new(); - s.push(7, "y1", 0.25, &grid, &[1.0, 2.0, 3.0, 2.0, 1.0]); - s.push(7, "y2", 0.75, &grid, &[0.0, 1.0, 2.0, 1.0, 0.0]); - s.push(7, "b3", 0.5, &offset, &[9.0; 5]); - s.push(7, "b4", 0.1, &[], &[]); - s.push(8, "y1", 0.6, &grid, &[4.0; 5]); + s.push(7, "y1", 0.25, &grid, &[1.0, 2.0, 3.0, 2.0, 1.0]) + .unwrap(); + s.push(7, "y2", 0.75, &grid, &[0.0, 1.0, 2.0, 1.0, 0.0]) + .unwrap(); + s.push(7, "b3", 0.5, &offset, &[9.0; 5]).unwrap(); + s.push(7, "b4", 0.1, &[], &[]).unwrap(); + s.push(8, "y1", 0.6, &grid, &[4.0; 5]).unwrap(); assert_eq!(s.nrows(), 5); assert_eq!(s.rt(0), grid.as_slice()); assert_eq!(s.inten(1), &[0.0, 1.0, 2.0, 1.0, 0.0]); @@ -2777,7 +2977,9 @@ mod tests { #[test] fn an_unsorted_or_duplicated_axis_refuses_the_fast_paths() { let mut descending = ChromStore::new(); - descending.push(1, "y1", 0.0, &[3.0, 1.0, 2.0], &[1.0, 2.0, 3.0]); + descending + .push(1, "y1", 0.0, &[3.0, 1.0, 2.0], &[1.0, 2.0, 3.0]) + .unwrap(); assert!( !descending.rt_sorted, "a descending step must send every nearest-sample search back to the scan" @@ -2785,7 +2987,9 @@ mod tests { assert!(!descending.axis_strict[0]); let mut duplicated = ChromStore::new(); - duplicated.push(1, "y1", 0.0, &[1.0, 1.0, 2.0], &[1.0, 2.0, 3.0]); + duplicated + .push(1, "y1", 0.0, &[1.0, 1.0, 2.0], &[1.0, 2.0, 3.0]) + .unwrap(); assert!( duplicated.rt_sorted, "non-decreasing is all the nearest-sample search needs" @@ -2797,7 +3001,9 @@ mod tests { ); let mut negative = ChromStore::new(); - negative.push(1, "y1", 0.0, &[-1.0, 2.0], &[1.0, 2.0]); + negative + .push(1, "y1", 0.0, &[-1.0, 2.0], &[1.0, 2.0]) + .unwrap(); assert!( !negative.rt_sorted, "bit order is value order only above zero" @@ -2842,7 +3048,8 @@ mod tests { ] { let mut s = ChromStore::new(); for (r, &c) in cids.iter().enumerate() { - s.push(c, &format!("y{r}"), r as f32, &[r as f32], &[1.0]); + s.push(c, &format!("y{r}"), r as f32, &[r as f32], &[1.0]) + .unwrap(); } let index = CandIndex::build(&s); let mut want: BTreeMap> = BTreeMap::new(); diff --git a/rust/mumdia/crates/mumdia/src/stages/search_seed.rs b/rust/mumdia/crates/mumdia/src/stages/search_seed.rs index 925c9b71..e9b6b938 100644 --- a/rust/mumdia/crates/mumdia/src/stages/search_seed.rs +++ b/rust/mumdia/crates/mumdia/src/stages/search_seed.rs @@ -424,32 +424,61 @@ fn label_column(is_decoy: &[bool]) -> Vec { .collect() } -/// Median and maximum candidate-window width over the window groups THIS search actually -/// serves, from each group's `(lo, hi)` candidate range. +/// The `(lower_mz, upper_mz)` `convert` writes for an MS2 scan whose reported isolation +/// window is zero-width or whose precursor is missing (`convert.rs`, "AIF / all-ion"). It is +/// a property of the SCAN, so it reads the same on every band of a grouped search. +const FULL_RANGE_WINDOW: (f64, f64) = (0.0, 1.0e6); + +/// True when the run mixes the full-range isolation window with narrow ones, which is the +/// signature of a scan that lost its isolation window: [`FULL_RANGE_WINDOW`] is what +/// `convert` substitutes. Every window full-range is a legitimate all-ion acquisition and +/// none is an ordinary run, so only the mixture is worth a warning. /// -/// Size each worker's scratch to the MEDIAN candidate window, not the widest one and -/// certainly not the whole library. The scratch is indexed window-relative, and it grows on -/// demand, so an underestimate costs a few reallocations while an overestimate costs -/// 16 B x width per rayon worker up front. +/// This test, unlike the candidate-width ratio below, reads the same under a grouped search. +/// A full-range window resolves to the WHOLE loaded band, and so does the band's own +/// isolation window, so on a band the two widths are equal and no ratio can separate them. +fn has_stray_full_range_window(windows: &[(f64, f64)]) -> bool { + let n = windows.iter().filter(|&&w| w == FULL_RANGE_WINDOW).count(); + n > 0 && n < windows.len() +} + +/// What the per-worker scratch sizing and the width warning read off the window groups. +struct WindowSurvey { + /// Median candidate-window width over the SERVED groups; 1 when none is served. + median: usize, + /// Widest served candidate window, or `None` when no group is served. + widest: Option, + /// Groups this search serves, of `total`. + served: usize, + total: usize, +} + +/// Candidate-window widths over the window groups THIS search actually serves, from each +/// group's `(lo, hi)` candidate range. +/// +/// The MEDIAN is the sizing statistic, not the widest and certainly not the library. The +/// scratch is indexed window-relative and grows on demand, so an underestimate costs a +/// reallocation while an overestimate costs 16 B x width per rayon worker up front. /// /// The max is the wrong statistic because one window can be the whole library: -/// `convert.rs` maps BOTH a zero-width reported isolation window and a missing precursor to -/// the full range (0, 1e6), which `candidate_range` resolves to every candidate. So a single -/// malformed or all-ion scan in an otherwise 50-window run sized every worker's scratch to -/// the library -- 877 MB per worker on the profiled 54.8M-candidate library, about 28 GB of -/// commit charge on 32 cores, for arrays a worker inside one narrow window never touches. +/// [`FULL_RANGE_WINDOW`] resolves to every candidate. So a single malformed or all-ion scan +/// in an otherwise 50-window run sized every worker's scratch to the library -- 877 MB per +/// worker on the profiled 54.8M-candidate library, about 28 GB of commit charge on 32 cores, +/// for arrays a worker inside one narrow window never touches. +/// +/// Groups that select nothing (`hi <= lo`) are EXCLUDED from the widths, which is the same +/// test the worker uses to skip a group, and counted instead: `served` against `total`. +/// Under an isolation-window-group search the library is one m/z band while the run's +/// windows are all of them, so with 63 bands 62 of every 63 groups select nothing and +/// contributed a width of 1. The median was then 1 on every band and `widest > 8 * median` +/// fired on every band, printing a warning about a zero-width or missing isolation window +/// once per band for a run that has no such scan. /// -/// Groups that select nothing (`hi <= lo`) are EXCLUDED, which is the same test the worker -/// uses to skip a group. Under an isolation-window-group search the library is one m/z band -/// while the run's windows are all of them, so with 63 bands 62 of every 63 groups select -/// nothing and contribute a width of 1: the median was then 1 on every band, the scratch was -/// sized to a single slot and grown on first use, and `widest > 8 * median` fired on every -/// band, printing a warning about a zero-width or missing isolation window once per band for -/// a run that has no such scan. Restricted to the in-band groups, both the sizing and the -/// warning describe this search. +/// `served`/`total` is what keeps the correction from becoming a memory regression: see +/// [`initial_scratch_width`]. /// /// Widths are collected in group order and then sorted, so the value is deterministic. -fn window_width_stats(ranges: &[(u32, u32)]) -> (usize, Option) { +fn window_survey(ranges: &[(u32, u32)]) -> WindowSurvey { let mut widths: Vec = ranges .iter() .filter(|(lo, hi)| hi > lo) @@ -457,7 +486,37 @@ fn window_width_stats(ranges: &[(u32, u32)]) -> (usize, Option) { .collect(); widths.sort_unstable(); let median = widths.get(widths.len() / 2).copied().unwrap_or(1).max(1); - (median, widths.last().copied()) + WindowSurvey { + median, + widest: widths.last().copied(), + served: widths.len(), + total: ranges.len(), + } +} + +/// Width each rayon leaf's [`SeedScratch`] is allocated with before it knows which groups it +/// drew. +/// +/// `map_init` runs its initialiser once per LEAF, not once per thread, and it runs before +/// the leaf sees a group, so a leaf that draws only unserved groups pays the allocation and +/// then returns empty. Where every group is served -- an ungrouped search -- the eager size +/// is what the leaf's first `accumulate` would have grown it to anyway, so sizing it up +/// front costs nothing and saves the reallocation. Where most groups are NOT served it is +/// pure waste, and that is exactly the banded search this survey was corrected for: the +/// library is one m/z band while the scans are all of them, so 62 of 63 groups are unserved +/// and eager sizing would commit 16 B x band width on every leaf (about 14 MB per worker on +/// the profiled 54.8M-candidate library over 63 bands, roughly 445 MB across 32 workers) +/// for arrays nearly all of them never touch. +/// +/// So: the served median when a majority of groups is served, one slot otherwise. The leaf +/// that does draw work still reaches the same size through `SeedScratch::ensure`, in one +/// amortised reallocation. +fn initial_scratch_width(survey: &WindowSurvey) -> usize { + if survey.served * 2 > survey.total { + survey.median + } else { + 1 + } } /// fragindex seed over isolation-window groups, in parallel. Each scan belongs to @@ -486,36 +545,62 @@ fn seed_fragindex_windows( .push(si); } let group_vec: Vec> = groups.into_values().collect(); - let ranges: Vec<(u32, u32)> = group_vec + let group_windows: Vec<(f64, f64)> = group_vec .iter() .filter_map(|ids| ids.first()) .map(|&si| { let w = &scans[si].window; - idx.candidate_range(w.lower_mz, w.upper_mz) + (w.lower_mz, w.upper_mz) }) .collect(); - let (median_window_width, widest) = window_width_stats(&ranges); - // Say so when the two disagree by a lot: a window spanning most of the library is - // either a genuine all-ion acquisition or a malformed scan, and both are worth - // knowing about before wondering why extraction is slow. - if let Some(widest) = widest { - if widest > 8 * median_window_width { + let ranges: Vec<(u32, u32)> = group_windows + .iter() + .map(|&(lo, hi)| idx.candidate_range(lo, hi)) + .collect(); + let survey = window_survey(&ranges); + // Two diagnostics, in order of how much each actually establishes. + // + // 1. A scan that lost its isolation window is recognisable from the WINDOW: convert + // substitutes the full m/z range, so a run mixing full-range groups with narrow ones + // contains such a scan whatever library happens to be loaded. This is the test that + // survives a grouped search, where the ratio below cannot fire at all -- a + // full-range window resolves to the whole loaded band, and so does the band's own + // isolation window, so on a band the two widths are equal. + // 2. Otherwise, a window covering far more candidates than the median is still worth a + // word even when the window itself is well formed (an unusually wide quadrupole + // setting, or a library whose m/z range sits inside one window), before anyone + // wonders why extraction is slow. Under a grouped search this second test sees only + // the in-band groups, so it reports on this band, which is all it can honestly say. + if has_stray_full_range_window(&group_windows) { + warn!( + full_range_windows = group_windows + .iter() + .filter(|&&w| w == FULL_RANGE_WINDOW) + .count(), + window_groups = group_windows.len(), + "search-seed: some isolation windows are the full m/z range that convert \ + writes for a scan with a zero-width or missing isolation window, while others \ + are narrow. An all-ion acquisition makes EVERY window full-range; a mixture \ + means those scans lost their isolation window" + ); + } else if let Some(widest) = survey.widest { + if widest > 8 * survey.median { warn!( - median_window_candidates = median_window_width, + median_window_candidates = survey.median, widest_window_candidates = widest, library_candidates = idx.n_cand(), "search-seed: one isolation window covers far more candidates than the \ - median. An all-ion acquisition does this legitimately; otherwise check \ - for a scan with a zero-width or missing isolation window, which convert \ - maps to the full m/z range" + median of the window groups this search serves. An all-ion acquisition \ + does this legitimately" ); } } + let scratch_width = initial_scratch_width(&survey); let partials: Vec> = group_vec .par_iter() .map_init( - || SeedScratch::new(median_window_width), + || SeedScratch::new(scratch_width), |scratch, ids| { if ids.is_empty() { return Vec::new(); @@ -600,7 +685,15 @@ fn hyperscore(matched: u32, sum_obs: f64) -> f64 { #[cfg(test)] mod survey_tests { - use super::{label_column, window_width_stats}; + use super::{ + has_stray_full_range_window, initial_scratch_width, label_column, window_survey, + FULL_RANGE_WINDOW, + }; + + fn window_width_stats(ranges: &[(u32, u32)]) -> (usize, Option) { + let s = window_survey(ranges); + (s.median, s.widest) + } #[test] fn window_width_stats_ignores_the_groups_this_band_does_not_serve() { @@ -633,6 +726,78 @@ mod survey_tests { assert_eq!(window_width_stats(&[]), (1, None)); } + #[test] + fn a_band_that_serves_one_group_of_many_is_not_sized_from_the_band_width() { + // `map_init` runs its initialiser once per rayon LEAF and before the leaf knows + // which groups it drew, so sizing the scratch from the served median would commit + // 16 B x band width on every leaf while only the leaf holding the one in-band group + // ever touches it. The correction to the median must not turn into that. + let mut banded = vec![(0u32, 0u32); 62]; + banded.push((1000, 871_000)); + let survey = window_survey(&banded); + assert_eq!((survey.served, survey.total), (1, 63)); + assert_eq!( + survey.median, 870_001, + "the warning still reads the band width" + ); + assert_eq!( + initial_scratch_width(&survey), + 1, + "62 of 63 leaves would allocate a band they never probe" + ); + + // An ungrouped search serves every group, and there the eager size is what the + // leaf's first `accumulate` would have grown it to anyway. + let plain = window_survey(&[(0, 10), (100, 110), (50, 60)]); + assert_eq!((plain.served, plain.total), (3, 3)); + assert_eq!(initial_scratch_width(&plain), 11); + // One dead window among live ones does not tip the majority. + let mostly = window_survey(&[(0, 10), (100, 110), (7, 7), (50, 60)]); + assert_eq!(initial_scratch_width(&mostly), 11); + // Nothing served at all: one slot, and `SeedScratch::ensure` never runs. + assert_eq!(initial_scratch_width(&window_survey(&[])), 1); + } + + #[test] + fn a_stray_full_range_window_is_flagged_from_the_window_and_not_from_the_width() { + // convert maps a zero-width or missing isolation window to the full m/z range, so + // the malformed scan is visible in the window itself. The candidate-width ratio + // that used to carry this warning cannot see it under a grouped search: the + // full-range window resolves to the WHOLE loaded band, and so does the band's own + // isolation window, so the two served widths are equal and `widest > 8 * median` + // is false. The band below is 5,001 candidates wide and serves exactly those two + // groups. + let banded_ranges = [(1000u32, 6000u32), (1000, 6000)]; + let banded = window_survey(&banded_ranges); + assert_eq!(banded.widest, Some(5001)); + assert!( + banded.widest.unwrap() <= 8 * banded.median, + "the width ratio is silent on a band, which is the false negative" + ); + // The window test is not, and it does not consult the library at all. + let mut banded_windows = vec![(500.0f64, 510.0f64); 62]; + banded_windows.push(FULL_RANGE_WINDOW); + assert!(has_stray_full_range_window(&banded_windows)); + + // A genuine all-ion acquisition is every window full-range: legitimate, silent. + assert!(!has_stray_full_range_window(&[ + FULL_RANGE_WINDOW, + FULL_RANGE_WINDOW + ])); + // An ordinary run has none. + assert!(!has_stray_full_range_window(&[ + (500.0, 510.0), + (510.0, 520.0) + ])); + assert!(!has_stray_full_range_window(&[])); + // A wide-but-not-full window is not this warning's business; the width ratio keeps + // it. + assert!(!has_stray_full_range_window(&[ + (500.0, 510.0), + (0.0, 2000.0) + ])); + } + #[test] fn label_column_is_the_boolean_the_old_code_parsed_back_out_of_it() { // The artifact text is unchanged, and the boolean the stage uses is now its source From bccf710b2505b22f62f59a156cb3395eacec20eb Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:17:33 +0200 Subject: [PATCH 24/35] style: run cargo fmt over the pool splice and the pool subcommand Co-Authored-By: Claude Opus 5 (1M context) --- rust/mumdia/crates/mumdia-io/src/table.rs | 15 ++++++++++--- rust/mumdia/crates/mumdia/src/main.rs | 4 +--- rust/mumdia/crates/mumdia/src/stages/pool.rs | 23 +++++++++++++++----- 3 files changed, 30 insertions(+), 12 deletions(-) diff --git a/rust/mumdia/crates/mumdia-io/src/table.rs b/rust/mumdia/crates/mumdia-io/src/table.rs index 88673baf..11d4f6fc 100644 --- a/rust/mumdia/crates/mumdia-io/src/table.rs +++ b/rust/mumdia/crates/mumdia-io/src/table.rs @@ -2058,11 +2058,20 @@ mod write_chunking_tests { let (chunked, once) = (p("big_chunked.parquet"), p("big_once.parquet")); write_table(&chunked, cols(n)).unwrap(); write_one_batch(&once, cols(n)); - let a = TableFile::open(&chunked).unwrap().row_group_stats("id").unwrap(); - let b = TableFile::open(&once).unwrap().row_group_stats("id").unwrap(); + let a = TableFile::open(&chunked) + .unwrap() + .row_group_stats("id") + .unwrap(); + let b = TableFile::open(&once) + .unwrap() + .row_group_stats("id") + .unwrap(); let rows_a: Vec = a.iter().map(|g| g.rows).collect(); let rows_b: Vec = b.iter().map(|g| g.rows).collect(); - assert_eq!(rows_a, rows_b, "row groups moved: {rows_a:?} against {rows_b:?}"); + assert_eq!( + rows_a, rows_b, + "row groups moved: {rows_a:?} against {rows_b:?}" + ); assert_eq!(rows_a, vec![1_048_576, 7]); assert_eq!( std::fs::read(&chunked).unwrap(), diff --git a/rust/mumdia/crates/mumdia/src/main.rs b/rust/mumdia/crates/mumdia/src/main.rs index fbb5566b..6ed87c2c 100644 --- a/rust/mumdia/crates/mumdia/src/main.rs +++ b/rust/mumdia/crates/mumdia/src/main.rs @@ -1349,9 +1349,7 @@ fn real_main() -> Result<()> { .filter(|b| std::path::Path::new(&b.competed).exists()) .collect(); if bands.is_empty() { - anyhow::bail!( - "no band directory under {groups_dir} holds a psms_competed.parquet" - ); + anyhow::bail!("no band directory under {groups_dir} holds a psms_competed.parquet"); } let out = |n: &str| format!("{out_dir}/{n}"); let (op, oc, ok) = ( diff --git a/rust/mumdia/crates/mumdia/src/stages/pool.rs b/rust/mumdia/crates/mumdia/src/stages/pool.rs index e8b0f25d..0e7ac4a3 100644 --- a/rust/mumdia/crates/mumdia/src/stages/pool.rs +++ b/rust/mumdia/crates/mumdia/src/stages/pool.rs @@ -280,7 +280,9 @@ mod tests { // to carry over from the tables it was spliced from. Col::LargeListF32( "trace".into(), - cids.iter().map(|c| vec![*c as f32, *c as f32 + 0.5]).collect(), + cids.iter() + .map(|c| vec![*c as f32, *c as f32 + 0.5]) + .collect(), ), ]; if extra { @@ -367,8 +369,16 @@ mod tests { let cids: Vec = ids.collect(); let n = cids.len(); let schema = std::sync::Arc::new(arrow::datatypes::Schema::new(vec![ - arrow::datatypes::Field::new("candidate_id", arrow::datatypes::DataType::UInt32, false), - arrow::datatypes::Field::new("prelim_score", arrow::datatypes::DataType::Float64, false), + arrow::datatypes::Field::new( + "candidate_id", + arrow::datatypes::DataType::UInt32, + false, + ), + arrow::datatypes::Field::new( + "prelim_score", + arrow::datatypes::DataType::Float64, + false, + ), ])); let mut w = BatchWriter::with_row_group_rows(&p, schema.clone(), 50).unwrap(); for chunk in cids.chunks(25) { @@ -376,9 +386,10 @@ mod tests { schema.clone(), vec![ std::sync::Arc::new(UInt32Array::from(chunk.to_vec())), - std::sync::Arc::new(arrow::array::Float64Array::from( - vec![score; chunk.len()], - )), + std::sync::Arc::new(arrow::array::Float64Array::from(vec![ + score; + chunk.len() + ])), ], ) .unwrap(); From faec501b05dbcb7756e2b4d75b156e6b93b41c31 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:20:45 +0200 Subject: [PATCH 25/35] fix(features): act on the review of the flat-buffer/writer-thread change An adversarial review of 1ac1b44 raised seven findings against the features stage. Five are real and fixed here, one is rejected on arithmetic, and one is answered with the evidence it said was missing. The stage's outputs are unchanged: a golden captured from 1ac1b44^ now proves that in the suite. FIXED - The confident-bounds pass re-opened the chromatogram file and re-parsed its whole footer once per span (`TableFile::open_rows`), and `ChromStream::open` opened it a second time for `column_names(path)`. Both are gone: `confident_global_bounds` takes the caller's already-open `&TableFile` and uses `TableFile::span` (an Arc clone of the parsed footer, and the call that method exists for and had no callers), and the stream answers "does this table carry frag_obs_mz?" from the handle's schema, which a span carries in full. The cost used to grow with the saving: the pruning's whole purpose is to produce many spans, and each one cost two footer parses. `the_chrom_stream_reads_its_optional_column_from_the_handle_it_is_given` pins the old answer across has-column x whole-file/span. - The PIN and the features table were committed in the opposite order. The in-line write ran `pin.finish()` then `writer.close()`, so a PIN that failed at the flush left the previous features table untouched; moving `close` onto the writer thread silently reversed it, publishing the table and then erroring. The PIN is now flushed inside the scope, before the commit marker, so a failure returns without the marker and the writer drops unpublished. `a_failed_pin_finish_does_not_publish_the_features_table` pins the old ordering through a `PinFinish` parameter (cfg(test) variant, not a global, so it cannot interfere with another test in the binary). - The pruning fixture wrote a fixed 6 rows per candidate into 12-row row groups, so no candidate ever straddled a boundary and the safety argument the pruning rests on was never exercised. The fixture now varies the fragment count (4..9); the test asserts that candidates do straddle boundaries, that a skipped group sits between two kept ones, and that more than one span comes back, and it adds a chunk size larger than a row group. - The claimed benefit of the pruning is a property of that fixture. Measured: extract writes 1<<16-row groups, so one group spans thousands of candidates while the confident set (~21,856 on AIF) is spread over the id range, and essentially every group holds an anchor. `row_group_pruning_saves_nothing_on_a _production_shaped_table` builds that distribution and asserts the pruning returns ONE span covering the whole table. The doc comment now says so instead of quoting the fixture's 40%, and says what it does buy (a clustered confident set reads 3 of 200 groups). Kept, not dropped: it is exact, it now costs no extra file open, and the sparse case is real. - `in_flight_peak` asserted `2 * max_matrix_bytes`, which was the in-line write's peak (one matrix plus its copy). With the writer thread there is no copy but the previous chunk's columns are held while the next chunk's matrix AND its extended buffer are live, so the overlap is nearer 2.9x at Extended -- above what it replaced. The report measures the three buffers now. `the_in_flight_peak_is_not_two_matrices` pins the arithmetic. ALSO, from the review's output-equality section - `align_shared_axis` decided "same axis" by `ptr::eq` plus a length check, so the fast path's correctness rested on a `ChromChunk::axis_for` invariant nothing in the function could see. It compares bit patterns now (pointer equality kept as the cheap accept), which keeps -0.0 and +0.0 distinct and so accepts exactly the set the union build reproduces. One existing assertion that pinned the old rejection was updated; the equality assertion it carried is unchanged. - `confident_row_spans` refuses statistics that cannot describe a u32 column (negative or inverted bounds) rather than skipping a group on them, which would drop an anchor and shift the half-widths that bound every candidate. - `charge_states_per_row` asserts the charge column covers every PSM row. Both builds zip, and a zip truncates where the `charge[i]` it replaced panicked. REJECTED - "The per-column `ValueMatrix` multiplies the live block count." The arithmetic does not support it. 387 columns of one chunk are 387 blocks; at most two matrices are in flight, so 774 -- 0.07% of the 1,048,576 mappings a process gets, against the 92,030 + 36,368 medium mappings measured on the process that hit the limit. What the layout removes is a whole-matrix copy at every write, hundreds of MB. The trade is lopsided in favour of the columns. The arithmetic is now in the type's doc comment and `value_matrix_columns_are_one_block_each` pins the count so a future widening cannot creep past it. The sub-point that `payload_bytes` became a formula rather than a measurement is granted: it sums the columns' capacities now, so a moved-out column counts as the zero it is. - The finding that `extended_features_are_chunk_invariant` compares new against new, and that `colix_covers_every_active_column` cannot see a transposed pair, is correct and is answered rather than argued with. Route taken: the golden, not the cheaper ColIx-name settlement (both are here). `extended_features_match_the_pre_permutation_build` pins a digest of every column's bit patterns, and the PIN hash, captured by running this same fixture against 1ac1b44^ -- the last commit whose assembly keyed each value by a feature-name string literal. Verified non-vacuous: swapping `ff.frag_corr` and `ff.frag_cosine` between their two columns fails it, while chunk invariance still passes, which is exactly the gap the review named. It also covers the uniform-shift case and the `align_shared_axis` widening above. `colix_resolves_each_field_to_the_column_its_own_name_denotes` is the second layer: `stringify!` from the same macro that builds the struct, so each field is pinned to the column its own name denotes without listing 387 names. Not touched, reported instead: `cargo fmt` reformats `mumdia-io/src/table.rs` and `stages/pool.rs`, which are committed unformatted on this branch and are outside this change's file set. cargo clippy -p mumdia --all-targets -D warnings and cargo test --workspace pass; rustfmt --check is clean on the files touched here. Co-Authored-By: Claude Opus 5 (1M context) --- .../crates/mumdia/src/stages/features.rs | 679 +++++++++++++++++- 1 file changed, 643 insertions(+), 36 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/features.rs b/rust/mumdia/crates/mumdia/src/stages/features.rs index a8bcb900..9cc7aa22 100644 --- a/rust/mumdia/crates/mumdia/src/stages/features.rs +++ b/rust/mumdia/crates/mumdia/src/stages/features.rs @@ -385,6 +385,15 @@ fn align_traces(rows: &[ChromRow]) -> TraceAlign { /// bit: sorting and deduplicating a strictly ascending axis returns it unchanged, and /// looking a row's own RT value up in its own map returns that row's intensity at the /// same position. A row with no trace at all gets zeros either way (its map is empty). +/// +/// "Same axis" is decided by VALUE, not by provenance. The pointer check is kept only as +/// the cheap accept for the case the store actually produces; a `ptr::eq` + length test +/// would otherwise make the fast path's correctness depend on an invariant of +/// [`ChromChunk::axis_for`] -- that a shared slice is exactly the row's own trace grid -- +/// which nothing here can see, and any future store that handed out a longer shared buffer +/// would move every extended feature of every PSM at once with no test failing. Comparing +/// bit patterns rather than `==` keeps `-0.0` and `+0.0` distinct, so the accepted set is +/// the one the union build reproduces exactly. fn align_shared_axis(rows: &[ChromRow]) -> Option { let mut shared: Option<&[f32]> = None; for r in rows { @@ -399,7 +408,12 @@ fn align_shared_axis(rows: &[ChromRow]) -> Option { match shared { None => shared = Some(r.rt), Some(a) => { - if a.len() != r.rt.len() || !std::ptr::eq(a.as_ptr(), r.rt.as_ptr()) { + if a.len() != r.rt.len() { + return None; + } + if !std::ptr::eq(a.as_ptr(), r.rt.as_ptr()) + && !a.iter().zip(r.rt).all(|(x, y)| x.to_bits() == y.to_bits()) + { return None; } } @@ -917,10 +931,15 @@ struct ChromStream { } impl ChromStream { - fn open(ch: &TableFile, path: &str) -> Result { - let has_obs_mz = mumdia_io::table::column_names(path)? - .iter() - .any(|c| c == "frag_obs_mz"); + /// `ch` may be a whole-file handle or a [`TableFile::span`] of one; a span carries the + /// whole file's schema, so the optional column is decided from the handle either way. + /// + /// This used to call `mumdia_io::table::column_names(path)`, which opens the file and + /// parses its footer a SECOND time for a question the handle can already answer. The + /// confident-bounds pass opens one stream per span, so that was two full footer parses + /// per span on a table whose footer is not small. + fn open(ch: &TableFile) -> Result { + let has_obs_mz = ch.has_column("frag_obs_mz"); let mut cols = vec![ "candidate_id", "frag_name", @@ -1149,15 +1168,22 @@ fn plan_chunks( /// flat layout it replaced had to copy every column (`column(c).to_vec()`), which meant /// the whole matrix existed twice at the moment of the write -- the memory report said /// `largest_chunk` and meant half of the real peak. +/// +/// One allocation per column rather than one for the matrix does cut against this stage's +/// stated problem (the NUMBER of live medium blocks, not bytes), so the arithmetic is +/// worth stating: 387 columns of one chunk are 387 blocks and at most two matrices are in +/// flight, so 774 -- 0.07% of the kernel's 1,048,576 mappings per process, against the +/// 92,030 + 36,368 medium mappings measured on the process that hit the limit. What it +/// buys is a whole matrix (hundreds of MB) not copied at every write. The trade is +/// lopsided in favour of the columns; `value_matrix_columns_are_one_block_each` pins the +/// count so a future widening cannot creep past it unnoticed. struct ValueMatrix { - rows: usize, cols: Vec>, } impl ValueMatrix { fn new(n_cols: usize, rows: usize) -> ValueMatrix { ValueMatrix { - rows, cols: (0..n_cols).map(|_| vec![0.0; rows]).collect(), } } @@ -1181,8 +1207,15 @@ impl ValueMatrix { std::mem::take(&mut self.cols[c]) } + /// What the columns actually hold, summed over them rather than computed from the + /// shape, so a column already moved out by [`ValueMatrix::take_column`] counts as the + /// zero it is. `vec![0.0; rows]` allocates capacity exactly `rows`, so on a full matrix + /// this equals `cols * rows * 8`. fn payload_bytes(&self) -> usize { - self.cols.len() * self.rows * std::mem::size_of::() + self.cols + .iter() + .map(|c| c.capacity() * std::mem::size_of::()) + .sum() } } @@ -1224,6 +1257,15 @@ macro_rules! col_ix { fn named(&self) -> Vec> { vec![$(self.$field,)*] } + + /// Each field paired with the feature name it is spelled as, so a test can + /// check that a field resolves to the column its OWN name denotes rather than + /// only that the resolved columns form a permutation. `stringify!` takes the + /// names from the field list itself, so this cannot drift from it. + #[cfg(test)] + fn named_pairs(&self) -> Vec<(&'static str, Option)> { + vec![$((stringify!($field), self.$field),)*] + } } }; } @@ -1343,6 +1385,14 @@ impl PinWriter { /// what this used to do for every run: one `HashSet` heap block per distinct /// peptidoform, about a million live for the whole run on a real library. fn charge_states_per_row(pf_of_row: &[u32], n_pf: usize, charge: &[i32]) -> Vec { + // Both builds zip the two columns, which TRUNCATES where the old `charge[i]` indexing + // panicked. A short charge column is an artifact-shape error, so it stays loud rather + // than becoming a silent partial reduction over the rows that happened to line up. + assert_eq!( + pf_of_row.len(), + charge.len(), + "features: the charge column does not cover every PSM row" + ); charge_states_by_mask(pf_of_row, n_pf, charge) .unwrap_or_else(|| charge_states_by_set(pf_of_row, n_pf, charge)) } @@ -1390,9 +1440,26 @@ fn charge_states_by_set(pf_of_row: &[u32], n_pf: usize, charge: &[i32]) -> Vec s.max || s.max > Some(u32::MAX as f64)) + { + return None; + } let mut spans: Vec<(usize, usize)> = Vec::new(); let mut start = 0usize; for s in stats { @@ -1429,7 +1506,7 @@ fn confident_row_spans( /// used, so a candidate that straddles one is split exactly where it was split before and /// the half-width samples, and therefore the percentiles, are unchanged. fn confident_global_bounds( - chrom_path: &str, + ch: &TableFile, spans: &[(usize, usize)], confident_rows: &HashMap>, apex_rt: &[f64], @@ -1442,9 +1519,15 @@ fn confident_global_bounds( let keep: std::collections::HashSet = confident_rows.keys().copied().collect(); // Chunk reading already groups rows by candidate, so this reuses it and keeps only // the confident candidates' rows long enough to bound their peak. + // + // `ch.span` rather than `TableFile::open_rows`: the caller already holds an open + // handle with the footer parsed, and `span` is an `Arc` clone of it. `open_rows` is + // `open` plus `span`, so it re-opened the file and re-parsed the whole footer once per + // span -- and the pruning's whole point is to produce MANY spans, so the cost grew + // with the saving. for &(first, n_rows) in spans { - let span = TableFile::open_rows(chrom_path, first, n_rows)?; - let mut stream = ChromStream::open(&span, chrom_path)?; + let span = ch.span(first, n_rows)?; + let mut stream = ChromStream::open(&span)?; let (mut abs, end) = (first, first + n_rows); while abs < end { let want = (chunk_rows - abs % chunk_rows).min(end - abs); @@ -1510,10 +1593,39 @@ pub fn run(p: FeaturesParams) -> Result { run_with_chunk_rows(p, CHUNK_CHROM_ROWS) } +/// How the PIN is closed. `Fail` exists only under `cfg(test)` and injects a failure at +/// the flush, which is the one step of the stage that can fail after every chunk has been +/// computed but before the features table is published. It is a parameter rather than a +/// global so that `a_failed_pin_finish_does_not_publish_the_features_table` cannot +/// interfere with any other test in this binary. +#[derive(Clone, Copy)] +enum PinFinish { + Normal, + #[cfg(test)] + Fail, +} + +impl PinFinish { + fn apply(self, w: PinWriter) -> Result<()> { + match self { + PinFinish::Normal => w.finish(), + #[cfg(test)] + PinFinish::Fail => { + drop(w); + Err(anyhow!("features: injected PIN flush failure")) + } + } + } +} + /// [`run`] with an explicit chunk size. Only the chunk boundaries change with it: the /// feature values, the row order and the PIN bytes do not, which is what the /// `features_chunking_is_value_preserving` test asserts by hashing both artifacts. pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result { + run_chunked(p, chunk_rows, PinFinish::Normal) +} + +fn run_chunked(p: FeaturesParams, chunk_rows: usize, pin_finish: PinFinish) -> Result { let t0 = Instant::now(); let ps = TableFile::open(p.psms)?; let cid = ps.u32("candidate_id")?; @@ -1651,7 +1763,7 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result "features: confident-bounds pass reads this much of the chromatogram table" ); confident_global_bounds( - p.chromatograms, + &ch, &spans, &confident_rows, &apex_rt, @@ -1739,6 +1851,9 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result // Accounting for the audit (docs/27): the largest chunk in flight against the run // total streamed is exactly the quantity chunking changes. let (mut max_frag_bytes, mut max_ms1_bytes, mut max_matrix_bytes) = (0usize, 0usize, 0usize); + // Value-buffer high-water mark across the overlap, and the bytes the writer thread is + // holding right now (the previous chunk's moved columns; 0 before the first send). + let (mut max_inflight_bytes, mut in_writer) = (0usize, 0usize); let (mut tot_frag_bytes, mut tot_ms1_bytes) = (0usize, 0usize); let (mut n_cand, mut n_frag_rows, mut n_ms1_rows) = (0usize, 0usize, 0usize); @@ -1751,7 +1866,6 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result // exactly what the serial code saw at that point. let chunk_rows: Vec = chunks.iter().map(|c| c.chrom_rows).collect(); let ch_ref = &ch; - let chrom_path = p.chromatograms; let rows = std::thread::scope(|sc| -> Result { // A RENDEZVOUS channel, not a one-deep buffer. `sync_channel(1)` let the loader // finish a chunk, park it in the buffer and start a third, so three chunks of @@ -1800,7 +1914,7 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result writer.close() }); sc.spawn(move || { - let mut stream = match ChromStream::open(ch_ref, chrom_path) { + let mut stream = match ChromStream::open(ch_ref) { Ok(s) => s, Err(e) => { let _ = tx.send(Err(e)); @@ -2018,8 +2132,16 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result - rt_err / gradient; } drop(frag_feats); + // `ext_vals` and the matrix are both live for the whole assembly above, so the + // stage's real high-water mark is not one matrix. Measure the three value-shaped + // buffers that overlap rather than asserting a multiple of one of them: this + // chunk's matrix, this chunk's extended buffer, and the PREVIOUS chunk's columns, + // which the writer thread still owns until it has encoded them. + let ext_bytes = crate::memlog::bytes_of(&ext_vals); drop(ext_vals); - max_matrix_bytes = max_matrix_bytes.max(m.payload_bytes()); + let matrix_bytes = m.payload_bytes(); + max_matrix_bytes = max_matrix_bytes.max(matrix_bytes); + max_inflight_bytes = max_inflight_bytes.max(matrix_bytes + ext_bytes + in_writer); // PIN first: it reads the same values the columns below move into Arrow. if let Some(w) = pin.as_mut() { @@ -2062,17 +2184,29 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result // The writer failed; its error is the real one, so surface that. break; } + // The send has returned, so the writer now owns this chunk's columns and holds + // them while the next chunk is assembled. They are the matrix's columns, moved. + in_writer = matrix_bytes; + } + // The PIN is flushed BEFORE the commit marker, because `writer.close()` publishes + // `features.parquet` over the previous good one and the PIN write is the last thing + // that can still fail. The in-line version ran `pin.finish()` and then + // `writer.close()`, so a PIN that failed at flush (a full disk, a lost permission on + // `out_pin`) left the previous features table untouched. Moving `close` onto the + // writer thread silently reversed that: the table was published and the stage then + // errored, leaving a published features table with no matching PIN. Failing here + // returns without sending the marker, so the writer drops unpublished exactly as it + // does for any other error in this closure. + if let Some(w) = pin.take() { + pin_finish.apply(w)?; } - // Commit only after every chunk went out; a `break` above leaves this send - // failing, and the writer's own error is what `join` then returns. + // Commit only after every chunk went out AND the PIN is on disk; a `break` above + // leaves this send failing, and the writer's own error is what `join` then returns. let _ = wtx.send(None); drop(wtx); wh.join() .map_err(|_| anyhow!("features: the parquet writer thread panicked"))? })?; - if let Some(w) = pin { - w.finish()?; - } crate::memlog::report( "features chromatogram store", @@ -2090,16 +2224,20 @@ pub fn run_with_chunk_rows(p: FeaturesParams, chunk_rows: usize) -> Result chrom_rows = chrom_rows_total, "mem: features chromatogram store shape" ); - // Two matrices can be alive at once and the report used to name only one, so it - // understated the peak by exactly 2x. It did so for the wrong reason as well: the - // second copy was `column(c).to_vec()` into the writer's columns, pure waste. The - // columns are moved now, and the second matrix that remains is the next chunk's, - // built while the writer thread encodes this one -- overlap, not a copy. + // `in_flight_peak` is MEASURED, not a multiple of `largest_chunk`. The in-line write + // this replaced held one matrix plus the `column(c).to_vec()` copy of it, so 2x was the + // right figure for it; the writer thread holds no copy but does hold the previous + // chunk's columns while this chunk is assembled, and the assembly also has the extended + // buffer live (337 of 387 columns wide at Extended, so ~0.87x a matrix). The overlap + // therefore peaks near 2.9x a matrix, ABOVE the 2x it replaced -- the encode came off + // the critical path at the cost of bytes, and the report says so rather than asserting + // the old figure. `max_inflight_bytes` covers only these three buffers; the chromatogram + // store is reported separately above and the Arrow encode's own buffers are the writer's. crate::memlog::report( "features value matrix", &[ ("largest_chunk", max_matrix_bytes), - ("in_flight_peak", 2 * max_matrix_bytes), + ("in_flight_peak", max_inflight_bytes), ], ); @@ -2821,11 +2959,13 @@ mod tests { assert_eq!(al.traces_full[0], vec![1.0, 0.0, 2.0, 3.0, 0.0]); assert_eq!(al.traces_full[1], vec![0.0, 4.0, 5.0, 0.0, 6.0]); - // Equal VALUES but stored separately are not the same slice: still correct, just - // through the general path. + // Equal VALUES stored separately ARE the same axis. This used to be rejected, + // because the precondition was `ptr::eq` and the fast path's correctness rested on + // the chunk store's provenance rather than on the values; the answer was the same + // then and is the same now, which is what `same` checks. let copy = g1; let rows = vec![row("y1", 1.0, &g1, &v1), row("y2", 1.0, ©, &v2)]; - assert!(align_shared_axis(&rows).is_none()); + assert!(align_shared_axis(&rows).is_some()); same(&align_traces(&rows), &align_union(&rows)); // A non-ascending axis would be reordered by the sort, and a repeated RT value @@ -2873,6 +3013,17 @@ mod tests { ); } + #[test] + #[should_panic(expected = "does not cover every PSM row")] + fn a_short_charge_column_is_an_error_not_a_partial_reduction() { + // Both builds zip the peptidoform ids with the charges, and a zip truncates where + // the `charge[i]` indexing it replaced panicked. A charge column that does not cover + // every PSM row is an artifact-shape error, so it must stay loud rather than + // silently reducing over a prefix. + let pf = [0u32, 0, 1, 1]; + charge_states_per_row(&pf, 2, &[2i32, 3]); + } + // --- the feature -> column permutation --- #[test] @@ -2903,6 +3054,34 @@ mod tests { } } + #[test] + fn colix_resolves_each_field_to_the_column_its_own_name_denotes() { + // Coverage alone only proves the permutation is TOTAL. It cannot see a field + // resolving to the wrong column, because a permutation stays a permutation under a + // swap. This pins each field to the index of the column spelled like it, which is + // what the string literal it replaced did, taking the names from the field list via + // `stringify!` rather than repeating them. + for set in [FeatureSet::Minimal, FeatureSet::Rich, FeatureSet::Extended] { + let cols = active_features(set); + let ext = extended_name_refs(); + let ix = ColIx::new(&cols, &ext); + for (name, got) in ix.named_pairs() { + let want = cols.iter().position(|c| c == name); + assert_eq!(got, want, "{set:?}: field '{name}' resolved to {got:?}"); + } + for (name, got) in ext.iter().zip(&ix.ext) { + let want = cols.iter().position(|c| c == name); + assert_eq!(*got, want, "{set:?}: extended '{name}' resolved to {got:?}"); + } + } + // This still cannot catch a transposition in the ASSEMBLY itself -- writing + // `ff.frag_corr` into the `frag_cosine` column and vice versa resolves both fields + // correctly and covers every column once. No structural check can: a swap leaves a + // permutation a permutation. `extended_features_match_the_pre_permutation_build` + // is the test that sees it, by comparing the values against a digest captured from + // the name-keyed assembly. + } + #[test] fn value_matrix_hands_its_columns_over_by_move() { let mut m = ValueMatrix::new(3, 2); @@ -2910,8 +3089,55 @@ mod tests { m.set(Some(1), 1, -1.0); m.set(None, 0, 99.0); // a feature outside the active set is dropped assert_eq!(m.column(1), &[7.5, -1.0]); + // A taken column leaves nothing behind: the write moves the values out rather than + // copying them, which is what stopped the matrix existing twice at the write. + let before = m.payload_bytes(); assert_eq!(m.take_column(1), vec![7.5, -1.0]); assert_eq!(m.column(0), &[0.0, 0.0]); + assert_eq!( + m.payload_bytes(), + before - 2 * std::mem::size_of::(), + "payload_bytes must measure what is still held, not the original shape" + ); + } + + #[test] + fn value_matrix_columns_are_one_block_each() { + // The per-column layout is what lets the writer take the columns by move, and it is + // also the one place this change adds heap blocks rather than removing them. The + // arithmetic that makes it acceptable: one block per column, at most two matrices in + // flight, against the 1,048,576 mappings a process gets. + let cols = active_features(FeatureSet::Extended).len(); + let m = ValueMatrix::new(cols, 4096); + assert_eq!(m.cols.len(), cols); + assert!( + 2 * cols < 1_048_576 / 1000, + "two matrices in flight must stay far under the mapping limit: {cols} columns" + ); + assert_eq!(m.payload_bytes(), cols * 4096 * std::mem::size_of::()); + } + + #[test] + fn the_in_flight_peak_is_not_two_matrices() { + // What the report used to assert (`2 * largest_chunk`) was the in-line write's peak: + // one matrix plus the copy of it. With the writer thread there is no copy, but the + // previous chunk's columns are still held while this chunk's matrix AND its extended + // buffer are live, so the overlap is wider than 2x. The report measures the three; + // this pins the arithmetic that makes the old constant wrong. + let (rows, n_cols) = (4096usize, active_features(FeatureSet::Extended).len()); + let n_ext = extended_name_refs().len(); + let matrix = n_cols * rows * std::mem::size_of::(); + let ext = n_ext * rows * std::mem::size_of::(); + let in_flight = matrix + ext + matrix; + assert!( + in_flight > 2 * matrix, + "the overlap must be reported as more than two matrices: {in_flight} vs {}", + 2 * matrix + ); + assert!( + n_ext * 100 / n_cols >= 80, + "the extended buffer is ~0.87x a matrix" + ); } // --- confident-bounds row-group pruning --- @@ -2962,6 +3188,13 @@ mod tests { /// A chromatogram table with several row groups, its confident candidates, and the /// PSM apex RTs, for the pruning-equality test below. + /// + /// The fragment count per candidate VARIES (4 to 9), so candidates straddle the 12-row + /// group boundaries. That is the configuration the pruning's safety argument is about + /// and the previous fixture never produced: at a fixed 6 rows per candidate every + /// boundary fell exactly between two candidates, so "a candidate spanning two row + /// groups puts its id inside BOTH ranges and both are kept" was never exercised. + /// The anchor pattern also leaves skipped groups between kept ones. fn craft_chrom_for_bounds(path: &str) -> (HashMap>, Vec) { let (mut cid, mut name): (Vec, Vec) = (Vec::new(), Vec::new()); let (mut fmz, mut pint): (Vec, Vec) = (Vec::new(), Vec::new()); @@ -2976,7 +3209,10 @@ mod tests { confident.insert(c, vec![c as usize]); } let grid: Vec = (0..9).map(|k| (apex - 4.0 + k as f64) as f32).collect(); - for f in 0..6u32 { + // 4..=9 fragments, so the running row total is not a multiple of the 12-row + // group size and candidates land across boundaries. + let nfrag = 4 + c % 6; + for f in 0..nfrag { cid.push(c); name.push(format!("y{}", f + 1)); fmz.push(200.0 + f as f64 * 30.0); @@ -3022,6 +3258,31 @@ mod tests { let mut sorted: Vec = confident.keys().copied().collect(); sorted.sort_unstable(); let rg = ch.row_group_stats("candidate_id").unwrap(); + + // The whole safety argument is about a candidate whose rows cross a row-group + // boundary: its id is then inside BOTH groups' [min, max] and neither can be + // skipped. Assert the fixture actually produces that, because the previous one -- + // a fixed 6 rows per candidate into 12-row groups -- never did, and the argument + // went untested while the test passed. + let straddlers = rg.windows(2).filter(|w| w[0].max == w[1].min).count(); + assert!( + straddlers > 5, + "the fixture must make candidates straddle group boundaries; only {straddlers} do" + ); + // And a SKIPPED group must sit between two kept ones, which is what forces the + // chunk grid inside a span to stay on absolute rows rather than restart per span. + let kept: Vec = rg + .iter() + .map(|s| { + let (lo, hi) = (s.min.unwrap(), s.max.unwrap()); + sorted.iter().any(|&c| (c as f64) >= lo && (c as f64) <= hi) + }) + .collect(); + assert!( + kept.windows(3).any(|w| w[0] && !w[1] && w[2]), + "the fixture must skip a group between two kept ones" + ); + let spans = confident_row_spans(&rg, &sorted).expect("the writer records statistics"); let read: usize = spans.iter().map(|s| s.1).sum(); assert!( @@ -3030,16 +3291,18 @@ mod tests { ch.nrows ); assert!(read > 0); + assert!(spans.len() > 1, "a single span would not exercise the grid"); // Same half-widths at several chunk sizes, pruned against the whole-table scan. // The chunk size matters: a candidate that straddles a chunk boundary is bounded // from each part separately, and the pruned pass has to split it in exactly the - // same places for the percentiles to come out the same. - for chunk in [7usize, 16, 64, 1 << 20] { + // same places for the percentiles to come out the same. Chunk sizes both below and + // above the 12-row group size, so a span covers several groups in one chunk. + for chunk in [7usize, 16, 64, 100, 1 << 20] { let pruned = - confident_global_bounds(&path, &spans, &confident, &apex_rt, &cfg, chunk).unwrap(); + confident_global_bounds(&ch, &spans, &confident, &apex_rt, &cfg, chunk).unwrap(); let whole = - confident_global_bounds(&path, &[(0, ch.nrows)], &confident, &apex_rt, &cfg, chunk) + confident_global_bounds(&ch, &[(0, ch.nrows)], &confident, &apex_rt, &cfg, chunk) .unwrap(); let bits = |v: Option<(f64, f64)>| v.map(|(l, r)| (l.to_bits(), r.to_bits())); assert_eq!( @@ -3051,6 +3314,79 @@ mod tests { } } + #[test] + fn row_group_pruning_saves_nothing_on_a_production_shaped_table() { + // The doc comment on `confident_row_spans` claims the pruning is worth close to + // nothing on real data. This measures that claim instead of asserting it, from the + // shape extract actually writes: `CHROM_ROW_GROUP_ROWS = 1 << 16` rows per group and + // about 10 chromatogram rows per candidate, so ~6,550 candidates per group, with the + // confident set (seed PSMs at spectrum_q <= 0.01) spread over the id range. + let rows_per_group = 1usize << 16; + let rows_per_candidate = 10usize; + let per_group = rows_per_group / rows_per_candidate; + let n_groups = 200usize; + let n_cand = (per_group * n_groups) as u32; + let s: Vec = (0..n_groups) + .map(|g| mumdia_io::table::RowGroupStats { + rows: rows_per_group, + min: Some((g * per_group) as f64), + max: Some(((g + 1) * per_group - 1) as f64), + }) + .collect(); + + // 1% of candidates confident, evenly spread: every group holds ~65 of them. + let spread: Vec = (0..n_cand).step_by(100).collect(); + let spans = confident_row_spans(&s, &spread).unwrap(); + let read: usize = spans.iter().map(|x| x.1).sum(); + let total = rows_per_group * n_groups; + assert_eq!( + (spans.len(), read), + (1, total), + "on a production-shaped table the pruning reads the whole file in one span" + ); + + // Where it does pay: a confident set clustered in a slice of the id range, which is + // the shape a per-group search or a re-run over a subset produces. + let clustered: Vec = (0..per_group as u32 * 3).collect(); + let read: usize = confident_row_spans(&s, &clustered) + .unwrap() + .iter() + .map(|x| x.1) + .sum(); + assert_eq!( + read, + 3 * rows_per_group, + "a clustered set reads 3 of 200 groups" + ); + } + + #[test] + fn row_group_pruning_refuses_statistics_that_cannot_describe_a_u32_column() { + // A writer that ordered an unsigned column's min/max under SIGNED comparison hands + // back a negative bound. Skipping a group on a range like that would drop a + // confident anchor and shift the global half-widths that bound every candidate, so + // the pruning must refuse and read the whole table instead. + let signed = vec![ + mumdia_io::table::RowGroupStats { + rows: 10, + min: Some(-4.0), + max: Some(9.0), + }, + mumdia_io::table::RowGroupStats { + rows: 10, + min: Some(10.0), + max: Some(19.0), + }, + ]; + assert!(confident_row_spans(&signed, &[15]).is_none()); + // An inverted range is equally unusable. + let inverted = stats(&[(10, 9, 0)]); + assert!(confident_row_spans(&inverted, &[5]).is_none()); + // And the same groups with honest bounds do prune. + let ok = stats(&[(10, 0, 9), (10, 10, 19)]); + assert_eq!(confident_row_spans(&ok, &[15]).unwrap(), vec![(10, 10)]); + } + // --- the whole chunked pass, end to end --- /// Crafted psms_extracted + chromatograms for the Extended feature set: shared window @@ -3168,6 +3504,123 @@ mod tests { (psms, chrom) } + fn fnv(h: &mut u64, b: &[u8]) { + for &x in b { + *h ^= x as u64; + *h = h.wrapping_mul(0x1000_0000_01b3); + } + } + + /// Order-sensitive digest of a features table: every column name in file order and + /// then its values as raw bit patterns. Writing a value into a different column than + /// before changes it, which is the property + /// `extended_features_match_the_pre_permutation_build` needs. + fn features_digest(path: &str) -> (usize, usize, u64) { + let t = mumdia_io::table::Table::read(path).unwrap(); + let names = t.column_names(); + let mut h = 0xcbf2_9ce4_8422_2325u64; + for name in &names { + fnv(&mut h, name.as_bytes()); + if let Ok(v) = t.f64(name) { + fnv(&mut h, b"f64"); + for x in &v { + fnv(&mut h, &x.to_bits().to_le_bytes()); + } + } else if let Ok(v) = t.u32(name) { + fnv(&mut h, b"u32"); + for x in &v { + fnv(&mut h, &x.to_le_bytes()); + } + } else if let Ok(v) = t.i32(name) { + fnv(&mut h, b"i32"); + for x in &v { + fnv(&mut h, &x.to_le_bytes()); + } + } else if let Ok(v) = t.str(name) { + fnv(&mut h, b"str"); + for s in &v { + fnv(&mut h, s.as_bytes()); + } + } + } + (t.nrows, names.len(), h) + } + + #[test] + fn extended_features_match_the_pre_permutation_build() { + // A GOLDEN, captured by running THIS fixture against commit 1ac1b44^ -- the last + // commit whose serial assembly still keyed each value by a feature-name string + // literal (`m.set("frag_corr", r, ff.frag_corr)`), before `ColIx` replaced the + // names with a precomputed column permutation. + // + // This is the only test that can see a transposed pair. `colix_covers_every_active + // _column` proves the permutation is total and + // `colix_resolves_each_field_to_the_column_its_own_name_denotes` proves each field + // resolves to its own name, but neither can see `m.set(ix.frag_cosine, r, + // ff.frag_corr)` written alongside `m.set(ix.frag_corr, r, ff.frag_cosine)`: the + // swap is still a permutation and each field still resolves correctly. It also + // catches what chunk invariance cannot -- a regression that shifts every row + // identically -- and the equality evidence the change was promoted on was a + // temporary capture module the commit deleted, so nothing in the repository + // reproduced it until now. + // + // REGENERATING: these two constants are a claim about the values of ANOTHER commit. + // If a deliberate change to the feature arithmetic breaks them, re-derive them by + // running this fixture at 1ac1b44^ (or at the last commit whose values are trusted) + // and changing them in one commit that says which values moved and why. Taking the + // new numbers from the current build makes the test pin nothing. + const GOLDEN_DIGEST: u64 = 0x4e43_7960_2b2a_a1cc; + const GOLDEN_PIN: &str = "806a126473eafce7a1567ac71b5253e8605af0062864a02f2ab5f0b77e851817"; + + let dir = std::env::temp_dir().join("mumdia_features_golden"); + std::fs::create_dir_all(&dir).unwrap(); + let (psms, chrom) = craft_extended_inputs(&dir, 61); + // Both chunk sizes, because the golden was captured at both and they agreed: the + // pin therefore covers the chunked path as well as the single-chunk one. + for (tag, chunk) in [("one", 1usize << 20), ("many", 1usize)] { + let out = dir + .join(format!("golden_{tag}.parquet")) + .to_string_lossy() + .to_string(); + let pin = dir + .join(format!("golden_{tag}.pin")) + .to_string_lossy() + .to_string(); + let mut cfg = FeaturesConfig { + set: FeatureSet::Extended, + emit_pin: true, + bound_from_confident: false, + ..Default::default() + }; + cfg.ms1_precursor_features = true; + run_with_chunk_rows( + FeaturesParams { + psms: &psms, + chromatograms: &chrom, + seed: None, + out: &out, + out_pin: &pin, + cfg: &cfg, + config_hash: "test", + }, + chunk, + ) + .unwrap(); + let (rows, ncols, digest) = features_digest(&out); + assert_eq!((rows, ncols), (61, 398), "chunk {tag}: table shape moved"); + assert_eq!( + digest, GOLDEN_DIGEST, + "chunk {tag}: a feature value or its column differs from the name-keyed \ + assembly at 1ac1b44^ (digest 0x{digest:016x})" + ); + assert_eq!( + mumdia_io::hash::blake3_file(&pin).unwrap(), + GOLDEN_PIN, + "chunk {tag}: the PIN bytes differ from the name-keyed assembly" + ); + } + } + #[test] fn extended_features_are_chunk_invariant() { // The `features_chunking_is_value_preserving` integration test runs the DEFAULT @@ -3350,4 +3803,158 @@ mod tests { "temp files left behind: {leftovers:?}" ); } + + #[test] + fn a_failed_pin_finish_does_not_publish_the_features_table() { + // The in-line write ran `pin.finish()` and THEN `writer.close()`, so a PIN that + // failed at the flush left the previous features table exactly as it was. Moving + // `close` onto the writer thread reversed the order silently: the table was + // published and the stage then errored, leaving a published features table with no + // matching PIN. This pins the OLD ordering. + let dir = std::env::temp_dir().join("mumdia_features_pin_order"); + std::fs::create_dir_all(&dir).unwrap(); + let (psms, chrom) = craft_extended_inputs(&dir, 11); + let out = dir + .join("features_pin.parquet") + .to_string_lossy() + .to_string(); + let pin = dir.join("features_pin.pin").to_string_lossy().to_string(); + let cfg = FeaturesConfig { + set: FeatureSet::Rich, + emit_pin: true, + bound_from_confident: false, + ..Default::default() + }; + let params = || FeaturesParams { + psms: &psms, + chromatograms: &chrom, + seed: None, + out: &out, + out_pin: &pin, + cfg: &cfg, + config_hash: "test", + }; + + run_chunked(params(), 4, PinFinish::Normal).expect("the good run must succeed"); + let good = mumdia_io::hash::blake3_file(&out).unwrap(); + + let err = run_chunked(params(), 4, PinFinish::Fail) + .expect_err("a PIN that cannot be flushed must fail the stage"); + assert!( + err.to_string().contains("PIN"), + "the PIN failure must be the reported error, not the writer's: {err}" + ); + assert_eq!( + mumdia_io::hash::blake3_file(&out).unwrap(), + good, + "a PIN failure must not publish the features table over the previous one" + ); + let leftovers: Vec<_> = std::fs::read_dir(&dir) + .unwrap() + .filter_map(|e| e.ok()) + .map(|e| e.file_name().to_string_lossy().to_string()) + .filter(|n| n.contains("features_pin.parquet.tmp-")) + .collect(); + assert!( + leftovers.is_empty(), + "temp files left behind: {leftovers:?}" + ); + } + + #[test] + fn the_chrom_stream_reads_its_optional_column_from_the_handle_it_is_given() { + // `ChromStream::open` used to answer "does this table carry frag_obs_mz?" by + // opening the file a second time and parsing its whole footer + // (`mumdia_io::table::column_names(path)`), once per stream -- and the + // confident-bounds pass opens one stream per span. It reads the handle's schema + // now, including on a `TableFile::span`, which carries the WHOLE file's schema. + // Same answer either way, which is what this pins: the observed m/z must still + // come from `frag_obs_mz` where the column exists and fall back to `frag_mz` where + // it does not. + let dir = std::env::temp_dir().join("mumdia_features_stream_schema"); + std::fs::create_dir_all(&dir).unwrap(); + let with = dir.join("with_obs.parquet").to_string_lossy().to_string(); + let without = dir.join("no_obs.parquet").to_string_lossy().to_string(); + let (cid, name): (Vec, Vec) = ( + (0..8u32).flat_map(|c| [c, c]).collect(), + (0..16).map(|i| format!("y{}", i % 2 + 1)).collect(), + ); + let fmz: Vec = (0..16).map(|i| 200.0 + i as f64).collect(); + let obs: Vec = fmz.iter().map(|v| v + 0.25).collect(); + let pint: Vec = vec![1.0; 16]; + let rt: Vec> = vec![vec![1.0, 2.0, 3.0]; 16]; + let inten: Vec> = vec![vec![4.0, 5.0, 6.0]; 16]; + let base = |extra: bool| { + let mut cols = vec![ + Col::U32("candidate_id".into(), cid.clone()), + Col::Str("frag_name".into(), name.clone()), + Col::F64("frag_mz".into(), fmz.clone()), + Col::F32("predicted_intensity".into(), pint.clone()), + Col::ListF32("rt".into(), rt.clone()), + Col::ListF32("intensity".into(), inten.clone()), + ]; + if extra { + cols.push(Col::F64("frag_obs_mz".into(), obs.clone())); + } + cols + }; + let mut w = TableWriter::new(&with).with_row_group_rows(4); + w.write_cols(base(true)).unwrap(); + w.close().unwrap(); + let mut w = TableWriter::new(&without).with_row_group_rows(4); + w.write_cols(base(false)).unwrap(); + w.close().unwrap(); + + // Every combination of "has the column" x "whole file / span of it". + for (path, offset) in [(&with, 0usize), (&with, 8), (&without, 0), (&without, 8)] { + let ch = TableFile::open(path).unwrap(); + let handle = if offset == 0 { + ch.span(0, 16).unwrap() + } else { + ch.span(offset, 8).unwrap() + }; + let mut stream = ChromStream::open(&handle).unwrap(); + let mut names = NameTab::default(); + let chunk = stream.read_chunk(handle.nrows, &mut names).unwrap(); + assert_eq!(chunk.cids.len(), handle.nrows / 2); + let rows = chunk.rows(&chunk.frag, 0, &names); + let (m, o) = (rows[0].frag_mz, rows[0].frag_obs_mz); + if path == &with { + assert_eq!(o, m + 0.25, "frag_obs_mz must be read where it exists"); + } else { + assert_eq!( + o, m, + "frag_obs_mz must fall back to frag_mz where it does not" + ); + } + } + } + + #[test] + fn the_shared_axis_fast_path_does_not_depend_on_where_the_axis_came_from() { + // The precondition used to be `ptr::eq` plus a length check, so the fast path was + // correct only because of an invariant of the chunk store that nothing here could + // see. It compares bit patterns now, so two rows carrying equal-valued axes from + // different buffers take the same path -- and, as always, must agree with the union + // build. A `-0.0` against a `+0.0` stays on the union path, because the two builds + // do not agree there. + let axis_a: Vec = vec![1.0, 2.0, 3.0]; + let axis_b: Vec = vec![1.0, 2.0, 3.0]; // equal values, different allocation + let (ia, ib) = (vec![10.0f32, 20.0, 30.0], vec![1.0f32, 2.0, 3.0]); + let rows = vec![row("y1", 1.0, &axis_a, &ia), row("y2", 0.5, &axis_b, &ib)]; + let fast = align_shared_axis(&rows).expect("equal axes must take the fast path"); + same(&fast, &align_union(&rows)); + + let signed_zero: Vec = vec![-0.0, 2.0, 3.0]; + let plus_zero: Vec = vec![0.0, 2.0, 3.0]; + let rows = vec![ + row("y1", 1.0, &signed_zero, &ia), + row("y2", 0.5, &plus_zero, &ib), + ]; + assert!( + align_shared_axis(&rows).is_none(), + "-0.0 and +0.0 are not the same axis for the union build, so not for this one" + ); + same(&align_traces(&rows), &align_union(&rows)); + } } From 96847ca06d418c89c229fc114dca5edd4ea600ea Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:22:41 +0200 Subject: [PATCH 26/35] docs(33): the measured pool numbers, and the standalone pool command 136 GB of one run's band artifacts: 3 min 12 s spliced at 2.6 GB resident, against 3.5 MB/s decoding and re-encoding on one core, which had reached 42.7 GB in 2 h 50 min when it was stopped. The disk does 221 MB/s. Co-Authored-By: Claude Opus 5 (1M context) --- docs/33_window_groups.md | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/docs/33_window_groups.md b/docs/33_window_groups.md index cba2acfb..68aed914 100644 --- a/docs/33_window_groups.md +++ b/docs/33_window_groups.md @@ -175,11 +175,22 @@ wants and already encoded, so `pool` splices each band's parquet row groups into without decoding them (`mumdia_io::table::SpliceWriter`, the column-chunk append the parquet writer exposes for concatenation). Only the row groups that hold a candidate the overlap dedup drops are decoded, filtered and re-encoded, and window overlap puts those at the two -ends of a band. Measured on the production run before the change: decoding and re-encoding -ran at 3.5 MB/s on one core, so one run's 68 GB of chromatograms would have taken about five -hours, on a disk that reads at 221 MB/s. The spliced output holds the same rows in the same -order with the same values; its row groups are the bands' own, so it is not byte-identical to -a re-encoded pool. The feature and competed schema companions +ends of a band. Measured on the production seven-file experiment, both arms on the same +63 bands of one run (537,047,953 chromatogram rows and 35,844,209 competed rows, 136 GB): + +| pooling | wall | peak RSS | +|---|---|---| +| decode and re-encode | 3.5 MB/s, 42.7 GB in 2 h 50 min (killed) | 4.3 GB | +| splice the row groups | 3 min 12 s for all 136 GB | 2.6 GB | + +The disk reads 221 MB/s (`dd`, direct), so the old path was two orders of magnitude off the +hardware and single-threaded: 110% CPU throughout. The spliced output holds the same rows in +the same order with the same values; its row groups are the bands' own, so it is not +byte-identical to a re-encoded pool. + +`mumdia pool --groups-dir /groups` runs the same stage standalone, which is how those +numbers were taken, and is what to reach for when a grouped search finished but the run did +not: the band directories hold everything, and pooling them is a copy. The feature and competed schema companions (`
.schema.json`, the classifier's column list) are copied from the first band; every band wrote the same one. Each pooled table gets a `.report.json` whose stage is `pool` and whose stats record the number of groups and the overlap duplicates removed. From a34d10b7453e2d592c60ce6819ca84a7d2540ee5 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:28:35 +0200 Subject: [PATCH 27/35] docs(changelog): the memory-mapping pass, the spliced pool, and the two byte changes Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6298c4a7..50b91197 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,45 @@ than a number. Both are recorded in every run's `manifest.json`. ### Changed +- **The engine holds far fewer heap blocks, because that, not memory, is what a banded + search runs out of.** A grouped search of a 203M-precursor library died at about 290 GB + resident with 1.7 TB free, reporting a failed 3 KB allocation. Sampled on the live + process at 180 GB it held 129,393 memory mappings of the kernel's 1,048,576 per-process + limit, 92,030 of them 64-256 KB and 36,368 of them 256 KB-1 MB: one heap block per item, + which mimalloc maps individually. Six subsystems now use flat buffers where they used a + block per row: the library's `label` column is read as one bit per row rather than one + `String` (203M blocks), compete groups PSMs through one sorted vector rather than a + `Vec` per competition group, features holds one buffer rather than a `Vec` per PSM + and dense ids rather than a `HashSet` per peptidoform, quant reads chromatograms into + flat buffers with an axis store rather than three blocks per row, and rescoring's + per-fold training matrix is one buffer rather than one per training row. Rescore also + stops materialising the engine's own feature matrix when a sidecar classifier will read + the handoff parquet instead, which at experiment scale is about 250 GB that need not + exist. Artifacts are unchanged except where noted below. +- **Pooling a grouped run's bands is a byte copy.** The rows are already in the pooled + table's order and already encoded, so `pool` splices the bands' parquet column chunks + into the output and decodes only the row groups whose `candidate_id` statistics say they + hold a candidate the overlap dedup drops. Measured on one run of the production + experiment, 63 bands and 136 GB: 3 minutes 12 seconds at 2.6 GB resident, against + 3.5 MB/s at 110% CPU before, which had reached 42.7 GB in 2 hours 50 minutes when it was + stopped, on a disk that reads 221 MB/s. The pooled tables hold the same rows in the same + order with the same values; their row groups are the bands' own, so their content hashes + differ from a re-encoded pool's. +- **A grouped run no longer pools `psms_extracted`.** Its only reader is the candidate + audit, so it is pooled when `extract.emit_candidate_audit` is set and left per band + otherwise. The band tables are written either way. +- **`groups.parallel` is clamped to one less than the thread count.** A band in flight + parks a rayon worker on its extraction's accumulation channel, so as many bands as + threads leaves no worker to feed them and the run deadlocks with no error and no output + (reproduced on the fixture at `parallel = 2, --threads 2`). A larger value now warns and + uses `threads - 1`. +- **`psms_extracted.parquet` and `run_windows.parquet` change content hash.** `write_table` + now encodes in 65,536-row chunks instead of building one Arrow copy of the whole table, + and on a column that is entirely NULL -- `apex_im` and the three ion-mobility columns + always are -- the definition levels are run-encoded, which moves the writer's internal + mini-batch size and so the page framing. Measured at 196,615 rows: 16 bytes, with every + row, value and row-group boundary unchanged. + - **Library writers emit fragment tables sorted by `candidate_id`.** `import_diann_lib.py`, `make_reverse_decoys.py` and `make_shift_decoys.py` finish with a streaming bucket sort (`_lib_io.sort_fragments_by_candidate`: partition by candidate-id range into temporary @@ -55,6 +94,9 @@ than a number. Both are recorded in every run's `manifest.json`. filtered scan, with a warning. ### Added +- **`mumdia pool` pools a grouped run's band artifacts from the command line.** `run` does + this itself at the end of a grouped search; standalone it is for the case where the + search finished and the run did not, which now costs a pool rather than a re-search. - **`groups.window_groups` searches a run one isolation-window group at a time.** The run's windows are cut, in ascending m/z, into contiguous groups whose library bands hold about the same number of precursors (planned from the precursor table's row-group From 0877839c92c7a8346fe29b2564e4f2b36a17aa19 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:29:16 +0200 Subject: [PATCH 28/35] docs(config): regenerate for the groups.parallel thread-count bound Co-Authored-By: Claude Opus 5 (1M context) --- configs/config-schema.json | 6 +++--- docs/24_config_reference.md | 39 +++++++++++++++++++------------------ 2 files changed, 23 insertions(+), 22 deletions(-) diff --git a/configs/config-schema.json b/configs/config-schema.json index 29621cff..47553b7a 100644 --- a/configs/config-schema.json +++ b/configs/config-schema.json @@ -33,7 +33,7 @@ "default_text": "0", "help": "", "gates": [], - "source_line": 1929 + "source_line": 1934 }, { "path": "convert.thermo_raw_parser", @@ -2456,9 +2456,9 @@ "optional": false, "default": 1, "default_text": "1", - "help": "Bands searched at the same time inside one run. `1` (the default) is one band at a time, which is what bounds the memory: each band in flight holds its own extraction working set, so the peak is this many bands' worth. Raise it to fill a large machine, after checking one band's peak RSS: on a 203M-precursor library at 63 bands the largest band took 39 GB and the median far less. Results do not depend on it; bands are independent and their artifacts are pooled in band order either way.", + "help": "Bands searched at the same time inside one run. `1` (the default) is one band at a time, which is what bounds the memory: each band in flight holds its own extraction working set, so the peak is this many bands' worth. Raise it to fill a large machine, after checking one band's peak RSS: on a 203M-precursor library at 63 bands the largest band took 39 GB and the median far less. Results do not depend on it; bands are independent and their artifacts are pooled in band order either way. It must stay below the thread count: a band in flight parks one worker on its accumulation channel, so as many bands as there are threads leaves nothing to do the probing and the run deadlocks. A larger value is clamped to `threads - 1` with a warning rather than hanging.", "gates": [], - "source_line": 1912 + "source_line": 1917 } ], "list_sections": [ diff --git a/docs/24_config_reference.md b/docs/24_config_reference.md index af4b85c7..74a90849 100644 --- a/docs/24_config_reference.md +++ b/docs/24_config_reference.md @@ -73,7 +73,7 @@ undocumented on purpose; those fields are counted under "Coverage". ## (top level) -`Config` (rust/mumdia/crates/mumdia-core/src/config.rs:1926). stage document: [docs/02_config_and_data_model.md](02_config_and_data_model.md). +`Config` (rust/mumdia/crates/mumdia-core/src/config.rs:1931). stage document: [docs/02_config_and_data_model.md](02_config_and_data_model.md). | Field | Type | Default | Gated | Description | |---|---|---|---|---| @@ -396,7 +396,7 @@ Searching a run one isolation-window group at a time. A group of isolation windo |---|---|---|---|---| | `window_groups` | `usize` | `1` | | Number of window groups. `1` (the default) is the ordinary single-library search. Groups are contiguous bands of isolation windows balanced by the number of library precursors they select, read from the precursor table's row-group statistics. | | `calibration` | `GroupCalibration` | `global` | | Anchors for the RT calibration of each group; see `GroupCalibration`. | -| `parallel` | `usize` | `1` | | Bands searched at the same time inside one run. `1` (the default) is one band at a time, which is what bounds the memory: each band in flight holds its own extraction working set, so the peak is this many bands' worth. Raise it to fill a large machine, after checking one band's peak RSS: on a 203M-precursor library at 63 bands the largest band took 39 GB and the median far less. Results do not depend on it; bands are independent and their artifacts are pooled in band order either way. | +| `parallel` | `usize` | `1` | | Bands searched at the same time inside one run. `1` (the default) is one band at a time, which is what bounds the memory: each band in flight holds its own extraction working set, so the peak is this many bands' worth. Raise it to fill a large machine, after checking one band's peak RSS: on a 203M-precursor library at 63 bands the largest band took 39 GB and the median far less. Results do not depend on it; bands are independent and their artifacts are pooled in band order either way. It must stay below the thread count: a band in flight parks one worker on its accumulation channel, so as many bands as there are threads leaves nothing to do the probing and the run deadlocks. A larger value is clamped to `threads - 1` with a warning rather than hanging. | ## peptidoforms.fixed_mods[] / peptidoforms.variable_mods[] @@ -783,6 +783,7 @@ listed with the file it is in. | `MUMDIA_NN_WARM_EPOCHS` | sidecar | `0` | `scripts/nn_rescore_worker.py:622` | | `MUMDIA_NN_WARM_START` | sidecar | `0` | `scripts/nn_rescore_worker.py:621` | | `MUMDIA_NN_WD` | sidecar | `1e-4` | `scripts/nn_rescore_worker.py:635` | +| `MUMDIA_PARQUET_COMPRESSION` | engine | none (unset means off) | `rust/mumdia/crates/mumdia-io/src/table.rs:37` | | `MUMDIA_PEPTDEEP_DEVICE` | sidecar | `"auto"` | `scripts/peptdeep_worker.py:214` | | `MUMDIA_PYTHON` | engine | none (unset means off) | `rust/mumdia/crates/mumdia/src/python.rs:233` | | `MUMDIA_PYTHON_DEEPLC` | engine | none (unset means off) | `rust/mumdia/crates/mumdia/src/python.rs:232` | @@ -790,7 +791,7 @@ listed with the file it is in. | `MUMDIA_PYTHON_MS2PIP` | engine | none (unset means off) | `rust/mumdia/crates/mumdia/src/python.rs:232` | | `MUMDIA_PYTHON_PEPTDEEP` | engine | none (unset means off) | `rust/mumdia/crates/mumdia/src/python.rs:232` | | `MUMDIA_PYTHON_RESCORE` | engine | none (unset means off) | `rust/mumdia/crates/mumdia/src/python.rs:232` | -| `MUMDIA_RESCORE_MODEL` | both | `"nn"` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:357`, `scripts/mokapot_worker.py:181`, `scripts/mokapot_worker.py:37` | +| `MUMDIA_RESCORE_MODEL` | both | `"nn"` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:488`, `scripts/mokapot_worker.py:181`, `scripts/mokapot_worker.py:37` | | `MUMDIA_SCRIPTS` | sidecar | `os.path.dirname(os.path.abspath(__file__` | `scripts/mh_shard_predict.py:23` | | `MUMDIA_THERMO_PARSER` | engine | none (unset means off) | `rust/mumdia/crates/mumdia/src/raw.rs:171` | | `MUMDIA_XGB_DEPTH` | sidecar | `"6"` | `scripts/mokapot_worker.py:63` | @@ -803,7 +804,7 @@ listed with the file it is in. | `ProgramFiles(x86)` | engine | none (unset means off) | `rust/mumdia/crates/mumdia/src/raw.rs:202` | | `VIRTUAL_ENV` | engine | none (unset means off) | `rust/mumdia/crates/mumdia/src/python.rs:242` | -66 variables are read: 16 engine-side, 53 sidecar-side, 3 on both sides. +67 variables are read: 17 engine-side, 53 sidecar-side, 3 on both sides. ### Variables the code sets @@ -816,23 +817,23 @@ one exception noted in its own help text: it sets `MUMDIA_NN_THREADS` and |---|---|---|---| | `KMP_DUPLICATE_LIB_OK` | sidecar | `"TRUE"` | `scripts/deeplc_finetune.py:28` | | `MKL_NUM_THREADS` | sidecar | `"1"` | `scripts/deeplc_finetune.py:31` | -| `MUMDIA_NN_FOLDS` | engine | `p.cfg.folds.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1391` | -| `MUMDIA_NN_FOLD_KEYS` | engine | `&foldkeys` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1418` | -| `MUMDIA_NN_ITERS` | engine | `p.cfg.num_iter.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1392` | -| `MUMDIA_NN_MARGIN_FRAC` | engine | `p.cfg.train_margin_frac.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1416` | -| `MUMDIA_NN_NEG_RATIO` | engine | `p.cfg.train_neg_ratio.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1397` | -| `MUMDIA_NN_NEG_SELECT` | engine | `match p.cfg.train_neg_select { mumdia_core::config::NegSelect::Random => "random", mumdia_core::config::NegSelect::Margin => "margin", mumdia_core::config::NegSelect::Hybrid => "hybrid", }` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1398` | -| `MUMDIA_NN_SEEDS` | engine | `p.cfg.seeds.max(1).to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1417` | +| `MUMDIA_NN_FOLDS` | engine | `p.cfg.folds.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1889` | +| `MUMDIA_NN_FOLD_KEYS` | engine | `foldkeys` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1916` | +| `MUMDIA_NN_ITERS` | engine | `p.cfg.num_iter.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1890` | +| `MUMDIA_NN_MARGIN_FRAC` | engine | `p.cfg.train_margin_frac.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1914` | +| `MUMDIA_NN_NEG_RATIO` | engine | `p.cfg.train_neg_ratio.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1895` | +| `MUMDIA_NN_NEG_SELECT` | engine | `match p.cfg.train_neg_select { mumdia_core::config::NegSelect::Random => "random", mumdia_core::config::NegSelect::Margin => "margin", mumdia_core::config::NegSelect::Hybrid => "hybrid", }` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1896` | +| `MUMDIA_NN_SEEDS` | engine | `p.cfg.seeds.max(1).to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1915` | | `MUMDIA_NN_THREADS` | engine | `n.to_string()` | `rust/mumdia/crates/mumdia/src/main.rs:94` | -| `MUMDIA_NN_TRAIN_FDR` | engine | `p.cfg.train_fdr.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1393` | -| `MUMDIA_NN_TRAIN_SUB` | engine | `p.cfg.train_subsample.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1406` | -| `MUMDIA_NN_WARM_EPOCHS` | engine | `p.cfg.train_warm_epochs.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1415` | -| `MUMDIA_NN_WARM_START` | engine | `if p.cfg.train_warm_epochs > 0 { "1" } else { "0" }` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1407` | +| `MUMDIA_NN_TRAIN_FDR` | engine | `p.cfg.train_fdr.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1891` | +| `MUMDIA_NN_TRAIN_SUB` | engine | `p.cfg.train_subsample.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1904` | +| `MUMDIA_NN_WARM_EPOCHS` | engine | `p.cfg.train_warm_epochs.to_string()` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1913` | +| `MUMDIA_NN_WARM_START` | engine | `if p.cfg.train_warm_epochs > 0 { "1" } else { "0" }` | `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1905` | | `NUMEXPR_NUM_THREADS` | sidecar | `"1"` | `scripts/deeplc_finetune.py:32` | | `OMP_NUM_THREADS` | both | `"1"` in deeplc_finetune.py; `n.to_string()` in main.rs | `rust/mumdia/crates/mumdia/src/main.rs:94`, `scripts/deeplc_finetune.py:29` | | `OPENBLAS_NUM_THREADS` | sidecar | `"1"` | `scripts/deeplc_finetune.py:30` | | `PYTHONIOENCODING` | engine | `"utf-8"` | `rust/mumdia/crates/mumdia/src/sidecar.rs:548` | -| `PYTHONUTF8` | engine | `"1"` | `rust/mumdia/crates/mumdia/src/sidecar.rs:548`, `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1117`, `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1385` | +| `PYTHONUTF8` | engine | `"1"` | `rust/mumdia/crates/mumdia/src/sidecar.rs:548`, `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1352`, `rust/mumdia/crates/mumdia/src/stages/rescore.rs:1883` | ## Unresolved by the generator @@ -848,11 +849,11 @@ Every field whose struct has an `impl Default` resolved from the source. 2 environment read(s) whose name is not a literal: -- `rust/mumdia/crates/mumdia/src/stages/extract.rs:2760: env read via closure of `&mut flushed`` -- `rust/mumdia/crates/mumdia/src/stages/extract.rs:2778: env read via closure of `&mut cand_hits`` +- `rust/mumdia/crates/mumdia/src/stages/extract.rs:2786: env read via closure of `&mut flushed`` +- `rust/mumdia/crates/mumdia/src/stages/extract.rs:2804: env read via closure of `&mut cand_hits`` ## Coverage -19 structs and 193 fields emitted from `rust/mumdia/crates/mumdia-core/src/config.rs`, plus 25 enumerations, 1 named profile(s), 66 environment variables read and 19 set. +19 structs and 193 fields emitted from `rust/mumdia/crates/mumdia-core/src/config.rs`, plus 25 enumerations, 1 named profile(s), 67 environment variables read and 19 set. 20 field(s) carry a gating marker in their doc comment. 48 field(s) carry no doc comment at all, so their description is empty above. 0 default(s) could not be resolved and 2 have none by design. From 36fdf597cd566e6b26fb166a48c42575f8caf079 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:20:03 +0200 Subject: [PATCH 29/35] perf(extract): cut the per-candidate allocations and the per-candidate window scan Four scratch structures in the per-candidate pass, plus a linear scan over every isolation window. - The distinct-fragment set was `hits.iter().map(|h| h.frag).collect::>()` followed by a sort and a dedup: one `u16` PER HIT, and a candidate with no `run_windows` row collects 10^4-10^5 hits, so that was a 20-200 KB allocation and an n log n sort to recover at most a few dozen distinct ordinals. It is now a bitmask over the candidate's local fragment ordinals (`FragSet`, 512 inline, spilling to a sorted vector above that, which no real library reaches). The sorted vector is materialised only where the co-elution score asks for it, which the default gate mode does not. - The grid alignment built a `HashMap` of the scan grid per candidate. Both sides are ascending, so it is a merge; the match is still on the exact bit pattern, so a group whose RT is not a grid RT is dropped exactly as before. - The chromatogram emission built a `HashMap>` over the whole candidate and then a `HashMap` per fragment to sample it onto the grid. In grid mode `groups` IS the grid -- it was rebuilt on it above, one entry per grid RT in the same order -- so a fragment's trace is its value in each group, read directly. The RT axis is the grid itself, the same for every fragment and for the MS1 XICs, so it is converted once instead of once per fragment. In sparse mode the trace is the fragment's entries in group order, which is what the per-fragment vector held before its already-sorted stable sort. - `wsum` (per-fragment intensity-weighted observed m/z) was a `HashMap` over ordinals that are `0..n_frag` by construction; it is now indexed. An untouched entry keeps `sum_w == 0`, which takes the same theoretical-m/z fallback the absent map entry took. - The scan grid scanned all ~150 windows per candidate to find the covering ones. `windows` is sorted by lower m/z, and a covering window has `lower_mz <= pm` and `lower_mz >= upper_mz - max_width >= pm - max_width`, so the span is two binary searches; the membership test inside it is unchanged. With one covering window the slice is already ascending, so the sort is skipped (the dedup still runs). No speed claim from the fixture available here. On the AIF benchmark (1,691,048 candidates, 465,806 scans, 152 windows, 728,814 candidates with evidence, 41,677 accepted) the extract compute phase is 2.12-2.16 s with and without these changes over five warm runs each, i.e. flat within noise. That fixture cannot resolve them: it is I/O-dominated (4.0 s total, ~1.8 s of it loading a 309 MB spectra table), its `run_windows` covers every candidate so the 10^4-10^5-hit population these target never occurs, and it averages about two hits per candidate. The first numbers measured for this change were cold-cache and are not real. Risk to output equality: none intended. Verified end to end against the base binary on that fixture, both `--release` from the same tree: `psms_extracted.parquet` is BYTE-IDENTICAL, and `chromatograms.parquet` is equal as a table (620,169 rows, every column, same order, `pyarrow` equality) -- its bytes differ only because the preceding commit moves the row-group boundaries, which the flush-unit change documents. Tests: `FragSet` is asserted against sort+dedup over several ordinal patterns, including the spill past the inline mask; the covering-window span is asserted to select exactly the windows a full scan selects, over uneven overlapping windows, a gap, and 2,000 precursor m/z probes plus every window edge. Co-Authored-By: Claude Opus 5 (1M context) --- .../crates/mumdia/src/stages/extract.rs | 311 +++++++++++++++--- 1 file changed, 257 insertions(+), 54 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/extract.rs b/rust/mumdia/crates/mumdia/src/stages/extract.rs index 98abe675..633aef1f 100644 --- a/rust/mumdia/crates/mumdia/src/stages/extract.rs +++ b/rust/mumdia/crates/mumdia/src/stages/extract.rs @@ -461,6 +461,69 @@ struct Contested { apportioned: f64, } +/// Ordinals a [`FragSet`] holds in its inline bitmask; above this it spills to a sorted +/// vector. A candidate carries a few dozen predicted fragments (6 by default, 12 in the +/// shipped DIA-NN library), so the spill is unreachable in practice and costs nothing when +/// it is not used (an empty `Vec` does not allocate). +const FRAG_SET_BITS: usize = 512; + +/// A set of a candidate's local fragment ordinals. +/// +/// This replaces `hits.iter().map(|h| h.frag).collect::>()` followed by a sort and +/// a dedup, which allocated one `u16` per HIT and sorted it: a candidate with no +/// `run_windows` row collects 10^4-10^5 hits, so that was a 20-200 KB allocation and an +/// `n log n` sort to recover at most a few dozen distinct values. +#[derive(Default)] +struct FragSet { + bits: [u64; FRAG_SET_BITS / 64], + spill: Vec, +} + +impl FragSet { + #[inline] + fn insert(&mut self, f: u16) { + let i = f as usize; + if i < FRAG_SET_BITS { + self.bits[i >> 6] |= 1u64 << (i & 63); + } else if let Err(p) = self.spill.binary_search(&f) { + self.spill.insert(p, f); + } + } + + #[inline] + fn contains(&self, f: u16) -> bool { + let i = f as usize; + if i < FRAG_SET_BITS { + self.bits[i >> 6] & (1u64 << (i & 63)) != 0 + } else { + self.spill.binary_search(&f).is_ok() + } + } + + fn len(&self) -> usize { + self.bits + .iter() + .map(|w| w.count_ones() as usize) + .sum::() + + self.spill.len() + } + + /// The ordinals, ascending: exactly what the sorted, deduplicated vector held. + fn to_vec(&self) -> Vec { + let mut v = Vec::with_capacity(self.len()); + for (w, &word) in self.bits.iter().enumerate() { + let mut word = word; + while word != 0 { + let b = word.trailing_zeros() as usize; + v.push((w * 64 + b) as u16); + word &= word - 1; + } + } + v.extend_from_slice(&self.spill); + v + } +} + /// Index of the value in ascending `rts` nearest to `t` (binary search). fn nearest_index(rts: &[f64], t: f64) -> usize { if rts.is_empty() { @@ -1960,6 +2023,15 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { } else { Vec::new() }; + // Widest isolation window, so the per-candidate scan over `windows` (sorted by lower + // m/z) can binary-search the span that can possibly cover a precursor instead of + // walking all of them: a window covering `pm` has `lower_mz <= pm` and + // `lower_mz >= upper_mz - max_width >= pm - max_width`. + let window_max_width = + windows + .iter() + .map(|w| w.1 - w.0) + .fold(0.0f64, |a, b| if b > a { b } else { a }); // Per-run mass recalibration (optional). Reads the scalar offset + learned // tolerance, plus an optional m/z-dependent correction grid (mass_cal_loess). @@ -2353,11 +2425,14 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { }; let per_candidate = |cid: u32, hits: &mut [Hit]| -> Vec { - // distinct matched fragments (tier b) - let mut distinct: Vec = hits.iter().map(|h| h.frag).collect(); - distinct.sort_unstable(); - distinct.dedup(); - if distinct.len() < p.cfg.presence_min_matched.max(1) { + // distinct matched fragments (tier b), as a bitmask over the candidate's local + // fragment ordinals rather than a sorted, deduplicated `u16` per hit. + let mut distinct = FragSet::default(); + for h in hits.iter() { + distinct.insert(h.frag); + } + let n_distinct = distinct.len(); + if n_distinct < p.cfg.presence_min_matched.max(1) { return Vec::new(); } @@ -2392,31 +2467,49 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { let grid: Vec = if !windows.is_empty() { let pm = lib.cands[cid as usize].precursor_mz; let (lo, hi) = (rt_lo[cid as usize], rt_hi[cid as usize]); + // `windows` is sorted by lower m/z, so a covering window has + // `lower_mz <= pm` AND, since its width is at most `window_max_width`, + // `lower_mz >= pm - window_max_width`. Binary-search that span instead of + // scanning all ~150 windows per candidate; the membership test inside the + // span is unchanged, so the same windows contribute in the same order. + let s = windows.partition_point(|w| w.0 < pm - window_max_width); + let e = windows.partition_point(|w| w.0 <= pm); let mut g: Vec = Vec::new(); - for (wl, wu, rts) in &windows { + let mut covering = 0usize; + for (wl, wu, rts) in &windows[s..e] { if *wl <= pm && pm <= *wu { let a = rts.partition_point(|&r| r < lo); let b = rts.partition_point(|&r| r <= hi); g.extend_from_slice(&rts[a..b]); + covering += 1; } } - g.sort_by(|a, b| a.total_cmp(b)); + // One covering window contributes one already-ascending slice of its own + // sorted scan RTs, which is the overwhelmingly common case; only an overlap + // of two windows needs the merge sorting. + if covering > 1 { + g.sort_by(|a, b| a.total_cmp(b)); + } g.dedup(); g } else { Vec::new() }; if !grid.is_empty() { - let g2i: HashMap = grid - .iter() - .enumerate() - .map(|(j, r)| (r.to_bits(), j)) - .collect(); let mut aligned: Vec<(f64, BTreeMap)> = grid.iter().map(|&r| (r, BTreeMap::new())).collect(); + // Both sides are ascending (`grid` is sorted and deduplicated; the scan groups + // were built from rt-sorted hits), so this is a merge rather than a per- + // candidate `HashMap` of the grid. The match is still on the exact bit + // pattern, so a group whose RT is not a grid RT is dropped exactly as before. + let mut j = 0usize; for (rt, map) in std::mem::take(&mut groups) { - if let Some(&j) = g2i.get(&rt.to_bits()) { + while j < grid.len() && grid[j] < rt { + j += 1; + } + if j < grid.len() && grid[j].to_bits() == rt.to_bits() { aligned[j].1 = map; + j += 1; } } groups = aligned; @@ -2551,8 +2644,8 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // Acceptance (tier c): presence, consecutive-scan run, and matched // fraction of the predicted fragments (symmetric discriminator). - let matched_fraction = distinct.len() as f64 / (fmzs0.len().max(1) as f64); - if distinct.len() < p.cfg.presence_min_fragments.max(1) + let matched_fraction = n_distinct as f64 / (fmzs0.len().max(1) as f64); + if n_distinct < p.cfg.presence_min_fragments.max(1) || best_run < scan_window || best_run < p.cfg.min_coelution_run || matched_fraction < p.cfg.min_matched_fraction @@ -2628,7 +2721,9 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { None => 1.0, }; let peak_spec = || peak_spectral_score(&groups, &sig, fints0); - let coel = || coelution_gate_score(&groups, &distinct, &sig, fints0); + // `to_vec` only where the co-elution score is actually asked for (a non-default + // gate mode, or the diagnostics), so the default chain never materialises it. + let coel = || coelution_gate_score(&groups, &distinct.to_vec(), &sig, fints0); if p.cfg.gate_min_score > 0.0 { // Acceptance gate. `gate_min_score` thresholds the ACTIVE gate_mode's @@ -2648,7 +2743,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { if rejected { let rescued = p.cfg.ms1_rescue && ms1_support - && distinct.len() >= p.cfg.presence_min_fragments.max(1); + && n_distinct >= p.cfg.presence_min_fragments.max(1); if !rescued { return Vec::new(); } @@ -2699,12 +2794,19 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { } }; - // Per-fragment intensity-weighted observed m/z (for mass accuracy). - let mut wsum: HashMap = HashMap::new(); // frag -> (sum w*mz, sum w) + let (fmzs, fints, fnames) = lib.cand_frags(cid); + + // Per-fragment intensity-weighted observed m/z (for mass accuracy). Indexed by + // the candidate-local ordinal rather than hashed: a fragment ordinal is + // `0..n_frag` by construction, so the map was a hash table over a dense range. + // An untouched entry keeps `sum_w == 0`, which takes the same theoretical-m/z + // fallback the absent map entry took. + let mut wsum: Vec<(f64, f64)> = vec![(0.0, 0.0); fmzs.len()]; // (sum w*mz, sum w) for h in hits.iter() { - let e = wsum.entry(h.frag).or_insert((0.0, 0.0)); - e.0 += h.obs_mz * h.inten as f64; - e.1 += h.inten as f64; + if let Some(e) = wsum.get_mut(h.frag as usize) { + e.0 += h.obs_mz * h.inten as f64; + e.1 += h.inten as f64; + } } let mut chrom_rows: Vec = Vec::new(); @@ -2713,14 +2815,21 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // on the full isolation-window scan grid (all scans of the covering window(s) // within the RT window), with 0.0 where the fragment is absent, so the elution // profile drops to zero between peaks (correct boundary calling downstream). - let (fmzs, fints, fnames) = lib.cand_frags(cid); - let mut per_frag: HashMap> = HashMap::new(); - for (rt, map) in &groups { - for (&frag, &inten) in map { - per_frag.entry(frag).or_default().push((*rt, inten)); - } - } - // (the acquisition-scan `grid` was computed above, before apex/co-elution) + // + // The traces are read straight off the scan groups. In grid mode `groups` IS the + // grid (it was rebuilt on it above, one entry per grid RT in the same order), so a + // fragment's grid-sampled trace is its value in each group, and the RT axis is the + // grid itself and the same for every fragment and for the MS1 XICs below: one + // conversion instead of one per fragment. That replaces a `HashMap>` over the whole candidate plus a `HashMap` per + // fragment, and produces the same values in the same order. + let mut observed = FragSet::default(); + for (_, m) in &groups { + for &f in m.keys() { + observed.insert(f); + } + } + let grid_rt: Vec = grid.iter().map(|&r| r as f32).collect(); // Emit a row for EVERY predicted transition so the feature families see the // full predicted set (a missing strong ion is penalized). An OBSERVED // fragment carries its grid-sampled (or sorted) trace; a NEVER-OBSERVED one @@ -2733,29 +2842,31 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // only fragments with obs_apex > 0. for fi in 0..fmzs.len() { let frag = fi as u16; - let obs_mz = wsum - .get(&frag) - .map(|(sm, sw)| if *sw > 0.0 { sm / sw } else { fmzs[fi] as f64 }) - .unwrap_or(fmzs[fi] as f64); - let (rts, ints): (Vec, Vec) = match per_frag.get(&frag) { - Some(v) if !grid.is_empty() => { - let m: HashMap = v.iter().map(|(r, i)| (r.to_bits(), *i)).collect(); - ( - grid.iter().map(|r| *r as f32).collect(), - grid.iter() - .map(|r| *m.get(&r.to_bits()).unwrap_or(&0.0)) - .collect(), - ) - } - Some(v) => { - let mut s = v.clone(); - s.sort_by(|a, b| a.0.total_cmp(&b.0)); - ( - s.iter().map(|(r, _)| *r as f32).collect(), - s.iter().map(|(_, i)| *i).collect(), - ) + let (sm, sw) = wsum[fi]; + let obs_mz = if sw > 0.0 { sm / sw } else { fmzs[fi] as f64 }; + let (rts, ints): (Vec, Vec) = if !observed.contains(frag) { + (Vec::new(), Vec::new()) // absent predicted transition + } else if !grid.is_empty() { + ( + grid_rt.clone(), + groups + .iter() + .map(|(_, m)| *m.get(&frag).unwrap_or(&0.0)) + .collect(), + ) + } else { + // Sparse mode: the groups are already ascending in RT, so the trace is + // the fragment's entries read in group order -- what the per-fragment + // vector held before its (already-sorted) stable sort. + let mut rts: Vec = Vec::new(); + let mut ints: Vec = Vec::new(); + for (rt, m) in &groups { + if let Some(&i) = m.get(&frag) { + rts.push(*rt as f32); + ints.push(i); + } } - None => (Vec::new(), Vec::new()), // absent predicted transition + (rts, ints) }; chrom_rows.push(( cid, @@ -2778,7 +2889,6 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { if !ms1_scans.is_empty() && !grid.is_empty() { let sp = ISOTOPE_SPACING / c.charge as f64; let tol = p.cfg.prec_tol_ppm; - let grid_rt: Vec = grid.iter().map(|&r| r as f32).collect(); for (nm, dmz) in [("ms1_mono", 0.0), ("ms1_iso1", sp), ("ms1_iso2", 2.0 * sp)] { let mz = c.precursor_mz + dmz; let ints: Vec = grid @@ -2832,7 +2942,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { peak_rank: 0, // selected apex; ranks >= 1 added when promote_top_peaks > 1 apex_rt, apex_int: apex_sum, - n_match: distinct.len() as i32, + n_match: n_distinct as i32, corun: best_run as i32, npred: fmzs0.len() as i32, calrt: rt_cal[cid as usize], @@ -3369,6 +3479,99 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { Ok((n_psms, n_chrom)) } +#[cfg(test)] +mod frag_set_tests { + use super::{FragSet, FRAG_SET_BITS}; + + /// The bitmask replaces "collect every hit's ordinal, sort, dedup", so it has to + /// produce that vector, and its length, for any input. + #[test] + fn reproduces_sort_and_dedup() { + for &case in &[ + &[][..], + &[0][..], + &[7, 7, 7][..], + &[5, 1, 9, 1, 63, 64, 0, 9][..], + &[511, 510, 0, 1, 64, 65, 128, FRAG_SET_BITS as u16 - 1][..], + ] { + let mut set = FragSet::default(); + for &f in case { + set.insert(f); + } + let mut reference: Vec = case.to_vec(); + reference.sort_unstable(); + reference.dedup(); + assert_eq!(set.len(), reference.len(), "{case:?}"); + assert_eq!(set.to_vec(), reference, "{case:?}"); + for f in 0..600u16 { + assert_eq!(set.contains(f), reference.contains(&f), "{case:?} frag {f}"); + } + } + } + + /// Ordinals past the inline mask spill to a sorted vector rather than panicking, and + /// the two halves still read back as one ascending set. + #[test] + fn ordinals_past_the_mask_spill_in_order() { + let case: Vec = vec![900, 3, 512, 900, 1200, 511, 512]; + let mut set = FragSet::default(); + for &f in &case { + set.insert(f); + } + let mut reference = case.clone(); + reference.sort_unstable(); + reference.dedup(); + assert_eq!(set.to_vec(), reference); + assert_eq!(set.len(), reference.len()); + assert!(set.contains(1200) && set.contains(3) && !set.contains(4)); + } +} + +#[cfg(test)] +mod covering_window_tests { + /// The per-candidate grid used to scan every isolation window. The binary-searched + /// span must select exactly the windows the scan did, for uneven widths and for + /// precursors falling in a gap, or a candidate silently loses part of its scan grid. + #[test] + fn the_searched_span_holds_every_covering_window() { + // Deliberately uneven: overlapping wide and narrow windows, plus a gap. + let mut windows: Vec<(f64, f64)> = vec![ + (400.0, 404.0), + (402.0, 426.0), + (404.0, 408.0), + (408.0, 412.0), + (412.0, 413.0), + (420.0, 430.0), + (430.0, 470.0), + (460.0, 461.0), + ]; + windows.sort_by(|a, b| a.0.total_cmp(&b.0)); + let max_width = windows + .iter() + .map(|w| w.1 - w.0) + .fold(0.0f64, |a, b| if b > a { b } else { a }); + let mut pms: Vec = Vec::new(); + for i in 0..2000 { + pms.push(395.0 + i as f64 * 0.05); + } + for &w in &windows { + pms.push(w.0); + pms.push(w.1); + } + for pm in pms { + let full: Vec = (0..windows.len()) + .filter(|&i| windows[i].0 <= pm && pm <= windows[i].1) + .collect(); + let s = windows.partition_point(|w| w.0 < pm - max_width); + let e = windows.partition_point(|w| w.0 <= pm); + let span: Vec = (s..e) + .filter(|&i| windows[i].0 <= pm && pm <= windows[i].1) + .collect(); + assert_eq!(span, full, "pm {pm}"); + } + } +} + #[cfg(test)] mod mass_offset_tests { use super::MassOffset; From 8b52941fa3af03621df5baaf7f532bc9b3c91723 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:37:09 +0200 Subject: [PATCH 30/35] ci: regenerate the CLI reference and licenses, and stop the generator trusting a stale in-tree binary The committed CLI reference was missing the whole `pool` subcommand: the generator prefers `rust/mumdia/target/release/mumdia` over the target directory cargo actually reports, and on a machine whose `build.target-dir` is redirected off the synced tree that path holds whatever was built before the redirect. Here it was a binary from 2026-09-09. CI has no such directory, which is the only reason it noticed. The generator now asks cargo first, from inside the workspace, the way ci/smoke.sh already does. Also: the third-party licence inventory picks up the zstd crates, and a public doc comment in mumdia-io no longer links to a private item, which cargo doc refuses under -D warnings. Co-Authored-By: Claude Opus 5 (1M context) --- THIRD_PARTY_LICENSES.md | 49 ++++++++++++++++++++--- ci/gen_cli_reference.py | 38 ++++++++++++++++-- docs/23_cli_reference.md | 33 +++++++++++++-- rust/mumdia/crates/mumdia-io/src/table.rs | 2 +- 4 files changed, 110 insertions(+), 12 deletions(-) diff --git a/THIRD_PARTY_LICENSES.md b/THIRD_PARTY_LICENSES.md index db3cf8ce..ae543258 100644 --- a/THIRD_PARTY_LICENSES.md +++ b/THIRD_PARTY_LICENSES.md @@ -9,7 +9,7 @@ accompanies it in every release archive and in the container image. ## Obligations -177 third-party crates. Every one declares an SPDX expression; none is +184 third-party crates. Every one declares an SPDX expression; none is unspecified. The following copyleft identifiers appear, in each case as one arm of a @@ -24,17 +24,17 @@ Licence identifiers by crate count: | SPDX identifier | crates | |---|---| -| `MIT` | 151 | -| `Apache-2.0` | 146 | +| `MIT` | 158 | +| `Apache-2.0` | 152 | | `Apache-2.0 WITH LLVM-exception` | 4 | | `Unlicense` | 4 | | `CC0-1.0` | 3 | | `BSD-2-Clause` | 2 | | `BSD-3-Clause` | 2 | +| `LGPL-2.1-or-later` | 2 | | `Zlib` | 2 | | `0BSD` | 1 | | `BSL-1.0` | 1 | -| `LGPL-2.1-or-later` | 1 | | `MIT-0` | 1 | | `Unicode-3.0` | 1 | @@ -112,6 +112,7 @@ Licence identifiers by crate count: | `futures-util` | 0.3.32 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-lang/futures-rs](https://github.com/rust-lang/futures-rs) | | `getrandom` | 0.2.17 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-random/getrandom](https://github.com/rust-random/getrandom) | | `getrandom` | 0.3.4 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-random/getrandom](https://github.com/rust-random/getrandom) | +| `getrandom` | 0.4.3 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-random/getrandom](https://github.com/rust-random/getrandom) | | `half` | 2.7.1 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/VoidStarKat/half-rs](https://github.com/VoidStarKat/half-rs) | | `hashbrown` | 0.17.1 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-lang/hashbrown](https://github.com/rust-lang/hashbrown) | | `heck` | 0.5.0 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/withoutboats/heck](https://github.com/withoutboats/heck) | @@ -123,6 +124,7 @@ Licence identifiers by crate count: | `indexmap` | 2.14.0 | Apache-2.0 OR MIT | `Apache-2.0` | [https://github.com/indexmap-rs/indexmap](https://github.com/indexmap-rs/indexmap) | | `is_terminal_polyfill` | 1.70.2 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/polyfill-rs/is_terminal_polyfill](https://github.com/polyfill-rs/is_terminal_polyfill) | | `itoa` | 1.0.18 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/dtolnay/itoa](https://github.com/dtolnay/itoa) | +| `jobserver` | 0.1.35 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-lang/jobserver-rs](https://github.com/rust-lang/jobserver-rs) | | `js-sys` | 0.3.103 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/wasm-bindgen/wasm-bindgen/tree/master/crates/js-sys](https://github.com/wasm-bindgen/wasm-bindgen/tree/master/crates/js-sys) | | `lazy_static` | 1.5.0 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-lang-nursery/lazy-static.rs](https://github.com/rust-lang-nursery/lazy-static.rs) | | `lexical-core` | 1.0.6 | MIT/Apache-2.0 | `Apache-2.0` | [https://github.com/Alexhuszagh/rust-lexical](https://github.com/Alexhuszagh/rust-lexical) | @@ -155,10 +157,12 @@ Licence identifiers by crate count: | `outref` | 0.5.2 | MIT | - | [https://github.com/Nugine/outref](https://github.com/Nugine/outref) | | `parquet` | 59.3.0 | Apache-2.0 | - | [https://github.com/apache/arrow-rs](https://github.com/apache/arrow-rs) | | `pin-project-lite` | 0.2.17 | Apache-2.0 OR MIT | `Apache-2.0` | [https://github.com/taiki-e/pin-project-lite](https://github.com/taiki-e/pin-project-lite) | +| `pkg-config` | 0.3.34 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-lang/pkg-config-rs](https://github.com/rust-lang/pkg-config-rs) | | `proc-macro2` | 1.0.106 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/dtolnay/proc-macro2](https://github.com/dtolnay/proc-macro2) | | `quick-xml` | 0.41.0 | MIT | - | [https://github.com/tafia/quick-xml](https://github.com/tafia/quick-xml) | | `quote` | 1.0.46 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/dtolnay/quote](https://github.com/dtolnay/quote) | | `r-efi` | 5.3.0 | MIT OR Apache-2.0 OR LGPL-2.1-or-later | `Apache-2.0` | [https://github.com/r-efi/r-efi](https://github.com/r-efi/r-efi) | +| `r-efi` | 6.0.0 | MIT OR Apache-2.0 OR LGPL-2.1-or-later | `Apache-2.0` | [https://github.com/r-efi/r-efi](https://github.com/r-efi/r-efi) | | `rayon` | 1.12.0 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rayon-rs/rayon](https://github.com/rayon-rs/rayon) | | `rayon-core` | 1.13.0 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rayon-rs/rayon](https://github.com/rayon-rs/rayon) | | `regex` | 1.12.4 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/rust-lang/regex](https://github.com/rust-lang/regex) | @@ -219,6 +223,9 @@ Licence identifiers by crate count: | `zerocopy` | 0.8.52 | BSD-2-Clause OR Apache-2.0 OR MIT | `Apache-2.0` | [https://github.com/google/zerocopy](https://github.com/google/zerocopy) | | `zerocopy-derive` | 0.8.52 | BSD-2-Clause OR Apache-2.0 OR MIT | `Apache-2.0` | [https://github.com/google/zerocopy](https://github.com/google/zerocopy) | | `zmij` | 1.0.21 | MIT | - | [https://github.com/dtolnay/zmij](https://github.com/dtolnay/zmij) | +| `zstd` | 0.13.3 | MIT | - | [https://github.com/gyscos/zstd-rs](https://github.com/gyscos/zstd-rs) | +| `zstd-safe` | 7.2.4 | MIT OR Apache-2.0 | `Apache-2.0` | [https://github.com/gyscos/zstd-rs](https://github.com/gyscos/zstd-rs) | +| `zstd-sys` | 2.0.16+zstd.1.5.7 | MIT/Apache-2.0 | `Apache-2.0` | [https://github.com/gyscos/zstd-rs](https://github.com/gyscos/zstd-rs) | ## Licence texts @@ -330,7 +337,7 @@ licence text alone does not discharge the obligation: MIT requires the copyright notice in all copies, BSD requires it retained, and Apache-2.0 section 4(d) requires the contents of a NOTICE file to be propagated. -Notices recovered for 126 of 177 crates (the remainder ship no copyright line in a licence file; see their repositories, linked above). +Notices recovered for 133 of 184 crates (the remainder ship no copyright line in a licence file; see their repositories, linked above). **adler2 2.0.1** @@ -551,6 +558,11 @@ Notices recovered for 126 of 177 crates (the remainder ship no copyright line in - Copyright (c) 2018-2025 The rust-random Project Developers - Copyright (c) 2014 The Rust Project Developers +**getrandom 0.4.3** + +- Copyright (c) 2018-2026 The rust-random Project Developers +- Copyright (c) 2014 The Rust Project Developers + **hashbrown 0.17.1** - Copyright (c) 2016 Amanieu d'Antras @@ -591,6 +603,10 @@ Notices recovered for 126 of 177 crates (the remainder ship no copyright line in - Copyright (c) Individual contributors +**jobserver 0.1.35** + +- Copyright (c) 2014 Alex Crichton + **js-sys 0.3.103** - Copyright (c) 2014 Alex Crichton @@ -701,6 +717,10 @@ Notices recovered for 126 of 177 crates (the remainder ship no copyright line in - Copyright (c) 2022 Nugine +**pkg-config 0.3.34** + +- Copyright (c) 2014 Alex Crichton + **quick-xml 0.41.0** - Copyright (c) 2016 Johann Tuffe @@ -711,6 +731,12 @@ Notices recovered for 126 of 177 crates (the remainder ship no copyright line in - Copyright (C) 2019-2023 Microsoft Corporation - Copyright (C) 2022-2023 David Rheinsberg +**r-efi 6.0.0** + +- Copyright (C) 2017-2023 Red Hat, Inc. +- Copyright (C) 2019-2023 Microsoft Corporation +- Copyright (C) 2022-2023 David Rheinsberg + **rayon 1.12.0** - Copyright (c) 2010 The Rust Project Developers @@ -880,6 +906,19 @@ Notices recovered for 126 of 177 crates (the remainder ship no copyright line in - Copyright 2023 The Fuchsia Authors - Copyright 2019 The Fuchsia Authors. +**zstd 0.13.3** + +- Copyright (c) 2016 Alexandre Bury + +**zstd-safe 7.2.4** + +- Copyright (c) 2016 Alexandre Bury + +**zstd-sys 2.0.16+zstd.1.5.7** + +- Copyright (c) 2016-present, Facebook, Inc. All rights reserved. +- Copyright (c) 2016 Alexandre Bury + ### NOTICE files, verbatim Apache-2.0 section 4(d) requires these to travel with the distribution. diff --git a/ci/gen_cli_reference.py b/ci/gen_cli_reference.py index 7ac901bb..b676e43a 100644 --- a/ci/gen_cli_reference.py +++ b/ci/gen_cli_reference.py @@ -48,6 +48,7 @@ import os import re import shutil +import json import subprocess import sys from pathlib import Path @@ -88,6 +89,31 @@ # --------------------------------------------------------------------------- +def cargo_target_binaries() -> list[Path]: + """`release/mumdia` under the target directory cargo reports, if cargo answers. + + Run from inside the workspace, not with `--manifest-path`: cargo discovers + `.cargo/config.toml` by walking up from the CURRENT directory, so from the repository + root it reports the in-tree `target/` and misses the redirect this exists to follow. + """ + ws = REPO_ROOT / "rust" / "mumdia" + if not shutil.which("cargo") or not ws.is_dir(): + return [] + try: + out = subprocess.run( + ["cargo", "metadata", "--format-version", "1", "--no-deps"], + cwd=ws, + capture_output=True, + text=True, + timeout=120, + check=True, + ).stdout + target = Path(json.loads(out)["target_directory"]) + except (subprocess.SubprocessError, OSError, ValueError, KeyError): + return [] + return [target / "release" / "mumdia", target / "release" / "mumdia.exe"] + + def find_binary(explicit: str | None) -> Path: """Resolve the binary to interrogate, or exit with the search order shown.""" if explicit: @@ -101,10 +127,16 @@ def find_binary(explicit: str | None) -> Path: if not p.is_file(): sys.exit(f"error: $MUMDIA_BIN={from_env} is not a file") return p - for rel in BIN_CANDIDATES: - p = Path(rel) + # The target directory cargo actually reports comes before the in-tree path. A + # redirected `build.target-dir` -- which `.cargo/config.toml.example` documents, to keep + # build output out of a synced folder -- leaves `rust/mumdia/target/` holding whatever + # was built before the redirect, and this generator then documents a months-old CLI. + # That is not hypothetical: it shipped a reference missing a whole subcommand, and CI + # caught it only because CI has no such stale directory. `ci/smoke.sh` avoids the same + # trap the same way. + for p in (cargo_target_binaries() + [Path(rel) for rel in BIN_CANDIDATES if rel]): if not p.is_absolute(): - p = REPO_ROOT / rel + p = REPO_ROOT / p if p.is_file(): return p on_path = shutil.which(PROG) diff --git a/docs/23_cli_reference.md b/docs/23_cli_reference.md index 2945d944..837665c8 100644 --- a/docs/23_cli_reference.md +++ b/docs/23_cli_reference.md @@ -45,6 +45,7 @@ Commands: rt-im-train Per-run RT calibration + windows -> run_windows.parquet, cal.json extract Targeted 3D extraction (peak-major cascade) -> psms_extracted, chromatograms features Compute the minimal feature set -> features.parquet + PIN + pool Pool a grouped run's band artifacts into the run-level tables compete Keep the best candidate per competition group -> psms_competed.parquet rescore Rescore + native target-decoy q-values -> psms_scored.parquet quant Quantify identified peptides + roll up to protein groups @@ -137,6 +138,7 @@ first sentence of the description, with the full text in the section below. | [`rt-im-train`](#rt-im-train) | yes | Per-run RT calibration + windows -> run_windows.parquet, cal.json | | [`extract`](#extract) | yes | Targeted 3D extraction (peak-major cascade) -> psms_extracted, chromatograms | | [`features`](#features) | yes | Compute the minimal feature set -> features.parquet + PIN | +| [`pool`](#pool) | no | Pool a grouped run's band artifacts into the run-level tables | | [`compete`](#compete) | yes | Keep the best candidate per competition group -> psms_competed.parquet | | [`rescore`](#rescore) | yes | Rescore + native target-decoy q-values -> psms_scored.parquet | | [`quant`](#quant) | yes | Quantify identified peptides + roll up to protein groups | @@ -152,11 +154,11 @@ first sentence of the description, with the full text in the section below. | [`doctor`](#doctor) | yes | Check that the configured Python sidecar environments are usable | | `help` | n/a | Print this message or the help of the given subcommand(s) | -19 of the 22 documented subcommands accept `--config`: +19 of the 23 documented subcommands accept `--config`: `align`, `compete`, `convert`, `digest`, `doctor`, `extract`, `features`, `mbr`, `peak-census`, `peptidoforms`, `predict-frag`, `prescan`, `quant`, `report`, `rescore`, `rt-im-train`, `run`, `run-experiment`, `search-seed`. -3 do not, so every setting they use comes from their own flags: - `audit`, `inspect`, `quant-lfq`. +4 do not, so every setting they use comes from their own flags: + `audit`, `inspect`, `pool`, `quant-lfq`. ## convert @@ -347,6 +349,9 @@ Options: --restrict-candidates Optional candidate allowlist (a prior run's psms.parquet): restrict extraction to these candidate_ids. For "gate first, then compete" - re-extract with a peak_claim strategy over only the gate-accepted survivors, keeping the two-pass profile map small + --fragment-offset + The library row that `--lib-precursors` row 0 came from, when that table is one isolation-window band of a larger library (`groups.window_groups` writes such bands, with ids rebased to `0..n`). The band's fragments are then read from the shared `--lib-fragments` table by that id range, selectively when the table is sorted by `candidate_id`, so a band needs no fragment table of its own + --config ``` @@ -376,6 +381,28 @@ Options: Plus the 5 repeated flags removed above: see "Global flags". +## pool + +```text +Pool a grouped run's band artifacts into the run-level tables. + +`run` does this itself at the end of a grouped search (`groups.window_groups`). Standalone it is for the case where the search finished and the run did not: the band directories hold everything, and pooling them is a byte copy of their parquet row groups, so a killed run costs a pool rather than a re-search. + +Usage: mumdia pool [OPTIONS] --groups-dir + +Options: + --groups-dir + The run's `groups/` directory, holding the `gNN/` band directories + + --out-dir + Where the pooled tables go. Default: the parent of `--groups-dir`, which is where a run writes them + + --psms + Also pool `psms_extracted`, which only the candidate audit reads +``` + +Plus the 5 repeated flags removed above: see "Global flags". + ## compete ```text diff --git a/rust/mumdia/crates/mumdia-io/src/table.rs b/rust/mumdia/crates/mumdia-io/src/table.rs index 11d4f6fc..1a89cf36 100644 --- a/rust/mumdia/crates/mumdia-io/src/table.rs +++ b/rust/mumdia/crates/mumdia-io/src/table.rs @@ -427,7 +427,7 @@ col_chunks!( /// Write columns to a Parquet file. Returns the row count. All columns must /// share the same length. /// -/// The columns are validated as a set, then encoded in [`WRITE_TABLE_CHUNK_ROWS`] chunks +/// The columns are validated as a set, then encoded in `WRITE_TABLE_CHUNK_ROWS` chunks /// through [`TableWriter`]. Building ONE record batch for the whole table first, as this /// used to, meant a second full Arrow copy of every column existed beside the source /// vectors: on a wide artifact that is a second copy of the whole table, and each column From ea169595e4e6ede32eef469edd35e138aa2dcb54 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:38:40 +0200 Subject: [PATCH 31/35] perf(extract): revert the per-window query-m/z hoist, which costs memory and buys nothing Item (3) of the brief was to hoist the per-peak setup recomputed once per sub-range task: `peak.mz / mass_off.factor_at(peak.mz)` and the bin computation (a `ln()`). Only the first half is reachable from this file -- the bin is computed inside `FragIndex::probe_peak_win` and using a precomputed one needs a new entry point on `FragIndex`, which is outside `extract.rs`. This reverts the first half as well, because it measured as a loss. The buffer is 8 bytes per peak of every window in flight. Measured peak RSS on the AIF fixture (39.6M MS2 peaks, 152 windows), two runs per arm on the same host: with the whole run in one batch (`windows_in_flight: 400`, which is the shape a grouped band search has, since `window_groups` leaves a band 1-3 windows and one batch) the buffer covers every peak of the run, 317 MB, and the process went to 2,913 and 2,921 MB against 2,693 and 2,719 MB for the baseline -- the flat accumulator's saving swallowed and then some. Under the default `windows_in_flight` it is a sixteenth of that, but still not free. It buys no measurable time. The extract compute phase is 2.12-2.16 s with the hoist and 2.12-2.16 s without, five warm runs each; the redundant work it removes is one f64 division per peak per task (the masscal grid, where the binary search would also be saved, is absent on this run: `mz_cal_grid=0`). Memory is the whole point of this branch, so a structure that is hundreds of MB in the shape that is failing in production and flat in time does not belong in it. Both halves are recorded in a comment at the call site so the next reader does not re-derive them. Also records, at the counting sort, that giving back the hit buffer's doubling slack with `shrink_to_fit` was tried and reverted too: 2,520-2,560 MB against 2,456-2,505 MB, i.e. the copy costs more than the slack it returns, and the slack is the same slack the per-candidate vectors carried before. Output equality unchanged: re-verified on the same fixture with the final binary, `psms_extracted.parquet` byte-identical to the baseline's and `chromatograms.parquet` equal as a table (620,169 rows, every column, same order). Co-Authored-By: Claude Opus 5 (1M context) --- .../crates/mumdia/src/stages/extract.rs | 51 ++++++------------- 1 file changed, 16 insertions(+), 35 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/extract.rs b/rust/mumdia/crates/mumdia/src/stages/extract.rs index 633aef1f..9ba66d8a 100644 --- a/rust/mumdia/crates/mumdia/src/stages/extract.rs +++ b/rust/mumdia/crates/mumdia/src/stages/extract.rs @@ -1158,7 +1158,6 @@ fn accumulate_groups( // windows must report before sub-range `k` is final. let mut tasks: Vec<(usize, usize, u32, u32)> = Vec::new(); let mut expected: Vec = vec![0; n_sub]; - let mut per_window: Vec = vec![0; groups.len()]; for (k, exp) in expected.iter_mut().enumerate() { let (s, e) = sub_bounds(k); for (gi, g) in groups.iter().enumerate() { @@ -1166,40 +1165,22 @@ fn accumulate_groups( if hi > lo { tasks.push((k, gi, lo, hi)); *exp += 1; - per_window[gi] += 1; } } } - // Offset-corrected query m/z per peak, computed ONCE per window instead of once per - // sub-range task. `MassOffset::factor_at` is a binary search into the m/z-dependent - // calibration grid when a masscal grid is present, and a division either way, and both - // were repeated once per task -- up to 2x threads -- for every peak of every scan of - // the window. Flat with one offset per scan, so it is two allocations per window rather - // than one per scan, and built only where the window is actually split. - let qmz: Vec, Vec)>> = groups - .iter() - .enumerate() - .map(|(gi, g)| { - (per_window[gi] > 1).then(|| { - let mut q: Vec = - Vec::with_capacity(g.scans.iter().map(|&si| scans[si].peaks.len()).sum()); - let mut off: Vec = Vec::with_capacity(g.scans.len() + 1); - off.push(0); - for &si in &g.scans { - for peak in &scans[si].peaks { - q.push(peak.mz / mass_off.factor_at(peak.mz)); - } - off.push(q.len()); - } - (q, off) - }) - }) - .collect(); + // Note on `peak.mz / mass_off.factor_at(peak.mz)` below: it is recomputed once per + // task for every peak of the window, so the tasks of one window repeat it. Hoisting + // it into a per-window buffer was tried and reverted: it is 8 bytes per peak of every + // window in flight, which measured +200 MB of peak RSS on the AIF fixture (and would + // be a third of a gigabyte with the whole run in one batch, which is the shape a + // grouped band search has), for no measurable time -- the extract compute phase was + // 2.12-2.16 s with and without it over five runs. The other half of that hoist, the + // `ln()` bin computation inside `probe_peak_win`, needs a `FragIndex` entry point + // taking a precomputed bin and cannot be done from this file at all. let (tx, rx) = std::sync::mpsc::channel::<(usize, usize, HitStore)>(); { let probe_range = |gi: usize, lo: u32, hi: u32| -> HitStore { let ids = &groups[gi].scans; - let qmz = &qmz[gi]; // Flat `(cid, hit)` pairs in probe order, grouped by candidate at the end of // the task with a stable counting sort. The task-local `HashMap>` this replaces was one of the two populations of medium heap @@ -1211,16 +1192,12 @@ fn accumulate_groups( // reprobes the same bins, so cache each bin's narrowed posting range once // instead of binary-searching it per peak. let mut nw = idx.window_narrow(lo, hi); - for (k, &si) in ids.iter().enumerate() { + for &si in ids { let scan = &scans[si]; let rt = scan.rt_seconds; - let qs: Option<&[f64]> = qmz.as_ref().map(|(q, off)| &q[off[k]..off[k + 1]]); - for (pk, peak) in scan.peaks.iter().enumerate() { + for peak in &scan.peaks { let inten = peak.intensity; - let q_mz = match qs { - Some(s) => s[pk], - None => peak.mz / mass_off.factor_at(peak.mz), - }; + let q_mz = peak.mz / mass_off.factor_at(peak.mz); let obs_mz = peak.mz; claimants.clear(); idx.probe_peak_win(&mut nw, q_mz, |cid, _pmz, pint, frag| { @@ -1293,6 +1270,10 @@ fn accumulate_groups( } let (cids, offs) = group_hits_by_candidate(lo, hi, &mut flat_cid, &mut flat_hit); drop(flat_cid); + // Not shrunk: the buffer's doubling slack is the same slack the per-candidate + // vectors carried, and `shrink_to_fit` here measured 2,520-2,560 MB of peak + // RSS against 2,456-2,505 MB without it on the AIF fixture, i.e. its copy + // costs more than the slack it returns. HitStore { cids, offs, From da1c4b75df633a3c64836ec637d365f746937ed2 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:40:46 +0200 Subject: [PATCH 32/35] test(features): the golden feature digest is per platform, because the values are The digest the review asked for was captured on Windows and fails on Linux. It is not nondeterminism: the same fixture gives 0x4e43..a1cc on Windows and 0x6137..c42a on Linux, stably, two runs each on the same binary. The extended battery reaches `ln`, `exp` and `powf`, which are the platform's libm and agree only to within the last bit, so the engine's f64 feature values are not bit-identical across operating systems. Worth knowing on its own: any comparison of feature values, or of any artifact hash downstream of them, has to be made on one operating system. The PIN is unaffected and its hash is the same on both, because it formats to fixed decimals. The test now carries a digest per platform and asserts the one for the platform it runs on; a platform with no entry checks the shape and the PIN and says so, rather than failing on a constant captured somewhere else. Co-Authored-By: Claude Opus 5 (1M context) --- .../crates/mumdia/src/stages/features.rs | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/features.rs b/rust/mumdia/crates/mumdia/src/stages/features.rs index 9cc7aa22..f7dc9dac 100644 --- a/rust/mumdia/crates/mumdia/src/stages/features.rs +++ b/rust/mumdia/crates/mumdia/src/stages/features.rs @@ -3569,7 +3569,22 @@ mod tests { // running this fixture at 1ac1b44^ (or at the last commit whose values are trusted) // and changing them in one commit that says which values moved and why. Taking the // new numbers from the current build makes the test pin nothing. - const GOLDEN_DIGEST: u64 = 0x4e43_7960_2b2a_a1cc; + // + // THE DIGEST IS PER PLATFORM, and that is a finding rather than a nuisance: the + // same fixture gives 0x4e43..a1cc on Windows and 0x6137..c42a on Linux, stably on + // both (two runs each). The extended battery reaches `ln`, `exp` and `powf`, which + // are the platform's libm and agree only to within the last bit, so the engine's + // f64 feature VALUES are not bit-identical across operating systems. The PIN is: it + // formats to fixed decimals, which is why the same PIN hash holds on both. A + // platform with no entry here checks the shape and the PIN and says it has no + // digest, rather than failing on a constant captured somewhere else. + const GOLDEN_DIGEST: Option = if cfg!(target_os = "windows") { + Some(0x4e43_7960_2b2a_a1cc) + } else if cfg!(target_os = "linux") { + Some(0x6137_b855_a2da_c42a) + } else { + None + }; const GOLDEN_PIN: &str = "806a126473eafce7a1567ac71b5253e8605af0062864a02f2ab5f0b77e851817"; let dir = std::env::temp_dir().join("mumdia_features_golden"); @@ -3608,11 +3623,15 @@ mod tests { .unwrap(); let (rows, ncols, digest) = features_digest(&out); assert_eq!((rows, ncols), (61, 398), "chunk {tag}: table shape moved"); - assert_eq!( - digest, GOLDEN_DIGEST, - "chunk {tag}: a feature value or its column differs from the name-keyed \ - assembly at 1ac1b44^ (digest 0x{digest:016x})" - ); + match GOLDEN_DIGEST { + Some(golden) => assert_eq!( + digest, golden, + "chunk {tag}: a feature value or its column differs from the name-keyed assembly at 1ac1b44^ (this platform computes 0x{digest:016x})" + ), + None => eprintln!( + "no feature digest is recorded for this platform; chunk {tag} computes 0x{digest:016x}. Add it above once it has been checked against a platform that has one." + ), + } assert_eq!( mumdia_io::hash::blake3_file(&pin).unwrap(), GOLDEN_PIN, From 35521e4c65853b8a9bf04224feb06143bece40bb Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:43:59 +0200 Subject: [PATCH 33/35] perf(extract): stop parking a pool worker on the consumer, and bitset the allowlist Two items from the brief, both scheduling/lookup only, neither touching a value. (7) The consumer runs on the thread that called `accumulate_groups`, and under `groups.parallel` that thread is itself a rayon worker, so a plain blocking `recv` parks one worker per band in flight: 24 bands parked 24 of 32 threads and the pool ran on what was left. It now calls `rayon::yield_now` while waiting, which executes one queued task instead of parking -- usually one of this batch's own probes -- and falls back to a 1 ms blocking wait when the pool reports `Idle`, so an empty pool does not become a spin. The results are still consumed in arrival order and reordered by sub-range, so nothing about the output depends on it. The per-candidate merge this thread used to do is gone with the flat accumulator; what is left is the gather and the flush. (6) The `--restrict-candidates` allowlist was a `HashSet` probed once per VERIFIED POSTING, inside the probe callback: a hash and a bucket load for every fragment match of every peak of every scan, on all six probe sites. Candidate ids are dense `0..ncand`, so it is now a bitset (`CandMask`): a shift and a load, in one allocation of `ncand / 8` bytes instead of a hash table -- smaller as soon as the allowlist holds more than about a sixty-fourth of the library, which is the case the allowlist exists for (a gate-on pass over a 35M-83M candidate immunopeptidomics library). Ids outside `0..ncand` are absent from the mask, which is not a behaviour change: the probe only ever asks about ids it produced. The other half of (6), slicing and zipping the posting arrays in the verify loop to drop bounds checks, is NOT done: that loop is `FragIndex::emit_range` in `matchers/fragindex.rs`, outside the one file this work may touch. The ppm predicate is untouched either way, as the brief requires. Verified on the AIF fixture (1,691,048 candidates, 465,806 scans, 152 windows) against the base binary, both `--release`, on the plain path AND on the allowlist path (`--restrict-candidates` over the plain run's 41,677 accepted candidates, which the mask counts identically): `psms_extracted.parquet` is BYTE-IDENTICAL in both, and `chromatograms.parquet` is equal as a table in both (620,169 rows, every column, same order), its bytes differing only through the row-group boundaries the sub-range flush moves. Co-Authored-By: Claude Opus 5 (1M context) --- .../crates/mumdia/src/stages/extract.rs | 153 ++++++++++++++++-- 1 file changed, 141 insertions(+), 12 deletions(-) diff --git a/rust/mumdia/crates/mumdia/src/stages/extract.rs b/rust/mumdia/crates/mumdia/src/stages/extract.rs index 9ba66d8a..1e772f5c 100644 --- a/rust/mumdia/crates/mumdia/src/stages/extract.rs +++ b/rust/mumdia/crates/mumdia/src/stages/extract.rs @@ -344,6 +344,41 @@ impl HitRun { } } +/// Receive one probing result, executing pending pool work while waiting instead of +/// parking. +/// +/// The consumer runs on the thread that called `accumulate_groups`, and under +/// `groups.parallel` that thread is itself a rayon worker: a plain blocking `recv` parks +/// one worker per band in flight, so 24 bands parked 24 of 32 threads and the pool ran on +/// what was left. `yield_now` executes one queued task instead, which is usually one of +/// this batch's own probes. When the pool has nothing queued it returns `Idle` (or `None` +/// off a worker thread), and a short blocking wait then avoids a spin. +/// +/// Scheduling only: the results are still consumed in arrival order and reordered by +/// sub-range, so nothing about the output depends on this. +fn recv_participating(rx: &std::sync::mpsc::Receiver) -> T { + use std::sync::mpsc::{RecvTimeoutError, TryRecvError}; + loop { + match rx.try_recv() { + Ok(v) => return v, + Err(TryRecvError::Disconnected) => { + panic!("every probing task sends exactly one store") + } + Err(TryRecvError::Empty) => {} + } + if let Some(rayon::Yield::Executed) = rayon::yield_now() { + continue; + } + match rx.recv_timeout(std::time::Duration::from_millis(1)) { + Ok(v) => return v, + Err(RecvTimeoutError::Timeout) => {} + Err(RecvTimeoutError::Disconnected) => { + panic!("every probing task sends exactly one store") + } + } + } +} + /// Gather up to `max_cands` candidates below `bound` out of `runs`, ascending by candidate /// id, appending each candidate's hits run by run. Runs are held in window order, so a /// candidate seen in several windows keeps exactly the hit sequence the serial path @@ -461,6 +496,49 @@ struct Contested { apportioned: f64, } +/// The `--restrict-candidates` allowlist, as a bitset over the dense candidate id range. +/// +/// The allowlist used to be a `HashSet` probed once per VERIFIED POSTING, inside the +/// probe callback: a hash and a bucket load for every fragment match of every peak of +/// every scan. Candidate ids are dense `0..ncand` (`index.rs`), so one bit per candidate +/// answers the same question with a shift and a load, and it is one allocation of +/// `ncand / 8` bytes instead of a hash table of the allowlist -- smaller as soon as the +/// allowlist holds more than about a sixty-fourth of the library, which is the case the +/// allowlist exists for (a prior gate-on pass over a 35M-83M candidate library). +struct CandMask { + bits: Vec, + n: usize, +} + +impl CandMask { + fn new(n: usize) -> CandMask { + CandMask { + bits: vec![0; n.div_ceil(64)], + n, + } + } + + fn insert(&mut self, c: u32) { + let i = c as usize; + if i < self.n { + self.bits[i >> 6] |= 1u64 << (i & 63); + } + } + + /// Ids outside `0..n` are absent. The probe only ever asks about ids it produced, so + /// they are in range by construction; an out-of-range id in the allowlist file could + /// not have matched the hash set either, because no probe would ask for it. + #[inline] + fn contains(&self, c: u32) -> bool { + let i = c as usize; + i < self.n && self.bits[i >> 6] & (1u64 << (i & 63)) != 0 + } + + fn len(&self) -> usize { + self.bits.iter().map(|w| w.count_ones() as usize).sum() + } +} + /// Ordinals a [`FragSet`] holds in its inline bitmask; above this it spills to a sorted /// vector. A candidate carries a few dozen predicted fragments (6 by default, 12 in the /// shipped DIA-NN library), so the spill is unreachable in practice and costs nothing when @@ -1119,7 +1197,7 @@ fn accumulate_groups( rt_hi: &[f64], mass_off: &MassOffset, cfg: &ExtractConfig, - restrict: Option<&std::collections::HashSet>, + restrict: Option<&CandMask>, // First candidate a window of a LATER batch can still add hits to; nothing at or above // it may be flushed here. bound: u32, @@ -1209,7 +1287,7 @@ fn accumulate_groups( // serial path applies it: a candidate outside the list neither // collects hits nor competes for a shared peak. if let Some(s) = restrict { - if !s.contains(&cid) { + if !s.contains(cid) { return; } } @@ -1301,7 +1379,7 @@ fn accumulate_groups( let mut next_k = 0usize; let mut stopped = false; for _ in 0..n { - let (k, gi, store) = rx.recv().expect("every task sends exactly one store"); + let (k, gi, store) = recv_participating(&rx); arrived[k] += 1; pending.insert((k, gi), store); while next_k < n_sub && arrived[next_k] == expected[next_k] { @@ -1381,7 +1459,7 @@ fn extract_twopass_windows( mass_off: &MassOffset, frag_tol: f64, cfg: &ExtractConfig, - restrict: Option<&std::collections::HashSet>, + restrict: Option<&CandMask>, reassign: bool, claim_margin: f32, ) -> (HashMap>, HashMap) { @@ -1436,7 +1514,7 @@ fn extract_twopass_windows( return; } if let Some(s) = restrict { - if !s.contains(&cid) { + if !s.contains(cid) { return; } } @@ -1490,7 +1568,7 @@ fn extract_twopass_windows( return; } if let Some(s) = restrict { - if !s.contains(&cid) { + if !s.contains(cid) { return; } } @@ -1561,7 +1639,7 @@ fn extract_twopass_windows( return; } if let Some(s) = restrict { - if !s.contains(&cid) { + if !s.contains(cid) { return; } } @@ -1672,7 +1750,7 @@ fn extract_twopass_windows( return; } if let Some(s) = restrict { - if !s.contains(&cid) { + if !s.contains(cid) { return; } } @@ -1744,7 +1822,7 @@ fn extract_twopass_windows( return; } if let Some(s) = restrict { - if !s.contains(&cid) { + if !s.contains(cid) { return; } } @@ -1915,10 +1993,13 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { // Optional candidate allowlist (gate-first-then-compete): restrict extraction to // the accepted survivors of a prior gate-on run so the two-pass peak-claim profile // map stays small. - let restrict: Option> = match p.restrict_candidates { + let restrict: Option = match p.restrict_candidates { Some(path) => { let t = TableFile::open(path)?; - let s: std::collections::HashSet = t.u32("candidate_id")?.into_iter().collect(); + let mut s = CandMask::new(lib.n_candidates()); + for c in t.u32("candidate_id")? { + s.insert(c); + } info!( restrict_candidates = s.len(), "extract: restricting to candidate allowlist" @@ -2144,7 +2225,7 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { return; } if let Some(s) = &restrict { - if !s.contains(&cid) { + if !s.contains(cid) { return; } } @@ -3460,6 +3541,54 @@ pub fn run(p: ExtractParams) -> Result<(u64, u64)> { Ok((n_psms, n_chrom)) } +#[cfg(test)] +mod cand_mask_tests { + use super::CandMask; + + /// The allowlist bitset must answer exactly what the `HashSet` answered for every + /// id the probe can produce, which is every id in `0..ncand`. + #[test] + fn answers_the_same_as_the_hash_set() { + let n = 1000usize; + let ids: Vec = vec![0, 1, 63, 64, 65, 127, 128, 500, 999, 500]; + let reference: std::collections::HashSet = ids.iter().copied().collect(); + let mut mask = CandMask::new(n); + for &c in &ids { + mask.insert(c); + } + assert_eq!( + mask.len(), + reference.len(), + "duplicates must not be counted" + ); + for c in 0..n as u32 { + assert_eq!(mask.contains(c), reference.contains(&c), "id {c}"); + } + } + + /// An id past the library's candidate count cannot be asked about by the probe, so it + /// must not panic and must not be stored. + #[test] + fn ids_outside_the_library_are_absent() { + let mut mask = CandMask::new(64); + mask.insert(5); + mask.insert(64); + mask.insert(u32::MAX); + assert_eq!(mask.len(), 1); + assert!(mask.contains(5)); + assert!(!mask.contains(64) && !mask.contains(u32::MAX)); + } + + /// An empty library gives an empty mask rather than an out-of-bounds write. + #[test] + fn an_empty_library_gives_an_empty_mask() { + let mut mask = CandMask::new(0); + mask.insert(0); + assert_eq!(mask.len(), 0); + assert!(!mask.contains(0)); + } +} + #[cfg(test)] mod frag_set_tests { use super::{FragSet, FRAG_SET_BITS}; From ca1185433b90b8e49fe548799c2abeac8de31147 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 18:45:26 +0200 Subject: [PATCH 34/35] ci: regenerate the SBOM for the zstd crates Co-Authored-By: Claude Opus 5 (1M context) --- sbom.cdx.json | 184 +++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 183 insertions(+), 1 deletion(-) diff --git a/sbom.cdx.json b/sbom.cdx.json index e51e21b6..446497f7 100644 --- a/sbom.cdx.json +++ b/sbom.cdx.json @@ -1390,6 +1390,25 @@ } ] }, + { + "type": "library", + "bom-ref": "pkg:cargo/getrandom@0.4.3", + "name": "getrandom", + "version": "0.4.3", + "purl": "pkg:cargo/getrandom@0.4.3", + "description": "A small cross-platform library for retrieving random data from system source", + "licenses": [ + { + "expression": "MIT OR Apache-2.0" + } + ], + "externalReferences": [ + { + "type": "vcs", + "url": "https://github.com/rust-random/getrandom" + } + ] + }, { "type": "library", "bom-ref": "pkg:cargo/half@2.7.1", @@ -1599,6 +1618,25 @@ } ] }, + { + "type": "library", + "bom-ref": "pkg:cargo/jobserver@0.1.35", + "name": "jobserver", + "version": "0.1.35", + "purl": "pkg:cargo/jobserver@0.1.35", + "description": "An implementation of the GNU Make jobserver for Rust.", + "licenses": [ + { + "expression": "MIT OR Apache-2.0" + } + ], + "externalReferences": [ + { + "type": "vcs", + "url": "https://github.com/rust-lang/jobserver-rs" + } + ] + }, { "type": "library", "bom-ref": "pkg:cargo/js-sys@0.3.103", @@ -2249,6 +2287,25 @@ } ] }, + { + "type": "library", + "bom-ref": "pkg:cargo/pkg-config@0.3.34", + "name": "pkg-config", + "version": "0.3.34", + "purl": "pkg:cargo/pkg-config@0.3.34", + "description": "A library to run the pkg-config system tool at build time in order to be used in Cargo build scripts.", + "licenses": [ + { + "expression": "MIT OR Apache-2.0" + } + ], + "externalReferences": [ + { + "type": "vcs", + "url": "https://github.com/rust-lang/pkg-config-rs" + } + ] + }, { "type": "library", "bom-ref": "pkg:cargo/proc-macro2@1.0.106", @@ -2329,6 +2386,29 @@ } ] }, + { + "type": "library", + "bom-ref": "pkg:cargo/r-efi@6.0.0", + "name": "r-efi", + "version": "6.0.0", + "purl": "pkg:cargo/r-efi@6.0.0", + "description": "UEFI Reference Specification Protocol Constants and Definitions", + "licenses": [ + { + "expression": "MIT OR Apache-2.0 OR LGPL-2.1-or-later" + } + ], + "externalReferences": [ + { + "type": "vcs", + "url": "https://github.com/r-efi/r-efi" + }, + { + "type": "website", + "url": "https://github.com/r-efi/r-efi/wiki" + } + ] + }, { "type": "library", "bom-ref": "pkg:cargo/rayon@1.12.0", @@ -3524,6 +3604,63 @@ "url": "https://github.com/dtolnay/zmij" } ] + }, + { + "type": "library", + "bom-ref": "pkg:cargo/zstd@0.13.3", + "name": "zstd", + "version": "0.13.3", + "purl": "pkg:cargo/zstd@0.13.3", + "description": "Binding for the zstd compression library.", + "licenses": [ + { + "expression": "MIT" + } + ], + "externalReferences": [ + { + "type": "vcs", + "url": "https://github.com/gyscos/zstd-rs" + } + ] + }, + { + "type": "library", + "bom-ref": "pkg:cargo/zstd-safe@7.2.4", + "name": "zstd-safe", + "version": "7.2.4", + "purl": "pkg:cargo/zstd-safe@7.2.4", + "description": "Safe low-level bindings for the zstd compression library.", + "licenses": [ + { + "expression": "MIT OR Apache-2.0" + } + ], + "externalReferences": [ + { + "type": "vcs", + "url": "https://github.com/gyscos/zstd-rs" + } + ] + }, + { + "type": "library", + "bom-ref": "pkg:cargo/zstd-sys@2.0.16+zstd.1.5.7", + "name": "zstd-sys", + "version": "2.0.16+zstd.1.5.7", + "purl": "pkg:cargo/zstd-sys@2.0.16+zstd.1.5.7", + "description": "Low-level bindings for the zstd compression library.", + "licenses": [ + { + "expression": "MIT/Apache-2.0" + } + ], + "externalReferences": [ + { + "type": "vcs", + "url": "https://github.com/gyscos/zstd-rs" + } + ] } ], "dependencies": [ @@ -3831,6 +3968,8 @@ "ref": "pkg:cargo/cc@1.2.65", "dependsOn": [ "pkg:cargo/find-msvc-tools@0.1.9", + "pkg:cargo/jobserver@0.1.35", + "pkg:cargo/libc@0.2.186", "pkg:cargo/shlex@2.0.1" ] }, @@ -4042,6 +4181,14 @@ "pkg:cargo/wasip2@1.0.4+wasi-0.2.12" ] }, + { + "ref": "pkg:cargo/getrandom@0.4.3", + "dependsOn": [ + "pkg:cargo/cfg-if@1.0.4", + "pkg:cargo/libc@0.2.186", + "pkg:cargo/r-efi@6.0.0" + ] + }, { "ref": "pkg:cargo/half@2.7.1", "dependsOn": [ @@ -4108,6 +4255,13 @@ "ref": "pkg:cargo/itoa@1.0.18", "dependsOn": [] }, + { + "ref": "pkg:cargo/jobserver@0.1.35", + "dependsOn": [ + "pkg:cargo/getrandom@0.4.3", + "pkg:cargo/libc@0.2.186" + ] + }, { "ref": "pkg:cargo/js-sys@0.3.103", "dependsOn": [ @@ -4387,13 +4541,18 @@ "pkg:cargo/num-traits@0.2.19", "pkg:cargo/seq-macro@0.3.6", "pkg:cargo/snap@1.1.1", - "pkg:cargo/twox-hash@2.1.2" + "pkg:cargo/twox-hash@2.1.2", + "pkg:cargo/zstd@0.13.3" ] }, { "ref": "pkg:cargo/pin-project-lite@0.2.17", "dependsOn": [] }, + { + "ref": "pkg:cargo/pkg-config@0.3.34", + "dependsOn": [] + }, { "ref": "pkg:cargo/proc-macro2@1.0.106", "dependsOn": [ @@ -4417,6 +4576,10 @@ "ref": "pkg:cargo/r-efi@5.3.0", "dependsOn": [] }, + { + "ref": "pkg:cargo/r-efi@6.0.0", + "dependsOn": [] + }, { "ref": "pkg:cargo/rayon-core@1.13.0", "dependsOn": [ @@ -4783,6 +4946,25 @@ { "ref": "pkg:cargo/zmij@1.0.21", "dependsOn": [] + }, + { + "ref": "pkg:cargo/zstd-safe@7.2.4", + "dependsOn": [ + "pkg:cargo/zstd-sys@2.0.16+zstd.1.5.7" + ] + }, + { + "ref": "pkg:cargo/zstd-sys@2.0.16+zstd.1.5.7", + "dependsOn": [ + "pkg:cargo/cc@1.2.65", + "pkg:cargo/pkg-config@0.3.34" + ] + }, + { + "ref": "pkg:cargo/zstd@0.13.3", + "dependsOn": [ + "pkg:cargo/zstd-safe@7.2.4" + ] } ] } From 2a991edaacb92b4aff1a107d7cf6231ab25069f3 Mon Sep 17 00:00:00 2001 From: MuMDIA Date: Tue, 22 Sep 2026 19:04:52 +0200 Subject: [PATCH 35/35] docs(config): regenerate after the accumulator rewrite The reference records the environment reads whose name is not a literal by source line; the CSR accumulator moved three of them. Co-Authored-By: Claude Opus 5 (1M context) --- docs/24_config_reference.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/24_config_reference.md b/docs/24_config_reference.md index 74a90849..447b2646 100644 --- a/docs/24_config_reference.md +++ b/docs/24_config_reference.md @@ -847,10 +847,11 @@ Every field whose struct has an `impl Default` resolved from the source. - `peptidoforms.fixed_mods[].name` (`String`) - `peptidoforms.fixed_mods[].residue` (`char`) -2 environment read(s) whose name is not a literal: +3 environment read(s) whose name is not a literal: -- `rust/mumdia/crates/mumdia/src/stages/extract.rs:2786: env read via closure of `&mut flushed`` -- `rust/mumdia/crates/mumdia/src/stages/extract.rs:2804: env read via closure of `&mut cand_hits`` +- `rust/mumdia/crates/mumdia/src/stages/extract.rs:1419: env read via closure of `chunk`` +- `rust/mumdia/crates/mumdia/src/stages/extract.rs:3333: env read via closure of `&mut chunk`` +- `rust/mumdia/crates/mumdia/src/stages/extract.rs:3369: env read via closure of `&mut chunk`` ## Coverage