use std::fs::File;
use std::path::{Path, PathBuf};
use crate::csv::parse_cell;
use crate::error::MattenDataError;
use crate::table::{CellValue, Table};
pub struct CsvBatchReader {
reader: ::csv::Reader<File>,
path: PathBuf,
headers: Vec<String>,
batch_rows: usize,
line: usize,
done: bool,
}
impl CsvBatchReader {
pub fn open(path: &Path, batch_rows: usize) -> Result<Self, MattenDataError> {
if batch_rows == 0 {
return Err(MattenDataError::InvalidBatchSize);
}
let file = File::open(path).map_err(|source| MattenDataError::Io {
path: path.to_path_buf(),
source,
})?;
let mut reader = ::csv::ReaderBuilder::new()
.has_headers(true)
.flexible(true)
.from_reader(file);
let headers: Vec<String> = reader
.headers()
.map_err(|e| map_csv_error(e, path))?
.iter()
.map(|h| h.trim().to_string())
.collect();
if headers.is_empty() {
return Err(MattenDataError::EmptyInput);
}
for (i, name) in headers.iter().enumerate() {
if name.is_empty() {
return Err(MattenDataError::Csv {
message: format!("header column {} is empty", i + 1),
});
}
}
for i in 0..headers.len() {
for j in (i + 1)..headers.len() {
if headers[i] == headers[j] {
return Err(MattenDataError::DuplicateColumn {
name: headers[i].clone(),
});
}
}
}
Ok(CsvBatchReader {
reader,
path: path.to_path_buf(),
headers,
batch_rows,
line: 1, done: false,
})
}
pub fn next_batch(&mut self) -> Result<Option<Table>, MattenDataError> {
if self.done {
return Ok(None);
}
let n = self.headers.len();
let mut rows: Vec<Vec<CellValue>> = Vec::with_capacity(self.batch_rows);
let mut record = ::csv::StringRecord::new();
while rows.len() < self.batch_rows {
let read = match self.reader.read_record(&mut record) {
Ok(read) => read,
Err(e) => {
self.done = true;
return Err(map_csv_error(e, &self.path));
}
};
if !read {
break; }
self.line += 1;
if record.is_empty() {
continue;
}
if record.len() != n {
self.done = true;
return Err(MattenDataError::RaggedRow {
row: self.line,
expected: n,
actual: record.len(),
});
}
rows.push(record.iter().map(parse_cell).collect());
}
if rows.is_empty() {
self.done = true;
return Ok(None);
}
Ok(Some(Table::from_parts(self.headers.clone(), rows)))
}
}
fn map_csv_error(e: ::csv::Error, path: &Path) -> MattenDataError {
let message = e.to_string();
match e.into_kind() {
::csv::ErrorKind::Io(source) => MattenDataError::Io {
path: path.to_path_buf(),
source,
},
_ => MattenDataError::Csv { message },
}
}