nautilus_serialization/arrow/
close.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, str::FromStr, sync::Arc};
17
18use arrow::{
19    array::{Decimal128Array, UInt64Array},
20    datatypes::{DataType, Field, Schema},
21    error::ArrowError,
22    record_batch::RecordBatch,
23};
24use nautilus_model::{
25    data::close::InstrumentClose, enums::InstrumentCloseType, identifiers::InstrumentId,
26};
27
28use super::{
29    DecodeDataFromRecordBatch, EncodingError, KEY_IDENTIFIER, KEY_INSTRUMENT_ID,
30    KEY_PRICE_PRECISION, decode_decimal_price, decode_required_timestamp, enum_dictionary_array,
31    enum_dictionary_data_type, extract_column, extract_column_string, fixed_decimal_data_type,
32    identifier_array_from_display, metadata_with_type_name, price_decimal_array,
33};
34use crate::arrow::{ArrowSchemaProvider, Data, DecodeFromRecordBatch, EncodeToRecordBatch};
35
36impl ArrowSchemaProvider for InstrumentClose {
37    fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
38        let fields = vec![
39            Field::new("close_price", fixed_decimal_data_type(), true),
40            Field::new("close_type", enum_dictionary_data_type(), false),
41            Field::new("ts_event", crate::arrow::timestamp_data_type(), false),
42            Field::new("ts_init", crate::arrow::timestamp_data_type(), false),
43            Field::new(KEY_IDENTIFIER, DataType::Utf8, true),
44        ];
45
46        Schema::new_with_metadata(fields, metadata_with_type_name("InstrumentClose", metadata))
47    }
48}
49
50fn parse_metadata(metadata: &HashMap<String, String>) -> Result<(InstrumentId, u8), EncodingError> {
51    let instrument_id_str = metadata
52        .get(KEY_INSTRUMENT_ID)
53        .ok_or_else(|| EncodingError::MissingMetadata(KEY_INSTRUMENT_ID))?;
54    let instrument_id = InstrumentId::from_str(instrument_id_str)
55        .map_err(|e| EncodingError::ParseError(KEY_INSTRUMENT_ID, e.to_string()))?;
56
57    let price_precision = metadata
58        .get(KEY_PRICE_PRECISION)
59        .ok_or_else(|| EncodingError::MissingMetadata(KEY_PRICE_PRECISION))?
60        .parse::<u8>()
61        .map_err(|e| EncodingError::ParseError(KEY_PRICE_PRECISION, e.to_string()))?;
62
63    Ok((instrument_id, price_precision))
64}
65
66impl EncodeToRecordBatch for InstrumentClose {
67    fn encode_batch<T>(
68        metadata: &HashMap<String, String>,
69        data: &[T],
70    ) -> Result<RecordBatch, ArrowError>
71    where
72        T: std::borrow::Borrow<Self>,
73    {
74        let mut ts_event_builder = UInt64Array::builder(data.len());
75        let mut ts_init_builder = UInt64Array::builder(data.len());
76
77        for item in data.iter().map(std::borrow::Borrow::borrow) {
78            ts_event_builder.append_value(item.ts_event.as_u64());
79            ts_init_builder.append_value(item.ts_init.as_u64());
80        }
81
82        crate::arrow::record_batch_with_timestamps(
83            Self::get_schema(Some(metadata.clone())).into(),
84            vec![
85                Arc::new(price_decimal_array(
86                    data.iter()
87                        .map(std::borrow::Borrow::borrow)
88                        .map(|item| item.close_price.raw()),
89                    "close_price",
90                )?),
91                Arc::new(enum_dictionary_array(
92                    data.iter()
93                        .map(std::borrow::Borrow::borrow)
94                        .map(|item| item.close_type),
95                )?),
96                Arc::new(ts_event_builder.finish()),
97                Arc::new(ts_init_builder.finish()),
98                Arc::new(identifier_array_from_display(
99                    data.iter()
100                        .map(std::borrow::Borrow::borrow)
101                        .map(|item| item.instrument_id),
102                )),
103            ],
104        )
105    }
106
107    fn metadata(&self) -> HashMap<String, String> {
108        Self::get_metadata(&self.instrument_id, self.close_price.precision)
109    }
110}
111
112impl DecodeFromRecordBatch for InstrumentClose {
113    fn decode_batch(
114        metadata: &HashMap<String, String>,
115        record_batch: RecordBatch,
116    ) -> Result<Vec<Self>, EncodingError> {
117        let (instrument_id, price_precision) = parse_metadata(metadata)?;
118        let record_batch = crate::arrow::record_batch_with_u64_timestamps(&record_batch)?;
119        let record_batch = &record_batch;
120        let cols = record_batch.columns();
121
122        let close_price_values =
123            extract_column::<Decimal128Array>(cols, "close_price", 0, fixed_decimal_data_type())?;
124        let close_type_values = extract_column_string(cols, "close_type", 1)?;
125        let ts_event_values = extract_column::<UInt64Array>(cols, "ts_event", 2, DataType::UInt64)?;
126        let ts_init_values = extract_column::<UInt64Array>(cols, "ts_init", 3, DataType::UInt64)?;
127
128        let result: Result<Vec<Self>, EncodingError> = (0..record_batch.num_rows())
129            .map(|row| {
130                let close_price =
131                    decode_decimal_price(close_price_values, price_precision, "close_price", row)?;
132                let close_type_value = close_type_values.value(row);
133                let close_type = InstrumentCloseType::from_str(close_type_value).map_err(|e| {
134                    EncodingError::ParseError(stringify!(InstrumentCloseType), e.to_string())
135                })?;
136                Ok(Self {
137                    instrument_id,
138                    close_price,
139                    close_type,
140                    ts_event: decode_required_timestamp(ts_event_values, "ts_event", row)?,
141                    ts_init: decode_required_timestamp(ts_init_values, "ts_init", row)?,
142                })
143            })
144            .collect();
145
146        result
147    }
148}
149
150impl DecodeDataFromRecordBatch for InstrumentClose {
151    fn decode_data_batch(
152        metadata: &HashMap<String, String>,
153        record_batch: RecordBatch,
154    ) -> Result<Vec<Data>, EncodingError> {
155        let items: Vec<Self> = Self::decode_batch(metadata, record_batch)?;
156        Ok(items.into_iter().map(Data::from).collect())
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use std::sync::Arc;
163
164    use arrow::array::{Array, TimestampNanosecondArray};
165    use nautilus_model::types::{Price, fixed::FIXED_SCALAR, price::PriceRaw};
166    use rstest::rstest;
167
168    use super::*;
169    use crate::arrow::{KEY_TYPE_NAME, get_raw_price};
170
171    #[rstest]
172    fn test_get_schema() {
173        let instrument_id = InstrumentId::from("AAPL.XNAS");
174        let metadata = HashMap::from([
175            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
176            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
177        ]);
178        let schema = InstrumentClose::get_schema(Some(metadata.clone()));
179
180        let expected_fields = vec![
181            Field::new("close_price", fixed_decimal_data_type(), true),
182            Field::new("close_type", enum_dictionary_data_type(), false),
183            Field::new("ts_event", crate::arrow::timestamp_data_type(), false),
184            Field::new("ts_init", crate::arrow::timestamp_data_type(), false),
185            Field::new(KEY_IDENTIFIER, DataType::Utf8, true),
186        ];
187
188        let mut expected_metadata = metadata;
189        expected_metadata.insert(KEY_TYPE_NAME.to_string(), "InstrumentClose".to_string());
190        let expected_schema = Schema::new_with_metadata(expected_fields, expected_metadata);
191        assert_eq!(schema, expected_schema);
192    }
193
194    #[rstest]
195    fn test_get_schema_map() {
196        let schema_map = InstrumentClose::get_schema_map();
197        let mut expected_map = HashMap::new();
198
199        let fixed_size_binary = "Decimal128(38, 16)".to_string();
200        expected_map.insert("close_price".to_string(), fixed_size_binary);
201        expected_map.insert(
202            "close_type".to_string(),
203            "Dictionary(Int8, Utf8)".to_string(),
204        );
205        expected_map.insert(
206            "ts_event".to_string(),
207            "Timestamp(Nanosecond, Some(\"UTC\"))".to_string(),
208        );
209        expected_map.insert(
210            "ts_init".to_string(),
211            "Timestamp(Nanosecond, Some(\"UTC\"))".to_string(),
212        );
213        expected_map.insert(KEY_IDENTIFIER.to_string(), "Utf8".to_string());
214        assert_eq!(schema_map, expected_map);
215    }
216
217    #[rstest]
218    fn test_encode_batch() {
219        let instrument_id = InstrumentId::from("AAPL.XNAS");
220        let metadata = HashMap::from([
221            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
222            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
223        ]);
224
225        let close1 = InstrumentClose {
226            instrument_id,
227            close_price: Price::from("150.50"),
228            close_type: InstrumentCloseType::EndOfSession,
229            ts_event: 1.into(),
230            ts_init: 3.into(),
231        };
232
233        let close2 = InstrumentClose {
234            instrument_id,
235            close_price: Price::from("151.25"),
236            close_type: InstrumentCloseType::ContractExpired,
237            ts_event: 2.into(),
238            ts_init: 4.into(),
239        };
240
241        let data = vec![close1, close2];
242        let record_batch = InstrumentClose::encode_batch(&metadata, &data).unwrap();
243
244        let columns = record_batch.columns();
245        let close_price_values = columns[0]
246            .as_any()
247            .downcast_ref::<Decimal128Array>()
248            .unwrap();
249        let close_type_values = extract_column_string(columns, "close_type", 1).unwrap();
250        let ts_event_values = columns[2]
251            .as_any()
252            .downcast_ref::<TimestampNanosecondArray>()
253            .unwrap();
254        let ts_init_values = columns[3]
255            .as_any()
256            .downcast_ref::<TimestampNanosecondArray>()
257            .unwrap();
258
259        assert_eq!(columns.len(), 5);
260        assert_eq!(close_price_values.len(), 2);
261        assert_eq!(
262            get_raw_price(close_price_values.value(0)),
263            (150.50 * FIXED_SCALAR) as PriceRaw
264        );
265        assert_eq!(
266            get_raw_price(close_price_values.value(1)),
267            (151.25 * FIXED_SCALAR) as PriceRaw
268        );
269        assert_eq!(close_type_values.len(), 2);
270        assert_eq!(close_type_values.value(0), "END_OF_SESSION");
271        assert_eq!(close_type_values.value(1), "CONTRACT_EXPIRED");
272        assert_eq!(ts_event_values.len(), 2);
273        assert_eq!(ts_event_values.value(0), 1);
274        assert_eq!(ts_event_values.value(1), 2);
275        assert_eq!(ts_init_values.len(), 2);
276        assert_eq!(ts_init_values.value(0), 3);
277        assert_eq!(ts_init_values.value(1), 4);
278    }
279
280    #[rstest]
281    fn test_decode_batch() {
282        let instrument_id = InstrumentId::from("AAPL.XNAS");
283        let metadata = HashMap::from([
284            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
285            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
286        ]);
287
288        let raw_price1 = (150.50 * FIXED_SCALAR) as PriceRaw;
289        let raw_price2 = (151.25 * FIXED_SCALAR) as PriceRaw;
290        let close_price = crate::arrow::test_support::decimal_array_from_bytes(vec![
291            &raw_price1.to_le_bytes(),
292            &raw_price2.to_le_bytes(),
293        ]);
294        let close_type = enum_dictionary_array([
295            InstrumentCloseType::EndOfSession,
296            InstrumentCloseType::ContractExpired,
297        ])
298        .unwrap();
299        let ts_event = UInt64Array::from(vec![1, 2]);
300        let ts_init = UInt64Array::from(vec![3, 4]);
301
302        let record_batch = crate::arrow::record_batch_with_timestamps(
303            crate::arrow::schema_without_identifier_column(&InstrumentClose::get_schema(Some(
304                metadata.clone(),
305            )))
306            .into(),
307            vec![
308                Arc::new(close_price),
309                Arc::new(close_type),
310                Arc::new(ts_event),
311                Arc::new(ts_init),
312            ],
313        )
314        .unwrap();
315
316        let decoded_data = InstrumentClose::decode_batch(&metadata, record_batch).unwrap();
317
318        assert_eq!(decoded_data.len(), 2);
319        assert_eq!(decoded_data[0].instrument_id, instrument_id);
320        assert_eq!(decoded_data[0].close_price, Price::from_raw(raw_price1, 2));
321        assert_eq!(
322            decoded_data[0].close_type,
323            InstrumentCloseType::EndOfSession
324        );
325        assert_eq!(decoded_data[0].ts_event.as_u64(), 1);
326        assert_eq!(decoded_data[0].ts_init.as_u64(), 3);
327
328        assert_eq!(decoded_data[1].instrument_id, instrument_id);
329        assert_eq!(decoded_data[1].close_price, Price::from_raw(raw_price2, 2));
330        assert_eq!(
331            decoded_data[1].close_type,
332            InstrumentCloseType::ContractExpired
333        );
334        assert_eq!(decoded_data[1].ts_event.as_u64(), 2);
335        assert_eq!(decoded_data[1].ts_init.as_u64(), 4);
336    }
337
338    #[rstest]
339    fn test_decode_batch_rejects_null_timestamp_with_field_and_row() {
340        let instrument_id = InstrumentId::from("AAPL.XNAS");
341        let metadata = HashMap::from([
342            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
343            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
344        ]);
345        let close = InstrumentClose {
346            instrument_id,
347            close_price: Price::from("150.50"),
348            close_type: InstrumentCloseType::EndOfSession,
349            ts_event: 1.into(),
350            ts_init: 2.into(),
351        };
352        let encoded = InstrumentClose::encode_batch(&metadata, &[close]).unwrap();
353        let mut columns = encoded.columns().to_vec();
354        columns[3] = Arc::new(TimestampNanosecondArray::from(vec![None]).with_timezone("UTC"));
355        let fields = encoded
356            .schema()
357            .fields()
358            .iter()
359            .map(|field| {
360                if field.name() == "ts_init" {
361                    Arc::new(field.as_ref().clone().with_nullable(true))
362                } else {
363                    Arc::clone(field)
364                }
365            })
366            .collect::<Vec<_>>();
367        let schema = Arc::new(Schema::new_with_metadata(fields, metadata.clone()));
368        let batch = RecordBatch::try_new(schema, columns).unwrap();
369
370        let error = InstrumentClose::decode_batch(&metadata, batch).unwrap_err();
371
372        assert!(error.to_string().contains("ts_init"));
373        assert!(error.to_string().contains("row 0"));
374    }
375
376    #[rstest]
377    fn test_decode_batch_invalid_close_price_returns_error() {
378        let instrument_id = InstrumentId::from("AAPL.XNAS");
379        let metadata = HashMap::from([
380            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
381            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
382        ]);
383
384        let invalid_price: PriceRaw = PriceRaw::MAX - 1000;
385        let close_price = crate::arrow::test_support::decimal_array_from_bytes(vec![
386            &invalid_price.to_le_bytes(),
387        ]);
388        let close_type = enum_dictionary_array([InstrumentCloseType::EndOfSession]).unwrap();
389        let ts_event = UInt64Array::from(vec![1]);
390        let ts_init = UInt64Array::from(vec![2]);
391
392        let record_batch = crate::arrow::record_batch_with_timestamps(
393            crate::arrow::schema_without_identifier_column(&InstrumentClose::get_schema(Some(
394                metadata.clone(),
395            )))
396            .into(),
397            vec![
398                Arc::new(close_price),
399                Arc::new(close_type),
400                Arc::new(ts_event),
401                Arc::new(ts_init),
402            ],
403        )
404        .unwrap();
405
406        let result = InstrumentClose::decode_batch(&metadata, record_batch);
407        assert!(result.is_err());
408        let err = result.unwrap_err();
409        assert!(
410            err.to_string().contains("close_price") && err.to_string().contains("row 0"),
411            "Expected close_price error at row 0, was: {err}"
412        );
413    }
414
415    #[rstest]
416    fn test_decode_batch_invalid_close_type_returns_error() {
417        let instrument_id = InstrumentId::from("AAPL.XNAS");
418        let metadata = HashMap::from([
419            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
420            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
421        ]);
422
423        let raw_price = (150.50 * FIXED_SCALAR) as PriceRaw;
424        let close_price =
425            crate::arrow::test_support::decimal_array_from_bytes(vec![&raw_price.to_le_bytes()]);
426        let close_type = enum_dictionary_array(["INVALID"]).unwrap();
427        let ts_event = UInt64Array::from(vec![1]);
428        let ts_init = UInt64Array::from(vec![2]);
429
430        let record_batch = crate::arrow::record_batch_with_timestamps(
431            crate::arrow::schema_without_identifier_column(&InstrumentClose::get_schema(Some(
432                metadata.clone(),
433            )))
434            .into(),
435            vec![
436                Arc::new(close_price),
437                Arc::new(close_type),
438                Arc::new(ts_event),
439                Arc::new(ts_init),
440            ],
441        )
442        .unwrap();
443
444        let result = InstrumentClose::decode_batch(&metadata, record_batch);
445        assert!(result.is_err());
446        let err = result.unwrap_err();
447        assert!(
448            err.to_string().contains("InstrumentCloseType"),
449            "Expected InstrumentCloseType error, was: {err}"
450        );
451    }
452
453    #[rstest]
454    fn test_decode_batch_missing_instrument_id_returns_error() {
455        let instrument_id = InstrumentId::from("AAPL.XNAS");
456        let mut metadata = HashMap::from([
457            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
458            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
459        ]);
460
461        let raw_price = (150.50 * FIXED_SCALAR) as PriceRaw;
462        let close_price =
463            crate::arrow::test_support::decimal_array_from_bytes(vec![&raw_price.to_le_bytes()]);
464        let close_type = enum_dictionary_array([InstrumentCloseType::EndOfSession]).unwrap();
465        let ts_event = UInt64Array::from(vec![1]);
466        let ts_init = UInt64Array::from(vec![2]);
467
468        let record_batch = crate::arrow::record_batch_with_timestamps(
469            crate::arrow::schema_without_identifier_column(&InstrumentClose::get_schema(Some(
470                metadata.clone(),
471            )))
472            .into(),
473            vec![
474                Arc::new(close_price),
475                Arc::new(close_type),
476                Arc::new(ts_event),
477                Arc::new(ts_init),
478            ],
479        )
480        .unwrap();
481
482        metadata.remove(KEY_INSTRUMENT_ID);
483
484        let result = InstrumentClose::decode_batch(&metadata, record_batch);
485        assert!(result.is_err());
486        let err = result.unwrap_err();
487        assert!(
488            err.to_string().contains("instrument_id"),
489            "Expected missing instrument_id error, was: {err}"
490        );
491    }
492
493    #[rstest]
494    fn test_encode_decode_round_trip() {
495        let instrument_id = InstrumentId::from("AAPL.XNAS");
496        let metadata = HashMap::from([
497            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
498            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
499        ]);
500
501        let close1 = InstrumentClose {
502            instrument_id,
503            close_price: Price::from("150.50"),
504            close_type: InstrumentCloseType::EndOfSession,
505            ts_event: 1_000_000_000.into(),
506            ts_init: 1_000_000_001.into(),
507        };
508
509        let close2 = InstrumentClose {
510            instrument_id,
511            close_price: Price::from("151.25"),
512            close_type: InstrumentCloseType::ContractExpired,
513            ts_event: 2_000_000_000.into(),
514            ts_init: 2_000_000_001.into(),
515        };
516
517        let original = vec![close1, close2];
518        let record_batch = InstrumentClose::encode_batch(&metadata, &original).unwrap();
519        let decoded = InstrumentClose::decode_batch(&metadata, record_batch).unwrap();
520
521        assert_eq!(decoded.len(), original.len());
522        for (orig, dec) in original.iter().zip(decoded.iter()) {
523            assert_eq!(dec.instrument_id, orig.instrument_id);
524            assert_eq!(dec.close_price, orig.close_price);
525            assert_eq!(dec.close_type, orig.close_type);
526            assert_eq!(dec.ts_event, orig.ts_event);
527            assert_eq!(dec.ts_init, orig.ts_init);
528        }
529    }
530}
nautilus_serialization/arrow/close.rs

nautilus_serialization/arrow/
close.rs
