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, CsvTypeVariants, 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 {
38 self.base_dir.join(format!("{}.csv", table_name))
39 }
40
41 #[cfg(feature = "vista")]
42 pub(crate) fn base_dir(&self) -> &std::path::Path {
43 &self.base_dir
44 }
45
46 pub(crate) fn read_csv(
54 &self,
55 table_name: &str,
56 columns: &IndexMap<String, Column<AnyCsvType>>,
57 ) -> Result<IndexMap<String, Record<AnyCsvType>>> {
58 let path = self.file_path(table_name);
59 let mut reader = csv::Reader::from_path(&path)
60 .map_err(|e| error!("Failed to open CSV file", path = path.display(), detail = e))?;
61
62 let headers = reader
63 .headers()
64 .map_err(|e| error!("Failed to read CSV headers", detail = e))?
65 .clone();
66
67 let mut key_for: Vec<String> = headers.iter().map(|h| h.to_string()).collect();
71 let mut variant_for: Vec<Option<CsvTypeVariants>> = vec![None; headers.len()];
72 for (name, col) in columns {
73 let csv_header = col.alias().unwrap_or(name);
74 if let Some(idx) = headers.iter().position(|h| h == csv_header) {
75 key_for[idx] = name.clone();
76 variant_for[idx] = variant_from_type_name(col.get_type());
77 }
78 }
79
80 let id_header = columns
82 .get(&self.id_column)
83 .and_then(|c| c.alias())
84 .unwrap_or(&self.id_column);
85 let id_col_index = headers.iter().position(|h| h == id_header);
86
87 let mut records = IndexMap::new();
88 for (row_idx, result) in reader.records().enumerate() {
89 let csv_record = result.map_err(|e| error!("Failed to read CSV row", detail = e))?;
90
91 let id = if let Some(idx) = id_col_index {
92 csv_record.get(idx).unwrap_or_default().to_string()
93 } else {
94 row_idx.to_string()
95 };
96
97 let mut record = Record::new();
98 for (i, field) in csv_record.iter().enumerate() {
99 if let Some(key) = key_for.get(i) {
100 let variant = variant_for.get(i).copied().flatten();
101 record.insert(key.clone(), parse_with_type(field, variant));
102 }
103 }
104
105 records.insert(id, record);
106 }
107
108 Ok(records)
109 }
110}