nautilus_serialization/arrow/
instrument_status.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use 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_with_metadata_fields, encode_batch_with_identifier,
26        metadata_for_type, schema_for_type_with_identifier,
27    },
28};
29
30const INSTRUMENT_STATUS_FIELDS: &[JsonFieldSpec] = &[
31    JsonFieldSpec::enum_dictionary("action", false),
32    JsonFieldSpec::timestamp("ts_event", false),
33    JsonFieldSpec::timestamp("ts_init", false),
34    JsonFieldSpec::utf8("reason", true),
35    JsonFieldSpec::utf8("trading_event", true),
36    JsonFieldSpec::boolean("is_trading", true),
37    JsonFieldSpec::boolean("is_quoting", true),
38    JsonFieldSpec::boolean("is_short_sell_restricted", true),
39];
40
41impl ArrowSchemaProvider for InstrumentStatus {
42    fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
43        schema_for_type_with_identifier("InstrumentStatus", metadata, INSTRUMENT_STATUS_FIELDS)
44    }
45}
46
47impl EncodeToRecordBatch for InstrumentStatus {
48    fn encode_batch<T>(
49        metadata: &HashMap<String, String>,
50        data: &[T],
51    ) -> Result<RecordBatch, ArrowError>
52    where
53        T: std::borrow::Borrow<Self>,
54    {
55        encode_batch_with_identifier(
56            "InstrumentStatus",
57            metadata,
58            data.iter().map(std::borrow::Borrow::borrow),
59            INSTRUMENT_STATUS_FIELDS,
60            data.iter()
61                .map(std::borrow::Borrow::borrow)
62                .map(|status| status.instrument_id),
63        )
64    }
65
66    fn metadata(&self) -> HashMap<String, String> {
67        let mut metadata = metadata_for_type("InstrumentStatus");
68        metadata.insert(
69            KEY_INSTRUMENT_ID.to_string(),
70            self.instrument_id.to_string(),
71        );
72        metadata
73    }
74}
75
76impl DecodeFromRecordBatch for InstrumentStatus {
77    fn decode_batch(
78        metadata: &HashMap<String, String>,
79        record_batch: RecordBatch,
80    ) -> Result<Vec<Self>, EncodingError> {
81        decode_batch_with_metadata_fields(
82            metadata,
83            &record_batch,
84            INSTRUMENT_STATUS_FIELDS,
85            &[KEY_INSTRUMENT_ID],
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 schema = record_batch.schema();
141        let decoded: Vec<Data> =
142            InstrumentStatus::decode_data_batch(schema.metadata(), record_batch).unwrap();
143
144        assert_eq!(
145            schema
146                .fields()
147                .iter()
148                .map(|field| field.name().as_str())
149                .collect::<Vec<_>>(),
150            vec![
151                "action",
152                "ts_event",
153                "ts_init",
154                "reason",
155                "trading_event",
156                "is_trading",
157                "is_quoting",
158                "is_short_sell_restricted",
159                "identifier",
160            ]
161        );
162        assert_eq!(
163            schema.field(0).data_type(),
164            &crate::arrow::enum_dictionary_data_type()
165        );
166        assert_eq!(
167            schema.metadata(),
168            &HashMap::from([
169                ("type_name".to_string(), "InstrumentStatus".to_string()),
170                (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
171            ])
172        );
173
174        assert_eq!(decoded.len(), original.len());
175        for (orig, dec) in original.iter().zip(decoded.iter()) {
176            match dec {
177                Data::InstrumentStatus(s) => assert_eq!(s, orig),
178                other => panic!("expected Data::InstrumentStatus, was {other:?}"),
179            }
180        }
181    }
182}
nautilus_serialization/arrow/instrument_status.rs

nautilus_serialization/arrow/
instrument_status.rs
