1use std::collections::HashMap;
10use std::sync::{Arc, LazyLock};
11use std::time::SystemTime;
12
13use arrow::array::{
14 Array, ArrayRef, RecordBatch, StringArray, StringBuilder, TimestampMicrosecondArray,
15 TimestampMicrosecondBuilder, UInt32Array, UInt32Builder, UInt64Array, UInt64Builder,
16};
17use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
18use serde::{Deserialize, Serialize};
19
20use graphforge_core::{GfError, TypeId};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
33pub struct RuntimeTypeId(pub u32);
34
35const RUNTIME_RELATION_TYPE_TAG: u32 = 1 << 31;
36
37#[must_use]
43pub fn runtime_relation_type_id(id: RuntimeTypeId) -> TypeId {
44 assert!(
45 id.0 < RUNTIME_RELATION_TYPE_TAG,
46 "runtime relation type ID exceeds the plan encoding range"
47 );
48 TypeId(id.0 | RUNTIME_RELATION_TYPE_TAG)
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
53pub struct RuntimePropId(pub u32);
54
55pub static RUNTIME_CATALOG_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
61 Arc::new(Schema::new(vec![
62 Field::new("entry_kind", DataType::Utf8, false),
63 Field::new("name", DataType::Utf8, false),
64 Field::new("runtime_id", DataType::UInt32, false),
65 Field::new("observation_count", DataType::UInt64, false),
66 Field::new(
67 "first_seen",
68 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
69 false,
70 ),
71 Field::new(
72 "last_seen",
73 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
74 false,
75 ),
76 Field::new("owner_label", DataType::Utf8, true),
77 ]))
78});
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85enum EntryKind {
86 EntityType,
87 RelationType,
88 Property,
89}
90
91#[derive(Debug, Clone)]
92struct CatalogEntry {
93 kind: EntryKind,
94 name: String,
95 runtime_id: u32,
96 observation_count: u64,
97 first_seen: i64,
99 last_seen: i64,
101 owner_label: Option<String>,
103}
104
105fn now_micros() -> i64 {
107 SystemTime::now()
108 .duration_since(SystemTime::UNIX_EPOCH)
109 .map_or(0, |d| i64::try_from(d.as_micros()).unwrap_or(i64::MAX))
110}
111
112#[derive(Debug, Clone, Default)]
125pub struct RuntimeCatalog {
126 entity_types: HashMap<String, usize>,
128 relation_types: HashMap<String, usize>,
130 properties: HashMap<(String, Option<String>), usize>,
132 entries: Vec<CatalogEntry>,
134 next_type_id: u32,
136 next_prop_id: u32,
138}
139
140impl RuntimeCatalog {
141 #[must_use]
143 pub fn new() -> Self {
144 Self::default()
145 }
146
147 pub fn intern_label(&mut self, name: &str) -> RuntimeTypeId {
152 let now = now_micros();
153 if let Some(&idx) = self.entity_types.get(name) {
154 let entry = &mut self.entries[idx];
155 entry.observation_count += 1;
156 entry.last_seen = now;
157 return RuntimeTypeId(entry.runtime_id);
158 }
159 let id = self.next_type_id;
160 self.next_type_id += 1;
161 let idx = self.entries.len();
162 self.entries.push(CatalogEntry {
163 kind: EntryKind::EntityType,
164 name: name.to_owned(),
165 runtime_id: id,
166 observation_count: 1,
167 first_seen: now,
168 last_seen: now,
169 owner_label: None,
170 });
171 self.entity_types.insert(name.to_owned(), idx);
172 RuntimeTypeId(id)
173 }
174
175 pub fn intern_relation_type(&mut self, name: &str) -> RuntimeTypeId {
177 let now = now_micros();
178 if let Some(&idx) = self.relation_types.get(name) {
179 let entry = &mut self.entries[idx];
180 entry.observation_count += 1;
181 entry.last_seen = now;
182 return RuntimeTypeId(entry.runtime_id);
183 }
184 let id = self.next_type_id;
185 self.next_type_id += 1;
186 let idx = self.entries.len();
187 self.entries.push(CatalogEntry {
188 kind: EntryKind::RelationType,
189 name: name.to_owned(),
190 runtime_id: id,
191 observation_count: 1,
192 first_seen: now,
193 last_seen: now,
194 owner_label: None,
195 });
196 self.relation_types.insert(name.to_owned(), idx);
197 RuntimeTypeId(id)
198 }
199
200 pub fn intern_property(&mut self, name: &str, owner_label: Option<&str>) -> RuntimePropId {
205 let now = now_micros();
206 let key = (name.to_owned(), owner_label.map(str::to_owned));
207 if let Some(&idx) = self.properties.get(&key) {
208 let entry = &mut self.entries[idx];
209 entry.observation_count += 1;
210 entry.last_seen = now;
211 return RuntimePropId(entry.runtime_id);
212 }
213 let id = self.next_prop_id;
214 self.next_prop_id += 1;
215 let idx = self.entries.len();
216 self.entries.push(CatalogEntry {
217 kind: EntryKind::Property,
218 name: name.to_owned(),
219 runtime_id: id,
220 observation_count: 1,
221 first_seen: now,
222 last_seen: now,
223 owner_label: owner_label.map(str::to_owned),
224 });
225 self.properties.insert(key, idx);
226 RuntimePropId(id)
227 }
228
229 #[must_use]
231 pub fn contains_entity_type(&self, name: &str) -> bool {
232 self.entity_types.contains_key(name)
233 }
234
235 #[must_use]
237 pub fn entity_types(&self) -> Vec<&str> {
238 self.entity_types.keys().map(String::as_str).collect()
239 }
240
241 #[must_use]
243 pub fn relation_types(&self) -> Vec<&str> {
244 self.relation_types.keys().map(String::as_str).collect()
245 }
246
247 #[must_use]
249 pub fn properties_for(&self, label: &str) -> Vec<&str> {
250 self.properties
251 .iter()
252 .filter(|((_, owner), _)| owner.as_deref() == Some(label))
253 .map(|((name, _), _)| name.as_str())
254 .collect()
255 }
256
257 #[must_use]
264 pub fn property_name(&self, id: RuntimePropId) -> Option<&str> {
265 self.entries
266 .iter()
267 .find(|e| e.kind == EntryKind::Property && e.runtime_id == id.0)
268 .map(|e| e.name.as_str())
269 }
270
271 pub fn property_names(&self) -> impl Iterator<Item = (RuntimePropId, &str)> + '_ {
274 self.entries
275 .iter()
276 .filter(|e| e.kind == EntryKind::Property)
277 .map(|e| (RuntimePropId(e.runtime_id), e.name.as_str()))
278 }
279
280 #[must_use]
286 pub fn relation_type_name(&self, id: RuntimeTypeId) -> Option<&str> {
287 self.entries
288 .iter()
289 .find(|e| e.kind == EntryKind::RelationType && e.runtime_id == id.0)
290 .map(|e| e.name.as_str())
291 }
292
293 pub fn relation_type_names_with_ids(&self) -> impl Iterator<Item = (RuntimeTypeId, &str)> + '_ {
296 self.entries
297 .iter()
298 .filter(|e| e.kind == EntryKind::RelationType)
299 .map(|e| (RuntimeTypeId(e.runtime_id), e.name.as_str()))
300 }
301
302 #[must_use]
309 pub fn entity_type_name(&self, id: RuntimeTypeId) -> Option<&str> {
310 self.entries
311 .iter()
312 .find(|e| e.kind == EntryKind::EntityType && e.runtime_id == id.0)
313 .map(|e| e.name.as_str())
314 }
315
316 pub fn entity_type_names_with_ids(&self) -> impl Iterator<Item = (RuntimeTypeId, &str)> + '_ {
320 self.entries
321 .iter()
322 .filter(|e| e.kind == EntryKind::EntityType)
323 .map(|e| (RuntimeTypeId(e.runtime_id), e.name.as_str()))
324 }
325
326 #[must_use]
331 pub fn to_record_batch(&self) -> RecordBatch {
332 let n = self.entries.len();
333 let mut kind_b = StringBuilder::with_capacity(n, n * 12);
334 let mut name_b = StringBuilder::with_capacity(n, n * 32);
335 let mut id_b = UInt32Builder::with_capacity(n);
336 let mut count_b = UInt64Builder::with_capacity(n);
337 let mut first_b = TimestampMicrosecondBuilder::with_capacity(n);
338 let mut last_b = TimestampMicrosecondBuilder::with_capacity(n);
339 let mut owner_b = StringBuilder::with_capacity(n, n * 16);
340
341 for entry in &self.entries {
342 kind_b.append_value(match entry.kind {
343 EntryKind::EntityType => "entity_type",
344 EntryKind::RelationType => "relation_type",
345 EntryKind::Property => "property",
346 });
347 name_b.append_value(&entry.name);
348 id_b.append_value(entry.runtime_id);
349 count_b.append_value(entry.observation_count);
350 first_b.append_value(entry.first_seen);
351 last_b.append_value(entry.last_seen);
352 match &entry.owner_label {
353 Some(label) => owner_b.append_value(label),
354 None => owner_b.append_null(),
355 }
356 }
357
358 let first_arr = first_b.finish().with_timezone_opt(Some(Arc::from("UTC")));
359 let last_arr = last_b.finish().with_timezone_opt(Some(Arc::from("UTC")));
360
361 let columns: Vec<ArrayRef> = vec![
362 Arc::new(kind_b.finish()),
363 Arc::new(name_b.finish()),
364 Arc::new(id_b.finish()),
365 Arc::new(count_b.finish()),
366 Arc::new(first_arr),
367 Arc::new(last_arr),
368 Arc::new(owner_b.finish()),
369 ];
370
371 RecordBatch::try_new(RUNTIME_CATALOG_SCHEMA.clone(), columns)
372 .expect("schema and array lengths must be consistent")
373 }
374
375 pub fn from_record_batch(batch: &RecordBatch) -> Result<Self, GfError> {
382 let storage_err = |msg: &str| GfError::Storage(msg.to_owned());
383
384 let kinds = batch
385 .column(0)
386 .as_any()
387 .downcast_ref::<StringArray>()
388 .ok_or_else(|| storage_err("runtime_catalog col 0 (entry_kind) not Utf8"))?;
389 let names = batch
390 .column(1)
391 .as_any()
392 .downcast_ref::<StringArray>()
393 .ok_or_else(|| storage_err("runtime_catalog col 1 (name) not Utf8"))?;
394 let ids = batch
395 .column(2)
396 .as_any()
397 .downcast_ref::<UInt32Array>()
398 .ok_or_else(|| storage_err("runtime_catalog col 2 (runtime_id) not UInt32"))?;
399 let counts = batch
400 .column(3)
401 .as_any()
402 .downcast_ref::<UInt64Array>()
403 .ok_or_else(|| storage_err("runtime_catalog col 3 (observation_count) not UInt64"))?;
404 let first_seens = batch
405 .column(4)
406 .as_any()
407 .downcast_ref::<TimestampMicrosecondArray>()
408 .ok_or_else(|| {
409 storage_err("runtime_catalog col 4 (first_seen) not TimestampMicrosecond")
410 })?;
411 let last_seens = batch
412 .column(5)
413 .as_any()
414 .downcast_ref::<TimestampMicrosecondArray>()
415 .ok_or_else(|| {
416 storage_err("runtime_catalog col 5 (last_seen) not TimestampMicrosecond")
417 })?;
418 let owners = batch
419 .column(6)
420 .as_any()
421 .downcast_ref::<StringArray>()
422 .ok_or_else(|| storage_err("runtime_catalog col 6 (owner_label) not Utf8"))?;
423
424 let mut catalog = Self::new();
425 let mut max_type_id: u32 = 0;
426 let mut max_prop_id: u32 = 0;
427
428 for row in 0..batch.num_rows() {
429 let kind = match kinds.value(row) {
430 "entity_type" => EntryKind::EntityType,
431 "relation_type" => EntryKind::RelationType,
432 "property" => EntryKind::Property,
433 other => {
434 return Err(GfError::Storage(format!(
435 "runtime_catalog: unknown entry_kind '{other}'"
436 )));
437 }
438 };
439 let name = names.value(row).to_owned();
440 let runtime_id = ids.value(row);
441 let observation_count = counts.value(row);
442 let first_seen = first_seens.value(row);
443 let last_seen = last_seens.value(row);
444 let owner_label = if owners.is_null(row) {
445 None
446 } else {
447 Some(owners.value(row).to_owned())
448 };
449
450 let idx = catalog.entries.len();
451 catalog.entries.push(CatalogEntry {
452 kind,
453 name: name.clone(),
454 runtime_id,
455 observation_count,
456 first_seen,
457 last_seen,
458 owner_label: owner_label.clone(),
459 });
460
461 match kind {
462 EntryKind::EntityType => {
463 catalog.entity_types.insert(name, idx);
464 max_type_id = max_type_id.max(runtime_id + 1);
465 }
466 EntryKind::RelationType => {
467 catalog.relation_types.insert(name, idx);
468 max_type_id = max_type_id.max(runtime_id + 1);
469 }
470 EntryKind::Property => {
471 catalog.properties.insert((name, owner_label), idx);
472 max_prop_id = max_prop_id.max(runtime_id + 1);
473 }
474 }
475 }
476
477 catalog.next_type_id = max_type_id;
478 catalog.next_prop_id = max_prop_id;
479 Ok(catalog)
480 }
481}
482
483#[cfg(test)]
488mod tests {
489 use std::collections::HashSet;
490
491 use super::*;
492
493 fn make_catalog() -> RuntimeCatalog {
494 let mut cat = RuntimeCatalog::new();
495 cat.intern_label("Person");
496 cat.intern_label("Company");
497 cat.intern_relation_type("KNOWS");
498 cat.intern_property("name", Some("Person"));
499 cat.intern_property("founded", Some("Company"));
500 cat
501 }
502
503 #[test]
504 fn intern_label_same_id_for_same_name() {
505 let mut cat = RuntimeCatalog::new();
506 let id1 = cat.intern_label("Person");
507 let id2 = cat.intern_label("Person");
508 assert_eq!(id1, id2);
509 }
510
511 #[test]
512 fn property_name_reverse_lookup() {
513 let mut cat = RuntimeCatalog::new();
514 let name_id = cat.intern_property("name", Some("Person"));
515 let founded_id = cat.intern_property("founded", Some("Company"));
516 assert_eq!(cat.property_name(name_id), Some("name"));
517 assert_eq!(cat.property_name(founded_id), Some("founded"));
518 assert_eq!(cat.property_name(RuntimePropId(9999)), None);
520 }
521
522 #[test]
523 fn property_names_lists_all_properties() {
524 let cat = make_catalog();
525 let names: HashSet<&str> = cat.property_names().map(|(_, n)| n).collect();
526 assert_eq!(names, HashSet::from(["name", "founded"]));
527 }
528
529 #[test]
530 fn intern_label_distinct_ids_for_distinct_names() {
531 let mut cat = RuntimeCatalog::new();
532 let person = cat.intern_label("Person");
533 let company = cat.intern_label("Company");
534 assert_ne!(person, company);
535 }
536
537 #[test]
538 fn intern_relation_type_same_id() {
539 let mut cat = RuntimeCatalog::new();
540 let id1 = cat.intern_relation_type("KNOWS");
541 let id2 = cat.intern_relation_type("KNOWS");
542 assert_eq!(id1, id2);
543 }
544
545 #[test]
546 fn intern_relation_type_distinct_from_entity_type() {
547 let mut cat = RuntimeCatalog::new();
550 let person_id = cat.intern_label("Person");
551 let knows_id = cat.intern_relation_type("KNOWS");
552 assert_ne!(person_id, knows_id);
553 }
554
555 #[test]
556 fn type_id_and_prop_id_are_independent() {
557 let mut cat = RuntimeCatalog::new();
558 let type_id = cat.intern_label("Person");
560 let prop_id = cat.intern_property("name", Some("Person"));
562 assert_eq!(type_id.0, 0);
563 assert_eq!(prop_id.0, 0);
564 }
565
566 #[test]
567 fn contains_entity_type() {
568 let mut cat = RuntimeCatalog::new();
569 cat.intern_label("Person");
570 assert!(cat.contains_entity_type("Person"));
571 assert!(!cat.contains_entity_type("Unknown"));
572 }
573
574 #[test]
575 fn entity_types_returns_all_interned() {
576 let cat = make_catalog();
577 let types: HashSet<&str> = cat.entity_types().into_iter().collect();
578 assert!(types.contains("Person"));
579 assert!(types.contains("Company"));
580 assert_eq!(types.len(), 2);
581 }
582
583 #[test]
584 fn properties_for_returns_correct_props() {
585 let cat = make_catalog();
586 let props: HashSet<&str> = cat.properties_for("Person").into_iter().collect();
587 assert!(props.contains("name"));
588 assert!(!props.contains("founded"));
589 }
590
591 #[test]
592 fn observation_count_increments() {
593 let mut cat = RuntimeCatalog::new();
594 cat.intern_label("Person");
595 cat.intern_label("Person");
596 cat.intern_label("Person");
597 let batch = cat.to_record_batch();
598 let counts = batch
599 .column(3)
600 .as_any()
601 .downcast_ref::<UInt64Array>()
602 .unwrap();
603 assert_eq!(counts.value(0), 3);
604 }
605
606 #[test]
607 fn empty_catalog_to_record_batch_has_correct_schema() {
608 let cat = RuntimeCatalog::new();
609 let batch = cat.to_record_batch();
610 assert_eq!(batch.num_rows(), 0);
611 assert_eq!(batch.schema(), *RUNTIME_CATALOG_SCHEMA);
612 }
613
614 #[test]
615 fn roundtrip_to_from_record_batch() {
616 let mut cat = make_catalog();
617 cat.intern_label("Person");
619 cat.intern_label("Person");
620
621 let batch = cat.to_record_batch();
622 let restored = RuntimeCatalog::from_record_batch(&batch).unwrap();
623
624 let et: HashSet<&str> = restored.entity_types().into_iter().collect();
626 assert!(et.contains("Person"));
627 assert!(et.contains("Company"));
628
629 let rt: HashSet<&str> = restored.relation_types().into_iter().collect();
630 assert!(rt.contains("KNOWS"));
631
632 let pp: HashSet<&str> = restored.properties_for("Person").into_iter().collect();
634 assert!(pp.contains("name"));
635
636 let person_id_orig = cat.intern_label("Person");
638 let person_id_rest = {
639 let mut r = restored.clone();
640 r.intern_label("Person")
641 };
642 assert_eq!(person_id_orig, person_id_rest);
643
644 let restored_batch = restored.to_record_batch();
646 let kinds = restored_batch
647 .column(0)
648 .as_any()
649 .downcast_ref::<StringArray>()
650 .unwrap();
651 let counts = restored_batch
652 .column(3)
653 .as_any()
654 .downcast_ref::<UInt64Array>()
655 .unwrap();
656 let person_row = (0..restored_batch.num_rows())
657 .find(|&r| kinds.value(r) == "entity_type")
658 .unwrap();
659 assert_eq!(counts.value(person_row), 3);
660 }
661
662 #[test]
663 fn roundtrip_empty_catalog() {
664 let cat = RuntimeCatalog::new();
665 let batch = cat.to_record_batch();
666 let restored = RuntimeCatalog::from_record_batch(&batch).unwrap();
667 assert_eq!(restored.entity_types().len(), 0);
668 assert_eq!(restored.relation_types().len(), 0);
669 }
670
671 #[test]
672 fn from_record_batch_unknown_entry_kind_returns_error() {
673 let mut cat = RuntimeCatalog::new();
678 cat.intern_label("X");
679 let good_batch = cat.to_record_batch();
680 assert_eq!(good_batch.num_rows(), 1);
681
682 let bad_kinds = Arc::new(StringArray::from(vec!["bogus_kind"])) as ArrayRef;
685 let mut cols: Vec<ArrayRef> = good_batch.columns().to_vec();
686 cols[0] = bad_kinds;
687
688 let batch = RecordBatch::try_new(RUNTIME_CATALOG_SCHEMA.clone(), cols).unwrap();
689 let result = RuntimeCatalog::from_record_batch(&batch);
690 assert!(
691 matches!(result, Err(GfError::Storage(_))),
692 "expected Storage error for unknown entry_kind"
693 );
694 }
695}