use std::{
path::{Path, PathBuf},
sync::Arc,
};
use crate::{LadduDataError, LadduDataResult, schema::Schema};
use super::source_error;
pub(crate) struct SourceBuild<C> {
pub(crate) files: Arc<[Arc<PathBuf>]>,
pub(crate) context: C,
pub(crate) schema: Arc<Schema>,
}
pub(crate) struct SourceBuildOptions<'a> {
pub(crate) pattern: &'a str,
pub(crate) sort: bool,
pub(crate) format: &'a str,
pub(crate) explicit_schema: Option<Arc<Schema>>,
pub(crate) infer_schema: bool,
pub(crate) validate_all_files: bool,
}
pub(crate) fn resolve_source_files(
pattern: &str,
sort: bool,
format: &str,
) -> LadduDataResult<Arc<[Arc<PathBuf>]>> {
let mut files: Vec<PathBuf> = glob::glob(pattern)
.map_err(|error| source_error("resolve source glob", pattern, error))?
.collect::<Result<_, _>>()
.map_err(|error| source_error("resolve source files", pattern, error))?;
if sort {
files.sort();
}
if files.is_empty() {
return Err(source_error(
"resolve source glob",
pattern,
format!("no {format} files matched glob"),
));
}
Ok(files.into_iter().map(Arc::new).collect())
}
pub(crate) fn build_source<C, Prepare, Infer, Validate>(
options: SourceBuildOptions<'_>,
prepare: Prepare,
infer: Infer,
mut validate: Validate,
) -> LadduDataResult<SourceBuild<C>>
where
Prepare: FnOnce(&Path) -> LadduDataResult<C>,
Infer: FnOnce(&Path, &C) -> LadduDataResult<Schema>,
Validate: FnMut(&Path, &Schema, &C) -> LadduDataResult<()>,
{
let files = resolve_source_files(options.pattern, options.sort, options.format)?;
let first_path = files[0].as_ref();
let context = prepare(first_path)?;
let schema = match options.explicit_schema {
Some(schema) => schema,
None if options.infer_schema => Arc::new(infer(first_path, &context)?),
None => return Err(LadduDataError::InvalidArgument("schema required")),
};
if options.validate_all_files {
for file in files.iter() {
validate(file.as_ref(), &schema, &context)?;
}
}
Ok(SourceBuild {
files,
context,
schema,
})
}
#[cfg(test)]
mod tests {
use std::{
fs,
path::PathBuf,
sync::atomic::{AtomicUsize, Ordering},
};
use super::*;
fn temp_dir() -> PathBuf {
let path = std::env::temp_dir().join(format!(
"laddu-source-builder-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
));
fs::create_dir_all(&path).unwrap();
path
}
#[test]
fn resolves_sorted_nonempty_globs_and_rejects_empty_matches() {
let dir = temp_dir();
let second = dir.join("b.parquet");
let first = dir.join("a.parquet");
fs::write(&second, []).unwrap();
fs::write(&first, []).unwrap();
let pattern = format!("{}/*.parquet", dir.display());
let files = resolve_source_files(&pattern, true, "parquet").unwrap();
assert_eq!(files.as_ref(), &[Arc::new(first), Arc::new(second)]);
let empty = format!("{}/*.root", dir.display());
assert!(matches!(
resolve_source_files(&empty, true, "ROOT"),
Err(LadduDataError::Source(message)) if message.contains("no ROOT files")
));
fs::remove_dir_all(dir).unwrap();
}
#[test]
fn builds_explicit_or_inferred_schema_and_validates_all_files() {
let dir = temp_dir();
let first = dir.join("a.data");
let second = dir.join("b.data");
fs::write(&first, []).unwrap();
fs::write(&second, []).unwrap();
let pattern = format!("{}/*.data", dir.display());
let validations = AtomicUsize::new(0);
let explicit = Arc::new(Schema::new(std::iter::empty::<&str>(), ["x"], false).unwrap());
let built = build_source(
SourceBuildOptions {
pattern: &pattern,
sort: true,
format: "data",
explicit_schema: Some(Arc::clone(&explicit)),
infer_schema: false,
validate_all_files: true,
},
|_| Ok(()),
|_, _| unreachable!("explicit schema skips inference"),
|_, schema, _| {
assert_eq!(schema, explicit.as_ref());
validations.fetch_add(1, Ordering::Relaxed);
Ok(())
},
)
.unwrap();
assert_eq!(built.schema.as_ref(), explicit.as_ref());
assert_eq!(built.files.len(), 2);
assert_eq!(validations.load(Ordering::Relaxed), 2);
let inferred = build_source(
SourceBuildOptions {
pattern: &pattern,
sort: true,
format: "data",
explicit_schema: None,
infer_schema: true,
validate_all_files: false,
},
|_| Ok(()),
|_, _| Schema::new(["p"], std::iter::empty::<&str>(), false),
|_, _, _| unreachable!("validation is disabled"),
)
.unwrap();
assert_eq!(inferred.schema.p4s().len(), 1);
fs::remove_dir_all(dir).unwrap();
}
}