1use std::collections::VecDeque;
4use std::fs::{self, File, Metadata};
5use std::io::{Seek, SeekFrom};
6#[cfg(unix)]
7use std::os::unix::fs::MetadataExt;
8#[cfg(windows)]
9use std::os::windows::fs::MetadataExt;
10use std::path::{Path, PathBuf};
11use std::time::SystemTime;
12
13use chrono::{NaiveDate, NaiveDateTime};
14use polars::prelude::*;
15
16use crate::convert::{dataframe_to_bars, dataframe_to_ticks};
17use crate::error::{DataError, Result};
18use crate::models::{Bar, Tick};
19use crate::parquet_store::ParquetStore;
20
21pub const DEFAULT_PARQUET_SCAN_ROWS: usize = 65_536;
23
24#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
26pub struct ParquetScanBounds {
27 pub from: Option<NaiveDateTime>,
28 pub to: Option<NaiveDateTime>,
29}
30
31impl ParquetScanBounds {
32 pub const fn new(from: Option<NaiveDateTime>, to: Option<NaiveDateTime>) -> Self {
33 Self { from, to }
34 }
35
36 fn validate(self) -> Result<Self> {
37 if let (Some(from), Some(to)) = (self.from, self.to)
38 && from > to
39 {
40 return Err(DataError::InvalidScanBounds { from, to });
41 }
42 Ok(self)
43 }
44
45 fn contains(self, ts: NaiveDateTime) -> bool {
46 self.from.is_none_or(|from| ts >= from) && self.to.is_none_or(|to| ts <= to)
47 }
48}
49
50#[derive(Debug, Clone, PartialEq)]
52pub struct ParquetScannedRow<T> {
53 pub row: T,
54 pub source_row_ordinal: u64,
55}
56
57trait Timestamped {
58 fn timestamp(&self) -> NaiveDateTime;
59}
60
61impl Timestamped for Tick {
62 fn timestamp(&self) -> NaiveDateTime {
63 self.ts
64 }
65}
66
67impl Timestamped for Bar {
68 fn timestamp(&self) -> NaiveDateTime {
69 self.ts
70 }
71}
72
73#[derive(Debug, Clone, PartialEq, Eq)]
74struct FileFingerprint {
75 len: u64,
76 modified: Option<SystemTime>,
77 created: Option<SystemTime>,
78 #[cfg(unix)]
79 device: u64,
80 #[cfg(unix)]
81 inode: u64,
82 #[cfg(windows)]
83 volume_serial_number: Option<u32>,
84 #[cfg(windows)]
85 file_index: Option<u64>,
86}
87
88impl FileFingerprint {
89 fn from_metadata(metadata: &Metadata) -> Self {
90 Self {
91 len: metadata.len(),
92 modified: metadata.modified().ok(),
93 created: metadata.created().ok(),
94 #[cfg(unix)]
95 device: metadata.dev(),
96 #[cfg(unix)]
97 inode: metadata.ino(),
98 #[cfg(windows)]
99 volume_serial_number: metadata.volume_serial_number(),
100 #[cfg(windows)]
101 file_index: metadata.file_index(),
102 }
103 }
104}
105
106#[derive(Debug, Clone)]
107struct PartitionDescriptor {
108 path: PathBuf,
109 fingerprint: FileFingerprint,
110 row_count: usize,
111 source_row_base: u64,
112}
113
114impl PartitionDescriptor {
115 fn describe(path: PathBuf, source_row_base: u64) -> Result<Self> {
116 let file = File::open(&path)?;
117 let fingerprint = FileFingerprint::from_metadata(&file.metadata()?);
118 ensure_path_matches(&path, &fingerprint)?;
119
120 let mut reader = ParquetReader::new(file.try_clone()?);
121 let row_count = reader.num_rows()?;
122 ensure_file_matches(&path, &file, &fingerprint)?;
123 ensure_path_matches(&path, &fingerprint)?;
124
125 Ok(Self {
126 path,
127 fingerprint,
128 row_count,
129 source_row_base,
130 })
131 }
132
133 fn validate(&self) -> Result<()> {
134 ensure_path_matches(&self.path, &self.fingerprint)
135 }
136
137 fn open_generation(&self) -> Result<OpenPartition> {
138 self.validate()?;
139 let file = match File::open(&self.path) {
140 Ok(file) => file,
141 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
142 return Err(partition_changed(&self.path));
143 }
144 Err(error) => return Err(error.into()),
145 };
146 ensure_file_matches(&self.path, &file, &self.fingerprint)?;
147 self.validate()?;
148 Ok(OpenPartition {
149 descriptor: self.clone(),
150 file,
151 })
152 }
153
154 fn ordinal(&self, row_offset: usize) -> Result<u64> {
155 self.source_row_base
156 .checked_add(
157 u64::try_from(row_offset).map_err(|_| {
158 DataError::Other("Parquet source row ordinal exceeds u64".into())
159 })?,
160 )
161 .ok_or_else(|| DataError::Other("Parquet source row ordinal overflow".into()))
162 }
163}
164
165struct OpenPartition {
166 descriptor: PartitionDescriptor,
167 file: File,
168}
169
170impl OpenPartition {
171 fn validate(&self) -> Result<()> {
172 ensure_file_matches(
173 &self.descriptor.path,
174 &self.file,
175 &self.descriptor.fingerprint,
176 )?;
177 self.descriptor.validate()
178 }
179
180 fn reader_file(&self) -> Result<File> {
181 self.validate()?;
182 let mut file = self.file.try_clone()?;
183 file.seek(SeekFrom::Start(0))?;
184 Ok(file)
185 }
186}
187
188fn partition_changed(path: &Path) -> DataError {
189 DataError::ParquetPartitionChanged {
190 path: path.display().to_string(),
191 }
192}
193
194fn ensure_file_matches(path: &Path, file: &File, expected: &FileFingerprint) -> Result<()> {
195 if FileFingerprint::from_metadata(&file.metadata()?) == *expected {
196 Ok(())
197 } else {
198 Err(partition_changed(path))
199 }
200}
201
202fn ensure_path_matches(path: &Path, expected: &FileFingerprint) -> Result<()> {
203 let metadata = match fs::metadata(path) {
204 Ok(metadata) => metadata,
205 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
206 return Err(partition_changed(path));
207 }
208 Err(error) => return Err(error.into()),
209 };
210 if FileFingerprint::from_metadata(&metadata) == *expected {
211 Ok(())
212 } else {
213 Err(partition_changed(path))
214 }
215}
216
217#[derive(Debug, Clone)]
218struct PartitionScan {
219 partitions: Vec<PartitionDescriptor>,
220 bounds: ParquetScanBounds,
221}
222
223impl PartitionScan {
224 fn describe(
225 directory: &Path,
226 bounds: ParquetScanBounds,
227 is_cancelled: &mut dyn FnMut() -> bool,
228 ) -> Result<Self> {
229 let bounds = bounds.validate()?;
230 let partitions = list_partitions(directory, bounds, is_cancelled)?;
231 Ok(Self { partitions, bounds })
232 }
233
234 fn validate(&self) -> Result<()> {
235 for partition in &self.partitions {
236 partition.validate()?;
237 }
238 Ok(())
239 }
240}
241
242type PartitionLoader<T> = fn(File, usize, usize) -> Result<Vec<T>>;
243
244struct PartitionCursor<T> {
245 partitions: VecDeque<PartitionDescriptor>,
246 current_partition: Option<OpenPartition>,
247 current_rows: std::vec::IntoIter<ParquetScannedRow<T>>,
248 bounds: ParquetScanBounds,
249 rows_per_read: usize,
250 partition_offset: usize,
251 finish_current_partition: bool,
252 last_read_ts: Option<NaiveDateTime>,
253 load_partition: PartitionLoader<T>,
254}
255
256impl<T: Timestamped> PartitionCursor<T> {
257 fn open(
258 scan: PartitionScan,
259 rows_per_read: usize,
260 load_partition: PartitionLoader<T>,
261 ) -> Result<Self> {
262 if rows_per_read == 0 {
263 return Err(DataError::InvalidScanReadSize);
264 }
265 scan.validate()?;
266 Ok(Self {
267 partitions: scan.partitions.into(),
268 current_partition: None,
269 current_rows: Vec::new().into_iter(),
270 bounds: scan.bounds,
271 rows_per_read,
272 partition_offset: 0,
273 finish_current_partition: false,
274 last_read_ts: None,
275 load_partition,
276 })
277 }
278
279 fn next_row(
280 &mut self,
281 is_cancelled: &mut dyn FnMut() -> bool,
282 ) -> Result<Option<ParquetScannedRow<T>>> {
283 loop {
284 ensure_not_cancelled(is_cancelled)?;
285 if !self.current_rows.as_slice().is_empty() {
286 self.current_partition
287 .as_ref()
288 .expect("buffered rows retain their open partition")
289 .validate()?;
290 return Ok(self.current_rows.next());
291 }
292
293 if self.finish_current_partition {
294 self.current_partition = None;
295 self.partition_offset = 0;
296 self.finish_current_partition = false;
297 }
298
299 if self.current_partition.is_none() {
300 let Some(descriptor) = self.partitions.pop_front() else {
301 return Ok(None);
302 };
303 self.current_partition = Some(descriptor.open_generation()?);
304 }
305
306 let partition = self
307 .current_partition
308 .as_ref()
309 .expect("current partition was opened");
310 if self.partition_offset >= partition.descriptor.row_count {
311 self.finish_current_partition = true;
312 continue;
313 }
314
315 ensure_not_cancelled(is_cancelled)?;
316 let rows_to_read = self
317 .rows_per_read
318 .min(partition.descriptor.row_count - self.partition_offset);
319 let reader = partition.reader_file()?;
320 let rows = (self.load_partition)(reader, self.partition_offset, rows_to_read)?;
321 ensure_not_cancelled(is_cancelled)?;
322 partition.validate()?;
323 if rows.len() != rows_to_read {
324 return Err(DataError::Other(format!(
325 "Parquet slice {} at offset {} returned {} rows instead of {}",
326 partition.descriptor.path.display(),
327 self.partition_offset,
328 rows.len(),
329 rows_to_read
330 )));
331 }
332
333 let read_offset = self.partition_offset;
334 let (keep_len, last_read_ts, exceeded_upper_bound) = validate_monotonic_prefix(
335 self.last_read_ts,
336 &rows,
337 &partition.descriptor.path,
338 self.bounds.to,
339 )?;
340 self.last_read_ts = last_read_ts;
341 self.partition_offset = self
342 .partition_offset
343 .checked_add(rows.len())
344 .ok_or_else(|| DataError::Other("Parquet scan row offset overflow".into()))?;
345 self.finish_current_partition =
346 exceeded_upper_bound || self.partition_offset >= partition.descriptor.row_count;
347 if exceeded_upper_bound {
348 self.partitions.clear();
349 }
350
351 let descriptor = &partition.descriptor;
352 let mut bounded_rows = Vec::with_capacity(keep_len);
353 for (index, row) in rows.into_iter().take(keep_len).enumerate() {
354 if self.bounds.contains(row.timestamp()) {
355 bounded_rows.push(ParquetScannedRow {
356 row,
357 source_row_ordinal: descriptor.ordinal(read_offset + index)?,
358 });
359 }
360 }
361 self.current_rows = bounded_rows.into_iter();
362 }
363 }
364
365 fn remaining_partitions(&self) -> usize {
366 self.partitions.len() + usize::from(self.current_partition.is_some())
367 }
368
369 fn rows_per_read(&self) -> usize {
370 self.rows_per_read
371 }
372}
373
374fn validate_monotonic_prefix<T: Timestamped>(
375 mut previous: Option<NaiveDateTime>,
376 rows: &[T],
377 path: &Path,
378 inclusive_to: Option<NaiveDateTime>,
379) -> Result<(usize, Option<NaiveDateTime>, bool)> {
380 for (index, row) in rows.iter().enumerate() {
381 let current = row.timestamp();
382 if let Some(previous) = previous
383 && current < previous
384 {
385 return Err(DataError::NonMonotonicParquetData {
386 path: path.display().to_string(),
387 previous,
388 current,
389 });
390 }
391 if inclusive_to.is_some_and(|to| current > to) {
392 return Ok((index, previous, true));
393 }
394 previous = Some(current);
395 }
396 Ok((rows.len(), previous, false))
397}
398
399fn find_latest_row<T, P>(
400 scan: &PartitionScan,
401 rows_per_read: usize,
402 load_partition: PartitionLoader<T>,
403 mut predicate: P,
404 is_cancelled: &mut dyn FnMut() -> bool,
405) -> Result<Option<ParquetScannedRow<T>>>
406where
407 T: Timestamped,
408 P: FnMut(&T) -> bool,
409{
410 if rows_per_read == 0 {
411 return Err(DataError::InvalidScanReadSize);
412 }
413 scan.validate()?;
414 let mut later_first_ts = None;
415
416 for descriptor in scan.partitions.iter().rev() {
417 ensure_not_cancelled(is_cancelled)?;
418 let partition = descriptor.open_generation()?;
419 let mut end = descriptor.row_count;
420 while end > 0 {
421 ensure_not_cancelled(is_cancelled)?;
422 let start = end.saturating_sub(rows_per_read);
423 let reader = partition.reader_file()?;
424 let rows = load_partition(reader, start, end - start)?;
425 ensure_not_cancelled(is_cancelled)?;
426 partition.validate()?;
427 if rows.len() != end - start {
428 return Err(DataError::Other(format!(
429 "Parquet reverse slice {} at offset {} returned {} rows instead of {}",
430 descriptor.path.display(),
431 start,
432 rows.len(),
433 end - start
434 )));
435 }
436 validate_monotonic_rows(None, &rows, &descriptor.path)?;
437 if let (Some(last), Some(later)) =
438 (rows.last().map(Timestamped::timestamp), later_first_ts)
439 && last > later
440 {
441 return Err(DataError::NonMonotonicParquetData {
442 path: descriptor.path.display().to_string(),
443 previous: last,
444 current: later,
445 });
446 }
447 if let Some(first) = rows.first().map(Timestamped::timestamp) {
448 later_first_ts = Some(first);
449 }
450
451 for (index, row) in rows.into_iter().enumerate().rev() {
452 let timestamp = row.timestamp();
453 if scan.bounds.from.is_some_and(|from| timestamp < from) {
454 return Ok(None);
455 }
456 if scan.bounds.contains(timestamp) && predicate(&row) {
457 return Ok(Some(ParquetScannedRow {
458 row,
459 source_row_ordinal: descriptor.ordinal(start + index)?,
460 }));
461 }
462 }
463 end = start;
464 }
465 }
466
467 ensure_not_cancelled(is_cancelled)?;
468 Ok(None)
469}
470
471#[derive(Debug, Clone)]
473pub struct ParquetTickScan {
474 inner: PartitionScan,
475}
476
477impl ParquetTickScan {
478 pub fn describe(
480 root: impl AsRef<Path>,
481 exchange: &str,
482 symbol: &str,
483 bounds: ParquetScanBounds,
484 ) -> Result<Self> {
485 Self::describe_cancellable(root, exchange, symbol, bounds, || false)
486 }
487
488 pub fn describe_cancellable<F>(
490 root: impl AsRef<Path>,
491 exchange: &str,
492 symbol: &str,
493 bounds: ParquetScanBounds,
494 mut is_cancelled: F,
495 ) -> Result<Self>
496 where
497 F: FnMut() -> bool,
498 {
499 let directory = root
500 .as_ref()
501 .join("ticks")
502 .join(format!("exchange={exchange}"))
503 .join(format!("symbol={symbol}"));
504 Ok(Self {
505 inner: PartitionScan::describe(&directory, bounds, &mut is_cancelled)?,
506 })
507 }
508
509 pub fn cursor(&self) -> Result<ParquetTickCursor> {
511 self.cursor_with_read_size(DEFAULT_PARQUET_SCAN_ROWS)
512 }
513
514 pub fn cursor_with_read_size(&self, rows_per_read: usize) -> Result<ParquetTickCursor> {
516 Ok(ParquetTickCursor {
517 inner: PartitionCursor::open(self.inner.clone(), rows_per_read, read_tick_partition)?,
518 })
519 }
520
521 pub fn latest_valid_tick_cancellable<F>(
523 &self,
524 mut is_cancelled: F,
525 ) -> Result<Option<ParquetScannedRow<Tick>>>
526 where
527 F: FnMut() -> bool,
528 {
529 find_latest_row(
530 &self.inner,
531 DEFAULT_PARQUET_SCAN_ROWS,
532 read_tick_partition,
533 tick_has_valid_quote,
534 &mut is_cancelled,
535 )
536 }
537
538 pub fn latest_valid_tick_before_cancellable<F>(
540 &self,
541 before: NaiveDateTime,
542 mut is_cancelled: F,
543 ) -> Result<Option<ParquetScannedRow<Tick>>>
544 where
545 F: FnMut() -> bool,
546 {
547 find_latest_row(
548 &self.inner,
549 DEFAULT_PARQUET_SCAN_ROWS,
550 read_tick_partition,
551 |tick| tick.ts < before && tick_has_valid_quote(tick),
552 &mut is_cancelled,
553 )
554 }
555}
556
557pub struct ParquetTickCursor {
559 inner: PartitionCursor<Tick>,
560}
561
562impl ParquetTickCursor {
563 pub fn open(
565 root: impl AsRef<Path>,
566 exchange: &str,
567 symbol: &str,
568 bounds: ParquetScanBounds,
569 ) -> Result<Self> {
570 Self::open_with_read_size(root, exchange, symbol, bounds, DEFAULT_PARQUET_SCAN_ROWS)
571 }
572
573 pub fn open_with_read_size(
575 root: impl AsRef<Path>,
576 exchange: &str,
577 symbol: &str,
578 bounds: ParquetScanBounds,
579 rows_per_read: usize,
580 ) -> Result<Self> {
581 Self::open_cancellable_with_read_size(root, exchange, symbol, bounds, rows_per_read, || {
582 false
583 })
584 }
585
586 pub fn open_cancellable<F>(
588 root: impl AsRef<Path>,
589 exchange: &str,
590 symbol: &str,
591 bounds: ParquetScanBounds,
592 is_cancelled: F,
593 ) -> Result<Self>
594 where
595 F: FnMut() -> bool,
596 {
597 Self::open_cancellable_with_read_size(
598 root,
599 exchange,
600 symbol,
601 bounds,
602 DEFAULT_PARQUET_SCAN_ROWS,
603 is_cancelled,
604 )
605 }
606
607 pub fn open_cancellable_with_read_size<F>(
609 root: impl AsRef<Path>,
610 exchange: &str,
611 symbol: &str,
612 bounds: ParquetScanBounds,
613 rows_per_read: usize,
614 is_cancelled: F,
615 ) -> Result<Self>
616 where
617 F: FnMut() -> bool,
618 {
619 ParquetTickScan::describe_cancellable(root, exchange, symbol, bounds, is_cancelled)?
620 .cursor_with_read_size(rows_per_read)
621 }
622
623 pub fn next_tick(&mut self) -> Result<Option<Tick>> {
625 self.next_tick_with_ordinal()
626 .map(|row| row.map(|row| row.row))
627 }
628
629 pub fn next_tick_with_ordinal(&mut self) -> Result<Option<ParquetScannedRow<Tick>>> {
631 self.next_tick_with_ordinal_cancellable(|| false)
632 }
633
634 pub fn next_tick_cancellable<F>(&mut self, is_cancelled: F) -> Result<Option<Tick>>
636 where
637 F: FnMut() -> bool,
638 {
639 self.next_tick_with_ordinal_cancellable(is_cancelled)
640 .map(|row| row.map(|row| row.row))
641 }
642
643 pub fn next_tick_with_ordinal_cancellable<F>(
645 &mut self,
646 mut is_cancelled: F,
647 ) -> Result<Option<ParquetScannedRow<Tick>>>
648 where
649 F: FnMut() -> bool,
650 {
651 self.inner.next_row(&mut is_cancelled)
652 }
653
654 pub fn remaining_partitions(&self) -> usize {
656 self.inner.remaining_partitions()
657 }
658
659 pub fn rows_per_read(&self) -> usize {
661 self.inner.rows_per_read()
662 }
663}
664
665#[derive(Debug, Clone)]
667pub struct ParquetBarScan {
668 inner: PartitionScan,
669}
670
671impl ParquetBarScan {
672 pub fn describe(
674 root: impl AsRef<Path>,
675 exchange: &str,
676 symbol: &str,
677 timeframe: &str,
678 bounds: ParquetScanBounds,
679 ) -> Result<Self> {
680 Self::describe_cancellable(root, exchange, symbol, timeframe, bounds, || false)
681 }
682
683 pub fn describe_cancellable<F>(
685 root: impl AsRef<Path>,
686 exchange: &str,
687 symbol: &str,
688 timeframe: &str,
689 bounds: ParquetScanBounds,
690 mut is_cancelled: F,
691 ) -> Result<Self>
692 where
693 F: FnMut() -> bool,
694 {
695 let directory = root
696 .as_ref()
697 .join("bars")
698 .join(format!("exchange={exchange}"))
699 .join(format!("symbol={symbol}"))
700 .join(format!("timeframe={timeframe}"));
701 Ok(Self {
702 inner: PartitionScan::describe(&directory, bounds, &mut is_cancelled)?,
703 })
704 }
705
706 pub fn cursor(&self) -> Result<ParquetBarCursor> {
708 self.cursor_with_read_size(DEFAULT_PARQUET_SCAN_ROWS)
709 }
710
711 pub fn cursor_with_read_size(&self, rows_per_read: usize) -> Result<ParquetBarCursor> {
713 Ok(ParquetBarCursor {
714 inner: PartitionCursor::open(self.inner.clone(), rows_per_read, read_bar_partition)?,
715 })
716 }
717
718 pub fn latest_valid_bar_cancellable<F>(
720 &self,
721 mut is_cancelled: F,
722 ) -> Result<Option<ParquetScannedRow<Bar>>>
723 where
724 F: FnMut() -> bool,
725 {
726 find_latest_row(
727 &self.inner,
728 DEFAULT_PARQUET_SCAN_ROWS,
729 read_bar_partition,
730 |bar| bar.close.is_finite() && bar.close > 0.0,
731 &mut is_cancelled,
732 )
733 }
734}
735
736pub struct ParquetBarCursor {
738 inner: PartitionCursor<Bar>,
739}
740
741impl ParquetBarCursor {
742 pub fn open(
744 root: impl AsRef<Path>,
745 exchange: &str,
746 symbol: &str,
747 timeframe: &str,
748 bounds: ParquetScanBounds,
749 ) -> Result<Self> {
750 Self::open_with_read_size(
751 root,
752 exchange,
753 symbol,
754 timeframe,
755 bounds,
756 DEFAULT_PARQUET_SCAN_ROWS,
757 )
758 }
759
760 pub fn open_with_read_size(
762 root: impl AsRef<Path>,
763 exchange: &str,
764 symbol: &str,
765 timeframe: &str,
766 bounds: ParquetScanBounds,
767 rows_per_read: usize,
768 ) -> Result<Self> {
769 Self::open_cancellable_with_read_size(
770 root,
771 exchange,
772 symbol,
773 timeframe,
774 bounds,
775 rows_per_read,
776 || false,
777 )
778 }
779
780 pub fn open_cancellable<F>(
782 root: impl AsRef<Path>,
783 exchange: &str,
784 symbol: &str,
785 timeframe: &str,
786 bounds: ParquetScanBounds,
787 is_cancelled: F,
788 ) -> Result<Self>
789 where
790 F: FnMut() -> bool,
791 {
792 Self::open_cancellable_with_read_size(
793 root,
794 exchange,
795 symbol,
796 timeframe,
797 bounds,
798 DEFAULT_PARQUET_SCAN_ROWS,
799 is_cancelled,
800 )
801 }
802
803 pub fn open_cancellable_with_read_size<F>(
805 root: impl AsRef<Path>,
806 exchange: &str,
807 symbol: &str,
808 timeframe: &str,
809 bounds: ParquetScanBounds,
810 rows_per_read: usize,
811 is_cancelled: F,
812 ) -> Result<Self>
813 where
814 F: FnMut() -> bool,
815 {
816 ParquetBarScan::describe_cancellable(
817 root,
818 exchange,
819 symbol,
820 timeframe,
821 bounds,
822 is_cancelled,
823 )?
824 .cursor_with_read_size(rows_per_read)
825 }
826
827 pub fn next_bar(&mut self) -> Result<Option<Bar>> {
829 self.next_bar_with_ordinal()
830 .map(|row| row.map(|row| row.row))
831 }
832
833 pub fn next_bar_with_ordinal(&mut self) -> Result<Option<ParquetScannedRow<Bar>>> {
835 self.next_bar_with_ordinal_cancellable(|| false)
836 }
837
838 pub fn next_bar_cancellable<F>(&mut self, is_cancelled: F) -> Result<Option<Bar>>
840 where
841 F: FnMut() -> bool,
842 {
843 self.next_bar_with_ordinal_cancellable(is_cancelled)
844 .map(|row| row.map(|row| row.row))
845 }
846
847 pub fn next_bar_with_ordinal_cancellable<F>(
849 &mut self,
850 mut is_cancelled: F,
851 ) -> Result<Option<ParquetScannedRow<Bar>>>
852 where
853 F: FnMut() -> bool,
854 {
855 self.inner.next_row(&mut is_cancelled)
856 }
857
858 pub fn remaining_partitions(&self) -> usize {
860 self.inner.remaining_partitions()
861 }
862
863 pub fn rows_per_read(&self) -> usize {
865 self.inner.rows_per_read()
866 }
867}
868
869impl ParquetStore {
870 pub fn scan_ticks(
872 &self,
873 exchange: &str,
874 symbol: &str,
875 bounds: ParquetScanBounds,
876 ) -> Result<ParquetTickCursor> {
877 ParquetTickCursor::open(self.root_path(), exchange, symbol, bounds)
878 }
879
880 pub fn scan_ticks_cancellable<F>(
882 &self,
883 exchange: &str,
884 symbol: &str,
885 bounds: ParquetScanBounds,
886 is_cancelled: F,
887 ) -> Result<ParquetTickCursor>
888 where
889 F: FnMut() -> bool,
890 {
891 ParquetTickCursor::open_cancellable(
892 self.root_path(),
893 exchange,
894 symbol,
895 bounds,
896 is_cancelled,
897 )
898 }
899
900 pub fn scan_bars(
902 &self,
903 exchange: &str,
904 symbol: &str,
905 timeframe: &str,
906 bounds: ParquetScanBounds,
907 ) -> Result<ParquetBarCursor> {
908 ParquetBarCursor::open(self.root_path(), exchange, symbol, timeframe, bounds)
909 }
910
911 pub fn scan_bars_cancellable<F>(
913 &self,
914 exchange: &str,
915 symbol: &str,
916 timeframe: &str,
917 bounds: ParquetScanBounds,
918 is_cancelled: F,
919 ) -> Result<ParquetBarCursor>
920 where
921 F: FnMut() -> bool,
922 {
923 ParquetBarCursor::open_cancellable(
924 self.root_path(),
925 exchange,
926 symbol,
927 timeframe,
928 bounds,
929 is_cancelled,
930 )
931 }
932}
933
934fn list_partitions(
935 directory: &Path,
936 bounds: ParquetScanBounds,
937 is_cancelled: &mut dyn FnMut() -> bool,
938) -> Result<Vec<PartitionDescriptor>> {
939 ensure_not_cancelled(is_cancelled)?;
940 if !directory.exists() {
941 return Ok(Vec::new());
942 }
943
944 let from_date = bounds.from.map(|ts| ts.date());
945 let to_date = bounds.to.map(|ts| ts.date());
946 let mut paths = Vec::new();
947 for entry in fs::read_dir(directory)? {
948 ensure_not_cancelled(is_cancelled)?;
949 let path = entry?.path();
950 if path
951 .extension()
952 .is_none_or(|extension| extension != "parquet")
953 {
954 continue;
955 }
956
957 let stem = path
958 .file_stem()
959 .and_then(|value| value.to_str())
960 .ok_or_else(|| DataError::InvalidDatePartition(path.display().to_string()))?;
961 let date = NaiveDate::parse_from_str(stem, "%Y-%m-%d")
962 .map_err(|_| DataError::InvalidDatePartition(path.display().to_string()))?;
963 if date.format("%Y-%m-%d").to_string() != stem {
964 return Err(DataError::InvalidDatePartition(path.display().to_string()));
965 }
966 if from_date.is_some_and(|from| date < from) || to_date.is_some_and(|to| date > to) {
967 continue;
968 }
969 paths.push((date, path));
970 }
971 paths.sort_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1)));
972
973 let mut partitions = Vec::with_capacity(paths.len());
974 let mut source_row_base = 0u64;
975 for (_, path) in paths {
976 ensure_not_cancelled(is_cancelled)?;
977 let descriptor = PartitionDescriptor::describe(path, source_row_base)?;
978 source_row_base =
979 source_row_base
980 .checked_add(u64::try_from(descriptor.row_count).map_err(|_| {
981 DataError::Other("Parquet partition row count exceeds u64".into())
982 })?)
983 .ok_or_else(|| DataError::Other("Parquet source row ordinal overflow".into()))?;
984 partitions.push(descriptor);
985 }
986 ensure_not_cancelled(is_cancelled)?;
987 Ok(partitions)
988}
989
990fn validate_monotonic_rows<T: Timestamped>(
991 mut previous: Option<NaiveDateTime>,
992 rows: &[T],
993 path: &Path,
994) -> Result<()> {
995 for row in rows {
996 let current = row.timestamp();
997 if let Some(previous) = previous
998 && current < previous
999 {
1000 return Err(DataError::NonMonotonicParquetData {
1001 path: path.display().to_string(),
1002 previous,
1003 current,
1004 });
1005 }
1006 previous = Some(current);
1007 }
1008 Ok(())
1009}
1010
1011fn ensure_not_cancelled(is_cancelled: &mut dyn FnMut() -> bool) -> Result<()> {
1012 if is_cancelled() {
1013 Err(DataError::Cancelled)
1014 } else {
1015 Ok(())
1016 }
1017}
1018
1019fn read_tick_partition(file: File, offset: usize, rows: usize) -> Result<Vec<Tick>> {
1020 let dataframe = ParquetReader::new(file)
1021 .with_slice(Some((offset, rows)))
1022 .finish()?;
1023 dataframe_to_ticks(&dataframe)
1024}
1025
1026fn read_bar_partition(file: File, offset: usize, rows: usize) -> Result<Vec<Bar>> {
1027 let dataframe = ParquetReader::new(file)
1028 .with_slice(Some((offset, rows)))
1029 .finish()?;
1030 dataframe_to_bars(&dataframe)
1031}
1032
1033fn tick_has_valid_quote(tick: &Tick) -> bool {
1034 match (tick.bid, tick.ask) {
1035 (Some(bid), Some(ask)) => {
1036 bid.is_finite() && ask.is_finite() && bid > 0.0 && ask > 0.0 && bid <= ask
1037 }
1038 _ => false,
1039 }
1040}
1041
1042#[cfg(test)]
1043mod tests {
1044 use std::time::{SystemTime, UNIX_EPOCH};
1045
1046 use super::*;
1047 use crate::convert::ticks_to_dataframe;
1048 use crate::models::Timeframe;
1049
1050 fn ts(day: u32, hour: u32) -> NaiveDateTime {
1051 NaiveDate::from_ymd_opt(2026, 1, day)
1052 .unwrap()
1053 .and_hms_opt(hour, 0, 0)
1054 .unwrap()
1055 }
1056
1057 fn temp_root(name: &str) -> PathBuf {
1058 let nonce = SystemTime::now()
1059 .duration_since(UNIX_EPOCH)
1060 .unwrap()
1061 .as_nanos();
1062 std::env::temp_dir().join(format!(
1063 "qs-data-preprocess-{name}-{}-{nonce}",
1064 std::process::id()
1065 ))
1066 }
1067
1068 fn tick(at: NaiveDateTime) -> Tick {
1069 Tick {
1070 exchange: "test".into(),
1071 symbol: "EURUSD".into(),
1072 ts: at,
1073 bid: Some(1.0),
1074 ask: Some(1.1),
1075 last: None,
1076 volume: None,
1077 flags: None,
1078 }
1079 }
1080
1081 fn bar(at: NaiveDateTime) -> Bar {
1082 Bar {
1083 exchange: "test".into(),
1084 symbol: "EURUSD".into(),
1085 timeframe: Timeframe::M1,
1086 ts: at,
1087 open: 1.0,
1088 high: 1.2,
1089 low: 0.9,
1090 close: 1.1,
1091 tick_vol: 10,
1092 volume: 10,
1093 spread: 1,
1094 }
1095 }
1096
1097 #[test]
1098 fn tick_cursor_scans_partitions_in_ascending_bounded_order() {
1099 let root = temp_root("ticks");
1100 let store = ParquetStore::open(&root).unwrap();
1101 store
1102 .insert_ticks(&[tick(ts(2, 1)), tick(ts(1, 2)), tick(ts(1, 1))])
1103 .unwrap();
1104
1105 let mut cursor = store
1106 .scan_ticks(
1107 "test",
1108 "EURUSD",
1109 ParquetScanBounds::new(Some(ts(1, 2)), Some(ts(2, 1))),
1110 )
1111 .unwrap();
1112 assert_eq!(cursor.remaining_partitions(), 2);
1113 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(1, 2));
1114 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(2, 1));
1115 assert!(cursor.next_tick().unwrap().is_none());
1116
1117 fs::remove_dir_all(root).ok();
1118 }
1119
1120 #[test]
1121 fn bar_cursor_scans_ordered_date_partitions() {
1122 let root = temp_root("bars");
1123 let store = ParquetStore::open(&root).unwrap();
1124 store.insert_bars(&[bar(ts(2, 1)), bar(ts(1, 1))]).unwrap();
1125
1126 let mut cursor = store
1127 .scan_bars("test", "EURUSD", "1m", ParquetScanBounds::default())
1128 .unwrap();
1129 assert_eq!(cursor.next_bar().unwrap().unwrap().ts, ts(1, 1));
1130 assert_eq!(cursor.next_bar().unwrap().unwrap().ts, ts(2, 1));
1131 assert!(cursor.next_bar().unwrap().is_none());
1132
1133 fs::remove_dir_all(root).ok();
1134 }
1135
1136 #[test]
1137 fn cursor_reads_large_partitions_through_bounded_slices() {
1138 let root = temp_root("slices");
1139 let store = ParquetStore::open(&root).unwrap();
1140 store
1141 .insert_ticks(&[tick(ts(1, 1)), tick(ts(1, 2)), tick(ts(1, 3))])
1142 .unwrap();
1143 let mut cursor = ParquetTickCursor::open_with_read_size(
1144 &root,
1145 "test",
1146 "EURUSD",
1147 ParquetScanBounds::default(),
1148 1,
1149 )
1150 .unwrap();
1151
1152 assert_eq!(cursor.rows_per_read(), 1);
1153 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(1, 1));
1154 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(1, 2));
1155 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(1, 3));
1156 assert!(cursor.next_tick().unwrap().is_none());
1157 assert!(matches!(
1158 ParquetTickCursor::open_with_read_size(
1159 &root,
1160 "test",
1161 "EURUSD",
1162 ParquetScanBounds::default(),
1163 0,
1164 ),
1165 Err(DataError::InvalidScanReadSize)
1166 ));
1167
1168 fs::remove_dir_all(root).ok();
1169 }
1170
1171 #[test]
1172 fn cursor_cancellation_does_not_consume_the_partition() {
1173 let root = temp_root("cancel");
1174 let store = ParquetStore::open(&root).unwrap();
1175 store.insert_ticks(&[tick(ts(1, 1))]).unwrap();
1176 let mut cursor = store
1177 .scan_ticks("test", "EURUSD", ParquetScanBounds::default())
1178 .unwrap();
1179 let mut checks = 0;
1180
1181 let error = cursor
1182 .next_tick_cancellable(|| {
1183 checks += 1;
1184 checks == 3
1185 })
1186 .unwrap_err();
1187 assert!(matches!(error, DataError::Cancelled));
1188 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(1, 1));
1189
1190 fs::remove_dir_all(root).ok();
1191 }
1192
1193 #[test]
1194 fn cursor_rejects_non_monotonic_partition_rows_before_emitting_them() {
1195 let root = temp_root("monotonic");
1196 let path = root
1197 .join("ticks")
1198 .join("exchange=test")
1199 .join("symbol=EURUSD")
1200 .join("2026-01-01.parquet");
1201 fs::create_dir_all(path.parent().unwrap()).unwrap();
1202 let mut dataframe = ticks_to_dataframe(&[tick(ts(1, 2)), tick(ts(1, 1))]).unwrap();
1203 ParquetWriter::new(File::create(&path).unwrap())
1204 .finish(&mut dataframe)
1205 .unwrap();
1206
1207 let mut cursor =
1208 ParquetTickCursor::open(&root, "test", "EURUSD", ParquetScanBounds::default()).unwrap();
1209 assert!(matches!(
1210 cursor.next_tick(),
1211 Err(DataError::NonMonotonicParquetData {
1212 previous,
1213 current,
1214 ..
1215 }) if previous == ts(1, 2) && current == ts(1, 1)
1216 ));
1217
1218 fs::remove_dir_all(root).ok();
1219 }
1220
1221 #[test]
1222 fn cursor_reports_physical_ordinals_across_filtered_rows_and_partitions() {
1223 let root = temp_root("ordinals");
1224 let store = ParquetStore::open(&root).unwrap();
1225 store
1226 .insert_ticks(&[
1227 tick(ts(1, 1)),
1228 tick(ts(1, 2)),
1229 tick(ts(1, 3)),
1230 tick(ts(2, 1)),
1231 ])
1232 .unwrap();
1233 let mut cursor = ParquetTickCursor::open_with_read_size(
1234 &root,
1235 "test",
1236 "EURUSD",
1237 ParquetScanBounds::new(Some(ts(1, 2)), None),
1238 2,
1239 )
1240 .unwrap();
1241
1242 let rows = [
1243 cursor.next_tick_with_ordinal().unwrap().unwrap(),
1244 cursor.next_tick_with_ordinal().unwrap().unwrap(),
1245 cursor.next_tick_with_ordinal().unwrap().unwrap(),
1246 ];
1247 assert_eq!(
1248 rows.map(|row| (row.row.ts, row.source_row_ordinal)),
1249 [(ts(1, 2), 1), (ts(1, 3), 2), (ts(2, 1), 3)]
1250 );
1251
1252 fs::remove_dir_all(root).ok();
1253 }
1254
1255 #[test]
1256 fn running_cursor_fails_if_atomic_replacement_changes_its_partition() {
1257 let root = temp_root("replace-running");
1258 let store = ParquetStore::open(&root).unwrap();
1259 store
1260 .insert_ticks(&[tick(ts(1, 1)), tick(ts(1, 2)), tick(ts(1, 3))])
1261 .unwrap();
1262 let mut cursor = ParquetTickCursor::open_with_read_size(
1263 &root,
1264 "test",
1265 "EURUSD",
1266 ParquetScanBounds::default(),
1267 1,
1268 )
1269 .unwrap();
1270 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(1, 1));
1271 let partition_path = root
1272 .join("ticks")
1273 .join("exchange=test")
1274 .join("symbol=EURUSD")
1275 .join("2026-01-01.parquet");
1276 #[cfg(unix)]
1277 let original_inode = {
1278 use std::os::unix::fs::MetadataExt as _;
1279 fs::metadata(&partition_path).unwrap().ino()
1280 };
1281
1282 store.insert_ticks(&[tick(ts(1, 4))]).unwrap();
1283 #[cfg(unix)]
1284 {
1285 use std::os::unix::fs::MetadataExt as _;
1286 assert_ne!(original_inode, fs::metadata(&partition_path).unwrap().ino());
1287 }
1288 assert!(
1289 fs::read_dir(partition_path.parent().unwrap())
1290 .unwrap()
1291 .all(|entry| !entry
1292 .unwrap()
1293 .file_name()
1294 .to_string_lossy()
1295 .ends_with(".tmp"))
1296 );
1297 assert!(matches!(
1298 cursor.next_tick(),
1299 Err(DataError::ParquetPartitionChanged { .. })
1300 ));
1301
1302 fs::remove_dir_all(root).ok();
1303 }
1304
1305 #[test]
1306 fn described_scan_rejects_replacement_before_reopen() {
1307 let root = temp_root("replace-described");
1308 let store = ParquetStore::open(&root).unwrap();
1309 store.insert_ticks(&[tick(ts(1, 1))]).unwrap();
1310 let scan = ParquetTickScan::describe(&root, "test", "EURUSD", ParquetScanBounds::default())
1311 .unwrap();
1312
1313 store.insert_ticks(&[tick(ts(1, 2))]).unwrap();
1314 assert!(matches!(
1315 scan.cursor(),
1316 Err(DataError::ParquetPartitionChanged { .. })
1317 ));
1318
1319 fs::remove_dir_all(root).ok();
1320 }
1321
1322 #[test]
1323 fn upper_bound_stops_at_the_first_later_monotonic_row() {
1324 let root = temp_root("upper-bound");
1325 let store = ParquetStore::open(&root).unwrap();
1326 store
1327 .insert_ticks(&[
1328 tick(ts(1, 1)),
1329 tick(ts(1, 2)),
1330 tick(ts(1, 3)),
1331 tick(ts(1, 4)),
1332 ])
1333 .unwrap();
1334 let mut cursor = ParquetTickCursor::open_with_read_size(
1335 &root,
1336 "test",
1337 "EURUSD",
1338 ParquetScanBounds::new(None, Some(ts(1, 2))),
1339 3,
1340 )
1341 .unwrap();
1342
1343 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(1, 1));
1344 assert_eq!(cursor.next_tick().unwrap().unwrap().ts, ts(1, 2));
1345 assert!(cursor.next_tick().unwrap().is_none());
1346 assert_eq!(cursor.remaining_partitions(), 0);
1347
1348 fs::remove_dir_all(root).ok();
1349 }
1350
1351 #[test]
1352 fn reverse_chunks_find_the_latest_valid_tick_without_materializing_the_day() {
1353 let root = temp_root("latest-reverse");
1354 let store = ParquetStore::open(&root).unwrap();
1355 let mut invalid = tick(ts(1, 4));
1356 invalid.ask = None;
1357 store
1358 .insert_ticks(&[tick(ts(1, 1)), tick(ts(1, 2)), tick(ts(1, 3)), invalid])
1359 .unwrap();
1360 let scan = ParquetTickScan::describe(&root, "test", "EURUSD", ParquetScanBounds::default())
1361 .unwrap();
1362
1363 let latest = find_latest_row(
1364 &scan.inner,
1365 2,
1366 read_tick_partition,
1367 tick_has_valid_quote,
1368 &mut || false,
1369 )
1370 .unwrap()
1371 .unwrap();
1372 assert_eq!(latest.row.ts, ts(1, 3));
1373 assert_eq!(latest.source_row_ordinal, 2);
1374
1375 fs::remove_dir_all(root).ok();
1376 }
1377}