yuzu_data/
format.rs

1//! On-the-wire format detection for data files. Objects can be stored as gzip
2//! CSV (`.csv.gz`, the default write format), plain CSV (`.csv`), or Apache
3//! Parquet (`.parquet`, requires the `parquet` feature). Detection is by content
4//! magic bytes, not the object key, so a file parses correctly regardless of its
5//! extension and existing gzip CSV data keeps working unchanged.
6
7use crate::error::DataError;
8use flate2::read::GzDecoder;
9use std::io::Read;
10
11/// The storage encoding of a data object, inferred from its leading bytes.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Format {
14    /// gzip-compressed CSV — the default per-symbol/panel write format (`.csv.gz`).
15    CsvGz,
16    /// Uncompressed CSV text (`.csv`).
17    Csv,
18    /// Apache Parquet (`.parquet`).
19    Parquet,
20}
21
22impl Format {
23    /// Infer the format from a buffer's magic bytes. gzip starts `1f 8b`,
24    /// Parquet files start (and end) with the ASCII marker `PAR1`; anything
25    /// else is treated as plain CSV text.
26    pub fn detect(bytes: &[u8]) -> Format {
27        if bytes.starts_with(&[0x1f, 0x8b]) {
28            Format::CsvGz
29        } else if bytes.starts_with(b"PAR1") {
30            Format::Parquet
31        } else {
32            Format::Csv
33        }
34    }
35}
36
37/// Candidate object-key extensions to probe for a per-symbol/panel object, in
38/// priority order. `.csv.gz` is first so existing deployments hit on the first
39/// GET; `.parquet` is only probed when the feature that can read it is enabled
40/// (otherwise the extra GET would fetch bytes we could not parse).
41#[cfg(feature = "parquet")]
42pub(crate) const CANDIDATE_EXTS: &[&str] = &[".csv.gz", ".parquet", ".csv"];
43#[cfg(not(feature = "parquet"))]
44pub(crate) const CANDIDATE_EXTS: &[&str] = &[".csv.gz", ".csv"];
45
46/// Decode a CSV-family buffer to text: gunzip when gzip, otherwise interpret the
47/// bytes as UTF-8 as-is (plain CSV). Parquet buffers are handled elsewhere and
48/// must not be passed here.
49pub(crate) fn read_csv_text(bytes: &[u8]) -> Result<String, DataError> {
50    match Format::detect(bytes) {
51        Format::CsvGz => {
52            let mut text = String::new();
53            GzDecoder::new(bytes)
54                .read_to_string(&mut text)
55                .map_err(|e| DataError::Io(e.to_string()))?;
56            Ok(text)
57        }
58        Format::Csv => {
59            String::from_utf8(bytes.to_vec()).map_err(|e| DataError::Parse(e.to_string()))
60        }
61        Format::Parquet => Err(DataError::Parse(
62            "expected CSV bytes, got Parquet".to_string(),
63        )),
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use flate2::write::GzEncoder;
71    use flate2::Compression;
72    use std::io::Write;
73
74    fn gz(text: &str) -> Vec<u8> {
75        let mut e = GzEncoder::new(Vec::new(), Compression::default());
76        e.write_all(text.as_bytes()).unwrap();
77        e.finish().unwrap()
78    }
79
80    #[test]
81    fn detects_each_format_from_magic_bytes() {
82        assert_eq!(Format::detect(&gz("day,close\n")), Format::CsvGz);
83        assert_eq!(Format::detect(b"day,close\n2024-01-02,10\n"), Format::Csv);
84        assert_eq!(Format::detect(b"PAR1\x00\x00PAR1"), Format::Parquet);
85        assert_eq!(Format::detect(b""), Format::Csv);
86    }
87
88    #[test]
89    fn read_csv_text_handles_gzip_and_plain() {
90        let plain = "day,close\n2024-01-02,10\n";
91        assert_eq!(read_csv_text(plain.as_bytes()).unwrap(), plain);
92        assert_eq!(read_csv_text(&gz(plain)).unwrap(), plain);
93        assert!(read_csv_text(b"PAR1junk").is_err());
94    }
95}
yuzu_data/format.rs

yuzu_data/
format.rs
