CSV column discernment

Author: KyGost

Diff for: src/databases/csv/discern.rs

@@ -0,0 +1,102 @@
+use {
+	super::string_record_to_row,
+	crate::{
+		data::Schema, CSVDatabaseError, CSVSettings, Cast, Column, Result, Value,
+		ValueType,
+	},
+	csv::{Reader, ReaderBuilder, StringRecord},
+	std::{
+		default::Default,
+		fs::{File},
+	},
+};
+
+impl CSVSettings {
+	pub(crate) fn new_reader(&self, file: File) -> Reader<File> {
+		ReaderBuilder::new()
+			.delimiter(self.delimiter)
+			.has_headers(self.has_header.unwrap_or(true))
+			.from_reader(file)
+	}
+	pub(crate) fn discern_header(&mut self, header: &StringRecord) -> Vec<String> {
+		let header = string_record_to_row(header);
+
+		let has_header = if let Some(has_header) = self.has_header {
+			has_header
+		} else {
+			let has_header = !header
+				.iter()
+				.map(ValueType::from)
+				.any(|vt| vt != ValueType::Str);
+			self.has_header = Some(has_header);
+			has_header
+		};
+
+		if has_header {
+			header
+				.into_iter()
+				.map(Cast::cast)
+				.collect::<Result<Vec<String>>>()
+				.unwrap()
+		} else {
+			header
+				.into_iter()
+				.enumerate()
+				.map(|(index, _)| format!("column_{}", index))
+				.collect()
+		}
+	}
+	pub(crate) fn discern_schema(&mut self, file: File) -> Result<Option<Schema>> {
+		let mut reader = self.new_reader(file);
+		let header = reader
+			.headers()
+			.map_err(|error| CSVDatabaseError::HeaderError(format!("{:?}", error)))?;
+		let header = self.discern_header(header);
+		if header.is_empty() {
+			return Ok(None);
+		}
+		let value_types = self.discern_types(reader);
+
+		let column_defs = header
+			.into_iter()
+			.zip(value_types)
+			.map(|(header, value_type)| {
+				let mut column = Column::default();
+				column.name = header;
+				column.data_type = value_type;
+				column
+			})
+			.collect();
+		Ok(Some(Schema {
+			table_name: String::new(),
+			column_defs,
+			indexes: vec![],
+		}))
+	}
+	pub(crate) fn discern_types(&self, reader: Reader<File>) -> Vec<ValueType> {
+		let sample = reader
+			.into_records()
+			.take(self.sample_rows)
+			.map(|record| string_record_to_row(&record.unwrap()))
+			.collect();
+		discern_types_from_sample(sample)
+	}
+}
+
+pub(crate) fn discern_types_from_sample(sample: Vec<Vec<Value>>) -> Vec<ValueType> {
+	let mut types = sample
+		.into_iter()
+		.map(|row| row.iter().map(ValueType::from).collect());
+	let first_types = types.next().unwrap();
+	types.fold(first_types, |mut out_types, row_types| {
+		out_types
+			.iter_mut()
+			.zip(row_types)
+			.for_each(|(out_type, row_type)| {
+				if out_type != &row_type {
+					*out_type = ValueType::Any
+				}
+			});
+		out_types
+	})
+}

Diff for: src/databases/csv/mod.rs

@@ -1,16 +1,19 @@
 mod auto_increment;
 mod base;
+mod discern;
 mod mutable;
+mod record;
 mod utils;
 
+pub use {discern::*, record::*};
+
 use {
-	crate::{data::Schema, Column, DBFull, Database, Result, ValueType, WIPError},
-	csv::ReaderBuilder,
+	crate::{data::Schema, DBFull, Database, Result, WIPError},
 	serde::{Deserialize, Serialize},
 	std::{
 		default::Default,
 		fmt::Debug,
-		fs::{File, OpenOptions},
+		fs::{OpenOptions},
 	},
 	thiserror::Error,
 };
@@ -19,6 +22,9 @@ use {
 pub enum CSVDatabaseError {
 	#[error("CSV storages only support one table at a time")]
 	OnlyOneTableAllowed,
+
+	#[error("Failed to open CSV because of a error with header: {0}")]
+	HeaderError(String),
 }
 
 pub struct CSVDatabase {
@@ -30,12 +36,16 @@ pub struct CSVDatabase {
 pub struct CSVSettings {
 	pub delimiter: u8,
 	pub quoting: bool,
+	pub has_header: Option<bool>,
+	pub sample_rows: usize,
 }
 impl Default for CSVSettings {
 	fn default() -> Self {
 		Self {
 			delimiter: b',',
 			quoting: true,
+			has_header: None,
+			sample_rows: 100,
 		}
 	}
 }
@@ -51,46 +61,19 @@ impl CSVDatabase {
 	pub fn new(path: &str) -> Result<Self> {
 		Self::new_with_settings(path, CSVSettings::default())
 	}
-	pub fn new_with_settings(path: &str, csv_settings: CSVSettings) -> Result<Self> {
+	pub fn new_with_settings(path: &str, mut csv_settings: CSVSettings) -> Result<Self> {
 		let file = OpenOptions::new()
 			.read(true)
 			.write(true)
 			.create(true)
 			.open(path)
 			.map_err(|error| WIPError::Debug(format!("{:?}", error)))?;
 
-		let schema = discern_schema(file, &csv_settings)?;
+		let schema = csv_settings.discern_schema(file)?;
 		Ok(Self {
 			schema,
 			path: path.to_string(),
 			csv_settings,
 		})
 	}
 }
-
-fn discern_schema(file: File, csv_settings: &CSVSettings) -> Result<Option<Schema>> {
-	let mut reader = ReaderBuilder::new()
-		.delimiter(csv_settings.delimiter)
-		.from_reader(file);
-	let headers = reader
-		.headers()
-		.map_err(|error| WIPError::Debug(format!("{:?}", error)))?;
-	let column_defs = headers
-		.iter()
-		.map(|header| {
-			let mut column = Column::default();
-			column.name = header.to_string();
-			column.data_type = ValueType::Str;
-			column
-		})
-		.collect();
-	if headers.is_empty() {
-		Ok(None)
-	} else {
-		Ok(Some(Schema {
-			table_name: String::new(),
-			column_defs,
-			indexes: vec![],
-		}))
-	}
-}

Diff for: src/databases/csv/record.rs

@@ -0,0 +1,15 @@
+use {crate::Value, csv::StringRecord};
+
+pub fn string_record_to_row(record: &StringRecord) -> Vec<Value> {
+	record.iter().map(csv_cell_to_value).collect()
+}
+
+fn csv_cell_to_value(cell: &str) -> Value {
+	let cell = cell.to_string();
+	cell.parse::<bool>()
+		.map(|v| Value::Bool(v))
+		.or_else(|_| cell.parse::<u64>().map(|v| Value::U64(v)))
+		.or_else(|_| cell.parse::<i64>().map(|v| Value::I64(v)))
+		.or_else(|_| cell.parse::<f64>().map(|v| Value::F64(v)))
+		.unwrap_or_else(|_| Value::Str(cell))
+}