nautilus_serialization/arrow/display/
bar.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::Bar;
27
28use super::{
29 float64_field, price_to_f64, quantity_to_f64, timestamp_field, unix_nanos_to_i64, utf8_field,
30};
31use crate::arrow::timestamp_data_type;
32
33#[must_use]
35pub fn bars_schema() -> Schema {
36 Schema::new(vec![
37 utf8_field("instrument_id", false),
38 utf8_field("bar_type", false),
39 float64_field("open", false),
40 float64_field("high", false),
41 float64_field("low", false),
42 float64_field("close", false),
43 float64_field("volume", false),
44 timestamp_field("ts_event", false),
45 timestamp_field("ts_init", false),
46 ])
47}
48
49pub fn encode_bars(data: &[Bar]) -> Result<RecordBatch, ArrowError> {
63 let mut instrument_id_builder = StringBuilder::new();
64 let mut bar_type_builder = StringBuilder::new();
65 let mut open_builder = Float64Builder::with_capacity(data.len());
66 let mut high_builder = Float64Builder::with_capacity(data.len());
67 let mut low_builder = Float64Builder::with_capacity(data.len());
68 let mut close_builder = Float64Builder::with_capacity(data.len());
69 let mut volume_builder = Float64Builder::with_capacity(data.len());
70 let mut ts_event_builder =
71 TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
72 let mut ts_init_builder =
73 TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
74
75 for bar in data {
76 instrument_id_builder.append_value(bar.instrument_id().to_string());
77 bar_type_builder.append_value(bar.bar_type.to_string());
78 open_builder.append_value(price_to_f64(&bar.open));
79 high_builder.append_value(price_to_f64(&bar.high));
80 low_builder.append_value(price_to_f64(&bar.low));
81 close_builder.append_value(price_to_f64(&bar.close));
82 volume_builder.append_value(quantity_to_f64(&bar.volume));
83 ts_event_builder.append_value(unix_nanos_to_i64(bar.ts_event.as_u64()));
84 ts_init_builder.append_value(unix_nanos_to_i64(bar.ts_init.as_u64()));
85 }
86
87 RecordBatch::try_new(
88 Arc::new(bars_schema()),
89 vec![
90 Arc::new(instrument_id_builder.finish()),
91 Arc::new(bar_type_builder.finish()),
92 Arc::new(open_builder.finish()),
93 Arc::new(high_builder.finish()),
94 Arc::new(low_builder.finish()),
95 Arc::new(close_builder.finish()),
96 Arc::new(volume_builder.finish()),
97 Arc::new(ts_event_builder.finish()),
98 Arc::new(ts_init_builder.finish()),
99 ],
100 )
101}
102
103#[cfg(test)]
104mod tests {
105 use std::str::FromStr;
106
107 use arrow::{
108 array::{Array, Float64Array, StringArray, TimestampNanosecondArray},
109 datatypes::{DataType, TimeUnit},
110 };
111 use nautilus_model::{
112 data::BarType,
113 types::{Price, Quantity},
114 };
115 use rstest::rstest;
116
117 use super::*;
118
119 fn make_bar(
120 bar_type_str: &str,
121 open: &str,
122 high: &str,
123 low: &str,
124 close: &str,
125 ts: u64,
126 ) -> Bar {
127 let bar_type = BarType::from_str(bar_type_str).unwrap();
128 Bar::new(
129 bar_type,
130 Price::from(open),
131 Price::from(high),
132 Price::from(low),
133 Price::from(close),
134 Quantity::from(1_100),
135 ts.into(),
136 (ts + 1).into(),
137 )
138 }
139
140 #[rstest]
141 fn test_encode_bars_schema() {
142 let batch = encode_bars(&[]).unwrap();
143 let fields = batch.schema().fields().clone();
144 assert_eq!(fields.len(), 9);
145 assert_eq!(fields[0].name(), "instrument_id");
146 assert_eq!(fields[0].data_type(), &DataType::Utf8);
147 assert_eq!(fields[1].name(), "bar_type");
148 assert_eq!(fields[1].data_type(), &DataType::Utf8);
149 assert_eq!(fields[2].name(), "open");
150 assert_eq!(fields[2].data_type(), &DataType::Float64);
151 assert_eq!(fields[3].name(), "high");
152 assert_eq!(fields[4].name(), "low");
153 assert_eq!(fields[5].name(), "close");
154 assert_eq!(fields[6].name(), "volume");
155 assert_eq!(fields[6].data_type(), &DataType::Float64);
156 assert_eq!(fields[7].name(), "ts_event");
157 assert_eq!(
158 fields[7].data_type(),
159 &DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
160 );
161 assert_eq!(fields[8].name(), "ts_init");
162 }
163
164 #[rstest]
165 fn test_encode_bars_values() {
166 let bars = vec![
167 make_bar(
168 "AAPL.XNAS-1-MINUTE-LAST-INTERNAL",
169 "100.10",
170 "102.00",
171 "100.00",
172 "101.00",
173 1_000,
174 ),
175 make_bar(
176 "AAPL.XNAS-1-MINUTE-LAST-INTERNAL",
177 "100.20",
178 "102.00",
179 "100.00",
180 "101.00",
181 2_000,
182 ),
183 ];
184 let batch = encode_bars(&bars).unwrap();
185
186 assert_eq!(batch.num_rows(), 2);
187
188 let instrument_id_col = batch
189 .column(0)
190 .as_any()
191 .downcast_ref::<StringArray>()
192 .unwrap();
193 let bar_type_col = batch
194 .column(1)
195 .as_any()
196 .downcast_ref::<StringArray>()
197 .unwrap();
198 let open_col = batch
199 .column(2)
200 .as_any()
201 .downcast_ref::<Float64Array>()
202 .unwrap();
203 let high_col = batch
204 .column(3)
205 .as_any()
206 .downcast_ref::<Float64Array>()
207 .unwrap();
208 let low_col = batch
209 .column(4)
210 .as_any()
211 .downcast_ref::<Float64Array>()
212 .unwrap();
213 let close_col = batch
214 .column(5)
215 .as_any()
216 .downcast_ref::<Float64Array>()
217 .unwrap();
218 let volume_col = batch
219 .column(6)
220 .as_any()
221 .downcast_ref::<Float64Array>()
222 .unwrap();
223 let ts_event_col = batch
224 .column(7)
225 .as_any()
226 .downcast_ref::<TimestampNanosecondArray>()
227 .unwrap();
228 let ts_init_col = batch
229 .column(8)
230 .as_any()
231 .downcast_ref::<TimestampNanosecondArray>()
232 .unwrap();
233
234 assert_eq!(instrument_id_col.value(0), "AAPL.XNAS");
235 assert_eq!(bar_type_col.value(0), "AAPL.XNAS-1-MINUTE-LAST-INTERNAL");
236 assert!((open_col.value(0) - 100.10).abs() < 1e-9);
237 assert!((open_col.value(1) - 100.20).abs() < 1e-9);
238 assert!((high_col.value(0) - 102.00).abs() < 1e-9);
239 assert!((low_col.value(0) - 100.00).abs() < 1e-9);
240 assert!((close_col.value(0) - 101.00).abs() < 1e-9);
241 assert!((volume_col.value(0) - 1_100.0).abs() < 1e-9);
242 assert_eq!(ts_event_col.value(0), 1_000);
243 assert_eq!(ts_init_col.value(1), 2_001);
244 }
245
246 #[rstest]
247 fn test_encode_bars_empty() {
248 let batch = encode_bars(&[]).unwrap();
249 assert_eq!(batch.num_rows(), 0);
250 }
251
252 #[rstest]
253 fn test_encode_bars_mixed_instruments() {
254 let bars = vec![
255 make_bar(
256 "AAPL.XNAS-1-MINUTE-LAST-INTERNAL",
257 "100.10",
258 "102.00",
259 "100.00",
260 "101.00",
261 1,
262 ),
263 make_bar(
264 "MSFT.XNAS-1-MINUTE-LAST-INTERNAL",
265 "250.00",
266 "251.00",
267 "249.00",
268 "250.50",
269 2,
270 ),
271 ];
272 let batch = encode_bars(&bars).unwrap();
273 let instrument_id_col = batch
274 .column(0)
275 .as_any()
276 .downcast_ref::<StringArray>()
277 .unwrap();
278 assert_eq!(instrument_id_col.value(0), "AAPL.XNAS");
279 assert_eq!(instrument_id_col.value(1), "MSFT.XNAS");
280 }
281}