Skip to main content

nautilus_serialization/arrow/display/
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
16//! Display-mode Arrow encoder for [`InstrumentClose`].
17
18use 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/// Returns the display-mode Arrow schema for [`InstrumentClose`].
32#[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
43/// Encodes instrument closes as a display-friendly Arrow [`RecordBatch`].
44///
45/// Emits a `Float64` `close_price` column, `Utf8` columns for the instrument
46/// ID and close type, and `Timestamp(Nanosecond)` columns for event and init
47/// times. Mixed-instrument batches are supported. Precision is lost on the
48/// conversion to `f64`; use
49/// [`crate::arrow::instrument_closes_to_arrow_record_batch_bytes`] for catalog
50/// storage.
51///
52/// Returns an empty [`RecordBatch`] with the correct schema when `data` is empty.
53///
54/// # Errors
55///
56/// Returns an [`ArrowError`] if the Arrow `RecordBatch` cannot be constructed.
57pub 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}