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 list of `OrderBookDelta` into Arrow IPC bytes for Python.
172///
173/// # Errors
174///
175/// Returns a `PyErr` if encoding fails.
176#[pyfunction(name = "book_deltas_to_arrow_record_batch_bytes")]
177#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
178#[expect(clippy::needless_pass_by_value)]
179pub fn py_book_deltas_to_arrow_record_batch_bytes(
180    py: Python,
181    data: Vec<OrderBookDelta>,
182) -> PyResult<Py<PyBytes>> {
183    match book_deltas_to_arrow_record_batch_bytes(&data) {
184        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
185        Err(e) => Err(to_pyvalue_err(e)),
186    }
187}
188
189/// Converts a list of `OrderBookDepth10` into Arrow IPC bytes for Python.
190///
191/// # Errors
192///
193/// Returns a `PyErr` if encoding fails.
194#[pyfunction(name = "book_depth10_to_arrow_record_batch_bytes")]
195#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
196#[expect(clippy::needless_pass_by_value)]
197pub fn py_book_depth10_to_arrow_record_batch_bytes(
198    py: Python,
199    data: Vec<OrderBookDepth10>,
200) -> PyResult<Py<PyBytes>> {
201    match book_depth10_to_arrow_record_batch_bytes(&data) {
202        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
203        Err(e) => Err(to_pyvalue_err(e)),
204    }
205}
206
207/// Converts a list of `QuoteTick` into Arrow IPC bytes for Python.
208///
209/// # Errors
210///
211/// Returns a `PyErr` if encoding fails.
212#[pyfunction(name = "quotes_to_arrow_record_batch_bytes")]
213#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
214#[expect(clippy::needless_pass_by_value)]
215pub fn py_quotes_to_arrow_record_batch_bytes(
216    py: Python,
217    data: Vec<QuoteTick>,
218) -> PyResult<Py<PyBytes>> {
219    match quotes_to_arrow_record_batch_bytes(&data) {
220        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
221        Err(e) => Err(to_pyvalue_err(e)),
222    }
223}
224
225/// Converts a list of `TradeTick` into Arrow IPC bytes for Python.
226///
227/// # Errors
228///
229/// Returns a `PyErr` if encoding fails.
230#[pyfunction(name = "trades_to_arrow_record_batch_bytes")]
231#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
232#[expect(clippy::needless_pass_by_value)]
233pub fn py_trades_to_arrow_record_batch_bytes(
234    py: Python,
235    data: Vec<TradeTick>,
236) -> PyResult<Py<PyBytes>> {
237    match trades_to_arrow_record_batch_bytes(&data) {
238        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
239        Err(e) => Err(to_pyvalue_err(e)),
240    }
241}
242
243/// Converts a list of `Bar` into Arrow IPC bytes for Python.
244///
245/// # Errors
246///
247/// Returns a `PyErr` if encoding fails.
248#[pyfunction(name = "bars_to_arrow_record_batch_bytes")]
249#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
250#[expect(clippy::needless_pass_by_value)]
251pub fn py_bars_to_arrow_record_batch_bytes(py: Python, data: Vec<Bar>) -> PyResult<Py<PyBytes>> {
252    match bars_to_arrow_record_batch_bytes(&data) {
253        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
254        Err(e) => Err(to_pyvalue_err(e)),
255    }
256}
257
258/// Converts a list of `MarkPriceUpdate` into Arrow IPC bytes for Python.
259///
260/// # Errors
261///
262/// Returns a `PyErr` if encoding fails.
263#[pyfunction(name = "mark_prices_to_arrow_record_batch_bytes")]
264#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
265#[expect(clippy::needless_pass_by_value)]
266pub fn py_mark_prices_to_arrow_record_batch_bytes(
267    py: Python,
268    data: Vec<MarkPriceUpdate>,
269) -> PyResult<Py<PyBytes>> {
270    match mark_prices_to_arrow_record_batch_bytes(&data) {
271        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
272        Err(e) => Err(to_pyvalue_err(e)),
273    }
274}
275
276/// Converts a list of `IndexPriceUpdate` into Arrow IPC bytes for Python.
277///
278/// # Errors
279///
280/// Returns a `PyErr` if encoding fails.
281#[pyfunction(name = "index_prices_to_arrow_record_batch_bytes")]
282#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
283#[expect(clippy::needless_pass_by_value)]
284pub fn py_index_prices_to_arrow_record_batch_bytes(
285    py: Python,
286    data: Vec<IndexPriceUpdate>,
287) -> PyResult<Py<PyBytes>> {
288    match index_prices_to_arrow_record_batch_bytes(&data) {
289        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
290        Err(e) => Err(to_pyvalue_err(e)),
291    }
292}
293
294/// Converts a list of `InstrumentStatus` into Arrow IPC bytes for Python.
295///
296/// # Errors
297///
298/// Returns a `PyErr` if encoding fails.
299#[pyfunction(name = "instrument_status_to_arrow_record_batch_bytes")]
300#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
301#[expect(clippy::needless_pass_by_value)]
302pub fn py_instrument_status_to_arrow_record_batch_bytes(
303    py: Python,
304    data: Vec<InstrumentStatus>,
305) -> PyResult<Py<PyBytes>> {
306    match instrument_status_to_arrow_record_batch_bytes(&data) {
307        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
308        Err(e) => Err(to_pyvalue_err(e)),
309    }
310}
311
312/// Converts a list of `OptionGreeks` into Arrow IPC bytes for Python.
313///
314/// # Errors
315///
316/// Returns a `PyErr` if encoding fails.
317#[pyfunction(name = "option_greeks_to_arrow_record_batch_bytes")]
318#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
319#[expect(clippy::needless_pass_by_value)]
320pub fn py_option_greeks_to_arrow_record_batch_bytes(
321    py: Python,
322    data: Vec<OptionGreeks>,
323) -> PyResult<Py<PyBytes>> {
324    match option_greeks_to_arrow_record_batch_bytes(&data) {
325        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
326        Err(e) => Err(to_pyvalue_err(e)),
327    }
328}
329
330/// Decodes Arrow IPC bytes into a list of `OptionGreeks`.
331///
332/// # Errors
333///
334/// Returns a `PyErr` if decoding fails.
335#[pyfunction(name = "option_greeks_from_arrow_record_batch_bytes")]
336#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
337pub fn py_option_greeks_from_arrow_record_batch_bytes(
338    _py: Python,
339    data: Vec<u8>,
340) -> PyResult<Vec<OptionGreeks>> {
341    let cursor = Cursor::new(data);
342    let reader = StreamReader::try_new(cursor, None).map_err(to_pyruntime_err)?;
343
344    let mut results = Vec::new();
345
346    for batch_result in reader {
347        let batch = batch_result.map_err(to_pyruntime_err)?;
348        let metadata = batch.schema().metadata().clone();
349        let decoded = OptionGreeks::decode_batch(&metadata, batch).map_err(to_pyvalue_err)?;
350        results.extend(decoded);
351    }
352
353    Ok(results)
354}
355
356/// Decodes Arrow IPC bytes into a list of `InstrumentStatus`.
357///
358/// # Errors
359///
360/// Returns a `PyErr` if decoding fails.
361#[pyfunction(name = "instrument_status_from_arrow_record_batch_bytes")]
362#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
363pub fn py_instrument_status_from_arrow_record_batch_bytes(
364    _py: Python,
365    data: Vec<u8>,
366) -> PyResult<Vec<InstrumentStatus>> {
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 =
376            InstrumentStatus::decode_typed_batch(&metadata, batch).map_err(to_pyvalue_err)?;
377        results.extend(decoded);
378    }
379
380    Ok(results)
381}
382
383/// Converts a list of `InstrumentClose` into Arrow IPC bytes for Python.
384///
385/// # Errors
386///
387/// Returns a `PyErr` if encoding fails.
388#[pyfunction(name = "instrument_closes_to_arrow_record_batch_bytes")]
389#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
390#[expect(clippy::needless_pass_by_value)]
391pub fn py_instrument_closes_to_arrow_record_batch_bytes(
392    py: Python,
393    data: Vec<InstrumentClose>,
394) -> PyResult<Py<PyBytes>> {
395    match instrument_closes_to_arrow_record_batch_bytes(&data) {
396        Ok(batch) => arrow_record_batch_to_pybytes(py, &batch),
397        Err(e) => Err(to_pyvalue_err(e)),
398    }
399}