From 784ef4b30b0540e030eabc187891b252aff45b40 Mon Sep 17 00:00:00 2001 From: "J. Sebastian Paez" Date: Thu, 27 Aug 2026 19:08:19 -0700 Subject: [PATCH 01/19] refactor: make a source id polymorphic, keeping the shape the file used An id is an id whether the source wrote an integer or a string, so `SourceId`/`OwnedSourceId` carry either and `#[serde(untagged)]` renders each as what it is. Carafe's `"id": 7` stays a JSON number; a DIA-NN `transition_group_id` will stay a string rather than being coerced into a u64 or replaced by one we made up. `SourceIds` gains a `Text` arm, stored blob+offsets like `seq_strip_blob`, and `SourceIds::owned` keeps an all-numeric library on the dense integer column rather than paying for a blob. The parquet `library_id` and `decoy_group_id` columns become Utf8, format version 3. `HasQueryData::id` is gone -- it returned a u64 nothing consumed. No behaviour change on its own: every id is still Numeric until the readers start propagating. --- python/timsquery_pyo3/src/chromatogram.rs | 4 +- python/timsquery_pyo3/src/elution_group.rs | 4 +- python/timsquery_pyo3/src/iterator.rs | 10 +- python/timsquery_pyo3/src/lib.rs | 15 ++ python/timsquery_pyo3/src/spectrum.rs | 8 +- rust/apex_sim/src/sim.rs | 2 +- .../models/aggregators/chromatogram_agg.rs | 10 +- .../src/models/aggregators/point_agg.rs | 8 +- .../src/models/aggregators/spectrum_agg.rs | 10 +- rust/timsquery/src/models/mod.rs | 2 + rust/timsquery/src/models/query_handle.rs | 7 +- rust/timsquery/src/models/source_id.rs | 197 ++++++++++++++++-- rust/timsquery/src/models/target.rs | 15 +- rust/timsquery/src/models/target_columns.rs | 43 +++- .../src/serde/chromatogram_output.rs | 4 +- rust/timsquery/src/serde/library_file.rs | 8 +- rust/timsquery/src/traits/queriable_data.rs | 5 - rust/timsquery/src/traits/query_geom.rs | 18 +- rust/timsquery/tests/carafe_contract/main.rs | 9 +- rust/timsquery_cli/src/processing.rs | 6 +- rust/timsquery_viewer/src/calibration.rs | 2 +- rust/timsquery_viewer/src/plot_renderer.rs | 4 +- .../src/data_sources/reference_library.rs | 23 +- rust/timsseek/src/ml/qvalues.rs | 25 +-- rust/timsseek/src/models/sequence.rs | 2 +- rust/timsseek/src/scoring/apex_finding.rs | 7 +- rust/timsseek/src/scoring/blocks/identity.rs | 31 +-- rust/timsseek/src/scoring/parquet_writer.rs | 6 +- rust/timsseek/src/scoring/pipeline.rs | 5 +- rust/timsseek_cli/src/processing.rs | 38 ++-- 30 files changed, 372 insertions(+), 156 deletions(-) diff --git a/python/timsquery_pyo3/src/chromatogram.rs b/python/timsquery_pyo3/src/chromatogram.rs index 99464447..471e8c81 100644 --- a/python/timsquery_pyo3/src/chromatogram.rs +++ b/python/timsquery_pyo3/src/chromatogram.rs @@ -63,8 +63,8 @@ impl PyChromatogramResult { } #[getter] - fn id(&self) -> u64 { - self.collector.id + fn id<'py>(&self, py: pyo3::Python<'py>) -> pyo3::PyResult> { + crate::source_id_to_py(py, &self.collector.id) } fn __repr__(&self) -> String { diff --git a/python/timsquery_pyo3/src/elution_group.rs b/python/timsquery_pyo3/src/elution_group.rs index 2fe68c6a..66ad74db 100644 --- a/python/timsquery_pyo3/src/elution_group.rs +++ b/python/timsquery_pyo3/src/elution_group.rs @@ -60,8 +60,8 @@ impl PyElutionGroup { } #[getter] - fn id(&self) -> u64 { - self.inner.id() + fn id<'py>(&self, py: pyo3::Python<'py>) -> pyo3::PyResult> { + crate::source_id_to_py(py, &self.inner.id().to_owned_id()) } #[getter] diff --git a/python/timsquery_pyo3/src/iterator.rs b/python/timsquery_pyo3/src/iterator.rs index 44c09c36..a6a65a22 100644 --- a/python/timsquery_pyo3/src/iterator.rs +++ b/python/timsquery_pyo3/src/iterator.rs @@ -26,8 +26,7 @@ pub enum ToleranceSource { /// collector pool is never exposed — it reuses Rust-side buffers across chunks. #[pyclass(frozen)] pub struct PyChromatogramArrays { - #[pyo3(get)] - id: u64, + id: timsquery::models::OwnedSourceId, precursor_intensities: Py, fragment_intensities: Py, #[pyo3(get)] @@ -42,6 +41,11 @@ pub struct PyChromatogramArrays { #[pymethods] impl PyChromatogramArrays { + #[getter] + fn id<'py>(&self, py: pyo3::Python<'py>) -> pyo3::PyResult> { + crate::source_id_to_py(py, &self.id) + } + #[getter] fn precursor_intensities<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> { self.precursor_intensities.clone_ref(py).into_bound(py) @@ -72,7 +76,7 @@ fn extract_arrays( let rt = collector.rt_range_milis(); Ok(PyChromatogramArrays { - id: collector.id, + id: collector.id.clone(), precursor_intensities: prec_np.into_any().unbind(), fragment_intensities: frag_np.into_any().unbind(), precursor_labels: collector diff --git a/python/timsquery_pyo3/src/lib.rs b/python/timsquery_pyo3/src/lib.rs index a74908a2..c12151b7 100644 --- a/python/timsquery_pyo3/src/lib.rs +++ b/python/timsquery_pyo3/src/lib.rs @@ -21,6 +21,21 @@ mod spectrum; mod tolerance; use pyo3::prelude::*; +use timsquery::models::OwnedSourceId; + +/// A source id keeps the shape the library gave it, so Python sees an `int` +/// for a numeric id and a `str` for a text one (DIA-NN's +/// `transition_group_id`) rather than one coerced into the other. +pub(crate) fn source_id_to_py<'py>( + py: Python<'py>, + id: &OwnedSourceId, +) -> PyResult> { + use pyo3::IntoPyObject; + match id { + OwnedSourceId::Numeric(n) => Ok(n.into_pyobject(py)?.into_any()), + OwnedSourceId::Text(s) => Ok(s.into_pyobject(py)?.into_any()), + } +} #[pymodule] fn timsquery_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/python/timsquery_pyo3/src/spectrum.rs b/python/timsquery_pyo3/src/spectrum.rs index 1661ecce..9a962a14 100644 --- a/python/timsquery_pyo3/src/spectrum.rs +++ b/python/timsquery_pyo3/src/spectrum.rs @@ -61,8 +61,8 @@ impl PySpectralResult { /// The elution group id. #[getter] - fn id(&self) -> u64 { - self.collector.id + fn id<'py>(&self, py: pyo3::Python<'py>) -> pyo3::PyResult> { + crate::source_id_to_py(py, &self.collector.id) } fn __repr__(&self) -> String { @@ -153,8 +153,8 @@ impl PyMzMobilityResult { /// The elution group id. #[getter] - fn id(&self) -> u64 { - self.collector.id + fn id<'py>(&self, py: pyo3::Python<'py>) -> pyo3::PyResult> { + crate::source_id_to_py(py, &self.collector.id) } fn __repr__(&self) -> String { diff --git a/rust/apex_sim/src/sim.rs b/rust/apex_sim/src/sim.rs index d1d1428e..6174ddda 100644 --- a/rust/apex_sim/src/sim.rs +++ b/rust/apex_sim/src/sim.rs @@ -340,7 +340,7 @@ pub fn build(params: &SimParams) -> SimData { TupleRange::try_new(map(0), map(n - 1)).expect("start < end for positive period"); let chromatograms = ChromatogramCollector:: { - id: 0, + id: Default::default(), mobility_ook0: 1.0, rt_seconds: (map((realized_apex as usize).min(n - 1)) as f32) / 1000.0, precursor_mono_mz: dummy_mz, diff --git a/rust/timsquery/src/models/aggregators/chromatogram_agg.rs b/rust/timsquery/src/models/aggregators/chromatogram_agg.rs index 80b186aa..35244c6a 100644 --- a/rust/timsquery/src/models/aggregators/chromatogram_agg.rs +++ b/rust/timsquery/src/models/aggregators/chromatogram_agg.rs @@ -26,7 +26,7 @@ use timscentroid::utils::TupleRange; #[derive(Debug, Clone, Serialize)] pub struct ChromatogramCollector { // Query scalars carried from the eg at reset time. - pub id: u64, + pub id: crate::models::OwnedSourceId, pub mobility_ook0: f32, pub rt_seconds: f32, pub precursor_mono_mz: f64, @@ -90,7 +90,7 @@ impl ChromatogramCollector { let fragments = MzMajorIntensityArray::try_new_empty(fragment_order, num_cycles, start.index())?; Ok(Self { - id: eg.output_id(), + id: eg.output_id().to_owned_id(), mobility_ook0: eg.mobility_ook0(), rt_seconds: eg.rt_seconds(), precursor_mono_mz: eg.mono_precursor_mz(), @@ -136,7 +136,7 @@ impl ChromatogramCollector { return Err(DataProcessingError::ExpectedNonEmptyData); } - self.id = eg.output_id(); + self.id = eg.output_id().to_owned_id(); self.mobility_ook0 = mobility_override.unwrap_or_else(|| eg.mobility_ook0()); self.rt_seconds = rt_override.unwrap_or_else(|| eg.rt_seconds()); self.precursor_mono_mz = eg.mono_precursor_mz(); @@ -219,10 +219,6 @@ impl ChromatogramCollector { } impl HasQueryData for ChromatogramCollector { - fn id(&self) -> u64 { - self.id - } - fn precursor_mz_limits(&self) -> (f64, f64) { self.precursor_mz_limits } diff --git a/rust/timsquery/src/models/aggregators/point_agg.rs b/rust/timsquery/src/models/aggregators/point_agg.rs index 672c55e4..17caac9f 100644 --- a/rust/timsquery/src/models/aggregators/point_agg.rs +++ b/rust/timsquery/src/models/aggregators/point_agg.rs @@ -10,7 +10,7 @@ const POINT_INLINE_CAP: usize = 13; #[derive(Debug, Clone, Serialize)] pub struct PointIntensityAggregator { - pub id: u64, + pub id: crate::models::OwnedSourceId, pub mobility_ook0: f32, pub rt_seconds: f32, pub precursor_mono_mz: f64, @@ -42,7 +42,7 @@ impl PointIntensityAggregator { fragment_mzs.push(mz); } Self { - id: eg.output_id(), + id: eg.output_id().to_owned_id(), mobility_ook0: eg.mobility_ook0(), rt_seconds: eg.rt_seconds(), precursor_mono_mz: eg.mono_precursor_mz(), @@ -58,10 +58,6 @@ impl PointIntensityAggregator { } impl HasQueryData for PointIntensityAggregator { - fn id(&self) -> u64 { - self.id - } - fn precursor_mz_limits(&self) -> (f64, f64) { self.precursor_mz_limits } diff --git a/rust/timsquery/src/models/aggregators/spectrum_agg.rs b/rust/timsquery/src/models/aggregators/spectrum_agg.rs index b30514df..38fa564e 100644 --- a/rust/timsquery/src/models/aggregators/spectrum_agg.rs +++ b/rust/timsquery/src/models/aggregators/spectrum_agg.rs @@ -33,7 +33,7 @@ const SPEC_INLINE_CAP: usize = 13; #[derive(Debug, Clone, Serialize)] pub struct SpectralCollector { // Query scalars carried from eg at construction / reset. - pub id: u64, + pub id: crate::models::OwnedSourceId, pub mobility_ook0: f32, pub rt_seconds: f32, pub precursor_mono_mz: f64, @@ -53,7 +53,7 @@ pub struct SpectralCollector { impl SpectralCollector { pub fn new(eg: &impl QueryGeom