openmassspec_arrow/
lib.rs

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

openmassspec_arrow/
lib.rs
