nautilus_serialization/arrow/
mod.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Defines the Apache Arrow schema for Nautilus types.
17
18pub mod account_state;
19pub mod bar;
20pub mod close;
21pub mod custom;
22pub mod delta;
23pub mod depth;
24pub mod funding;
25pub mod index_price;
26pub mod instrument;
27pub mod instrument_status;
28pub mod json;
29pub mod mark_price;
30pub mod option_greeks;
31pub mod order_event;
32pub mod position_event;
33pub mod quote;
34pub mod report;
35pub mod snapshot;
36pub mod trade;
37
38#[cfg(feature = "display")]
39pub mod display;
40
41use std::{
42    collections::HashMap,
43    io::{self, Write},
44};
45
46use arrow::{
47    array::{
48        Array, ArrayRef, BinaryArray, BinaryViewArray, FixedSizeBinaryArray, StringArray,
49        StringViewArray,
50    },
51    datatypes::{DataType, Schema},
52    error::ArrowError,
53    ipc::writer::StreamWriter,
54    record_batch::RecordBatch,
55};
56use nautilus_model::{
57    data::{
58        Data, IndexPriceUpdate, InstrumentStatus, MarkPriceUpdate, bar::Bar,
59        close::InstrumentClose, delta::OrderBookDelta, depth::OrderBookDepth10,
60        option_chain::OptionGreeks, quote::QuoteTick, trade::TradeTick,
61    },
62    enums::BookAction,
63    types::{
64        PRICE_ERROR, PRICE_UNDEF, Price, QUANTITY_UNDEF, Quantity,
65        fixed::{PRECISION_BYTES, correct_price_raw, correct_quantity_raw},
66        price::PriceRaw,
67        quantity::QuantityRaw,
68    },
69};
70#[cfg(feature = "python")]
71use pyo3::prelude::*;
72use ustr::Ustr;
73
74// Define metadata key constants constants
75const KEY_BAR_TYPE: &str = "bar_type";
76pub const KEY_INSTRUMENT_ID: &str = "instrument_id";
77pub const KEY_PRICE_PRECISION: &str = "price_precision";
78pub const KEY_SIZE_PRECISION: &str = "size_precision";
79
80#[derive(thiserror::Error, Debug)]
81pub enum DataStreamingError {
82    #[error("I/O error: {0}")]
83    IoError(#[from] io::Error),
84    #[error("Arrow error: {0}")]
85    ArrowError(#[from] arrow::error::ArrowError),
86    #[cfg(feature = "python")]
87    #[error("Python error: {0}")]
88    PythonError(#[from] PyErr),
89}
90
91#[derive(thiserror::Error, Debug)]
92pub enum EncodingError {
93    #[error("Empty data")]
94    EmptyData,
95    #[error(
96        "Mixed metadata at row {index}; encode each instrument, bar type, or precision separately"
97    )]
98    MixedMetadata { index: usize },
99    #[error("Missing metadata key: `{0}`")]
100    MissingMetadata(&'static str),
101    #[error("Missing data column: `{0}` at index {1}")]
102    MissingColumn(&'static str, usize),
103    #[error("Error parsing `{0}`: {1}")]
104    ParseError(&'static str, String),
105    #[error("Invalid column type `{0}` at index {1}: expected {2}, found {3}")]
106    InvalidColumnType(&'static str, usize, DataType, DataType),
107    #[error(
108        "Precision mode mismatch for `{field}`: catalog data has {actual_bytes} byte values, \
109         but this build expects {expected_bytes} bytes. The catalog was created with a different \
110         precision mode (standard=8 bytes, high=16 bytes). Rebuild the catalog or change your \
111         build's precision mode. See: https://nautilustrader.io/docs/latest/getting_started/installation#precision-mode"
112    )]
113    PrecisionMismatch {
114        field: &'static str,
115        expected_bytes: i32,
116        actual_bytes: i32,
117    },
118    #[error("Arrow error: {0}")]
119    ArrowError(#[from] arrow::error::ArrowError),
120}
121
122#[inline]
123fn get_raw_price(bytes: &[u8]) -> PriceRaw {
124    PriceRaw::from_le_bytes(
125        bytes
126            .try_into()
127            .expect("Price raw bytes must be exactly the size of PriceRaw"),
128    )
129}
130
131#[inline]
132fn get_raw_quantity(bytes: &[u8]) -> QuantityRaw {
133    QuantityRaw::from_le_bytes(
134        bytes
135            .try_into()
136            .expect("Quantity raw bytes must be exactly the size of QuantityRaw"),
137    )
138}
139
140/// Gets raw price bytes and corrects for floating-point precision errors in stored data.
141///
142/// Data from catalogs may have been created with `int(value * FIXED_SCALAR)` which
143/// introduces floating-point errors. This corrects the raw value to the nearest valid
144/// multiple of the scale factor for the given precision.
145///
146/// Sentinel values (`PRICE_UNDEF`, `PRICE_ERROR`) are preserved unchanged.
147#[inline]
148fn get_corrected_raw_price(bytes: &[u8], precision: u8) -> PriceRaw {
149    let raw = get_raw_price(bytes);
150
151    // Preserve sentinel values unchanged
152    if raw == PRICE_UNDEF || raw == PRICE_ERROR {
153        return raw;
154    }
155
156    correct_price_raw(raw, precision)
157}
158
159/// Gets raw quantity bytes and corrects for floating-point precision errors in stored data.
160///
161/// Data from catalogs may have been created with `int(value * FIXED_SCALAR)` which
162/// introduces floating-point errors. This corrects the raw value to the nearest valid
163/// multiple of the scale factor for the given precision.
164///
165/// Sentinel values (`QUANTITY_UNDEF`) are preserved unchanged.
166#[inline]
167fn get_corrected_raw_quantity(bytes: &[u8], precision: u8) -> QuantityRaw {
168    let raw = get_raw_quantity(bytes);
169
170    // Preserve sentinel values unchanged
171    if raw == QUANTITY_UNDEF {
172        return raw;
173    }
174
175    correct_quantity_raw(raw, precision)
176}
177
178/// Decodes a [`Price`] from raw bytes with bounds validation.
179///
180/// Uses corrected raw values to handle floating-point precision errors in stored data.
181/// Sentinel values (`PRICE_UNDEF`, `PRICE_ERROR`) are preserved unchanged.
182///
183/// # Errors
184///
185/// Returns an [`EncodingError::ParseError`] if the price value is out of bounds.
186pub fn decode_price(
187    bytes: &[u8],
188    precision: u8,
189    field: &'static str,
190    row: usize,
191) -> Result<Price, EncodingError> {
192    let raw = get_corrected_raw_price(bytes, precision);
193    Price::from_raw_checked(raw, precision)
194        .map_err(|e| EncodingError::ParseError(field, format!("row {row}: {e}")))
195}
196
197/// Decodes a [`Quantity`] from raw bytes with bounds validation.
198///
199/// Uses corrected raw values to handle floating-point precision errors in stored data.
200/// Sentinel values (`QUANTITY_UNDEF`) are preserved unchanged.
201///
202/// # Errors
203///
204/// Returns an [`EncodingError::ParseError`] if the quantity value is out of bounds.
205pub fn decode_quantity(
206    bytes: &[u8],
207    precision: u8,
208    field: &'static str,
209    row: usize,
210) -> Result<Quantity, EncodingError> {
211    let raw = get_corrected_raw_quantity(bytes, precision);
212    Quantity::from_raw_checked(raw, precision)
213        .map_err(|e| EncodingError::ParseError(field, format!("row {row}: {e}")))
214}
215
216/// Decodes a [`Price`] from raw bytes, using precision 0 for sentinel values.
217///
218/// For order book data where sentinel values indicate empty levels.
219///
220/// # Errors
221///
222/// Returns an [`EncodingError::ParseError`] if the price value is out of bounds.
223pub fn decode_price_with_sentinel(
224    bytes: &[u8],
225    precision: u8,
226    field: &'static str,
227    row: usize,
228) -> Result<Price, EncodingError> {
229    let raw = get_raw_price(bytes);
230    let (final_raw, final_precision) = if raw == PRICE_UNDEF {
231        (raw, 0)
232    } else {
233        (get_corrected_raw_price(bytes, precision), precision)
234    };
235    Price::from_raw_checked(final_raw, final_precision)
236        .map_err(|e| EncodingError::ParseError(field, format!("row {row}: {e}")))
237}
238
239/// Decodes a [`Quantity`] from raw bytes, using precision 0 for sentinel values.
240///
241/// For order book data where sentinel values indicate empty levels.
242///
243/// # Errors
244///
245/// Returns an [`EncodingError::ParseError`] if the quantity value is out of bounds.
246pub fn decode_quantity_with_sentinel(
247    bytes: &[u8],
248    precision: u8,
249    field: &'static str,
250    row: usize,
251) -> Result<Quantity, EncodingError> {
252    let raw = get_raw_quantity(bytes);
253    let (final_raw, final_precision) = if raw == QUANTITY_UNDEF {
254        (raw, 0)
255    } else {
256        (get_corrected_raw_quantity(bytes, precision), precision)
257    };
258    Quantity::from_raw_checked(final_raw, final_precision)
259        .map_err(|e| EncodingError::ParseError(field, format!("row {row}: {e}")))
260}
261
262/// Provides Apache Arrow schema definitions for data types.
263pub trait ArrowSchemaProvider {
264    /// Returns the Arrow schema for this type with optional metadata.
265    fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema;
266
267    /// Returns a map of field names to their Arrow data types.
268    #[must_use]
269    fn get_schema_map() -> HashMap<String, String> {
270        let schema = Self::get_schema(None);
271        let mut map = HashMap::new();
272
273        for field in schema.fields() {
274            let name = field.name().clone();
275            let data_type = format!("{:?}", field.data_type());
276            map.insert(name, data_type);
277        }
278        map
279    }
280}
281
282/// Encodes data types to Apache Arrow RecordBatch format.
283pub trait EncodeToRecordBatch
284where
285    Self: Sized + ArrowSchemaProvider,
286{
287    /// Encodes a batch of values into an Arrow `RecordBatch` using the provided metadata.
288    ///
289    /// # Errors
290    ///
291    /// Returns an `ArrowError` if the encoding fails.
292    fn encode_batch(
293        metadata: &HashMap<String, String>,
294        data: &[Self],
295    ) -> Result<RecordBatch, ArrowError>;
296
297    /// Returns the metadata for this data element.
298    fn metadata(&self) -> HashMap<String, String>;
299
300    /// Returns the metadata selected for a chunk.
301    ///
302    /// The default uses the first element. Implementations may override this when leading sentinel
303    /// values do not carry meaningful metadata.
304    ///
305    /// # Panics
306    ///
307    /// Panics if `chunk` is empty.
308    fn chunk_metadata(chunk: &[Self]) -> HashMap<String, String> {
309        chunk
310            .first()
311            .map(Self::metadata)
312            .expect("Chunk must have at least one element to encode")
313    }
314}
315
316/// Decodes data types from Apache Arrow RecordBatch format.
317pub trait DecodeFromRecordBatch
318where
319    Self: Sized + Into<Data> + ArrowSchemaProvider,
320{
321    /// Decodes a `RecordBatch` into a vector of values of the implementing type, using the provided metadata.
322    ///
323    /// # Errors
324    ///
325    /// Returns an `EncodingError` if the decoding fails.
326    fn decode_batch(
327        metadata: &HashMap<String, String>,
328        record_batch: RecordBatch,
329    ) -> Result<Vec<Self>, EncodingError>;
330}
331
332/// Decodes strongly typed values from Apache Arrow RecordBatch format.
333pub trait DecodeTypedFromRecordBatch
334where
335    Self: Sized + ArrowSchemaProvider,
336{
337    /// Decodes a `RecordBatch` into a vector of values of the implementing type.
338    ///
339    /// # Errors
340    ///
341    /// Returns an `EncodingError` if the decoding fails.
342    fn decode_typed_batch(
343        metadata: &HashMap<String, String>,
344        record_batch: RecordBatch,
345    ) -> Result<Vec<Self>, EncodingError>;
346}
347
348impl<T> DecodeTypedFromRecordBatch for T
349where
350    T: DecodeFromRecordBatch,
351{
352    fn decode_typed_batch(
353        metadata: &HashMap<String, String>,
354        record_batch: RecordBatch,
355    ) -> Result<Vec<Self>, EncodingError> {
356        Self::decode_batch(metadata, record_batch)
357    }
358}
359
360/// Decodes raw Data objects from Apache Arrow RecordBatch format.
361pub trait DecodeDataFromRecordBatch
362where
363    Self: Sized + ArrowSchemaProvider,
364{
365    /// Decodes a `RecordBatch` into raw `Data` values, using the provided metadata.
366    ///
367    /// # Errors
368    ///
369    /// Returns an `EncodingError` if the decoding fails.
370    fn decode_data_batch(
371        metadata: &HashMap<String, String>,
372        record_batch: RecordBatch,
373    ) -> Result<Vec<Data>, EncodingError>;
374}
375
376/// Writes RecordBatch data to output streams.
377pub trait WriteStream {
378    /// Writes a `RecordBatch` to the implementing output stream.
379    ///
380    /// # Errors
381    ///
382    /// Returns a `DataStreamingError` if writing or finishing the stream fails.
383    fn write(&mut self, record_batch: &RecordBatch) -> Result<(), DataStreamingError>;
384}
385
386impl<T: Write> WriteStream for T {
387    fn write(&mut self, record_batch: &RecordBatch) -> Result<(), DataStreamingError> {
388        let mut writer = StreamWriter::try_new(self, &record_batch.schema())?;
389        writer.write(record_batch)?;
390        writer.finish()?;
391        Ok(())
392    }
393}
394
395/// Extracts a string column, accepting both Utf8 (`StringArray`) and Utf8View (`StringViewArray`).
396/// Parquet may return Utf8View when reading, so this handles both formats.
397///
398/// # Errors
399///
400/// Returns an error if:
401/// - `column_index` is out of range: `EncodingError::MissingColumn`.
402/// - The column type is neither Utf8 nor Utf8View: `EncodingError::InvalidColumnType`.
403pub fn extract_column_string<'a>(
404    cols: &'a [ArrayRef],
405    column_key: &'static str,
406    column_index: usize,
407) -> Result<StringColumnRef<'a>, EncodingError> {
408    let column_values = cols
409        .get(column_index)
410        .ok_or(EncodingError::MissingColumn(column_key, column_index))?;
411    let dt = column_values.data_type();
412    if let Some(arr) = column_values.as_any().downcast_ref::<StringArray>() {
413        Ok(StringColumnRef::Utf8(arr))
414    } else if let Some(arr) = column_values.as_any().downcast_ref::<StringViewArray>() {
415        Ok(StringColumnRef::Utf8View(arr))
416    } else {
417        Err(EncodingError::InvalidColumnType(
418            column_key,
419            column_index,
420            DataType::Utf8,
421            dt.clone(),
422        ))
423    }
424}
425
426/// Reference to a string column, either Utf8 or Utf8View.
427#[derive(Debug)]
428pub enum StringColumnRef<'a> {
429    Utf8(&'a StringArray),
430    Utf8View(&'a StringViewArray),
431}
432
433impl StringColumnRef<'_> {
434    /// Returns the string value at row `i`.
435    #[inline]
436    #[must_use]
437    pub fn value(&self, i: usize) -> &str {
438        match self {
439            Self::Utf8(arr) => arr.value(i),
440            Self::Utf8View(arr) => arr.value(i),
441        }
442    }
443}
444
445/// Extracts a binary column, accepting both Binary (`BinaryArray`) and BinaryView
446/// (`BinaryViewArray`).
447/// DataFusion may return BinaryView when reading Parquet, so this handles both formats.
448///
449/// # Errors
450///
451/// Returns an error if:
452/// - `column_index` is out of range: `EncodingError::MissingColumn`.
453/// - The column type is neither Binary nor BinaryView: `EncodingError::InvalidColumnType`.
454pub fn extract_column_binary<'a>(
455    cols: &'a [ArrayRef],
456    column_key: &'static str,
457    column_index: usize,
458) -> Result<BinaryColumnRef<'a>, EncodingError> {
459    let column_values = cols
460        .get(column_index)
461        .ok_or(EncodingError::MissingColumn(column_key, column_index))?;
462    let dt = column_values.data_type();
463    if let Some(arr) = column_values.as_any().downcast_ref::<BinaryArray>() {
464        Ok(BinaryColumnRef::Binary(arr))
465    } else if let Some(arr) = column_values.as_any().downcast_ref::<BinaryViewArray>() {
466        Ok(BinaryColumnRef::BinaryView(arr))
467    } else {
468        Err(EncodingError::InvalidColumnType(
469            column_key,
470            column_index,
471            DataType::Binary,
472            dt.clone(),
473        ))
474    }
475}
476
477/// Reference to a binary column, either Binary or BinaryView.
478#[derive(Debug)]
479pub enum BinaryColumnRef<'a> {
480    Binary(&'a BinaryArray),
481    BinaryView(&'a BinaryViewArray),
482}
483
484impl BinaryColumnRef<'_> {
485    /// Returns the bytes at row `i`.
486    #[inline]
487    #[must_use]
488    pub fn value(&self, i: usize) -> &[u8] {
489        match self {
490            Self::Binary(arr) => arr.value(i),
491            Self::BinaryView(arr) => arr.value(i),
492        }
493    }
494}
495
496/// Extracts and downcasts the specified `column_key` column from an Arrow array slice.
497///
498/// # Errors
499///
500/// Returns an error if:
501/// - `column_index` is out of range: `EncodingError::MissingColumn`.
502/// - The column type does not match `expected_type`: `EncodingError::InvalidColumnType`.
503pub fn extract_column<'a, T: Array + 'static>(
504    cols: &'a [ArrayRef],
505    column_key: &'static str,
506    column_index: usize,
507    expected_type: DataType,
508) -> Result<&'a T, EncodingError> {
509    let column_values = cols
510        .get(column_index)
511        .ok_or(EncodingError::MissingColumn(column_key, column_index))?;
512    let downcasted_values =
513        column_values
514            .as_any()
515            .downcast_ref::<T>()
516            .ok_or(EncodingError::InvalidColumnType(
517                column_key,
518                column_index,
519                expected_type,
520                column_values.data_type().clone(),
521            ))?;
522    Ok(downcasted_values)
523}
524
525/// Extracts a column by name when present, falling back to an index for older schemas.
526///
527/// # Errors
528///
529/// Returns an error if the resolved column is missing or has the wrong type.
530pub fn extract_column_by_name_or_index<'a, T: Array + 'static>(
531    record_batch: &'a RecordBatch,
532    column_key: &'static str,
533    fallback_index: usize,
534    expected_type: DataType,
535) -> Result<&'a T, EncodingError> {
536    let column_index = record_batch
537        .schema()
538        .index_of(column_key)
539        .unwrap_or(fallback_index);
540    extract_column::<T>(
541        record_batch.columns(),
542        column_key,
543        column_index,
544        expected_type,
545    )
546}
547
548/// Extracts an optional UTF-8 column by name.
549///
550/// # Errors
551///
552/// Returns an error if the column exists but is not UTF-8.
553pub fn extract_optional_string_column_by_name<'a>(
554    record_batch: &'a RecordBatch,
555    column_key: &'static str,
556) -> Result<Option<&'a StringArray>, EncodingError> {
557    let Ok(column_index) = record_batch.schema().index_of(column_key) else {
558        return Ok(None);
559    };
560    let column_values = record_batch
561        .columns()
562        .get(column_index)
563        .ok_or(EncodingError::MissingColumn(column_key, column_index))?;
564    let downcasted_values = column_values.as_any().downcast_ref::<StringArray>().ok_or(
565        EncodingError::InvalidColumnType(
566            column_key,
567            column_index,
568            DataType::Utf8,
569            column_values.data_type().clone(),
570        ),
571    )?;
572    Ok(Some(downcasted_values))
573}
574
575/// Returns an optional [`Ustr`] value from an optional string column.
576#[must_use]
577pub fn optional_ustr_value(values: Option<&StringArray>, row: usize) -> Option<Ustr> {
578    values.and_then(|column| (!column.is_null(row)).then(|| Ustr::from(column.value(row))))
579}
580
581/// Validates that a [`FixedSizeBinaryArray`] has the expected precision byte width.
582///
583/// This detects precision mode mismatches that occur when catalog data was encoded
584/// with a different precision mode (64-bit standard vs 128-bit high-precision).
585///
586/// # Errors
587///
588/// Returns [`EncodingError::PrecisionMismatch`] if the actual byte width doesn't
589/// match [`PRECISION_BYTES`].
590pub fn validate_precision_bytes(
591    array: &FixedSizeBinaryArray,
592    field: &'static str,
593) -> Result<(), EncodingError> {
594    let actual = array.value_length();
595    if actual != PRECISION_BYTES {
596        return Err(EncodingError::PrecisionMismatch {
597            field,
598            expected_bytes: PRECISION_BYTES,
599            actual_bytes: actual,
600        });
601    }
602    Ok(())
603}
604
605/// Converts a vector of `OrderBookDelta` into an Arrow `RecordBatch`.
606///
607/// # Errors
608///
609/// Returns an error if:
610/// - `data` is empty: `EncodingError::EmptyData`.
611/// - Instrument IDs differ, or non-clear precision metadata differs:
612///   `EncodingError::MixedMetadata`.
613/// - Encoding fails: `EncodingError::ArrowError`.
614pub fn book_deltas_to_arrow_record_batch_bytes(
615    data: &[OrderBookDelta],
616) -> Result<RecordBatch, EncodingError> {
617    let Some(first) = data.first() else {
618        return Err(EncodingError::EmptyData);
619    };
620
621    let metadata = OrderBookDelta::chunk_metadata(data);
622    let instrument_id = data
623        .iter()
624        .find(|delta| delta.action != BookAction::Clear)
625        .unwrap_or(first)
626        .instrument_id;
627
628    if let Some(index) = data.iter().position(|delta| {
629        delta.instrument_id != instrument_id
630            || (delta.action != BookAction::Clear && delta.metadata() != metadata)
631    }) {
632        return Err(EncodingError::MixedMetadata { index });
633    }
634
635    OrderBookDelta::encode_batch(&metadata, data).map_err(EncodingError::ArrowError)
636}
637
638/// Converts a vector of `OrderBookDepth10` into an Arrow `RecordBatch`.
639///
640/// # Errors
641///
642/// Returns an error if:
643/// - `data` is empty: `EncodingError::EmptyData`.
644/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
645/// - Encoding fails: `EncodingError::ArrowError`.
646pub fn book_depth10_to_arrow_record_batch_bytes(
647    data: &[OrderBookDepth10],
648) -> Result<RecordBatch, EncodingError> {
649    let Some(first) = data.first() else {
650        return Err(EncodingError::EmptyData);
651    };
652    let precision = data
653        .iter()
654        .flat_map(|depth| depth.bids.iter().chain(&depth.asks))
655        .find(|order| !order.price.is_undefined() && !order.size.is_undefined())
656        .map_or(
657            (first.bids[0].price.precision, first.bids[0].size.precision),
658            |order| (order.price.precision, order.size.precision),
659        );
660
661    if let Some(index) = data.iter().position(|depth| {
662        depth.instrument_id != first.instrument_id || !depth_precision_is_uniform(depth, precision)
663    }) {
664        return Err(EncodingError::MixedMetadata { index });
665    }
666
667    let metadata = OrderBookDepth10::get_metadata(&first.instrument_id, precision.0, precision.1);
668    OrderBookDepth10::encode_batch(&metadata, data).map_err(EncodingError::ArrowError)
669}
670
671fn depth_precision_is_uniform(depth: &OrderBookDepth10, precision: (u8, u8)) -> bool {
672    depth.bids.iter().chain(&depth.asks).all(|order| {
673        match (order.price.is_undefined(), order.size.is_undefined()) {
674            (true, true) => true,
675            (false, false) => {
676                order.price.precision == precision.0 && order.size.precision == precision.1
677            }
678            _ => false,
679        }
680    })
681}
682
683/// Converts a vector of `QuoteTick` into an Arrow `RecordBatch`.
684///
685/// # Errors
686///
687/// Returns an error if:
688/// - `data` is empty: `EncodingError::EmptyData`.
689/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
690/// - Encoding fails: `EncodingError::ArrowError`.
691pub fn quotes_to_arrow_record_batch_bytes(
692    data: &[QuoteTick],
693) -> Result<RecordBatch, EncodingError> {
694    encode_batch_with_metadata(data)
695}
696
697/// Converts a vector of `TradeTick` into an Arrow `RecordBatch`.
698///
699/// # Errors
700///
701/// Returns an error if:
702/// - `data` is empty: `EncodingError::EmptyData`.
703/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
704/// - Encoding fails: `EncodingError::ArrowError`.
705pub fn trades_to_arrow_record_batch_bytes(
706    data: &[TradeTick],
707) -> Result<RecordBatch, EncodingError> {
708    encode_batch_with_metadata(data)
709}
710
711/// Converts a vector of `Bar` into an Arrow `RecordBatch`.
712///
713/// # Errors
714///
715/// Returns an error if:
716/// - `data` is empty: `EncodingError::EmptyData`.
717/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
718/// - Encoding fails: `EncodingError::ArrowError`.
719pub fn bars_to_arrow_record_batch_bytes(data: &[Bar]) -> Result<RecordBatch, EncodingError> {
720    encode_batch_with_metadata(data)
721}
722
723/// Converts a vector of `MarkPriceUpdate` into an Arrow `RecordBatch`.
724///
725/// # Errors
726///
727/// Returns an error if:
728/// - `data` is empty: `EncodingError::EmptyData`.
729/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
730/// - Encoding fails: `EncodingError::ArrowError`.
731pub fn mark_prices_to_arrow_record_batch_bytes(
732    data: &[MarkPriceUpdate],
733) -> Result<RecordBatch, EncodingError> {
734    encode_batch_with_metadata(data)
735}
736
737/// Converts a vector of `IndexPriceUpdate` into an Arrow `RecordBatch`.
738///
739/// # Errors
740///
741/// Returns an error if:
742/// - `data` is empty: `EncodingError::EmptyData`.
743/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
744/// - Encoding fails: `EncodingError::ArrowError`.
745pub fn index_prices_to_arrow_record_batch_bytes(
746    data: &[IndexPriceUpdate],
747) -> Result<RecordBatch, EncodingError> {
748    encode_batch_with_metadata(data)
749}
750
751/// Converts a vector of `InstrumentStatus` into an Arrow `RecordBatch`.
752///
753/// # Errors
754///
755/// Returns an error if:
756/// - `data` is empty: `EncodingError::EmptyData`.
757/// - Encoding fails: `EncodingError::ArrowError`.
758#[expect(clippy::missing_panics_doc)] // Guarded by empty check
759pub fn instrument_status_to_arrow_record_batch_bytes(
760    data: &[InstrumentStatus],
761) -> Result<RecordBatch, EncodingError> {
762    if data.is_empty() {
763        return Err(EncodingError::EmptyData);
764    }
765
766    let first = data.first().unwrap();
767    let metadata = first.metadata();
768    InstrumentStatus::encode_batch(&metadata, data).map_err(EncodingError::ArrowError)
769}
770
771/// Converts a vector of `OptionGreeks` into an Arrow `RecordBatch`.
772///
773/// # Errors
774///
775/// Returns an error if:
776/// - `data` is empty: `EncodingError::EmptyData`.
777/// - Encoding fails: `EncodingError::ArrowError`.
778#[expect(clippy::missing_panics_doc)] // Guarded by empty check
779pub fn option_greeks_to_arrow_record_batch_bytes(
780    data: &[OptionGreeks],
781) -> Result<RecordBatch, EncodingError> {
782    if data.is_empty() {
783        return Err(EncodingError::EmptyData);
784    }
785
786    let first = data.first().unwrap();
787    let metadata = first.metadata();
788    OptionGreeks::encode_batch(&metadata, data).map_err(EncodingError::ArrowError)
789}
790
791/// Converts a vector of `InstrumentClose` into an Arrow `RecordBatch`.
792///
793/// # Errors
794///
795/// Returns an error if:
796/// - `data` is empty: `EncodingError::EmptyData`.
797/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
798/// - Encoding fails: `EncodingError::ArrowError`.
799pub fn instrument_closes_to_arrow_record_batch_bytes(
800    data: &[InstrumentClose],
801) -> Result<RecordBatch, EncodingError> {
802    encode_batch_with_metadata(data)
803}
804
805fn encode_batch_with_metadata<T>(data: &[T]) -> Result<RecordBatch, EncodingError>
806where
807    T: EncodeToRecordBatch,
808{
809    if data.is_empty() {
810        return Err(EncodingError::EmptyData);
811    }
812
813    let metadata = T::chunk_metadata(data);
814    if let Some(index) = data.iter().position(|value| value.metadata() != metadata) {
815        return Err(EncodingError::MixedMetadata { index });
816    }
817
818    T::encode_batch(&metadata, data).map_err(EncodingError::ArrowError)
819}
820
821#[cfg(test)]
822mod tests {
823    use nautilus_model::{
824        data::{
825            Bar, BarSpecification, BarType, BookOrder, OrderBookDelta, OrderBookDepth10, QuoteTick,
826            depth::DEPTH10_LEN,
827        },
828        enums::{AggregationSource, BarAggregation, BookAction, OrderSide, PriceType},
829        identifiers::InstrumentId,
830        types::{PRICE_UNDEF, Price, QUANTITY_UNDEF, Quantity},
831    };
832    use rstest::rstest;
833
834    use super::*;
835
836    #[rstest]
837    fn test_quotes_to_arrow_record_batch_rejects_mixed_instruments() {
838        let first = QuoteTick::new(
839            InstrumentId::from("AAPL.XNAS"),
840            Price::from("100.01"),
841            Price::from("100.02"),
842            Quantity::from("10"),
843            Quantity::from("11"),
844            1.into(),
845            1.into(),
846        );
847        let second = QuoteTick::new(
848            InstrumentId::from("MSFT.XNAS"),
849            Price::from("200.01"),
850            Price::from("200.02"),
851            Quantity::from("20"),
852            Quantity::from("21"),
853            2.into(),
854            2.into(),
855        );
856
857        let result = quotes_to_arrow_record_batch_bytes(&[first, second]);
858
859        assert!(matches!(
860            result,
861            Err(EncodingError::MixedMetadata { index: 1 })
862        ));
863    }
864
865    #[rstest]
866    fn test_quotes_to_arrow_record_batch_rejects_mixed_precision() {
867        let instrument_id = InstrumentId::from("AAPL.XNAS");
868        let first = QuoteTick::new(
869            instrument_id,
870            Price::from("100.01"),
871            Price::from("100.02"),
872            Quantity::from("10.00"),
873            Quantity::from("11.00"),
874            1.into(),
875            1.into(),
876        );
877        let second = QuoteTick::new(
878            instrument_id,
879            Price::from("100.010"),
880            Price::from("100.020"),
881            Quantity::from("10.000"),
882            Quantity::from("11.000"),
883            2.into(),
884            2.into(),
885        );
886
887        let result = quotes_to_arrow_record_batch_bytes(&[first, second]);
888
889        assert!(matches!(
890            result,
891            Err(EncodingError::MixedMetadata { index: 1 })
892        ));
893    }
894
895    #[rstest]
896    fn test_bars_to_arrow_record_batch_rejects_mixed_bar_types() {
897        let instrument_id = InstrumentId::from("AAPL.XNAS");
898        let first_type = BarType::new(
899            instrument_id,
900            BarSpecification::new(1, BarAggregation::Minute, PriceType::Last),
901            AggregationSource::Internal,
902        );
903        let second_type = BarType::new(
904            instrument_id,
905            BarSpecification::new(5, BarAggregation::Minute, PriceType::Last),
906            AggregationSource::Internal,
907        );
908        let first = Bar::new(
909            first_type,
910            Price::from("100.01"),
911            Price::from("100.02"),
912            Price::from("100.00"),
913            Price::from("100.01"),
914            Quantity::from("10"),
915            1.into(),
916            1.into(),
917        );
918        let second = Bar::new(
919            second_type,
920            Price::from("100.01"),
921            Price::from("100.02"),
922            Price::from("100.00"),
923            Price::from("100.01"),
924            Quantity::from("11"),
925            2.into(),
926            2.into(),
927        );
928
929        let result = bars_to_arrow_record_batch_bytes(&[first, second]);
930
931        assert!(matches!(
932            result,
933            Err(EncodingError::MixedMetadata { index: 1 })
934        ));
935    }
936
937    #[rstest]
938    fn test_depth10_to_arrow_record_batch_rejects_mixed_level_price_precision() {
939        let instrument_id = InstrumentId::from("AUD/USD.SIM");
940        let bid = BookOrder::new(
941            OrderSide::Buy,
942            Price::from("1.23"),
943            Quantity::from("100.00"),
944            1,
945        );
946        let ask = BookOrder::new(
947            OrderSide::Sell,
948            Price::from("1.24"),
949            Quantity::from("100.00"),
950            2,
951        );
952        let mut asks = [ask; DEPTH10_LEN];
953        asks[1].price = Price::from("1.241");
954        let depth = OrderBookDepth10::new(
955            instrument_id,
956            [bid; DEPTH10_LEN],
957            asks,
958            [1; DEPTH10_LEN],
959            [1; DEPTH10_LEN],
960            0,
961            1,
962            1.into(),
963            1.into(),
964        );
965
966        let result = book_depth10_to_arrow_record_batch_bytes(&[depth]);
967
968        assert!(matches!(
969            result,
970            Err(EncodingError::MixedMetadata { index: 0 })
971        ));
972    }
973
974    #[rstest]
975    fn test_depth10_to_arrow_record_batch_rejects_mixed_level_size_precision() {
976        let instrument_id = InstrumentId::from("AUD/USD.SIM");
977        let bid = BookOrder::new(
978            OrderSide::Buy,
979            Price::from("1.23"),
980            Quantity::from("100.00"),
981            1,
982        );
983        let ask = BookOrder::new(
984            OrderSide::Sell,
985            Price::from("1.24"),
986            Quantity::from("100.00"),
987            2,
988        );
989        let mut bids = [bid; DEPTH10_LEN];
990        bids[1].size = Quantity::from("100.000");
991        let depth = OrderBookDepth10::new(
992            instrument_id,
993            bids,
994            [ask; DEPTH10_LEN],
995            [1; DEPTH10_LEN],
996            [1; DEPTH10_LEN],
997            0,
998            1,
999            1.into(),
1000            1.into(),
1001        );
1002
1003        let result = book_depth10_to_arrow_record_batch_bytes(&[depth]);
1004
1005        assert!(matches!(
1006            result,
1007            Err(EncodingError::MixedMetadata { index: 0 })
1008        ));
1009    }
1010
1011    #[rstest]
1012    fn test_depth10_to_arrow_record_batch_uses_first_defined_level_precision() {
1013        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1014        let bid = BookOrder::new(
1015            OrderSide::Buy,
1016            Price::from("1.23"),
1017            Quantity::from("100.00"),
1018            1,
1019        );
1020        let ask = BookOrder::new(
1021            OrderSide::Sell,
1022            Price::from("1.24"),
1023            Quantity::from("100.00"),
1024            2,
1025        );
1026        let mut bids = [bid; DEPTH10_LEN];
1027        bids[0].price = Price::from_raw(PRICE_UNDEF, 0);
1028        bids[0].size = Quantity::from_raw(QUANTITY_UNDEF, 0);
1029        let depth = OrderBookDepth10::new(
1030            instrument_id,
1031            bids,
1032            [ask; DEPTH10_LEN],
1033            [0; DEPTH10_LEN],
1034            [1; DEPTH10_LEN],
1035            0,
1036            1,
1037            1.into(),
1038            1.into(),
1039        );
1040
1041        let result = book_depth10_to_arrow_record_batch_bytes(&[depth]).unwrap();
1042
1043        assert_eq!(
1044            result.schema().metadata().get(KEY_PRICE_PRECISION).unwrap(),
1045            "2"
1046        );
1047        assert_eq!(
1048            result.schema().metadata().get(KEY_SIZE_PRECISION).unwrap(),
1049            "2"
1050        );
1051    }
1052
1053    #[rstest]
1054    #[case::price(true)]
1055    #[case::size(false)]
1056    fn test_depth10_to_arrow_record_batch_rejects_partial_undefined_level(
1057        #[case] price_undefined: bool,
1058    ) {
1059        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1060        let bid = BookOrder::new(
1061            OrderSide::Buy,
1062            Price::from("1.23"),
1063            Quantity::from("100.00"),
1064            1,
1065        );
1066        let ask = BookOrder::new(
1067            OrderSide::Sell,
1068            Price::from("1.24"),
1069            Quantity::from("100.00"),
1070            2,
1071        );
1072        let mut asks = [ask; DEPTH10_LEN];
1073        if price_undefined {
1074            asks[1].price = Price::from_raw(PRICE_UNDEF, 0);
1075        } else {
1076            asks[1].size = Quantity::from_raw(QUANTITY_UNDEF, 0);
1077        }
1078        let depth = OrderBookDepth10::new(
1079            instrument_id,
1080            [bid; DEPTH10_LEN],
1081            asks,
1082            [1; DEPTH10_LEN],
1083            [1; DEPTH10_LEN],
1084            0,
1085            1,
1086            1.into(),
1087            1.into(),
1088        );
1089
1090        let result = book_depth10_to_arrow_record_batch_bytes(&[depth]);
1091
1092        assert!(matches!(
1093            result,
1094            Err(EncodingError::MixedMetadata { index: 0 })
1095        ));
1096    }
1097
1098    #[rstest]
1099    fn test_deltas_to_arrow_record_batch_skips_leading_clears_for_precision() {
1100        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1101        let first = OrderBookDelta::clear(instrument_id, 0, 1.into(), 1.into());
1102        let second = OrderBookDelta::clear(instrument_id, 1, 2.into(), 2.into());
1103        let third = OrderBookDelta::new(
1104            instrument_id,
1105            BookAction::Add,
1106            BookOrder::new(
1107                OrderSide::Buy,
1108                Price::from("1.23"),
1109                Quantity::from("100.000000"),
1110                1,
1111            ),
1112            0,
1113            2,
1114            3.into(),
1115            3.into(),
1116        );
1117        let expected = vec![first, second, third];
1118
1119        let batch = book_deltas_to_arrow_record_batch_bytes(&expected).unwrap();
1120        let metadata = batch.schema().metadata().clone();
1121        assert_eq!(
1122            metadata.get(KEY_PRICE_PRECISION).map(String::as_str),
1123            Some("2")
1124        );
1125        assert_eq!(
1126            metadata.get(KEY_SIZE_PRECISION).map(String::as_str),
1127            Some("6")
1128        );
1129
1130        let decoded = OrderBookDelta::decode_batch(&metadata, batch).unwrap();
1131
1132        assert_eq!(decoded, expected);
1133        assert_eq!(decoded[2].order.price.precision, 2);
1134        assert_eq!(decoded[2].order.size.precision, 6);
1135    }
1136
1137    #[rstest]
1138    fn test_deltas_to_arrow_record_batch_all_clear_roundtrip() {
1139        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1140        let expected = vec![
1141            OrderBookDelta::clear(instrument_id, 0, 1.into(), 1.into()),
1142            OrderBookDelta::clear(instrument_id, 1, 2.into(), 2.into()),
1143        ];
1144
1145        let batch = book_deltas_to_arrow_record_batch_bytes(&expected).unwrap();
1146        let metadata = batch.schema().metadata().clone();
1147        let decoded = OrderBookDelta::decode_batch(&metadata, batch).unwrap();
1148
1149        assert_eq!(decoded, expected);
1150    }
1151
1152    #[rstest]
1153    fn test_deltas_to_arrow_record_batch_rejects_mixed_precision() {
1154        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1155        let first = OrderBookDelta::new(
1156            instrument_id,
1157            BookAction::Add,
1158            BookOrder::new(
1159                OrderSide::Buy,
1160                Price::from("1.23"),
1161                Quantity::from("100.00"),
1162                1,
1163            ),
1164            0,
1165            1,
1166            1.into(),
1167            1.into(),
1168        );
1169        let second = OrderBookDelta::new(
1170            instrument_id,
1171            BookAction::Update,
1172            BookOrder::new(
1173                OrderSide::Buy,
1174                Price::from("1.234"),
1175                Quantity::from("100.000"),
1176                1,
1177            ),
1178            0,
1179            2,
1180            2.into(),
1181            2.into(),
1182        );
1183
1184        let result = book_deltas_to_arrow_record_batch_bytes(&[first, second]);
1185
1186        assert!(matches!(
1187            result,
1188            Err(EncodingError::MixedMetadata { index: 1 })
1189        ));
1190    }
1191
1192    #[rstest]
1193    fn test_deltas_to_arrow_record_batch_rejects_mixed_instruments() {
1194        let first = OrderBookDelta::clear(InstrumentId::from("AUD/USD.SIM"), 0, 1.into(), 1.into());
1195        let second = OrderBookDelta::new(
1196            InstrumentId::from("EUR/USD.SIM"),
1197            BookAction::Add,
1198            BookOrder::new(
1199                OrderSide::Buy,
1200                Price::from("1.23"),
1201                Quantity::from("100.00"),
1202                1,
1203            ),
1204            0,
1205            1,
1206            2.into(),
1207            2.into(),
1208        );
1209
1210        let result = book_deltas_to_arrow_record_batch_bytes(&[first, second]);
1211
1212        // The first non-clear delta supplies metadata, so the leading clear is the mismatched row.
1213        assert!(matches!(
1214            result,
1215            Err(EncodingError::MixedMetadata { index: 0 })
1216        ));
1217    }
1218}
nautilus_serialization/arrow/mod.rs

nautilus_serialization/arrow/
mod.rs
