yuzu_data/
loader.rs

1use crate::csv_io::{parse_series, Field};
2use crate::error::DataError;
3use crate::source::ObjectSource;
4use ndarray::Array2;
5use std::collections::{BTreeSet, HashMap};
6use yuzu_core::panel::Panel;
7
8/// Default object-key directory for per-symbol price files. Override at the entry
9/// point (e.g. `YUZU_PRICES_DIR`) for a custom layout — the engine stays generic.
10pub const PRICES_DIR: &str = "prices";
11
12/// Read the per-symbol price file for each symbol (probing `.csv.gz`, then
13/// `.parquet` when that feature is on, then `.csv`; format detected from
14/// content), extract `field`, filter to `[from, to]` (inclusive, YYYYMMDD), and
15/// assemble a Panel: rows = sorted union of all kept days, columns = `symbols` in
16/// the given order. Missing cells (and symbols with no file) are NaN. `dir`
17/// defaults to [`PRICES_DIR`] at call sites.
18pub fn load_panel<S: ObjectSource + Sync>(
19    source: &S,
20    symbols: &[String],
21    field: Field,
22    from: i32,
23    to: i32,
24    dir: &str,
25) -> Result<Panel, DataError> {
26    // Fetch + parse every symbol concurrently (network-bound). A corrupt file is
27    // treated as missing (its column stays NaN) rather than sinking the batch.
28    let per_symbol =
29        crate::parallel::fetch_series(source, symbols, dir, from, to, |b| parse_series(b, field))?;
30
31    let mut date_set: BTreeSet<i32> = BTreeSet::new();
32    for map in &per_symbol {
33        date_set.extend(map.keys().copied());
34    }
35    let dates: Vec<i32> = date_set.into_iter().collect();
36    let row_of: HashMap<i32, usize> = dates.iter().enumerate().map(|(i, d)| (*d, i)).collect();
37    let mut data = Array2::from_elem((dates.len(), symbols.len()), f64::NAN);
38    for (c, map) in per_symbol.iter().enumerate() {
39        for (d, v) in map {
40            data[[row_of[d], c]] = *v;
41        }
42    }
43
44    Panel::new(dates, symbols.to_vec(), data).map_err(|e| DataError::Parse(e.to_string()))
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50    use crate::csv_io::{write_series, OhlcvRow};
51    use crate::source::LocalSource;
52    use std::fs;
53
54    fn r(day: i32, c: f64) -> OhlcvRow {
55        OhlcvRow {
56            day,
57            adj_open: c,
58            adj_high: c + 1.0,
59            adj_low: c - 1.0,
60            adj_close: c,
61            volume: 100.0,
62        }
63    }
64
65    fn fixture_dir(tag: &str) -> std::path::PathBuf {
66        // per-test dir (tests run in parallel) — a shared name races on remove/write.
67        let dir = std::env::temp_dir().join(format!("yuzu_data_loader_{tag}"));
68        // Clean up any existing files
69        let _ = fs::remove_dir_all(&dir);
70        fs::create_dir_all(dir.join("prices")).unwrap();
71        fs::write(
72            dir.join("prices/AAPL.csv.gz"),
73            write_series(&[r(20240102, 10.0), r(20240103, 11.0), r(20240104, 12.0)]).unwrap(),
74        )
75        .unwrap();
76        fs::write(
77            dir.join("prices/MSFT.csv.gz"),
78            write_series(&[r(20240103, 50.0), r(20240105, 52.0)]).unwrap(),
79        )
80        .unwrap();
81        dir
82    }
83
84    #[test]
85    fn loads_chosen_field_with_union_dates_and_nan() {
86        let dir = fixture_dir("union");
87        let src = LocalSource::new(&dir);
88        let syms = vec!["AAPL".to_string(), "MSFT".to_string(), "ZZZ".to_string()];
89
90        let close =
91            load_panel(&src, &syms, Field::AdjClose, 20240102, 20240104, PRICES_DIR).unwrap();
92        assert_eq!(close.dates, vec![20240102, 20240103, 20240104]);
93        assert_eq!(close.data[[0, 0]], 10.0); // AAPL close
94        assert!(close.data[[0, 1]].is_nan()); // MSFT absent on 0102
95        assert_eq!(close.data[[1, 1]], 50.0); // MSFT close on 0103
96        assert!(close.data[[2, 2]].is_nan()); // ZZZ no file
97
98        // a different field comes from the same files
99        let high = load_panel(&src, &syms, Field::AdjHigh, 20240102, 20240104, PRICES_DIR).unwrap();
100        assert_eq!(high.data[[0, 0]], 11.0); // AAPL high = close + 1
101    }
102
103    #[test]
104    fn loads_plain_csv_files_alongside_gzip() {
105        // A `.csv` (uncompressed) file must load via the same path (probed after
106        // `.csv.gz`), so a mixed mirror works.
107        let dir = std::env::temp_dir().join("yuzu_data_loader_plaincsv");
108        let _ = fs::remove_dir_all(&dir);
109        fs::create_dir_all(dir.join("prices")).unwrap();
110        fs::write(
111            dir.join("prices/AAPL.csv"),
112            "day,adj_open,adj_high,adj_low,adj_close,volume\n\
113             2024-01-02,9,11,9,10,100\n\
114             2024-01-03,10,12,10,11,100\n",
115        )
116        .unwrap();
117        let src = LocalSource::new(&dir);
118        let syms = vec!["AAPL".to_string()];
119        let close =
120            load_panel(&src, &syms, Field::AdjClose, 20240102, 20240103, PRICES_DIR).unwrap();
121        assert_eq!(close.dates, vec![20240102, 20240103]);
122        assert_eq!(close.data[[0, 0]], 10.0);
123        assert_eq!(close.data[[1, 0]], 11.0);
124    }
125
126    #[cfg(feature = "parquet")]
127    #[test]
128    fn loads_parquet_files_through_the_same_path() {
129        use arrow_array::{ArrayRef, Float64Array, Int32Array, RecordBatch};
130        use arrow_schema::{DataType, Field as AField, Schema};
131        use parquet::arrow::ArrowWriter;
132        use std::sync::Arc;
133
134        let dir = std::env::temp_dir().join("yuzu_data_loader_parquet");
135        let _ = fs::remove_dir_all(&dir);
136        fs::create_dir_all(dir.join("prices")).unwrap();
137
138        let schema = Arc::new(Schema::new(vec![
139            AField::new("day", DataType::Int32, false),
140            AField::new("adj_close", DataType::Float64, true),
141            AField::new("volume", DataType::Float64, true),
142        ]));
143        let batch = RecordBatch::try_new(
144            schema.clone(),
145            vec![
146                Arc::new(Int32Array::from(vec![20240102, 20240103])) as ArrayRef,
147                Arc::new(Float64Array::from(vec![10.0, 11.0])) as ArrayRef,
148                Arc::new(Float64Array::from(vec![100.0, 200.0])) as ArrayRef,
149            ],
150        )
151        .unwrap();
152        let mut buf = Vec::new();
153        let mut w = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
154        w.write(&batch).unwrap();
155        w.close().unwrap();
156        fs::write(dir.join("prices/NVDA.parquet"), buf).unwrap();
157
158        let src = LocalSource::new(&dir);
159        let syms = vec!["NVDA".to_string()];
160        let close =
161            load_panel(&src, &syms, Field::AdjClose, 20240102, 20240103, PRICES_DIR).unwrap();
162        assert_eq!(close.dates, vec![20240102, 20240103]);
163        assert_eq!(close.data[[0, 0]], 10.0);
164        assert_eq!(close.data[[1, 0]], 11.0);
165        // a different field comes from the same parquet file
166        let vol = load_panel(&src, &syms, Field::Volume, 20240102, 20240103, PRICES_DIR).unwrap();
167        assert_eq!(vol.data[[1, 0]], 200.0);
168    }
169
170    #[test]
171    fn skips_a_corrupt_file_instead_of_failing_the_batch() {
172        let dir = fixture_dir("corrupt");
173        // A truncated/non-gzip file (e.g. a half-finished R2 sync) must not sink the batch.
174        fs::write(dir.join("prices/BAD.csv.gz"), b"not gzip at all").unwrap();
175        let src = LocalSource::new(&dir);
176        let syms = vec!["AAPL".to_string(), "BAD".to_string()];
177
178        let close =
179            load_panel(&src, &syms, Field::AdjClose, 20240102, 20240104, PRICES_DIR).unwrap();
180        assert_eq!(close.data[[0, 0]], 10.0); // AAPL still loads
181        assert!(close.data[[0, 1]].is_nan()); // BAD column stays NaN — skipped, not an error
182        assert!(close.data[[2, 1]].is_nan());
183    }
184}
yuzu_data/loader.rs

yuzu_data/
loader.rs
