Skip to main content

data_preprocess/
parquet_store.rs

1//! Parquet-based storage backend for tick and bar data.
2//!
3//! Uses Hive-style directory partitioning:
4//!   {root}/ticks/exchange={ex}/symbol={sym}/{date}.parquet
5//!   {root}/bars/exchange={ex}/symbol={sym}/timeframe={tf}/{date}.parquet
6
7use 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
24/// Parquet-based storage backend for tick and bar data.
25pub struct ParquetStore {
26    root: PathBuf,
27}
28
29impl ParquetStore {
30    /// Open a Parquet data store rooted at the given directory, creating it if needed.
31    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    // ── Import ──────────────────────────────────────────────────
38
39    /// Import ticks, deduplicating against existing data per date partition.
40    /// Returns the number of rows actually inserted (after dedup).
41    pub fn insert_ticks(&self, ticks: &[Tick]) -> Result<usize> {
42        if ticks.is_empty() {
43            return Ok(0);
44        }
45
46        // Group ticks by (exchange, symbol, date)
47        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    /// Import bars, deduplicating against existing data per date partition.
81    /// Returns the number of rows actually inserted (after dedup).
82    pub fn insert_bars(&self, bars: &[Bar]) -> Result<usize> {
83        if bars.is_empty() {
84            return Ok(0);
85        }
86
87        // Group bars by (exchange, symbol, timeframe, date)
88        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    // ── Query ───────────────────────────────────────────────────
127
128    /// Query ticks for a given exchange+symbol, optionally filtered by date range.
129    /// Returns (ticks, total_count_matching_filters).
130    pub fn query_ticks(&self, opts: &QueryOpts) -> Result<(Vec<Tick>, u64)> {
131        self.query_ticks_cancellable(opts, || false)
132    }
133
134    /// Cancellable tick query.
135    ///
136    /// Cancellation is checked while traversing directory entries, before and
137    /// after every Parquet file, and between DataFrame processing stages. The
138    /// low-level Polars collect/read for one file remains atomic because Polars
139    /// does not expose an interruption hook for that operation.
140    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    /// Return the latest tick with a valid quote strictly before `before`.
184    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    /// Cancellable strict-before lookup for the latest tick with a valid quote.
194    ///
195    /// Date partitions and their rows are searched newest-to-oldest.
196    /// Ticks at `before` are excluded, and ticks with missing, non-finite, non-positive, or crossed bid/ask prices are skipped.
197    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    /// Query bars for a given exchange+symbol+timeframe, optionally filtered by date range.
223    /// Returns (bars, total_count_matching_filters).
224    pub fn query_bars(&self, opts: &BarQueryOpts) -> Result<(Vec<Bar>, u64)> {
225        self.query_bars_cancellable(opts, || false)
226    }
227
228    /// Cancellable bar query with the same cooperative boundaries as
229    /// [`Self::query_ticks_cancellable`].
230    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    // ── Delete ──────────────────────────────────────────────────
274
275    /// Delete ticks matching exchange+symbol, optionally within a date range.
276    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    /// Delete bars matching exchange+symbol+timeframe, optionally within a date range.
291    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    /// Delete ALL data (ticks + bars) for an exchange+symbol pair.
307    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        // Remove tick directory
312        let tick_dir = self.tick_dir(exchange, symbol);
313        if tick_dir.exists() {
314            fs::remove_dir_all(&tick_dir)?;
315        }
316
317        // Remove bar directories for all timeframes
318        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    /// Delete ALL data for an entire exchange.
331    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    // ── Stats ───────────────────────────────────────────────────
349
350    /// Summary statistics across all data, optionally filtered by exchange and/or symbol.
351    pub fn stats(&self, exchange: Option<&str>, symbol: Option<&str>) -> Result<Vec<StatRow>> {
352        let mut rows = Vec::new();
353
354        // Collect tick stats
355        self.collect_tick_stats(&mut rows, exchange, symbol)?;
356
357        // Collect bar stats
358        self.collect_bar_stats(&mut rows, exchange, symbol)?;
359
360        // Sort by exchange, symbol, data_type
361        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    /// Total size of all Parquet files under the data root (bytes).
372    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    // ── Private helpers ─────────────────────────────────────────
385
386    pub(crate) fn root_path(&self) -> &Path {
387        &self.root
388    }
389
390    /// Build tick directory path for a given exchange+symbol.
391    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    /// Build bar directory path for a given exchange+symbol+timeframe.
399    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    /// Count total rows across all parquet files in a directory.
408    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    /// Count total rows recursively across all parquet files under a directory.
427    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    /// Count all bar rows for a given exchange+symbol across all timeframes.
443    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    /// Collect tick stats from the directory tree.
453    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    /// Collect bar stats from the directory tree.
493    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    /// Iterate over exchange/symbol directories under a top-level dir.
533    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    /// Iterate over exchange/symbol/timeframe directories under a top-level dir.
568    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    /// Read all parquet files in a directory and aggregate row count + min/max ts.
618    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
659// ── Free functions ──────────────────────────────────────────────
660
661/// Parse a Hive partition value from a directory name like "exchange=ctrader".
662fn 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
677/// List parquet files in a directory, optionally filtered by date range in filename.
678fn 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            // Filename-level date pruning
704            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
718/// Read a single Parquet file into a DataFrame.
719fn 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
727/// Write a DataFrame to a temporary file and atomically replace the partition.
728fn 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
824/// Concat two tick DataFrames, dedup on (exchange, symbol, ts), sort by ts.
825fn concat_and_dedup_ticks(existing: DataFrame, new: DataFrame) -> Result<DataFrame> {
826    let combined = concat_dataframes(vec![existing, new])?;
827    dedup_ticks(combined)
828}
829
830/// Dedup a tick DataFrame on (exchange, symbol, ts) keeping first, sort by ts.
831fn 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
839/// Concat two bar DataFrames, dedup on (exchange, symbol, timeframe, ts), sort by ts.
840fn concat_and_dedup_bars(existing: DataFrame, new: DataFrame) -> Result<DataFrame> {
841    let combined = concat_dataframes(vec![existing, new])?;
842    dedup_bars(combined)
843}
844
845/// Dedup a bar DataFrame on (exchange, symbol, timeframe, ts) keeping first, sort by ts.
846fn 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
859/// Vertically concatenate multiple DataFrames.
860fn 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
872/// Apply timestamp range filter to a DataFrame with a "ts" datetime column.
873fn 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
901/// Apply limit, tail, and descending pagination to a sorted DataFrame.
902fn apply_pagination(
903    df: DataFrame,
904    limit: usize,
905    tail: bool,
906    descending: bool,
907) -> Result<DataFrame> {
908    // limit == 0 means "no limit" — return all rows.
909    let result = if tail && limit > 0 {
910        // Take last N rows, then optionally reverse for descending
911        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
939/// Delete rows from a date-partitioned directory, optionally within a date range.
940fn 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        // Delete everything in the directory
947        let count = count_all_rows_in_dir(dir);
948        // Remove all parquet files but keep the directory
949        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        // Filter to keep rows OUTSIDE the delete range
966        let filtered = apply_ts_filter_inverted(df, from, to)?;
967
968        if filtered.height() == 0 {
969            // All rows deleted — remove the file
970            fs::remove_file(file_path)?;
971            total_deleted += original_count;
972        } else if filtered.height() < original_count {
973            // Partial deletion — rewrite the file
974            total_deleted += original_count - filtered.height();
975            write_parquet_file(file_path, &mut filtered.clone())?;
976        }
977        // else: no rows matched the range in this file
978    }
979
980    Ok(total_deleted)
981}
982
983/// Filter to keep rows OUTSIDE a timestamp range (inverse of apply_ts_filter).
984fn 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            // Keep rows where ts < from OR ts > to
998            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
1016/// Count all rows across parquet files in a directory (non-recursive).
1017fn 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
1032/// Recursively walk a directory and collect all file paths.
1033fn 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
1051/// Convert microsecond epoch to NaiveDateTime.
1052fn 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}