nautilus_serialization/arrow/display/
close.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::InstrumentClose;
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 instrument_closes_schema() -> Schema {
34 Schema::new(vec![
35 utf8_field("instrument_id", false),
36 float64_field("close_price", false),
37 utf8_field("close_type", false),
38 timestamp_field("ts_event", false),
39 timestamp_field("ts_init", false),
40 ])
41}
42
43pub fn encode_instrument_closes(data: &[InstrumentClose]) -> Result<RecordBatch, ArrowError> {
58 let mut instrument_id_builder = StringBuilder::new();
59 let mut close_price_builder = Float64Builder::with_capacity(data.len());
60 let mut close_type_builder = StringBuilder::new();
61 let mut ts_event_builder =
62 TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
63 let mut ts_init_builder =
64 TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
65
66 for close in data {
67 instrument_id_builder.append_value(close.instrument_id.to_string());
68 close_price_builder.append_value(price_to_f64(&close.close_price));
69 close_type_builder.append_value(format!("{}", close.close_type));
70 ts_event_builder.append_value(unix_nanos_to_i64(close.ts_event.as_u64()));
71 ts_init_builder.append_value(unix_nanos_to_i64(close.ts_init.as_u64()));
72 }
73
74 RecordBatch::try_new(
75 Arc::new(instrument_closes_schema()),
76 vec![
77 Arc::new(instrument_id_builder.finish()),
78 Arc::new(close_price_builder.finish()),
79 Arc::new(close_type_builder.finish()),
80 Arc::new(ts_event_builder.finish()),
81 Arc::new(ts_init_builder.finish()),
82 ],
83 )
84}
85
86#[cfg(test)]
87mod tests {
88 use arrow::{
89 array::{Array, Float64Array, StringArray, TimestampNanosecondArray},
90 datatypes::{DataType, TimeUnit},
91 };
92 use nautilus_model::{enums::InstrumentCloseType, identifiers::InstrumentId, types::Price};
93 use rstest::rstest;
94
95 use super::*;
96
97 fn make_close(
98 instrument_id: &str,
99 price: &str,
100 close_type: InstrumentCloseType,
101 ts: u64,
102 ) -> InstrumentClose {
103 InstrumentClose {
104 instrument_id: InstrumentId::from(instrument_id),
105 close_price: Price::from(price),
106 close_type,
107 ts_event: ts.into(),
108 ts_init: (ts + 1).into(),
109 }
110 }
111
112 #[rstest]
113 fn test_encode_instrument_closes_schema() {
114 let batch = encode_instrument_closes(&[]).unwrap();
115 let fields = batch.schema().fields().clone();
116 assert_eq!(fields.len(), 5);
117 assert_eq!(fields[0].name(), "instrument_id");
118 assert_eq!(fields[0].data_type(), &DataType::Utf8);
119 assert_eq!(fields[1].name(), "close_price");
120 assert_eq!(fields[1].data_type(), &DataType::Float64);
121 assert_eq!(fields[2].name(), "close_type");
122 assert_eq!(fields[2].data_type(), &DataType::Utf8);
123 assert_eq!(fields[3].name(), "ts_event");
124 assert_eq!(
125 fields[3].data_type(),
126 &DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
127 );
128 assert_eq!(fields[4].name(), "ts_init");
129 }
130
131 #[rstest]
132 fn test_encode_instrument_closes_values() {
133 let closes = vec![
134 make_close(
135 "AAPL.XNAS",
136 "150.50",
137 InstrumentCloseType::EndOfSession,
138 1_000,
139 ),
140 make_close(
141 "AAPL.XNAS",
142 "151.25",
143 InstrumentCloseType::ContractExpired,
144 2_000,
145 ),
146 ];
147 let batch = encode_instrument_closes(&closes).unwrap();
148
149 assert_eq!(batch.num_rows(), 2);
150
151 let close_price_col = batch
152 .column(1)
153 .as_any()
154 .downcast_ref::<Float64Array>()
155 .unwrap();
156 let close_type_col = batch
157 .column(2)
158 .as_any()
159 .downcast_ref::<StringArray>()
160 .unwrap();
161 let ts_event_col = batch
162 .column(3)
163 .as_any()
164 .downcast_ref::<TimestampNanosecondArray>()
165 .unwrap();
166
167 assert!((close_price_col.value(0) - 150.50).abs() < 1e-9);
168 assert!((close_price_col.value(1) - 151.25).abs() < 1e-9);
169 assert_eq!(
170 close_type_col.value(0),
171 format!("{}", InstrumentCloseType::EndOfSession)
172 );
173 assert_eq!(
174 close_type_col.value(1),
175 format!("{}", InstrumentCloseType::ContractExpired)
176 );
177 assert_eq!(ts_event_col.value(0), 1_000);
178 }
179
180 #[rstest]
181 fn test_encode_instrument_closes_empty() {
182 let batch = encode_instrument_closes(&[]).unwrap();
183 assert_eq!(batch.num_rows(), 0);
184 }
185
186 #[rstest]
187 fn test_encode_instrument_closes_mixed_instruments() {
188 let closes = vec![
189 make_close("AAPL.XNAS", "150.50", InstrumentCloseType::EndOfSession, 1),
190 make_close("MSFT.XNAS", "300.00", InstrumentCloseType::EndOfSession, 2),
191 ];
192 let batch = encode_instrument_closes(&closes).unwrap();
193 let instrument_id_col = batch
194 .column(0)
195 .as_any()
196 .downcast_ref::<StringArray>()
197 .unwrap();
198 assert_eq!(instrument_id_col.value(0), "AAPL.XNAS");
199 assert_eq!(instrument_id_col.value(1), "MSFT.XNAS");
200 }
201}