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

yuzu_data/
combined.rs
