nautilus_serialization/arrow/
account_state.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;
17
18use arrow::{datatypes::Schema, error::ArrowError, record_batch::RecordBatch};
19use nautilus_model::events::AccountState;
20
21use super::{
22    ArrowSchemaProvider, DecodeTypedFromRecordBatch, EncodeToRecordBatch, EncodingError,
23    json::{JsonFieldSpec, decode_batch, encode_batch, metadata_for_type, schema_for_type},
24};
25
26const ACCOUNT_STATE_FIELDS: &[JsonFieldSpec] = &[
27    JsonFieldSpec::utf8("account_id", false),
28    JsonFieldSpec::utf8("account_type", false),
29    JsonFieldSpec::utf8("base_currency", true),
30    JsonFieldSpec::utf8_json("balances", false),
31    JsonFieldSpec::utf8_json("margins", false),
32    JsonFieldSpec::boolean("is_reported", false),
33    JsonFieldSpec::utf8("event_id", false),
34    JsonFieldSpec::timestamp("ts_event", false),
35    JsonFieldSpec::timestamp("ts_init", false),
36    JsonFieldSpec::utf8_json("info", true),
37];
38
39impl ArrowSchemaProvider for AccountState {
40    fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
41        schema_for_type("AccountState", metadata, ACCOUNT_STATE_FIELDS)
42    }
43}
44
45impl EncodeToRecordBatch for AccountState {
46    fn encode_batch<T>(
47        metadata: &HashMap<String, String>,
48        data: &[T],
49    ) -> Result<RecordBatch, ArrowError>
50    where
51        T: std::borrow::Borrow<Self>,
52    {
53        encode_batch(
54            "AccountState",
55            metadata,
56            data.iter().map(std::borrow::Borrow::borrow),
57            ACCOUNT_STATE_FIELDS,
58        )
59    }
60
61    fn metadata(&self) -> HashMap<String, String> {
62        metadata_for_type("AccountState")
63    }
64}
65
66impl DecodeTypedFromRecordBatch for AccountState {
67    fn decode_typed_batch(
68        metadata: &HashMap<String, String>,
69        record_batch: RecordBatch,
70    ) -> Result<Vec<Self>, EncodingError> {
71        let fields = if record_batch.schema().index_of("info").is_ok() {
72            ACCOUNT_STATE_FIELDS
73        } else {
74            &ACCOUNT_STATE_FIELDS[..ACCOUNT_STATE_FIELDS.len() - 1]
75        };
76        decode_batch(metadata, &record_batch, fields, Some("AccountState"))
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use nautilus_core::Params;
83    use nautilus_model::{
84        events::account::stubs::cash_account_state,
85        types::{AccountBalance, Money},
86    };
87    use rstest::rstest;
88    use serde_json::json;
89
90    use super::*;
91
92    #[rstest]
93    fn test_account_state_round_trip(cash_account_state: AccountState) {
94        let info = Params::new(
95            [
96                (
97                    "total_wallet_balance".to_string(),
98                    json!("1525000.00000001"),
99                ),
100                ("can_trade".to_string(), json!(true)),
101            ]
102            .into(),
103        );
104        let state = cash_account_state.with_info(Some(info));
105        let metadata = state.metadata();
106        let batch = AccountState::encode_batch(&metadata, std::slice::from_ref(&state)).unwrap();
107        let decoded = AccountState::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
108
109        assert_eq!(decoded.len(), 1);
110        assert_eq!(decoded[0].account_id, state.account_id);
111        assert_eq!(decoded[0].balances, state.balances);
112        assert_eq!(decoded[0].margins, state.margins);
113        assert_eq!(decoded[0].base_currency, state.base_currency);
114        assert_eq!(decoded[0].info, state.info);
115    }
116
117    #[rstest]
118    fn test_account_state_exact_usd_round_trip(mut cash_account_state: AccountState) {
119        cash_account_state.balances = vec![AccountBalance::new(
120            Money::from_str_exact("1.0001 USD").unwrap(),
121            Money::from_str_exact("0.000001 USD").unwrap(),
122            Money::from_str_exact("1.000099 USD").unwrap(),
123        )];
124        let metadata = cash_account_state.metadata();
125        let batch =
126            AccountState::encode_batch(&metadata, std::slice::from_ref(&cash_account_state))
127                .unwrap();
128        let decoded = AccountState::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
129
130        assert_eq!(decoded.len(), 1);
131        assert_eq!(decoded[0].balances, cash_account_state.balances);
132        assert_eq!(
133            serde_json::to_value(&decoded[0]).unwrap(),
134            serde_json::to_value(&cash_account_state).unwrap(),
135        );
136    }
137
138    #[rstest]
139    fn test_account_state_decodes_legacy_batch_without_info(cash_account_state: AccountState) {
140        let metadata = cash_account_state.metadata();
141        let legacy_fields = &ACCOUNT_STATE_FIELDS[..ACCOUNT_STATE_FIELDS.len() - 1];
142        let batch = encode_batch(
143            "AccountState",
144            &metadata,
145            std::slice::from_ref(&cash_account_state),
146            legacy_fields,
147        )
148        .unwrap();
149        let decoded = AccountState::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
150
151        assert_eq!(decoded.len(), 1);
152        assert!(decoded[0].info.is_none());
153    }
154}
nautilus_serialization/arrow/account_state.rs

nautilus_serialization/arrow/
account_state.rs
