use std::io::Write;
use dataprof::{
ColumnStats, CsvParserConfig, DataType, EngineType, ProfileReport, Profiler, analyze_column,
analyze_csv_file, infer_schema,
};
use tempfile::NamedTempFile;
const NAMES: [&str; 6] = [
"leading_zero",
"es_postal",
"compact_date",
"departure_time",
"quantity",
"amount_cents",
];
const ROWS: [[&str; 6]; 6] = [
["00123", "28013", "20240115", "0930", "3", "12345678"],
["00456", "08001", "20240216", "1415", "12", "23456789"],
["00789", "41001", "20240317", "0805", "0", "34567891"],
["01234", "46001", "20240418", "2210", "25", "45678912"],
["05678", "48001", "20240519", "1130", "7", "56789123"],
["09999", "50001", "20240620", "0645", "100", "67891234"],
];
const EXPECTED: [(&str, DataType); 6] = [
("leading_zero", DataType::String),
("es_postal", DataType::String),
("compact_date", DataType::String),
("departure_time", DataType::String),
("quantity", DataType::Integer),
("amount_cents", DataType::Integer),
];
fn csv_fixture() -> NamedTempFile {
let mut file = NamedTempFile::with_suffix(".csv").unwrap();
writeln!(file, "{}", NAMES.join(",")).unwrap();
for row in ROWS {
writeln!(file, "{}", row.join(",")).unwrap();
}
file.flush().unwrap();
file
}
fn json_records() -> Vec<String> {
ROWS.iter()
.map(|row| {
let fields: Vec<String> = NAMES
.iter()
.zip(row)
.map(|(name, value)| format!(r#""{name}":"{value}""#))
.collect();
format!("{{{}}}", fields.join(","))
})
.collect()
}
fn json_fixture() -> NamedTempFile {
let mut file = NamedTempFile::with_suffix(".json").unwrap();
write!(file, "[{}]", json_records().join(",")).unwrap();
file.flush().unwrap();
file
}
fn jsonl_fixture() -> NamedTempFile {
let mut file = NamedTempFile::with_suffix(".jsonl").unwrap();
writeln!(file, "{}", json_records().join("\n")).unwrap();
file.flush().unwrap();
file
}
#[cfg(feature = "parquet")]
fn parquet_fixture() -> NamedTempFile {
use std::sync::Arc;
use arrow::array::{ArrayRef, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use parquet::arrow::ArrowWriter;
let schema = Arc::new(Schema::new(
NAMES
.iter()
.map(|name| Field::new(*name, DataType::Utf8, false))
.collect::<Vec<_>>(),
));
let columns: Vec<ArrayRef> = (0..NAMES.len())
.map(|index| {
Arc::new(StringArray::from(
ROWS.iter().map(|row| row[index]).collect::<Vec<_>>(),
)) as ArrayRef
})
.collect();
let batch = RecordBatch::try_new(schema.clone(), columns).unwrap();
let file = NamedTempFile::with_suffix(".parquet").unwrap();
let mut writer = ArrowWriter::try_new(file.reopen().unwrap(), schema, None).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();
file
}
fn reports() -> Vec<(String, ProfileReport)> {
let csv = csv_fixture();
let mut reports = vec![(
"standard".to_string(),
analyze_csv_file(csv.path(), &CsvParserConfig::default()).expect("standard CSV"),
)];
for engine in [
EngineType::Auto,
EngineType::Incremental,
EngineType::Columnar,
] {
let report = Profiler::new()
.engine(engine)
.analyze_file(csv.path())
.unwrap_or_else(|e| panic!("[{engine:?}] {e}"));
reports.push((format!("{engine:?}"), report));
}
let mut files = vec![("json", json_fixture()), ("jsonl", jsonl_fixture())];
#[cfg(feature = "parquet")]
files.push(("parquet", parquet_fixture()));
for (label, file) in files {
let report = Profiler::new()
.analyze_file(file.path())
.unwrap_or_else(|e| panic!("[{label}] {e}"));
reports.push((label.to_string(), report));
}
reports
}
#[test]
fn every_path_types_digit_codes_as_text() {
let mut wrong = Vec::new();
for (label, report) in reports() {
let document = serde_json::to_value(&report).expect("serializes");
let restored: ProfileReport =
serde_json::from_value(document).expect("the report reads back");
for (name, expected) in EXPECTED {
for (side, report) in [("", &report), (" read back", &restored)] {
let column = report
.column_profiles
.iter()
.find(|column| column.name == name)
.unwrap_or_else(|| panic!("[{label}] {name} column{side}"));
let stats_match = match expected {
DataType::String => matches!(column.stats, ColumnStats::Text(_)),
_ => matches!(column.stats, ColumnStats::Numeric(_)),
};
if column.data_type != expected || !stats_match {
wrong.push(format!(
"[{label}] {name}{side}: {:?} with {:?}",
column.data_type, column.stats
));
}
}
}
let consistency = report
.quality
.as_ref()
.and_then(|quality| quality.metrics.consistency.as_ref())
.unwrap_or_else(|| panic!("[{label}] consistency assessed"));
if consistency.data_type_consistency != 100.0 {
wrong.push(format!(
"[{label}] data_type_consistency {}",
consistency.data_type_consistency
));
}
}
assert!(wrong.is_empty(), "{}", wrong.join("\n"));
}
#[test]
fn schema_inference_types_them_the_same_way() {
let mut files = vec![
("csv", csv_fixture()),
("json", json_fixture()),
("jsonl", jsonl_fixture()),
];
#[cfg(feature = "parquet")]
files.push(("parquet", parquet_fixture()));
for (label, file) in files {
let schema = infer_schema(file.path()).unwrap_or_else(|e| panic!("[{label}] {e}"));
for (name, expected) in EXPECTED {
let column = schema
.columns
.iter()
.find(|column| column.name == name)
.unwrap_or_else(|| panic!("[{label}] {name} column"));
assert_eq!(column.data_type, expected, "[{label}] {name}");
}
}
}
#[test]
fn column_analysis_types_them_the_same_way() {
for (index, (name, expected)) in EXPECTED.into_iter().enumerate() {
let values: Vec<String> = ROWS.iter().map(|row| row[index].to_string()).collect();
assert_eq!(analyze_column(name, &values).data_type, expected, "{name}");
}
}