1use std::{borrow::Borrow, collections::HashMap, str::FromStr, sync::Arc};
19
20use arrow::{
21 array::{Array, 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 identifiers::{InstrumentId, Symbol},
29 instruments::equity::Equity,
30 types::{price::Price, quantity::Quantity},
31};
32use rust_decimal::Decimal;
33use ustr::Ustr;
34
35use super::validate_required_instrument_columns;
36use crate::arrow::{
37 ArrowSchemaProvider, EncodeToRecordBatch, EncodingError, KEY_INSTRUMENT_ID,
38 KEY_PRICE_PRECISION, extract_column, extract_column_by_name,
39 extract_optional_string_column_by_name, json_string_field, metadata_with_type_name,
40 optional_ustr_value, record_batch_with_timestamps, record_batch_with_u64_timestamps,
41 timestamp_data_type,
42};
43
44impl ArrowSchemaProvider for Equity {
45 fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
46 let fields = vec![
47 Field::new("id", DataType::Utf8, false),
48 Field::new("raw_symbol", DataType::Utf8, false),
49 Field::new("currency", DataType::Utf8, false),
50 Field::new("price_precision", DataType::UInt8, false),
51 Field::new("price_increment", DataType::Utf8, false),
52 Field::new("lot_size", DataType::Utf8, true), Field::new("isin", DataType::Utf8, true), 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("margin_init", DataType::Utf8, false),
59 Field::new("margin_maint", DataType::Utf8, false),
60 Field::new("tick_scheme", DataType::Utf8, true),
61 json_string_field("info", true),
62 Field::new("ts_event", timestamp_data_type(), false),
63 Field::new("ts_init", timestamp_data_type(), false),
64 ];
65
66 Schema::new_with_metadata(fields, metadata_with_type_name("Equity", metadata))
67 }
68}
69
70impl EncodeToRecordBatch for Equity {
71 fn encode_batch<T>(
72 #[allow(unused)] metadata: &HashMap<String, String>,
73 data: &[T],
74 ) -> Result<RecordBatch, ArrowError>
75 where
76 T: std::borrow::Borrow<Self>,
77 {
78 let mut id_builder = StringBuilder::new();
79 let mut raw_symbol_builder = StringBuilder::new();
80 let mut currency_builder = StringBuilder::new();
81 let mut price_precision_builder = UInt8Array::builder(data.len());
82 let mut price_increment_builder = StringBuilder::new();
83 let mut lot_size_builder = StringBuilder::new();
84 let mut isin_builder = StringBuilder::new();
85 let mut max_quantity_builder = StringBuilder::new();
86 let mut min_quantity_builder = StringBuilder::new();
87 let mut max_price_builder = StringBuilder::new();
88 let mut min_price_builder = StringBuilder::new();
89 let mut margin_init_builder = StringBuilder::new();
90 let mut margin_maint_builder = StringBuilder::new();
91 let mut tick_scheme_builder = StringBuilder::new();
92 let mut info_builder = StringBuilder::new();
93 let mut ts_event_builder = UInt64Array::builder(data.len());
94 let mut ts_init_builder = UInt64Array::builder(data.len());
95
96 for equity in data.iter().map(Borrow::borrow) {
97 id_builder.append_value(equity.id.to_string());
98 raw_symbol_builder.append_value(equity.raw_symbol);
99 currency_builder.append_value(equity.currency.to_string());
100 price_precision_builder.append_value(equity.price_precision);
101 price_increment_builder.append_value(equity.price_increment.to_string());
102
103 if let Some(lot_size) = equity.lot_size {
104 lot_size_builder.append_value(lot_size.to_string());
105 } else {
106 lot_size_builder.append_null();
107 }
108
109 if let Some(isin) = equity.isin {
110 isin_builder.append_value(isin);
111 } else {
112 isin_builder.append_null();
113 }
114
115 if let Some(max_qty) = equity.max_quantity {
116 max_quantity_builder.append_value(max_qty.to_string());
117 } else {
118 max_quantity_builder.append_null();
119 }
120
121 if let Some(min_qty) = equity.min_quantity {
122 min_quantity_builder.append_value(min_qty.to_string());
123 } else {
124 min_quantity_builder.append_null();
125 }
126
127 if let Some(max_p) = equity.max_price {
128 max_price_builder.append_value(max_p.to_string());
129 } else {
130 max_price_builder.append_null();
131 }
132
133 if let Some(min_p) = equity.min_price {
134 min_price_builder.append_value(min_p.to_string());
135 } else {
136 min_price_builder.append_null();
137 }
138
139 margin_init_builder.append_value(equity.margin_init.to_string());
140 margin_maint_builder.append_value(equity.margin_maint.to_string());
141
142 if let Some(tick_scheme) = equity.tick_scheme {
143 tick_scheme_builder.append_value(tick_scheme);
144 } else {
145 tick_scheme_builder.append_null();
146 }
147
148 if let Some(ref info) = equity.info {
149 match serde_json::to_string(info) {
150 Ok(json) => {
151 info_builder.append_value(json);
152 }
153 Err(e) => {
154 return Err(ArrowError::InvalidArgumentError(format!(
155 "Failed to serialize info dict to JSON: {e}"
156 )));
157 }
158 }
159 } else {
160 info_builder.append_null();
161 }
162
163 ts_event_builder.append_value(equity.ts_event.as_u64());
164 ts_init_builder.append_value(equity.ts_init.as_u64());
165 }
166
167 record_batch_with_timestamps(
168 Self::get_schema(Some(metadata.clone())).into(),
169 vec![
170 Arc::new(id_builder.finish()),
171 Arc::new(raw_symbol_builder.finish()),
172 Arc::new(currency_builder.finish()),
173 Arc::new(price_precision_builder.finish()),
174 Arc::new(price_increment_builder.finish()),
175 Arc::new(lot_size_builder.finish()),
176 Arc::new(isin_builder.finish()),
177 Arc::new(max_quantity_builder.finish()),
178 Arc::new(min_quantity_builder.finish()),
179 Arc::new(max_price_builder.finish()),
180 Arc::new(min_price_builder.finish()),
181 Arc::new(margin_init_builder.finish()),
182 Arc::new(margin_maint_builder.finish()),
183 Arc::new(tick_scheme_builder.finish()),
184 Arc::new(info_builder.finish()),
185 Arc::new(ts_event_builder.finish()),
186 Arc::new(ts_init_builder.finish()),
187 ],
188 )
189 }
190
191 fn metadata(&self) -> HashMap<String, String> {
192 let mut metadata = HashMap::new();
193 metadata.insert(KEY_INSTRUMENT_ID.to_string(), self.id.to_string());
194 metadata.insert(
195 KEY_PRICE_PRECISION.to_string(),
196 self.price_precision.to_string(),
197 );
198 metadata
199 }
200}
201
202pub fn decode_equity_batch(
212 metadata: &HashMap<String, String>,
213 record_batch: &RecordBatch,
214) -> Result<Vec<Equity>, EncodingError> {
215 validate_required_instrument_columns::<Equity>(metadata, record_batch)?;
216 let record_batch = record_batch_with_u64_timestamps(record_batch)?;
217 let record_batch = &record_batch;
218 let cols = record_batch.columns();
219 let num_rows = record_batch.num_rows();
220
221 let id_values = extract_column::<StringArray>(cols, "id", 0, DataType::Utf8)?;
223 let raw_symbol_values = extract_column::<StringArray>(cols, "raw_symbol", 1, DataType::Utf8)?;
224 let currency_values = extract_column::<StringArray>(cols, "currency", 2, DataType::Utf8)?;
225 let price_precision_values =
226 extract_column::<UInt8Array>(cols, "price_precision", 3, DataType::UInt8)?;
227 let price_increment_values =
228 extract_column::<StringArray>(cols, "price_increment", 4, DataType::Utf8)?;
229 let lot_size_values = cols
230 .get(5)
231 .ok_or_else(|| EncodingError::MissingColumn("lot_size", 5))?;
232 let isin_values = cols
233 .get(6)
234 .ok_or_else(|| EncodingError::MissingColumn("isin", 6))?;
235 let max_quantity_values = cols
236 .get(7)
237 .ok_or_else(|| EncodingError::MissingColumn("max_quantity", 7))?;
238 let min_quantity_values = cols
239 .get(8)
240 .ok_or_else(|| EncodingError::MissingColumn("min_quantity", 8))?;
241 let max_price_values = cols
242 .get(9)
243 .ok_or_else(|| EncodingError::MissingColumn("max_price", 9))?;
244 let min_price_values = cols
245 .get(10)
246 .ok_or_else(|| EncodingError::MissingColumn("min_price", 10))?;
247 let margin_init_values =
248 extract_column::<StringArray>(cols, "margin_init", 11, DataType::Utf8)?;
249 let margin_maint_values =
250 extract_column::<StringArray>(cols, "margin_maint", 12, DataType::Utf8)?;
251 let tick_scheme_values = extract_optional_string_column_by_name(record_batch, "tick_scheme")?;
252 let info_values = extract_column_by_name::<StringArray>(record_batch, "info", DataType::Utf8)?;
253 let ts_event_values =
254 extract_column_by_name::<UInt64Array>(record_batch, "ts_event", DataType::UInt64)?;
255 let ts_init_values =
256 extract_column_by_name::<UInt64Array>(record_batch, "ts_init", DataType::UInt64)?;
257
258 let mut result = Vec::with_capacity(num_rows);
259
260 for i in 0..num_rows {
261 let id = InstrumentId::from_str(id_values.value(i))
262 .map_err(|e| EncodingError::ParseError("id", format!("row {i}: {e}")))?;
263 let raw_symbol = Symbol::from(raw_symbol_values.value(i));
264 let currency =
265 super::decode_currency(currency_values.value(i), "currency", "equity.currency", i)?;
266 let price_prec = price_precision_values.value(i);
267
268 let price_increment = Price::from_str(price_increment_values.value(i))
269 .map_err(|e| EncodingError::ParseError("price_increment", format!("row {i}: {e}")))?;
270
271 let lot_size = if lot_size_values.is_null(i) {
272 None
273 } else {
274 let lot_size_str = lot_size_values
275 .as_any()
276 .downcast_ref::<StringArray>()
277 .ok_or_else(|| {
278 EncodingError::ParseError("lot_size", format!("row {i}: invalid type"))
279 })?
280 .value(i);
281 Some(
282 Quantity::from_str(lot_size_str)
283 .map_err(|e| EncodingError::ParseError("lot_size", format!("row {i}: {e}")))?,
284 )
285 };
286
287 let isin = if isin_values.is_null(i) {
288 None
289 } else {
290 let isin_str = isin_values
291 .as_any()
292 .downcast_ref::<StringArray>()
293 .ok_or_else(|| EncodingError::ParseError("isin", format!("row {i}: invalid type")))?
294 .value(i);
295 Some(Ustr::from(isin_str))
296 };
297
298 let max_quantity =
299 if max_quantity_values.is_null(i) {
300 None
301 } else {
302 let max_qty_str = max_quantity_values
303 .as_any()
304 .downcast_ref::<StringArray>()
305 .ok_or_else(|| {
306 EncodingError::ParseError("max_quantity", format!("row {i}: invalid type"))
307 })?
308 .value(i);
309 Some(Quantity::from_str(max_qty_str).map_err(|e| {
310 EncodingError::ParseError("max_quantity", format!("row {i}: {e}"))
311 })?)
312 };
313
314 let min_quantity =
315 if min_quantity_values.is_null(i) {
316 None
317 } else {
318 let min_qty_str = min_quantity_values
319 .as_any()
320 .downcast_ref::<StringArray>()
321 .ok_or_else(|| {
322 EncodingError::ParseError("min_quantity", format!("row {i}: invalid type"))
323 })?
324 .value(i);
325 Some(Quantity::from_str(min_qty_str).map_err(|e| {
326 EncodingError::ParseError("min_quantity", format!("row {i}: {e}"))
327 })?)
328 };
329
330 let max_price = if max_price_values.is_null(i) {
331 None
332 } else {
333 let max_p_str = max_price_values
334 .as_any()
335 .downcast_ref::<StringArray>()
336 .ok_or_else(|| {
337 EncodingError::ParseError("max_price", format!("row {i}: invalid type"))
338 })?
339 .value(i);
340 Some(
341 Price::from_str(max_p_str)
342 .map_err(|e| EncodingError::ParseError("max_price", format!("row {i}: {e}")))?,
343 )
344 };
345
346 let min_price = if min_price_values.is_null(i) {
347 None
348 } else {
349 let min_p_str = min_price_values
350 .as_any()
351 .downcast_ref::<StringArray>()
352 .ok_or_else(|| {
353 EncodingError::ParseError("min_price", format!("row {i}: invalid type"))
354 })?
355 .value(i);
356 Some(
357 Price::from_str(min_p_str)
358 .map_err(|e| EncodingError::ParseError("min_price", format!("row {i}: {e}")))?,
359 )
360 };
361
362 let margin_init = Decimal::from_str(margin_init_values.value(i))
363 .map_err(|e| EncodingError::ParseError("margin_init", format!("row {i}: {e}")))?;
364 let margin_maint = Decimal::from_str(margin_maint_values.value(i))
365 .map_err(|e| EncodingError::ParseError("margin_maint", format!("row {i}: {e}")))?;
366
367 let info = if info_values.is_null(i) {
368 None
369 } else {
370 let info_json = info_values
371 .as_any()
372 .downcast_ref::<StringArray>()
373 .ok_or_else(|| EncodingError::ParseError("info", format!("row {i}: invalid type")))?
374 .value(i);
375
376 match serde_json::from_str::<Params>(info_json) {
377 Ok(info_dict) => Some(info_dict),
378 Err(e) => {
379 return Err(EncodingError::ParseError(
380 "info",
381 format!("row {i}: failed to deserialize JSON: {e}"),
382 ));
383 }
384 }
385 };
386
387 let ts_event = nautilus_core::UnixNanos::from(ts_event_values.value(i));
388 let ts_init = nautilus_core::UnixNanos::from(ts_init_values.value(i));
389
390 let tick_scheme = optional_ustr_value(tick_scheme_values, i);
391
392 let equity = Equity::builder()
393 .instrument_id(id)
394 .raw_symbol(raw_symbol)
395 .maybe_isin(isin)
396 .currency(currency)
397 .price_precision(price_prec)
398 .price_increment(price_increment)
399 .maybe_lot_size(lot_size)
400 .maybe_max_quantity(max_quantity)
401 .maybe_min_quantity(min_quantity)
402 .maybe_max_price(max_price)
403 .maybe_min_price(min_price)
404 .margin_init(margin_init)
405 .margin_maint(margin_maint)
406 .maybe_tick_scheme(tick_scheme)
407 .maybe_info(info)
408 .ts_event(ts_event)
409 .ts_init(ts_init)
410 .build()
411 .map_err(|e| super::instrument_validation_error::<Equity>(i, e))?;
412
413 result.push(equity);
414 }
415
416 Ok(result)
417}