use std::collections::HashMap;
use std::sync::{Arc, LazyLock};
use std::time::SystemTime;
use arrow::array::{
Array, ArrayRef, RecordBatch, StringArray, StringBuilder, TimestampMicrosecondArray,
TimestampMicrosecondBuilder, UInt32Array, UInt32Builder, UInt64Array, UInt64Builder,
};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
use serde::{Deserialize, Serialize};
use graphforge_core::{GfError, TypeId};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct RuntimeTypeId(pub u32);
const RUNTIME_RELATION_TYPE_TAG: u32 = 1 << 31;
#[must_use]
pub fn runtime_relation_type_id(id: RuntimeTypeId) -> TypeId {
assert!(
id.0 < RUNTIME_RELATION_TYPE_TAG,
"runtime relation type ID exceeds the plan encoding range"
);
TypeId(id.0 | RUNTIME_RELATION_TYPE_TAG)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct RuntimePropId(pub u32);
pub static RUNTIME_CATALOG_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
Arc::new(Schema::new(vec![
Field::new("entry_kind", DataType::Utf8, false),
Field::new("name", DataType::Utf8, false),
Field::new("runtime_id", DataType::UInt32, false),
Field::new("observation_count", DataType::UInt64, false),
Field::new(
"first_seen",
DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
false,
),
Field::new(
"last_seen",
DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
false,
),
Field::new("owner_label", DataType::Utf8, true),
]))
});
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum EntryKind {
EntityType,
RelationType,
Property,
}
#[derive(Debug, Clone)]
struct CatalogEntry {
kind: EntryKind,
name: String,
runtime_id: u32,
observation_count: u64,
first_seen: i64,
last_seen: i64,
owner_label: Option<String>,
}
fn now_micros() -> i64 {
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map_or(0, |d| i64::try_from(d.as_micros()).unwrap_or(i64::MAX))
}
#[derive(Debug, Clone, Default)]
pub struct RuntimeCatalog {
entity_types: HashMap<String, usize>,
relation_types: HashMap<String, usize>,
properties: HashMap<(String, Option<String>), usize>,
entries: Vec<CatalogEntry>,
next_type_id: u32,
next_prop_id: u32,
}
impl RuntimeCatalog {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn intern_label(&mut self, name: &str) -> RuntimeTypeId {
let now = now_micros();
if let Some(&idx) = self.entity_types.get(name) {
let entry = &mut self.entries[idx];
entry.observation_count += 1;
entry.last_seen = now;
return RuntimeTypeId(entry.runtime_id);
}
let id = self.next_type_id;
self.next_type_id += 1;
let idx = self.entries.len();
self.entries.push(CatalogEntry {
kind: EntryKind::EntityType,
name: name.to_owned(),
runtime_id: id,
observation_count: 1,
first_seen: now,
last_seen: now,
owner_label: None,
});
self.entity_types.insert(name.to_owned(), idx);
RuntimeTypeId(id)
}
pub fn intern_relation_type(&mut self, name: &str) -> RuntimeTypeId {
let now = now_micros();
if let Some(&idx) = self.relation_types.get(name) {
let entry = &mut self.entries[idx];
entry.observation_count += 1;
entry.last_seen = now;
return RuntimeTypeId(entry.runtime_id);
}
let id = self.next_type_id;
self.next_type_id += 1;
let idx = self.entries.len();
self.entries.push(CatalogEntry {
kind: EntryKind::RelationType,
name: name.to_owned(),
runtime_id: id,
observation_count: 1,
first_seen: now,
last_seen: now,
owner_label: None,
});
self.relation_types.insert(name.to_owned(), idx);
RuntimeTypeId(id)
}
pub fn intern_property(&mut self, name: &str, owner_label: Option<&str>) -> RuntimePropId {
let now = now_micros();
let key = (name.to_owned(), owner_label.map(str::to_owned));
if let Some(&idx) = self.properties.get(&key) {
let entry = &mut self.entries[idx];
entry.observation_count += 1;
entry.last_seen = now;
return RuntimePropId(entry.runtime_id);
}
let id = self.next_prop_id;
self.next_prop_id += 1;
let idx = self.entries.len();
self.entries.push(CatalogEntry {
kind: EntryKind::Property,
name: name.to_owned(),
runtime_id: id,
observation_count: 1,
first_seen: now,
last_seen: now,
owner_label: owner_label.map(str::to_owned),
});
self.properties.insert(key, idx);
RuntimePropId(id)
}
#[must_use]
pub fn contains_entity_type(&self, name: &str) -> bool {
self.entity_types.contains_key(name)
}
#[must_use]
pub fn entity_types(&self) -> Vec<&str> {
self.entity_types.keys().map(String::as_str).collect()
}
#[must_use]
pub fn relation_types(&self) -> Vec<&str> {
self.relation_types.keys().map(String::as_str).collect()
}
#[must_use]
pub fn properties_for(&self, label: &str) -> Vec<&str> {
self.properties
.iter()
.filter(|((_, owner), _)| owner.as_deref() == Some(label))
.map(|((name, _), _)| name.as_str())
.collect()
}
#[must_use]
pub fn property_name(&self, id: RuntimePropId) -> Option<&str> {
self.entries
.iter()
.find(|e| e.kind == EntryKind::Property && e.runtime_id == id.0)
.map(|e| e.name.as_str())
}
pub fn property_names(&self) -> impl Iterator<Item = (RuntimePropId, &str)> + '_ {
self.entries
.iter()
.filter(|e| e.kind == EntryKind::Property)
.map(|e| (RuntimePropId(e.runtime_id), e.name.as_str()))
}
#[must_use]
pub fn relation_type_name(&self, id: RuntimeTypeId) -> Option<&str> {
self.entries
.iter()
.find(|e| e.kind == EntryKind::RelationType && e.runtime_id == id.0)
.map(|e| e.name.as_str())
}
pub fn relation_type_names_with_ids(&self) -> impl Iterator<Item = (RuntimeTypeId, &str)> + '_ {
self.entries
.iter()
.filter(|e| e.kind == EntryKind::RelationType)
.map(|e| (RuntimeTypeId(e.runtime_id), e.name.as_str()))
}
#[must_use]
pub fn entity_type_name(&self, id: RuntimeTypeId) -> Option<&str> {
self.entries
.iter()
.find(|e| e.kind == EntryKind::EntityType && e.runtime_id == id.0)
.map(|e| e.name.as_str())
}
pub fn entity_type_names_with_ids(&self) -> impl Iterator<Item = (RuntimeTypeId, &str)> + '_ {
self.entries
.iter()
.filter(|e| e.kind == EntryKind::EntityType)
.map(|e| (RuntimeTypeId(e.runtime_id), e.name.as_str()))
}
#[must_use]
pub fn to_record_batch(&self) -> RecordBatch {
let n = self.entries.len();
let mut kind_b = StringBuilder::with_capacity(n, n * 12);
let mut name_b = StringBuilder::with_capacity(n, n * 32);
let mut id_b = UInt32Builder::with_capacity(n);
let mut count_b = UInt64Builder::with_capacity(n);
let mut first_b = TimestampMicrosecondBuilder::with_capacity(n);
let mut last_b = TimestampMicrosecondBuilder::with_capacity(n);
let mut owner_b = StringBuilder::with_capacity(n, n * 16);
for entry in &self.entries {
kind_b.append_value(match entry.kind {
EntryKind::EntityType => "entity_type",
EntryKind::RelationType => "relation_type",
EntryKind::Property => "property",
});
name_b.append_value(&entry.name);
id_b.append_value(entry.runtime_id);
count_b.append_value(entry.observation_count);
first_b.append_value(entry.first_seen);
last_b.append_value(entry.last_seen);
match &entry.owner_label {
Some(label) => owner_b.append_value(label),
None => owner_b.append_null(),
}
}
let first_arr = first_b.finish().with_timezone_opt(Some(Arc::from("UTC")));
let last_arr = last_b.finish().with_timezone_opt(Some(Arc::from("UTC")));
let columns: Vec<ArrayRef> = vec![
Arc::new(kind_b.finish()),
Arc::new(name_b.finish()),
Arc::new(id_b.finish()),
Arc::new(count_b.finish()),
Arc::new(first_arr),
Arc::new(last_arr),
Arc::new(owner_b.finish()),
];
RecordBatch::try_new(RUNTIME_CATALOG_SCHEMA.clone(), columns)
.expect("schema and array lengths must be consistent")
}
pub fn from_record_batch(batch: &RecordBatch) -> Result<Self, GfError> {
let storage_err = |msg: &str| GfError::Storage(msg.to_owned());
let kinds = batch
.column(0)
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| storage_err("runtime_catalog col 0 (entry_kind) not Utf8"))?;
let names = batch
.column(1)
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| storage_err("runtime_catalog col 1 (name) not Utf8"))?;
let ids = batch
.column(2)
.as_any()
.downcast_ref::<UInt32Array>()
.ok_or_else(|| storage_err("runtime_catalog col 2 (runtime_id) not UInt32"))?;
let counts = batch
.column(3)
.as_any()
.downcast_ref::<UInt64Array>()
.ok_or_else(|| storage_err("runtime_catalog col 3 (observation_count) not UInt64"))?;
let first_seens = batch
.column(4)
.as_any()
.downcast_ref::<TimestampMicrosecondArray>()
.ok_or_else(|| {
storage_err("runtime_catalog col 4 (first_seen) not TimestampMicrosecond")
})?;
let last_seens = batch
.column(5)
.as_any()
.downcast_ref::<TimestampMicrosecondArray>()
.ok_or_else(|| {
storage_err("runtime_catalog col 5 (last_seen) not TimestampMicrosecond")
})?;
let owners = batch
.column(6)
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| storage_err("runtime_catalog col 6 (owner_label) not Utf8"))?;
let mut catalog = Self::new();
let mut max_type_id: u32 = 0;
let mut max_prop_id: u32 = 0;
for row in 0..batch.num_rows() {
let kind = match kinds.value(row) {
"entity_type" => EntryKind::EntityType,
"relation_type" => EntryKind::RelationType,
"property" => EntryKind::Property,
other => {
return Err(GfError::Storage(format!(
"runtime_catalog: unknown entry_kind '{other}'"
)));
}
};
let name = names.value(row).to_owned();
let runtime_id = ids.value(row);
let observation_count = counts.value(row);
let first_seen = first_seens.value(row);
let last_seen = last_seens.value(row);
let owner_label = if owners.is_null(row) {
None
} else {
Some(owners.value(row).to_owned())
};
let idx = catalog.entries.len();
catalog.entries.push(CatalogEntry {
kind,
name: name.clone(),
runtime_id,
observation_count,
first_seen,
last_seen,
owner_label: owner_label.clone(),
});
match kind {
EntryKind::EntityType => {
catalog.entity_types.insert(name, idx);
max_type_id = max_type_id.max(runtime_id + 1);
}
EntryKind::RelationType => {
catalog.relation_types.insert(name, idx);
max_type_id = max_type_id.max(runtime_id + 1);
}
EntryKind::Property => {
catalog.properties.insert((name, owner_label), idx);
max_prop_id = max_prop_id.max(runtime_id + 1);
}
}
}
catalog.next_type_id = max_type_id;
catalog.next_prop_id = max_prop_id;
Ok(catalog)
}
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use super::*;
fn make_catalog() -> RuntimeCatalog {
let mut cat = RuntimeCatalog::new();
cat.intern_label("Person");
cat.intern_label("Company");
cat.intern_relation_type("KNOWS");
cat.intern_property("name", Some("Person"));
cat.intern_property("founded", Some("Company"));
cat
}
#[test]
fn intern_label_same_id_for_same_name() {
let mut cat = RuntimeCatalog::new();
let id1 = cat.intern_label("Person");
let id2 = cat.intern_label("Person");
assert_eq!(id1, id2);
}
#[test]
fn property_name_reverse_lookup() {
let mut cat = RuntimeCatalog::new();
let name_id = cat.intern_property("name", Some("Person"));
let founded_id = cat.intern_property("founded", Some("Company"));
assert_eq!(cat.property_name(name_id), Some("name"));
assert_eq!(cat.property_name(founded_id), Some("founded"));
assert_eq!(cat.property_name(RuntimePropId(9999)), None);
}
#[test]
fn property_names_lists_all_properties() {
let cat = make_catalog();
let names: HashSet<&str> = cat.property_names().map(|(_, n)| n).collect();
assert_eq!(names, HashSet::from(["name", "founded"]));
}
#[test]
fn intern_label_distinct_ids_for_distinct_names() {
let mut cat = RuntimeCatalog::new();
let person = cat.intern_label("Person");
let company = cat.intern_label("Company");
assert_ne!(person, company);
}
#[test]
fn intern_relation_type_same_id() {
let mut cat = RuntimeCatalog::new();
let id1 = cat.intern_relation_type("KNOWS");
let id2 = cat.intern_relation_type("KNOWS");
assert_eq!(id1, id2);
}
#[test]
fn intern_relation_type_distinct_from_entity_type() {
let mut cat = RuntimeCatalog::new();
let person_id = cat.intern_label("Person");
let knows_id = cat.intern_relation_type("KNOWS");
assert_ne!(person_id, knows_id);
}
#[test]
fn type_id_and_prop_id_are_independent() {
let mut cat = RuntimeCatalog::new();
let type_id = cat.intern_label("Person");
let prop_id = cat.intern_property("name", Some("Person"));
assert_eq!(type_id.0, 0);
assert_eq!(prop_id.0, 0);
}
#[test]
fn contains_entity_type() {
let mut cat = RuntimeCatalog::new();
cat.intern_label("Person");
assert!(cat.contains_entity_type("Person"));
assert!(!cat.contains_entity_type("Unknown"));
}
#[test]
fn entity_types_returns_all_interned() {
let cat = make_catalog();
let types: HashSet<&str> = cat.entity_types().into_iter().collect();
assert!(types.contains("Person"));
assert!(types.contains("Company"));
assert_eq!(types.len(), 2);
}
#[test]
fn properties_for_returns_correct_props() {
let cat = make_catalog();
let props: HashSet<&str> = cat.properties_for("Person").into_iter().collect();
assert!(props.contains("name"));
assert!(!props.contains("founded"));
}
#[test]
fn observation_count_increments() {
let mut cat = RuntimeCatalog::new();
cat.intern_label("Person");
cat.intern_label("Person");
cat.intern_label("Person");
let batch = cat.to_record_batch();
let counts = batch
.column(3)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
assert_eq!(counts.value(0), 3);
}
#[test]
fn empty_catalog_to_record_batch_has_correct_schema() {
let cat = RuntimeCatalog::new();
let batch = cat.to_record_batch();
assert_eq!(batch.num_rows(), 0);
assert_eq!(batch.schema(), *RUNTIME_CATALOG_SCHEMA);
}
#[test]
fn roundtrip_to_from_record_batch() {
let mut cat = make_catalog();
cat.intern_label("Person");
cat.intern_label("Person");
let batch = cat.to_record_batch();
let restored = RuntimeCatalog::from_record_batch(&batch).unwrap();
let et: HashSet<&str> = restored.entity_types().into_iter().collect();
assert!(et.contains("Person"));
assert!(et.contains("Company"));
let rt: HashSet<&str> = restored.relation_types().into_iter().collect();
assert!(rt.contains("KNOWS"));
let pp: HashSet<&str> = restored.properties_for("Person").into_iter().collect();
assert!(pp.contains("name"));
let person_id_orig = cat.intern_label("Person");
let person_id_rest = {
let mut r = restored.clone();
r.intern_label("Person")
};
assert_eq!(person_id_orig, person_id_rest);
let restored_batch = restored.to_record_batch();
let kinds = restored_batch
.column(0)
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
let counts = restored_batch
.column(3)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let person_row = (0..restored_batch.num_rows())
.find(|&r| kinds.value(r) == "entity_type")
.unwrap();
assert_eq!(counts.value(person_row), 3);
}
#[test]
fn roundtrip_empty_catalog() {
let cat = RuntimeCatalog::new();
let batch = cat.to_record_batch();
let restored = RuntimeCatalog::from_record_batch(&batch).unwrap();
assert_eq!(restored.entity_types().len(), 0);
assert_eq!(restored.relation_types().len(), 0);
}
#[test]
fn from_record_batch_unknown_entry_kind_returns_error() {
let mut cat = RuntimeCatalog::new();
cat.intern_label("X");
let good_batch = cat.to_record_batch();
assert_eq!(good_batch.num_rows(), 1);
let bad_kinds = Arc::new(StringArray::from(vec!["bogus_kind"])) as ArrayRef;
let mut cols: Vec<ArrayRef> = good_batch.columns().to_vec();
cols[0] = bad_kinds;
let batch = RecordBatch::try_new(RUNTIME_CATALOG_SCHEMA.clone(), cols).unwrap();
let result = RuntimeCatalog::from_record_batch(&batch);
assert!(
matches!(result, Err(GfError::Storage(_))),
"expected Storage error for unknown entry_kind"
);
}
}