1use std::path::PathBuf;
2
3use indexmap::IndexMap;
4use vantage_core::error;
5use vantage_dataset::traits::Result;
6use vantage_table::column::core::Column;
7use vantage_table::traits::column_like::ColumnLike;
8use vantage_types::Record;
9
10use crate::type_system::{AnyCsvType, parse_with_type, variant_from_type_name};
11
12#[derive(Clone, Debug)]
17pub struct Csv {
18 base_dir: PathBuf,
19 pub(crate) id_column: String,
20}
21
22impl Csv {
23 pub fn new(base_dir: impl Into<PathBuf>) -> Self {
25 Self {
26 base_dir: base_dir.into(),
27 id_column: "id".to_string(),
28 }
29 }
30
31 pub fn with_id_column(mut self, column: impl Into<String>) -> Self {
33 self.id_column = column.into();
34 self
35 }
36
37 fn file_path(&self, table_name: &str) -> PathBuf {
39 self.base_dir.join(format!("{}.csv", table_name))
40 }
41
42 pub(crate) fn read_csv(
47 &self,
48 table_name: &str,
49 columns: &IndexMap<String, Column<AnyCsvType>>,
50 ) -> Result<IndexMap<String, Record<AnyCsvType>>> {
51 let path = self.file_path(table_name);
52 let mut reader = csv::Reader::from_path(&path)
53 .map_err(|e| error!("Failed to open CSV file", path = path.display(), detail = e))?;
54
55 let headers = reader
56 .headers()
57 .map_err(|e| error!("Failed to read CSV headers", detail = e))?
58 .clone();
59
60 let id_col_index = headers.iter().position(|h| h == self.id_column);
61
62 let mut records = IndexMap::new();
63
64 for (row_idx, result) in reader.records().enumerate() {
65 let csv_record = result.map_err(|e| error!("Failed to read CSV row", detail = e))?;
66
67 let id = if let Some(idx) = id_col_index {
68 csv_record.get(idx).unwrap_or_default().to_string()
69 } else {
70 row_idx.to_string()
71 };
72
73 let mut record = Record::new();
74 for (i, field) in csv_record.iter().enumerate() {
75 if let Some(header) = headers.get(i) {
76 let variant = columns
77 .get(header)
78 .and_then(|col| variant_from_type_name(col.get_type()));
79 record.insert(header.to_string(), parse_with_type(field, variant));
80 }
81 }
82
83 records.insert(id, record);
84 }
85
86 Ok(records)
87 }
88}