openmassspec_core/
arrow.rs

1//! Apache Arrow bridge for [`crate::SpectrumRecord`].
2//!
3//! This module is gated behind the `arrow` Cargo feature. It exposes:
4//!
5//! * [`spectrum_record_schema`]: the canonical Arrow [`Schema`] for a
6//!   stream of `SpectrumRecord`s.
7//! * [`SpectrumBatchBuilder`]: a streaming builder that accumulates rows
8//!   and produces a [`RecordBatch`] when finalized.
9//!
10//! The schema is intentionally flat (one row per spectrum). Peak arrays
11//! are stored as `LargeList<Float64>` / `LargeList<Float32>` so that
12//! individual spectra can exceed `i32::MAX` peaks without truncation.
13//! Precursor fields are inlined as nullable scalar columns; an MS1
14//! spectrum has all precursor columns null. This shape is friendly to
15//! Polars / DataFusion / DuckDB consumers.
16//!
17//! The Arrow schema is part of this crate's stable surface for the
18//! purposes of consumers that pin `openmassspec-core` directly; any column
19//! addition is a minor-version bump, any removal or rename is breaking.
20
21use std::sync::Arc;
22
23use arrow_array::builder::{
24    ArrayBuilder, Float32Builder, Float64Builder, Int32Builder, LargeListBuilder, StringBuilder,
25    UInt32Builder, UInt8Builder,
26};
27use arrow_array::{ArrayRef, RecordBatch};
28use arrow_schema::{DataType, Field, Schema, SchemaRef};
29
30use crate::{Activation, Analyzer, MobilityArrayKind, Polarity, ScanMode, SpectrumRecord};
31
32/// Return the canonical [`Schema`] for a `RecordBatch` of spectra.
33pub fn spectrum_record_schema() -> SchemaRef {
34    let mz_item = Arc::new(Field::new("item", DataType::Float64, false));
35    let int_item = Arc::new(Field::new("item", DataType::Float32, false));
36    let mob_item = Arc::new(Field::new("item", DataType::Float32, false));
37    Arc::new(Schema::new(vec![
38        Field::new("index", DataType::UInt32, false),
39        Field::new("scan_number", DataType::UInt32, false),
40        Field::new("native_id", DataType::Utf8, false),
41        Field::new("ms_level", DataType::UInt8, false),
42        Field::new("polarity", DataType::Utf8, true),
43        Field::new("scan_mode", DataType::Utf8, true),
44        Field::new("analyzer", DataType::Utf8, true),
45        Field::new("filter", DataType::Utf8, true),
46        Field::new("retention_time_sec", DataType::Float64, false),
47        Field::new("total_ion_current", DataType::Float64, true),
48        Field::new("base_peak_mz", DataType::Float64, true),
49        Field::new("base_peak_intensity", DataType::Float64, true),
50        Field::new("low_mz", DataType::Float64, true),
51        Field::new("high_mz", DataType::Float64, true),
52        Field::new("ion_injection_time_ms", DataType::Float64, true),
53        Field::new("inv_mobility", DataType::Float64, true),
54        Field::new("faims_cv", DataType::Float64, true),
55        Field::new("precursor_target_mz", DataType::Float64, true),
56        Field::new("precursor_selected_mz", DataType::Float64, true),
57        Field::new("precursor_isolation_width", DataType::Float64, true),
58        Field::new("precursor_charge", DataType::Int32, true),
59        Field::new("precursor_intensity", DataType::Float64, true),
60        Field::new("precursor_collision_energy", DataType::Float64, true),
61        Field::new("precursor_ce_is_nce", DataType::UInt8, true),
62        Field::new("precursor_native_id", DataType::Utf8, true),
63        Field::new("precursor_activation", DataType::Utf8, true),
64        Field::new("precursor_analyzer", DataType::Utf8, true),
65        Field::new("precursor_ccs", DataType::Float64, true),
66        Field::new_large_list("mz", mz_item, false),
67        Field::new_large_list("intensity", int_item, false),
68        Field::new_large_list("inv_mobility_per_peak", mob_item, true),
69        Field::new("mobility_array_kind", DataType::Utf8, true),
70    ]))
71}
72
73fn polarity_str(p: Polarity) -> &'static str {
74    match p {
75        Polarity::Positive => "positive",
76        Polarity::Negative => "negative",
77    }
78}
79
80fn scan_mode_str(m: ScanMode) -> &'static str {
81    match m {
82        ScanMode::Profile => "profile",
83        ScanMode::Centroid => "centroid",
84    }
85}
86
87fn analyzer_str(a: Analyzer) -> &'static str {
88    match a {
89        Analyzer::ITMS => "itms",
90        Analyzer::TQMS => "tqms",
91        Analyzer::SQMS => "sqms",
92        Analyzer::TOFMS => "tof",
93        Analyzer::FTMS => "ftms",
94        Analyzer::Sector => "sector",
95    }
96}
97
98fn activation_str(a: Activation) -> &'static str {
99    match a {
100        Activation::CID => "cid",
101        Activation::HCD => "hcd",
102        Activation::ETD => "etd",
103        Activation::ECD => "ecd",
104        Activation::UVPD => "uvpd",
105        Activation::PQD => "pqd",
106        Activation::PD => "pd",
107        Activation::SID => "sid",
108        Activation::EThcD => "ethcd",
109        Activation::IRMPD => "irmpd",
110        Activation::MPID => "mpid",
111    }
112}
113
114fn mobility_kind_str(k: MobilityArrayKind) -> &'static str {
115    match k {
116        MobilityArrayKind::InverseReducedVsPerCm2 => "inverse_reduced_k0",
117        MobilityArrayKind::DriftTimeMilliseconds => "drift_time_ms",
118    }
119}
120
121/// Streaming builder that accumulates `SpectrumRecord`s and produces a
122/// single Arrow [`RecordBatch`] when finalized.
123///
124/// All rows in a batch share one `mobility_array_kind` value, recorded
125/// once at construction. Push rows with [`push`](Self::push); call
126/// [`finish`](Self::finish) to materialize the batch.
127pub struct SpectrumBatchBuilder {
128    schema: SchemaRef,
129    mobility_kind: Option<MobilityArrayKind>,
130    index: UInt32Builder,
131    scan_number: UInt32Builder,
132    native_id: StringBuilder,
133    ms_level: UInt8Builder,
134    polarity: StringBuilder,
135    scan_mode: StringBuilder,
136    analyzer: StringBuilder,
137    filter: StringBuilder,
138    retention_time_sec: Float64Builder,
139    total_ion_current: Float64Builder,
140    base_peak_mz: Float64Builder,
141    base_peak_intensity: Float64Builder,
142    low_mz: Float64Builder,
143    high_mz: Float64Builder,
144    ion_injection_time_ms: Float64Builder,
145    inv_mobility: Float64Builder,
146    faims_cv: Float64Builder,
147    precursor_target_mz: Float64Builder,
148    precursor_selected_mz: Float64Builder,
149    precursor_isolation_width: Float64Builder,
150    precursor_charge: Int32Builder,
151    precursor_intensity: Float64Builder,
152    precursor_collision_energy: Float64Builder,
153    precursor_ce_is_nce: UInt8Builder,
154    precursor_native_id: StringBuilder,
155    precursor_activation: StringBuilder,
156    precursor_analyzer: StringBuilder,
157    precursor_ccs: Float64Builder,
158    mz: LargeListBuilder<Float64Builder>,
159    intensity: LargeListBuilder<Float32Builder>,
160    inv_mobility_per_peak: LargeListBuilder<Float32Builder>,
161    mobility_array_kind_col: StringBuilder,
162}
163
164impl SpectrumBatchBuilder {
165    /// Create a new builder. Pass the `mobility_array_kind` from the
166    /// source's [`crate::RunMetadata`] so the resulting Arrow batch
167    /// carries the unit/CV interpretation alongside the data.
168    pub fn new(mobility_kind: Option<MobilityArrayKind>) -> Self {
169        Self {
170            schema: spectrum_record_schema(),
171            mobility_kind,
172            index: UInt32Builder::new(),
173            scan_number: UInt32Builder::new(),
174            native_id: StringBuilder::new(),
175            ms_level: UInt8Builder::new(),
176            polarity: StringBuilder::new(),
177            scan_mode: StringBuilder::new(),
178            analyzer: StringBuilder::new(),
179            filter: StringBuilder::new(),
180            retention_time_sec: Float64Builder::new(),
181            total_ion_current: Float64Builder::new(),
182            base_peak_mz: Float64Builder::new(),
183            base_peak_intensity: Float64Builder::new(),
184            low_mz: Float64Builder::new(),
185            high_mz: Float64Builder::new(),
186            ion_injection_time_ms: Float64Builder::new(),
187            inv_mobility: Float64Builder::new(),
188            faims_cv: Float64Builder::new(),
189            precursor_target_mz: Float64Builder::new(),
190            precursor_selected_mz: Float64Builder::new(),
191            precursor_isolation_width: Float64Builder::new(),
192            precursor_charge: Int32Builder::new(),
193            precursor_intensity: Float64Builder::new(),
194            precursor_collision_energy: Float64Builder::new(),
195            precursor_ce_is_nce: UInt8Builder::new(),
196            precursor_native_id: StringBuilder::new(),
197            precursor_activation: StringBuilder::new(),
198            precursor_analyzer: StringBuilder::new(),
199            precursor_ccs: Float64Builder::new(),
200            mz: LargeListBuilder::new(Float64Builder::new()).with_field(Arc::new(Field::new(
201                "item",
202                DataType::Float64,
203                false,
204            ))),
205            intensity: LargeListBuilder::new(Float32Builder::new())
206                .with_field(Arc::new(Field::new("item", DataType::Float32, false))),
207            inv_mobility_per_peak: LargeListBuilder::new(Float32Builder::new())
208                .with_field(Arc::new(Field::new("item", DataType::Float32, false))),
209            mobility_array_kind_col: StringBuilder::new(),
210        }
211    }
212
213    /// Schema for the batch produced by this builder.
214    pub fn schema(&self) -> SchemaRef {
215        self.schema.clone()
216    }
217
218    /// Append one spectrum row.
219    pub fn push(&mut self, rec: &SpectrumRecord) {
220        self.index.append_value(rec.index as u32);
221        self.scan_number.append_value(rec.scan_number);
222        self.native_id.append_value(&rec.native_id);
223        self.ms_level.append_value(rec.ms_level as u8);
224        self.polarity.append_option(rec.polarity.map(polarity_str));
225        self.scan_mode
226            .append_option(rec.scan_mode.map(scan_mode_str));
227        self.analyzer.append_option(rec.analyzer.map(analyzer_str));
228        self.filter.append_option(rec.filter.as_deref());
229        self.retention_time_sec.append_value(rec.retention_time_sec);
230        self.total_ion_current.append_option(rec.total_ion_current);
231        self.base_peak_mz.append_option(rec.base_peak_mz);
232        self.base_peak_intensity
233            .append_option(rec.base_peak_intensity);
234        self.low_mz.append_option(rec.low_mz);
235        self.high_mz.append_option(rec.high_mz);
236        self.ion_injection_time_ms
237            .append_option(rec.ion_injection_time_ms);
238        self.inv_mobility.append_option(rec.inv_mobility);
239        self.faims_cv.append_option(rec.faims_cv);
240
241        match &rec.precursor {
242            Some(p) => {
243                self.precursor_target_mz.append_option(p.target_mz);
244                self.precursor_selected_mz.append_option(p.selected_mz);
245                self.precursor_isolation_width
246                    .append_option(p.isolation_width);
247                self.precursor_charge.append_option(p.charge);
248                self.precursor_intensity.append_option(p.intensity);
249                self.precursor_collision_energy
250                    .append_option(p.collision_energy);
251                self.precursor_ce_is_nce.append_value(u8::from(p.ce_is_nce));
252                self.precursor_native_id
253                    .append_option(p.precursor_native_id.as_deref());
254                self.precursor_activation
255                    .append_option(p.activation.map(activation_str));
256                self.precursor_analyzer
257                    .append_option(p.analyzer.map(analyzer_str));
258                self.precursor_ccs.append_option(p.ccs);
259            }
260            None => {
261                self.precursor_target_mz.append_null();
262                self.precursor_selected_mz.append_null();
263                self.precursor_isolation_width.append_null();
264                self.precursor_charge.append_null();
265                self.precursor_intensity.append_null();
266                self.precursor_collision_energy.append_null();
267                self.precursor_ce_is_nce.append_null();
268                self.precursor_native_id.append_null();
269                self.precursor_activation.append_null();
270                self.precursor_analyzer.append_null();
271                self.precursor_ccs.append_null();
272            }
273        }
274
275        for &v in &rec.mz {
276            self.mz.values().append_value(v);
277        }
278        self.mz.append(true);
279        for &v in &rec.intensity {
280            self.intensity.values().append_value(v);
281        }
282        self.intensity.append(true);
283        match &rec.inv_mobility_per_peak {
284            Some(mob) => {
285                for &v in mob {
286                    self.inv_mobility_per_peak.values().append_value(v);
287                }
288                self.inv_mobility_per_peak.append(true);
289            }
290            None => self.inv_mobility_per_peak.append(false),
291        }
292        self.mobility_array_kind_col
293            .append_option(self.mobility_kind.map(mobility_kind_str));
294    }
295
296    /// Number of rows accumulated so far.
297    pub fn len(&self) -> usize {
298        self.index.len()
299    }
300
301    /// `true` if no rows have been pushed.
302    pub fn is_empty(&self) -> bool {
303        self.len() == 0
304    }
305
306    /// Materialize the accumulated rows into a [`RecordBatch`].
307    pub fn finish(mut self) -> Result<RecordBatch, arrow_schema::ArrowError> {
308        let arrays: Vec<ArrayRef> = vec![
309            Arc::new(self.index.finish()),
310            Arc::new(self.scan_number.finish()),
311            Arc::new(self.native_id.finish()),
312            Arc::new(self.ms_level.finish()),
313            Arc::new(self.polarity.finish()),
314            Arc::new(self.scan_mode.finish()),
315            Arc::new(self.analyzer.finish()),
316            Arc::new(self.filter.finish()),
317            Arc::new(self.retention_time_sec.finish()),
318            Arc::new(self.total_ion_current.finish()),
319            Arc::new(self.base_peak_mz.finish()),
320            Arc::new(self.base_peak_intensity.finish()),
321            Arc::new(self.low_mz.finish()),
322            Arc::new(self.high_mz.finish()),
323            Arc::new(self.ion_injection_time_ms.finish()),
324            Arc::new(self.inv_mobility.finish()),
325            Arc::new(self.faims_cv.finish()),
326            Arc::new(self.precursor_target_mz.finish()),
327            Arc::new(self.precursor_selected_mz.finish()),
328            Arc::new(self.precursor_isolation_width.finish()),
329            Arc::new(self.precursor_charge.finish()),
330            Arc::new(self.precursor_intensity.finish()),
331            Arc::new(self.precursor_collision_energy.finish()),
332            Arc::new(self.precursor_ce_is_nce.finish()),
333            Arc::new(self.precursor_native_id.finish()),
334            Arc::new(self.precursor_activation.finish()),
335            Arc::new(self.precursor_analyzer.finish()),
336            Arc::new(self.precursor_ccs.finish()),
337            Arc::new(self.mz.finish()),
338            Arc::new(self.intensity.finish()),
339            Arc::new(self.inv_mobility_per_peak.finish()),
340            Arc::new(self.mobility_array_kind_col.finish()),
341        ];
342        RecordBatch::try_new(self.schema.clone(), arrays)
343    }
344}
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349    use crate::{PrecursorInfo, SpectrumRecord};
350    use arrow_array::Array;
351
352    fn rec(index: usize, ms_level: u32, n_peaks: usize, with_mob: bool) -> SpectrumRecord {
353        let mz: Vec<f64> = (0..n_peaks).map(|i| 100.0 + i as f64).collect();
354        let intensity: Vec<f32> = (0..n_peaks).map(|i| 10.0 + i as f32).collect();
355        let mobility = if with_mob {
356            Some((0..n_peaks).map(|i| 0.5 + i as f32 * 0.01).collect())
357        } else {
358            None
359        };
360        SpectrumRecord {
361            index,
362            scan_number: (index + 1) as u32,
363            native_id: format!("scan={}", index + 1),
364            ms_level,
365            polarity: Some(Polarity::Positive),
366            scan_mode: Some(ScanMode::Centroid),
367            analyzer: Some(Analyzer::TOFMS),
368            filter: None,
369            retention_time_sec: index as f64,
370            total_ion_current: Some(intensity.iter().map(|&v| v as f64).sum()),
371            base_peak_mz: mz.last().copied(),
372            base_peak_intensity: intensity.last().map(|&v| v as f64),
373            low_mz: mz.first().copied(),
374            high_mz: mz.last().copied(),
375            ion_injection_time_ms: None,
376            inv_mobility: None,
377            faims_cv: if index == 0 { Some(-40.0) } else { None },
378            precursor: if ms_level >= 2 {
379                Some(PrecursorInfo {
380                    target_mz: Some(500.0),
381                    selected_mz: Some(500.5),
382                    isolation_width: Some(2.0),
383                    charge: Some(2),
384                    ..Default::default()
385                })
386            } else {
387                None
388            },
389            mz,
390            intensity,
391            inv_mobility_per_peak: mobility,
392        }
393    }
394
395    #[test]
396    fn schema_round_trip() {
397        let mut b = SpectrumBatchBuilder::new(Some(MobilityArrayKind::DriftTimeMilliseconds));
398        b.push(&rec(0, 1, 3, true));
399        b.push(&rec(1, 2, 4, false));
400        let batch = b.finish().unwrap();
401        assert_eq!(batch.num_rows(), 2);
402        assert_eq!(batch.schema().fields().len(), 32);
403        let mz_col = batch
404            .column_by_name("mz")
405            .unwrap()
406            .as_any()
407            .downcast_ref::<arrow_array::LargeListArray>()
408            .unwrap();
409        assert_eq!(mz_col.value_length(0), 3);
410        assert_eq!(mz_col.value_length(1), 4);
411        let mob_col = batch
412            .column_by_name("inv_mobility_per_peak")
413            .unwrap()
414            .as_any()
415            .downcast_ref::<arrow_array::LargeListArray>()
416            .unwrap();
417        assert!(mob_col.is_valid(0));
418        assert!(mob_col.is_null(1));
419        let faims_col = batch
420            .column_by_name("faims_cv")
421            .unwrap()
422            .as_any()
423            .downcast_ref::<arrow_array::Float64Array>()
424            .unwrap();
425        assert_eq!(faims_col.value(0), -40.0);
426        assert!(faims_col.is_null(1));
427    }
428}
openmassspec_core/arrow.rs

openmassspec_core/
arrow.rs
