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)]
822fn fixed_size_binary<const N: usize>(values: Vec<&[u8; N]>) -> FixedSizeBinaryArray {
823    FixedSizeBinaryArray::try_from_iter(values.into_iter()).unwrap()
824}
825
826#[cfg(test)]
827mod tests {
828    use nautilus_model::{
829        data::{
830            Bar, BarSpecification, BarType, BookOrder, OrderBookDelta, OrderBookDepth10, QuoteTick,
831            depth::DEPTH10_LEN,
832        },
833        enums::{AggregationSource, BarAggregation, BookAction, OrderSide, PriceType},
834        identifiers::InstrumentId,
835        types::{PRICE_UNDEF, Price, QUANTITY_UNDEF, Quantity},
836    };
837    use rstest::rstest;
838
839    use super::*;
840
841    #[rstest]
842    fn test_quotes_to_arrow_record_batch_rejects_mixed_instruments() {
843        let first = QuoteTick::new(
844            InstrumentId::from("AAPL.XNAS"),
845            Price::from("100.01"),
846            Price::from("100.02"),
847            Quantity::from("10"),
848            Quantity::from("11"),
849            1.into(),
850            1.into(),
851        );
852        let second = QuoteTick::new(
853            InstrumentId::from("MSFT.XNAS"),
854            Price::from("200.01"),
855            Price::from("200.02"),
856            Quantity::from("20"),
857            Quantity::from("21"),
858            2.into(),
859            2.into(),
860        );
861
862        let result = quotes_to_arrow_record_batch_bytes(&[first, second]);
863
864        assert!(matches!(
865            result,
866            Err(EncodingError::MixedMetadata { index: 1 })
867        ));
868    }
869
870    #[rstest]
871    fn test_quotes_to_arrow_record_batch_rejects_mixed_precision() {
872        let instrument_id = InstrumentId::from("AAPL.XNAS");
873        let first = QuoteTick::new(
874            instrument_id,
875            Price::from("100.01"),
876            Price::from("100.02"),
877            Quantity::from("10.00"),
878            Quantity::from("11.00"),
879            1.into(),
880            1.into(),
881        );
882        let second = QuoteTick::new(
883            instrument_id,
884            Price::from("100.010"),
885            Price::from("100.020"),
886            Quantity::from("10.000"),
887            Quantity::from("11.000"),
888            2.into(),
889            2.into(),
890        );
891
892        let result = quotes_to_arrow_record_batch_bytes(&[first, second]);
893
894        assert!(matches!(
895            result,
896            Err(EncodingError::MixedMetadata { index: 1 })
897        ));
898    }
899
900    #[rstest]
901    fn test_bars_to_arrow_record_batch_rejects_mixed_bar_types() {
902        let instrument_id = InstrumentId::from("AAPL.XNAS");
903        let first_type = BarType::new(
904            instrument_id,
905            BarSpecification::new(1, BarAggregation::Minute, PriceType::Last),
906            AggregationSource::Internal,
907        );
908        let second_type = BarType::new(
909            instrument_id,
910            BarSpecification::new(5, BarAggregation::Minute, PriceType::Last),
911            AggregationSource::Internal,
912        );
913        let first = Bar::new(
914            first_type,
915            Price::from("100.01"),
916            Price::from("100.02"),
917            Price::from("100.00"),
918            Price::from("100.01"),
919            Quantity::from("10"),
920            1.into(),
921            1.into(),
922        );
923        let second = Bar::new(
924            second_type,
925            Price::from("100.01"),
926            Price::from("100.02"),
927            Price::from("100.00"),
928            Price::from("100.01"),
929            Quantity::from("11"),
930            2.into(),
931            2.into(),
932        );
933
934        let result = bars_to_arrow_record_batch_bytes(&[first, second]);
935
936        assert!(matches!(
937            result,
938            Err(EncodingError::MixedMetadata { index: 1 })
939        ));
940    }
941
942    #[rstest]
943    fn test_depth10_to_arrow_record_batch_rejects_mixed_level_price_precision() {
944        let instrument_id = InstrumentId::from("AUD/USD.SIM");
945        let bid = BookOrder::new(
946            OrderSide::Buy,
947            Price::from("1.23"),
948            Quantity::from("100.00"),
949            1,
950        );
951        let ask = BookOrder::new(
952            OrderSide::Sell,
953            Price::from("1.24"),
954            Quantity::from("100.00"),
955            2,
956        );
957        let mut asks = [ask; DEPTH10_LEN];
958        asks[1].price = Price::from("1.241");
959        let depth = OrderBookDepth10::new(
960            instrument_id,
961            [bid; DEPTH10_LEN],
962            asks,
963            [1; DEPTH10_LEN],
964            [1; DEPTH10_LEN],
965            0,
966            1,
967            1.into(),
968            1.into(),
969        );
970
971        let result = book_depth10_to_arrow_record_batch_bytes(&[depth]);
972
973        assert!(matches!(
974            result,
975            Err(EncodingError::MixedMetadata { index: 0 })
976        ));
977    }
978
979    #[rstest]
980    fn test_depth10_to_arrow_record_batch_rejects_mixed_level_size_precision() {
981        let instrument_id = InstrumentId::from("AUD/USD.SIM");
982        let bid = BookOrder::new(
983            OrderSide::Buy,
984            Price::from("1.23"),
985            Quantity::from("100.00"),
986            1,
987        );
988        let ask = BookOrder::new(
989            OrderSide::Sell,
990            Price::from("1.24"),
991            Quantity::from("100.00"),
992            2,
993        );
994        let mut bids = [bid; DEPTH10_LEN];
995        bids[1].size = Quantity::from("100.000");
996        let depth = OrderBookDepth10::new(
997            instrument_id,
998            bids,
999            [ask; DEPTH10_LEN],
1000            [1; DEPTH10_LEN],
1001            [1; DEPTH10_LEN],
1002            0,
1003            1,
1004            1.into(),
1005            1.into(),
1006        );
1007
1008        let result = book_depth10_to_arrow_record_batch_bytes(&[depth]);
1009
1010        assert!(matches!(
1011            result,
1012            Err(EncodingError::MixedMetadata { index: 0 })
1013        ));
1014    }
1015
1016    #[rstest]
1017    fn test_depth10_to_arrow_record_batch_uses_first_defined_level_precision() {
1018        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1019        let bid = BookOrder::new(
1020            OrderSide::Buy,
1021            Price::from("1.23"),
1022            Quantity::from("100.00"),
1023            1,
1024        );
1025        let ask = BookOrder::new(
1026            OrderSide::Sell,
1027            Price::from("1.24"),
1028            Quantity::from("100.00"),
1029            2,
1030        );
1031        let mut bids = [bid; DEPTH10_LEN];
1032        bids[0].price = Price::from_raw(PRICE_UNDEF, 0);
1033        bids[0].size = Quantity::from_raw(QUANTITY_UNDEF, 0);
1034        let depth = OrderBookDepth10::new(
1035            instrument_id,
1036            bids,
1037            [ask; DEPTH10_LEN],
1038            [0; DEPTH10_LEN],
1039            [1; DEPTH10_LEN],
1040            0,
1041            1,
1042            1.into(),
1043            1.into(),
1044        );
1045
1046        let result = book_depth10_to_arrow_record_batch_bytes(&[depth]).unwrap();
1047
1048        assert_eq!(
1049            result.schema().metadata().get(KEY_PRICE_PRECISION).unwrap(),
1050            "2"
1051        );
1052        assert_eq!(
1053            result.schema().metadata().get(KEY_SIZE_PRECISION).unwrap(),
1054            "2"
1055        );
1056    }
1057
1058    #[rstest]
1059    #[case::price(true)]
1060    #[case::size(false)]
1061    fn test_depth10_to_arrow_record_batch_rejects_partial_undefined_level(
1062        #[case] price_undefined: bool,
1063    ) {
1064        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1065        let bid = BookOrder::new(
1066            OrderSide::Buy,
1067            Price::from("1.23"),
1068            Quantity::from("100.00"),
1069            1,
1070        );
1071        let ask = BookOrder::new(
1072            OrderSide::Sell,
1073            Price::from("1.24"),
1074            Quantity::from("100.00"),
1075            2,
1076        );
1077        let mut asks = [ask; DEPTH10_LEN];
1078        if price_undefined {
1079            asks[1].price = Price::from_raw(PRICE_UNDEF, 0);
1080        } else {
1081            asks[1].size = Quantity::from_raw(QUANTITY_UNDEF, 0);
1082        }
1083        let depth = OrderBookDepth10::new(
1084            instrument_id,
1085            [bid; DEPTH10_LEN],
1086            asks,
1087            [1; DEPTH10_LEN],
1088            [1; DEPTH10_LEN],
1089            0,
1090            1,
1091            1.into(),
1092            1.into(),
1093        );
1094
1095        let result = book_depth10_to_arrow_record_batch_bytes(&[depth]);
1096
1097        assert!(matches!(
1098            result,
1099            Err(EncodingError::MixedMetadata { index: 0 })
1100        ));
1101    }
1102
1103    #[rstest]
1104    fn test_deltas_to_arrow_record_batch_skips_leading_clears_for_precision() {
1105        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1106        let first = OrderBookDelta::clear(instrument_id, 0, 1.into(), 1.into());
1107        let second = OrderBookDelta::clear(instrument_id, 1, 2.into(), 2.into());
1108        let third = OrderBookDelta::new(
1109            instrument_id,
1110            BookAction::Add,
1111            BookOrder::new(
1112                OrderSide::Buy,
1113                Price::from("1.23"),
1114                Quantity::from("100.000000"),
1115                1,
1116            ),
1117            0,
1118            2,
1119            3.into(),
1120            3.into(),
1121        );
1122        let expected = vec![first, second, third];
1123
1124        let batch = book_deltas_to_arrow_record_batch_bytes(&expected).unwrap();
1125        let metadata = batch.schema().metadata().clone();
1126        assert_eq!(
1127            metadata.get(KEY_PRICE_PRECISION).map(String::as_str),
1128            Some("2")
1129        );
1130        assert_eq!(
1131            metadata.get(KEY_SIZE_PRECISION).map(String::as_str),
1132            Some("6")
1133        );
1134
1135        let decoded = OrderBookDelta::decode_batch(&metadata, batch).unwrap();
1136
1137        assert_eq!(decoded, expected);
1138        assert_eq!(decoded[2].order.price.precision, 2);
1139        assert_eq!(decoded[2].order.size.precision, 6);
1140    }
1141
1142    #[rstest]
1143    fn test_deltas_to_arrow_record_batch_all_clear_roundtrip() {
1144        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1145        let expected = vec![
1146            OrderBookDelta::clear(instrument_id, 0, 1.into(), 1.into()),
1147            OrderBookDelta::clear(instrument_id, 1, 2.into(), 2.into()),
1148        ];
1149
1150        let batch = book_deltas_to_arrow_record_batch_bytes(&expected).unwrap();
1151        let metadata = batch.schema().metadata().clone();
1152        let decoded = OrderBookDelta::decode_batch(&metadata, batch).unwrap();
1153
1154        assert_eq!(decoded, expected);
1155    }
1156
1157    #[rstest]
1158    fn test_deltas_to_arrow_record_batch_rejects_mixed_precision() {
1159        let instrument_id = InstrumentId::from("AUD/USD.SIM");
1160        let first = OrderBookDelta::new(
1161            instrument_id,
1162            BookAction::Add,
1163            BookOrder::new(
1164                OrderSide::Buy,
1165                Price::from("1.23"),
1166                Quantity::from("100.00"),
1167                1,
1168            ),
1169            0,
1170            1,
1171            1.into(),
1172            1.into(),
1173        );
1174        let second = OrderBookDelta::new(
1175            instrument_id,
1176            BookAction::Update,
1177            BookOrder::new(
1178                OrderSide::Buy,
1179                Price::from("1.234"),
1180                Quantity::from("100.000"),
1181                1,
1182            ),
1183            0,
1184            2,
1185            2.into(),
1186            2.into(),
1187        );
1188
1189        let result = book_deltas_to_arrow_record_batch_bytes(&[first, second]);
1190
1191        assert!(matches!(
1192            result,
1193            Err(EncodingError::MixedMetadata { index: 1 })
1194        ));
1195    }
1196
1197    #[rstest]
1198    fn test_deltas_to_arrow_record_batch_rejects_mixed_instruments() {
1199        let first = OrderBookDelta::clear(InstrumentId::from("AUD/USD.SIM"), 0, 1.into(), 1.into());
1200        let second = OrderBookDelta::new(
1201            InstrumentId::from("EUR/USD.SIM"),
1202            BookAction::Add,
1203            BookOrder::new(
1204                OrderSide::Buy,
1205                Price::from("1.23"),
1206                Quantity::from("100.00"),
1207                1,
1208            ),
1209            0,
1210            1,
1211            2.into(),
1212            2.into(),
1213        );
1214
1215        let result = book_deltas_to_arrow_record_batch_bytes(&[first, second]);
1216
1217        // The first non-clear delta supplies metadata, so the leading clear is the mismatched row.
1218        assert!(matches!(
1219            result,
1220            Err(EncodingError::MixedMetadata { index: 0 })
1221        ));
1222    }
1223}
nautilus_serialization/arrow/mod.rs

nautilus_serialization/arrow/
mod.rs
