nautilus_serialization/arrow/
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
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::{data::prices::IndexPriceUpdate, identifiers::InstrumentId};
25
26use super::{
27    DecodeDataFromRecordBatch, EncodingError, KEY_IDENTIFIER, KEY_INSTRUMENT_ID,
28    KEY_PRICE_PRECISION, decode_decimal_price, decode_required_timestamp, extract_column,
29    fixed_decimal_data_type, identifier_array_from_display, metadata_with_type_name,
30    price_decimal_array,
31};
32use crate::arrow::{ArrowSchemaProvider, Data, DecodeFromRecordBatch, EncodeToRecordBatch};
33
34impl ArrowSchemaProvider for IndexPriceUpdate {
35    fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
36        let fields = vec![
37            Field::new("value", fixed_decimal_data_type(), true),
38            Field::new("ts_event", crate::arrow::timestamp_data_type(), false),
39            Field::new("ts_init", crate::arrow::timestamp_data_type(), false),
40            Field::new(KEY_IDENTIFIER, DataType::Utf8, true),
41        ];
42
43        Schema::new_with_metadata(
44            fields,
45            metadata_with_type_name("IndexPriceUpdate", metadata),
46        )
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 IndexPriceUpdate {
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 update in data.iter().map(std::borrow::Borrow::borrow) {
78            ts_event_builder.append_value(update.ts_event.as_u64());
79            ts_init_builder.append_value(update.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(|update| update.value.raw()),
89                    "value",
90                )?),
91                Arc::new(ts_event_builder.finish()),
92                Arc::new(ts_init_builder.finish()),
93                Arc::new(identifier_array_from_display(
94                    data.iter()
95                        .map(std::borrow::Borrow::borrow)
96                        .map(|update| update.instrument_id),
97                )),
98            ],
99        )
100    }
101
102    fn metadata(&self) -> HashMap<String, String> {
103        let mut metadata = HashMap::new();
104        metadata.insert(
105            KEY_INSTRUMENT_ID.to_string(),
106            self.instrument_id.to_string(),
107        );
108        metadata.insert(
109            KEY_PRICE_PRECISION.to_string(),
110            self.value.precision.to_string(),
111        );
112        metadata
113    }
114}
115
116impl DecodeFromRecordBatch for IndexPriceUpdate {
117    fn decode_batch(
118        metadata: &HashMap<String, String>,
119        record_batch: RecordBatch,
120    ) -> Result<Vec<Self>, EncodingError> {
121        let (instrument_id, price_precision) = parse_metadata(metadata)?;
122        let record_batch = crate::arrow::record_batch_with_u64_timestamps(&record_batch)?;
123        let record_batch = &record_batch;
124        let cols = record_batch.columns();
125
126        let value_values =
127            extract_column::<Decimal128Array>(cols, "value", 0, fixed_decimal_data_type())?;
128        let ts_event_values = extract_column::<UInt64Array>(cols, "ts_event", 1, DataType::UInt64)?;
129        let ts_init_values = extract_column::<UInt64Array>(cols, "ts_init", 2, DataType::UInt64)?;
130
131        let result: Result<Vec<Self>, EncodingError> = (0..record_batch.num_rows())
132            .map(|row| {
133                let value = decode_decimal_price(value_values, price_precision, "value", row)?;
134                Ok(Self {
135                    instrument_id,
136                    value,
137                    ts_event: decode_required_timestamp(ts_event_values, "ts_event", row)?,
138                    ts_init: decode_required_timestamp(ts_init_values, "ts_init", row)?,
139                })
140            })
141            .collect();
142
143        result
144    }
145}
146
147impl DecodeDataFromRecordBatch for IndexPriceUpdate {
148    fn decode_data_batch(
149        metadata: &HashMap<String, String>,
150        record_batch: RecordBatch,
151    ) -> Result<Vec<Data>, EncodingError> {
152        let updates: Vec<Self> = Self::decode_batch(metadata, record_batch)?;
153        Ok(updates.into_iter().map(Data::from).collect())
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use std::sync::Arc;
160
161    use arrow::array::{Array, TimestampNanosecondArray};
162    use nautilus_model::types::{Price, fixed::FIXED_SCALAR, price::PriceRaw};
163    use rstest::rstest;
164    use rust_decimal_macros::dec;
165
166    use super::*;
167    use crate::arrow::{KEY_TYPE_NAME, get_raw_price};
168
169    #[rstest]
170    fn test_get_schema() {
171        let instrument_id = InstrumentId::from("BTC-USDT.BINANCE");
172        let metadata = HashMap::from([
173            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
174            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
175        ]);
176        let schema = IndexPriceUpdate::get_schema(Some(metadata.clone()));
177
178        let expected_fields = vec![
179            Field::new("value", fixed_decimal_data_type(), true),
180            Field::new("ts_event", crate::arrow::timestamp_data_type(), false),
181            Field::new("ts_init", crate::arrow::timestamp_data_type(), false),
182            Field::new(KEY_IDENTIFIER, DataType::Utf8, true),
183        ];
184
185        let mut expected_metadata = metadata;
186        expected_metadata.insert(KEY_TYPE_NAME.to_string(), "IndexPriceUpdate".to_string());
187        let expected_schema = Schema::new_with_metadata(expected_fields, expected_metadata);
188        assert_eq!(schema, expected_schema);
189    }
190
191    #[rstest]
192    fn test_get_schema_map() {
193        let schema_map = IndexPriceUpdate::get_schema_map();
194        let mut expected_map = HashMap::new();
195
196        let fixed_size_binary = "Decimal128(38, 16)".to_string();
197        expected_map.insert("value".to_string(), fixed_size_binary);
198        expected_map.insert(
199            "ts_event".to_string(),
200            "Timestamp(Nanosecond, Some(\"UTC\"))".to_string(),
201        );
202        expected_map.insert(
203            "ts_init".to_string(),
204            "Timestamp(Nanosecond, Some(\"UTC\"))".to_string(),
205        );
206        expected_map.insert(KEY_IDENTIFIER.to_string(), "Utf8".to_string());
207        assert_eq!(schema_map, expected_map);
208    }
209
210    #[rstest]
211    fn test_encode_batch() {
212        let instrument_id = InstrumentId::from("BTC-USDT.BINANCE");
213        let metadata = HashMap::from([
214            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
215            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
216        ]);
217
218        let update1 = IndexPriceUpdate {
219            instrument_id,
220            value: Price::from("50000.00"),
221            ts_event: 1.into(),
222            ts_init: 3.into(),
223        };
224
225        let update2 = IndexPriceUpdate {
226            instrument_id,
227            value: Price::from("51000.00"),
228            ts_event: 2.into(),
229            ts_init: 4.into(),
230        };
231
232        let data = vec![update1, update2];
233        let record_batch = IndexPriceUpdate::encode_batch(&metadata, &data).unwrap();
234
235        let columns = record_batch.columns();
236        let value_values = columns[0]
237            .as_any()
238            .downcast_ref::<Decimal128Array>()
239            .unwrap();
240        let ts_event_values = columns[1]
241            .as_any()
242            .downcast_ref::<TimestampNanosecondArray>()
243            .unwrap();
244        let ts_init_values = columns[2]
245            .as_any()
246            .downcast_ref::<TimestampNanosecondArray>()
247            .unwrap();
248
249        assert_eq!(columns.len(), 4);
250        assert_eq!(value_values.len(), 2);
251        assert_eq!(
252            get_raw_price(value_values.value(0)),
253            Price::from(dec!(50000.00).to_string()).raw()
254        );
255        assert_eq!(
256            get_raw_price(value_values.value(1)),
257            Price::from(dec!(51000.00).to_string()).raw()
258        );
259        assert_eq!(ts_event_values.len(), 2);
260        assert_eq!(ts_event_values.value(0), 1);
261        assert_eq!(ts_event_values.value(1), 2);
262        assert_eq!(ts_init_values.len(), 2);
263        assert_eq!(ts_init_values.value(0), 3);
264        assert_eq!(ts_init_values.value(1), 4);
265    }
266
267    #[rstest]
268    fn test_decode_batch() {
269        let instrument_id = InstrumentId::from("BTC-USDT.BINANCE");
270        let metadata = HashMap::from([
271            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
272            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
273        ]);
274
275        let raw_price1 = (50.00 * FIXED_SCALAR) as PriceRaw;
276        let raw_price2 = (51.00 * FIXED_SCALAR) as PriceRaw;
277        let value = crate::arrow::test_support::decimal_array_from_bytes(vec![
278            &raw_price1.to_le_bytes(),
279            &raw_price2.to_le_bytes(),
280        ]);
281        let ts_event = UInt64Array::from(vec![1, 2]);
282        let ts_init = UInt64Array::from(vec![3, 4]);
283
284        let record_batch = crate::arrow::record_batch_with_timestamps(
285            crate::arrow::schema_without_identifier_column(&IndexPriceUpdate::get_schema(Some(
286                metadata.clone(),
287            )))
288            .into(),
289            vec![Arc::new(value), Arc::new(ts_event), Arc::new(ts_init)],
290        )
291        .unwrap();
292
293        let decoded_data = IndexPriceUpdate::decode_batch(&metadata, record_batch).unwrap();
294
295        assert_eq!(decoded_data.len(), 2);
296        assert_eq!(decoded_data[0].instrument_id, instrument_id);
297        assert_eq!(decoded_data[0].value, Price::from_raw(raw_price1, 2));
298        assert_eq!(decoded_data[0].ts_event.as_u64(), 1);
299        assert_eq!(decoded_data[0].ts_init.as_u64(), 3);
300
301        assert_eq!(decoded_data[1].instrument_id, instrument_id);
302        assert_eq!(decoded_data[1].value, Price::from_raw(raw_price2, 2));
303        assert_eq!(decoded_data[1].ts_event.as_u64(), 2);
304        assert_eq!(decoded_data[1].ts_init.as_u64(), 4);
305    }
306
307    #[rstest]
308    fn test_decode_batch_rejects_null_timestamp_with_field_and_row() {
309        let instrument_id = InstrumentId::from("BTC-USDT.BINANCE");
310        let metadata = HashMap::from([
311            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
312            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
313        ]);
314        let update = IndexPriceUpdate {
315            instrument_id,
316            value: Price::from("50000.00"),
317            ts_event: 1.into(),
318            ts_init: 2.into(),
319        };
320        let encoded = IndexPriceUpdate::encode_batch(&metadata, &[update]).unwrap();
321        let mut columns = encoded.columns().to_vec();
322        columns[2] = Arc::new(TimestampNanosecondArray::from(vec![None]).with_timezone("UTC"));
323        let fields = encoded
324            .schema()
325            .fields()
326            .iter()
327            .map(|field| {
328                if field.name() == "ts_init" {
329                    Arc::new(field.as_ref().clone().with_nullable(true))
330                } else {
331                    Arc::clone(field)
332                }
333            })
334            .collect::<Vec<_>>();
335        let schema = Arc::new(Schema::new_with_metadata(fields, metadata.clone()));
336        let batch = RecordBatch::try_new(schema, columns).unwrap();
337
338        let error = IndexPriceUpdate::decode_batch(&metadata, batch).unwrap_err();
339
340        assert!(error.to_string().contains("ts_init"));
341        assert!(error.to_string().contains("row 0"));
342    }
343
344    #[rstest]
345    fn test_decode_batch_invalid_value_returns_error() {
346        let instrument_id = InstrumentId::from("BTC-USDT.BINANCE");
347        let metadata = HashMap::from([
348            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
349            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
350        ]);
351
352        let invalid_price: PriceRaw = PriceRaw::MAX - 1000;
353        let value = crate::arrow::test_support::decimal_array_from_bytes(vec![
354            &invalid_price.to_le_bytes(),
355        ]);
356        let ts_event = UInt64Array::from(vec![1]);
357        let ts_init = UInt64Array::from(vec![2]);
358
359        let record_batch = crate::arrow::record_batch_with_timestamps(
360            crate::arrow::schema_without_identifier_column(&IndexPriceUpdate::get_schema(Some(
361                metadata.clone(),
362            )))
363            .into(),
364            vec![Arc::new(value), Arc::new(ts_event), Arc::new(ts_init)],
365        )
366        .unwrap();
367
368        let result = IndexPriceUpdate::decode_batch(&metadata, record_batch);
369        assert!(result.is_err());
370        let err = result.unwrap_err();
371        assert!(
372            err.to_string().contains("value") && err.to_string().contains("row 0"),
373            "Expected value error at row 0, was: {err}"
374        );
375    }
376
377    #[rstest]
378    fn test_decode_batch_missing_instrument_id_returns_error() {
379        let mut metadata = HashMap::from([
380            (
381                KEY_INSTRUMENT_ID.to_string(),
382                "BTC-USDT.BINANCE".to_string(),
383            ),
384            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
385        ]);
386
387        let raw_price = (50.00 * FIXED_SCALAR) as PriceRaw;
388        let value =
389            crate::arrow::test_support::decimal_array_from_bytes(vec![&raw_price.to_le_bytes()]);
390        let ts_event = UInt64Array::from(vec![1]);
391        let ts_init = UInt64Array::from(vec![2]);
392
393        let record_batch = crate::arrow::record_batch_with_timestamps(
394            crate::arrow::schema_without_identifier_column(&IndexPriceUpdate::get_schema(Some(
395                metadata.clone(),
396            )))
397            .into(),
398            vec![Arc::new(value), Arc::new(ts_event), Arc::new(ts_init)],
399        )
400        .unwrap();
401
402        metadata.remove(KEY_INSTRUMENT_ID);
403
404        let result = IndexPriceUpdate::decode_batch(&metadata, record_batch);
405        assert!(result.is_err());
406        let err = result.unwrap_err();
407        assert!(
408            err.to_string().contains("instrument_id"),
409            "Expected missing instrument_id error, was: {err}"
410        );
411    }
412
413    #[rstest]
414    fn test_encode_decode_round_trip() {
415        let instrument_id = InstrumentId::from("BTC-USDT.BINANCE");
416        let metadata = HashMap::from([
417            (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
418            (KEY_PRICE_PRECISION.to_string(), "2".to_string()),
419        ]);
420
421        let update1 = IndexPriceUpdate {
422            instrument_id,
423            value: Price::from("50000.00"),
424            ts_event: 1_000_000_000.into(),
425            ts_init: 1_000_000_001.into(),
426        };
427
428        let update2 = IndexPriceUpdate {
429            instrument_id,
430            value: Price::from("51000.00"),
431            ts_event: 2_000_000_000.into(),
432            ts_init: 2_000_000_001.into(),
433        };
434
435        let original = vec![update1, update2];
436        let record_batch = IndexPriceUpdate::encode_batch(&metadata, &original).unwrap();
437        let decoded = IndexPriceUpdate::decode_batch(&metadata, record_batch).unwrap();
438
439        assert_eq!(decoded.len(), original.len());
440        for (orig, dec) in original.iter().zip(decoded.iter()) {
441            assert_eq!(dec.instrument_id, orig.instrument_id);
442            assert_eq!(dec.value, orig.value);
443            assert_eq!(dec.ts_event, orig.ts_event);
444            assert_eq!(dec.ts_init, orig.ts_init);
445        }
446    }
447}
nautilus_serialization/arrow/index_price.rs

nautilus_serialization/arrow/
index_price.rs
