nautilus_serialization/python/
arrow.rs1use 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
49pub fn arrow_record_batch_to_pybytes(py: Python, batch: &RecordBatch) -> PyResult<Py<PyBytes>> {
55 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#[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#[pyfunction]
104#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.serialization")]
105#[expect(clippy::missing_panics_doc)] pub 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() .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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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}