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

nautilus_serialization/arrow/
mod.rs
