Skip to main content

nautilus_serialization/arrow/display/
index_price.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 [`IndexPriceUpdate`].
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::IndexPriceUpdate;
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 [`IndexPriceUpdate`].
32#[must_use]
33pub fn index_prices_schema() -> Schema {
34    Schema::new(vec![
35        utf8_field("instrument_id", false),
36        float64_field("value", false),
37        timestamp_field("ts_event", false),
38        timestamp_field("ts_init", false),
39    ])
40}
41
42/// Encodes index price updates as a display-friendly Arrow [`RecordBatch`].
43///
44/// Emits a `Float64` `value` column, a `Utf8` `instrument_id` column, and
45/// `Timestamp(Nanosecond)` columns for event and init times. Mixed-instrument
46/// batches are supported. Precision is lost on the conversion to `f64`; use
47/// [`crate::arrow::index_prices_to_arrow_record_batch_bytes`] for catalog storage.
48///
49/// Returns an empty [`RecordBatch`] with the correct schema when `data` is empty.
50///
51/// # Errors
52///
53/// Returns an [`ArrowError`] if the Arrow `RecordBatch` cannot be constructed.
54pub fn encode_index_prices(data: &[IndexPriceUpdate]) -> Result<RecordBatch, ArrowError> {
55    let mut instrument_id_builder = StringBuilder::new();
56    let mut value_builder = Float64Builder::with_capacity(data.len());
57    let mut ts_event_builder =
58        TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
59    let mut ts_init_builder =
60        TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
61
62    for update in data {
63        instrument_id_builder.append_value(update.instrument_id.to_string());
64        value_builder.append_value(price_to_f64(&update.value));
65        ts_event_builder.append_value(unix_nanos_to_i64(update.ts_event.as_u64()));
66        ts_init_builder.append_value(unix_nanos_to_i64(update.ts_init.as_u64()));
67    }
68
69    RecordBatch::try_new(
70        Arc::new(index_prices_schema()),
71        vec![
72            Arc::new(instrument_id_builder.finish()),
73            Arc::new(value_builder.finish()),
74            Arc::new(ts_event_builder.finish()),
75            Arc::new(ts_init_builder.finish()),
76        ],
77    )
78}
79
80#[cfg(test)]
81mod tests {
82    use arrow::{
83        array::{Array, Float64Array, StringArray, TimestampNanosecondArray},
84        datatypes::{DataType, TimeUnit},
85    };
86    use nautilus_model::{identifiers::InstrumentId, types::Price};
87    use rstest::rstest;
88
89    use super::*;
90
91    fn make_update(instrument_id: &str, value: &str, ts: u64) -> IndexPriceUpdate {
92        IndexPriceUpdate {
93            instrument_id: InstrumentId::from(instrument_id),
94            value: Price::from(value),
95            ts_event: ts.into(),
96            ts_init: (ts + 1).into(),
97        }
98    }
99
100    #[rstest]
101    fn test_encode_index_prices_schema() {
102        let batch = encode_index_prices(&[]).unwrap();
103        let fields = batch.schema().fields().clone();
104        assert_eq!(fields.len(), 4);
105        assert_eq!(fields[0].name(), "instrument_id");
106        assert_eq!(fields[0].data_type(), &DataType::Utf8);
107        assert_eq!(fields[1].name(), "value");
108        assert_eq!(fields[1].data_type(), &DataType::Float64);
109        assert_eq!(fields[2].name(), "ts_event");
110        assert_eq!(
111            fields[2].data_type(),
112            &DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
113        );
114        assert_eq!(fields[3].name(), "ts_init");
115    }
116
117    #[rstest]
118    fn test_encode_index_prices_values() {
119        let updates = vec![
120            make_update("BTC-USDT.BINANCE", "50200.00", 1_000),
121            make_update("BTC-USDT.BINANCE", "50300.00", 2_000),
122        ];
123        let batch = encode_index_prices(&updates).unwrap();
124
125        assert_eq!(batch.num_rows(), 2);
126
127        let value_col = batch
128            .column(1)
129            .as_any()
130            .downcast_ref::<Float64Array>()
131            .unwrap();
132        let ts_event_col = batch
133            .column(2)
134            .as_any()
135            .downcast_ref::<TimestampNanosecondArray>()
136            .unwrap();
137
138        assert!((value_col.value(0) - 50_200.00).abs() < 1e-9);
139        assert!((value_col.value(1) - 50_300.00).abs() < 1e-9);
140        assert_eq!(ts_event_col.value(0), 1_000);
141    }
142
143    #[rstest]
144    fn test_encode_index_prices_empty() {
145        let batch = encode_index_prices(&[]).unwrap();
146        assert_eq!(batch.num_rows(), 0);
147    }
148
149    #[rstest]
150    fn test_encode_index_prices_mixed_instruments() {
151        let updates = vec![
152            make_update("BTC-USDT.BINANCE", "50200.00", 1),
153            make_update("ETH-USDT.BINANCE", "2500.00", 2),
154        ];
155        let batch = encode_index_prices(&updates).unwrap();
156        let instrument_id_col = batch
157            .column(0)
158            .as_any()
159            .downcast_ref::<StringArray>()
160            .unwrap();
161        assert_eq!(instrument_id_col.value(0), "BTC-USDT.BINANCE");
162        assert_eq!(instrument_id_col.value(1), "ETH-USDT.BINANCE");
163    }
164}