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

yuzu_data/
csv_io.rs
