1use nautilus_core::UnixNanos;
19use nautilus_model::data::{Bar, QuoteTick, TradeTick};
20use nautilus_persistence::{
21 backend::parquet::{catalog::ParquetDataCatalog, paths::parse_filename_timestamps},
22 catalog::types::{CatalogDataType, data_type_from_data_path_prefix},
23};
24
25use super::{
26 CatalogReplayData, CatalogReplayRecord, CatalogSliceCoverage, CatalogSlicePlan,
27 CatalogSliceQuery, ReplayCatalog, ReplayTimeRange,
28};
29
30#[derive(Debug)]
32pub struct ParquetReplayCatalog<'a> {
33 catalog: &'a mut ParquetDataCatalog,
34}
35
36impl<'a> ParquetReplayCatalog<'a> {
37 pub const fn new(catalog: &'a mut ParquetDataCatalog) -> Self {
39 Self { catalog }
40 }
41}
42
43impl ReplayCatalog for ParquetReplayCatalog<'_> {
44 type Error = anyhow::Error;
45
46 fn plan_slice(
47 &mut self,
48 query: &CatalogSliceQuery,
49 ) -> Result<CatalogSliceCoverage, Self::Error> {
50 let data_type = CatalogDataType::Data(data_type_from_data_path_prefix(&query.data_cls)?);
51 let mut files = self.catalog.query_files(
52 &data_type,
53 query.identifiers_option(),
54 Some(query.start),
55 Some(query.end),
56 )?;
57 files.sort();
58
59 let intervals = files
60 .iter()
61 .filter_map(|file| {
62 parse_filename_timestamps(file).map(|(start, end)| {
63 ReplayTimeRange::new(UnixNanos::from(start), UnixNanos::from(end))
64 })
65 })
66 .collect();
67
68 Ok(CatalogSliceCoverage { files, intervals })
69 }
70
71 fn load_slice(
72 &mut self,
73 plan: &CatalogSlicePlan,
74 ) -> Result<Vec<CatalogReplayRecord>, Self::Error> {
75 let identifiers = plan.query.identifiers_option();
76 let start = Some(plan.query.start);
77 let end = Some(plan.query.end);
78 let files = Some(plan.coverage.files.clone());
79
80 match plan.query.data_cls.as_str() {
81 "quotes" => Ok(catalog_replay_records(self.catalog.query::<QuoteTick>(
82 identifiers,
83 start,
84 end,
85 None,
86 files,
87 false,
88 )?)),
89 "trades" => Ok(catalog_replay_records(self.catalog.query::<TradeTick>(
90 identifiers,
91 start,
92 end,
93 None,
94 files,
95 false,
96 )?)),
97 "bars" => Ok(catalog_replay_records(self.catalog.query::<Bar>(
98 identifiers,
99 start,
100 end,
101 None,
102 files,
103 false,
104 )?)),
105 data_cls => {
106 anyhow::bail!("catalog replay loading for {data_cls} is not supported")
107 }
108 }
109 }
110}
111
112fn catalog_replay_records<T>(records: Vec<T>) -> Vec<CatalogReplayRecord>
113where
114 T: Into<CatalogReplayData>,
115{
116 records
117 .into_iter()
118 .map(Into::into)
119 .map(CatalogReplayRecord::from_data)
120 .collect()
121}
122
123#[cfg(test)]
124mod tests {
125 use std::{
126 fs::{self, File},
127 path::Path,
128 };
129
130 use nautilus_model::{
131 data::{Bar, BarSpecification, BarType, QuoteTick, TradeTick},
132 enums::{AggregationSource, AggressorSide, BarAggregation, PriceType},
133 identifiers::{InstrumentId, TradeId},
134 types::{Price, Quantity},
135 };
136 use nautilus_persistence::backend::parquet::{
137 catalog::ParquetDataCatalog, paths::timestamps_to_filename,
138 };
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::Buy,
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::Sell,
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::Buy,
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}