1use std::{collections::HashMap, sync::Arc};
17
18use arrow::{
19 array::{FixedSizeBinaryArray, FixedSizeBinaryBuilder, UInt64Array},
20 datatypes::{DataType, Field, Schema},
21 error::ArrowError,
22 record_batch::RecordBatch,
23};
24use nautilus_model::{data::QuoteTick, types::fixed::PRECISION_BYTES};
25
26use super::{
27 DecodeDataFromRecordBatch, EncodingError, decode_price, decode_quantity, extract_column,
28 parse_price_size_metadata, validate_precision_bytes,
29};
30use crate::arrow::{ArrowSchemaProvider, Data, DecodeFromRecordBatch, EncodeToRecordBatch};
31
32impl ArrowSchemaProvider for QuoteTick {
33 fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
34 let fields = vec![
35 Field::new(
36 "bid_price",
37 DataType::FixedSizeBinary(PRECISION_BYTES),
38 false,
39 ),
40 Field::new(
41 "ask_price",
42 DataType::FixedSizeBinary(PRECISION_BYTES),
43 false,
44 ),
45 Field::new(
46 "bid_size",
47 DataType::FixedSizeBinary(PRECISION_BYTES),
48 false,
49 ),
50 Field::new(
51 "ask_size",
52 DataType::FixedSizeBinary(PRECISION_BYTES),
53 false,
54 ),
55 Field::new("ts_event", DataType::UInt64, false),
56 Field::new("ts_init", DataType::UInt64, false),
57 ];
58
59 match metadata {
60 Some(metadata) => Schema::new_with_metadata(fields, metadata),
61 None => Schema::new(fields),
62 }
63 }
64}
65
66impl EncodeToRecordBatch for QuoteTick {
67 fn encode_batch(
68 metadata: &HashMap<String, String>,
69 data: &[Self],
70 ) -> Result<RecordBatch, ArrowError> {
71 let mut bid_price_builder =
72 FixedSizeBinaryBuilder::with_capacity(data.len(), PRECISION_BYTES);
73 let mut ask_price_builder =
74 FixedSizeBinaryBuilder::with_capacity(data.len(), PRECISION_BYTES);
75 let mut bid_size_builder =
76 FixedSizeBinaryBuilder::with_capacity(data.len(), PRECISION_BYTES);
77 let mut ask_size_builder =
78 FixedSizeBinaryBuilder::with_capacity(data.len(), PRECISION_BYTES);
79 let mut ts_event_builder = UInt64Array::builder(data.len());
80 let mut ts_init_builder = UInt64Array::builder(data.len());
81
82 for quote in data {
83 bid_price_builder
84 .append_value(quote.bid_price.raw().to_le_bytes())
85 .unwrap();
86 ask_price_builder
87 .append_value(quote.ask_price.raw().to_le_bytes())
88 .unwrap();
89 bid_size_builder
90 .append_value(quote.bid_size.raw().to_le_bytes())
91 .unwrap();
92 ask_size_builder
93 .append_value(quote.ask_size.raw().to_le_bytes())
94 .unwrap();
95 ts_event_builder.append_value(quote.ts_event.as_u64());
96 ts_init_builder.append_value(quote.ts_init.as_u64());
97 }
98
99 RecordBatch::try_new(
100 Self::get_schema(Some(metadata.clone())).into(),
101 vec![
102 Arc::new(bid_price_builder.finish()),
103 Arc::new(ask_price_builder.finish()),
104 Arc::new(bid_size_builder.finish()),
105 Arc::new(ask_size_builder.finish()),
106 Arc::new(ts_event_builder.finish()),
107 Arc::new(ts_init_builder.finish()),
108 ],
109 )
110 }
111
112 fn metadata(&self) -> HashMap<String, String> {
113 Self::get_metadata(
114 &self.instrument_id,
115 self.bid_price.precision,
116 self.bid_size.precision,
117 )
118 }
119}
120
121impl DecodeFromRecordBatch for QuoteTick {
122 fn decode_batch(
123 metadata: &HashMap<String, String>,
124 record_batch: RecordBatch,
125 ) -> Result<Vec<Self>, EncodingError> {
126 let (instrument_id, price_precision, size_precision) = parse_price_size_metadata(metadata)?;
127 let cols = record_batch.columns();
128
129 let bid_price_values = extract_column::<FixedSizeBinaryArray>(
130 cols,
131 "bid_price",
132 0,
133 DataType::FixedSizeBinary(PRECISION_BYTES),
134 )?;
135 let ask_price_values = extract_column::<FixedSizeBinaryArray>(
136 cols,
137 "ask_price",
138 1,
139 DataType::FixedSizeBinary(PRECISION_BYTES),
140 )?;
141 let bid_size_values = extract_column::<FixedSizeBinaryArray>(
142 cols,
143 "bid_size",
144 2,
145 DataType::FixedSizeBinary(PRECISION_BYTES),
146 )?;
147 let ask_size_values = extract_column::<FixedSizeBinaryArray>(
148 cols,
149 "ask_size",
150 3,
151 DataType::FixedSizeBinary(PRECISION_BYTES),
152 )?;
153 let ts_event_values = extract_column::<UInt64Array>(cols, "ts_event", 4, DataType::UInt64)?;
154 let ts_init_values = extract_column::<UInt64Array>(cols, "ts_init", 5, DataType::UInt64)?;
155
156 validate_precision_bytes(bid_price_values, "bid_price")?;
157 validate_precision_bytes(ask_price_values, "ask_price")?;
158 validate_precision_bytes(bid_size_values, "bid_size")?;
159 validate_precision_bytes(ask_size_values, "ask_size")?;
160
161 let result: Result<Vec<Self>, EncodingError> = (0..record_batch.num_rows())
162 .map(|row| {
163 let bid_price = decode_price(
164 bid_price_values.value(row),
165 price_precision,
166 "bid_price",
167 row,
168 )?;
169 let ask_price = decode_price(
170 ask_price_values.value(row),
171 price_precision,
172 "ask_price",
173 row,
174 )?;
175 let bid_size =
176 decode_quantity(bid_size_values.value(row), size_precision, "bid_size", row)?;
177 let ask_size =
178 decode_quantity(ask_size_values.value(row), size_precision, "ask_size", row)?;
179 Ok(Self {
180 instrument_id,
181 bid_price,
182 ask_price,
183 bid_size,
184 ask_size,
185 ts_event: ts_event_values.value(row).into(),
186 ts_init: ts_init_values.value(row).into(),
187 })
188 })
189 .collect();
190
191 result
192 }
193}
194
195impl DecodeDataFromRecordBatch for QuoteTick {
196 fn decode_data_batch(
197 metadata: &HashMap<String, String>,
198 record_batch: RecordBatch,
199 ) -> Result<Vec<Data>, EncodingError> {
200 let ticks: Vec<Self> = Self::decode_batch(metadata, record_batch)?;
201 Ok(ticks.into_iter().map(Data::from).collect())
202 }
203}
204
205#[cfg(test)]
206mod tests {
207 use std::{collections::HashMap, sync::Arc};
208
209 use arrow::{array::Array, record_batch::RecordBatch};
210 use nautilus_model::{
211 identifiers::InstrumentId,
212 types::{Price, Quantity, fixed::FIXED_SCALAR, price::PriceRaw, quantity::QuantityRaw},
213 };
214 use rstest::rstest;
215
216 use super::*;
217 use crate::arrow::{
218 KEY_INSTRUMENT_ID, KEY_PRICE_PRECISION, fixed_size_binary, get_raw_price, get_raw_quantity,
219 };
220
221 #[rstest]
222 fn test_get_schema() {
223 let instrument_id = InstrumentId::from("AAPL.XNAS");
224 let metadata = QuoteTick::get_metadata(&instrument_id, 2, 0);
225 let schema = QuoteTick::get_schema(Some(metadata.clone()));
226
227 let mut expected_fields = Vec::with_capacity(6);
228
229 expected_fields.push(Field::new(
230 "bid_price",
231 DataType::FixedSizeBinary(PRECISION_BYTES),
232 false,
233 ));
234 expected_fields.push(Field::new(
235 "ask_price",
236 DataType::FixedSizeBinary(PRECISION_BYTES),
237 false,
238 ));
239
240 expected_fields.extend(vec![
241 Field::new(
242 "bid_size",
243 DataType::FixedSizeBinary(PRECISION_BYTES),
244 false,
245 ),
246 Field::new(
247 "ask_size",
248 DataType::FixedSizeBinary(PRECISION_BYTES),
249 false,
250 ),
251 Field::new("ts_event", DataType::UInt64, false),
252 Field::new("ts_init", DataType::UInt64, false),
253 ]);
254
255 let expected_schema = Schema::new_with_metadata(expected_fields, metadata);
256 assert_eq!(schema, expected_schema);
257 }
258
259 #[rstest]
260 fn test_get_schema_map() {
261 let arrow_schema = QuoteTick::get_schema_map();
262 let mut expected_map = HashMap::new();
263
264 let fixed_size_binary = format!("FixedSizeBinary({PRECISION_BYTES})");
265 expected_map.insert("bid_price".to_string(), fixed_size_binary.clone());
266 expected_map.insert("ask_price".to_string(), fixed_size_binary.clone());
267 expected_map.insert("bid_size".to_string(), fixed_size_binary.clone());
268 expected_map.insert("ask_size".to_string(), fixed_size_binary);
269 expected_map.insert("ts_event".to_string(), "UInt64".to_string());
270 expected_map.insert("ts_init".to_string(), "UInt64".to_string());
271 assert_eq!(arrow_schema, expected_map);
272 }
273
274 #[rstest]
275 fn test_encode_quote_tick() {
276 let instrument_id = InstrumentId::from("AAPL.XNAS");
278 let tick1 = QuoteTick {
279 instrument_id,
280 bid_price: Price::from("100.10"),
281 ask_price: Price::from("101.50"),
282 bid_size: Quantity::from(1000),
283 ask_size: Quantity::from(500),
284 ts_event: 1.into(),
285 ts_init: 3.into(),
286 };
287
288 let tick2 = QuoteTick {
289 instrument_id,
290 bid_price: Price::from("100.75"),
291 ask_price: Price::from("100.20"),
292 bid_size: Quantity::from(750),
293 ask_size: Quantity::from(300),
294 ts_event: 2.into(),
295 ts_init: 4.into(),
296 };
297
298 let data = vec![tick1, tick2];
299 let metadata = QuoteTick::get_metadata(&instrument_id, 2, 0);
300 let record_batch = QuoteTick::encode_batch(&metadata, &data).unwrap();
301
302 let columns = record_batch.columns();
304
305 let bid_price_values = columns[0]
306 .as_any()
307 .downcast_ref::<FixedSizeBinaryArray>()
308 .unwrap();
309 let ask_price_values = columns[1]
310 .as_any()
311 .downcast_ref::<FixedSizeBinaryArray>()
312 .unwrap();
313 assert_eq!(
314 get_raw_price(bid_price_values.value(0)),
315 (100.10 * FIXED_SCALAR) as PriceRaw
316 );
317 assert_eq!(
318 get_raw_price(bid_price_values.value(1)),
319 (100.75 * FIXED_SCALAR) as PriceRaw
320 );
321 assert_eq!(
322 get_raw_price(ask_price_values.value(0)),
323 (101.50 * FIXED_SCALAR) as PriceRaw
324 );
325 assert_eq!(
326 get_raw_price(ask_price_values.value(1)),
327 (100.20 * FIXED_SCALAR) as PriceRaw
328 );
329
330 let bid_size_values = columns[2]
331 .as_any()
332 .downcast_ref::<FixedSizeBinaryArray>()
333 .unwrap();
334 let ask_size_values = columns[3]
335 .as_any()
336 .downcast_ref::<FixedSizeBinaryArray>()
337 .unwrap();
338 let ts_event_values = columns[4].as_any().downcast_ref::<UInt64Array>().unwrap();
339 let ts_init_values = columns[5].as_any().downcast_ref::<UInt64Array>().unwrap();
340
341 assert_eq!(columns.len(), 6);
342 assert_eq!(bid_size_values.len(), 2);
343 assert_eq!(
344 get_raw_quantity(bid_size_values.value(0)),
345 (1000.0 * FIXED_SCALAR) as QuantityRaw
346 );
347 assert_eq!(
348 get_raw_quantity(bid_size_values.value(1)),
349 (750.0 * FIXED_SCALAR) as QuantityRaw
350 );
351 assert_eq!(ask_size_values.len(), 2);
352 assert_eq!(
353 get_raw_quantity(ask_size_values.value(0)),
354 (500.0 * FIXED_SCALAR) as QuantityRaw
355 );
356 assert_eq!(
357 get_raw_quantity(ask_size_values.value(1)),
358 (300.0 * FIXED_SCALAR) as QuantityRaw
359 );
360 assert_eq!(ts_event_values.len(), 2);
361 assert_eq!(ts_event_values.value(0), 1);
362 assert_eq!(ts_event_values.value(1), 2);
363 assert_eq!(ts_init_values.len(), 2);
364 assert_eq!(ts_init_values.value(0), 3);
365 assert_eq!(ts_init_values.value(1), 4);
366 }
367
368 #[rstest]
369 fn test_decode_batch() {
370 let instrument_id = InstrumentId::from("AAPL.XNAS");
371 let metadata = QuoteTick::get_metadata(&instrument_id, 2, 0);
372
373 let raw_bid1 = (100.00 * FIXED_SCALAR) as PriceRaw;
374 let raw_bid2 = (99.00 * FIXED_SCALAR) as PriceRaw;
375 let raw_ask1 = (101.00 * FIXED_SCALAR) as PriceRaw;
376 let raw_ask2 = (100.00 * FIXED_SCALAR) as PriceRaw;
377
378 let (bid_price, ask_price) = (
379 fixed_size_binary(vec![&raw_bid1.to_le_bytes(), &raw_bid2.to_le_bytes()]),
380 fixed_size_binary(vec![&raw_ask1.to_le_bytes(), &raw_ask2.to_le_bytes()]),
381 );
382
383 let bid_size = fixed_size_binary(vec![
384 &((100.0 * FIXED_SCALAR) as PriceRaw).to_le_bytes(),
385 &((90.0 * FIXED_SCALAR) as PriceRaw).to_le_bytes(),
386 ]);
387 let ask_size = fixed_size_binary(vec![
388 &((110.0 * FIXED_SCALAR) as PriceRaw).to_le_bytes(),
389 &((100.0 * FIXED_SCALAR) as PriceRaw).to_le_bytes(),
390 ]);
391 let ts_event = UInt64Array::from(vec![1, 2]);
392 let ts_init = UInt64Array::from(vec![3, 4]);
393
394 let record_batch = RecordBatch::try_new(
395 QuoteTick::get_schema(Some(metadata.clone())).into(),
396 vec![
397 Arc::new(bid_price),
398 Arc::new(ask_price),
399 Arc::new(bid_size),
400 Arc::new(ask_size),
401 Arc::new(ts_event),
402 Arc::new(ts_init),
403 ],
404 )
405 .unwrap();
406
407 let decoded_data = QuoteTick::decode_batch(&metadata, record_batch).unwrap();
408 assert_eq!(decoded_data.len(), 2);
409
410 assert_eq!(decoded_data[0].bid_price, Price::from_raw(raw_bid1, 2));
412 assert_eq!(decoded_data[0].ask_price, Price::from_raw(raw_ask1, 2));
413 assert_eq!(decoded_data[1].bid_price, Price::from_raw(raw_bid2, 2));
414 assert_eq!(decoded_data[1].ask_price, Price::from_raw(raw_ask2, 2));
415 }
416
417 #[rstest]
418 fn test_decode_batch_invalid_bid_price_returns_error() {
419 let instrument_id = InstrumentId::from("AAPL.XNAS");
420 let metadata = QuoteTick::get_metadata(&instrument_id, 2, 0);
421
422 let invalid_price: PriceRaw = PriceRaw::MAX - 1000;
423 let valid_price = (100.00 * FIXED_SCALAR) as PriceRaw;
424
425 let bid_price = fixed_size_binary(vec![&invalid_price.to_le_bytes()]);
426 let ask_price = fixed_size_binary(vec![&valid_price.to_le_bytes()]);
427 let bid_size =
428 fixed_size_binary(vec![&((100.0 * FIXED_SCALAR) as QuantityRaw).to_le_bytes()]);
429 let ask_size =
430 fixed_size_binary(vec![&((100.0 * FIXED_SCALAR) as QuantityRaw).to_le_bytes()]);
431 let ts_event = UInt64Array::from(vec![1]);
432 let ts_init = UInt64Array::from(vec![2]);
433
434 let record_batch = RecordBatch::try_new(
435 QuoteTick::get_schema(Some(metadata.clone())).into(),
436 vec![
437 Arc::new(bid_price),
438 Arc::new(ask_price),
439 Arc::new(bid_size),
440 Arc::new(ask_size),
441 Arc::new(ts_event),
442 Arc::new(ts_init),
443 ],
444 )
445 .unwrap();
446
447 let result = QuoteTick::decode_batch(&metadata, record_batch);
448 assert!(result.is_err());
449 let err = result.unwrap_err();
450 assert!(
451 err.to_string().contains("bid_price") && err.to_string().contains("row 0"),
452 "Expected bid_price error at row 0, was: {err}"
453 );
454 }
455
456 #[rstest]
457 fn test_decode_batch_invalid_ask_size_returns_error() {
458 use nautilus_model::types::{fixed::FIXED_PRECISION, quantity::QUANTITY_RAW_MAX};
459
460 let instrument_id = InstrumentId::from("AAPL.XNAS");
461 let metadata = QuoteTick::get_metadata(&instrument_id, 2, FIXED_PRECISION);
464
465 let valid_price = (100.00 * FIXED_SCALAR) as PriceRaw;
466 let bid_price = fixed_size_binary(vec![&valid_price.to_le_bytes()]);
467 let ask_price = fixed_size_binary(vec![&valid_price.to_le_bytes()]);
468 let bid_size =
469 fixed_size_binary(vec![&((100.0 * FIXED_SCALAR) as QuantityRaw).to_le_bytes()]);
470
471 let invalid_size = QUANTITY_RAW_MAX + 1;
472 let ask_size = fixed_size_binary(vec![&invalid_size.to_le_bytes()]);
473 let ts_event = UInt64Array::from(vec![1]);
474 let ts_init = UInt64Array::from(vec![2]);
475
476 let record_batch = RecordBatch::try_new(
477 QuoteTick::get_schema(Some(metadata.clone())).into(),
478 vec![
479 Arc::new(bid_price),
480 Arc::new(ask_price),
481 Arc::new(bid_size),
482 Arc::new(ask_size),
483 Arc::new(ts_event),
484 Arc::new(ts_init),
485 ],
486 )
487 .unwrap();
488
489 let result = QuoteTick::decode_batch(&metadata, record_batch);
490 assert!(result.is_err());
491 let err = result.unwrap_err();
492 assert!(
493 err.to_string().contains("ask_size") && err.to_string().contains("row 0"),
494 "Expected ask_size error at row 0, was: {err}"
495 );
496 }
497
498 #[rstest]
499 fn test_decode_batch_missing_instrument_id_returns_error() {
500 let instrument_id = InstrumentId::from("AAPL.XNAS");
501 let mut metadata = QuoteTick::get_metadata(&instrument_id, 2, 0);
502 metadata.remove(KEY_INSTRUMENT_ID);
503
504 let valid_price = (100.00 * FIXED_SCALAR) as PriceRaw;
505 let bid_price = fixed_size_binary(vec![&valid_price.to_le_bytes()]);
506 let ask_price = fixed_size_binary(vec![&valid_price.to_le_bytes()]);
507 let bid_size =
508 fixed_size_binary(vec![&((100.0 * FIXED_SCALAR) as QuantityRaw).to_le_bytes()]);
509 let ask_size =
510 fixed_size_binary(vec![&((100.0 * FIXED_SCALAR) as QuantityRaw).to_le_bytes()]);
511 let ts_event = UInt64Array::from(vec![1]);
512 let ts_init = UInt64Array::from(vec![2]);
513
514 let record_batch = RecordBatch::try_new(
515 QuoteTick::get_schema(Some(metadata.clone())).into(),
516 vec![
517 Arc::new(bid_price),
518 Arc::new(ask_price),
519 Arc::new(bid_size),
520 Arc::new(ask_size),
521 Arc::new(ts_event),
522 Arc::new(ts_init),
523 ],
524 )
525 .unwrap();
526
527 let result = QuoteTick::decode_batch(&metadata, record_batch);
528 assert!(result.is_err());
529 let err = result.unwrap_err();
530 assert!(
531 err.to_string().contains("instrument_id"),
532 "Expected missing instrument_id error, was: {err}"
533 );
534 }
535
536 #[rstest]
537 fn test_decode_batch_missing_price_precision_returns_error() {
538 let instrument_id = InstrumentId::from("AAPL.XNAS");
539 let mut metadata = QuoteTick::get_metadata(&instrument_id, 2, 0);
540 metadata.remove(KEY_PRICE_PRECISION);
541
542 let valid_price = (100.00 * FIXED_SCALAR) as PriceRaw;
543 let bid_price = fixed_size_binary(vec![&valid_price.to_le_bytes()]);
544 let ask_price = fixed_size_binary(vec![&valid_price.to_le_bytes()]);
545 let bid_size =
546 fixed_size_binary(vec![&((100.0 * FIXED_SCALAR) as QuantityRaw).to_le_bytes()]);
547 let ask_size =
548 fixed_size_binary(vec![&((100.0 * FIXED_SCALAR) as QuantityRaw).to_le_bytes()]);
549 let ts_event = UInt64Array::from(vec![1]);
550 let ts_init = UInt64Array::from(vec![2]);
551
552 let record_batch = RecordBatch::try_new(
553 QuoteTick::get_schema(Some(metadata.clone())).into(),
554 vec![
555 Arc::new(bid_price),
556 Arc::new(ask_price),
557 Arc::new(bid_size),
558 Arc::new(ask_size),
559 Arc::new(ts_event),
560 Arc::new(ts_init),
561 ],
562 )
563 .unwrap();
564
565 let result = QuoteTick::decode_batch(&metadata, record_batch);
566 assert!(result.is_err());
567 let err = result.unwrap_err();
568 assert!(
569 err.to_string().contains("price_precision"),
570 "Expected missing price_precision error, was: {err}"
571 );
572 }
573
574 #[rstest]
575 fn test_encode_decode_round_trip() {
576 let instrument_id = InstrumentId::from("AAPL.XNAS");
577 let metadata = QuoteTick::get_metadata(&instrument_id, 2, 0);
578
579 let tick1 = QuoteTick {
580 instrument_id,
581 bid_price: Price::from("100.10"),
582 ask_price: Price::from("100.20"),
583 bid_size: Quantity::from(1000),
584 ask_size: Quantity::from(500),
585 ts_event: 1_000_000_000.into(),
586 ts_init: 1_000_000_001.into(),
587 };
588
589 let tick2 = QuoteTick {
590 instrument_id,
591 bid_price: Price::from("100.15"),
592 ask_price: Price::from("100.25"),
593 bid_size: Quantity::from(750),
594 ask_size: Quantity::from(250),
595 ts_event: 2_000_000_000.into(),
596 ts_init: 2_000_000_001.into(),
597 };
598
599 let original = vec![tick1, tick2];
600 let record_batch = QuoteTick::encode_batch(&metadata, &original).unwrap();
601 let decoded = QuoteTick::decode_batch(&metadata, record_batch).unwrap();
602
603 assert_eq!(decoded.len(), original.len());
604 for (orig, dec) in original.iter().zip(decoded.iter()) {
605 assert_eq!(dec.instrument_id, orig.instrument_id);
606 assert_eq!(dec.bid_price, orig.bid_price);
607 assert_eq!(dec.ask_price, orig.ask_price);
608 assert_eq!(dec.bid_size, orig.bid_size);
609 assert_eq!(dec.ask_size, orig.ask_size);
610 assert_eq!(dec.ts_event, orig.ts_event);
611 assert_eq!(dec.ts_init, orig.ts_init);
612 }
613 }
614}