lance_graph/
source_catalog.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Context-free source catalog for DataFusion logical planning.
5
6use std::any::Any;
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use arrow_schema::{Schema, SchemaRef};
11use datafusion::logical_expr::TableSource;
12
13/// A minimal catalog to resolve node labels and relationship types to logical table sources.
14pub trait GraphSourceCatalog: Send + Sync {
15    fn node_source(&self, label: &str) -> Option<Arc<dyn TableSource>>;
16    fn relationship_source(&self, rel_type: &str) -> Option<Arc<dyn TableSource>>;
17}
18
19/// A simple in-memory catalog useful for tests and bootstrap wiring.
20pub struct InMemoryCatalog {
21    node_sources: HashMap<String, Arc<dyn TableSource>>,
22    rel_sources: HashMap<String, Arc<dyn TableSource>>,
23}
24
25impl InMemoryCatalog {
26    pub fn new() -> Self {
27        Self {
28            node_sources: HashMap::new(),
29            rel_sources: HashMap::new(),
30        }
31    }
32
33    pub fn with_node_source(
34        mut self,
35        label: impl Into<String>,
36        source: Arc<dyn TableSource>,
37    ) -> Self {
38        // Normalize key to lowercase for case-insensitive lookup
39        self.node_sources
40            .insert(label.into().to_lowercase(), source);
41        self
42    }
43
44    pub fn with_relationship_source(
45        mut self,
46        rel_type: impl Into<String>,
47        source: Arc<dyn TableSource>,
48    ) -> Self {
49        // Normalize key to lowercase for case-insensitive lookup
50        self.rel_sources
51            .insert(rel_type.into().to_lowercase(), source);
52        self
53    }
54}
55
56impl Default for InMemoryCatalog {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl GraphSourceCatalog for InMemoryCatalog {
63    /// Get node source with case-insensitive label lookup
64    ///
65    /// Note: Keys are stored as lowercase, so this is an O(1) operation.
66    fn node_source(&self, label: &str) -> Option<Arc<dyn TableSource>> {
67        self.node_sources.get(&label.to_lowercase()).cloned()
68    }
69
70    /// Get relationship source with case-insensitive type lookup
71    ///
72    /// Note: Keys are stored as lowercase, so this is an O(1) operation.
73    fn relationship_source(&self, rel_type: &str) -> Option<Arc<dyn TableSource>> {
74        self.rel_sources.get(&rel_type.to_lowercase()).cloned()
75    }
76}
77
78/// A trivial logical table source with a fixed schema.
79pub struct SimpleTableSource {
80    schema: SchemaRef,
81}
82
83impl SimpleTableSource {
84    pub fn new(schema: SchemaRef) -> Self {
85        Self { schema }
86    }
87    pub fn empty() -> Self {
88        Self {
89            schema: Arc::new(Schema::empty()),
90        }
91    }
92}
93
94impl TableSource for SimpleTableSource {
95    fn as_any(&self) -> &dyn Any {
96        self
97    }
98    fn schema(&self) -> SchemaRef {
99        self.schema.clone()
100    }
101}
lance_graph/source_catalog.rs

lance_graph/
source_catalog.rs
