Skip to main content

arrow_graph_core/
graph_factory.rs

1//! Graph store factory — select and create the right store backend.
2//!
3//! # Backends
4//!
5//! | Backend | Store Type | Use Case |
6//! |---------|-----------|----------|
7//! | `InMemory` | `ArrowGraphStore` | Tests, short-lived queries |
8//! | `KnowledgeGraph` | `KgStore` | Prefix-aware KG operations |
9//! | `Simple` | `SimpleTripleStore` | Default namespace + layer |
10
11use crate::kg_store::KgStore;
12use crate::store::ArrowGraphStore;
13use crate::triple_store::SimpleTripleStore;
14
15/// Available graph store backends.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum GraphBackend {
18    /// Raw Arrow graph store — no defaults, no extras.
19    InMemory,
20    /// Simple triple store with default namespace + layer.
21    Simple,
22    /// Full knowledge graph with prefix management and keyword search.
23    KnowledgeGraph,
24}
25
26/// Configuration for graph store creation.
27#[derive(Debug, Clone)]
28pub struct GraphStoreConfig {
29    /// Which backend to use.
30    pub backend: GraphBackend,
31    /// Namespace partitions (for InMemory and KnowledgeGraph).
32    pub namespaces: Vec<String>,
33    /// Default namespace (for Simple and KnowledgeGraph).
34    pub default_namespace: String,
35    /// Default layer (for Simple and KnowledgeGraph). None = 0.
36    pub default_layer: Option<u8>,
37}
38
39impl Default for GraphStoreConfig {
40    fn default() -> Self {
41        Self {
42            backend: GraphBackend::KnowledgeGraph,
43            namespaces: vec!["default".to_string()],
44            default_namespace: "default".to_string(),
45            default_layer: Some(0),
46        }
47    }
48}
49
50impl GraphStoreConfig {
51    /// Create config for a specific backend with default namespace.
52    pub fn new(backend: GraphBackend) -> Self {
53        Self {
54            backend,
55            ..Default::default()
56        }
57    }
58
59    /// Set the namespace partitions.
60    pub fn with_namespaces(mut self, namespaces: &[&str]) -> Self {
61        self.namespaces = namespaces.iter().map(|s| s.to_string()).collect();
62        if !self.namespaces.contains(&self.default_namespace) && !self.namespaces.is_empty() {
63            self.default_namespace = self.namespaces[0].clone();
64        }
65        self
66    }
67
68    /// Set the default namespace.
69    pub fn with_default_namespace(mut self, ns: &str) -> Self {
70        self.default_namespace = ns.to_string();
71        self
72    }
73
74    /// Set the default layer.
75    pub fn with_layer(mut self, layer: Option<u8>) -> Self {
76        self.default_layer = layer;
77        self
78    }
79}
80
81/// Created graph store — enum dispatch over backend types.
82pub enum CreatedStore {
83    InMemory(ArrowGraphStore),
84    Simple(SimpleTripleStore),
85    KnowledgeGraph(KgStore),
86}
87
88impl CreatedStore {
89    /// Get the number of triples in the store.
90    pub fn len(&self) -> usize {
91        match self {
92            Self::InMemory(s) => s.len(),
93            Self::Simple(s) => s.len(),
94            Self::KnowledgeGraph(s) => s.len(),
95        }
96    }
97
98    /// Whether the store is empty.
99    pub fn is_empty(&self) -> bool {
100        self.len() == 0
101    }
102
103    /// Unwrap as ArrowGraphStore (panics if wrong variant).
104    pub fn into_arrow(self) -> ArrowGraphStore {
105        match self {
106            Self::InMemory(s) => s,
107            _ => panic!("expected InMemory variant"),
108        }
109    }
110
111    /// Unwrap as SimpleTripleStore (panics if wrong variant).
112    pub fn into_simple(self) -> SimpleTripleStore {
113        match self {
114            Self::Simple(s) => s,
115            _ => panic!("expected Simple variant"),
116        }
117    }
118
119    /// Unwrap as KgStore (panics if wrong variant).
120    pub fn into_kg(self) -> KgStore {
121        match self {
122            Self::KnowledgeGraph(s) => s,
123            _ => panic!("expected KnowledgeGraph variant"),
124        }
125    }
126
127    /// Try to unwrap as ArrowGraphStore.
128    pub fn try_into_arrow(self) -> Option<ArrowGraphStore> {
129        match self {
130            Self::InMemory(s) => Some(s),
131            _ => None,
132        }
133    }
134
135    /// Try to unwrap as SimpleTripleStore.
136    pub fn try_into_simple(self) -> Option<SimpleTripleStore> {
137        match self {
138            Self::Simple(s) => Some(s),
139            _ => None,
140        }
141    }
142
143    /// Try to unwrap as KgStore.
144    pub fn try_into_kg(self) -> Option<KgStore> {
145        match self {
146            Self::KnowledgeGraph(s) => Some(s),
147            _ => None,
148        }
149    }
150}
151
152/// List available graph backends. All are always available (CPU Arrow).
153pub fn available_backends() -> Vec<GraphBackend> {
154    vec![
155        GraphBackend::InMemory,
156        GraphBackend::Simple,
157        GraphBackend::KnowledgeGraph,
158    ]
159}
160
161/// Create a graph store based on configuration.
162pub fn create_graph_store(config: &GraphStoreConfig) -> CreatedStore {
163    let ns_refs: Vec<&str> = config.namespaces.iter().map(|s| s.as_str()).collect();
164    match config.backend {
165        GraphBackend::InMemory => CreatedStore::InMemory(ArrowGraphStore::new(&ns_refs)),
166        GraphBackend::Simple => CreatedStore::Simple(SimpleTripleStore::with_defaults(
167            &config.default_namespace,
168            config.default_layer,
169        )),
170        GraphBackend::KnowledgeGraph => CreatedStore::KnowledgeGraph(KgStore::with_config(
171            &ns_refs,
172            &config.default_namespace,
173            config.default_layer,
174        )),
175    }
176}
177
178/// Create a graph store with default configuration (KnowledgeGraph backend).
179pub fn create_default_store() -> KgStore {
180    KgStore::new()
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    #[test]
188    fn test_available_backends() {
189        let backends = available_backends();
190        assert_eq!(backends.len(), 3);
191    }
192
193    #[test]
194    fn test_create_in_memory() {
195        let config = GraphStoreConfig::new(GraphBackend::InMemory);
196        let store = create_graph_store(&config);
197        assert!(store.is_empty());
198        let _arrow = store.into_arrow();
199    }
200
201    #[test]
202    fn test_create_simple() {
203        let config = GraphStoreConfig::new(GraphBackend::Simple);
204        let store = create_graph_store(&config);
205        assert!(store.is_empty());
206        let _simple = store.into_simple();
207    }
208
209    #[test]
210    fn test_create_knowledge_graph() {
211        let config = GraphStoreConfig::new(GraphBackend::KnowledgeGraph);
212        let store = create_graph_store(&config);
213        assert!(store.is_empty());
214        let kg = store.into_kg();
215        assert!(!kg.prefixes().is_empty());
216    }
217
218    #[test]
219    fn test_create_default_store() {
220        let store = create_default_store();
221        assert!(store.is_empty());
222        assert!(!store.prefixes().is_empty());
223    }
224
225    #[test]
226    fn test_config_with_namespaces() {
227        let config = GraphStoreConfig::new(GraphBackend::InMemory)
228            .with_namespaces(&["world", "code", "self"]);
229        let store = create_graph_store(&config);
230        let arrow = store.into_arrow();
231        assert_eq!(arrow.namespaces().len(), 3);
232    }
233
234    #[test]
235    fn test_try_into_wrong_variant() {
236        let store = create_graph_store(&GraphStoreConfig::new(GraphBackend::InMemory));
237        assert!(store.try_into_kg().is_none());
238    }
239}