metadata: read only the proteins a search hits

.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

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(

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<Vec<u32>>) -> PyResult<matching::MetaColumns> {
+    matching::run_metadata(pepidx_path, protein_nums).map_err(pyo3::exceptions::PyValueError::new_err)
 }
 
 #[pyfunction]

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<Vec<u32>>) -> Result<MetaColumns, String> {
     let n = index.num_proteins;
-    let mut pnum = Vec::with_capacity(n);
+    let proteins: Vec<u32> = 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<String>; 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<Vec<u32>>) -> Result<MetaColumns, String> {
     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 {

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()