nautilus_databento/arrow/
mod.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//! Apache Arrow schema and encoding/decoding for Databento types.
17
18pub mod imbalance;
19pub mod statistics;
20
21use std::{collections::HashMap, fmt::Display, str::FromStr, sync::Arc};
22
23use arrow::array::{Array, UInt8Array};
24use nautilus_model::{enums::FromU8, identifiers::InstrumentId};
25use nautilus_serialization::arrow::{
26    EncodingError, KEY_INSTRUMENT_ID, KEY_PRICE_PRECISION, KEY_SIZE_PRECISION, StringColumnRef,
27    enum_dictionary_data_type,
28};
29
30fn parse_metadata(
31    metadata: &HashMap<String, String>,
32) -> Result<(InstrumentId, u8, u8), EncodingError> {
33    let instrument_id_str = metadata
34        .get(KEY_INSTRUMENT_ID)
35        .ok_or_else(|| EncodingError::MissingMetadata(KEY_INSTRUMENT_ID))?;
36    let instrument_id = InstrumentId::from_str(instrument_id_str)
37        .map_err(|e| EncodingError::ParseError(KEY_INSTRUMENT_ID, e.to_string()))?;
38
39    let price_precision = metadata
40        .get(KEY_PRICE_PRECISION)
41        .ok_or_else(|| EncodingError::MissingMetadata(KEY_PRICE_PRECISION))?
42        .parse::<u8>()
43        .map_err(|e| EncodingError::ParseError(KEY_PRICE_PRECISION, e.to_string()))?;
44
45    let size_precision = metadata
46        .get(KEY_SIZE_PRECISION)
47        .ok_or_else(|| EncodingError::MissingMetadata(KEY_SIZE_PRECISION))?
48        .parse::<u8>()
49        .map_err(|e| EncodingError::ParseError(KEY_SIZE_PRECISION, e.to_string()))?;
50
51    Ok((instrument_id, price_precision, size_precision))
52}
53
54enum EnumColumn<'a> {
55    Strings(StringColumnRef<'a>, &'static str),
56    Codes(&'a UInt8Array, &'static str),
57}
58
59impl<'a> EnumColumn<'a> {
60    fn try_from_column(
61        column: &'a Arc<dyn Array>,
62        field: &'static str,
63        index: usize,
64    ) -> Result<Self, EncodingError> {
65        if let Some(values) = StringColumnRef::try_from_array(column.as_ref()) {
66            return Ok(Self::Strings(values, field));
67        }
68        column
69            .as_any()
70            .downcast_ref::<UInt8Array>()
71            .map(|values| Self::Codes(values, field))
72            .ok_or_else(|| {
73                EncodingError::InvalidColumnType(
74                    field,
75                    index,
76                    enum_dictionary_data_type(),
77                    column.data_type().clone(),
78                )
79            })
80    }
81
82    fn decode<T>(&self, row: usize) -> Result<T, EncodingError>
83    where
84        T: FromStr + FromU8,
85        T::Err: Display,
86    {
87        match self {
88            Self::Strings(values, field) => values
89                .value(row)
90                .parse::<T>()
91                .map_err(|e| EncodingError::ParseError(field, format!("row {row}: {e}"))),
92            Self::Codes(values, field) => {
93                let value = values.value(row);
94                T::from_u8(value).ok_or_else(|| {
95                    EncodingError::ParseError(field, format!("Invalid enum value, was {value}"))
96                })
97            }
98        }
99    }
100
101    fn decode_optional<T, F>(
102        &self,
103        row: usize,
104        legacy_none: &str,
105        decode_code: F,
106    ) -> Result<Option<T>, EncodingError>
107    where
108        T: FromStr,
109        T::Err: Display,
110        F: Fn(u8) -> Option<T>,
111    {
112        match self {
113            Self::Strings(values, field) => {
114                let value = values.value(row);
115                if value.eq_ignore_ascii_case(legacy_none) {
116                    Ok(None)
117                } else {
118                    value
119                        .parse::<T>()
120                        .map(Some)
121                        .map_err(|e| EncodingError::ParseError(field, format!("row {row}: {e}")))
122                }
123            }
124            Self::Codes(values, field) => {
125                let value = values.value(row);
126                if value == 0 {
127                    Ok(None)
128                } else {
129                    decode_code(value).map(Some).ok_or_else(|| {
130                        EncodingError::ParseError(field, format!("Invalid enum value, was {value}"))
131                    })
132                }
133            }
134        }
135    }
136}
nautilus_databento/arrow/mod.rs

nautilus_databento/arrow/
mod.rs
