yuzu_data/
csv_io.rs

1//! Per-symbol gzip CSV I/O. Files hold full adjusted OHLCV
2//! (`day,adj_open,adj_high,adj_low,adj_close,volume`, oldest-first); `parse_series`
3//! extracts one chosen [`Field`] column into `(day, value)` rows.
4
5use crate::date::{date_to_i32, i32_to_date};
6use crate::error::DataError;
7use flate2::write::GzEncoder;
8use flate2::Compression;
9use std::io::Write;
10
11/// One row of adjusted OHLCV for a trading day.
12#[derive(Debug, Clone, PartialEq)]
13pub struct OhlcvRow {
14    pub day: i32,
15    pub adj_open: f64,
16    pub adj_high: f64,
17    pub adj_low: f64,
18    pub adj_close: f64,
19    pub volume: f64,
20}
21
22/// Which OHLCV column to pull into a Panel.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Field {
25    AdjOpen,
26    AdjHigh,
27    AdjLow,
28    AdjClose,
29    Volume,
30}
31
32impl Field {
33    /// Column index in the CSV (day is column 0).
34    fn col(self) -> usize {
35        match self {
36            Field::AdjOpen => 1,
37            Field::AdjHigh => 2,
38            Field::AdjLow => 3,
39            Field::AdjClose => 4,
40            Field::Volume => 5,
41        }
42    }
43
44    /// Column name (matches the CSV header and the Parquet column name).
45    pub fn name(self) -> &'static str {
46        match self {
47            Field::AdjOpen => "adj_open",
48            Field::AdjHigh => "adj_high",
49            Field::AdjLow => "adj_low",
50            Field::AdjClose => "adj_close",
51            Field::Volume => "volume",
52        }
53    }
54}
55
56/// Parse the OHLCV data for `field` into `(YYYYMMDD, value)` rows. The buffer's
57/// format is detected from its content: gzip CSV, plain CSV, or — with the
58/// `parquet` feature — Apache Parquet. CSV is header-optional, oldest-first.
59pub fn parse_series(bytes: &[u8], field: Field) -> Result<Vec<(i32, f64)>, DataError> {
60    #[cfg(feature = "parquet")]
61    if crate::format::Format::detect(bytes) == crate::format::Format::Parquet {
62        return crate::parquet_io::read_series(bytes, field.name());
63    }
64    let text = crate::format::read_csv_text(bytes)?;
65    parse_series_csv(&text, field.col())
66}
67
68/// Extract column `col` (0 = `day`) from decoded OHLCV CSV text.
69fn parse_series_csv(text: &str, col: usize) -> Result<Vec<(i32, f64)>, DataError> {
70    let mut out = Vec::new();
71    for line in text.lines() {
72        let line = line.trim();
73        if line.is_empty() || line.starts_with("day") {
74            continue;
75        }
76        let cells: Vec<&str> = line.split(',').collect();
77        if cells.len() <= col {
78            return Err(DataError::Parse(format!(
79                "row has {} cells, need > {col}",
80                cells.len()
81            )));
82        }
83        let day = date_to_i32(cells[0].trim())?;
84        let val: f64 = cells[col]
85            .trim()
86            .parse()
87            .map_err(|_| DataError::Parse(format!("bad value '{}'", cells[col])))?;
88        out.push((day, val));
89    }
90    Ok(out)
91}
92
93/// Serialize OHLCV rows (oldest-first) to gzip CSV with the standard header.
94pub fn write_series(rows: &[OhlcvRow]) -> Result<Vec<u8>, DataError> {
95    let mut buf = String::from("day,adj_open,adj_high,adj_low,adj_close,volume\n");
96    for r in rows {
97        buf.push_str(&format!(
98            "{},{},{},{},{},{}\n",
99            i32_to_date(r.day),
100            r.adj_open,
101            r.adj_high,
102            r.adj_low,
103            r.adj_close,
104            r.volume
105        ));
106    }
107    let mut enc = GzEncoder::new(Vec::new(), Compression::default());
108    enc.write_all(buf.as_bytes())
109        .map_err(|e| DataError::Io(e.to_string()))?;
110    enc.finish().map_err(|e| DataError::Io(e.to_string()))
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    fn row(day: i32, c: f64) -> OhlcvRow {
118        OhlcvRow {
119            day,
120            adj_open: c - 0.5,
121            adj_high: c + 1.0,
122            adj_low: c - 1.0,
123            adj_close: c,
124            volume: 1000.0,
125        }
126    }
127
128    #[test]
129    fn round_trips_and_selects_fields() {
130        let rows = vec![row(20240102, 10.5), row(20240103, 11.25)];
131        let gz = write_series(&rows).unwrap();
132        assert_eq!(
133            parse_series(&gz, Field::AdjClose).unwrap(),
134            vec![(20240102, 10.5), (20240103, 11.25)]
135        );
136        assert_eq!(
137            parse_series(&gz, Field::AdjHigh).unwrap(),
138            vec![(20240102, 11.5), (20240103, 12.25)]
139        );
140        assert_eq!(
141            parse_series(&gz, Field::AdjLow).unwrap(),
142            vec![(20240102, 9.5), (20240103, 10.25)]
143        );
144        assert_eq!(
145            parse_series(&gz, Field::AdjOpen).unwrap(),
146            vec![(20240102, 10.0), (20240103, 10.75)]
147        );
148        assert_eq!(
149            parse_series(&gz, Field::Volume).unwrap(),
150            vec![(20240102, 1000.0), (20240103, 1000.0)]
151        );
152    }
153
154    fn gz(text: &str) -> Vec<u8> {
155        let mut e = GzEncoder::new(Vec::new(), Compression::default());
156        e.write_all(text.as_bytes()).unwrap();
157        e.finish().unwrap()
158    }
159
160    #[test]
161    fn parse_errors_on_non_gzip_bad_value_bad_date_and_short_row() {
162        let hdr = "day,adj_open,adj_high,adj_low,adj_close,volume\n";
163        assert!(parse_series(b"not gzip", Field::AdjClose).is_err());
164        assert!(parse_series(
165            &gz(&format!("{hdr}2024-01-02,1,2,3,bad,5\n")),
166            Field::AdjClose
167        )
168        .is_err());
169        assert!(parse_series(&gz(&format!("{hdr}bad-date,1,2,3,4,5\n")), Field::AdjClose).is_err());
170        // row missing the volume column, but Volume requested
171        assert!(parse_series(&gz(&format!("{hdr}2024-01-02,1,2,3,4\n")), Field::Volume).is_err());
172    }
173}
yuzu_data/csv_io.rs

yuzu_data/
csv_io.rs
