nautilus_serialization/arrow/
instrument_status.rs1use std::collections::HashMap;
17
18use arrow::{datatypes::Schema, error::ArrowError, record_batch::RecordBatch};
19use nautilus_model::data::{Data, InstrumentStatus};
20
21use super::{
22 ArrowSchemaProvider, DecodeDataFromRecordBatch, DecodeFromRecordBatch, EncodeToRecordBatch,
23 EncodingError, KEY_INSTRUMENT_ID,
24 json::{
25 JsonFieldSpec, decode_batch, encode_batch_with_identifier, metadata_for_type,
26 schema_for_type_with_identifier,
27 },
28};
29
30const INSTRUMENT_STATUS_FIELDS: &[JsonFieldSpec] = &[
31 JsonFieldSpec::utf8("instrument_id", false),
32 JsonFieldSpec::utf8("action", false),
33 JsonFieldSpec::timestamp("ts_event", false),
34 JsonFieldSpec::timestamp("ts_init", false),
35 JsonFieldSpec::utf8("reason", true),
36 JsonFieldSpec::utf8("trading_event", true),
37 JsonFieldSpec::boolean("is_trading", true),
38 JsonFieldSpec::boolean("is_quoting", true),
39 JsonFieldSpec::boolean("is_short_sell_restricted", true),
40];
41
42impl ArrowSchemaProvider for InstrumentStatus {
43 fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
44 schema_for_type_with_identifier("InstrumentStatus", metadata, INSTRUMENT_STATUS_FIELDS)
45 }
46}
47
48impl EncodeToRecordBatch for InstrumentStatus {
49 fn encode_batch<T>(
50 metadata: &HashMap<String, String>,
51 data: &[T],
52 ) -> Result<RecordBatch, ArrowError>
53 where
54 T: std::borrow::Borrow<Self>,
55 {
56 encode_batch_with_identifier(
57 "InstrumentStatus",
58 metadata,
59 data.iter().map(std::borrow::Borrow::borrow),
60 INSTRUMENT_STATUS_FIELDS,
61 data.iter()
62 .map(std::borrow::Borrow::borrow)
63 .map(|status| status.instrument_id),
64 )
65 }
66
67 fn metadata(&self) -> HashMap<String, String> {
68 let mut metadata = metadata_for_type("InstrumentStatus");
69 metadata.insert(
70 KEY_INSTRUMENT_ID.to_string(),
71 self.instrument_id.to_string(),
72 );
73 metadata
74 }
75}
76
77impl DecodeFromRecordBatch for InstrumentStatus {
78 fn decode_batch(
79 metadata: &HashMap<String, String>,
80 record_batch: RecordBatch,
81 ) -> Result<Vec<Self>, EncodingError> {
82 decode_batch(
83 metadata,
84 &record_batch,
85 INSTRUMENT_STATUS_FIELDS,
86 Some("InstrumentStatus"),
87 )
88 }
89}
90
91impl DecodeDataFromRecordBatch for InstrumentStatus {
92 fn decode_data_batch(
93 metadata: &HashMap<String, String>,
94 record_batch: RecordBatch,
95 ) -> Result<Vec<Data>, EncodingError> {
96 let items: Vec<Self> = Self::decode_batch(metadata, record_batch)?;
97 Ok(items.into_iter().map(Data::from).collect())
98 }
99}
100
101#[cfg(test)]
102mod tests {
103 use nautilus_model::{enums::MarketStatusAction, identifiers::InstrumentId};
104 use rstest::rstest;
105 use ustr::Ustr;
106
107 use super::*;
108
109 #[rstest]
110 fn test_encode_decode_round_trip() {
111 let instrument_id = InstrumentId::from("AAPL.XNAS");
112 let metadata = HashMap::from([(KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string())]);
113
114 let status1 = InstrumentStatus::new(
115 instrument_id,
116 MarketStatusAction::Trading,
117 1_000_000_000.into(),
118 1_000_000_001.into(),
119 Some(Ustr::from("Normal trading")),
120 Some(Ustr::from("MARKET_OPEN")),
121 Some(true),
122 Some(true),
123 Some(false),
124 );
125
126 let status2 = InstrumentStatus::new(
127 instrument_id,
128 MarketStatusAction::Halt,
129 2_000_000_000.into(),
130 2_000_000_001.into(),
131 None,
132 None,
133 None,
134 None,
135 None,
136 );
137
138 let original = vec![status1, status2];
139 let record_batch = InstrumentStatus::encode_batch(&metadata, &original).unwrap();
140 let decoded: Vec<Data> =
141 InstrumentStatus::decode_data_batch(&metadata, record_batch).unwrap();
142
143 assert_eq!(decoded.len(), original.len());
144 for (orig, dec) in original.iter().zip(decoded.iter()) {
145 match dec {
146 Data::InstrumentStatus(s) => assert_eq!(s, orig),
147 other => panic!("expected Data::InstrumentStatus, was {other:?}"),
148 }
149 }
150 }
151}