Skip to main content

nautilus_persistence/
test_data.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//! Rust custom data types used for catalog roundtrip testing.
17//!
18//! Exposed to Python via the persistence PyO3 module so Python tests can exercise
19//! custom data write/query roundtrips.
20
21use std::collections::HashMap;
22
23use indexmap::IndexMap;
24use nautilus_core::{Params, UnixNanos};
25use nautilus_model::{
26    custom_data,
27    data::BarType,
28    enums::AggressorSide,
29    identifiers::{AccountId, InstrumentId},
30    types::{Currency, Money, Price, Quantity},
31};
32use nautilus_serialization::arrow_custom_data;
33
34/// A simple Rust custom data type for roundtrip testing.
35///
36/// Used in persistence integration tests (`test_catalog.rs`) and Python roundtrip tests.
37/// Tests call `ensure_custom_data_registered::<RustTestCustomData>()` before using the catalog.
38#[cfg_attr(
39    feature = "python",
40    expect(
41        clippy::unsafe_derive_deserialize,
42        reason = "test data uses the custom data macro output under test"
43    )
44)]
45#[arrow_custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
46#[custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
47pub struct RustTestCustomData {
48    pub instrument_id: InstrumentId,
49    pub value: f64,
50    pub flag: bool,
51    pub ts_event: UnixNanos,
52    pub ts_init: UnixNanos,
53}
54
55/// Rust custom data type that exercises raw byte field support.
56#[arrow_custom_data]
57#[custom_data]
58pub struct RustTestBytesCustomData {
59    pub value: Vec<u8>,
60    pub ts_event: UnixNanos,
61    pub ts_init: UnixNanos,
62}
63
64/// Rust custom data type with native fixed-point Arrow fields.
65#[cfg_attr(
66    feature = "python",
67    expect(
68        clippy::unsafe_derive_deserialize,
69        reason = "test data uses the custom data macro output under test"
70    )
71)]
72#[arrow_custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
73#[custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
74pub struct RustTestFixedCustomData {
75    pub instrument_id: InstrumentId,
76    pub price: Price,
77    pub quantity: Quantity,
78    #[custom_data_field(native_enum)]
79    pub aggressor_side: AggressorSide,
80    pub notional: Money,
81    #[custom_data_field(native_enum)]
82    pub nullable_aggressor_side: Option<AggressorSide>,
83    pub nullable_notional: Option<Money>,
84    pub ts_event: UnixNanos,
85    pub ts_init: UnixNanos,
86}
87
88/// YieldCurveData-equivalent custom data type using the macro with `Vec<f64>` fields.
89///
90/// Tests `Vec<f64>` / `ListFloat64` support. Exposed to Python for roundtrip tests.
91#[cfg_attr(
92    feature = "python",
93    expect(
94        clippy::unsafe_derive_deserialize,
95        reason = "test data uses the custom data macro output under test"
96    )
97)]
98#[arrow_custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
99#[custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
100pub struct MacroYieldCurveData {
101    pub curve_name: String,
102    pub tenors: Vec<f64>,
103    pub interest_rates: Vec<f64>,
104    pub ts_event: UnixNanos,
105    pub ts_init: UnixNanos,
106}
107
108/// Rust custom data type that exercises `Params` field support in the macro.
109#[cfg_attr(
110    feature = "python",
111    expect(
112        clippy::unsafe_derive_deserialize,
113        reason = "test data uses the custom data macro output under test"
114    )
115)]
116#[arrow_custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
117#[custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
118pub struct RustTestParamsCustomData {
119    pub name: String,
120    pub params: Params,
121    pub ts_event: UnixNanos,
122    pub ts_init: UnixNanos,
123}
124
125/// Rust custom data type that exercises typed map field support in the macro.
126#[cfg_attr(
127    feature = "python",
128    expect(
129        clippy::unsafe_derive_deserialize,
130        reason = "test data uses the custom data macro output under test"
131    )
132)]
133#[arrow_custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
134#[custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
135pub struct RustTestPriceMapCustomData {
136    pub name: String,
137    #[custom_data_field(serde)]
138    pub prices: IndexMap<InstrumentId, Price>,
139    pub ts_event: UnixNanos,
140    pub ts_init: UnixNanos,
141}
142
143/// Rust custom data type that exercises typed JSON map values across PyO3-supported types.
144#[cfg_attr(
145    feature = "python",
146    expect(
147        clippy::unsafe_derive_deserialize,
148        reason = "test data uses the custom data macro output under test"
149    )
150)]
151#[arrow_custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
152#[custom_data(pyo3, stub_module = "nautilus_trader.persistence")]
153pub struct RustTestTypedMapCustomData {
154    pub name: String,
155    #[custom_data_field(serde)]
156    pub instrument_ids: IndexMap<String, InstrumentId>,
157    #[custom_data_field(serde)]
158    pub account_ids: IndexMap<String, AccountId>,
159    #[custom_data_field(serde)]
160    pub currencies: IndexMap<String, Currency>,
161    #[custom_data_field(serde)]
162    pub bar_types: IndexMap<String, BarType>,
163    #[custom_data_field(serde)]
164    pub prices: IndexMap<String, Price>,
165    #[custom_data_field(serde)]
166    pub quantities: IndexMap<String, Quantity>,
167    #[custom_data_field(serde)]
168    pub monies: IndexMap<String, Money>,
169    #[custom_data_field(serde)]
170    pub prices_by_instrument: IndexMap<InstrumentId, Price>,
171    #[custom_data_field(serde)]
172    pub quantities_by_account: IndexMap<AccountId, Quantity>,
173    #[custom_data_field(serde)]
174    pub monies_by_currency: IndexMap<Currency, Money>,
175    #[custom_data_field(serde)]
176    pub prices_by_bar_type: IndexMap<BarType, Price>,
177    #[custom_data_field(serde)]
178    pub hash_prices_by_instrument: HashMap<InstrumentId, Price>,
179    #[custom_data_field(serde)]
180    pub strings: HashMap<String, String>,
181    #[custom_data_field(serde)]
182    pub floats_64: HashMap<String, f64>,
183    #[custom_data_field(serde)]
184    pub floats_32: HashMap<String, f32>,
185    #[custom_data_field(serde)]
186    pub booleans: HashMap<String, bool>,
187    #[custom_data_field(serde)]
188    pub integers_u64: HashMap<String, u64>,
189    #[custom_data_field(serde)]
190    pub integers_i64: HashMap<String, i64>,
191    #[custom_data_field(serde)]
192    pub integers_u32: HashMap<String, u32>,
193    #[custom_data_field(serde)]
194    pub integers_i32: HashMap<String, i32>,
195    pub ts_event: UnixNanos,
196    pub ts_init: UnixNanos,
197}
198
199/// Rust custom data type that exercises generic JSON map field support.
200#[arrow_custom_data]
201#[custom_data]
202pub struct RustTestHashMapCustomData {
203    pub name: String,
204    #[custom_data_field(serde)]
205    pub prices: HashMap<String, Price>,
206    pub ts_event: UnixNanos,
207    pub ts_init: UnixNanos,
208}
209
210/// Plain Serde enum stored inside custom data as a field.
211#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
212pub enum RustTestSerdeFieldKind {
213    Alpha,
214    Beta { count: u64 },
215}
216
217/// Plain Serde payload stored inside custom data without its own timestamps.
218#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
219pub struct RustTestSerdeFieldPayload {
220    pub kind: RustTestSerdeFieldKind,
221    pub label: String,
222    pub values: Vec<f64>,
223}
224
225/// Rust custom data type that exercises arbitrary Serde field support.
226#[arrow_custom_data]
227#[custom_data]
228pub struct RustTestSerdeFieldCustomData {
229    pub name: String,
230    #[custom_data_field(serde)]
231    pub payload: RustTestSerdeFieldPayload,
232    pub ts_event: UnixNanos,
233    pub ts_init: UnixNanos,
234}
235
236#[cfg(test)]
237mod tests {
238    use std::sync::Arc;
239
240    use arrow::datatypes::{DataType, Field, Schema};
241    use nautilus_model::data::get_arrow_schema;
242    use nautilus_serialization::{
243        arrow::{
244            ArrowSchemaProvider, DecodeDataFromRecordBatch, EncodeToRecordBatch,
245            timestamp_data_type,
246        },
247        ensure_custom_data_registered,
248    };
249    use rstest::rstest;
250
251    use super::*;
252
253    #[rstest]
254    fn registered_macro_custom_schemas_are_open() {
255        ensure_custom_data_registered::<RustTestCustomData>();
256        ensure_custom_data_registered::<RustTestFixedCustomData>();
257        ensure_custom_data_registered::<MacroYieldCurveData>();
258        ensure_custom_data_registered::<RustTestParamsCustomData>();
259        ensure_custom_data_registered::<RustTestPriceMapCustomData>();
260        ensure_custom_data_registered::<RustTestTypedMapCustomData>();
261        ensure_custom_data_registered::<RustTestHashMapCustomData>();
262        ensure_custom_data_registered::<RustTestSerdeFieldCustomData>();
263
264        let schemas = [
265            registered_schema::<RustTestCustomData>(),
266            registered_schema::<RustTestFixedCustomData>(),
267            registered_schema::<MacroYieldCurveData>(),
268            registered_schema::<RustTestParamsCustomData>(),
269            registered_schema::<RustTestPriceMapCustomData>(),
270            registered_schema::<RustTestTypedMapCustomData>(),
271            registered_schema::<RustTestHashMapCustomData>(),
272            registered_schema::<RustTestSerdeFieldCustomData>(),
273        ];
274
275        for (name, schema) in schemas {
276            for field in schema.fields() {
277                assert_open_custom_field(name, field);
278            }
279        }
280    }
281
282    #[rstest]
283    fn test_macro_yield_curve_data_schema_has_ts_init() {
284        let schema = MacroYieldCurveData::get_schema(None);
285        let field_names: Vec<_> = schema.fields().iter().map(|f| f.name().clone()).collect();
286        assert!(
287            field_names.iter().any(|f| f == "ts_init"),
288            "Schema must have ts_init for DataFusion ORDER BY; got: {field_names:?}",
289        );
290        assert!(
291            field_names.iter().any(|f| f == "ts_event"),
292            "Schema must have ts_event; got: {field_names:?}",
293        );
294    }
295
296    #[rstest]
297    fn test_rust_test_params_custom_data_schema_uses_utf8_for_params() {
298        let schema = RustTestParamsCustomData::get_schema(None);
299        let params_field = schema.field_with_name("params").unwrap();
300
301        assert_eq!(params_field.data_type(), &DataType::Utf8);
302    }
303
304    #[rstest]
305    fn test_rust_test_price_map_custom_data_schema_uses_utf8_for_prices() {
306        let schema = RustTestPriceMapCustomData::get_schema(None);
307        let prices_field = schema.field_with_name("prices").unwrap();
308
309        assert_eq!(prices_field.data_type(), &DataType::Utf8);
310    }
311
312    #[rstest]
313    fn test_rust_test_hash_map_custom_data_schema_uses_utf8_for_prices() {
314        let schema = RustTestHashMapCustomData::get_schema(None);
315        let prices_field = schema.field_with_name("prices").unwrap();
316
317        assert_eq!(prices_field.data_type(), &DataType::Utf8);
318    }
319
320    fn registered_schema<T>() -> (&'static str, Arc<Schema>)
321    where
322        T: ArrowSchemaProvider,
323    {
324        let name = std::any::type_name::<T>();
325        (
326            name,
327            get_arrow_schema(name.rsplit("::").next().unwrap()).unwrap(),
328        )
329    }
330
331    fn assert_open_custom_field(type_name: &str, field: &Field) {
332        if field.name().starts_with("ts_") {
333            assert_eq!(
334                field.data_type(),
335                &timestamp_data_type(),
336                "registered custom schema {type_name} timestamp field {} is not a UTC nanosecond timestamp",
337                field.name(),
338            );
339        }
340
341        match field.data_type() {
342            DataType::Binary
343            | DataType::LargeBinary
344            | DataType::BinaryView
345            | DataType::FixedSizeBinary(_) => {
346                panic!(
347                    "registered custom schema {type_name} contains opaque field {}: {}",
348                    field.name(),
349                    field.data_type(),
350                );
351            }
352            DataType::List(child)
353            | DataType::LargeList(child)
354            | DataType::ListView(child)
355            | DataType::LargeListView(child)
356            | DataType::FixedSizeList(child, _)
357            | DataType::Map(child, _) => assert_open_custom_field(type_name, child),
358            DataType::Struct(children) => {
359                for child in children {
360                    assert_open_custom_field(type_name, child);
361                }
362            }
363            _ => {}
364        }
365    }
366
367    #[rstest]
368    fn test_rust_test_serde_field_custom_data_schema_uses_utf8_for_payload() {
369        let schema = RustTestSerdeFieldCustomData::get_schema(None);
370        let payload_field = schema.field_with_name("payload").unwrap();
371
372        assert_eq!(payload_field.data_type(), &DataType::Utf8);
373    }
374
375    #[rstest]
376    fn test_rust_test_serde_field_custom_data_roundtrip_decodes_exact_payload() {
377        let original = RustTestSerdeFieldCustomData {
378            name: "serde-field".to_string(),
379            payload: RustTestSerdeFieldPayload {
380                kind: RustTestSerdeFieldKind::Beta { count: 7 },
381                label: "payload".to_string(),
382                values: vec![1.0, 2.0, 3.0],
383            },
384            ts_event: UnixNanos::from(10),
385            ts_init: UnixNanos::from(11),
386        };
387        let metadata = original.metadata();
388        let batch =
389            RustTestSerdeFieldCustomData::encode_batch(&metadata, std::slice::from_ref(&original))
390                .unwrap();
391        let decoded = RustTestSerdeFieldCustomData::decode_data_batch(&metadata, batch).unwrap();
392
393        assert_eq!(decoded.len(), 1);
394        let decoded =
395            RustTestSerdeFieldCustomData::try_from(decoded.into_iter().next().unwrap()).unwrap();
396        assert_eq!(decoded, original);
397    }
398
399    #[rstest]
400    fn test_rust_test_typed_map_custom_data_schema_uses_utf8_for_json_maps() {
401        let schema = RustTestTypedMapCustomData::get_schema(None);
402
403        for field_name in [
404            "instrument_ids",
405            "account_ids",
406            "currencies",
407            "bar_types",
408            "prices",
409            "quantities",
410            "monies",
411            "prices_by_instrument",
412            "quantities_by_account",
413            "monies_by_currency",
414            "prices_by_bar_type",
415            "hash_prices_by_instrument",
416            "strings",
417            "floats_64",
418            "floats_32",
419            "booleans",
420            "integers_u64",
421            "integers_i64",
422            "integers_u32",
423            "integers_i32",
424        ] {
425            let field = schema.field_with_name(field_name).unwrap();
426            assert_eq!(field.data_type(), &DataType::Utf8);
427        }
428    }
429}