1use std::collections::HashMap;
17
18use arrow::{datatypes::Schema, error::ArrowError, record_batch::RecordBatch};
19use nautilus_model::events::{OrderSnapshot, PositionSnapshot};
20
21use super::{
22 ArrowSchemaProvider, DecodeTypedFromRecordBatch, EncodeToRecordBatch, EncodingError,
23 KEY_INSTRUMENT_ID,
24 json::{JsonFieldSpec, decode_batch, encode_batch, metadata_for_type, schema_for_type},
25};
26
27const ORDER_SNAPSHOT_FIELDS: &[JsonFieldSpec] = &[
28 JsonFieldSpec::utf8("trader_id", false),
29 JsonFieldSpec::utf8("strategy_id", false),
30 JsonFieldSpec::utf8("instrument_id", false),
31 JsonFieldSpec::utf8("client_order_id", false),
32 JsonFieldSpec::utf8("venue_order_id", true),
33 JsonFieldSpec::utf8("position_id", true),
34 JsonFieldSpec::utf8("account_id", true),
35 JsonFieldSpec::utf8("last_trade_id", true),
36 JsonFieldSpec::utf8("order_type", false),
37 JsonFieldSpec::utf8("order_side", false),
38 JsonFieldSpec::utf8("quantity", false),
39 JsonFieldSpec::utf8("price", true),
40 JsonFieldSpec::utf8("trigger_price", true),
41 JsonFieldSpec::utf8("trigger_type", true),
42 JsonFieldSpec::utf8("limit_offset", true),
43 JsonFieldSpec::utf8("trailing_offset", true),
44 JsonFieldSpec::utf8("trailing_offset_type", true),
45 JsonFieldSpec::utf8("time_in_force", false),
46 JsonFieldSpec::u64("expire_time", true),
47 JsonFieldSpec::utf8("filled_qty", false),
48 JsonFieldSpec::utf8("liquidity_side", true),
49 JsonFieldSpec::decimal_str("avg_px", true),
50 JsonFieldSpec::decimal_str("slippage", true),
51 JsonFieldSpec::utf8_json("commissions", false),
52 JsonFieldSpec::utf8("status", false),
53 JsonFieldSpec::boolean("is_post_only", false),
54 JsonFieldSpec::boolean("is_reduce_only", false),
55 JsonFieldSpec::boolean("is_quote_quantity", false),
56 JsonFieldSpec::utf8("display_qty", true),
57 JsonFieldSpec::utf8("emulation_trigger", true),
58 JsonFieldSpec::utf8("trigger_instrument_id", true),
59 JsonFieldSpec::utf8("contingency_type", true),
60 JsonFieldSpec::utf8("order_list_id", true),
61 JsonFieldSpec::utf8_json("linked_order_ids", true),
62 JsonFieldSpec::utf8("parent_order_id", true),
63 JsonFieldSpec::utf8("exec_algorithm_id", true),
64 JsonFieldSpec::utf8_json("exec_algorithm_params", true),
65 JsonFieldSpec::utf8("exec_spawn_id", true),
66 JsonFieldSpec::utf8_json("tags", true),
67 JsonFieldSpec::utf8("init_id", false),
68 JsonFieldSpec::u64("ts_init", false),
69 JsonFieldSpec::u64("ts_last", false),
70 JsonFieldSpec::utf8("activation_price", true),
73];
74
75const POSITION_SNAPSHOT_FIELDS: &[JsonFieldSpec] = &[
76 JsonFieldSpec::utf8("trader_id", false),
77 JsonFieldSpec::utf8("strategy_id", false),
78 JsonFieldSpec::utf8("instrument_id", false),
79 JsonFieldSpec::utf8("position_id", false),
80 JsonFieldSpec::utf8("account_id", false),
81 JsonFieldSpec::utf8("opening_order_id", false),
82 JsonFieldSpec::utf8("closing_order_id", true),
83 JsonFieldSpec::utf8("entry", false),
84 JsonFieldSpec::utf8("side", false),
85 JsonFieldSpec::f64("signed_qty", false),
86 JsonFieldSpec::utf8("quantity", false),
87 JsonFieldSpec::utf8("peak_qty", false),
88 JsonFieldSpec::utf8("quote_currency", false),
89 JsonFieldSpec::utf8("base_currency", true),
90 JsonFieldSpec::utf8("settlement_currency", false),
91 JsonFieldSpec::f64("avg_px_open", false),
92 JsonFieldSpec::f64("avg_px_close", true),
93 JsonFieldSpec::f64("realized_return", true),
94 JsonFieldSpec::utf8("realized_pnl", true),
95 JsonFieldSpec::utf8("unrealized_pnl", true),
96 JsonFieldSpec::utf8_json("commissions", false),
97 JsonFieldSpec::u64("duration_ns", true),
98 JsonFieldSpec::u64("ts_opened", false),
99 JsonFieldSpec::u64("ts_closed", true),
100 JsonFieldSpec::u64("ts_init", false),
101 JsonFieldSpec::u64("ts_last", false),
102 JsonFieldSpec::utf8_json("replay_state", true),
103];
104
105fn instrument_metadata(type_name: &'static str, instrument_id: &str) -> HashMap<String, String> {
106 let mut metadata = metadata_for_type(type_name);
107 metadata.insert(KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string());
108 metadata
109}
110
111macro_rules! impl_snapshot_arrow {
112 ($type:ty, $type_name:expr, $fields:expr) => {
113 impl ArrowSchemaProvider for $type {
114 fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
115 schema_for_type($type_name, metadata, $fields)
116 }
117 }
118
119 impl EncodeToRecordBatch for $type {
120 fn encode_batch(
121 metadata: &HashMap<String, String>,
122 data: &[Self],
123 ) -> Result<RecordBatch, ArrowError> {
124 encode_batch($type_name, metadata, data, $fields)
125 }
126
127 fn metadata(&self) -> HashMap<String, String> {
128 instrument_metadata($type_name, &self.instrument_id.to_string())
129 }
130 }
131
132 impl DecodeTypedFromRecordBatch for $type {
133 fn decode_typed_batch(
134 metadata: &HashMap<String, String>,
135 record_batch: RecordBatch,
136 ) -> Result<Vec<Self>, EncodingError> {
137 decode_batch(metadata, &record_batch, $fields, Some($type_name))
138 }
139 }
140 };
141}
142
143impl_snapshot_arrow!(OrderSnapshot, "OrderSnapshot", ORDER_SNAPSHOT_FIELDS);
144impl_snapshot_arrow!(
145 PositionSnapshot,
146 "PositionSnapshot",
147 POSITION_SNAPSHOT_FIELDS
148);
149
150#[cfg(test)]
151mod tests {
152 use std::str::FromStr;
153
154 use arrow::datatypes::DataType;
155 use nautilus_core::UnixNanos;
156 use nautilus_model::{
157 enums::{OrderSide, OrderType, PositionSide},
158 identifiers::{AccountId, ClientOrderId, InstrumentId, PositionId, StrategyId, TraderId},
159 orders::OrderTestBuilder,
160 types::{Currency, Money, Price, Quantity},
161 };
162 use rstest::rstest;
163 use rust_decimal::Decimal;
164 use rust_decimal_macros::dec;
165
166 use super::*;
167
168 #[rstest]
169 fn test_order_snapshot_round_trip_preserves_decimal_precision() {
170 let order = OrderTestBuilder::new(OrderType::TrailingStopLimit)
171 .instrument_id(InstrumentId::from("BTCUSDT.BINANCE"))
172 .side(OrderSide::Buy)
173 .price(Price::from("50000"))
174 .trigger_price(Price::from("50500"))
175 .limit_offset(Decimal::from_str("0.123456789123456789").unwrap())
176 .trailing_offset(Decimal::from_str("0.987654321987654321").unwrap())
177 .quantity(Quantity::from("0.5"))
178 .build();
179 let snapshot = OrderSnapshot::from(order);
180 let metadata = snapshot.metadata();
181 let batch =
182 OrderSnapshot::encode_batch(&metadata, std::slice::from_ref(&snapshot)).unwrap();
183 let decoded = OrderSnapshot::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
184
185 assert_eq!(decoded, vec![snapshot]);
186 }
187
188 fn make_order_snapshot(avg_px: Option<Decimal>, slippage: Option<Decimal>) -> OrderSnapshot {
189 let order = OrderTestBuilder::new(OrderType::Limit)
190 .instrument_id(InstrumentId::from("BTCUSDT.BINANCE"))
191 .side(OrderSide::Buy)
192 .price(Price::from("50000"))
193 .quantity(Quantity::from("0.5"))
194 .build();
195 let mut snapshot = OrderSnapshot::from(order);
196 snapshot.avg_px = avg_px;
197 snapshot.slippage = slippage;
198 snapshot
199 }
200
201 fn legacy_float64_fields() -> Vec<JsonFieldSpec> {
203 ORDER_SNAPSHOT_FIELDS
204 .iter()
205 .map(|spec| match spec.name {
206 "avg_px" | "slippage" => JsonFieldSpec::f64(spec.name, spec.nullable),
207 _ => *spec,
208 })
209 .collect()
210 }
211
212 #[rstest]
213 fn test_order_snapshot_round_trip_preserves_exact_avg_px_and_slippage() {
214 let snapshot = make_order_snapshot(
217 Some(Decimal::from_str("1.6666666666666666666666666667").unwrap()),
218 Some(Decimal::from_str("0.0000000000000000000000000001").unwrap()),
219 );
220 let metadata = snapshot.metadata();
221 let batch =
222 OrderSnapshot::encode_batch(&metadata, std::slice::from_ref(&snapshot)).unwrap();
223
224 let avg_px_field = batch.schema().field_with_name("avg_px").unwrap().clone();
225 let slippage_field = batch.schema().field_with_name("slippage").unwrap().clone();
226 let decoded = OrderSnapshot::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
227
228 assert_eq!(avg_px_field.data_type(), &DataType::Utf8);
229 assert_eq!(slippage_field.data_type(), &DataType::Utf8);
230 assert_eq!(decoded, vec![snapshot]);
231 }
232
233 #[rstest]
234 fn test_order_snapshot_round_trip_null_avg_px_and_slippage() {
235 let snapshot = make_order_snapshot(None, None);
236 let metadata = snapshot.metadata();
237 let batch =
238 OrderSnapshot::encode_batch(&metadata, std::slice::from_ref(&snapshot)).unwrap();
239 let decoded = OrderSnapshot::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
240
241 assert_eq!(decoded, vec![snapshot]);
242 }
243
244 #[rstest]
245 fn test_order_snapshot_decodes_legacy_float64_columns() {
246 let snapshot = make_order_snapshot(Some(dec!(1.07)), Some(dec!(0.07)));
249 let metadata = snapshot.metadata();
250 let legacy_batch = encode_batch(
251 "OrderSnapshot",
252 &metadata,
253 std::slice::from_ref(&snapshot),
254 &legacy_float64_fields(),
255 )
256 .unwrap();
257
258 assert_eq!(
259 legacy_batch
260 .schema()
261 .field_with_name("avg_px")
262 .unwrap()
263 .data_type(),
264 &DataType::Float64
265 );
266
267 let decoded =
268 OrderSnapshot::decode_typed_batch(legacy_batch.schema().metadata(), legacy_batch)
269 .unwrap();
270
271 assert_eq!(decoded, vec![snapshot]);
272 }
273
274 #[rstest]
275 fn test_order_snapshot_decodes_legacy_float64_null_columns() {
276 let snapshot = make_order_snapshot(None, None);
277 let metadata = snapshot.metadata();
278 let legacy_batch = encode_batch(
279 "OrderSnapshot",
280 &metadata,
281 std::slice::from_ref(&snapshot),
282 &legacy_float64_fields(),
283 )
284 .unwrap();
285 let decoded =
286 OrderSnapshot::decode_typed_batch(legacy_batch.schema().metadata(), legacy_batch)
287 .unwrap();
288
289 assert_eq!(decoded, vec![snapshot]);
290 }
291
292 fn make_position_snapshot() -> PositionSnapshot {
293 PositionSnapshot {
294 trader_id: TraderId::from("TRADER-001"),
295 strategy_id: StrategyId::from("EMA-CROSS"),
296 instrument_id: InstrumentId::from("EURUSD.SIM"),
297 position_id: PositionId::from("P-001"),
298 account_id: AccountId::from("SIM-001"),
299 opening_order_id: ClientOrderId::from("O-1"),
300 closing_order_id: Some(ClientOrderId::from("O-2")),
301 entry: OrderSide::Buy,
302 side: PositionSide::Long,
303 signed_qty: 100.0,
304 quantity: Quantity::from("100"),
305 peak_qty: Quantity::from("100"),
306 quote_currency: Currency::USD(),
307 base_currency: Some(Currency::EUR()),
308 settlement_currency: Currency::USD(),
309 avg_px_open: 1.0500,
310 avg_px_close: Some(1.0600),
311 realized_return: Some(0.0095),
312 realized_pnl: Some(Money::new(100.0, Currency::USD())),
313 unrealized_pnl: Some(Money::new(50.0, Currency::USD())),
314 commissions: vec![Money::new(2.0, Currency::USD())],
315 duration_ns: Some(3_600_000_000_000),
316 ts_opened: UnixNanos::from(1_000_000_000),
317 ts_closed: Some(UnixNanos::from(4_600_000_000)),
318 ts_init: UnixNanos::from(2_000_000_000),
319 ts_last: UnixNanos::from(4_600_000_000),
320 replay_state: None,
321 }
322 }
323
324 #[rstest]
325 fn test_position_snapshot_round_trip() {
326 let mut snapshot = make_position_snapshot();
327 snapshot.replay_state = Some(serde_json::json!({"fill_voids": []}));
328 let metadata = snapshot.metadata();
329 let batch =
330 PositionSnapshot::encode_batch(&metadata, std::slice::from_ref(&snapshot)).unwrap();
331 let decoded =
332 PositionSnapshot::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
333
334 assert_eq!(decoded, vec![snapshot]);
335 }
336
337 #[rstest]
338 fn test_position_snapshot_round_trip_null_optionals() {
339 let mut snapshot = make_position_snapshot();
340 snapshot.closing_order_id = None;
341 snapshot.base_currency = None;
342 snapshot.avg_px_close = None;
343 snapshot.realized_return = None;
344 snapshot.realized_pnl = None;
345 snapshot.unrealized_pnl = None;
346 snapshot.duration_ns = None;
347 snapshot.ts_closed = None;
348
349 let metadata = snapshot.metadata();
350 let batch =
351 PositionSnapshot::encode_batch(&metadata, std::slice::from_ref(&snapshot)).unwrap();
352 let decoded =
353 PositionSnapshot::decode_typed_batch(batch.schema().metadata(), batch).unwrap();
354
355 assert_eq!(decoded, vec![snapshot]);
356 }
357}