use rudb_common::{Error, Field, LogicalType, Result, Value};
use rudb_io::File;
use rudb_kernels::cast_value;
use rudb_vector::{Chunk, VECTOR_SIZE, Vector};
use crate::dialect::{self, Dialect, Given};
use crate::infer;
const BLOCK: usize = 1 << 20;
#[derive(Debug)]
pub struct Reader {
file: Box<dyn File>,
path: String,
given: Given,
dialect: Dialect,
fields: Vec<Field>,
projection: Vec<usize>,
buffer: Vec<u8>,
at: usize,
offset: u64,
drained: bool,
line: u64,
scratch: Vec<String>,
}
impl Reader {
pub fn open(file: Box<dyn File>, path: &str) -> Result<Self> {
Self::open_with(file, path, Given::default())
}
pub fn open_with(file: Box<dyn File>, path: &str, given: Given) -> Result<Self> {
let mut reader = Self {
file,
path: path.to_string(),
given,
dialect: Dialect::comma_separated(),
fields: Vec::new(),
projection: Vec::new(),
buffer: Vec::new(),
at: 0,
offset: 0,
drained: false,
line: 1,
scratch: Vec::new(),
};
reader.fill()?;
let sample = reader.buffer.clone();
let quote = given.quote.or_else(|| dialect::quote(&sample));
let delimiter = match given.delimiter {
Some(byte) => byte,
None => dialect::delimiter(&sample, quote)?,
};
let escape = given.escape.or(quote);
reader.dialect = Dialect { delimiter, quote, escape, header: false };
let rows = reader.sample_rows(&sample)?;
let (header, fields) = describe(&rows, given.header);
reader.dialect.header = header;
reader.fields = fields;
reader.projection = (0..reader.fields.len()).collect();
if header {
reader.skip_record()?;
}
Ok(reader)
}
#[must_use]
pub fn fields(&self) -> Vec<Field> {
self.projection.iter().map(|&at| self.fields[at].clone()).collect()
}
pub fn project(&mut self, columns: &[usize]) -> Result<()> {
for &column in columns {
if column >= self.fields.len() {
return Err(Error::io(format!(
"column {column} is past the {} the file has",
self.fields.len()
)));
}
}
self.projection = columns.to_vec();
Ok(())
}
pub fn retype(&mut self, types: &[LogicalType]) -> Result<()> {
if types.len() != self.projection.len() {
return Err(Error::io(format!(
"{} types for a projection of {} columns",
types.len(),
self.projection.len()
)));
}
for (&at, ty) in self.projection.iter().zip(types) {
self.fields[at].ty = ty.clone();
}
Ok(())
}
#[must_use]
pub const fn dialect(&self) -> Dialect {
self.dialect
}
pub fn next_chunk(&mut self) -> Result<Option<Chunk>> {
let mut rows: Vec<Vec<Option<String>>> = Vec::new();
while rows.len() < VECTOR_SIZE {
match self.next_record()? {
Some(fields) => rows.push(fields),
None => break,
}
}
if rows.is_empty() {
return Ok(None);
}
let mut columns = Vec::with_capacity(self.projection.len());
for &at in &self.projection {
let field = &self.fields[at];
let mut values = Vec::with_capacity(rows.len());
for (row, held) in rows.iter().enumerate() {
let text = held.get(at).and_then(Option::as_deref);
values.push(self.convert(
text,
field,
self.line - rows.len() as u64 + row as u64,
)?);
}
columns.push(Vector::from_values(field.ty.clone(), &values)?);
}
Ok(Some(Chunk::with_rows(columns, rows.len())?))
}
fn convert(&self, text: Option<&str>, field: &Field, line: u64) -> Result<Value> {
let Some(text) = text else { return Ok(Value::Null) };
if field.ty == LogicalType::Varchar {
return Ok(Value::Varchar(text.to_string()));
}
let value = Value::Varchar(text.to_string());
match cast_value(&value, &field.ty, false) {
Ok(converted) => Ok(converted),
Err(_) => Err(Error::conversion(self.conversion_error(text, field, line))),
}
}
fn conversion_error(&self, text: &str, field: &Field, line: u64) -> String {
format!(
"CSV Error on Line: {line}\nOriginal Line: {text}\nError when converting column \
\"{}\". Could not convert string \"{text}\" to '{}'\n\nColumn {} is being converted \
as type {}\nThis type was auto-detected from the CSV file.\nPossible solutions:\n* \
Override the type for this column manually by setting the type explicitly, e.g., \
types={{'{}': 'VARCHAR'}}\n* Set the sample size to a larger value to enable the \
auto-detection to scan more values, e.g., sample_size=-1\n* Use a COPY statement to \
automatically derive types from an existing table.\n* Check whether the null string \
value is set correctly (e.g., nullstr = 'N/A')\n\n file = {}\n delimiter = {}\n \
quote = {}\n escape = {}\n header = {} {}\n sample_size = {}\n",
field.name,
field.ty,
field.name,
field.ty,
field.name,
self.path,
Given::shown(self.given.delimiter, Some(self.dialect.delimiter)),
Given::shown(self.given.quote, self.dialect.quote),
Given::shown(self.given.escape, self.dialect.escape),
self.dialect.header,
Given::source(self.given.header.is_some()),
infer::SAMPLE,
)
}
fn next_record(&mut self) -> Result<Option<Vec<Option<String>>>> {
let Some(()) = self.advance()? else { return Ok(None) };
Ok(Some(
self.scratch
.iter()
.map(|text| if text.is_empty() { None } else { Some(text.clone()) })
.collect(),
))
}
fn advance(&mut self) -> Result<Option<()>> {
loop {
let mut scratch = std::mem::take(&mut self.scratch);
let outcome = crate::scan::record(
&self.buffer,
self.at,
self.dialect,
self.drained,
&mut scratch,
);
self.scratch = scratch;
match outcome? {
Some(next) => {
self.at = next;
self.line += 1;
return Ok(Some(()));
}
None if self.drained => return Ok(None),
None => self.fill()?,
}
}
}
fn skip_record(&mut self) -> Result<()> {
self.advance()?;
Ok(())
}
fn fill(&mut self) -> Result<()> {
self.buffer.drain(..self.at);
self.at = 0;
let held = self.buffer.len();
self.buffer.resize(held + BLOCK, 0);
let read = self.file.read_at(self.offset, &mut self.buffer[held..])?;
self.buffer.truncate(held + read);
self.offset += read as u64;
if read == 0 {
self.drained = true;
}
Ok(())
}
fn sample_rows(&self, sample: &[u8]) -> Result<Vec<Vec<Option<String>>>> {
let mut rows = Vec::new();
let mut fields = Vec::new();
let mut at = 0;
while rows.len() <= infer::SAMPLE {
let Some(next) = crate::scan::record(sample, at, self.dialect, false, &mut fields)?
else {
break;
};
at = next;
rows.push(
fields
.iter()
.map(|text| if text.is_empty() { None } else { Some(text.clone()) })
.collect(),
);
}
Ok(rows)
}
}
fn describe(rows: &[Vec<Option<String>>], told: Option<bool>) -> (bool, Vec<Field>) {
let width = rows.iter().map(Vec::len).max().unwrap_or(0);
let body = types(&rows[1.min(rows.len())..], width);
let all_text = body.iter().all(|ty| *ty == LogicalType::Varchar);
let first_fits = rows.first().is_some_and(|first| {
first.iter().zip(&body).all(|(text, ty)| match text {
None => true,
Some(text) => infer::fits(text, ty),
})
});
let header = !rows.is_empty() && told.unwrap_or(rows.len() > 1 && (all_text || !first_fits));
if !header {
let types = types(rows, width);
let fields = types
.into_iter()
.enumerate()
.map(|(at, ty)| Field::new(format!("column{at}"), ty))
.collect();
return (false, fields);
}
let names = unique(&rows[0], width);
let fields = body.into_iter().zip(names).map(|(ty, name)| Field::new(name, ty)).collect();
(true, fields)
}
fn unique(header: &[Option<String>], width: usize) -> Vec<String> {
let mut taken: Vec<String> = Vec::with_capacity(width);
for at in 0..width {
let base = match header.get(at).and_then(Option::as_deref) {
Some(written) => written.to_string(),
None => format!("column{at}"),
};
let mut name = base.clone();
let mut next = 1;
while taken.iter().any(|held| held.eq_ignore_ascii_case(&name)) {
name = format!("{base}_{next}");
next += 1;
}
taken.push(name);
}
taken
}
fn types(rows: &[Vec<Option<String>>], width: usize) -> Vec<LogicalType> {
(0..width)
.map(|at| {
let values: Vec<Option<&str>> =
rows.iter().map(|row| row.get(at).and_then(Option::as_deref)).collect();
infer::column(&values)
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use rudb_io::{Filesystem, OpenMode, SimFilesystem};
use std::path::Path;
fn read(text: &str) -> Reader {
let filesystem = SimFilesystem::new();
let path = Path::new("/t.csv");
let file = filesystem.open(path, OpenMode::Create).expect("creates");
file.write_at(0, text.as_bytes()).expect("writes");
drop(file);
let file = filesystem.open(path, OpenMode::Read).expect("opens");
Reader::open(file, "/t.csv").expect("sniffs")
}
fn read_with(text: &str, given: Given) -> Reader {
let filesystem = SimFilesystem::new();
let path = Path::new("/t.csv");
let file = filesystem.open(path, OpenMode::Create).expect("creates");
file.write_at(0, text.as_bytes()).expect("writes");
drop(file);
let file = filesystem.open(path, OpenMode::Read).expect("opens");
Reader::open_with(file, "/t.csv", given).expect("reads")
}
fn names_and_types(reader: &Reader) -> Vec<(String, String)> {
reader.fields().iter().map(|f| (f.name.clone(), f.ty.to_string())).collect()
}
fn all(reader: &mut Reader) -> Vec<Vec<Value>> {
let mut rows = Vec::new();
while let Some(chunk) = reader.next_chunk().expect("reads") {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|at| chunk.value_at(row, at)).collect());
}
}
rows
}
#[test]
fn a_header_that_names_two_columns_the_same_thing_counts_the_second_one_up() {
let names: Vec<String> =
read("a,a,A,a_1\n1,2,3,4\nx,y,z,w\n").fields().into_iter().map(|f| f.name).collect();
assert_eq!(names, ["a", "a_1", "A_2", "a_1_1"]);
}
#[test]
fn a_header_and_three_types_are_what_duckdb_sniffs_for_the_same_bytes() {
let reader = read("a,b,c\n1,x,2.5\n2,y,3.5\n");
assert_eq!(
names_and_types(&reader),
[
("a".to_string(), "BIGINT".to_string()),
("b".to_string(), "VARCHAR".to_string()),
("c".to_string(), "DOUBLE".to_string()),
]
);
}
#[test]
fn a_file_with_no_header_gets_the_names_duckdb_gives_it() {
let reader = read("1,x\n2,y\n");
assert_eq!(
names_and_types(&reader),
[
("column0".to_string(), "BIGINT".to_string()),
("column1".to_string(), "VARCHAR".to_string()),
]
);
}
#[test]
fn two_rows_of_words_are_a_header_and_a_row() {
let reader = read("a,b\nc,d\n");
assert_eq!(reader.fields().iter().map(|f| f.name.clone()).collect::<Vec<_>>(), ["a", "b"]);
}
#[test]
fn one_column_of_words_under_a_row_of_numbers_is_still_a_header() {
let reader = read("a,2\n3,4\n");
assert_eq!(reader.fields().iter().map(|f| f.name.clone()).collect::<Vec<_>>(), ["a", "2"]);
}
#[test]
fn the_rows_are_the_rows_of_the_file() {
let mut reader = read("a,b\n1,x\n2,y\n");
assert_eq!(
all(&mut reader),
[
vec![Value::BigInt(1), Value::Varchar("x".into())],
vec![Value::BigInt(2), Value::Varchar("y".into())],
]
);
}
#[test]
fn an_empty_field_is_a_null_whether_it_was_quoted_or_not() {
let mut reader = read("a,b\n1,\n\"\",y\n");
assert_eq!(
all(&mut reader),
[vec![Value::BigInt(1), Value::Null], vec![Value::Null, Value::Varchar("y".into())],]
);
}
#[test]
fn a_projection_picks_columns_out_by_position_and_can_reorder_them() {
let mut reader = read("a,b,c\n1,x,2.5\n");
reader.project(&[2, 0]).expect("projects");
assert_eq!(reader.fields().iter().map(|f| f.name.clone()).collect::<Vec<_>>(), ["c", "a"]);
assert_eq!(all(&mut reader), [vec![Value::Double(2.5), Value::BigInt(1)]]);
}
#[test]
fn a_projection_of_nothing_still_counts_the_rows() {
let mut reader = read("a,b\n1,x\n2,y\n3,z\n");
reader.project(&[]).expect("projects");
let chunk = reader.next_chunk().expect("reads").expect("a chunk");
assert_eq!(chunk.len(), 3);
assert_eq!(chunk.width(), 0);
}
#[test]
fn a_pipe_separated_file_reads_as_one() {
let mut reader = read("a|b\n1|x\n");
assert_eq!(reader.dialect().delimiter, b'|');
assert_eq!(all(&mut reader), [vec![Value::BigInt(1), Value::Varchar("x".into())]]);
}
#[test]
fn a_quoted_field_with_a_delimiter_in_it_is_one_value() {
let mut reader = read("a,b\n1,\"x,y\"\n");
assert_eq!(all(&mut reader), [vec![Value::BigInt(1), Value::Varchar("x,y".into())]]);
}
#[test]
fn more_rows_than_fit_one_chunk_arrive_as_more_than_one_chunk() {
let mut text = String::from("a\n");
for row in 0..VECTOR_SIZE + 5 {
text.push_str(&format!("{row}\n"));
}
let mut reader = read(&text);
let first = reader.next_chunk().expect("reads").expect("a chunk");
assert_eq!(first.len(), VECTOR_SIZE);
let second = reader.next_chunk().expect("reads").expect("a second chunk");
assert_eq!(second.len(), 5);
assert!(reader.next_chunk().expect("reads").is_none());
}
#[test]
fn a_value_the_sniffer_never_saw_is_an_error_rather_than_a_wider_column() {
let mut text = String::from("c\n");
for row in 0..infer::SAMPLE {
text.push_str(&format!("{row}\n"));
}
text.push_str("oops\n");
let mut reader = read(&text);
assert_eq!(reader.fields()[0].ty, LogicalType::BigInt);
let error = all_or_error(&mut reader).unwrap_err();
let line = infer::SAMPLE + 2;
assert!(error.message().starts_with(&format!("CSV Error on Line: {line}")), "{error}");
assert!(
error.message().contains("Could not convert string \"oops\" to 'BIGINT'"),
"{error}"
);
assert!(error.message().contains("sample_size = 20480"), "{error}");
}
#[test]
fn a_file_told_it_has_no_header_reads_its_first_line_as_a_row() {
let mut reader =
read_with("a,b\n1,x\n2,y\n", Given { header: Some(false), ..Given::default() });
assert_eq!(
names_and_types(&reader),
[
("column0".to_string(), "VARCHAR".to_string()),
("column1".to_string(), "VARCHAR".to_string()),
]
);
assert_eq!(all(&mut reader).len(), 3);
}
#[test]
fn a_file_told_it_has_a_header_takes_its_first_line_as_the_names() {
let reader = read_with("1,2\n3,4\n", Given { header: Some(true), ..Given::default() });
assert_eq!(reader.fields().iter().map(|f| f.name.clone()).collect::<Vec<_>>(), ["1", "2"]);
}
#[test]
fn a_given_delimiter_is_the_delimiter_whatever_the_file_looks_like() {
let reader = read_with("a,b\n1,x\n", Given { delimiter: Some(b';'), ..Given::default() });
assert_eq!(reader.dialect().delimiter, b';');
assert_eq!(reader.fields().len(), 1);
}
#[test]
fn a_given_quote_makes_a_field_that_holds_the_delimiter_one_value() {
let mut reader =
read_with("a,b\n1,'x,y'\n", Given { quote: Some(b'\''), ..Given::default() });
assert_eq!(all(&mut reader), [vec![Value::BigInt(1), Value::Varchar("x,y".into())]]);
}
#[test]
fn the_block_says_set_by_user_for_what_the_call_gave_it() {
let mut text = String::from("c;d\n");
for row in 0..infer::SAMPLE {
text.push_str(&format!("{row};x\n"));
}
text.push_str("oops;x\n");
let given = Given { delimiter: Some(b';'), ..Given::default() };
let mut reader = read_with(&text, given);
let error = all_or_error(&mut reader).unwrap_err();
assert!(error.message().contains("delimiter = ; (Set By User)"), "{error}");
assert!(error.message().contains("header = true (Auto-Detected)"), "{error}");
}
#[test]
fn a_file_told_a_wider_type_than_it_sniffed_reads_its_whole_numbers_as_that_type() {
let mut reader = read("a\n1\n2\n");
assert_eq!(reader.fields()[0].ty, LogicalType::BigInt);
reader.retype(&[LogicalType::Double]).expect("one type for one column");
assert_eq!(reader.fields()[0].ty, LogicalType::Double);
assert_eq!(all(&mut reader), [[Value::Double(1.0)], [Value::Double(2.0)]]);
}
#[test]
fn a_type_list_that_is_not_as_long_as_the_projection_is_refused() {
let mut reader = read("a,b\n1,two\n");
let error = reader.retype(&[LogicalType::Double]).unwrap_err();
assert!(error.message().contains("1 types for a projection of 2 columns"), "{error}");
}
fn all_or_error(reader: &mut Reader) -> Result<Vec<Vec<Value>>> {
let mut rows = Vec::new();
while let Some(chunk) = reader.next_chunk()? {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|at| chunk.value_at(row, at)).collect());
}
}
Ok(rows)
}
}