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

nautilus_serialization/arrow/
close.rs
