Skip to main content

nautilus_tardis/
replay.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use 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    /// Cursor date UTC.
50    date_utc: Date,
51    /// Cursor end timestamp UNIX nanoseconds.
52    end_ns: UnixNanos,
53}
54
55impl DateCursor {
56    /// Creates a new [`DateCursor`] instance.
57    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        // Calculate end of the current UTC day
62        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
75/// Runs the Tardis Machine replay from a JSON configuration file.
76///
77/// # Errors
78///
79/// Returns an error if reading or parsing the config file fails,
80/// or if any downstream replay operation fails.
81///
82/// # Panics
83///
84/// Panics if unable to determine the output path (current directory fallback fails).
85pub 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    // Load and parse the replay configuration
90    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    // Initialize date cursors
166    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    // Initialize date collection maps
174    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    // Iterate through every remaining type and instrument sequentially
272
273    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 &quotes_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        // Update cursor
332        *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        // Update cursor
362        *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(&quote.instrument_id) {
383            batch_and_write_quotes(
384                &quotes_vec,
385                &quote.instrument_id,
386                cursor.date_utc,
387                path,
388                compression,
389            );
390        }
391        // Update cursor
392        *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        // Update cursor
422        *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        // Update cursor
446        *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        // Update cursor
476        *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
617/// Asserts that the given date is on or after the UNIX epoch (1970-01-01).
618///
619/// # Panics
620///
621/// Panics if the date is before 1970-01-01, as pre-epoch dates cannot be
622/// reliably represented as UnixNanos without overflow issues.
623fn 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
631/// Converts an ISO 8601 timestamp to a filesystem-safe format.
632///
633/// This function replaces colons and dots with hyphens to make the timestamp
634/// safe for use in filenames across different filesystems.
635fn iso_timestamp_to_file_timestamp(iso_timestamp: &str) -> String {
636    iso_timestamp.replace([':', '.'], "-")
637}
638
639/// Converts timestamps to a filename using ISO 8601 format.
640///
641/// This function converts two Unix nanosecond timestamps to a filename that uses
642/// ISO 8601 format with filesystem-safe characters, matching the catalog convention.
643fn 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    // Calculate start and end timestamps for the day (UTC)
675    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    // Start of day: 2024-01-01 00:00:00 UTC
769    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    // Midday: 2024-01-01 12:00:00 UTC
775    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    // End of day: 2024-01-01 23:59:59.999999999 UTC
781    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    // Start of new day: 2024-01-02 00:00:00 UTC
787    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 &quotes_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        // Trades must be written under the catalog convention path ("trades")
948        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        // Bars must be written under the catalog convention path ("bars"), consistent
1014        // with the other data types so they are discoverable by `ParquetDataCatalog`.
1015        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}