yuzu_data/
combined.rs

1//! Combined per-field panel files: one wide gzip CSV per series
2//! (`panels/{name}.csv.gz`, header `day,SYM1,SYM2,…`, one row per trading day,
3//! empty cell = NaN), holding the full universe. The container reads ONE object
4//! per series instead of hundreds of per-symbol files. Same gzip-CSV conventions
5//! as `csv_io`/`fundamentals`.
6
7use crate::csv_io::{parse_series, Field};
8use crate::date::{date_to_i32, i32_to_date};
9use crate::error::DataError;
10use crate::format::CANDIDATE_EXTS;
11use crate::fundamentals::{parse_fundamentals, FUNDAMENTAL_FIELDS, REPORT_EVENT_FIELD};
12use crate::source::{ObjectSink, ObjectSource};
13use flate2::write::GzEncoder;
14use flate2::Compression;
15use ndarray::Array2;
16use std::collections::{BTreeSet, HashMap};
17use std::io::Write;
18use yuzu_core::panel::Panel;
19
20/// Default object-key directory for combined per-field panel files.
21pub const PANELS_DIR: &str = "panels";
22
23/// Serialize a Panel to combined gzip CSV (`day,<sym…>`, empty cell for NaN).
24pub fn write_combined_panel(panel: &Panel) -> Result<Vec<u8>, DataError> {
25    let mut buf = String::from("day");
26    for s in &panel.symbols {
27        buf.push(',');
28        buf.push_str(s);
29    }
30    buf.push('\n');
31    for (r, day) in panel.dates.iter().enumerate() {
32        buf.push_str(&i32_to_date(*day));
33        for c in 0..panel.symbols.len() {
34            buf.push(',');
35            let v = panel.data[[r, c]];
36            if !v.is_nan() {
37                buf.push_str(&v.to_string());
38            }
39        }
40        buf.push('\n');
41    }
42    let mut enc = GzEncoder::new(Vec::new(), Compression::default());
43    enc.write_all(buf.as_bytes())
44        .map_err(|e| DataError::Io(e.to_string()))?;
45    enc.finish().map_err(|e| DataError::Io(e.to_string()))
46}
47
48/// Read the combined `{dir}/{name}` file (probing `.csv.gz`/`.parquet`/`.csv`)
49/// and assemble a Panel for `symbols` (in the given order; a symbol absent from
50/// the file gets a NaN column) over `[from, to]` inclusive. The format is
51/// detected from content. `Ok(None)` if the combined file does not exist in any
52/// supported format.
53pub fn load_combined_panel<S: ObjectSource>(
54    source: &S,
55    name: &str,
56    symbols: &[String],
57    from: i32,
58    to: i32,
59    dir: &str,
60) -> Result<Option<Panel>, DataError> {
61    let mut bytes = None;
62    for ext in CANDIDATE_EXTS {
63        if let Some(b) = source.get(&format!("{dir}/{name}{ext}"))? {
64            bytes = Some(b);
65            break;
66        }
67    }
68    let Some(bytes) = bytes else {
69        return Ok(None);
70    };
71
72    #[cfg(feature = "parquet")]
73    if crate::format::Format::detect(&bytes) == crate::format::Format::Parquet {
74        return load_combined_parquet(&bytes, symbols, from, to).map(Some);
75    }
76
77    let text = crate::format::read_csv_text(&bytes)?;
78    let mut lines = text.lines();
79    let header = lines.next().unwrap_or("");
80    // file column index (>=1; col 0 is "day") for every symbol in the file
81    let file_cols: HashMap<&str, usize> = header
82        .split(',')
83        .enumerate()
84        .skip(1)
85        .map(|(i, s)| (s.trim(), i))
86        .collect();
87    // requested symbol → file column (None ⇒ NaN column)
88    let col_of: Vec<Option<usize>> = symbols
89        .iter()
90        .map(|s| file_cols.get(s.as_str()).copied())
91        .collect();
92
93    let mut dates: Vec<i32> = Vec::new();
94    let mut rows: Vec<Vec<f64>> = Vec::new();
95    for line in lines {
96        let line = line.trim();
97        if line.is_empty() {
98            continue;
99        }
100        let cells: Vec<&str> = line.split(',').collect();
101        let day = date_to_i32(cells[0].trim())?;
102        if day < from || day > to {
103            continue;
104        }
105        let mut row = Vec::with_capacity(symbols.len());
106        for col in &col_of {
107            let v = match col {
108                Some(c) if *c < cells.len() => {
109                    let cell = cells[*c].trim();
110                    if cell.is_empty() {
111                        f64::NAN
112                    } else {
113                        cell.parse()
114                            .map_err(|_| DataError::Parse(format!("bad value '{cell}'")))?
115                    }
116                }
117                _ => f64::NAN,
118            };
119            row.push(v);
120        }
121        dates.push(day);
122        rows.push(row);
123    }
124    let panel = Panel::from_rows(dates, symbols.to_vec(), rows)
125        .map_err(|e| DataError::Parse(e.to_string()))?;
126    Ok(Some(panel))
127}
128
129/// Assemble a Panel from a wide combined-panel Parquet buffer: one column per
130/// symbol (absent symbol ⇒ NaN column), windowed to `[from, to]` inclusive.
131#[cfg(feature = "parquet")]
132fn load_combined_parquet(
133    bytes: &[u8],
134    symbols: &[String],
135    from: i32,
136    to: i32,
137) -> Result<Panel, DataError> {
138    let (all_dates, all_rows) = crate::parquet_io::read_wide(bytes, symbols)?;
139    let mut dates = Vec::new();
140    let mut rows = Vec::new();
141    for (day, row) in all_dates.into_iter().zip(all_rows) {
142        if day >= from && day <= to {
143            dates.push(day);
144            rows.push(row);
145        }
146    }
147    Panel::from_rows(dates, symbols.to_vec(), rows).map_err(|e| DataError::Parse(e.to_string()))
148}
149
150/// What a rebuild wrote: number of series files and the max day-count across them.
151pub struct RebuildSummary {
152    pub fields: usize,
153    pub days: usize,
154}
155
156/// Build `panels/{series}.csv.gz` for every OHLCV field + fundamental field from
157/// the per-symbol archives. Reads each per-symbol file once (concurrently), then
158/// transposes per field. Full overwrite, idempotent. A missing/corrupt per-symbol
159/// file leaves NaN cells (same fail-soft as the per-symbol loaders).
160pub fn rebuild_combined_panels<S: ObjectSource + ObjectSink + Sync>(
161    source: &S,
162    symbols: &[String],
163    prices_dir: &str,
164    fundamentals_dir: &str,
165    panels_dir: &str,
166) -> Result<RebuildSummary, DataError> {
167    let mut fields = 0usize;
168    let mut max_days = 0usize;
169
170    // --- OHLCV: read every per-symbol price file once, transpose per field.
171    // Probes .csv.gz/.parquet/.csv so a mirror in any supported format rebuilds. ---
172    let price_bytes = crate::parallel::fetch_symbols(source, prices_dir, symbols)?;
173    let price_series: &[(&str, Field)] = &[
174        ("open", Field::AdjOpen),
175        ("high", Field::AdjHigh),
176        ("low", Field::AdjLow),
177        ("close", Field::AdjClose),
178        ("volume", Field::Volume),
179    ];
180    for (name, field) in price_series {
181        let per_symbol: Vec<Vec<(i32, f64)>> = price_bytes
182            .iter()
183            .map(|b| match b {
184                Some(bytes) => parse_series(bytes, *field).unwrap_or_default(),
185                None => Vec::new(),
186            })
187            .collect();
188        let panel = assemble(symbols, &per_symbol)?;
189        max_days = max_days.max(panel.dates.len());
190        source.put(
191            &format!("{panels_dir}/{name}.csv.gz"),
192            &write_combined_panel(&panel)?,
193        )?;
194        fields += 1;
195    }
196
197    drop(price_bytes); // free the price bytes before holding the fundamentals bytes (halves peak RAM)
198
199    // --- Fundamentals: read every per-symbol fundamentals file once, transpose ---
200    let fund_bytes = crate::parallel::fetch_symbols(source, fundamentals_dir, symbols)?;
201    for name in FUNDAMENTAL_FIELDS
202        .iter()
203        .chain(std::iter::once(&REPORT_EVENT_FIELD))
204    {
205        let per_symbol: Vec<Vec<(i32, f64)>> = fund_bytes
206            .iter()
207            .map(|b| match b {
208                Some(bytes) => parse_fundamentals(bytes, name).unwrap_or_default(),
209                None => Vec::new(),
210            })
211            .collect();
212        let panel = assemble(symbols, &per_symbol)?;
213        max_days = max_days.max(panel.dates.len());
214        source.put(
215            &format!("{panels_dir}/{name}.csv.gz"),
216            &write_combined_panel(&panel)?,
217        )?;
218        fields += 1;
219    }
220
221    Ok(RebuildSummary {
222        fields,
223        days: max_days,
224    })
225}
226
227/// Union-dates assembly: rows = sorted union of all symbols' days, cols = symbols
228/// in order, missing cells NaN. (Mirrors the per-symbol loaders' assembly.)
229fn assemble(symbols: &[String], per_symbol: &[Vec<(i32, f64)>]) -> Result<Panel, DataError> {
230    let mut date_set: BTreeSet<i32> = BTreeSet::new();
231    for rows in per_symbol {
232        for (d, _) in rows {
233            date_set.insert(*d);
234        }
235    }
236    let dates: Vec<i32> = date_set.into_iter().collect();
237    let row_of: HashMap<i32, usize> = dates.iter().enumerate().map(|(i, d)| (*d, i)).collect();
238    let mut data = Array2::from_elem((dates.len(), symbols.len()), f64::NAN);
239    for (c, rows) in per_symbol.iter().enumerate() {
240        for (d, v) in rows {
241            data[[row_of[d], c]] = *v;
242        }
243    }
244    Panel::new(dates, symbols.to_vec(), data).map_err(|e| DataError::Parse(e.to_string()))
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use crate::source::LocalSource;
251    use ndarray::array;
252    use std::fs;
253    use yuzu_core::panel::Panel;
254
255    #[test]
256    fn rebuild_writes_combined_files_loadable_with_same_values() {
257        use crate::csv_io::{write_series, OhlcvRow};
258        use crate::fundamentals::{write_fundamentals, FundamentalRow, FUNDAMENTAL_FIELDS};
259        use crate::source::LocalSource;
260        use std::fs;
261
262        let dir = std::env::temp_dir().join("yuzu_combined_rebuild");
263        let _ = fs::remove_dir_all(&dir);
264        for d in ["prices", "fundamentals", "panels"] {
265            fs::create_dir_all(dir.join(d)).unwrap();
266        }
267        let ohlcv = |day, c| OhlcvRow {
268            day,
269            adj_open: c,
270            adj_high: c,
271            adj_low: c,
272            adj_close: c,
273            volume: 100.0,
274        };
275        fs::write(
276            dir.join("prices/AAA.csv.gz"),
277            write_series(&[ohlcv(20240102, 10.0), ohlcv(20240103, 11.0)]).unwrap(),
278        )
279        .unwrap();
280        fs::write(
281            dir.join("prices/BBB.csv.gz"),
282            write_series(&[ohlcv(20240103, 20.0)]).unwrap(),
283        )
284        .unwrap();
285        let frow = |day, pe| {
286            let mut values = vec![f64::NAN; FUNDAMENTAL_FIELDS.len()];
287            values[0] = pe; // "pe" is column 0
288            FundamentalRow {
289                day,
290                values,
291                report_event: 0.0,
292            }
293        };
294        fs::write(
295            dir.join("fundamentals/AAA.csv.gz"),
296            write_fundamentals(&[frow(20240102, 8.0), frow(20240103, 8.0)]).unwrap(),
297        )
298        .unwrap();
299        fs::write(
300            dir.join("fundamentals/BBB.csv.gz"),
301            write_fundamentals(&[frow(20240103, 15.0)]).unwrap(),
302        )
303        .unwrap();
304
305        let src = LocalSource::new(&dir);
306        let syms = vec!["AAA".to_string(), "BBB".to_string(), "CCC".to_string()];
307        let summary =
308            rebuild_combined_panels(&src, &syms, "prices", "fundamentals", "panels").unwrap();
309        assert_eq!(summary.fields, 5 + FUNDAMENTAL_FIELDS.len() + 1); // OHLCV + factors + report_event
310        assert_eq!(summary.days, 2);
311
312        let close = load_combined_panel(&src, "close", &syms, 20240102, 20240103, "panels")
313            .unwrap()
314            .unwrap();
315        assert_eq!(close.dates, vec![20240102, 20240103]);
316        assert_eq!(close.data[[0, 0]], 10.0); // AAA 0102
317        assert!(close.data[[0, 1]].is_nan()); // BBB absent 0102
318        assert!(close.data[[0, 2]].is_nan()); // CCC absent (no file)
319        assert!(close.data[[1, 2]].is_nan()); // CCC absent (no file)
320        assert_eq!(close.data[[1, 1]], 20.0); // BBB 0103
321        let pe = load_combined_panel(&src, "pe", &syms, 20240102, 20240103, "panels")
322            .unwrap()
323            .unwrap();
324        assert_eq!(pe.data[[0, 0]], 8.0); // AAA pe
325    }
326
327    #[cfg(feature = "parquet")]
328    #[test]
329    fn loads_a_wide_parquet_combined_panel() {
330        use arrow_array::{ArrayRef, Float64Array, Int32Array, RecordBatch};
331        use arrow_schema::{DataType, Field as AField, Schema};
332        use parquet::arrow::ArrowWriter;
333        use std::sync::Arc;
334
335        let schema = Arc::new(Schema::new(vec![
336            AField::new("day", DataType::Int32, false),
337            AField::new("AAA", DataType::Float64, true),
338            AField::new("BBB", DataType::Float64, true),
339        ]));
340        let batch = RecordBatch::try_new(
341            schema.clone(),
342            vec![
343                Arc::new(Int32Array::from(vec![20240102, 20240103, 20240104])) as ArrayRef,
344                Arc::new(Float64Array::from(vec![Some(10.0), Some(11.0), Some(12.0)])) as ArrayRef,
345                Arc::new(Float64Array::from(vec![None, Some(20.0), Some(21.0)])) as ArrayRef,
346            ],
347        )
348        .unwrap();
349        let mut buf = Vec::new();
350        let mut w = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
351        w.write(&batch).unwrap();
352        w.close().unwrap();
353
354        let dir = std::env::temp_dir().join("yuzu_combined_parquet");
355        let _ = fs::remove_dir_all(&dir);
356        fs::create_dir_all(dir.join("panels")).unwrap();
357        fs::write(dir.join("panels/close.parquet"), buf).unwrap();
358        let src = LocalSource::new(&dir);
359
360        // reorder + windowing + absent symbol, same contract as the CSV path
361        let syms = vec!["BBB".to_string(), "AAA".to_string(), "ZZZ".to_string()];
362        let p = load_combined_panel(&src, "close", &syms, 20240103, 20240104, PANELS_DIR)
363            .unwrap()
364            .unwrap();
365        assert_eq!(p.dates, vec![20240103, 20240104]);
366        assert_eq!(p.symbols, syms);
367        assert_eq!(p.data[[0, 0]], 20.0); // BBB 0103
368        assert_eq!(p.data[[0, 1]], 11.0); // AAA 0103
369        assert!(p.data[[0, 2]].is_nan()); // ZZZ absent column
370    }
371
372    #[test]
373    fn write_then_load_selects_subset_reorders_and_windows() {
374        let data = array![
375            [10.0, 20.0, f64::NAN],
376            [11.0, 21.0, 31.0],
377            [12.0, 22.0, 32.0],
378        ];
379        let panel = Panel::new(
380            vec![20240102, 20240103, 20240104],
381            vec!["AAA".into(), "BBB".into(), "CCC".into()],
382            data,
383        )
384        .unwrap();
385        let bytes = write_combined_panel(&panel).unwrap();
386
387        let dir = std::env::temp_dir().join("yuzu_combined_rw");
388        let _ = fs::remove_dir_all(&dir);
389        fs::create_dir_all(dir.join("panels")).unwrap();
390        fs::write(dir.join("panels/close.csv.gz"), bytes).unwrap();
391        let src = LocalSource::new(&dir);
392
393        // request CCC, AAA (reordered) + ZZZ (absent) over the inner window
394        let syms = vec!["CCC".to_string(), "AAA".to_string(), "ZZZ".to_string()];
395        let p = load_combined_panel(&src, "close", &syms, 20240103, 20240104, PANELS_DIR)
396            .unwrap()
397            .unwrap();
398        assert_eq!(p.dates, vec![20240103, 20240104]);
399        assert_eq!(p.symbols, syms);
400        assert_eq!(p.data[[0, 0]], 31.0); // CCC 0103
401        assert_eq!(p.data[[0, 1]], 11.0); // AAA 0103
402        assert!(p.data[[0, 2]].is_nan()); // ZZZ absent from file
403                                          // absent combined file → Ok(None) (caller falls back)
404        assert!(
405            load_combined_panel(&src, "missing", &syms, 0, 99999999, PANELS_DIR)
406                .unwrap()
407                .is_none()
408        );
409    }
410}
yuzu_data/combined.rs

yuzu_data/
combined.rs
