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")]
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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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}