Skip to main content

nautilus_serialization/arrow/display/
bar.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 [`Bar`].
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::Bar;
27
28use super::{
29    float64_field, price_to_f64, quantity_to_f64, timestamp_field, unix_nanos_to_i64, utf8_field,
30};
31use crate::arrow::timestamp_data_type;
32
33/// Returns the display-mode Arrow schema for [`Bar`].
34#[must_use]
35pub fn bars_schema() -> Schema {
36    Schema::new(vec![
37        utf8_field("instrument_id", false),
38        utf8_field("bar_type", false),
39        float64_field("open", false),
40        float64_field("high", false),
41        float64_field("low", false),
42        float64_field("close", false),
43        float64_field("volume", false),
44        timestamp_field("ts_event", false),
45        timestamp_field("ts_init", false),
46    ])
47}
48
49/// Encodes bars as a display-friendly Arrow [`RecordBatch`].
50///
51/// Emits `Float64` columns for OHLCV values, `Utf8` columns for the
52/// instrument ID and bar type, and `Timestamp(Nanosecond)` columns for
53/// event and init times. Mixed-instrument batches are supported. Precision
54/// is lost on the conversion to `f64`; use
55/// [`crate::arrow::bars_to_arrow_record_batch_bytes`] for catalog storage.
56///
57/// Returns an empty [`RecordBatch`] with the correct schema when `data` is empty.
58///
59/// # Errors
60///
61/// Returns an [`ArrowError`] if the Arrow `RecordBatch` cannot be constructed.
62pub fn encode_bars(data: &[Bar]) -> Result<RecordBatch, ArrowError> {
63    let mut instrument_id_builder = StringBuilder::new();
64    let mut bar_type_builder = StringBuilder::new();
65    let mut open_builder = Float64Builder::with_capacity(data.len());
66    let mut high_builder = Float64Builder::with_capacity(data.len());
67    let mut low_builder = Float64Builder::with_capacity(data.len());
68    let mut close_builder = Float64Builder::with_capacity(data.len());
69    let mut volume_builder = Float64Builder::with_capacity(data.len());
70    let mut ts_event_builder =
71        TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
72    let mut ts_init_builder =
73        TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
74
75    for bar in data {
76        instrument_id_builder.append_value(bar.instrument_id().to_string());
77        bar_type_builder.append_value(bar.bar_type.to_string());
78        open_builder.append_value(price_to_f64(&bar.open));
79        high_builder.append_value(price_to_f64(&bar.high));
80        low_builder.append_value(price_to_f64(&bar.low));
81        close_builder.append_value(price_to_f64(&bar.close));
82        volume_builder.append_value(quantity_to_f64(&bar.volume));
83        ts_event_builder.append_value(unix_nanos_to_i64(bar.ts_event.as_u64()));
84        ts_init_builder.append_value(unix_nanos_to_i64(bar.ts_init.as_u64()));
85    }
86
87    RecordBatch::try_new(
88        Arc::new(bars_schema()),
89        vec![
90            Arc::new(instrument_id_builder.finish()),
91            Arc::new(bar_type_builder.finish()),
92            Arc::new(open_builder.finish()),
93            Arc::new(high_builder.finish()),
94            Arc::new(low_builder.finish()),
95            Arc::new(close_builder.finish()),
96            Arc::new(volume_builder.finish()),
97            Arc::new(ts_event_builder.finish()),
98            Arc::new(ts_init_builder.finish()),
99        ],
100    )
101}
102
103#[cfg(test)]
104mod tests {
105    use std::str::FromStr;
106
107    use arrow::{
108        array::{Array, Float64Array, StringArray, TimestampNanosecondArray},
109        datatypes::{DataType, TimeUnit},
110    };
111    use nautilus_model::{
112        data::BarType,
113        types::{Price, Quantity},
114    };
115    use rstest::rstest;
116
117    use super::*;
118
119    fn make_bar(
120        bar_type_str: &str,
121        open: &str,
122        high: &str,
123        low: &str,
124        close: &str,
125        ts: u64,
126    ) -> Bar {
127        let bar_type = BarType::from_str(bar_type_str).unwrap();
128        Bar::new(
129            bar_type,
130            Price::from(open),
131            Price::from(high),
132            Price::from(low),
133            Price::from(close),
134            Quantity::from(1_100),
135            ts.into(),
136            (ts + 1).into(),
137        )
138    }
139
140    #[rstest]
141    fn test_encode_bars_schema() {
142        let batch = encode_bars(&[]).unwrap();
143        let fields = batch.schema().fields().clone();
144        assert_eq!(fields.len(), 9);
145        assert_eq!(fields[0].name(), "instrument_id");
146        assert_eq!(fields[0].data_type(), &DataType::Utf8);
147        assert_eq!(fields[1].name(), "bar_type");
148        assert_eq!(fields[1].data_type(), &DataType::Utf8);
149        assert_eq!(fields[2].name(), "open");
150        assert_eq!(fields[2].data_type(), &DataType::Float64);
151        assert_eq!(fields[3].name(), "high");
152        assert_eq!(fields[4].name(), "low");
153        assert_eq!(fields[5].name(), "close");
154        assert_eq!(fields[6].name(), "volume");
155        assert_eq!(fields[6].data_type(), &DataType::Float64);
156        assert_eq!(fields[7].name(), "ts_event");
157        assert_eq!(
158            fields[7].data_type(),
159            &DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
160        );
161        assert_eq!(fields[8].name(), "ts_init");
162    }
163
164    #[rstest]
165    fn test_encode_bars_values() {
166        let bars = vec![
167            make_bar(
168                "AAPL.XNAS-1-MINUTE-LAST-INTERNAL",
169                "100.10",
170                "102.00",
171                "100.00",
172                "101.00",
173                1_000,
174            ),
175            make_bar(
176                "AAPL.XNAS-1-MINUTE-LAST-INTERNAL",
177                "100.20",
178                "102.00",
179                "100.00",
180                "101.00",
181                2_000,
182            ),
183        ];
184        let batch = encode_bars(&bars).unwrap();
185
186        assert_eq!(batch.num_rows(), 2);
187
188        let instrument_id_col = batch
189            .column(0)
190            .as_any()
191            .downcast_ref::<StringArray>()
192            .unwrap();
193        let bar_type_col = batch
194            .column(1)
195            .as_any()
196            .downcast_ref::<StringArray>()
197            .unwrap();
198        let open_col = batch
199            .column(2)
200            .as_any()
201            .downcast_ref::<Float64Array>()
202            .unwrap();
203        let high_col = batch
204            .column(3)
205            .as_any()
206            .downcast_ref::<Float64Array>()
207            .unwrap();
208        let low_col = batch
209            .column(4)
210            .as_any()
211            .downcast_ref::<Float64Array>()
212            .unwrap();
213        let close_col = batch
214            .column(5)
215            .as_any()
216            .downcast_ref::<Float64Array>()
217            .unwrap();
218        let volume_col = batch
219            .column(6)
220            .as_any()
221            .downcast_ref::<Float64Array>()
222            .unwrap();
223        let ts_event_col = batch
224            .column(7)
225            .as_any()
226            .downcast_ref::<TimestampNanosecondArray>()
227            .unwrap();
228        let ts_init_col = batch
229            .column(8)
230            .as_any()
231            .downcast_ref::<TimestampNanosecondArray>()
232            .unwrap();
233
234        assert_eq!(instrument_id_col.value(0), "AAPL.XNAS");
235        assert_eq!(bar_type_col.value(0), "AAPL.XNAS-1-MINUTE-LAST-INTERNAL");
236        assert!((open_col.value(0) - 100.10).abs() < 1e-9);
237        assert!((open_col.value(1) - 100.20).abs() < 1e-9);
238        assert!((high_col.value(0) - 102.00).abs() < 1e-9);
239        assert!((low_col.value(0) - 100.00).abs() < 1e-9);
240        assert!((close_col.value(0) - 101.00).abs() < 1e-9);
241        assert!((volume_col.value(0) - 1_100.0).abs() < 1e-9);
242        assert_eq!(ts_event_col.value(0), 1_000);
243        assert_eq!(ts_init_col.value(1), 2_001);
244    }
245
246    #[rstest]
247    fn test_encode_bars_empty() {
248        let batch = encode_bars(&[]).unwrap();
249        assert_eq!(batch.num_rows(), 0);
250    }
251
252    #[rstest]
253    fn test_encode_bars_mixed_instruments() {
254        let bars = vec![
255            make_bar(
256                "AAPL.XNAS-1-MINUTE-LAST-INTERNAL",
257                "100.10",
258                "102.00",
259                "100.00",
260                "101.00",
261                1,
262            ),
263            make_bar(
264                "MSFT.XNAS-1-MINUTE-LAST-INTERNAL",
265                "250.00",
266                "251.00",
267                "249.00",
268                "250.50",
269                2,
270            ),
271        ];
272        let batch = encode_bars(&bars).unwrap();
273        let instrument_id_col = batch
274            .column(0)
275            .as_any()
276            .downcast_ref::<StringArray>()
277            .unwrap();
278        assert_eq!(instrument_id_col.value(0), "AAPL.XNAS");
279        assert_eq!(instrument_id_col.value(1), "MSFT.XNAS");
280    }
281}