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, replay selections contain
80/// duplicates, or 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    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    // Initialize date cursors
167    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    // Initialize date collection maps
175    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    // Iterate through every remaining type and instrument sequentially
273
274    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 &quotes_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        // Update cursor
333        *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        // Update cursor
363        *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(&quote.instrument_id) {
384            batch_and_write_quotes(
385                &quotes_vec,
386                &quote.instrument_id,
387                cursor.date_utc,
388                path,
389                compression,
390            );
391        }
392        // Update cursor
393        *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        // Update cursor
423        *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        // Update cursor
447        *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        // Update cursor
477        *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
618/// Asserts that the given date is on or after the UNIX epoch (1970-01-01).
619///
620/// # Panics
621///
622/// Panics if the date is before 1970-01-01, as pre-epoch dates cannot be
623/// reliably represented as UnixNanos without overflow issues.
624fn 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
632/// Converts an ISO 8601 timestamp to a filesystem-safe format.
633///
634/// This function replaces colons and dots with hyphens to make the timestamp
635/// safe for use in filenames across different filesystems.
636fn iso_timestamp_to_file_timestamp(iso_timestamp: &str) -> String {
637    iso_timestamp.replace([':', '.'], "-")
638}
639
640/// Converts timestamps to a filename using ISO 8601 format.
641///
642/// This function converts two Unix nanosecond timestamps to a filename that uses
643/// ISO 8601 format with filesystem-safe characters, matching the catalog convention.
644fn 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    // Calculate start and end timestamps for the day (UTC)
676    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    // Start of day: 2024-01-01 00:00:00 UTC
770    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    // Midday: 2024-01-01 12:00:00 UTC
776    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    // End of day: 2024-01-01 23:59:59.999999999 UTC
782    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    // Start of new day: 2024-01-02 00:00:00 UTC
788    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 &quotes_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        // Trades must be written under the catalog convention path ("trades")
949        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        // Bars must be written under the catalog convention path ("bars"), consistent
1015        // with the other data types so they are discoverable by `ParquetDataCatalog`.
1016        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}