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, MapBuilder,
25    MapFieldNames, StringBuilder, 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// Arrow-spec map field names for the `extra` column.
33const EXTRA_ENTRIES_FIELD: &str = "entries";
34const EXTRA_KEY_FIELD: &str = "key";
35const EXTRA_VALUE_FIELD: &str = "value";
36
37/// Return the canonical [`Schema`] for a `RecordBatch` of spectra.
38pub fn spectrum_record_schema() -> SchemaRef {
39    let mz_item = Arc::new(Field::new("item", DataType::Float64, false));
40    let int_item = Arc::new(Field::new("item", DataType::Float32, false));
41    let mob_item = Arc::new(Field::new("item", DataType::Float32, false));
42    Arc::new(Schema::new(vec![
43        Field::new("index", DataType::UInt32, false),
44        Field::new("scan_number", DataType::UInt32, false),
45        Field::new("native_id", DataType::Utf8, false),
46        Field::new("ms_level", DataType::UInt8, false),
47        Field::new("polarity", DataType::Utf8, true),
48        Field::new("scan_mode", DataType::Utf8, true),
49        Field::new("analyzer", DataType::Utf8, true),
50        Field::new("filter", DataType::Utf8, true),
51        Field::new("retention_time_sec", DataType::Float64, false),
52        Field::new("total_ion_current", DataType::Float64, true),
53        Field::new("base_peak_mz", DataType::Float64, true),
54        Field::new("base_peak_intensity", DataType::Float64, true),
55        Field::new("low_mz", DataType::Float64, true),
56        Field::new("high_mz", DataType::Float64, true),
57        Field::new("ion_injection_time_ms", DataType::Float64, true),
58        Field::new("inv_mobility", DataType::Float64, true),
59        Field::new("faims_cv", DataType::Float64, true),
60        Field::new("precursor_target_mz", DataType::Float64, true),
61        Field::new("precursor_selected_mz", DataType::Float64, true),
62        Field::new("precursor_isolation_width", DataType::Float64, true),
63        Field::new("precursor_charge", DataType::Int32, true),
64        Field::new("precursor_intensity", DataType::Float64, true),
65        Field::new("precursor_collision_energy", DataType::Float64, true),
66        Field::new("precursor_ce_is_nce", DataType::UInt8, true),
67        Field::new("precursor_native_id", DataType::Utf8, true),
68        Field::new("precursor_activation", DataType::Utf8, true),
69        Field::new("precursor_analyzer", DataType::Utf8, true),
70        Field::new("precursor_ccs", DataType::Float64, true),
71        Field::new_large_list("mz", mz_item, false),
72        Field::new_large_list("intensity", int_item, false),
73        Field::new_large_list("inv_mobility_per_peak", mob_item, true),
74        Field::new("mobility_array_kind", DataType::Utf8, true),
75        Field::new("acquisition_event_id", DataType::UInt32, true),
76        Field::new_map(
77            "extra",
78            EXTRA_ENTRIES_FIELD,
79            Field::new(EXTRA_KEY_FIELD, DataType::Utf8, false),
80            Field::new(EXTRA_VALUE_FIELD, DataType::Utf8, true),
81            false,
82            false,
83        ),
84    ]))
85}
86
87fn polarity_str(p: Polarity) -> &'static str {
88    match p {
89        Polarity::Positive => "positive",
90        Polarity::Negative => "negative",
91    }
92}
93
94fn scan_mode_str(m: ScanMode) -> &'static str {
95    match m {
96        ScanMode::Profile => "profile",
97        ScanMode::Centroid => "centroid",
98    }
99}
100
101fn analyzer_str(a: Analyzer) -> &'static str {
102    match a {
103        Analyzer::ITMS => "itms",
104        Analyzer::TQMS => "tqms",
105        Analyzer::SQMS => "sqms",
106        Analyzer::TOFMS => "tof",
107        Analyzer::FTMS => "ftms",
108        Analyzer::Sector => "sector",
109    }
110}
111
112fn activation_str(a: Activation) -> &'static str {
113    match a {
114        Activation::CID => "cid",
115        Activation::HCD => "hcd",
116        Activation::ETD => "etd",
117        Activation::ECD => "ecd",
118        Activation::UVPD => "uvpd",
119        Activation::PQD => "pqd",
120        Activation::PD => "pd",
121        Activation::SID => "sid",
122        Activation::EThcD => "ethcd",
123        Activation::IRMPD => "irmpd",
124        Activation::MPID => "mpid",
125    }
126}
127
128fn mobility_kind_str(k: MobilityArrayKind) -> &'static str {
129    match k {
130        MobilityArrayKind::InverseReducedVsPerCm2 => "inverse_reduced_k0",
131        MobilityArrayKind::DriftTimeMilliseconds => "drift_time_ms",
132    }
133}
134
135/// Streaming builder that accumulates `SpectrumRecord`s and produces a
136/// single Arrow [`RecordBatch`] when finalized.
137///
138/// All rows in a batch share one `mobility_array_kind` value, recorded
139/// once at construction. Push rows with [`push`](Self::push); call
140/// [`finish`](Self::finish) to materialize the batch.
141pub struct SpectrumBatchBuilder {
142    schema: SchemaRef,
143    mobility_kind: Option<MobilityArrayKind>,
144    index: UInt32Builder,
145    scan_number: UInt32Builder,
146    native_id: StringBuilder,
147    ms_level: UInt8Builder,
148    polarity: StringBuilder,
149    scan_mode: StringBuilder,
150    analyzer: StringBuilder,
151    filter: StringBuilder,
152    retention_time_sec: Float64Builder,
153    total_ion_current: Float64Builder,
154    base_peak_mz: Float64Builder,
155    base_peak_intensity: Float64Builder,
156    low_mz: Float64Builder,
157    high_mz: Float64Builder,
158    ion_injection_time_ms: Float64Builder,
159    inv_mobility: Float64Builder,
160    faims_cv: Float64Builder,
161    precursor_target_mz: Float64Builder,
162    precursor_selected_mz: Float64Builder,
163    precursor_isolation_width: Float64Builder,
164    precursor_charge: Int32Builder,
165    precursor_intensity: Float64Builder,
166    precursor_collision_energy: Float64Builder,
167    precursor_ce_is_nce: UInt8Builder,
168    precursor_native_id: StringBuilder,
169    precursor_activation: StringBuilder,
170    precursor_analyzer: StringBuilder,
171    precursor_ccs: Float64Builder,
172    mz: LargeListBuilder<Float64Builder>,
173    intensity: LargeListBuilder<Float32Builder>,
174    inv_mobility_per_peak: LargeListBuilder<Float32Builder>,
175    mobility_array_kind_col: StringBuilder,
176    acquisition_event_id: UInt32Builder,
177    extra: MapBuilder<StringBuilder, StringBuilder>,
178}
179
180impl SpectrumBatchBuilder {
181    /// Create a new builder. Pass the `mobility_array_kind` from the
182    /// source's [`crate::RunMetadata`] so the resulting Arrow batch
183    /// carries the unit/CV interpretation alongside the data.
184    pub fn new(mobility_kind: Option<MobilityArrayKind>) -> Self {
185        Self {
186            schema: spectrum_record_schema(),
187            mobility_kind,
188            index: UInt32Builder::new(),
189            scan_number: UInt32Builder::new(),
190            native_id: StringBuilder::new(),
191            ms_level: UInt8Builder::new(),
192            polarity: StringBuilder::new(),
193            scan_mode: StringBuilder::new(),
194            analyzer: StringBuilder::new(),
195            filter: StringBuilder::new(),
196            retention_time_sec: Float64Builder::new(),
197            total_ion_current: Float64Builder::new(),
198            base_peak_mz: Float64Builder::new(),
199            base_peak_intensity: Float64Builder::new(),
200            low_mz: Float64Builder::new(),
201            high_mz: Float64Builder::new(),
202            ion_injection_time_ms: Float64Builder::new(),
203            inv_mobility: Float64Builder::new(),
204            faims_cv: Float64Builder::new(),
205            precursor_target_mz: Float64Builder::new(),
206            precursor_selected_mz: Float64Builder::new(),
207            precursor_isolation_width: Float64Builder::new(),
208            precursor_charge: Int32Builder::new(),
209            precursor_intensity: Float64Builder::new(),
210            precursor_collision_energy: Float64Builder::new(),
211            precursor_ce_is_nce: UInt8Builder::new(),
212            precursor_native_id: StringBuilder::new(),
213            precursor_activation: StringBuilder::new(),
214            precursor_analyzer: StringBuilder::new(),
215            precursor_ccs: Float64Builder::new(),
216            mz: LargeListBuilder::new(Float64Builder::new()).with_field(Arc::new(Field::new(
217                "item",
218                DataType::Float64,
219                false,
220            ))),
221            intensity: LargeListBuilder::new(Float32Builder::new())
222                .with_field(Arc::new(Field::new("item", DataType::Float32, false))),
223            inv_mobility_per_peak: LargeListBuilder::new(Float32Builder::new())
224                .with_field(Arc::new(Field::new("item", DataType::Float32, false))),
225            mobility_array_kind_col: StringBuilder::new(),
226            acquisition_event_id: UInt32Builder::new(),
227            // Field names are pinned so the builder always matches the
228            // schema, whatever arrow's `MapBuilder` defaults are.
229            extra: MapBuilder::new(
230                Some(MapFieldNames {
231                    entry: EXTRA_ENTRIES_FIELD.to_string(),
232                    key: EXTRA_KEY_FIELD.to_string(),
233                    value: EXTRA_VALUE_FIELD.to_string(),
234                }),
235                StringBuilder::new(),
236                StringBuilder::new(),
237            ),
238        }
239    }
240
241    /// Schema for the batch produced by this builder.
242    pub fn schema(&self) -> SchemaRef {
243        self.schema.clone()
244    }
245
246    /// Append one spectrum row.
247    pub fn push(&mut self, rec: &SpectrumRecord) {
248        self.index.append_value(rec.index as u32);
249        self.scan_number.append_value(rec.scan_number);
250        self.native_id.append_value(&rec.native_id);
251        self.ms_level.append_value(rec.ms_level as u8);
252        self.polarity.append_option(rec.polarity.map(polarity_str));
253        self.scan_mode
254            .append_option(rec.scan_mode.map(scan_mode_str));
255        self.analyzer.append_option(rec.analyzer.map(analyzer_str));
256        self.filter.append_option(rec.filter.as_deref());
257        self.retention_time_sec.append_value(rec.retention_time_sec);
258        self.total_ion_current.append_option(rec.total_ion_current);
259        self.base_peak_mz.append_option(rec.base_peak_mz);
260        self.base_peak_intensity
261            .append_option(rec.base_peak_intensity);
262        self.low_mz.append_option(rec.low_mz);
263        self.high_mz.append_option(rec.high_mz);
264        self.ion_injection_time_ms
265            .append_option(rec.ion_injection_time_ms);
266        self.inv_mobility.append_option(rec.inv_mobility);
267        self.faims_cv.append_option(rec.faims_cv);
268
269        match &rec.precursor {
270            Some(p) => {
271                self.precursor_target_mz.append_option(p.target_mz);
272                self.precursor_selected_mz.append_option(p.selected_mz);
273                self.precursor_isolation_width
274                    .append_option(p.isolation_width);
275                self.precursor_charge.append_option(p.charge);
276                self.precursor_intensity.append_option(p.intensity);
277                self.precursor_collision_energy
278                    .append_option(p.collision_energy);
279                self.precursor_ce_is_nce.append_value(u8::from(p.ce_is_nce));
280                self.precursor_native_id
281                    .append_option(p.precursor_native_id.as_deref());
282                self.precursor_activation
283                    .append_option(p.activation.map(activation_str));
284                self.precursor_analyzer
285                    .append_option(p.analyzer.map(analyzer_str));
286                self.precursor_ccs.append_option(p.ccs);
287            }
288            None => {
289                self.precursor_target_mz.append_null();
290                self.precursor_selected_mz.append_null();
291                self.precursor_isolation_width.append_null();
292                self.precursor_charge.append_null();
293                self.precursor_intensity.append_null();
294                self.precursor_collision_energy.append_null();
295                self.precursor_ce_is_nce.append_null();
296                self.precursor_native_id.append_null();
297                self.precursor_activation.append_null();
298                self.precursor_analyzer.append_null();
299                self.precursor_ccs.append_null();
300            }
301        }
302
303        for &v in &rec.mz {
304            self.mz.values().append_value(v);
305        }
306        self.mz.append(true);
307        for &v in &rec.intensity {
308            self.intensity.values().append_value(v);
309        }
310        self.intensity.append(true);
311        match &rec.inv_mobility_per_peak {
312            Some(mob) => {
313                for &v in mob {
314                    self.inv_mobility_per_peak.values().append_value(v);
315                }
316                self.inv_mobility_per_peak.append(true);
317            }
318            None => self.inv_mobility_per_peak.append(false),
319        }
320        self.mobility_array_kind_col
321            .append_option(self.mobility_kind.map(mobility_kind_str));
322        self.acquisition_event_id
323            .append_option(rec.acquisition_event_id);
324        for (key, value) in &rec.extra {
325            self.extra.keys().append_value(key);
326            self.extra.values().append_value(value);
327        }
328        self.extra
329            .append(true)
330            .expect("extra keys and values are appended together");
331    }
332
333    /// Number of rows accumulated so far.
334    pub fn len(&self) -> usize {
335        self.index.len()
336    }
337
338    /// `true` if no rows have been pushed.
339    pub fn is_empty(&self) -> bool {
340        self.len() == 0
341    }
342
343    /// Materialize the accumulated rows into a [`RecordBatch`].
344    pub fn finish(mut self) -> Result<RecordBatch, arrow_schema::ArrowError> {
345        let arrays: Vec<ArrayRef> = vec![
346            Arc::new(self.index.finish()),
347            Arc::new(self.scan_number.finish()),
348            Arc::new(self.native_id.finish()),
349            Arc::new(self.ms_level.finish()),
350            Arc::new(self.polarity.finish()),
351            Arc::new(self.scan_mode.finish()),
352            Arc::new(self.analyzer.finish()),
353            Arc::new(self.filter.finish()),
354            Arc::new(self.retention_time_sec.finish()),
355            Arc::new(self.total_ion_current.finish()),
356            Arc::new(self.base_peak_mz.finish()),
357            Arc::new(self.base_peak_intensity.finish()),
358            Arc::new(self.low_mz.finish()),
359            Arc::new(self.high_mz.finish()),
360            Arc::new(self.ion_injection_time_ms.finish()),
361            Arc::new(self.inv_mobility.finish()),
362            Arc::new(self.faims_cv.finish()),
363            Arc::new(self.precursor_target_mz.finish()),
364            Arc::new(self.precursor_selected_mz.finish()),
365            Arc::new(self.precursor_isolation_width.finish()),
366            Arc::new(self.precursor_charge.finish()),
367            Arc::new(self.precursor_intensity.finish()),
368            Arc::new(self.precursor_collision_energy.finish()),
369            Arc::new(self.precursor_ce_is_nce.finish()),
370            Arc::new(self.precursor_native_id.finish()),
371            Arc::new(self.precursor_activation.finish()),
372            Arc::new(self.precursor_analyzer.finish()),
373            Arc::new(self.precursor_ccs.finish()),
374            Arc::new(self.mz.finish()),
375            Arc::new(self.intensity.finish()),
376            Arc::new(self.inv_mobility_per_peak.finish()),
377            Arc::new(self.mobility_array_kind_col.finish()),
378            Arc::new(self.acquisition_event_id.finish()),
379            Arc::new(self.extra.finish()),
380        ];
381        RecordBatch::try_new(self.schema.clone(), arrays)
382    }
383}
384
385#[cfg(test)]
386mod tests {
387    use super::*;
388    use crate::{PrecursorInfo, SpectrumRecord};
389    use arrow_array::Array;
390
391    fn rec(index: usize, ms_level: u32, n_peaks: usize, with_mob: bool) -> SpectrumRecord {
392        let mz: Vec<f64> = (0..n_peaks).map(|i| 100.0 + i as f64).collect();
393        let intensity: Vec<f32> = (0..n_peaks).map(|i| 10.0 + i as f32).collect();
394        let mobility = if with_mob {
395            Some((0..n_peaks).map(|i| 0.5 + i as f32 * 0.01).collect())
396        } else {
397            None
398        };
399        SpectrumRecord {
400            index,
401            scan_number: (index + 1) as u32,
402            native_id: format!("scan={}", index + 1),
403            ms_level,
404            polarity: Some(Polarity::Positive),
405            scan_mode: Some(ScanMode::Centroid),
406            analyzer: Some(Analyzer::TOFMS),
407            acquisition_event_id: (index == 0).then_some(7),
408            filter: None,
409            retention_time_sec: index as f64,
410            total_ion_current: Some(intensity.iter().map(|&v| v as f64).sum()),
411            base_peak_mz: mz.last().copied(),
412            base_peak_intensity: intensity.last().map(|&v| v as f64),
413            low_mz: mz.first().copied(),
414            high_mz: mz.last().copied(),
415            ion_injection_time_ms: None,
416            inv_mobility: None,
417            faims_cv: if index == 0 { Some(-40.0) } else { None },
418            precursor: if ms_level >= 2 {
419                Some(PrecursorInfo {
420                    target_mz: Some(500.0),
421                    selected_mz: Some(500.5),
422                    isolation_width: Some(2.0),
423                    charge: Some(2),
424                    ..Default::default()
425                })
426            } else {
427                None
428            },
429            mz,
430            intensity,
431            extra: if index == 0 {
432                [("opentfraw.resolution".to_string(), "120000".to_string())]
433                    .into_iter()
434                    .collect()
435            } else {
436                Default::default()
437            },
438            inv_mobility_per_peak: mobility,
439        }
440    }
441
442    #[test]
443    fn schema_round_trip() {
444        let mut b = SpectrumBatchBuilder::new(Some(MobilityArrayKind::DriftTimeMilliseconds));
445        b.push(&rec(0, 1, 3, true));
446        b.push(&rec(1, 2, 4, false));
447        let batch = b.finish().unwrap();
448        assert_eq!(batch.num_rows(), 2);
449        assert_eq!(batch.schema().fields().len(), 34);
450        let mz_col = batch
451            .column_by_name("mz")
452            .unwrap()
453            .as_any()
454            .downcast_ref::<arrow_array::LargeListArray>()
455            .unwrap();
456        assert_eq!(mz_col.value_length(0), 3);
457        assert_eq!(mz_col.value_length(1), 4);
458        let mob_col = batch
459            .column_by_name("inv_mobility_per_peak")
460            .unwrap()
461            .as_any()
462            .downcast_ref::<arrow_array::LargeListArray>()
463            .unwrap();
464        assert!(mob_col.is_valid(0));
465        assert!(mob_col.is_null(1));
466        let faims_col = batch
467            .column_by_name("faims_cv")
468            .unwrap()
469            .as_any()
470            .downcast_ref::<arrow_array::Float64Array>()
471            .unwrap();
472        assert_eq!(faims_col.value(0), -40.0);
473        assert!(faims_col.is_null(1));
474        let event_col = batch
475            .column_by_name("acquisition_event_id")
476            .unwrap()
477            .as_any()
478            .downcast_ref::<arrow_array::UInt32Array>()
479            .unwrap();
480        assert_eq!(event_col.value(0), 7);
481        assert!(event_col.is_null(1));
482        let extra_col = batch
483            .column_by_name("extra")
484            .unwrap()
485            .as_any()
486            .downcast_ref::<arrow_array::MapArray>()
487            .unwrap();
488        assert_eq!(extra_col.value_length(0), 1);
489        assert_eq!(extra_col.value_length(1), 0);
490        let keys = extra_col
491            .keys()
492            .as_any()
493            .downcast_ref::<arrow_array::StringArray>()
494            .unwrap();
495        let values = extra_col
496            .values()
497            .as_any()
498            .downcast_ref::<arrow_array::StringArray>()
499            .unwrap();
500        assert_eq!(keys.value(0), "opentfraw.resolution");
501        assert_eq!(values.value(0), "120000");
502    }
503}
openmassspec_core/arrow.rs

openmassspec_core/
arrow.rs
