pomelo_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.)
229/// Public so producers (e.g. the CLI's FMP snapshot-factor writer) can build a
230/// combined [`Panel`] from per-symbol `(day, value)` rows before serializing it
231/// with [`write_combined_panel`].
232pub fn assemble(symbols: &[String], per_symbol: &[Vec<(i32, f64)>]) -> Result<Panel, DataError> {
233    let mut date_set: BTreeSet<i32> = BTreeSet::new();
234    for rows in per_symbol {
235        for (d, _) in rows {
236            date_set.insert(*d);
237        }
238    }
239    let dates: Vec<i32> = date_set.into_iter().collect();
240    let row_of: HashMap<i32, usize> = dates.iter().enumerate().map(|(i, d)| (*d, i)).collect();
241    let mut data = Array2::from_elem((dates.len(), symbols.len()), f64::NAN);
242    for (c, rows) in per_symbol.iter().enumerate() {
243        for (d, v) in rows {
244            data[[row_of[d], c]] = *v;
245        }
246    }
247    Panel::new(dates, symbols.to_vec(), data).map_err(|e| DataError::Parse(e.to_string()))
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use crate::source::LocalSource;
254    use ndarray::array;
255    use std::fs;
256    use yuzu_core::panel::Panel;
257
258    #[test]
259    fn rebuild_writes_combined_files_loadable_with_same_values() {
260        use crate::csv_io::{write_series, OhlcvRow};
261        use crate::fundamentals::{write_fundamentals, FundamentalRow, FUNDAMENTAL_FIELDS};
262        use crate::source::LocalSource;
263        use std::fs;
264
265        let dir = std::env::temp_dir().join("yuzu_combined_rebuild");
266        let _ = fs::remove_dir_all(&dir);
267        for d in ["prices", "fundamentals", "panels"] {
268            fs::create_dir_all(dir.join(d)).unwrap();
269        }
270        let ohlcv = |day, c| OhlcvRow {
271            day,
272            adj_open: c,
273            adj_high: c,
274            adj_low: c,
275            adj_close: c,
276            volume: 100.0,
277        };
278        fs::write(
279            dir.join("prices/AAA.csv.gz"),
280            write_series(&[ohlcv(20240102, 10.0), ohlcv(20240103, 11.0)]).unwrap(),
281        )
282        .unwrap();
283        fs::write(
284            dir.join("prices/BBB.csv.gz"),
285            write_series(&[ohlcv(20240103, 20.0)]).unwrap(),
286        )
287        .unwrap();
288        let frow = |day, pe| {
289            let mut values = vec![f64::NAN; FUNDAMENTAL_FIELDS.len()];
290            values[0] = pe; // "pe" is column 0
291            FundamentalRow {
292                day,
293                values,
294                report_event: 0.0,
295            }
296        };
297        fs::write(
298            dir.join("fundamentals/AAA.csv.gz"),
299            write_fundamentals(&[frow(20240102, 8.0), frow(20240103, 8.0)]).unwrap(),
300        )
301        .unwrap();
302        fs::write(
303            dir.join("fundamentals/BBB.csv.gz"),
304            write_fundamentals(&[frow(20240103, 15.0)]).unwrap(),
305        )
306        .unwrap();
307
308        let src = LocalSource::new(&dir);
309        let syms = vec!["AAA".to_string(), "BBB".to_string(), "CCC".to_string()];
310        let summary =
311            rebuild_combined_panels(&src, &syms, "prices", "fundamentals", "panels").unwrap();
312        assert_eq!(summary.fields, 5 + FUNDAMENTAL_FIELDS.len() + 1); // OHLCV + factors + report_event
313        assert_eq!(summary.days, 2);
314
315        let close = load_combined_panel(&src, "close", &syms, 20240102, 20240103, "panels")
316            .unwrap()
317            .unwrap();
318        assert_eq!(close.dates, vec![20240102, 20240103]);
319        assert_eq!(close.data[[0, 0]], 10.0); // AAA 0102
320        assert!(close.data[[0, 1]].is_nan()); // BBB absent 0102
321        assert!(close.data[[0, 2]].is_nan()); // CCC absent (no file)
322        assert!(close.data[[1, 2]].is_nan()); // CCC absent (no file)
323        assert_eq!(close.data[[1, 1]], 20.0); // BBB 0103
324        let pe = load_combined_panel(&src, "pe", &syms, 20240102, 20240103, "panels")
325            .unwrap()
326            .unwrap();
327        assert_eq!(pe.data[[0, 0]], 8.0); // AAA pe
328    }
329
330    #[cfg(feature = "parquet")]
331    #[test]
332    fn loads_a_wide_parquet_combined_panel() {
333        use arrow_array::{ArrayRef, Float64Array, Int32Array, RecordBatch};
334        use arrow_schema::{DataType, Field as AField, Schema};
335        use parquet::arrow::ArrowWriter;
336        use std::sync::Arc;
337
338        let schema = Arc::new(Schema::new(vec![
339            AField::new("day", DataType::Int32, false),
340            AField::new("AAA", DataType::Float64, true),
341            AField::new("BBB", DataType::Float64, true),
342        ]));
343        let batch = RecordBatch::try_new(
344            schema.clone(),
345            vec![
346                Arc::new(Int32Array::from(vec![20240102, 20240103, 20240104])) as ArrayRef,
347                Arc::new(Float64Array::from(vec![Some(10.0), Some(11.0), Some(12.0)])) as ArrayRef,
348                Arc::new(Float64Array::from(vec![None, Some(20.0), Some(21.0)])) as ArrayRef,
349            ],
350        )
351        .unwrap();
352        let mut buf = Vec::new();
353        let mut w = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
354        w.write(&batch).unwrap();
355        w.close().unwrap();
356
357        let dir = std::env::temp_dir().join("yuzu_combined_parquet");
358        let _ = fs::remove_dir_all(&dir);
359        fs::create_dir_all(dir.join("panels")).unwrap();
360        fs::write(dir.join("panels/close.parquet"), buf).unwrap();
361        let src = LocalSource::new(&dir);
362
363        // reorder + windowing + absent symbol, same contract as the CSV path
364        let syms = vec!["BBB".to_string(), "AAA".to_string(), "ZZZ".to_string()];
365        let p = load_combined_panel(&src, "close", &syms, 20240103, 20240104, PANELS_DIR)
366            .unwrap()
367            .unwrap();
368        assert_eq!(p.dates, vec![20240103, 20240104]);
369        assert_eq!(p.symbols, syms);
370        assert_eq!(p.data[[0, 0]], 20.0); // BBB 0103
371        assert_eq!(p.data[[0, 1]], 11.0); // AAA 0103
372        assert!(p.data[[0, 2]].is_nan()); // ZZZ absent column
373    }
374
375    #[test]
376    fn write_then_load_selects_subset_reorders_and_windows() {
377        let data = array![
378            [10.0, 20.0, f64::NAN],
379            [11.0, 21.0, 31.0],
380            [12.0, 22.0, 32.0],
381        ];
382        let panel = Panel::new(
383            vec![20240102, 20240103, 20240104],
384            vec!["AAA".into(), "BBB".into(), "CCC".into()],
385            data,
386        )
387        .unwrap();
388        let bytes = write_combined_panel(&panel).unwrap();
389
390        let dir = std::env::temp_dir().join("yuzu_combined_rw");
391        let _ = fs::remove_dir_all(&dir);
392        fs::create_dir_all(dir.join("panels")).unwrap();
393        fs::write(dir.join("panels/close.csv.gz"), bytes).unwrap();
394        let src = LocalSource::new(&dir);
395
396        // request CCC, AAA (reordered) + ZZZ (absent) over the inner window
397        let syms = vec!["CCC".to_string(), "AAA".to_string(), "ZZZ".to_string()];
398        let p = load_combined_panel(&src, "close", &syms, 20240103, 20240104, PANELS_DIR)
399            .unwrap()
400            .unwrap();
401        assert_eq!(p.dates, vec![20240103, 20240104]);
402        assert_eq!(p.symbols, syms);
403        assert_eq!(p.data[[0, 0]], 31.0); // CCC 0103
404        assert_eq!(p.data[[0, 1]], 11.0); // AAA 0103
405        assert!(p.data[[0, 2]].is_nan()); // ZZZ absent from file
406                                          // absent combined file → Ok(None) (caller falls back)
407        assert!(
408            load_combined_panel(&src, "missing", &syms, 0, 99999999, PANELS_DIR)
409                .unwrap()
410                .is_none()
411        );
412    }
413}
pomelo_data/combined.rs

pomelo_data/
combined.rs
