Skip to main content

nautilus_serialization/python/
arrow.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
16use std::io::Cursor;
17
18use arrow::{
19    ipc::{reader::StreamReader, writer::StreamWriter},
20    record_batch::RecordBatch,
21};
22use nautilus_core::python::{to_pyruntime_err, to_pytype_err, to_pyvalue_err};
23use nautilus_model::{
24    data::{
25        Bar, IndexPriceUpdate, InstrumentStatus, MarkPriceUpdate, OptionGreeks, OrderBookDelta,
26        OrderBookDepth10, QuoteTick, TradeTick, close::InstrumentClose,
27    },
28    python::data::{
29        pyobjects_to_bars, pyobjects_to_book_deltas, pyobjects_to_index_prices,
30        pyobjects_to_instrument_closes, pyobjects_to_instrument_statuses, pyobjects_to_mark_prices,
31        pyobjects_to_option_greeks, pyobjects_to_quotes, pyobjects_to_trades,
32    },
33};
34use pyo3::{
35    conversion::IntoPyObjectExt,
36    prelude::*,
37    types::{PyBytes, PyType},
38};
39
40use crate::arrow::{
41    ArrowSchemaProvider, DecodeFromRecordBatch, DecodeTypedFromRecordBatch,
42    bars_to_arrow_record_batch_bytes, book_deltas_to_arrow_record_batch_bytes,
43    book_depth10_to_arrow_record_batch_bytes, index_prices_to_arrow_record_batch_bytes,
44    instrument_closes_to_arrow_record_batch_bytes, instrument_status_to_arrow_record_batch_bytes,
45    mark_prices_to_arrow_record_batch_bytes, option_greeks_to_arrow_record_batch_bytes,
46    quotes_to_arrow_record_batch_bytes, trades_to_arrow_record_batch_bytes,
47};
48
49/// Transforms the given record `batch` into Python `bytes`.
50///
51/// # Errors
52///
53/// Returns a `PyErr` if writing the Arrow IPC stream fails.
54pub fn arrow_record_batch_to_pybytes(py: Python, batch: &RecordBatch) -> PyResult<Py<PyBytes>> {
55    // Create a cursor to write to a byte array in memory
56    let mut cursor = Cursor::new(Vec::new());
57    {
58        let mut writer =
59            StreamWriter::try_new(&mut cursor, &batch.schema()).map_err(to_pyruntime_err)?;
60
61        writer.write(batch).map_err(to_pyruntime_err)?;
62
63        writer.finish().map_err(to_pyruntime_err)?;
64    }
65
66    let buffer = cursor.into_inner();
67    let pybytes = PyBytes::new(py, &buffer);
68
69    Ok(pybytes.into())
70}
71
72/// Returns a mapping from field names to Arrow data types for the given Rust data class.
73///
74/// # Errors
75///
76/// Returns a `PyErr` if the class name is not recognized or schema extraction fails.
77#[pyfunction]
78#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
79pub fn get_arrow_schema_map(py: Python<'_>, cls: &Bound<'_, PyType>) -> PyResult<Py<PyAny>> {
80    let cls_str: String = cls.getattr("__name__")?.extract()?;
81    let result_map = match cls_str.as_str() {
82        stringify!(OrderBookDelta) => OrderBookDelta::get_schema_map(),
83        stringify!(OrderBookDepth10) => OrderBookDepth10::get_schema_map(),
84        stringify!(QuoteTick) => QuoteTick::get_schema_map(),
85        stringify!(TradeTick) => TradeTick::get_schema_map(),
86        stringify!(Bar) => Bar::get_schema_map(),
87        stringify!(MarkPriceUpdate) => MarkPriceUpdate::get_schema_map(),
88        stringify!(IndexPriceUpdate) => IndexPriceUpdate::get_schema_map(),
89        stringify!(InstrumentStatus) => InstrumentStatus::get_schema_map(),
90        stringify!(OptionGreeks) => OptionGreeks::get_schema_map(),
91        stringify!(InstrumentClose) => InstrumentClose::get_schema_map(),
92        _ => {
93            return Err(to_pytype_err(format!(
94                "Arrow schema for `{cls_str}` is not currently implemented in Rust."
95            )));
96        }
97    };
98
99    result_map.into_py_any(py)
100}
101
102/// Converts a vector of `OrderBookDelta` into an Arrow `RecordBatch`.
103#[pyfunction]
104#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
105#[expect(clippy::missing_panics_doc)] // Guarded by empty check
106pub fn pyobjects_to_arrow_record_batch_bytes(
107    py: Python,
108    data: Vec<Bound<'_, PyAny>>,
109) -> PyResult<Py<PyBytes>> {
110    if data.is_empty() {
111        return Err(to_pyvalue_err("Empty data"));
112    }
113
114    let data_type: String = data
115        .first()
116        .unwrap() // SAFETY: Unwrap safe as already checked that `data` not empty
117        .getattr("__class__")?
118        .getattr("__name__")?
119        .extract()?;
120
121    match data_type.as_str() {
122        stringify!(OrderBookDelta) => {
123            let deltas = pyobjects_to_book_deltas(data)?;
124            py_book_deltas_to_arrow_record_batch_bytes(py, deltas)
125        }
126        stringify!(OrderBookDepth10) => {
127            let depth_snapshots: Vec<OrderBookDepth10> = data
128                .into_iter()
129                .map(|obj| obj.extract::<OrderBookDepth10>().map_err(Into::into))
130                .collect::<PyResult<Vec<OrderBookDepth10>>>()?;
131            py_book_depth10_to_arrow_record_batch_bytes(py, depth_snapshots)
132        }
133        stringify!(QuoteTick) => {
134            let quotes = pyobjects_to_quotes(data)?;
135            py_quotes_to_arrow_record_batch_bytes(py, quotes)
136        }
137        stringify!(TradeTick) => {
138            let trades = pyobjects_to_trades(data)?;
139            py_trades_to_arrow_record_batch_bytes(py, trades)
140        }
141        stringify!(Bar) => {
142            let bars = pyobjects_to_bars(data)?;
143            py_bars_to_arrow_record_batch_bytes(py, bars)
144        }
145        stringify!(MarkPriceUpdate) => {
146            let updates = pyobjects_to_mark_prices(data)?;
147            py_mark_prices_to_arrow_record_batch_bytes(py, updates)
148        }
149        stringify!(IndexPriceUpdate) => {
150            let index_prices = pyobjects_to_index_prices(data)?;
151            py_index_prices_to_arrow_record_batch_bytes(py, index_prices)
152        }
153        stringify!(InstrumentStatus) => {
154            let statuses = pyobjects_to_instrument_statuses(data)?;
155            py_instrument_status_to_arrow_record_batch_bytes(py, statuses)
156        }
157        stringify!(OptionGreeks) => {
158            let greeks = pyobjects_to_option_greeks(data)?;
159            py_option_greeks_to_arrow_record_batch_bytes(py, greeks)
160        }
161        stringify!(InstrumentClose) => {
162            let closes = pyobjects_to_instrument_closes(data)?;
163            py_instrument_closes_to_arrow_record_batch_bytes(py, closes)
164        }
165        _ => Err(to_pyvalue_err(format!(
166            "unsupported data type: {data_type}"
167        ))),
168    }
169}
170
171/// Converts a vector of `OrderBookDelta` into an Arrow `RecordBatch`.
172///
173/// # Errors
174///
175/// Returns an error if:
176/// - `data` is empty: `EncodingError::EmptyData`.
177/// - Instrument IDs differ, or non-clear precision metadata differs:
178///   `EncodingError::MixedMetadata`.
179/// - Encoding fails: `EncodingError::ArrowError`.
180#[pyfunction(name = "book_deltas_to_arrow_record_batch_bytes")]
181#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
182#[expect(clippy::needless_pass_by_value)]
183pub fn py_book_deltas_to_arrow_record_batch_bytes(
184    py: Python,
185    data: Vec<OrderBookDelta>,
186) -> PyResult<Py<PyBytes>> {
187    match book_deltas_to_arrow_record_batch_bytes(&data) {
188        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
189        Err(e) => Err(to_pyvalue_err(e)),
190    }
191}
192
193/// Converts a vector of `OrderBookDepth10` into an Arrow `RecordBatch`.
194///
195/// # Errors
196///
197/// Returns an error if:
198/// - `data` is empty: `EncodingError::EmptyData`.
199/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
200/// - Encoding fails: `EncodingError::ArrowError`.
201#[pyfunction(name = "book_depth10_to_arrow_record_batch_bytes")]
202#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
203#[expect(clippy::needless_pass_by_value)]
204pub fn py_book_depth10_to_arrow_record_batch_bytes(
205    py: Python,
206    data: Vec<OrderBookDepth10>,
207) -> PyResult<Py<PyBytes>> {
208    match book_depth10_to_arrow_record_batch_bytes(&data) {
209        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
210        Err(e) => Err(to_pyvalue_err(e)),
211    }
212}
213
214/// Converts a vector of `QuoteTick` into an Arrow `RecordBatch`.
215///
216/// # Errors
217///
218/// Returns an error if:
219/// - `data` is empty: `EncodingError::EmptyData`.
220/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
221/// - Encoding fails: `EncodingError::ArrowError`.
222#[pyfunction(name = "quotes_to_arrow_record_batch_bytes")]
223#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
224#[expect(clippy::needless_pass_by_value)]
225pub fn py_quotes_to_arrow_record_batch_bytes(
226    py: Python,
227    data: Vec<QuoteTick>,
228) -> PyResult<Py<PyBytes>> {
229    match quotes_to_arrow_record_batch_bytes(&data) {
230        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
231        Err(e) => Err(to_pyvalue_err(e)),
232    }
233}
234
235/// Converts a vector of `TradeTick` into an Arrow `RecordBatch`.
236///
237/// # Errors
238///
239/// Returns an error if:
240/// - `data` is empty: `EncodingError::EmptyData`.
241/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
242/// - Encoding fails: `EncodingError::ArrowError`.
243#[pyfunction(name = "trades_to_arrow_record_batch_bytes")]
244#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
245#[expect(clippy::needless_pass_by_value)]
246pub fn py_trades_to_arrow_record_batch_bytes(
247    py: Python,
248    data: Vec<TradeTick>,
249) -> PyResult<Py<PyBytes>> {
250    match trades_to_arrow_record_batch_bytes(&data) {
251        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
252        Err(e) => Err(to_pyvalue_err(e)),
253    }
254}
255
256/// Converts a vector of `Bar` into an Arrow `RecordBatch`.
257///
258/// # Errors
259///
260/// Returns an error if:
261/// - `data` is empty: `EncodingError::EmptyData`.
262/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
263/// - Encoding fails: `EncodingError::ArrowError`.
264#[pyfunction(name = "bars_to_arrow_record_batch_bytes")]
265#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
266#[expect(clippy::needless_pass_by_value)]
267pub fn py_bars_to_arrow_record_batch_bytes(py: Python, data: Vec<Bar>) -> PyResult<Py<PyBytes>> {
268    match bars_to_arrow_record_batch_bytes(&data) {
269        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
270        Err(e) => Err(to_pyvalue_err(e)),
271    }
272}
273
274/// Converts a vector of `MarkPriceUpdate` into an Arrow `RecordBatch`.
275///
276/// # Errors
277///
278/// Returns an error if:
279/// - `data` is empty: `EncodingError::EmptyData`.
280/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
281/// - Encoding fails: `EncodingError::ArrowError`.
282#[pyfunction(name = "mark_prices_to_arrow_record_batch_bytes")]
283#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
284#[expect(clippy::needless_pass_by_value)]
285pub fn py_mark_prices_to_arrow_record_batch_bytes(
286    py: Python,
287    data: Vec<MarkPriceUpdate>,
288) -> PyResult<Py<PyBytes>> {
289    match mark_prices_to_arrow_record_batch_bytes(&data) {
290        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
291        Err(e) => Err(to_pyvalue_err(e)),
292    }
293}
294
295/// Converts a vector of `IndexPriceUpdate` into an Arrow `RecordBatch`.
296///
297/// # Errors
298///
299/// Returns an error if:
300/// - `data` is empty: `EncodingError::EmptyData`.
301/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
302/// - Encoding fails: `EncodingError::ArrowError`.
303#[pyfunction(name = "index_prices_to_arrow_record_batch_bytes")]
304#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
305#[expect(clippy::needless_pass_by_value)]
306pub fn py_index_prices_to_arrow_record_batch_bytes(
307    py: Python,
308    data: Vec<IndexPriceUpdate>,
309) -> PyResult<Py<PyBytes>> {
310    match index_prices_to_arrow_record_batch_bytes(&data) {
311        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
312        Err(e) => Err(to_pyvalue_err(e)),
313    }
314}
315
316/// Converts a vector of `InstrumentStatus` into an Arrow `RecordBatch`.
317///
318/// # Errors
319///
320/// Returns an error if:
321/// - `data` is empty: `EncodingError::EmptyData`.
322/// - Encoding fails: `EncodingError::ArrowError`.
323#[pyfunction(name = "instrument_status_to_arrow_record_batch_bytes")]
324#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
325#[expect(clippy::needless_pass_by_value)]
326pub fn py_instrument_status_to_arrow_record_batch_bytes(
327    py: Python,
328    data: Vec<InstrumentStatus>,
329) -> PyResult<Py<PyBytes>> {
330    match instrument_status_to_arrow_record_batch_bytes(&data) {
331        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
332        Err(e) => Err(to_pyvalue_err(e)),
333    }
334}
335
336/// Converts a vector of `OptionGreeks` into an Arrow `RecordBatch`.
337///
338/// # Errors
339///
340/// Returns an error if:
341/// - `data` is empty: `EncodingError::EmptyData`.
342/// - Encoding fails: `EncodingError::ArrowError`.
343#[pyfunction(name = "option_greeks_to_arrow_record_batch_bytes")]
344#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
345#[expect(clippy::needless_pass_by_value)]
346pub fn py_option_greeks_to_arrow_record_batch_bytes(
347    py: Python,
348    data: Vec<OptionGreeks>,
349) -> PyResult<Py<PyBytes>> {
350    match option_greeks_to_arrow_record_batch_bytes(&data) {
351        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
352        Err(e) => Err(to_pyvalue_err(e)),
353    }
354}
355
356/// Decodes Arrow IPC bytes into a list of `OptionGreeks`.
357///
358/// # Errors
359///
360/// Returns a `PyErr` if decoding fails.
361#[pyfunction(name = "option_greeks_from_arrow_record_batch_bytes")]
362#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
363pub fn py_option_greeks_from_arrow_record_batch_bytes(
364    _py: Python,
365    data: Vec<u8>,
366) -> PyResult<Vec<OptionGreeks>> {
367    let cursor = Cursor::new(data);
368    let reader = StreamReader::try_new(cursor, None).map_err(to_pyruntime_err)?;
369
370    let mut results = Vec::new();
371
372    for batch_result in reader {
373        let batch = batch_result.map_err(to_pyruntime_err)?;
374        let metadata = batch.schema().metadata().clone();
375        let decoded = OptionGreeks::decode_batch(&metadata, batch).map_err(to_pyvalue_err)?;
376        results.extend(decoded);
377    }
378
379    Ok(results)
380}
381
382/// Decodes Arrow IPC bytes into a list of `InstrumentStatus`.
383///
384/// # Errors
385///
386/// Returns a `PyErr` if decoding fails.
387#[pyfunction(name = "instrument_status_from_arrow_record_batch_bytes")]
388#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
389pub fn py_instrument_status_from_arrow_record_batch_bytes(
390    _py: Python,
391    data: Vec<u8>,
392) -> PyResult<Vec<InstrumentStatus>> {
393    let cursor = Cursor::new(data);
394    let reader = StreamReader::try_new(cursor, None).map_err(to_pyruntime_err)?;
395
396    let mut results = Vec::new();
397
398    for batch_result in reader {
399        let batch = batch_result.map_err(to_pyruntime_err)?;
400        let metadata = batch.schema().metadata().clone();
401        let decoded =
402            InstrumentStatus::decode_typed_batch(&metadata, batch).map_err(to_pyvalue_err)?;
403        results.extend(decoded);
404    }
405
406    Ok(results)
407}
408
409/// Converts a vector of `InstrumentClose` into an Arrow `RecordBatch`.
410///
411/// # Errors
412///
413/// Returns an error if:
414/// - `data` is empty: `EncodingError::EmptyData`.
415/// - Metadata differs between rows: `EncodingError::MixedMetadata`.
416/// - Encoding fails: `EncodingError::ArrowError`.
417#[pyfunction(name = "instrument_closes_to_arrow_record_batch_bytes")]
418#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
419#[expect(clippy::needless_pass_by_value)]
420pub fn py_instrument_closes_to_arrow_record_batch_bytes(
421    py: Python,
422    data: Vec<InstrumentClose>,
423) -> PyResult<Py<PyBytes>> {
424    match instrument_closes_to_arrow_record_batch_bytes(&data) {
425        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
426        Err(e) => Err(to_pyvalue_err(e)),
427    }
428}