1use std::{collections::HashMap, str::FromStr, sync::Arc};
19
20use arrow::{
21 array::{BinaryArray, BinaryBuilder, StringArray, StringBuilder, UInt8Array, UInt64Array},
22 datatypes::{DataType, Field, Schema},
23 error::ArrowError,
24 record_batch::RecordBatch,
25};
26use nautilus_core::Params;
27use nautilus_model::{
28 enums::AssetClass,
29 identifiers::{InstrumentId, Symbol},
30 instruments::option_spread::OptionSpread,
31 types::{price::Price, quantity::Quantity},
32};
33#[allow(unused)]
34use rust_decimal::Decimal;
35#[allow(unused)]
36use serde_json::Value;
37use ustr::Ustr;
38
39use crate::arrow::{
40 ArrowSchemaProvider, EncodeToRecordBatch, EncodingError, KEY_INSTRUMENT_ID,
41 KEY_PRICE_PRECISION, KEY_SIZE_PRECISION, extract_column,
42};
43
44fn asset_class_to_string(ac: AssetClass) -> String {
46 match ac {
47 AssetClass::FX => "FX".to_string(),
48 AssetClass::Equity => "Equity".to_string(),
49 AssetClass::Commodity => "Commodity".to_string(),
50 AssetClass::Debt => "Debt".to_string(),
51 AssetClass::Index => "Index".to_string(),
52 AssetClass::Cryptocurrency => "Cryptocurrency".to_string(),
53 AssetClass::Alternative => "Alternative".to_string(),
54 }
55}
56
57fn asset_class_from_str(s: &str) -> Result<AssetClass, EncodingError> {
59 match s {
60 "FX" => Ok(AssetClass::FX),
61 "Equity" => Ok(AssetClass::Equity),
62 "Commodity" => Ok(AssetClass::Commodity),
63 "Debt" => Ok(AssetClass::Debt),
64 "Index" => Ok(AssetClass::Index),
65 "Cryptocurrency" => Ok(AssetClass::Cryptocurrency),
66 "Alternative" => Ok(AssetClass::Alternative),
67 _ => Err(EncodingError::ParseError(
68 "asset_class",
69 format!("Unknown asset class: {s}"),
70 )),
71 }
72}
73
74impl ArrowSchemaProvider for OptionSpread {
75 fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
76 let fields = vec![
77 Field::new("id", DataType::Utf8, false),
78 Field::new("raw_symbol", DataType::Utf8, false),
79 Field::new("underlying", DataType::Utf8, false),
80 Field::new("strategy_type", DataType::Utf8, false),
81 Field::new("asset_class", DataType::Utf8, false),
82 Field::new("exchange", DataType::Utf8, true), Field::new("currency", DataType::Utf8, false),
84 Field::new("price_precision", DataType::UInt8, false),
85 Field::new("size_precision", DataType::UInt8, false),
86 Field::new("price_increment", DataType::Utf8, false),
87 Field::new("size_increment", DataType::Utf8, false),
88 Field::new("multiplier", DataType::Utf8, false),
89 Field::new("lot_size", DataType::Utf8, false),
90 Field::new("max_quantity", DataType::Utf8, true), Field::new("min_quantity", DataType::Utf8, true), Field::new("max_price", DataType::Utf8, true), Field::new("min_price", DataType::Utf8, true), Field::new("activation_ns", DataType::UInt64, false),
95 Field::new("expiration_ns", DataType::UInt64, false),
96 Field::new("margin_init", DataType::Utf8, false),
97 Field::new("margin_maint", DataType::Utf8, false),
98 Field::new("maker_fee", DataType::Utf8, false),
99 Field::new("taker_fee", DataType::Utf8, false),
100 Field::new("info", DataType::Binary, true), Field::new("ts_event", DataType::UInt64, false),
102 Field::new("ts_init", DataType::UInt64, false),
103 ];
104
105 let mut final_metadata = HashMap::new();
106 final_metadata.insert("class".to_string(), "OptionSpread".to_string());
107
108 if let Some(meta) = metadata {
109 final_metadata.extend(meta);
110 }
111
112 Schema::new_with_metadata(fields, final_metadata)
113 }
114}
115
116impl EncodeToRecordBatch for OptionSpread {
117 fn encode_batch(
118 #[allow(unused)] metadata: &HashMap<String, String>,
119 data: &[Self],
120 ) -> Result<RecordBatch, ArrowError> {
121 let mut id_builder = StringBuilder::new();
122 let mut raw_symbol_builder = StringBuilder::new();
123 let mut underlying_builder = StringBuilder::new();
124 let mut strategy_type_builder = StringBuilder::new();
125 let mut asset_class_builder = StringBuilder::new();
126 let mut exchange_builder = StringBuilder::new();
127 let mut currency_builder = StringBuilder::new();
128 let mut price_precision_builder = UInt8Array::builder(data.len());
129 let mut size_precision_builder = UInt8Array::builder(data.len());
130 let mut price_increment_builder = StringBuilder::new();
131 let mut size_increment_builder = StringBuilder::new();
132 let mut multiplier_builder = StringBuilder::new();
133 let mut lot_size_builder = StringBuilder::new();
134 let mut max_quantity_builder = StringBuilder::new();
135 let mut min_quantity_builder = StringBuilder::new();
136 let mut max_price_builder = StringBuilder::new();
137 let mut min_price_builder = StringBuilder::new();
138 let mut activation_ns_builder = UInt64Array::builder(data.len());
139 let mut expiration_ns_builder = UInt64Array::builder(data.len());
140 let mut margin_init_builder = StringBuilder::new();
141 let mut margin_maint_builder = StringBuilder::new();
142 let mut maker_fee_builder = StringBuilder::new();
143 let mut taker_fee_builder = StringBuilder::new();
144 let mut info_builder = BinaryBuilder::new();
145 let mut ts_event_builder = UInt64Array::builder(data.len());
146 let mut ts_init_builder = UInt64Array::builder(data.len());
147
148 for os in data {
149 id_builder.append_value(os.id.to_string());
150 raw_symbol_builder.append_value(os.raw_symbol);
151 underlying_builder.append_value(os.underlying);
152 strategy_type_builder.append_value(os.strategy_type);
153 asset_class_builder.append_value(asset_class_to_string(os.asset_class));
154
155 if let Some(exchange) = os.exchange {
156 exchange_builder.append_value(exchange);
157 } else {
158 exchange_builder.append_null();
159 }
160
161 currency_builder.append_value(os.currency.to_string());
162 price_precision_builder.append_value(os.price_precision);
163 size_precision_builder.append_value(os.size_precision);
164 price_increment_builder.append_value(os.price_increment.to_string());
165 size_increment_builder.append_value(os.size_increment.to_string());
166 multiplier_builder.append_value(os.multiplier.to_string());
167 lot_size_builder.append_value(os.lot_size.to_string());
168
169 if let Some(max_qty) = os.max_quantity {
170 max_quantity_builder.append_value(max_qty.to_string());
171 } else {
172 max_quantity_builder.append_null();
173 }
174
175 if let Some(min_qty) = os.min_quantity {
176 min_quantity_builder.append_value(min_qty.to_string());
177 } else {
178 min_quantity_builder.append_null();
179 }
180
181 if let Some(max_p) = os.max_price {
182 max_price_builder.append_value(max_p.to_string());
183 } else {
184 max_price_builder.append_null();
185 }
186
187 if let Some(min_p) = os.min_price {
188 min_price_builder.append_value(min_p.to_string());
189 } else {
190 min_price_builder.append_null();
191 }
192
193 activation_ns_builder.append_value(os.activation_ns.as_u64());
194 expiration_ns_builder.append_value(os.expiration_ns.as_u64());
195 margin_init_builder.append_value(os.margin_init.to_string());
196 margin_maint_builder.append_value(os.margin_maint.to_string());
197 maker_fee_builder.append_value(os.maker_fee.to_string());
198 taker_fee_builder.append_value(os.taker_fee.to_string());
199
200 if let Some(ref info) = os.info {
202 match serde_json::to_vec(info) {
203 Ok(json_bytes) => {
204 info_builder.append_value(json_bytes);
205 }
206 Err(e) => {
207 return Err(ArrowError::InvalidArgumentError(format!(
208 "Failed to serialize info dict to JSON: {e}"
209 )));
210 }
211 }
212 } else {
213 info_builder.append_null();
214 }
215
216 ts_event_builder.append_value(os.ts_event.as_u64());
217 ts_init_builder.append_value(os.ts_init.as_u64());
218 }
219
220 let mut final_metadata = metadata.clone();
221 final_metadata.insert("class".to_string(), "OptionSpread".to_string());
222
223 RecordBatch::try_new(
224 Self::get_schema(Some(final_metadata)).into(),
225 vec![
226 Arc::new(id_builder.finish()),
227 Arc::new(raw_symbol_builder.finish()),
228 Arc::new(underlying_builder.finish()),
229 Arc::new(strategy_type_builder.finish()),
230 Arc::new(asset_class_builder.finish()),
231 Arc::new(exchange_builder.finish()),
232 Arc::new(currency_builder.finish()),
233 Arc::new(price_precision_builder.finish()),
234 Arc::new(size_precision_builder.finish()),
235 Arc::new(price_increment_builder.finish()),
236 Arc::new(size_increment_builder.finish()),
237 Arc::new(multiplier_builder.finish()),
238 Arc::new(lot_size_builder.finish()),
239 Arc::new(max_quantity_builder.finish()),
240 Arc::new(min_quantity_builder.finish()),
241 Arc::new(max_price_builder.finish()),
242 Arc::new(min_price_builder.finish()),
243 Arc::new(activation_ns_builder.finish()),
244 Arc::new(expiration_ns_builder.finish()),
245 Arc::new(margin_init_builder.finish()),
246 Arc::new(margin_maint_builder.finish()),
247 Arc::new(maker_fee_builder.finish()),
248 Arc::new(taker_fee_builder.finish()),
249 Arc::new(info_builder.finish()),
250 Arc::new(ts_event_builder.finish()),
251 Arc::new(ts_init_builder.finish()),
252 ],
253 )
254 }
255
256 fn metadata(&self) -> HashMap<String, String> {
257 let mut metadata = HashMap::new();
258 metadata.insert(KEY_INSTRUMENT_ID.to_string(), self.id.to_string());
259 metadata.insert(
260 KEY_PRICE_PRECISION.to_string(),
261 self.price_precision.to_string(),
262 );
263 metadata.insert(
264 KEY_SIZE_PRECISION.to_string(),
265 self.size_precision.to_string(),
266 );
267 metadata
268 }
269}
270
271pub fn decode_option_spread_batch(
278 #[allow(unused)] metadata: &HashMap<String, String>,
279 record_batch: &RecordBatch,
280) -> Result<Vec<OptionSpread>, EncodingError> {
281 let cols = record_batch.columns();
282 let num_rows = record_batch.num_rows();
283
284 let id_values = extract_column::<StringArray>(cols, "id", 0, DataType::Utf8)?;
285 let raw_symbol_values = extract_column::<StringArray>(cols, "raw_symbol", 1, DataType::Utf8)?;
286 let underlying_values = extract_column::<StringArray>(cols, "underlying", 2, DataType::Utf8)?;
287 let strategy_type_values =
288 extract_column::<StringArray>(cols, "strategy_type", 3, DataType::Utf8)?;
289 let asset_class_values = extract_column::<StringArray>(cols, "asset_class", 4, DataType::Utf8)?;
290 let exchange_values = cols
291 .get(5)
292 .ok_or_else(|| EncodingError::MissingColumn("exchange", 5))?;
293 let currency_values = extract_column::<StringArray>(cols, "currency", 6, DataType::Utf8)?;
294 let price_precision_values =
295 extract_column::<UInt8Array>(cols, "price_precision", 7, DataType::UInt8)?;
296 let size_precision_values =
297 extract_column::<UInt8Array>(cols, "size_precision", 8, DataType::UInt8)?;
298 let price_increment_values =
299 extract_column::<StringArray>(cols, "price_increment", 9, DataType::Utf8)?;
300 let size_increment_values =
301 extract_column::<StringArray>(cols, "size_increment", 10, DataType::Utf8)?;
302 let multiplier_values = extract_column::<StringArray>(cols, "multiplier", 11, DataType::Utf8)?;
303 let lot_size_values = extract_column::<StringArray>(cols, "lot_size", 12, DataType::Utf8)?;
304 let max_quantity_values = cols
305 .get(13)
306 .ok_or_else(|| EncodingError::MissingColumn("max_quantity", 13))?;
307 let min_quantity_values = cols
308 .get(14)
309 .ok_or_else(|| EncodingError::MissingColumn("min_quantity", 14))?;
310 let max_price_values = cols
311 .get(15)
312 .ok_or_else(|| EncodingError::MissingColumn("max_price", 15))?;
313 let min_price_values = cols
314 .get(16)
315 .ok_or_else(|| EncodingError::MissingColumn("min_price", 16))?;
316 let activation_ns_values =
317 extract_column::<UInt64Array>(cols, "activation_ns", 17, DataType::UInt64)?;
318 let expiration_ns_values =
319 extract_column::<UInt64Array>(cols, "expiration_ns", 18, DataType::UInt64)?;
320 let margin_init_values =
321 extract_column::<StringArray>(cols, "margin_init", 19, DataType::Utf8)?;
322 let margin_maint_values =
323 extract_column::<StringArray>(cols, "margin_maint", 20, DataType::Utf8)?;
324 let maker_fee_values = extract_column::<StringArray>(cols, "maker_fee", 21, DataType::Utf8)?;
325 let taker_fee_values = extract_column::<StringArray>(cols, "taker_fee", 22, DataType::Utf8)?;
326 let info_values = cols
327 .get(23)
328 .ok_or_else(|| EncodingError::MissingColumn("info", 23))?;
329 let ts_event_values = extract_column::<UInt64Array>(cols, "ts_event", 24, DataType::UInt64)?;
330 let ts_init_values = extract_column::<UInt64Array>(cols, "ts_init", 25, DataType::UInt64)?;
331
332 let mut result = Vec::with_capacity(num_rows);
333
334 for i in 0..num_rows {
335 let id = InstrumentId::from_str(id_values.value(i))
336 .map_err(|e| EncodingError::ParseError("id", format!("row {i}: {e}")))?;
337 let raw_symbol = Symbol::from(raw_symbol_values.value(i));
338 let underlying = Ustr::from(underlying_values.value(i));
339 let strategy_type = Ustr::from(strategy_type_values.value(i));
340 let asset_class = asset_class_from_str(asset_class_values.value(i))?;
341
342 let exchange = if exchange_values.is_null(i) {
343 None
344 } else {
345 let exchange_str = exchange_values
346 .as_any()
347 .downcast_ref::<StringArray>()
348 .ok_or_else(|| {
349 EncodingError::ParseError("exchange", format!("row {i}: invalid type"))
350 })?
351 .value(i);
352 Some(Ustr::from(exchange_str))
353 };
354
355 let currency = super::decode_currency(
356 currency_values.value(i),
357 "currency",
358 "option_spread.currency",
359 i,
360 )?;
361 let price_prec = price_precision_values.value(i);
362 let _size_prec = size_precision_values.value(i); let price_increment = Price::from_str(price_increment_values.value(i))
365 .map_err(|e| EncodingError::ParseError("price_increment", format!("row {i}: {e}")))?;
366 let _size_increment = Quantity::from_str(size_increment_values.value(i))
367 .map_err(|e| EncodingError::ParseError("size_increment", format!("row {i}: {e}")))?; let multiplier = Quantity::from_str(multiplier_values.value(i))
369 .map_err(|e| EncodingError::ParseError("multiplier", format!("row {i}: {e}")))?;
370 let lot_size = Quantity::from_str(lot_size_values.value(i))
371 .map_err(|e| EncodingError::ParseError("lot_size", format!("row {i}: {e}")))?;
372
373 let max_quantity =
374 if max_quantity_values.is_null(i) {
375 None
376 } else {
377 let max_qty_str = max_quantity_values
378 .as_any()
379 .downcast_ref::<StringArray>()
380 .ok_or_else(|| {
381 EncodingError::ParseError("max_quantity", format!("row {i}: invalid type"))
382 })?
383 .value(i);
384 Some(Quantity::from_str(max_qty_str).map_err(|e| {
385 EncodingError::ParseError("max_quantity", format!("row {i}: {e}"))
386 })?)
387 };
388
389 let min_quantity =
390 if min_quantity_values.is_null(i) {
391 None
392 } else {
393 let min_qty_str = min_quantity_values
394 .as_any()
395 .downcast_ref::<StringArray>()
396 .ok_or_else(|| {
397 EncodingError::ParseError("min_quantity", format!("row {i}: invalid type"))
398 })?
399 .value(i);
400 Some(Quantity::from_str(min_qty_str).map_err(|e| {
401 EncodingError::ParseError("min_quantity", format!("row {i}: {e}"))
402 })?)
403 };
404
405 let max_price = if max_price_values.is_null(i) {
406 None
407 } else {
408 let max_p_str = max_price_values
409 .as_any()
410 .downcast_ref::<StringArray>()
411 .ok_or_else(|| {
412 EncodingError::ParseError("max_price", format!("row {i}: invalid type"))
413 })?
414 .value(i);
415 Some(
416 Price::from_str(max_p_str)
417 .map_err(|e| EncodingError::ParseError("max_price", format!("row {i}: {e}")))?,
418 )
419 };
420
421 let min_price = if min_price_values.is_null(i) {
422 None
423 } else {
424 let min_p_str = min_price_values
425 .as_any()
426 .downcast_ref::<StringArray>()
427 .ok_or_else(|| {
428 EncodingError::ParseError("min_price", format!("row {i}: invalid type"))
429 })?
430 .value(i);
431 Some(
432 Price::from_str(min_p_str)
433 .map_err(|e| EncodingError::ParseError("min_price", format!("row {i}: {e}")))?,
434 )
435 };
436
437 let activation_ns = nautilus_core::UnixNanos::from(activation_ns_values.value(i));
438 let expiration_ns = nautilus_core::UnixNanos::from(expiration_ns_values.value(i));
439
440 let margin_init = Decimal::from_str(margin_init_values.value(i))
441 .map_err(|e| EncodingError::ParseError("margin_init", format!("row {i}: {e}")))?;
442 let margin_maint = Decimal::from_str(margin_maint_values.value(i))
443 .map_err(|e| EncodingError::ParseError("margin_maint", format!("row {i}: {e}")))?;
444 let maker_fee = Decimal::from_str(maker_fee_values.value(i))
445 .map_err(|e| EncodingError::ParseError("maker_fee", format!("row {i}: {e}")))?;
446 let taker_fee = Decimal::from_str(taker_fee_values.value(i))
447 .map_err(|e| EncodingError::ParseError("taker_fee", format!("row {i}: {e}")))?;
448
449 let info = if info_values.is_null(i) {
451 None
452 } else {
453 let info_bytes = info_values
454 .as_any()
455 .downcast_ref::<BinaryArray>()
456 .ok_or_else(|| EncodingError::ParseError("info", format!("row {i}: invalid type")))?
457 .value(i);
458
459 match serde_json::from_slice::<Params>(info_bytes) {
460 Ok(info_dict) => Some(info_dict),
461 Err(e) => {
462 return Err(EncodingError::ParseError(
463 "info",
464 format!("row {i}: failed to deserialize JSON: {e}"),
465 ));
466 }
467 }
468 };
469
470 let ts_event = nautilus_core::UnixNanos::from(ts_event_values.value(i));
471 let ts_init = nautilus_core::UnixNanos::from(ts_init_values.value(i));
472
473 let option_spread = OptionSpread::new(
474 id,
475 raw_symbol,
476 asset_class,
477 exchange,
478 underlying,
479 strategy_type,
480 activation_ns,
481 expiration_ns,
482 currency,
483 price_prec,
484 price_increment,
485 multiplier,
486 lot_size,
487 max_quantity,
488 min_quantity,
489 max_price,
490 min_price,
491 Some(margin_init),
492 Some(margin_maint),
493 Some(maker_fee),
494 Some(taker_fee),
495 info,
496 ts_event,
497 ts_init,
498 );
499
500 result.push(option_spread);
501 }
502
503 Ok(result)
504}