1use nautilus_core::UnixNanos;
19use nautilus_model::data::{Bar, QuoteTick, TradeTick};
20use nautilus_persistence::backend::catalog::{ParquetDataCatalog, parse_filename_timestamps};
21
22use super::{
23 CatalogReplayData, CatalogReplayRecord, CatalogSliceCoverage, CatalogSlicePlan,
24 CatalogSliceQuery, ReplayCatalog, ReplayTimeRange,
25};
26
27#[derive(Debug)]
29pub struct ParquetReplayCatalog<'a> {
30 catalog: &'a mut ParquetDataCatalog,
31}
32
33impl<'a> ParquetReplayCatalog<'a> {
34 pub const fn new(catalog: &'a mut ParquetDataCatalog) -> Self {
36 Self { catalog }
37 }
38}
39
40impl ReplayCatalog for ParquetReplayCatalog<'_> {
41 type Error = anyhow::Error;
42
43 fn plan_slice(
44 &mut self,
45 query: &CatalogSliceQuery,
46 ) -> Result<CatalogSliceCoverage, Self::Error> {
47 let mut files = self.catalog.query_files(
48 &query.data_cls,
49 query.identifiers_option(),
50 Some(query.start),
51 Some(query.end),
52 )?;
53 files.sort();
54
55 let intervals = files
56 .iter()
57 .filter_map(|file| {
58 parse_filename_timestamps(file).map(|(start, end)| {
59 ReplayTimeRange::new(UnixNanos::from(start), UnixNanos::from(end))
60 })
61 })
62 .collect();
63
64 Ok(CatalogSliceCoverage { files, intervals })
65 }
66
67 fn load_slice(
68 &mut self,
69 plan: &CatalogSlicePlan,
70 ) -> Result<Vec<CatalogReplayRecord>, Self::Error> {
71 let identifiers = plan.query.identifiers_option();
72 let start = Some(plan.query.start);
73 let end = Some(plan.query.end);
74 let files = Some(plan.coverage.files.clone());
75
76 match plan.query.data_cls.as_str() {
77 "quotes" => Ok(catalog_replay_records(
78 self.catalog.query_typed_data::<QuoteTick>(
79 identifiers,
80 start,
81 end,
82 None,
83 files,
84 false,
85 )?,
86 )),
87 "trades" => Ok(catalog_replay_records(
88 self.catalog.query_typed_data::<TradeTick>(
89 identifiers,
90 start,
91 end,
92 None,
93 files,
94 false,
95 )?,
96 )),
97 "bars" => Ok(catalog_replay_records(
98 self.catalog.query_typed_data::<Bar>(
99 identifiers,
100 start,
101 end,
102 None,
103 files,
104 false,
105 )?,
106 )),
107 data_cls => {
108 anyhow::bail!("catalog replay loading for {data_cls} is not supported")
109 }
110 }
111 }
112}
113
114fn catalog_replay_records<T>(records: Vec<T>) -> Vec<CatalogReplayRecord>
115where
116 T: Into<CatalogReplayData>,
117{
118 records
119 .into_iter()
120 .map(Into::into)
121 .map(CatalogReplayRecord::from_data)
122 .collect()
123}
124
125#[cfg(test)]
126mod tests {
127 use std::{
128 fs::{self, File},
129 path::Path,
130 };
131
132 use nautilus_model::{
133 data::{Bar, BarSpecification, BarType, QuoteTick, TradeTick},
134 enums::{AggregationSource, AggressorSide, BarAggregation, PriceType},
135 identifiers::{InstrumentId, TradeId},
136 types::{Price, Quantity},
137 };
138 use nautilus_persistence::backend::catalog::{ParquetDataCatalog, timestamps_to_filename};
139 use rstest::rstest;
140 use tempfile::TempDir;
141
142 use super::*;
143
144 #[rstest]
145 fn parquet_replay_catalog_plans_selected_slice_files() {
146 let temp_dir = TempDir::new().unwrap();
147 let mut catalog = ParquetDataCatalog::new(temp_dir.path(), None, None, None, None);
148
149 create_catalog_file(temp_dir.path(), "quotes", "AUDUSD.SIM", 1_000, 2_000);
150 create_catalog_file(temp_dir.path(), "quotes", "AUDUSD.SIM", 10_000, 11_000);
151 create_catalog_file(temp_dir.path(), "quotes", "ETHUSDT.BINANCE", 5_000, 6_000);
152
153 let query = CatalogSliceQuery {
154 data_cls: "quotes".to_string(),
155 identifiers: vec!["AUD/USD.SIM".to_string()],
156 start: UnixNanos::from(1_500),
157 end: UnixNanos::from(2_500),
158 required: true,
159 };
160 let coverage = ParquetReplayCatalog::new(&mut catalog)
161 .plan_slice(&query)
162 .unwrap();
163
164 assert_eq!(coverage.files.len(), 1);
165 assert!(
166 coverage.files[0].contains("data/quotes/AUDUSD.SIM/"),
167 "planned file should come from AUD/USD.SIM partition, was {}",
168 coverage.files[0],
169 );
170 assert_eq!(
171 coverage.intervals,
172 vec![ReplayTimeRange::new(
173 UnixNanos::from(1_000),
174 UnixNanos::from(2_000)
175 )]
176 );
177
178 let full_window_query = CatalogSliceQuery {
179 start: UnixNanos::from(0),
180 end: UnixNanos::from(12_000),
181 ..query.clone()
182 };
183 let full_window_coverage = ParquetReplayCatalog::new(&mut catalog)
184 .plan_slice(&full_window_query)
185 .unwrap();
186
187 assert_eq!(full_window_coverage.files.len(), 2);
188 assert_eq!(
189 full_window_coverage.intervals,
190 vec![
191 ReplayTimeRange::new(UnixNanos::from(1_000), UnixNanos::from(2_000)),
192 ReplayTimeRange::new(UnixNanos::from(10_000), UnixNanos::from(11_000)),
193 ]
194 );
195
196 let missing_query = CatalogSliceQuery {
197 start: UnixNanos::from(20_000),
198 end: UnixNanos::from(21_000),
199 ..query
200 };
201 let missing_coverage = ParquetReplayCatalog::new(&mut catalog)
202 .plan_slice(&missing_query)
203 .unwrap();
204
205 assert!(missing_coverage.is_missing());
206 assert!(missing_coverage.intervals.is_empty());
207 }
208
209 #[rstest]
210 fn parquet_replay_catalog_loads_selected_quote_records() {
211 let temp_dir = TempDir::new().unwrap();
212 let mut catalog = ParquetDataCatalog::new(temp_dir.path(), None, None, None, None);
213 let instrument_id = InstrumentId::from("AUD/USD.SIM");
214 let quotes = vec![
215 QuoteTick::new(
216 instrument_id,
217 Price::from("1.0001"),
218 Price::from("1.0002"),
219 Quantity::from("100"),
220 Quantity::from("100"),
221 UnixNanos::from(1_000),
222 UnixNanos::from(1_000),
223 ),
224 QuoteTick::new(
225 instrument_id,
226 Price::from("1.0003"),
227 Price::from("1.0004"),
228 Quantity::from("200"),
229 Quantity::from("200"),
230 UnixNanos::from(2_000),
231 UnixNanos::from(2_000),
232 ),
233 QuoteTick::new(
234 instrument_id,
235 Price::from("1.0005"),
236 Price::from("1.0006"),
237 Quantity::from("300"),
238 Quantity::from("300"),
239 UnixNanos::from(3_000),
240 UnixNanos::from(3_000),
241 ),
242 ];
243 catalog
244 .write_to_parquet("es, None, None, None)
245 .expect("write quotes");
246
247 let query = CatalogSliceQuery {
248 data_cls: "quotes".to_string(),
249 identifiers: vec!["AUD/USD.SIM".to_string()],
250 start: UnixNanos::from(1_500),
251 end: UnixNanos::from(2_500),
252 required: true,
253 };
254 let mut replay_catalog = ParquetReplayCatalog::new(&mut catalog);
255 let coverage = replay_catalog.plan_slice(&query).expect("plan slice");
256 let plan = catalog_slice_plan(query, coverage);
257
258 let records = replay_catalog.load_slice(&plan).expect("load slice");
259
260 assert_eq!(
261 records,
262 vec![CatalogReplayRecord::from_data(CatalogReplayData::Quote(
263 quotes[1]
264 ))],
265 );
266 }
267
268 #[rstest]
269 fn parquet_replay_catalog_loads_selected_trade_records() {
270 let temp_dir = TempDir::new().unwrap();
271 let mut catalog = ParquetDataCatalog::new(temp_dir.path(), None, None, None, None);
272 let instrument_id = InstrumentId::from("AUD/USD.SIM");
273 let trades = vec![
274 TradeTick::new(
275 instrument_id,
276 Price::from("1.0001"),
277 Quantity::from("100"),
278 AggressorSide::Buyer,
279 TradeId::from("T-1"),
280 UnixNanos::from(1_000),
281 UnixNanos::from(1_000),
282 ),
283 TradeTick::new(
284 instrument_id,
285 Price::from("1.0002"),
286 Quantity::from("200"),
287 AggressorSide::Seller,
288 TradeId::from("T-2"),
289 UnixNanos::from(2_000),
290 UnixNanos::from(2_000),
291 ),
292 TradeTick::new(
293 instrument_id,
294 Price::from("1.0003"),
295 Quantity::from("300"),
296 AggressorSide::Buyer,
297 TradeId::from("T-3"),
298 UnixNanos::from(3_000),
299 UnixNanos::from(3_000),
300 ),
301 ];
302 catalog
303 .write_to_parquet(&trades, None, None, None)
304 .expect("write trades");
305
306 let query = CatalogSliceQuery {
307 data_cls: "trades".to_string(),
308 identifiers: vec!["AUD/USD.SIM".to_string()],
309 start: UnixNanos::from(1_500),
310 end: UnixNanos::from(2_500),
311 required: true,
312 };
313 let mut replay_catalog = ParquetReplayCatalog::new(&mut catalog);
314 let coverage = replay_catalog.plan_slice(&query).expect("plan slice");
315 let plan = catalog_slice_plan(query, coverage);
316
317 let records = replay_catalog.load_slice(&plan).expect("load slice");
318
319 assert_eq!(
320 records,
321 vec![CatalogReplayRecord::from_data(CatalogReplayData::Trade(
322 trades[1]
323 ))],
324 );
325 }
326
327 #[rstest]
328 fn parquet_replay_catalog_loads_selected_bar_records() {
329 let temp_dir = TempDir::new().unwrap();
330 let mut catalog = ParquetDataCatalog::new(temp_dir.path(), None, None, None, None);
331 let instrument_id = InstrumentId::from("AUD/USD.SIM");
332 let bar_type = BarType::new(
333 instrument_id,
334 BarSpecification::new(1, BarAggregation::Minute, PriceType::Last),
335 AggregationSource::External,
336 );
337 let bars = vec![
338 Bar::new(
339 bar_type,
340 Price::from("1.0000"),
341 Price::from("1.0002"),
342 Price::from("1.0000"),
343 Price::from("1.0001"),
344 Quantity::from("100"),
345 UnixNanos::from(1_000),
346 UnixNanos::from(1_000),
347 ),
348 Bar::new(
349 bar_type,
350 Price::from("1.0001"),
351 Price::from("1.0004"),
352 Price::from("1.0001"),
353 Price::from("1.0003"),
354 Quantity::from("200"),
355 UnixNanos::from(2_000),
356 UnixNanos::from(2_000),
357 ),
358 Bar::new(
359 bar_type,
360 Price::from("1.0003"),
361 Price::from("1.0006"),
362 Price::from("1.0003"),
363 Price::from("1.0005"),
364 Quantity::from("300"),
365 UnixNanos::from(3_000),
366 UnixNanos::from(3_000),
367 ),
368 ];
369 catalog
370 .write_to_parquet(&bars, None, None, None)
371 .expect("write bars");
372
373 let query = CatalogSliceQuery {
374 data_cls: "bars".to_string(),
375 identifiers: vec!["AUD/USD.SIM".to_string()],
376 start: UnixNanos::from(1_500),
377 end: UnixNanos::from(2_500),
378 required: true,
379 };
380 let mut replay_catalog = ParquetReplayCatalog::new(&mut catalog);
381 let coverage = replay_catalog.plan_slice(&query).expect("plan slice");
382 let plan = catalog_slice_plan(query, coverage);
383
384 let records = replay_catalog.load_slice(&plan).expect("load slice");
385
386 assert_eq!(
387 records,
388 vec![CatalogReplayRecord::from_data(CatalogReplayData::Bar(
389 bars[1]
390 ))],
391 );
392 }
393
394 #[rstest]
395 fn parquet_replay_catalog_rejects_unsupported_load_slice() {
396 let temp_dir = TempDir::new().unwrap();
397 let mut catalog = ParquetDataCatalog::new(temp_dir.path(), None, None, None, None);
398 let plan = CatalogSlicePlan {
399 query: CatalogSliceQuery {
400 data_cls: "order_book_deltas".to_string(),
401 identifiers: vec!["AUD/USD.SIM".to_string()],
402 start: UnixNanos::from(1_000),
403 end: UnixNanos::from(2_000),
404 required: true,
405 },
406 coverage: CatalogSliceCoverage::from_files(vec![
407 "data/order_book_deltas/AUDUSD.SIM/1000_2000.parquet".to_string(),
408 ]),
409 };
410
411 let err = ParquetReplayCatalog::new(&mut catalog)
412 .load_slice(&plan)
413 .expect_err("unsupported data class must fail");
414
415 assert_eq!(
416 err.to_string(),
417 "catalog replay loading for order_book_deltas is not supported",
418 );
419 }
420
421 fn catalog_slice_plan(
422 query: CatalogSliceQuery,
423 coverage: CatalogSliceCoverage,
424 ) -> CatalogSlicePlan {
425 CatalogSlicePlan { query, coverage }
426 }
427
428 fn create_catalog_file(
429 base_path: &Path,
430 data_cls: &str,
431 identifier: &str,
432 start: u64,
433 end: u64,
434 ) {
435 let directory = base_path.join("data").join(data_cls).join(identifier);
436 fs::create_dir_all(&directory).unwrap();
437
438 let filename = timestamps_to_filename(UnixNanos::from(start), UnixNanos::from(end));
439 File::create(directory.join(filename)).unwrap();
440 }
441}