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

nautilus_serialization/arrow/
mod.rs
