1use std::collections::HashMap;
8use std::ffi::OsString;
9use std::fs::{self, File, OpenOptions};
10use std::path::{Path, PathBuf};
11use std::sync::atomic::{AtomicU64, Ordering};
12
13use chrono::NaiveDateTime;
14use polars::prelude::*;
15
16use crate::convert::{
17 bars_to_dataframe, dataframe_to_bars, dataframe_to_ticks, ndt_to_date_string,
18 ticks_to_dataframe,
19};
20use crate::error::{DataError, Result};
21use crate::models::{Bar, BarQueryOpts, QueryOpts, StatRow, Tick};
22use crate::scanner::{ParquetScanBounds, ParquetTickScan};
23
24pub struct ParquetStore {
26 root: PathBuf,
27}
28
29impl ParquetStore {
30 pub fn open(root: impl AsRef<Path>) -> Result<Self> {
32 let root = root.as_ref().to_path_buf();
33 fs::create_dir_all(&root)?;
34 Ok(Self { root })
35 }
36
37 pub fn insert_ticks(&self, ticks: &[Tick]) -> Result<usize> {
42 if ticks.is_empty() {
43 return Ok(0);
44 }
45
46 let mut groups: HashMap<(String, String, String), Vec<&Tick>> = HashMap::new();
48 for tick in ticks {
49 let date = ndt_to_date_string(&tick.ts);
50 let key = (tick.exchange.clone(), tick.symbol.clone(), date);
51 groups.entry(key).or_default().push(tick);
52 }
53
54 let mut total_inserted = 0usize;
55
56 for ((exchange, symbol, date), group_ticks) in &groups {
57 let dir = self.tick_dir(exchange, symbol);
58 fs::create_dir_all(&dir)?;
59 let file_path = dir.join(format!("{date}.parquet"));
60
61 let owned: Vec<Tick> = group_ticks.iter().map(|t| (*t).clone()).collect();
62 let new_df = ticks_to_dataframe(&owned)?;
63
64 if file_path.exists() {
65 let existing_df = read_parquet_file(&file_path)?;
66 let existing_count = existing_df.height();
67 let combined = concat_and_dedup_ticks(existing_df, new_df)?;
68 total_inserted += combined.height().saturating_sub(existing_count);
69 write_parquet_file(&file_path, &mut combined.clone())?;
70 } else {
71 let deduped = dedup_ticks(new_df)?;
72 total_inserted += deduped.height();
73 write_parquet_file(&file_path, &mut deduped.clone())?;
74 }
75 }
76
77 Ok(total_inserted)
78 }
79
80 pub fn insert_bars(&self, bars: &[Bar]) -> Result<usize> {
83 if bars.is_empty() {
84 return Ok(0);
85 }
86
87 let mut groups: HashMap<(String, String, String, String), Vec<&Bar>> = HashMap::new();
89 for bar in bars {
90 let date = ndt_to_date_string(&bar.ts);
91 let key = (
92 bar.exchange.clone(),
93 bar.symbol.clone(),
94 bar.timeframe.as_str().to_string(),
95 date,
96 );
97 groups.entry(key).or_default().push(bar);
98 }
99
100 let mut total_inserted = 0usize;
101
102 for ((exchange, symbol, timeframe, date), group_bars) in &groups {
103 let dir = self.bar_dir(exchange, symbol, timeframe);
104 fs::create_dir_all(&dir)?;
105 let file_path = dir.join(format!("{date}.parquet"));
106
107 let owned: Vec<Bar> = group_bars.iter().map(|b| (*b).clone()).collect();
108 let new_df = bars_to_dataframe(&owned)?;
109
110 if file_path.exists() {
111 let existing_df = read_parquet_file(&file_path)?;
112 let existing_count = existing_df.height();
113 let combined = concat_and_dedup_bars(existing_df, new_df)?;
114 total_inserted += combined.height().saturating_sub(existing_count);
115 write_parquet_file(&file_path, &mut combined.clone())?;
116 } else {
117 let deduped = dedup_bars(new_df)?;
118 total_inserted += deduped.height();
119 write_parquet_file(&file_path, &mut deduped.clone())?;
120 }
121 }
122
123 Ok(total_inserted)
124 }
125
126 pub fn query_ticks(&self, opts: &QueryOpts) -> Result<(Vec<Tick>, u64)> {
131 self.query_ticks_cancellable(opts, || false)
132 }
133
134 pub fn query_ticks_cancellable<F>(
141 &self,
142 opts: &QueryOpts,
143 mut is_cancelled: F,
144 ) -> Result<(Vec<Tick>, u64)>
145 where
146 F: FnMut() -> bool,
147 {
148 ensure_not_cancelled(&mut is_cancelled)?;
149 let dir = self.tick_dir(&opts.exchange, &opts.symbol);
150 if !dir.exists() {
151 return Ok((Vec::new(), 0));
152 }
153
154 let files = list_date_files_cancellable(&dir, opts.from, opts.to, &mut is_cancelled)?;
155 if files.is_empty() {
156 return Ok((Vec::new(), 0));
157 }
158
159 let mut all_dfs: Vec<DataFrame> = Vec::with_capacity(files.len());
160 for file in &files {
161 ensure_not_cancelled(&mut is_cancelled)?;
162 let df = read_parquet_file(file)?;
163 ensure_not_cancelled(&mut is_cancelled)?;
164 all_dfs.push(df);
165 }
166 let mut combined = concat_dataframes(all_dfs)?;
167 ensure_not_cancelled(&mut is_cancelled)?;
168
169 combined = apply_ts_filter(combined, opts.from, opts.to)?;
170 ensure_not_cancelled(&mut is_cancelled)?;
171 combined = combined.sort(["ts"], SortMultipleOptions::default())?;
172 ensure_not_cancelled(&mut is_cancelled)?;
173
174 let total = combined.height() as u64;
175 combined = apply_pagination(combined, opts.limit, opts.tail, opts.descending)?;
176 ensure_not_cancelled(&mut is_cancelled)?;
177
178 let ticks = dataframe_to_ticks(&combined)?;
179 ensure_not_cancelled(&mut is_cancelled)?;
180 Ok((ticks, total))
181 }
182
183 pub fn latest_valid_tick_before(
185 &self,
186 exchange: &str,
187 symbol: &str,
188 before: NaiveDateTime,
189 ) -> Result<Option<Tick>> {
190 self.latest_valid_tick_before_cancellable(exchange, symbol, before, || false)
191 }
192
193 pub fn latest_valid_tick_before_cancellable<F>(
198 &self,
199 exchange: &str,
200 symbol: &str,
201 before: NaiveDateTime,
202 mut is_cancelled: F,
203 ) -> Result<Option<Tick>>
204 where
205 F: FnMut() -> bool,
206 {
207 ensure_not_cancelled(&mut is_cancelled)?;
208 let scan = ParquetTickScan::describe_cancellable(
209 &self.root,
210 exchange,
211 symbol,
212 ParquetScanBounds::new(None, Some(before)),
213 &mut is_cancelled,
214 )?;
215 let latest = scan
216 .latest_valid_tick_before_cancellable(before, &mut is_cancelled)?
217 .map(|row| row.row);
218 ensure_not_cancelled(&mut is_cancelled)?;
219 Ok(latest)
220 }
221
222 pub fn query_bars(&self, opts: &BarQueryOpts) -> Result<(Vec<Bar>, u64)> {
225 self.query_bars_cancellable(opts, || false)
226 }
227
228 pub fn query_bars_cancellable<F>(
231 &self,
232 opts: &BarQueryOpts,
233 mut is_cancelled: F,
234 ) -> Result<(Vec<Bar>, u64)>
235 where
236 F: FnMut() -> bool,
237 {
238 ensure_not_cancelled(&mut is_cancelled)?;
239 let dir = self.bar_dir(&opts.exchange, &opts.symbol, &opts.timeframe);
240 if !dir.exists() {
241 return Ok((Vec::new(), 0));
242 }
243
244 let files = list_date_files_cancellable(&dir, opts.from, opts.to, &mut is_cancelled)?;
245 if files.is_empty() {
246 return Ok((Vec::new(), 0));
247 }
248
249 let mut all_dfs: Vec<DataFrame> = Vec::with_capacity(files.len());
250 for file in &files {
251 ensure_not_cancelled(&mut is_cancelled)?;
252 let df = read_parquet_file(file)?;
253 ensure_not_cancelled(&mut is_cancelled)?;
254 all_dfs.push(df);
255 }
256 let mut combined = concat_dataframes(all_dfs)?;
257 ensure_not_cancelled(&mut is_cancelled)?;
258
259 combined = apply_ts_filter(combined, opts.from, opts.to)?;
260 ensure_not_cancelled(&mut is_cancelled)?;
261 combined = combined.sort(["ts"], SortMultipleOptions::default())?;
262 ensure_not_cancelled(&mut is_cancelled)?;
263
264 let total = combined.height() as u64;
265 combined = apply_pagination(combined, opts.limit, opts.tail, opts.descending)?;
266 ensure_not_cancelled(&mut is_cancelled)?;
267
268 let bars = dataframe_to_bars(&combined)?;
269 ensure_not_cancelled(&mut is_cancelled)?;
270 Ok((bars, total))
271 }
272
273 pub fn delete_ticks(
277 &self,
278 exchange: &str,
279 symbol: &str,
280 from: Option<NaiveDateTime>,
281 to: Option<NaiveDateTime>,
282 ) -> Result<usize> {
283 let dir = self.tick_dir(exchange, symbol);
284 if !dir.exists() {
285 return Ok(0);
286 }
287 delete_from_partition(&dir, from, to)
288 }
289
290 pub fn delete_bars(
292 &self,
293 exchange: &str,
294 symbol: &str,
295 timeframe: &str,
296 from: Option<NaiveDateTime>,
297 to: Option<NaiveDateTime>,
298 ) -> Result<usize> {
299 let dir = self.bar_dir(exchange, symbol, timeframe);
300 if !dir.exists() {
301 return Ok(0);
302 }
303 delete_from_partition(&dir, from, to)
304 }
305
306 pub fn delete_symbol(&self, exchange: &str, symbol: &str) -> Result<(usize, usize)> {
308 let tick_count = self.count_rows_in_dir(&self.tick_dir(exchange, symbol));
309 let bar_count = self.count_all_bars_for_symbol(exchange, symbol);
310
311 let tick_dir = self.tick_dir(exchange, symbol);
313 if tick_dir.exists() {
314 fs::remove_dir_all(&tick_dir)?;
315 }
316
317 let bar_sym_dir = self
319 .root
320 .join("bars")
321 .join(format!("exchange={exchange}"))
322 .join(format!("symbol={symbol}"));
323 if bar_sym_dir.exists() {
324 fs::remove_dir_all(&bar_sym_dir)?;
325 }
326
327 Ok((tick_count, bar_count))
328 }
329
330 pub fn delete_exchange(&self, exchange: &str) -> Result<(usize, usize)> {
332 let tick_ex_dir = self.root.join("ticks").join(format!("exchange={exchange}"));
333 let bar_ex_dir = self.root.join("bars").join(format!("exchange={exchange}"));
334
335 let tick_count = self.count_rows_recursive(&tick_ex_dir);
336 let bar_count = self.count_rows_recursive(&bar_ex_dir);
337
338 if tick_ex_dir.exists() {
339 fs::remove_dir_all(&tick_ex_dir)?;
340 }
341 if bar_ex_dir.exists() {
342 fs::remove_dir_all(&bar_ex_dir)?;
343 }
344
345 Ok((tick_count, bar_count))
346 }
347
348 pub fn stats(&self, exchange: Option<&str>, symbol: Option<&str>) -> Result<Vec<StatRow>> {
352 let mut rows = Vec::new();
353
354 self.collect_tick_stats(&mut rows, exchange, symbol)?;
356
357 self.collect_bar_stats(&mut rows, exchange, symbol)?;
359
360 rows.sort_by(|a, b| {
362 a.exchange
363 .cmp(&b.exchange)
364 .then(a.symbol.cmp(&b.symbol))
365 .then(a.data_type.cmp(&b.data_type))
366 });
367
368 Ok(rows)
369 }
370
371 pub fn total_size(&self) -> Option<u64> {
373 let mut total = 0u64;
374 for entry in walkdir(&self.root) {
375 if entry.extension().is_some_and(|e| e == "parquet")
376 && let Ok(meta) = fs::metadata(&entry)
377 {
378 total += meta.len();
379 }
380 }
381 if total == 0 { None } else { Some(total) }
382 }
383
384 pub(crate) fn root_path(&self) -> &Path {
387 &self.root
388 }
389
390 fn tick_dir(&self, exchange: &str, symbol: &str) -> PathBuf {
392 self.root
393 .join("ticks")
394 .join(format!("exchange={exchange}"))
395 .join(format!("symbol={symbol}"))
396 }
397
398 fn bar_dir(&self, exchange: &str, symbol: &str, timeframe: &str) -> PathBuf {
400 self.root
401 .join("bars")
402 .join(format!("exchange={exchange}"))
403 .join(format!("symbol={symbol}"))
404 .join(format!("timeframe={timeframe}"))
405 }
406
407 fn count_rows_in_dir(&self, dir: &Path) -> usize {
409 if !dir.exists() {
410 return 0;
411 }
412 let mut count = 0;
413 if let Ok(entries) = fs::read_dir(dir) {
414 for entry in entries.flatten() {
415 let path = entry.path();
416 if path.extension().is_some_and(|e| e == "parquet")
417 && let Ok(df) = read_parquet_file(&path)
418 {
419 count += df.height();
420 }
421 }
422 }
423 count
424 }
425
426 fn count_rows_recursive(&self, dir: &Path) -> usize {
428 if !dir.exists() {
429 return 0;
430 }
431 let mut count = 0;
432 for path in walkdir(dir) {
433 if path.extension().is_some_and(|e| e == "parquet")
434 && let Ok(df) = read_parquet_file(&path)
435 {
436 count += df.height();
437 }
438 }
439 count
440 }
441
442 fn count_all_bars_for_symbol(&self, exchange: &str, symbol: &str) -> usize {
444 let bar_sym_dir = self
445 .root
446 .join("bars")
447 .join(format!("exchange={exchange}"))
448 .join(format!("symbol={symbol}"));
449 self.count_rows_recursive(&bar_sym_dir)
450 }
451
452 fn collect_tick_stats(
454 &self,
455 rows: &mut Vec<StatRow>,
456 exchange_filter: Option<&str>,
457 symbol_filter: Option<&str>,
458 ) -> Result<()> {
459 let ticks_dir = self.root.join("ticks");
460 if !ticks_dir.exists() {
461 return Ok(());
462 }
463
464 for (exchange, symbol, dir) in self.iter_exchange_symbol_dirs(&ticks_dir)? {
465 if let Some(ef) = exchange_filter
466 && exchange != ef
467 {
468 continue;
469 }
470 if let Some(sf) = symbol_filter
471 && symbol != sf
472 {
473 continue;
474 }
475
476 let (count, ts_min, ts_max) = self.aggregate_parquet_stats(&dir)?;
477 if count > 0 {
478 rows.push(StatRow {
479 exchange,
480 symbol,
481 data_type: "tick".to_string(),
482 count,
483 ts_min: ts_min.unwrap_or_default(),
484 ts_max: ts_max.unwrap_or_default(),
485 });
486 }
487 }
488
489 Ok(())
490 }
491
492 fn collect_bar_stats(
494 &self,
495 rows: &mut Vec<StatRow>,
496 exchange_filter: Option<&str>,
497 symbol_filter: Option<&str>,
498 ) -> Result<()> {
499 let bars_dir = self.root.join("bars");
500 if !bars_dir.exists() {
501 return Ok(());
502 }
503
504 for (exchange, symbol, timeframe, dir) in self.iter_exchange_symbol_tf_dirs(&bars_dir)? {
505 if let Some(ef) = exchange_filter
506 && exchange != ef
507 {
508 continue;
509 }
510 if let Some(sf) = symbol_filter
511 && symbol != sf
512 {
513 continue;
514 }
515
516 let (count, ts_min, ts_max) = self.aggregate_parquet_stats(&dir)?;
517 if count > 0 {
518 rows.push(StatRow {
519 exchange,
520 symbol,
521 data_type: format!("bar ({timeframe})"),
522 count,
523 ts_min: ts_min.unwrap_or_default(),
524 ts_max: ts_max.unwrap_or_default(),
525 });
526 }
527 }
528
529 Ok(())
530 }
531
532 fn iter_exchange_symbol_dirs(&self, base: &Path) -> Result<Vec<(String, String, PathBuf)>> {
534 let mut result = Vec::new();
535 if !base.exists() {
536 return Ok(result);
537 }
538
539 for ex_entry in fs::read_dir(base)?.flatten() {
540 let ex_path = ex_entry.path();
541 if !ex_path.is_dir() {
542 continue;
543 }
544 let exchange =
545 parse_partition_value(ex_path.file_name().unwrap().to_str().unwrap_or(""));
546 if exchange.is_empty() {
547 continue;
548 }
549
550 for sym_entry in fs::read_dir(&ex_path)?.flatten() {
551 let sym_path = sym_entry.path();
552 if !sym_path.is_dir() {
553 continue;
554 }
555 let symbol =
556 parse_partition_value(sym_path.file_name().unwrap().to_str().unwrap_or(""));
557 if symbol.is_empty() {
558 continue;
559 }
560 result.push((exchange.clone(), symbol, sym_path));
561 }
562 }
563
564 Ok(result)
565 }
566
567 fn iter_exchange_symbol_tf_dirs(
569 &self,
570 base: &Path,
571 ) -> Result<Vec<(String, String, String, PathBuf)>> {
572 let mut result = Vec::new();
573 if !base.exists() {
574 return Ok(result);
575 }
576
577 for ex_entry in fs::read_dir(base)?.flatten() {
578 let ex_path = ex_entry.path();
579 if !ex_path.is_dir() {
580 continue;
581 }
582 let exchange =
583 parse_partition_value(ex_path.file_name().unwrap().to_str().unwrap_or(""));
584 if exchange.is_empty() {
585 continue;
586 }
587
588 for sym_entry in fs::read_dir(&ex_path)?.flatten() {
589 let sym_path = sym_entry.path();
590 if !sym_path.is_dir() {
591 continue;
592 }
593 let symbol =
594 parse_partition_value(sym_path.file_name().unwrap().to_str().unwrap_or(""));
595 if symbol.is_empty() {
596 continue;
597 }
598
599 for tf_entry in fs::read_dir(&sym_path)?.flatten() {
600 let tf_path = tf_entry.path();
601 if !tf_path.is_dir() {
602 continue;
603 }
604 let timeframe =
605 parse_partition_value(tf_path.file_name().unwrap().to_str().unwrap_or(""));
606 if timeframe.is_empty() {
607 continue;
608 }
609 result.push((exchange.clone(), symbol.clone(), timeframe, tf_path));
610 }
611 }
612 }
613
614 Ok(result)
615 }
616
617 fn aggregate_parquet_stats(
619 &self,
620 dir: &Path,
621 ) -> Result<(u64, Option<NaiveDateTime>, Option<NaiveDateTime>)> {
622 let mut total_count = 0u64;
623 let mut global_min: Option<i64> = None;
624 let mut global_max: Option<i64> = None;
625
626 if !dir.exists() {
627 return Ok((0, None, None));
628 }
629
630 for entry in fs::read_dir(dir)?.flatten() {
631 let path = entry.path();
632 if path.extension().is_some_and(|e| e == "parquet") {
633 let df = read_parquet_file(&path)?;
634 total_count += df.height() as u64;
635
636 if df.height() > 0 {
637 let ts_col = df.column("ts").ok().and_then(|c| c.datetime().ok());
638 if let Some(ts) = ts_col {
639 if let Some(min_val) = ts.min() {
640 global_min =
641 Some(global_min.map_or(min_val, |cur: i64| cur.min(min_val)));
642 }
643 if let Some(max_val) = ts.max() {
644 global_max =
645 Some(global_max.map_or(max_val, |cur: i64| cur.max(max_val)));
646 }
647 }
648 }
649 }
650 }
651
652 let ts_min = global_min.map(micros_to_ndt);
653 let ts_max = global_max.map(micros_to_ndt);
654
655 Ok((total_count, ts_min, ts_max))
656 }
657}
658
659fn parse_partition_value(dir_name: &str) -> String {
663 dir_name
664 .split_once('=')
665 .map(|(_, v)| v.to_string())
666 .unwrap_or_default()
667}
668
669fn ensure_not_cancelled(is_cancelled: &mut dyn FnMut() -> bool) -> Result<()> {
670 if is_cancelled() {
671 Err(DataError::Cancelled)
672 } else {
673 Ok(())
674 }
675}
676
677fn list_date_files(
679 dir: &Path,
680 from: Option<NaiveDateTime>,
681 to: Option<NaiveDateTime>,
682) -> Result<Vec<PathBuf>> {
683 list_date_files_cancellable(dir, from, to, &mut || false)
684}
685
686fn list_date_files_cancellable(
687 dir: &Path,
688 from: Option<NaiveDateTime>,
689 to: Option<NaiveDateTime>,
690 is_cancelled: &mut dyn FnMut() -> bool,
691) -> Result<Vec<PathBuf>> {
692 ensure_not_cancelled(is_cancelled)?;
693 let mut files = Vec::new();
694 let from_date = from.map(|d| d.format("%Y-%m-%d").to_string());
695 let to_date = to.map(|d| d.format("%Y-%m-%d").to_string());
696
697 for entry in fs::read_dir(dir)? {
698 ensure_not_cancelled(is_cancelled)?;
699 let path = entry?.path();
700 if path.extension().is_some_and(|e| e == "parquet") {
701 let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("");
702
703 let dominated_by_from = from_date.as_ref().is_some_and(|fd| stem < fd.as_str());
705 let past_to = to_date.as_ref().is_some_and(|td| stem > td.as_str());
706
707 if !dominated_by_from && !past_to {
708 files.push(path);
709 }
710 }
711 }
712
713 ensure_not_cancelled(is_cancelled)?;
714 files.sort();
715 Ok(files)
716}
717
718fn read_parquet_file(path: &Path) -> Result<DataFrame> {
720 let file = std::fs::File::open(path)?;
721 let df = ParquetReader::new(file).finish()?;
722 Ok(df)
723}
724
725static NEXT_TEMP_FILE_ID: AtomicU64 = AtomicU64::new(0);
726
727fn write_parquet_file(path: &Path, df: &mut DataFrame) -> Result<()> {
729 let (temp_path, mut file) = create_partition_temp_file(path)?;
730 let write_result = (|| -> Result<()> {
731 ParquetWriter::new(&mut file)
732 .with_compression(ParquetCompression::Zstd(None))
733 .finish(df)?;
734 file.sync_all()?;
735 Ok(())
736 })();
737 drop(file);
738
739 if let Err(error) = write_result {
740 fs::remove_file(&temp_path).ok();
741 return Err(error);
742 }
743 if let Err(error) = atomic_replace(&temp_path, path) {
744 fs::remove_file(&temp_path).ok();
745 return Err(error.into());
746 }
747 Ok(())
748}
749
750fn create_partition_temp_file(path: &Path) -> Result<(PathBuf, File)> {
751 let parent = path.parent().ok_or_else(|| {
752 DataError::Other(format!("partition path has no parent: {}", path.display()))
753 })?;
754 let file_name = path.file_name().ok_or_else(|| {
755 DataError::Other(format!(
756 "partition path has no file name: {}",
757 path.display()
758 ))
759 })?;
760
761 loop {
762 let id = NEXT_TEMP_FILE_ID.fetch_add(1, Ordering::Relaxed);
763 let mut temp_name = OsString::from(".");
764 temp_name.push(file_name);
765 temp_name.push(format!(".{}.{}.tmp", std::process::id(), id));
766 let temp_path = parent.join(temp_name);
767 match OpenOptions::new()
768 .write(true)
769 .create_new(true)
770 .open(&temp_path)
771 {
772 Ok(file) => return Ok((temp_path, file)),
773 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue,
774 Err(error) => return Err(error.into()),
775 }
776 }
777}
778
779#[cfg(not(windows))]
780fn atomic_replace(from: &Path, to: &Path) -> std::io::Result<()> {
781 fs::rename(from, to)
782}
783
784#[cfg(windows)]
785fn atomic_replace(from: &Path, to: &Path) -> std::io::Result<()> {
786 use std::os::windows::ffi::OsStrExt;
787
788 const MOVEFILE_REPLACE_EXISTING: u32 = 0x1;
789 const MOVEFILE_WRITE_THROUGH: u32 = 0x8;
790
791 #[link(name = "kernel32")]
792 unsafe extern "system" {
793 fn MoveFileExW(
794 existing_file_name: *const u16,
795 new_file_name: *const u16,
796 flags: u32,
797 ) -> i32;
798 }
799
800 let from = from
801 .as_os_str()
802 .encode_wide()
803 .chain(Some(0))
804 .collect::<Vec<_>>();
805 let to = to
806 .as_os_str()
807 .encode_wide()
808 .chain(Some(0))
809 .collect::<Vec<_>>();
810 let replaced = unsafe {
811 MoveFileExW(
812 from.as_ptr(),
813 to.as_ptr(),
814 MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH,
815 )
816 };
817 if replaced == 0 {
818 Err(std::io::Error::last_os_error())
819 } else {
820 Ok(())
821 }
822}
823
824fn concat_and_dedup_ticks(existing: DataFrame, new: DataFrame) -> Result<DataFrame> {
826 let combined = concat_dataframes(vec![existing, new])?;
827 dedup_ticks(combined)
828}
829
830fn dedup_ticks(df: DataFrame) -> Result<DataFrame> {
832 let cols: Vec<String> = vec!["exchange".into(), "symbol".into(), "ts".into()];
833 let deduped = df
834 .unique_stable(Some(&cols), UniqueKeepStrategy::First, None)?
835 .sort(["ts"], SortMultipleOptions::default())?;
836 Ok(deduped)
837}
838
839fn concat_and_dedup_bars(existing: DataFrame, new: DataFrame) -> Result<DataFrame> {
841 let combined = concat_dataframes(vec![existing, new])?;
842 dedup_bars(combined)
843}
844
845fn dedup_bars(df: DataFrame) -> Result<DataFrame> {
847 let cols: Vec<String> = vec![
848 "exchange".into(),
849 "symbol".into(),
850 "timeframe".into(),
851 "ts".into(),
852 ];
853 let deduped = df
854 .unique_stable(Some(&cols), UniqueKeepStrategy::First, None)?
855 .sort(["ts"], SortMultipleOptions::default())?;
856 Ok(deduped)
857}
858
859fn concat_dataframes(dfs: Vec<DataFrame>) -> Result<DataFrame> {
861 if dfs.is_empty() {
862 return Err(DataError::Other("no dataframes to concat".into()));
863 }
864 if dfs.len() == 1 {
865 return Ok(dfs.into_iter().next().unwrap());
866 }
867 let lazy_frames: Vec<LazyFrame> = dfs.into_iter().map(|df| df.lazy()).collect();
868 let combined = polars::prelude::concat(lazy_frames, Default::default())?.collect()?;
869 Ok(combined)
870}
871
872fn apply_ts_filter(
874 df: DataFrame,
875 from: Option<NaiveDateTime>,
876 to: Option<NaiveDateTime>,
877) -> Result<DataFrame> {
878 if from.is_none() && to.is_none() {
879 return Ok(df);
880 }
881
882 let mut lf = df.lazy();
883
884 if let Some(f) = from {
885 let from_micros = f.and_utc().timestamp_micros();
886 lf = lf.filter(
887 col("ts")
888 .gt_eq(lit(from_micros).cast(DataType::Datetime(TimeUnit::Microseconds, None))),
889 );
890 }
891 if let Some(t) = to {
892 let to_micros = t.and_utc().timestamp_micros();
893 lf = lf.filter(
894 col("ts").lt_eq(lit(to_micros).cast(DataType::Datetime(TimeUnit::Microseconds, None))),
895 );
896 }
897
898 Ok(lf.collect()?)
899}
900
901fn apply_pagination(
903 df: DataFrame,
904 limit: usize,
905 tail: bool,
906 descending: bool,
907) -> Result<DataFrame> {
908 let result = if tail && limit > 0 {
910 let n = limit.min(df.height());
912 let tailed = df.tail(Some(n));
913 if descending {
914 tailed.sort(
915 ["ts"],
916 SortMultipleOptions::default().with_order_descending(true),
917 )?
918 } else {
919 tailed
920 }
921 } else if descending {
922 let sorted = df.sort(
923 ["ts"],
924 SortMultipleOptions::default().with_order_descending(true),
925 )?;
926 if limit > 0 {
927 sorted.head(Some(limit))
928 } else {
929 sorted
930 }
931 } else if limit > 0 {
932 df.head(Some(limit))
933 } else {
934 df
935 };
936 Ok(result)
937}
938
939fn delete_from_partition(
941 dir: &Path,
942 from: Option<NaiveDateTime>,
943 to: Option<NaiveDateTime>,
944) -> Result<usize> {
945 if from.is_none() && to.is_none() {
946 let count = count_all_rows_in_dir(dir);
948 for entry in fs::read_dir(dir)?.flatten() {
950 let path = entry.path();
951 if path.extension().is_some_and(|e| e == "parquet") {
952 fs::remove_file(&path)?;
953 }
954 }
955 return Ok(count);
956 }
957
958 let files = list_date_files(dir, from, to)?;
959 let mut total_deleted = 0usize;
960
961 for file_path in &files {
962 let df = read_parquet_file(file_path)?;
963 let original_count = df.height();
964
965 let filtered = apply_ts_filter_inverted(df, from, to)?;
967
968 if filtered.height() == 0 {
969 fs::remove_file(file_path)?;
971 total_deleted += original_count;
972 } else if filtered.height() < original_count {
973 total_deleted += original_count - filtered.height();
975 write_parquet_file(file_path, &mut filtered.clone())?;
976 }
977 }
979
980 Ok(total_deleted)
981}
982
983fn apply_ts_filter_inverted(
985 df: DataFrame,
986 from: Option<NaiveDateTime>,
987 to: Option<NaiveDateTime>,
988) -> Result<DataFrame> {
989 let mut lf = df.lazy();
990
991 match (from, to) {
992 (Some(f), Some(t)) => {
993 let from_micros = f.and_utc().timestamp_micros();
994 let to_micros = t.and_utc().timestamp_micros();
995 let from_lit = lit(from_micros).cast(DataType::Datetime(TimeUnit::Microseconds, None));
996 let to_lit = lit(to_micros).cast(DataType::Datetime(TimeUnit::Microseconds, None));
997 lf = lf.filter(col("ts").lt(from_lit).or(col("ts").gt(to_lit)));
999 }
1000 (Some(f), None) => {
1001 let from_micros = f.and_utc().timestamp_micros();
1002 let from_lit = lit(from_micros).cast(DataType::Datetime(TimeUnit::Microseconds, None));
1003 lf = lf.filter(col("ts").lt(from_lit));
1004 }
1005 (None, Some(t)) => {
1006 let to_micros = t.and_utc().timestamp_micros();
1007 let to_lit = lit(to_micros).cast(DataType::Datetime(TimeUnit::Microseconds, None));
1008 lf = lf.filter(col("ts").gt(to_lit));
1009 }
1010 (None, None) => {}
1011 }
1012
1013 Ok(lf.collect()?)
1014}
1015
1016fn count_all_rows_in_dir(dir: &Path) -> usize {
1018 let mut count = 0;
1019 if let Ok(entries) = fs::read_dir(dir) {
1020 for entry in entries.flatten() {
1021 let path = entry.path();
1022 if path.extension().is_some_and(|e| e == "parquet")
1023 && let Ok(df) = read_parquet_file(&path)
1024 {
1025 count += df.height();
1026 }
1027 }
1028 }
1029 count
1030}
1031
1032fn walkdir(dir: &Path) -> Vec<PathBuf> {
1034 let mut result = Vec::new();
1035 if !dir.exists() {
1036 return result;
1037 }
1038 if let Ok(entries) = fs::read_dir(dir) {
1039 for entry in entries.flatten() {
1040 let path = entry.path();
1041 if path.is_dir() {
1042 result.extend(walkdir(&path));
1043 } else {
1044 result.push(path);
1045 }
1046 }
1047 }
1048 result
1049}
1050
1051fn micros_to_ndt(micros: i64) -> NaiveDateTime {
1053 let secs = micros / 1_000_000;
1054 let nsecs = ((micros % 1_000_000) * 1_000) as u32;
1055 chrono::DateTime::from_timestamp(secs, nsecs)
1056 .map(|dt| dt.naive_utc())
1057 .unwrap_or_default()
1058}