pub mod csv;
pub mod delimited;
pub mod frame;
pub mod knobs;
#[cfg(feature = "xlsx")]
pub mod xlsx;
use super::table::RawCsv;
use crate::datatypes::values::ColumnType;
use indexmap::IndexMap;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
pub struct FormatSpec {
pub name: &'static str,
pub accepted_keys: &'static [&'static str],
pub knob_keys: &'static [&'static str],
pub validate_entry: ValidateEntry,
}
pub type ValidateEntry = fn(&str, &super::schema::FileSpec) -> Result<(), String>;
pub fn accept_any_entry(_name: &str, _file: &super::schema::FileSpec) -> Result<(), String> {
Ok(())
}
pub const INPUT_FORMATS: &[FormatSpec] = &[
csv::FORMAT,
delimited::FORMAT,
frame::FORMAT,
#[cfg(feature = "xlsx")]
xlsx::FORMAT,
];
pub fn input_format(name: &str) -> Option<&'static FormatSpec> {
INPUT_FORMATS.iter().find(|f| f.name == name)
}
pub fn input_format_names() -> String {
INPUT_FORMATS
.iter()
.map(|f| format!("'{}'", f.name))
.collect::<Vec<_>>()
.join(", ")
}
pub(crate) fn resolve_input_path(root: &Path, path: &str) -> PathBuf {
if Path::new(path).is_absolute() {
PathBuf::from(path)
} else {
root.join(path)
}
}
pub trait Source: Send + Sync {
fn display_name(&self) -> &str;
fn size_hint(&self) -> Option<u64>;
fn can_chunk(&self) -> bool;
fn read_all(&self) -> Result<RawCsv, String>;
fn chunks(
&self,
chunk_size: usize,
) -> Result<Box<dyn Iterator<Item = Result<RawCsv, String>> + '_>, String>;
fn known_column_types(&self) -> HashMap<String, ColumnType> {
HashMap::new()
}
fn read_warnings(&self) -> Vec<String> {
Vec::new()
}
fn scan_columns(
&self,
columns: &[&str],
visit: &mut dyn FnMut(usize, &str) -> bool,
) -> Result<(), String> {
for chunk in self.chunks(SCAN_CHUNK_ROWS)? {
let raw = chunk?;
let indices: Vec<Option<usize>> =
columns.iter().map(|name| raw.col_index(name)).collect();
for r in 0..raw.row_count() {
for (slot, idx) in indices.iter().enumerate() {
let Some(idx) = idx else { continue };
if raw.nulls[r][*idx] {
continue;
}
if !visit(slot, &raw.rows[r][*idx]) {
return Ok(());
}
}
}
}
Ok(())
}
}
const SCAN_CHUNK_ROWS: usize = 65_536;
#[derive(Default)]
pub struct InputRegistry {
sources: IndexMap<String, Box<dyn Source>>,
}
impl InputRegistry {
pub fn insert(&mut self, name: impl Into<String>, source: Box<dyn Source>) {
self.sources.entry(name.into()).or_insert(source);
}
pub fn read_warnings(&self) -> Vec<String> {
self.sources
.values()
.flat_map(|s| s.read_warnings())
.collect()
}
pub fn get(&self, name: &str) -> Result<&dyn Source, String> {
match self.sources.get(name) {
Some(s) => Ok(s.as_ref()),
None if self.sources.is_empty() => Err(format!(
"input '{name}' is not declared (no inputs declared)"
)),
None => Err(format!(
"input '{name}' is not declared; declared inputs: {}",
self.sources
.keys()
.map(String::as_str)
.collect::<Vec<_>>()
.join(", ")
)),
}
}
}
#[cfg(test)]
pub mod test_double {
use super::{RawCsv, Source};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
pub struct CountingSource {
inner: Box<dyn Source>,
opens: Arc<AtomicUsize>,
}
impl CountingSource {
pub fn new(inner: Box<dyn Source>) -> (Self, Arc<AtomicUsize>) {
let opens = Arc::new(AtomicUsize::new(0));
(
Self {
inner,
opens: opens.clone(),
},
opens,
)
}
}
pub struct DefaultScanSource(pub Box<dyn Source>);
impl Source for DefaultScanSource {
fn display_name(&self) -> &str {
self.0.display_name()
}
fn size_hint(&self) -> Option<u64> {
self.0.size_hint()
}
fn can_chunk(&self) -> bool {
self.0.can_chunk()
}
fn read_all(&self) -> Result<RawCsv, String> {
self.0.read_all()
}
fn chunks(
&self,
chunk_size: usize,
) -> Result<Box<dyn Iterator<Item = Result<RawCsv, String>> + '_>, String> {
self.0.chunks(chunk_size)
}
}
impl Source for CountingSource {
fn display_name(&self) -> &str {
self.inner.display_name()
}
fn size_hint(&self) -> Option<u64> {
self.inner.size_hint()
}
fn can_chunk(&self) -> bool {
self.inner.can_chunk()
}
fn read_all(&self) -> Result<RawCsv, String> {
self.opens.fetch_add(1, Ordering::SeqCst);
self.inner.read_all()
}
fn chunks(
&self,
chunk_size: usize,
) -> Result<Box<dyn Iterator<Item = Result<RawCsv, String>> + '_>, String> {
self.opens.fetch_add(1, Ordering::SeqCst);
self.inner.chunks(chunk_size)
}
}
}
#[cfg(test)]
mod registry_tests {
use super::csv::CsvFile;
use super::*;
use std::path::PathBuf;
fn registry_with(names: &[&str]) -> InputRegistry {
let mut reg = InputRegistry::default();
for n in names {
reg.insert(*n, Box::new(CsvFile::new(PathBuf::from(*n), n.to_string())));
}
reg
}
#[test]
fn unknown_name_error_lists_the_declared_ones() {
let reg = registry_with(&["a.csv", "b.csv"]);
let err = reg
.get("c.csv")
.err()
.expect("an undeclared name is an error");
assert_eq!(
err,
"input 'c.csv' is not declared; declared inputs: a.csv, b.csv"
);
}
#[test]
fn unknown_name_on_an_empty_registry_says_so() {
let reg = InputRegistry::default();
let err = reg
.get("c.csv")
.err()
.expect("an undeclared name is an error");
assert_eq!(err, "input 'c.csv' is not declared (no inputs declared)");
}
#[test]
fn a_name_declared_twice_keeps_the_first_source() {
let mut reg = InputRegistry::default();
reg.insert(
"a.csv",
Box::new(CsvFile::new(PathBuf::from("one/a.csv"), "first".into())),
);
reg.insert(
"a.csv",
Box::new(CsvFile::new(PathBuf::from("two/a.csv"), "second".into())),
);
assert_eq!(reg.get("a.csv").unwrap().display_name(), "first");
}
}
#[cfg(test)]
mod scan_tests {
use super::csv::CsvFile;
use super::test_double::DefaultScanSource;
use super::Source;
use std::io::Write;
fn write(content: &[u8]) -> tempfile::NamedTempFile {
let mut f = tempfile::NamedTempFile::new().unwrap();
f.write_all(content).unwrap();
f.flush().unwrap();
f
}
fn csv_file(f: &tempfile::NamedTempFile) -> CsvFile {
CsvFile::new(f.path().to_path_buf(), "sample.csv".to_string())
}
fn collect(source: &dyn Source, columns: &[&str]) -> Vec<(usize, String)> {
let mut seen = Vec::new();
source
.scan_columns(columns, &mut |slot, cell| {
seen.push((slot, cell.to_string()));
true
})
.unwrap();
seen
}
#[test]
fn the_csv_override_sees_exactly_what_the_default_does() {
let mut content = String::from("a,b,c\n");
for i in 0..300 {
let b = if i % 5 == 0 {
String::new()
} else {
format!("b{i}")
};
content.push_str(&format!("{i},{b}, \n"));
}
let f = write(content.as_bytes());
let fast = collect(&csv_file(&f), &["a", "b", "missing"]);
let slow = collect(
&DefaultScanSource(Box::new(csv_file(&f))),
&["a", "b", "missing"],
);
assert_eq!(fast, slow);
assert_eq!(fast.len(), 300 + 240);
assert!(fast.iter().all(|(slot, _)| *slot < 2));
}
#[test]
fn the_csv_override_stops_before_reading_the_next_row() {
let mut content = b"a\nfirst\n".to_vec();
content.extend_from_slice(&[0xff, 0xfe, b'\n']);
let f = write(&content);
let mut visits = 0usize;
let result = csv_file(&f).scan_columns(&["a"], &mut |_, _| {
visits += 1;
false
});
assert!(result.is_ok(), "{result:?}");
assert_eq!(visits, 1);
}
#[test]
fn the_default_stops_calling_the_sink_when_it_says_stop() {
let f = write(b"a\n1\n2\n3\n4\n");
let source = DefaultScanSource(Box::new(csv_file(&f)));
let mut visits = 0usize;
source
.scan_columns(&["a"], &mut |_, _| {
visits += 1;
false
})
.unwrap();
assert_eq!(visits, 1);
}
}