use indexmap::{IndexMap, IndexSet};
use crate::document::{Doc, Scalar as DocScalar};
use crate::error::SchemaError;
use crate::schema::{Field, FieldType, Record, Ref, Scalar, ScalarKind, Schema};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AnyFallback {
pub location: String,
pub reason: String,
}
pub fn infer(samples: &[Doc], root_name: &str) -> Result<Schema, SchemaError> {
infer_with_report(samples, root_name, false).map(|(schema, _)| schema)
}
pub fn infer_with_report(
samples: &[Doc],
root_name: &str,
allow_any: bool,
) -> Result<(Schema, Vec<AnyFallback>), SchemaError> {
if samples.is_empty() {
return Err(SchemaError::new(
"$",
"algebra.infer-no-samples",
"cannot infer a schema from zero samples",
));
}
for s in samples {
if s.root().is_leaf() {
return Err(SchemaError::new(
"$",
"algebra.infer-scalar-root",
"infer expects object (record) samples at the root",
));
}
}
let mut env: IndexMap<String, Record> = IndexMap::new();
let mut used: IndexSet<String> = IndexSet::new();
let mut fallbacks: Vec<AnyFallback> = Vec::new();
let roots: Vec<_> = samples.iter().map(Doc::root).collect();
infer_record(
&roots,
root_name,
&mut env,
&mut used,
allow_any,
&mut fallbacks,
)?;
let schema = Schema::new(Ref::new(root_name), env)?;
Ok((schema, fallbacks))
}
fn infer_record(
nodes: &[crate::document::Cursor<'_>],
name: &str,
env: &mut IndexMap<String, Record>,
used: &mut IndexSet<String>,
allow_any: bool,
fallbacks: &mut Vec<AnyFallback>,
) -> Result<(), SchemaError> {
used.insert(name.to_string());
let mut order: Vec<String> = Vec::new();
let mut seen_labels: IndexSet<String> = IndexSet::new();
for node in nodes {
for label in node.labels() {
if seen_labels.insert(label.clone()) {
order.push(label);
}
}
}
let mut children: IndexMap<String, Vec<crate::document::Cursor<'_>>> =
order.iter().map(|l| (l.clone(), Vec::new())).collect();
let mut per_sample_counts: IndexMap<String, Vec<usize>> =
order.iter().map(|l| (l.clone(), Vec::new())).collect();
for node in nodes {
let edges = node.edges().expect("root already confirmed non-leaf");
let mut counts_here: IndexMap<&str, usize> = IndexMap::new();
for (label, child) in &edges {
*counts_here.entry(label.as_str()).or_insert(0) += 1;
children.get_mut(label).unwrap().push(child.clone());
}
for label in &order {
let c = counts_here.get(label.as_str()).copied().unwrap_or(0);
per_sample_counts.get_mut(label).unwrap().push(c);
}
}
let mut fields: Vec<Field> = Vec::with_capacity(order.len());
for label in &order {
let counts = &per_sample_counts[label];
let lo = *counts.iter().min().unwrap();
let hi = *counts.iter().max().unwrap();
let (cmin, cmax) = if hi > 1 { (0, None) } else { (lo, Some(1)) };
let ty = infer_type(
&children[label],
label,
name,
env,
used,
allow_any,
fallbacks,
)?;
fields.push(Field::new(label.clone(), ty, cmin, cmax)?);
}
env.insert(name.to_string(), Record::new(fields)?);
Ok(())
}
fn infer_type(
child_nodes: &[crate::document::Cursor<'_>],
label: &str,
record_name: &str,
env: &mut IndexMap<String, Record>,
used: &mut IndexSet<String>,
allow_any: bool,
fallbacks: &mut Vec<AnyFallback>,
) -> Result<FieldType, SchemaError> {
let is_obj: Vec<bool> = child_nodes.iter().map(|c| !c.is_leaf()).collect();
if is_obj.iter().all(|&b| b) {
let rec_name = unique_name(label, used);
infer_record(child_nodes, &rec_name, env, used, allow_any, fallbacks)?;
return Ok(FieldType::Ref(Ref::new(rec_name)));
}
if is_obj.iter().any(|&b| b) {
if allow_any {
fallbacks.push(AnyFallback {
location: format!("{record_name}.{label}"),
reason: "mixes objects and values".to_string(),
});
return Ok(FieldType::Any);
}
return Err(SchemaError::new(
format!("$.{label}"),
"algebra.infer-mixed-shape",
format!("label {label:?} mixes objects and values; cannot infer one type"),
));
}
let mut names: IndexSet<&'static str> = IndexSet::new();
let mut null = false;
for c in child_nodes {
let v = c.value().expect("scalar node confirmed by is_obj check");
match v {
DocScalar::Null => null = true,
DocScalar::Bool(_) => {
names.insert("boolean");
}
DocScalar::Int(_) => {
names.insert("integer");
}
DocScalar::Float(_) => {
names.insert("number");
}
DocScalar::Str(_) => {
names.insert("string");
}
DocScalar::Date(_) => {
names.insert("date");
}
DocScalar::Time(_) => {
names.insert("time");
}
DocScalar::Datetime(_) => {
names.insert("datetime");
}
}
}
if names.contains("number") {
names.shift_remove("integer"); }
if names.is_empty() {
return Ok(FieldType::Scalar(Scalar::new(ScalarKind::String, null)));
}
if names.len() > 1 {
let mut sorted: Vec<&str> = names.into_iter().collect();
sorted.sort_unstable();
if allow_any {
fallbacks.push(AnyFallback {
location: format!("{record_name}.{label}"),
reason: format!(
"values of more than one scalar kind ({})",
sorted.join(", ")
),
});
return Ok(FieldType::Any);
}
return Err(SchemaError::new(
format!("$.{label}"),
"algebra.infer-conflicting-scalars",
format!(
"label {label:?} has values of more than one scalar ({}); cannot infer one scalar type",
sorted.join(", ")
),
));
}
let kind = ScalarKind::parse(names.iter().next().unwrap())
.expect("names only ever holds known scalar kind names, inserted above");
Ok(FieldType::Scalar(Scalar::new(kind, null)))
}
fn unique_name(base: &str, used: &mut IndexSet<String>) -> String {
let ident = identifier(base);
let name = if ident.is_empty() {
"Rec".to_string()
} else {
ident
};
let mut chars = name.chars();
let first = chars
.next()
.expect("name is never empty: identifier()'s fallback or the \"Rec\" default guarantees a first char");
let name = first.to_uppercase().collect::<String>() + chars.as_str();
let mut cand = name.clone();
let mut i = 2;
while used.contains(&cand) {
cand = format!("{name}{i}");
i += 1;
}
used.insert(cand.clone());
cand
}
fn identifier(s: &str) -> String {
let out: String = s
.chars()
.map(|c| {
if c.is_alphanumeric() || c == '_' {
c
} else {
'_'
}
})
.collect();
let trimmed = out.trim_start_matches(|c: char| c.is_ascii_digit() || c == '_');
if trimmed.is_empty() {
out
} else {
trimmed.to_string()
}
}
#[cfg(test)]
mod tests;