nautilus_serialization/arrow/display/
mark_price.rs1use std::sync::Arc;
19
20use arrow::{
21 array::{Float64Builder, StringBuilder, TimestampNanosecondBuilder},
22 datatypes::Schema,
23 error::ArrowError,
24 record_batch::RecordBatch,
25};
26use nautilus_model::data::MarkPriceUpdate;
27
28use super::{float64_field, price_to_f64, timestamp_field, unix_nanos_to_i64, utf8_field};
29use crate::arrow::timestamp_data_type;
30
31#[must_use]
33pub fn mark_prices_schema() -> Schema {
34 Schema::new(vec![
35 utf8_field("instrument_id", false),
36 float64_field("value", false),
37 timestamp_field("ts_event", false),
38 timestamp_field("ts_init", false),
39 ])
40}
41
42pub fn encode_mark_prices(data: &[MarkPriceUpdate]) -> Result<RecordBatch, ArrowError> {
55 let mut instrument_id_builder = StringBuilder::new();
56 let mut value_builder = Float64Builder::with_capacity(data.len());
57 let mut ts_event_builder =
58 TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
59 let mut ts_init_builder =
60 TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
61
62 for update in data {
63 instrument_id_builder.append_value(update.instrument_id.to_string());
64 value_builder.append_value(price_to_f64(&update.value));
65 ts_event_builder.append_value(unix_nanos_to_i64(update.ts_event.as_u64()));
66 ts_init_builder.append_value(unix_nanos_to_i64(update.ts_init.as_u64()));
67 }
68
69 RecordBatch::try_new(
70 Arc::new(mark_prices_schema()),
71 vec![
72 Arc::new(instrument_id_builder.finish()),
73 Arc::new(value_builder.finish()),
74 Arc::new(ts_event_builder.finish()),
75 Arc::new(ts_init_builder.finish()),
76 ],
77 )
78}
79
80#[cfg(test)]
81mod tests {
82 use arrow::{
83 array::{Array, Float64Array, StringArray, TimestampNanosecondArray},
84 datatypes::{DataType, TimeUnit},
85 };
86 use nautilus_model::{identifiers::InstrumentId, types::Price};
87 use rstest::rstest;
88
89 use super::*;
90
91 fn make_update(instrument_id: &str, value: &str, ts: u64) -> MarkPriceUpdate {
92 MarkPriceUpdate {
93 instrument_id: InstrumentId::from(instrument_id),
94 value: Price::from(value),
95 ts_event: ts.into(),
96 ts_init: (ts + 1).into(),
97 }
98 }
99
100 #[rstest]
101 fn test_encode_mark_prices_schema() {
102 let batch = encode_mark_prices(&[]).unwrap();
103 let fields = batch.schema().fields().clone();
104 assert_eq!(fields.len(), 4);
105 assert_eq!(fields[0].name(), "instrument_id");
106 assert_eq!(fields[0].data_type(), &DataType::Utf8);
107 assert_eq!(fields[1].name(), "value");
108 assert_eq!(fields[1].data_type(), &DataType::Float64);
109 assert_eq!(fields[2].name(), "ts_event");
110 assert_eq!(
111 fields[2].data_type(),
112 &DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
113 );
114 assert_eq!(fields[3].name(), "ts_init");
115 }
116
117 #[rstest]
118 fn test_encode_mark_prices_values() {
119 let updates = vec![
120 make_update("BTC-USDT.BINANCE", "50200.00", 1_000),
121 make_update("BTC-USDT.BINANCE", "50300.00", 2_000),
122 ];
123 let batch = encode_mark_prices(&updates).unwrap();
124
125 assert_eq!(batch.num_rows(), 2);
126
127 let instrument_id_col = batch
128 .column(0)
129 .as_any()
130 .downcast_ref::<StringArray>()
131 .unwrap();
132 let value_col = batch
133 .column(1)
134 .as_any()
135 .downcast_ref::<Float64Array>()
136 .unwrap();
137 let ts_event_col = batch
138 .column(2)
139 .as_any()
140 .downcast_ref::<TimestampNanosecondArray>()
141 .unwrap();
142
143 assert_eq!(instrument_id_col.value(0), "BTC-USDT.BINANCE");
144 assert!((value_col.value(0) - 50_200.00).abs() < 1e-9);
145 assert!((value_col.value(1) - 50_300.00).abs() < 1e-9);
146 assert_eq!(ts_event_col.value(0), 1_000);
147 }
148
149 #[rstest]
150 fn test_encode_mark_prices_empty() {
151 let batch = encode_mark_prices(&[]).unwrap();
152 assert_eq!(batch.num_rows(), 0);
153 }
154
155 #[rstest]
156 fn test_encode_mark_prices_mixed_instruments() {
157 let updates = vec![
158 make_update("BTC-USDT.BINANCE", "50200.00", 1),
159 make_update("ETH-USDT.BINANCE", "2500.00", 2),
160 ];
161 let batch = encode_mark_prices(&updates).unwrap();
162 let instrument_id_col = batch
163 .column(0)
164 .as_any()
165 .downcast_ref::<StringArray>()
166 .unwrap();
167 assert_eq!(instrument_id_col.value(0), "BTC-USDT.BINANCE");
168 assert_eq!(instrument_id_col.value(1), "ETH-USDT.BINANCE");
169 }
170}