use indexmap::IndexMap;
use crate::graph::GraphStorage;
use crate::graph::schema::{
EdgeConnection, EdgeSchema, GraphSchema, JsonType, PropertySchema, SchemaFormat, VertexSchema,
};
use crate::value::{ElementProperty, PropertyMap};
#[derive(Debug, Default, Clone)]
struct Observed {
null: bool,
boolean: bool,
integer: bool,
number: bool,
string: bool,
uuid: bool,
array: Option<Option<Box<Observed>>>,
object: Option<ObservedObject>,
}
#[derive(Debug, Default, Clone)]
struct ObservedObject {
count: usize,
keys: IndexMap<String, (usize, Observed)>,
}
impl ObservedObject {
fn observe_map<'a>(&mut self, entries: impl Iterator<Item = (&'a str, &'a ElementProperty)>) {
self.count += 1;
for (key, value) in entries {
if !is_property_value(value) {
continue;
}
let entry = self
.keys
.entry(key.to_string())
.or_insert_with(|| (0, Observed::default()));
entry.0 += 1;
entry.1.observe(value);
}
}
fn into_schema(self) -> PropertySchema {
let mut node = PropertySchema::object();
for (key, (count, observed)) in self.keys {
if count == self.count {
node.required.push(key.clone());
}
node.properties.insert(key, observed.into_schema());
}
node
}
}
fn is_property_value(value: &ElementProperty) -> bool {
crate::graph::schema::type_of(value).is_some()
}
impl Observed {
fn observe(&mut self, value: &ElementProperty) {
match value {
ElementProperty::Null => self.null = true,
ElementProperty::Boolean(_) => self.boolean = true,
ElementProperty::Number(n) if n.is_i64() => self.integer = true,
ElementProperty::Number(_) => self.number = true,
ElementProperty::String(_) | ElementProperty::Compressed(_) => self.string = true,
ElementProperty::Uuid(_) => self.uuid = true,
ElementProperty::Array(items) => {
let slot = self.array.get_or_insert(None);
for item in items {
slot.get_or_insert_with(Box::default).observe(item);
}
}
ElementProperty::Object(map) => self
.object
.get_or_insert_with(ObservedObject::default)
.observe_map(map.iter().map(|(k, v)| (k.as_str(), v))),
ElementProperty::Map(_) => {
self.object
.get_or_insert_with(ObservedObject::default)
.count += 1;
}
ElementProperty::Vertex(_) | ElementProperty::Edge(_) => {}
}
}
fn members(self) -> Vec<PropertySchema> {
let mut members = Vec::new();
if self.boolean {
members.push(PropertySchema::boolean());
}
if self.integer {
members.push(PropertySchema::integer());
}
if self.number {
members.push(PropertySchema::number());
}
if self.string {
members.push(PropertySchema::string());
}
if self.uuid {
members.push(PropertySchema::uuid());
}
if let Some(items) = self.array {
members.push(match items {
Some(items) => PropertySchema::array(items.into_schema()),
None => PropertySchema::of(JsonType::Array),
});
}
if let Some(object) = self.object {
members.push(object.into_schema());
}
if self.null {
members.push(PropertySchema::null());
}
members
}
fn into_schema(self) -> PropertySchema {
let split = self.string && self.uuid;
let members = self.members();
if members.len() == 1 {
return members.into_iter().next().unwrap_or_default();
}
if split {
return PropertySchema::any_of(members);
}
let mut node = PropertySchema {
types: Some(Vec::new()),
..PropertySchema::default()
};
for member in members {
if let Some(types) = &mut node.types {
types.extend(member.types.into_iter().flatten());
}
if member.format == Some(SchemaFormat::Uuid) {
node.format = Some(SchemaFormat::Uuid);
}
if member.items.is_some() {
node.items = member.items;
}
if !member.properties.is_empty() || !member.required.is_empty() {
node.properties = member.properties;
node.required = member.required;
}
}
node
}
}
#[derive(Debug, Default)]
pub struct SchemaInferer {
vertices: IndexMap<String, ObservedObject>,
edges: IndexMap<String, (ObservedObject, Vec<EdgeConnection>)>,
}
impl SchemaInferer {
pub fn new() -> Self {
Self::default()
}
pub fn observe_vertex(
&mut self,
labels: &[compact_str::CompactString],
properties: &PropertyMap,
) {
for label in labels {
self.vertices
.entry(label.to_string())
.or_default()
.observe_map(properties.iter().map(|(k, v)| (k.as_str(), v)));
}
}
pub fn observe_edge(
&mut self,
label: Option<&str>,
source_labels: &[compact_str::CompactString],
target_labels: &[compact_str::CompactString],
properties: &PropertyMap,
) {
let Some(label) = label.filter(|l| !l.is_empty()) else {
return;
};
let (observed, connections) = self.edges.entry(label.to_string()).or_default();
observed.observe_map(properties.iter().map(|(k, v)| (k.as_str(), v)));
if let (Some(from), Some(to)) = (source_labels.first(), target_labels.first())
&& !connections
.iter()
.any(|c| c.from == from.as_str() && c.to == to.as_str())
{
connections.push(EdgeConnection::new(from.as_str(), to.as_str()));
}
}
pub fn finish(self) -> GraphSchema {
let mut schema = GraphSchema::new();
for (label, observed) in self.vertices {
schema
.vertices
.insert(label, VertexSchema::new(observed.into_schema()));
}
for (label, (observed, connections)) in self.edges {
let mut edge = EdgeSchema::new(observed.into_schema());
edge.connections = connections;
schema.edges.insert(label, edge);
}
schema
}
pub fn infer_from_storage<S: GraphStorage>(storage: &S) -> GraphSchema {
let mut inferer = SchemaInferer::new();
let empty = PropertyMap::default();
for vertex in storage.vertices() {
inferer.observe_vertex(
storage.vertex_labels(vertex),
storage.vertex_properties(vertex).unwrap_or(&empty),
);
}
for edge in storage.edges() {
let Some((from, to)) = storage.edge_endpoints(edge) else {
continue;
};
inferer.observe_edge(
storage.edge_label(edge),
storage.vertex_labels(from),
storage.vertex_labels(to),
storage.edge_properties(edge).unwrap_or(&empty),
);
}
inferer.finish()
}
}