From 4731b8d4616fa98a98b49f818a655d3322340dea Mon Sep 17 00:00:00 2001 From: Roman Young Date: Wed, 14 Oct 2026 01:22:02 +0000 Subject: [PATCH] perf: read metadata only for the proteins a search hit Since #25, hits get their protein metadata from one join against a per-protein table instead of a copy in every hit row, but that table holds the whole proteome: 82k proteins with nine text fields each for human, read from the index and converted however few proteins the search hit. On small searches that cost more than the search itself (an exact search of 1,000 MHC ligands took 0.29 s, 0.26 s of it building the table). rs_metadata now takes an optional list of protein numbers and returns just those rows, ascending and deduplicated, raising ValueError for a number the index doesn't have; without the list it returns the whole table as before. _to_dataframe passes the distinct proteins in the hits, and the join asks for the hits' row order (maintain_order='left') instead of relying on polars' unspecified default. A hit on a protein the index doesn't have means the search read stale .pepidx files, so that error now says to rebuild them. Output is unchanged. test_metadata_rows.py builds each mode's frame twice from the same raw hits, once reading only the hit proteins and once reading every protein, and requires the two to be identical; it also covers duplicate, unsorted, empty and out-of-range requests, the column types of a search that hits nothing, and a stale index. The new file has its CI matrix row. --- .github/workflows/tests.yml | 2 + pepmatch/matcher.py | 32 ++++-- pepmatch/rs-engine/src/lib.rs | 7 +- pepmatch/rs-engine/src/match.rs | 34 ++++-- pepmatch/tests/test_metadata_rows.py | 160 +++++++++++++++++++++++++++ 5 files changed, 214 insertions(+), 21 deletions(-) create mode 100644 pepmatch/tests/test_metadata_rows.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2825be1..af56da9 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -36,6 +36,8 @@ jobs: test-file: pepmatch/tests/test_indel_seed_reduction.py - name: Indel Positions Equivalence test-file: pepmatch/tests/test_indel_positions_equivalence.py + - name: Metadata Rows + test-file: pepmatch/tests/test_metadata_rows.py steps: - name: Checkout Code diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index 0847ba5..95ae25c 100755 --- a/pepmatch/matcher.py +++ b/pepmatch/matcher.py @@ -524,19 +524,28 @@ def _best_match_filter(self, df): return pl.concat([matched_df, unmatched_df], how="vertical_relaxed") - def _metadata_table(self, pepidx_path: str) -> pl.DataFrame: - """Per-protein metadata read from `pepidx_path`, the index the search just used. - Protein numbers are that index's own FASTA ordinals, so any other file -- including - one whose name merely shares this proteome's prefix -- would misattribute every hit.""" - pnum, pid, name, species, taxon, gene, exist, seqver, geneprio, swiss = rs_metadata(pepidx_path) + def _metadata_table(self, pepidx_path: str, protein_nums=None) -> pl.DataFrame: + """Metadata for `protein_nums` (every protein when None), read from `pepidx_path`, the + index the search just used. Protein numbers are that index's own FASTA ordinals, so any + other file -- including one whose name merely shares this proteome's prefix -- would + misattribute every hit. Searches pass only the proteins they hit: building the whole + proteome's table (82k rows for human) cost more than the search itself on small queries. + A hit on a protein this index lacks means the search's indexes disagree, which raises.""" + try: + pnum, pid, name, species, taxon, gene, exist, seqver, geneprio, swiss = rs_metadata(pepidx_path, protein_nums) + except ValueError as e: # .pepidx files are reused by name, so a FASTA edited in place leaves stale ones + raise ValueError( + f"{e}. This proteome's .pepidx files in {self.preprocessed_files_path} were most likely built " + f"from different versions of {self.proteome_file}; delete them to rebuild." + ) from e + str_cols = ['Protein ID','Protein Name','Species','Taxon ID','Gene', + 'Protein Existence Level','Sequence Version','Gene Priority','SwissProt Reviewed'] m = pl.DataFrame({ 'protein_num': pl.Series(pnum, dtype=pl.UInt32), 'Protein ID': pid, 'Protein Name': name, 'Species': species, 'Taxon ID': taxon, 'Gene': gene, 'Protein Existence Level': exist, 'Sequence Version': seqver, 'Gene Priority': geneprio, 'SwissProt Reviewed': swiss, - }) - str_cols = ['Protein ID','Protein Name','Species','Taxon ID','Gene', - 'Protein Existence Level','Sequence Version','Gene Priority','SwissProt Reviewed'] + }, schema_overrides={c: pl.Utf8 for c in str_cols}) # String even when nothing was hit m = m.with_columns( pl.when(pl.col(c) == '').then(None).otherwise(pl.col(c)).alias(c) for c in str_cols ) @@ -561,7 +570,8 @@ def _final_columns(self, is_indels): def _to_dataframe(self, cols, pepidx_path, is_indels=False): """Build the results frame from the columnar Rust output, joining protein metadata - from `pepidx_path` once instead of cloning it into every hit row. rs_indel_match + from `pepidx_path` once instead of cloning it into every hit row. Only the proteins + the search hit are read, and the join keeps the hits' row order. rs_indel_match reuses the mismatch slot for its edit count, so only the column name differs by mode; indel positions are derived here from (query, matched) rather than in Rust, once per distinct pair since a matched peptide often recurs across proteins, and @@ -600,7 +610,9 @@ def _to_dataframe(self, cols, pepidx_path, is_indels=False): 'Index end': pl.Series(iend, dtype=pl.Int64), }) - df = base.join(self._metadata_table(pepidx_path), on='protein_num', how='left').drop('protein_num') + hit_proteins = base.get_column('protein_num').drop_nulls().unique().to_list() + metadata = self._metadata_table(pepidx_path, hit_proteins) + df = base.join(metadata, on='protein_num', how='left', maintain_order='left').drop('protein_num') if self.sequence_version: df = df.with_columns( diff --git a/pepmatch/rs-engine/src/lib.rs b/pepmatch/rs-engine/src/lib.rs index 1655a0b..d3b3da6 100644 --- a/pepmatch/rs-engine/src/lib.rs +++ b/pepmatch/rs-engine/src/lib.rs @@ -29,9 +29,12 @@ fn rs_discontinuous(pepidx_path: &str, epitopes: Vec<(String, Vec<(char, usize)> matching::run_discontinuous(pepidx_path, epitopes, max_mismatches) } +/// Metadata for `protein_nums` (1-based FASTA ordinals in this index), or for every protein +/// when it is omitted. A number outside the index raises `ValueError`. #[pyfunction] -fn rs_metadata(pepidx_path: &str) -> matching::MetaColumns { - matching::run_metadata(pepidx_path) +#[pyo3(signature = (pepidx_path, protein_nums=None))] +fn rs_metadata(pepidx_path: &str, protein_nums: Option>) -> PyResult { + matching::run_metadata(pepidx_path, protein_nums).map_err(pyo3::exceptions::PyValueError::new_err) } #[pyfunction] diff --git a/pepmatch/rs-engine/src/match.rs b/pepmatch/rs-engine/src/match.rs index c878d8c..420f958 100644 --- a/pepmatch/rs-engine/src/match.rs +++ b/pepmatch/rs-engine/src/match.rs @@ -576,22 +576,38 @@ pub(crate) fn run_discontinuous( unzip_records(records) } -fn metadata_columns(index: &PepIndex) -> MetaColumns { +/// One row per protein in `wanted` (ascending, duplicates collapsed), or per protein in the +/// index when `wanted` is None. A search passes only the proteins it hit, so the cost follows +/// the hits rather than the proteome (82k proteins for human, most of which a search never +/// touches). Duplicates are collapsed here because a repeated row would double every hit +/// joined to it. +fn metadata_columns(index: &PepIndex, wanted: Option>) -> Result { let n = index.num_proteins; - let mut pnum = Vec::with_capacity(n); + let proteins: Vec = match wanted { + Some(mut nums) => { + nums.sort_unstable(); + nums.dedup(); + if let Some(&bad) = nums.iter().find(|&&pn| pn == 0 || pn as usize > n) { + return Err(format!("protein number {bad} is not in this index (1..={n})")); + } + nums + } + None => (1..=n as u32).collect(), + }; + let mut pnum = Vec::with_capacity(proteins.len()); let mut f: [Vec; 9] = Default::default(); for v in f.iter_mut() { - v.reserve(n); + v.reserve(proteins.len()); } - for pn in 1..=n { - pnum.push(pn as u32); - let m = index.get_metadata(pn); + for pn in proteins { + pnum.push(pn); + let m = index.get_metadata(pn as usize); for (i, field) in m.into_iter().enumerate() { f[i].push(field); } } let [a, b, c, d, e, g, h, i, j] = f; - (pnum, a, b, c, d, e, g, h, i, j) + Ok((pnum, a, b, c, d, e, g, h, i, j)) } /// Seeds for indel search: the n+1 disjoint query tiles with the FEWEST proteome @@ -862,9 +878,9 @@ pub(crate) fn run_indel( unzip_records(records) } -pub(crate) fn run_metadata(pepidx_path: &str) -> MetaColumns { +pub(crate) fn run_metadata(pepidx_path: &str, protein_nums: Option>) -> Result { let index = PepIndex::open(pepidx_path); - metadata_columns(&index) + metadata_columns(&index, protein_nums) } pub(crate) fn run(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize, max_mismatches: usize) -> Columns { diff --git a/pepmatch/tests/test_metadata_rows.py b/pepmatch/tests/test_metadata_rows.py new file mode 100644 index 0000000..50ff73b --- /dev/null +++ b/pepmatch/tests/test_metadata_rows.py @@ -0,0 +1,160 @@ +import pytest +import polars.testing as plt + +from pathlib import Path + +import pepmatch.matcher as matcher_module +from pepmatch import Matcher, Preprocessor +from pepmatch._rs import rs_metadata + +DATA = Path(__file__).parent / 'data' + +# Three proteins. The second has no gene name (no GN=), so its Gene field is empty and must +# come back null; the third is never hit by the searches below, so its metadata is never read. +PROTEOME_FASTA = ( + '>sp|P11111|AAAA_HUMAN Alpha protein OS=Homo sapiens OX=9606 GN=AAA PE=1 SV=1\n' + 'MKVLATGGWQMKRRDEYH\n' + '>sp|P22222|BBBB_HUMAN Beta protein OS=Homo sapiens OX=9606 PE=2 SV=3\n' + 'RRHFDSAMFFLLNNMKVLAT\n' + '>tr|A0A333|A0A333_HUMAN Gamma protein OS=Homo sapiens OX=9606 GN=CCC PE=4 SV=1\n' + 'PQRSTWYVDEGHIKLMNPQ\n' +) + + +@pytest.fixture +def proteome(tmp_path) -> Path: + """The three-protein FASTA, preprocessed at k=5 next to it.""" + fasta = tmp_path / 'three.fasta' + fasta.write_text(PROTEOME_FASTA) + Preprocessor(str(fasta), preprocessed_files_path=str(tmp_path)).preprocess(5) + return fasta + + +def index_of(proteome): + return str(proteome.with_name('three_5mers.pepidx')) + + +def rows(columns): + return list(zip(*columns)) + + +def search(proteome, query): + return Matcher( + query=query, + proteome_file=str(proteome), + k=5, + preprocessed_files_path=str(proteome.parent), + output_format='dataframe', + ).match() + + +def test_requested_rows_equal_the_full_tables_rows(proteome): + full = rows(rs_metadata(index_of(proteome))) + assert [row[0] for row in full] == [1, 2, 3] + assert rows(rs_metadata(index_of(proteome), [3, 1])) == [full[0], full[2]] # ascending + + +def test_repeated_protein_numbers_collapse(proteome): + # A repeated row would double every hit joined to it. + assert [row[0] for row in rows(rs_metadata(index_of(proteome), [2, 2, 2]))] == [2] + + +def test_no_protein_numbers_return_empty_columns(proteome): + assert rs_metadata(index_of(proteome), []) == tuple([] for _ in range(10)) + + +@pytest.mark.parametrize('protein', [0, 4]) +def test_a_protein_number_outside_the_index_raises(proteome, protein): + with pytest.raises(ValueError, match=f'protein number {protein} '): + rs_metadata(index_of(proteome), [protein]) + + +def test_hits_carry_their_own_proteins_metadata(proteome): + df = search(proteome, ['MKVLAT', 'FFLLNN']) + got = { + (row['Query Sequence'], row['Protein ID'], row['Gene'], row['Protein Existence Level'], row['SwissProt Reviewed']) + for row in df.iter_rows(named=True) + } + assert got == { + ('MKVLAT', 'P11111.1', 'AAA', 1, True), + ('MKVLAT', 'P22222.3', None, 2, True), + ('FFLLNN', 'P22222.3', None, 2, True), + } + + +METADATA_COLUMNS = [ + 'Protein ID', 'Protein Name', 'Species', 'Taxon ID', 'Gene', + 'Protein Existence Level', 'Gene Priority', 'SwissProt Reviewed', +] + + +def test_a_search_that_hits_nothing_keeps_the_metadata_column_types(proteome): + # No protein is hit, so no metadata row is read. The metadata columns must still have the + # types a search with hits gives them, not turn into untyped nulls. + hit, miss = search(proteome, ['MKVLAT']), search(proteome, ['WWWWWW']) + assert miss['Matched Sequence'].is_null().all() + assert {c: miss.schema[c] for c in METADATA_COLUMNS} == {c: hit.schema[c] for c in METADATA_COLUMNS} + + +SEARCHES = { + 'exact': dict(query=DATA / 'exact_match_query.fasta'), + 'mismatch': dict(query=DATA / 'mismatching_query.fasta', max_mismatches=3), + 'best match': dict(query=DATA / 'best_match_query.fasta', best_match=True), + 'discontinuous with linear': dict( + query=['L354, V420, G461, Q468, E486, K499, D501, M503, G509', 'T819, L822, A835, F840, S871, Y884, P886', + 'CNERFERLMKELSE', 'FAKEPEPTIDE'], + max_mismatches=1), + 'one indel': dict(query=DATA / 'indel_query.fasta', max_indels=1), + 'two indels': dict(query=DATA / 'indel_query.fasta', max_indels=2), + 'nothing hit': dict(query=['WWWWWWWWW']), +} + + +@pytest.mark.parametrize('scenario', SEARCHES) +def test_reading_only_the_hit_proteins_gives_the_frame_reading_every_protein_gives(monkeypatch, tmp_path, scenario): + # Capture the raw columns each search hands _to_dataframe, then build the frame from those + # same columns twice: reading only the hit proteins' metadata, as searches now do, and + # reading every protein's, as they did before. The input is identical, so the comparison + # is exact even for searches whose hit order changes from run to run. + captured = [] + to_dataframe = Matcher._to_dataframe + + def capture(self, cols, pepidx_path, is_indels=False): + captured.append((self, cols, pepidx_path, is_indels)) + return to_dataframe(self, cols, pepidx_path, is_indels) + + monkeypatch.setattr(Matcher, '_to_dataframe', capture) + Matcher(proteome_file=DATA / 'proteome.fasta', preprocessed_files_path=str(tmp_path), **SEARCHES[scenario]).match() + monkeypatch.undo() + assert captured + + hits = 0 + for matcher, cols, pepidx_path, is_indels in captured: + hit_proteins_only = matcher._to_dataframe(cols, pepidx_path, is_indels) + monkeypatch.setattr(matcher_module, 'rs_metadata', lambda path, protein_nums=None: rs_metadata(path)) + every_protein = matcher._to_dataframe(cols, pepidx_path, is_indels) + monkeypatch.undo() + plt.assert_frame_equal(hit_proteins_only, every_protein) + hits += hit_proteins_only['Protein ID'].is_not_null().sum() + assert (hits > 0) == (scenario != 'nothing hit') # the comparison covered real hits + + +def test_a_hit_from_a_stale_index_raises_instead_of_losing_its_metadata(tmp_path): + # .pepidx files are reused by name, so editing a FASTA in place can leave its 5-mer index + # three proteins long while a rebuilt 2-mer index has four. An indel batch of short and long + # queries searches both and reads metadata from the 5-mer index, so a hit on the fourth + # protein has no metadata row: it must fail loudly, not come back with blank metadata. + fasta = tmp_path / 'three.fasta' + fasta.write_text(PROTEOME_FASTA) + Preprocessor(str(fasta), preprocessed_files_path=str(tmp_path)).preprocess(5) + fasta.write_text(PROTEOME_FASTA + '>sp|P44444|DDDD_HUMAN Delta protein OS=Homo sapiens OX=9606 GN=DDD PE=1 SV=1\nWWYYWWCCHH\n') + Preprocessor(str(fasta), preprocessed_files_path=str(tmp_path)).preprocess(2) + matcher = Matcher( + query=['YYWWC', 'MKVLATGGWQ'], + proteome_file=str(fasta), + max_indels=1, + preprocessed_files_path=str(tmp_path), + output_format='dataframe', + ) + with pytest.raises(ValueError, match='protein number 4 .* different versions'): + matcher.match()