1use std::{
17 fs,
18 path::{Path, PathBuf},
19};
20
21use ahash::{AHashMap, AHashSet};
22use anyhow::Context;
23use arrow::record_batch::RecordBatch;
24use futures_util::{StreamExt, pin_mut};
25use jiff::{Timestamp, civil::Date, tz::Offset};
26use nautilus_core::{UnixNanos, datetime::unix_nanos_to_iso8601, string::formatting::Separable};
27use nautilus_model::{
28 data::{
29 Bar, BarType, Data, OptionGreeks, OrderBookDelta, OrderBookDeltas, OrderBookDepth,
30 QuoteTick, TradeTick,
31 },
32 identifiers::InstrumentId,
33};
34use nautilus_persistence::common::paths::CatalogPathPrefix;
35use nautilus_serialization::arrow::{
36 bars_to_arrow_record_batch_bytes, book_deltas_to_arrow_record_batch_bytes,
37 book_depths_to_arrow_record_batch_bytes, option_greeks_to_arrow_record_batch_bytes,
38 quotes_to_arrow_record_batch_bytes, trades_to_arrow_record_batch_bytes,
39};
40use parquet::{arrow::ArrowWriter, basic::Compression, file::properties::WriterProperties};
41
42use crate::{
43 config::{BookSnapshotOutput, ParquetCompression, TardisReplayConfig},
44 http::TardisHttpClient,
45 machine::TardisMachineClient,
46};
47
48struct DateCursor {
49 date_utc: Date,
51 end_ns: UnixNanos,
53}
54
55impl DateCursor {
56 fn new(current_ns: UnixNanos) -> Self {
58 let current_utc = current_ns.to_datetime_utc();
59 let date_utc = Offset::UTC.to_datetime(current_utc).date();
60
61 let end_utc = utc_timestamp(date_utc, 23, 59, 59, 999_999_999);
63 let end_ns = UnixNanos::from(u64::try_from(end_utc.as_nanosecond()).unwrap_or(u64::MAX));
64
65 Self { date_utc, end_ns }
66 }
67}
68
69fn utc_timestamp(date: Date, hour: i8, minute: i8, second: i8, nanosecond: i32) -> Timestamp {
70 Offset::UTC
71 .to_timestamp(date.at(hour, minute, second, nanosecond))
72 .expect("valid UTC civil datetime")
73}
74
75pub async fn run_tardis_machine_replay_from_config(config_filepath: &Path) -> anyhow::Result<()> {
86 log::debug!("Starting replay");
87 log::debug!("Config filepath: {}", config_filepath.display());
88
89 let config_data = fs::read_to_string(config_filepath)
91 .with_context(|| format!("Failed to read config file: {}", config_filepath.display()))?;
92 let config: TardisReplayConfig = serde_json::from_str(&config_data)
93 .context("failed to parse config JSON into TardisReplayConfig")?;
94 config.validate()?;
95
96 let path = config
97 .output_path
98 .as_deref()
99 .map(Path::new)
100 .map(Path::to_path_buf)
101 .or_else(|| {
102 std::env::var("NAUTILUS_PATH")
103 .ok()
104 .map(|env_path| PathBuf::from(env_path).join("catalog").join("data"))
105 })
106 .unwrap_or_else(|| std::env::current_dir().expect("Failed to get current directory"));
107
108 log::debug!("Output path: {}", path.display());
109
110 let normalize_symbols = config.normalize_symbols.unwrap_or(true);
111 log::debug!("normalize_symbols={normalize_symbols}");
112
113 let book_snapshot_output = config
114 .book_snapshot_output
115 .clone()
116 .unwrap_or(BookSnapshotOutput::Deltas);
117 log::debug!("book_snapshot_output={book_snapshot_output:?}");
118
119 let extract_bbo_as_quotes = config.extract_bbo_as_quotes.unwrap_or(false);
120 log::debug!("extract_bbo_as_quotes={extract_bbo_as_quotes}");
121
122 let compression = config
123 .compression
124 .clone()
125 .unwrap_or(ParquetCompression::Zstd);
126 log::debug!("compression={compression:?}");
127 let compression = compression.as_parquet_compression();
128
129 let http_client = TardisHttpClient::new(
130 None,
131 config
132 .tardis_http_url
133 .as_ref()
134 .map(|value| value.expose_secret()),
135 None,
136 normalize_symbols,
137 config
138 .proxy_url
139 .as_ref()
140 .map(|value| value.expose_secret().to_owned()),
141 )?;
142 let mut machine_client = TardisMachineClient::new(
143 config
144 .tardis_ws_url
145 .as_ref()
146 .map(|value| value.expose_secret()),
147 normalize_symbols,
148 book_snapshot_output,
149 )?;
150 machine_client.extract_bbo_as_quotes = extract_bbo_as_quotes;
151
152 let exchanges: AHashSet<_> = config.options.iter().map(|opt| opt.exchange).collect();
153 let (instrument_map, _instruments) = http_client
154 .bootstrap_instruments(&exchanges)
155 .await
156 .context("failed to bootstrap instruments")?;
157
158 for (_, info) in &instrument_map {
159 machine_client.add_instrument_info((**info).clone());
160 }
161
162 log::debug!("Starting tardis-machine stream");
163 let stream = machine_client.replay(config.options).await?;
164 pin_mut!(stream);
165
166 let mut deltas_cursors: AHashMap<InstrumentId, DateCursor> = AHashMap::new();
168 let mut depths_cursors: AHashMap<InstrumentId, DateCursor> = AHashMap::new();
169 let mut quotes_cursors: AHashMap<InstrumentId, DateCursor> = AHashMap::new();
170 let mut trades_cursors: AHashMap<InstrumentId, DateCursor> = AHashMap::new();
171 let mut bars_cursors: AHashMap<BarType, DateCursor> = AHashMap::new();
172 let mut greeks_cursors: AHashMap<InstrumentId, DateCursor> = AHashMap::new();
173
174 let mut deltas_map: AHashMap<InstrumentId, Vec<OrderBookDelta>> = AHashMap::new();
176 let mut depths_map: AHashMap<InstrumentId, Vec<OrderBookDepth>> = AHashMap::new();
177 let mut quotes_map: AHashMap<InstrumentId, Vec<QuoteTick>> = AHashMap::new();
178 let mut trades_map: AHashMap<InstrumentId, Vec<TradeTick>> = AHashMap::new();
179 let mut bars_map: AHashMap<BarType, Vec<Bar>> = AHashMap::new();
180 let mut greeks_map: AHashMap<InstrumentId, Vec<OptionGreeks>> = AHashMap::new();
181
182 let mut msg_count = 0;
183
184 while let Some(result) = stream.next().await {
185 match result {
186 Ok(msg) => {
187 match msg {
188 Data::BookDelta(delta) => {
189 log::warn!(
190 "Skipping individual delta message for {} (use Deltas batch instead)",
191 delta.instrument_id
192 );
193 }
194 Data::BookDeltas(msg) => {
195 handle_deltas_msg(
196 &msg,
197 &mut deltas_map,
198 &mut deltas_cursors,
199 &path,
200 compression,
201 );
202 }
203 Data::BookDepth(msg) => {
204 handle_depth_msg(
205 *msg,
206 &mut depths_map,
207 &mut depths_cursors,
208 &path,
209 compression,
210 );
211 }
212 Data::Quote(msg) => {
213 handle_quote_msg(
214 msg,
215 &mut quotes_map,
216 &mut quotes_cursors,
217 &path,
218 compression,
219 );
220 }
221 Data::Trade(msg) => {
222 handle_trade_msg(
223 msg,
224 &mut trades_map,
225 &mut trades_cursors,
226 &path,
227 compression,
228 );
229 }
230 Data::Bar(msg) => {
231 handle_bar_msg(msg, &mut bars_map, &mut bars_cursors, &path, compression);
232 }
233 Data::MarkPrice(_) | Data::IndexPrice(_) | Data::FundingRate(_) => {
234 log::debug!(
235 "Skipping unsupported data type for instrument {}",
236 msg.instrument_id()
237 );
238 }
239 Data::OptionGreeks(msg) => {
240 handle_option_greeks_msg(
241 msg,
242 &mut greeks_map,
243 &mut greeks_cursors,
244 &path,
245 compression,
246 );
247 }
248 Data::InstrumentStatus(_) | Data::InstrumentClose(_) | Data::Custom(_) => {
249 log::debug!(
250 "Skipping unsupported data type for instrument {}",
251 msg.instrument_id()
252 );
253 }
254 #[allow(unreachable_patterns)]
255 _ => {
256 log::debug!("Skipping unsupported data type");
257 }
258 }
259
260 msg_count += 1;
261 if msg_count % 100_000 == 0 {
262 log::debug!("Processed {} messages", msg_count.separate_with_commas());
263 }
264 }
265 Err(e) => {
266 e.log_stream_error();
267 break;
268 }
269 }
270 }
271
272 for (instrument_id, deltas) in &deltas_map {
275 let cursor = deltas_cursors.get(instrument_id).expect("Expected cursor");
276 batch_and_write_deltas(deltas, instrument_id, cursor.date_utc, &path, compression);
277 }
278
279 for (instrument_id, depths) in &depths_map {
280 let cursor = depths_cursors.get(instrument_id).expect("Expected cursor");
281 batch_and_write_depths(depths, instrument_id, cursor.date_utc, &path, compression);
282 }
283
284 for (instrument_id, quotes) in "es_map {
285 let cursor = quotes_cursors.get(instrument_id).expect("Expected cursor");
286 batch_and_write_quotes(quotes, instrument_id, cursor.date_utc, &path, compression);
287 }
288
289 for (instrument_id, trades) in &trades_map {
290 let cursor = trades_cursors.get(instrument_id).expect("Expected cursor");
291 batch_and_write_trades(trades, instrument_id, cursor.date_utc, &path, compression);
292 }
293
294 for (bar_type, bars) in &bars_map {
295 let cursor = bars_cursors.get(bar_type).expect("Expected cursor");
296 batch_and_write_bars(bars, bar_type, cursor.date_utc, &path, compression);
297 }
298
299 for (instrument_id, greeks) in &greeks_map {
300 let cursor = greeks_cursors.get(instrument_id).expect("Expected cursor");
301 batch_and_write_greeks(greeks, instrument_id, cursor.date_utc, &path, compression);
302 }
303
304 log::debug!(
305 "Replay completed after {} messages",
306 msg_count.separate_with_commas()
307 );
308 Ok(())
309}
310
311fn handle_deltas_msg(
312 deltas: &OrderBookDeltas,
313 map: &mut AHashMap<InstrumentId, Vec<OrderBookDelta>>,
314 cursors: &mut AHashMap<InstrumentId, DateCursor>,
315 path: &Path,
316 compression: Compression,
317) {
318 let cursor = cursors
319 .entry(deltas.instrument_id)
320 .or_insert_with(|| DateCursor::new(deltas.ts_init));
321
322 if deltas.ts_init > cursor.end_ns {
323 if let Some(deltas_vec) = map.remove(&deltas.instrument_id) {
324 batch_and_write_deltas(
325 &deltas_vec,
326 &deltas.instrument_id,
327 cursor.date_utc,
328 path,
329 compression,
330 );
331 }
332 *cursor = DateCursor::new(deltas.ts_init);
334 }
335
336 map.entry(deltas.instrument_id)
337 .or_insert_with(|| Vec::with_capacity(100_000))
338 .extend(&*deltas.deltas);
339}
340
341fn handle_depth_msg(
342 depth: OrderBookDepth,
343 map: &mut AHashMap<InstrumentId, Vec<OrderBookDepth>>,
344 cursors: &mut AHashMap<InstrumentId, DateCursor>,
345 path: &Path,
346 compression: Compression,
347) {
348 let cursor = cursors
349 .entry(depth.instrument_id)
350 .or_insert_with(|| DateCursor::new(depth.ts_init));
351
352 if depth.ts_init > cursor.end_ns {
353 if let Some(depths_vec) = map.remove(&depth.instrument_id) {
354 batch_and_write_depths(
355 &depths_vec,
356 &depth.instrument_id,
357 cursor.date_utc,
358 path,
359 compression,
360 );
361 }
362 *cursor = DateCursor::new(depth.ts_init);
364 }
365
366 map.entry(depth.instrument_id)
367 .or_insert_with(|| Vec::with_capacity(100_000))
368 .push(depth);
369}
370
371fn handle_quote_msg(
372 quote: QuoteTick,
373 map: &mut AHashMap<InstrumentId, Vec<QuoteTick>>,
374 cursors: &mut AHashMap<InstrumentId, DateCursor>,
375 path: &Path,
376 compression: Compression,
377) {
378 let cursor = cursors
379 .entry(quote.instrument_id)
380 .or_insert_with(|| DateCursor::new(quote.ts_init));
381
382 if quote.ts_init > cursor.end_ns {
383 if let Some(quotes_vec) = map.remove("e.instrument_id) {
384 batch_and_write_quotes(
385 "es_vec,
386 "e.instrument_id,
387 cursor.date_utc,
388 path,
389 compression,
390 );
391 }
392 *cursor = DateCursor::new(quote.ts_init);
394 }
395
396 map.entry(quote.instrument_id)
397 .or_insert_with(|| Vec::with_capacity(100_000))
398 .push(quote);
399}
400
401fn handle_trade_msg(
402 trade: TradeTick,
403 map: &mut AHashMap<InstrumentId, Vec<TradeTick>>,
404 cursors: &mut AHashMap<InstrumentId, DateCursor>,
405 path: &Path,
406 compression: Compression,
407) {
408 let cursor = cursors
409 .entry(trade.instrument_id)
410 .or_insert_with(|| DateCursor::new(trade.ts_init));
411
412 if trade.ts_init > cursor.end_ns {
413 if let Some(trades_vec) = map.remove(&trade.instrument_id) {
414 batch_and_write_trades(
415 &trades_vec,
416 &trade.instrument_id,
417 cursor.date_utc,
418 path,
419 compression,
420 );
421 }
422 *cursor = DateCursor::new(trade.ts_init);
424 }
425
426 map.entry(trade.instrument_id)
427 .or_insert_with(|| Vec::with_capacity(100_000))
428 .push(trade);
429}
430
431fn handle_bar_msg(
432 bar: Bar,
433 map: &mut AHashMap<BarType, Vec<Bar>>,
434 cursors: &mut AHashMap<BarType, DateCursor>,
435 path: &Path,
436 compression: Compression,
437) {
438 let cursor = cursors
439 .entry(bar.bar_type)
440 .or_insert_with(|| DateCursor::new(bar.ts_init));
441
442 if bar.ts_init > cursor.end_ns {
443 if let Some(bars_vec) = map.remove(&bar.bar_type) {
444 batch_and_write_bars(&bars_vec, &bar.bar_type, cursor.date_utc, path, compression);
445 }
446 *cursor = DateCursor::new(bar.ts_init);
448 }
449
450 map.entry(bar.bar_type)
451 .or_insert_with(|| Vec::with_capacity(100_000))
452 .push(bar);
453}
454
455fn handle_option_greeks_msg(
456 greeks: OptionGreeks,
457 map: &mut AHashMap<InstrumentId, Vec<OptionGreeks>>,
458 cursors: &mut AHashMap<InstrumentId, DateCursor>,
459 path: &Path,
460 compression: Compression,
461) {
462 let cursor = cursors
463 .entry(greeks.instrument_id)
464 .or_insert_with(|| DateCursor::new(greeks.ts_init));
465
466 if greeks.ts_init > cursor.end_ns {
467 if let Some(greeks_vec) = map.remove(&greeks.instrument_id) {
468 batch_and_write_greeks(
469 &greeks_vec,
470 &greeks.instrument_id,
471 cursor.date_utc,
472 path,
473 compression,
474 );
475 }
476 *cursor = DateCursor::new(greeks.ts_init);
478 }
479
480 map.entry(greeks.instrument_id)
481 .or_insert_with(|| Vec::with_capacity(100_000))
482 .push(greeks);
483}
484
485fn batch_and_write_deltas(
486 deltas: &[OrderBookDelta],
487 instrument_id: &InstrumentId,
488 date: Date,
489 path: &Path,
490 compression: Compression,
491) {
492 match book_deltas_to_arrow_record_batch_bytes(deltas) {
493 Ok(batch) => write_batch(
494 &batch,
495 OrderBookDelta::path_prefix(),
496 instrument_id,
497 date,
498 path,
499 compression,
500 ),
501 Err(e) => {
502 log::error!("Error converting OrderBookDeltas to Arrow: {e:?}");
503 }
504 }
505}
506
507fn batch_and_write_depths(
508 depths: &[OrderBookDepth],
509 instrument_id: &InstrumentId,
510 date: Date,
511 path: &Path,
512 compression: Compression,
513) {
514 match book_depths_to_arrow_record_batch_bytes(depths) {
515 Ok(batch) => write_batch(
516 &batch,
517 OrderBookDepth::path_prefix(),
518 instrument_id,
519 date,
520 path,
521 compression,
522 ),
523 Err(e) => {
524 log::error!("Error converting OrderBookDepth to Arrow: {e:?}");
525 }
526 }
527}
528
529fn batch_and_write_quotes(
530 quotes: &[QuoteTick],
531 instrument_id: &InstrumentId,
532 date: Date,
533 path: &Path,
534 compression: Compression,
535) {
536 match quotes_to_arrow_record_batch_bytes(quotes) {
537 Ok(batch) => write_batch(
538 &batch,
539 QuoteTick::path_prefix(),
540 instrument_id,
541 date,
542 path,
543 compression,
544 ),
545 Err(e) => {
546 log::error!("Error converting QuoteTick to Arrow: {e:?}");
547 }
548 }
549}
550
551fn batch_and_write_trades(
552 trades: &[TradeTick],
553 instrument_id: &InstrumentId,
554 date: Date,
555 path: &Path,
556 compression: Compression,
557) {
558 match trades_to_arrow_record_batch_bytes(trades) {
559 Ok(batch) => write_batch(
560 &batch,
561 TradeTick::path_prefix(),
562 instrument_id,
563 date,
564 path,
565 compression,
566 ),
567 Err(e) => {
568 log::error!("Error converting TradeTick to Arrow: {e:?}");
569 }
570 }
571}
572
573fn batch_and_write_bars(
574 bars: &[Bar],
575 bar_type: &BarType,
576 date: Date,
577 path: &Path,
578 compression: Compression,
579) {
580 let batch = match bars_to_arrow_record_batch_bytes(bars) {
581 Ok(batch) => batch,
582 Err(e) => {
583 log::error!("Error converting Bar to Arrow: {e:?}");
584 return;
585 }
586 };
587
588 let filepath = path.join(parquet_filepath_bars(bar_type, date));
589 if let Err(e) = write_parquet_local(&batch, &filepath, compression) {
590 log::error!("Error writing {}: {e}", filepath.display());
591 } else {
592 log::debug!("File written: {}", filepath.display());
593 }
594}
595
596fn batch_and_write_greeks(
597 greeks: &[OptionGreeks],
598 instrument_id: &InstrumentId,
599 date: Date,
600 path: &Path,
601 compression: Compression,
602) {
603 match option_greeks_to_arrow_record_batch_bytes(greeks) {
604 Ok(batch) => write_batch(
605 &batch,
606 OptionGreeks::path_prefix(),
607 instrument_id,
608 date,
609 path,
610 compression,
611 ),
612 Err(e) => {
613 log::error!("Error converting OptionGreeks to Arrow: {e:?}");
614 }
615 }
616}
617
618fn assert_post_epoch(date: Date) {
625 let epoch = Date::constant(1970, 1, 1);
626 assert!(
627 date >= epoch,
628 "Tardis replay filenames require dates on or after 1970-01-01; received {date}"
629 );
630}
631
632fn iso_timestamp_to_file_timestamp(iso_timestamp: &str) -> String {
637 iso_timestamp.replace([':', '.'], "-")
638}
639
640fn timestamps_to_filename(timestamp_1: UnixNanos, timestamp_2: UnixNanos) -> String {
645 let datetime_1 = iso_timestamp_to_file_timestamp(&unix_nanos_to_iso8601(timestamp_1));
646 let datetime_2 = iso_timestamp_to_file_timestamp(&unix_nanos_to_iso8601(timestamp_2));
647
648 format!("{datetime_1}_{datetime_2}.parquet")
649}
650
651fn parquet_filepath(typename: &str, instrument_id: &InstrumentId, date: Date) -> PathBuf {
652 assert_post_epoch(date);
653
654 let instrument_id_str = instrument_id.to_string().replace('/', "");
655
656 let start_nanos = utc_timestamp(date, 0, 0, 0, 0).as_nanosecond();
657 let end_nanos = utc_timestamp(date, 23, 59, 59, 999_999_999).as_nanosecond();
658
659 let filename = timestamps_to_filename(
660 UnixNanos::from(u64::try_from(start_nanos).expect("date fits UnixNanos")),
661 UnixNanos::from(u64::try_from(end_nanos).expect("date fits UnixNanos")),
662 );
663
664 PathBuf::new()
665 .join(typename)
666 .join(instrument_id_str)
667 .join(filename)
668}
669
670fn parquet_filepath_bars(bar_type: &BarType, date: Date) -> PathBuf {
671 assert_post_epoch(date);
672
673 let bar_type_str = bar_type.to_string().replace('/', "");
674
675 let start_nanos = utc_timestamp(date, 0, 0, 0, 0).as_nanosecond();
677 let end_nanos = utc_timestamp(date, 23, 59, 59, 999_999_999).as_nanosecond();
678
679 let filename = timestamps_to_filename(
680 UnixNanos::from(u64::try_from(start_nanos).expect("date fits UnixNanos")),
681 UnixNanos::from(u64::try_from(end_nanos).expect("date fits UnixNanos")),
682 );
683
684 PathBuf::new()
685 .join(Bar::path_prefix())
686 .join(bar_type_str)
687 .join(filename)
688}
689
690fn write_batch(
691 batch: &RecordBatch,
692 typename: &str,
693 instrument_id: &InstrumentId,
694 date: Date,
695 path: &Path,
696 compression: Compression,
697) {
698 let filepath = path.join(parquet_filepath(typename, instrument_id, date));
699 if let Err(e) = write_parquet_local(batch, &filepath, compression) {
700 log::error!("Error writing {}: {e}", filepath.display());
701 } else {
702 log::debug!("File written: {}", filepath.display());
703 }
704}
705
706fn write_parquet_local(
707 batch: &RecordBatch,
708 file_path: &Path,
709 compression: Compression,
710) -> anyhow::Result<()> {
711 if let Some(parent) = file_path.parent() {
712 std::fs::create_dir_all(parent)?;
713 }
714
715 let file = std::fs::File::create(file_path)?;
716 let props = WriterProperties::builder()
717 .set_compression(compression)
718 .build();
719
720 let mut writer = ArrowWriter::try_new(file, batch.schema(), Some(props))?;
721 writer.write(batch)?;
722 writer.close()?;
723 Ok(())
724}
725
726#[cfg(test)]
727mod tests {
728 use std::sync::Arc;
729
730 use nautilus_core::DurationNanos;
731 use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
732 use rstest::rstest;
733
734 use super::*;
735 use crate::{
736 common::{enums::TardisExchange, testing::load_test_json},
737 config::BookSnapshotOutput,
738 machine::{
739 message::{BookSnapshotMsg, OptionSummaryMsg, TradeMsg, WsMessage},
740 parse::parse_tardis_ws_message,
741 types::TardisInstrumentMiniInfo,
742 },
743 };
744
745 fn utc_nanos(
746 year: i16,
747 month: i8,
748 day: i8,
749 hour: i8,
750 minute: i8,
751 second: i8,
752 nanosecond: i32,
753 ) -> u64 {
754 u64::try_from(
755 utc_timestamp(
756 Date::new(year, month, day).unwrap(),
757 hour,
758 minute,
759 second,
760 nanosecond,
761 )
762 .as_nanosecond(),
763 )
764 .unwrap()
765 }
766
767 #[rstest]
768 #[case(
769 utc_nanos(2024, 1, 1, 0, 0, 0, 0),
771 Date::new(2024, 1, 1).unwrap(),
772 utc_nanos(2024, 1, 1, 23, 59, 59, 999_999_999)
773)]
774 #[case(
775 utc_nanos(2024, 1, 1, 12, 0, 0, 0),
777 Date::new(2024, 1, 1).unwrap(),
778 utc_nanos(2024, 1, 1, 23, 59, 59, 999_999_999)
779)]
780 #[case(
781 utc_nanos(2024, 1, 1, 23, 59, 59, 999_999_999),
783 Date::new(2024, 1, 1).unwrap(),
784 utc_nanos(2024, 1, 1, 23, 59, 59, 999_999_999)
785)]
786 #[case(
787 utc_nanos(2024, 1, 2, 0, 0, 0, 0),
789 Date::new(2024, 1, 2).unwrap(),
790 utc_nanos(2024, 1, 2, 23, 59, 59, 999_999_999)
791)]
792 fn test_date_cursor(
793 #[case] timestamp: u64,
794 #[case] expected_date: Date,
795 #[case] expected_end_ns: u64,
796 ) {
797 let unix_nanos = UnixNanos::from(timestamp);
798 let cursor = DateCursor::new(unix_nanos);
799
800 assert_eq!(cursor.date_utc, expected_date);
801 assert_eq!(cursor.end_ns, UnixNanos::from(expected_end_ns));
802 }
803
804 #[rstest]
805 fn test_option_greeks_replay_catalog_round_trip() {
806 let instrument_id = InstrumentId::from("BTC-28JUN24-70000-C.DERIBIT");
807 let compression = ParquetCompression::Zstd.as_parquet_compression();
808 let info = Arc::new(TardisInstrumentMiniInfo::new(
809 instrument_id,
810 None,
811 TardisExchange::Deribit,
812 4,
813 1,
814 ));
815
816 let option_summary: OptionSummaryMsg =
817 serde_json::from_str(&load_test_json("option_summary.json")).unwrap();
818 let Some(Data::OptionGreeks(greeks_1)) = parse_tardis_ws_message(
819 WsMessage::OptionSummary(option_summary),
820 &info,
821 &BookSnapshotOutput::Deltas,
822 ) else {
823 panic!("Expected option_summary to route to Data::OptionGreeks");
824 };
825
826 let mut greeks_2 = greeks_1;
827 greeks_2.ts_event = greeks_1.ts_event + DurationNanos::from_secs(1);
828 greeks_2.ts_init = greeks_1.ts_init + DurationNanos::from_secs(1);
829 greeks_2.greeks.delta = 0.26;
830
831 let option_quote: BookSnapshotMsg =
832 serde_json::from_str(&load_test_json("option_book_snapshot.json")).unwrap();
833 let Some(Data::Quote(quote_1)) = parse_tardis_ws_message(
834 WsMessage::BookSnapshot(option_quote),
835 &info,
836 &BookSnapshotOutput::Deltas,
837 ) else {
838 panic!("Expected depth-1 option book snapshot to route to Data::Quote");
839 };
840
841 let mut quote_2 = quote_1;
842 quote_2.ts_event = quote_1.ts_event + DurationNanos::from_secs(1);
843 quote_2.ts_init = quote_1.ts_init + DurationNanos::from_secs(1);
844
845 let temp_dir = tempfile::tempdir().unwrap();
846 let data_path = temp_dir.path().join("data");
847
848 let mut quotes_map: AHashMap<InstrumentId, Vec<QuoteTick>> = AHashMap::new();
849 let mut quotes_cursors: AHashMap<InstrumentId, DateCursor> = AHashMap::new();
850 let mut greeks_map: AHashMap<InstrumentId, Vec<OptionGreeks>> = AHashMap::new();
851 let mut greeks_cursors: AHashMap<InstrumentId, DateCursor> = AHashMap::new();
852
853 for quote in [quote_1, quote_2] {
854 handle_quote_msg(
855 quote,
856 &mut quotes_map,
857 &mut quotes_cursors,
858 &data_path,
859 compression,
860 );
861 }
862
863 for greeks in [greeks_1, greeks_2] {
864 handle_option_greeks_msg(
865 greeks,
866 &mut greeks_map,
867 &mut greeks_cursors,
868 &data_path,
869 compression,
870 );
871 }
872
873 for (id, quotes) in "es_map {
874 let cursor = quotes_cursors.get(id).expect("Expected cursor");
875 batch_and_write_quotes(quotes, id, cursor.date_utc, &data_path, compression);
876 }
877
878 for (id, greeks) in &greeks_map {
879 let cursor = greeks_cursors.get(id).expect("Expected cursor");
880 batch_and_write_greeks(greeks, id, cursor.date_utc, &data_path, compression);
881 }
882
883 let mut catalog = ParquetDataCatalog::new(temp_dir.path(), None, None, None, None);
884 let identifiers = Some(vec![instrument_id.to_string()]);
885
886 let quotes_out = catalog
887 .quote_ticks(identifiers.clone(), None, None)
888 .unwrap();
889 let greeks_out = catalog.option_greeks(identifiers, None, None).unwrap();
890
891 assert_eq!(greeks_out, vec![greeks_1, greeks_2]);
892 assert_eq!(greeks_out[0].instrument_id, instrument_id);
893 assert_eq!(greeks_out[0].mark_iv, Some(0.565));
894 assert_eq!(greeks_out[0].underlying_price, Some(63_500.0));
895 assert!(greeks_out[0].ts_init < greeks_out[1].ts_init);
896
897 assert_eq!(quotes_out, vec![quote_1, quote_2]);
898 assert_eq!(quotes_out[0].instrument_id, instrument_id);
899 assert!(quotes_out[0].ts_init < quotes_out[1].ts_init);
900 }
901
902 #[rstest]
903 fn test_trades_replay_catalog_round_trip() {
904 let instrument_id = InstrumentId::from("XBTUSD.BITMEX");
905 let compression = ParquetCompression::Zstd.as_parquet_compression();
906 let info = Arc::new(TardisInstrumentMiniInfo::new(
907 instrument_id,
908 None,
909 TardisExchange::Bitmex,
910 1,
911 0,
912 ));
913
914 let trade_msg: TradeMsg = serde_json::from_str(&load_test_json("trade.json")).unwrap();
915 let Some(Data::Trade(trade_1)) = parse_tardis_ws_message(
916 WsMessage::Trade(trade_msg),
917 &info,
918 &BookSnapshotOutput::Deltas,
919 ) else {
920 panic!("Expected trade message to route to Data::Trade");
921 };
922
923 let mut trade_2 = trade_1;
924 trade_2.ts_event = trade_1.ts_event + DurationNanos::from_secs(1);
925 trade_2.ts_init = trade_1.ts_init + DurationNanos::from_secs(1);
926
927 let temp_dir = tempfile::tempdir().unwrap();
928 let data_path = temp_dir.path().join("data");
929
930 let mut trades_map: AHashMap<InstrumentId, Vec<TradeTick>> = AHashMap::new();
931 let mut trades_cursors: AHashMap<InstrumentId, DateCursor> = AHashMap::new();
932
933 for trade in [trade_1, trade_2] {
934 handle_trade_msg(
935 trade,
936 &mut trades_map,
937 &mut trades_cursors,
938 &data_path,
939 compression,
940 );
941 }
942
943 for (id, trades) in &trades_map {
944 let cursor = trades_cursors.get(id).expect("Expected cursor");
945 batch_and_write_trades(trades, id, cursor.date_utc, &data_path, compression);
946 }
947
948 assert!(data_path.join(TradeTick::path_prefix()).exists());
950 assert!(!data_path.join("trade_tick").exists());
951
952 let mut catalog = ParquetDataCatalog::new(temp_dir.path(), None, None, None, None);
953 let trades_out = catalog
954 .trade_ticks(Some(vec![instrument_id.to_string()]), None, None)
955 .unwrap();
956
957 assert_eq!(trades_out, vec![trade_1, trade_2]);
958 assert_eq!(trades_out[0].instrument_id, instrument_id);
959 assert!(trades_out[0].ts_init < trades_out[1].ts_init);
960 }
961
962 #[rstest]
963 fn test_bars_replay_catalog_round_trip() {
964 use nautilus_model::types::{Price, Quantity};
965
966 let compression = ParquetCompression::Zstd.as_parquet_compression();
967 let bar_type = BarType::from("BTCUSDT-PERP.BINANCE-1-MINUTE-LAST-EXTERNAL");
968
969 let ts_1 = UnixNanos::from(utc_nanos(2024, 1, 1, 0, 0, 0, 0));
970 let ts_2 = ts_1 + DurationNanos::from_secs(1);
971
972 let bar_1 = Bar::new(
973 bar_type,
974 Price::from("100.00"),
975 Price::from("110.00"),
976 Price::from("90.00"),
977 Price::from("105.00"),
978 Quantity::from("1000"),
979 ts_1,
980 ts_1,
981 );
982 let bar_2 = Bar::new(
983 bar_type,
984 Price::from("105.00"),
985 Price::from("115.00"),
986 Price::from("95.00"),
987 Price::from("110.00"),
988 Quantity::from("1000"),
989 ts_2,
990 ts_2,
991 );
992
993 let temp_dir = tempfile::tempdir().unwrap();
994 let data_path = temp_dir.path().join("data");
995
996 let mut bars_map: AHashMap<BarType, Vec<Bar>> = AHashMap::new();
997 let mut bars_cursors: AHashMap<BarType, DateCursor> = AHashMap::new();
998
999 for bar in [bar_1, bar_2] {
1000 handle_bar_msg(
1001 bar,
1002 &mut bars_map,
1003 &mut bars_cursors,
1004 &data_path,
1005 compression,
1006 );
1007 }
1008
1009 for (bar_type, bars) in &bars_map {
1010 let cursor = bars_cursors.get(bar_type).expect("Expected cursor");
1011 batch_and_write_bars(bars, bar_type, cursor.date_utc, &data_path, compression);
1012 }
1013
1014 assert!(data_path.join(Bar::path_prefix()).exists());
1017 assert!(!data_path.join("bar").exists());
1018
1019 let mut catalog = ParquetDataCatalog::new(temp_dir.path(), None, None, None, None);
1020 let bars_out = catalog.bars(None, None, None).unwrap();
1021
1022 assert_eq!(bars_out, vec![bar_1, bar_2]);
1023 assert_eq!(bars_out[0].bar_type, bar_type);
1024 assert!(bars_out[0].ts_init < bars_out[1].ts_init);
1025 }
1026}