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