Skip to main content

graphforge_ir/
catalog.rs

1//! [`RuntimeCatalog`] — auto-growing registry of observed entity types,
2//! relation types, and property names for GraphForge's progressive ontology model.
3//!
4//! In `exploratory` and `advisory` modes the binder delegates unknown label/type
5//! resolution here.  The catalog auto-assigns integer IDs and records observations
6//! so that cross-session stability is achieved by persisting to
7//! `topology/runtime_catalog.parquet`.
8
9use 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// ---------------------------------------------------------------------------
23// Public newtypes
24// ---------------------------------------------------------------------------
25
26/// Identifies an entity or relation type discovered at runtime.
27///
28/// Runtime IDs are local to a session (or persisted in
29/// `topology/runtime_catalog.parquet` for cross-session stability).  They are
30/// distinct from ontology [`TypeId`](graphforge_core::TypeId)s, which are assigned by
31/// the compiler from a formal ontology definition.
32#[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/// Encode one runtime-catalog relation ID in the plan TypeId space.
38///
39/// Ontology and runtime IDs both begin at zero. Tagging runtime relation IDs
40/// keeps advisory/exploratory misses disjoint from ontology relation IDs while
41/// the bound plan crosses the IR/lowering boundary.
42#[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/// Identifies a property discovered at runtime (not formally declared in an ontology).
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
53pub struct RuntimePropId(pub u32);
54
55// ---------------------------------------------------------------------------
56// Arrow schema
57// ---------------------------------------------------------------------------
58
59/// Arrow schema for `topology/runtime_catalog.parquet`.
60pub 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// ---------------------------------------------------------------------------
81// Private internals
82// ---------------------------------------------------------------------------
83
84#[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    /// Microseconds since Unix epoch (UTC).
98    first_seen: i64,
99    /// Microseconds since Unix epoch (UTC).
100    last_seen: i64,
101    /// For `Property` entries: the label this property was observed on.
102    owner_label: Option<String>,
103}
104
105/// Returns current time as microseconds since Unix epoch (UTC).
106fn 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// ---------------------------------------------------------------------------
113// RuntimeCatalog
114// ---------------------------------------------------------------------------
115
116/// Auto-growing registry of observed entity types, relation types, and property names.
117///
118/// In `exploratory` and `advisory` ontology modes, the binder delegates
119/// unknown label/type/property resolution here rather than returning an error.
120/// The catalog auto-assigns stable integer IDs and tracks observation statistics.
121///
122/// The catalog is not internally synchronised — callers that share it across
123/// threads must wrap it in `Arc<RwLock<RuntimeCatalog>>` (wired in issue #583).
124#[derive(Debug, Clone, Default)]
125pub struct RuntimeCatalog {
126    /// name → index into `entries`
127    entity_types: HashMap<String, usize>,
128    /// name → index into `entries`
129    relation_types: HashMap<String, usize>,
130    /// (name, owner_label) → index into `entries`
131    properties: HashMap<(String, Option<String>), usize>,
132    /// All entries in insertion order.
133    entries: Vec<CatalogEntry>,
134    /// Next ID to assign for entity types and relation types (shared space).
135    next_type_id: u32,
136    /// Next ID to assign for properties.
137    next_prop_id: u32,
138}
139
140impl RuntimeCatalog {
141    /// Creates an empty catalog.
142    #[must_use]
143    pub fn new() -> Self {
144        Self::default()
145    }
146
147    /// Interns an entity label and returns a stable [`RuntimeTypeId`].
148    ///
149    /// If `name` has been seen before the observation count is incremented and
150    /// the existing ID is returned unchanged.
151    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    /// Interns a relation type name and returns a stable [`RuntimeTypeId`].
176    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    /// Interns a property name and returns a stable [`RuntimePropId`].
201    ///
202    /// Properties are keyed by `(name, owner_label)` so the same property name
203    /// observed on different labels gets independent IDs.
204    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    /// Returns `true` if `name` has been interned as an entity type.
230    #[must_use]
231    pub fn contains_entity_type(&self, name: &str) -> bool {
232        self.entity_types.contains_key(name)
233    }
234
235    /// Returns all interned entity type names (order unspecified).
236    #[must_use]
237    pub fn entity_types(&self) -> Vec<&str> {
238        self.entity_types.keys().map(String::as_str).collect()
239    }
240
241    /// Returns all interned relation type names (order unspecified).
242    #[must_use]
243    pub fn relation_types(&self) -> Vec<&str> {
244        self.relation_types.keys().map(String::as_str).collect()
245    }
246
247    /// Returns all property names observed on `label` (order unspecified).
248    #[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    /// Resolves a [`RuntimePropId`] back to the property name it was interned
258    /// under, or `None` if no property entry carries that ID.
259    ///
260    /// Used by the relational lowering layer to turn a numeric `PropertyAccess`
261    /// ID back into the real column name when reading exploratory property
262    /// tables.
263    #[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    /// Returns `(RuntimePropId, name)` for every interned property (order
272    /// unspecified). Convenient for building a `PropId → name` map in one pass.
273    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    /// Resolves a relation-type [`RuntimeTypeId`] back to the name it was
281    /// interned under, or `None` if no relation-type entry carries that ID.
282    ///
283    /// Used by the relational lowering layer to resolve a `TypedEdgeScan`'s
284    /// relation name when reading exploratory edge tables (no ontology present).
285    #[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    /// Returns `(RuntimeTypeId, name)` for every interned relation type (order
294    /// unspecified). Convenient for building a `TypeId → relation-name` map.
295    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    /// Resolves an entity-type (node label) [`RuntimeTypeId`] back to the name
303    /// it was interned under, or `None` if no entity-type entry carries that ID.
304    ///
305    /// Mirror of [`relation_type_name`](Self::relation_type_name) for labels —
306    /// used to render a real label name for an unlabelled `MATCH (n) RETURN n`
307    /// in exploratory mode, where the ontology map is empty (#889).
308    #[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    /// Returns `(RuntimeTypeId, name)` for every interned entity type (node
317    /// label), order unspecified. Convenient for building a
318    /// `TypeId → label-name` map.
319    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    /// Serialises the catalog to an Arrow [`RecordBatch`] using [`RUNTIME_CATALOG_SCHEMA`].
327    ///
328    /// The resulting batch can be written to `topology/runtime_catalog.parquet`
329    /// and later restored via [`from_record_batch`](Self::from_record_batch).
330    #[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    /// Restores a `RuntimeCatalog` from an Arrow [`RecordBatch`] previously
376    /// produced by [`to_record_batch`](Self::to_record_batch).
377    ///
378    /// # Errors
379    /// Returns [`GfError::Storage`] if any column has the wrong type or an
380    /// unknown `entry_kind` value is encountered.
381    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// ---------------------------------------------------------------------------
484// Tests
485// ---------------------------------------------------------------------------
486
487#[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        // An unknown ID resolves to None.
519        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        // Entity "Person" (id=0) and relation "KNOWS" (id=1) share the type ID
548        // counter, so they get different IDs.
549        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        // First entity type gets type ID 0.
559        let type_id = cat.intern_label("Person");
560        // First property gets prop ID 0 — independent counter.
561        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        // Intern "Person" twice more so observation_count = 3.
618        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        // Entity types and relation types preserved.
625        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        // Properties preserved with owner.
633        let pp: HashSet<&str> = restored.properties_for("Person").into_iter().collect();
634        assert!(pp.contains("name"));
635
636        // IDs are preserved (stable after round-trip).
637        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        // Observation count for "Person" should be 3 (original 1 + 2 extra interns).
645        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        // Build a single-row catalog with a valid entry first, then manually
674        // produce a batch via to_record_batch and overwrite the entry_kind column
675        // with an invalid value.  We do this by constructing a fresh batch from
676        // scratch using the schema-aware builders.
677        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        // Replace the entry_kind column with an invalid value while keeping all
683        // other columns (and their types) intact.
684        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}