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        self.node_sources.insert(label.into(), source);
39        self
40    }
41
42    pub fn with_relationship_source(
43        mut self,
44        rel_type: impl Into<String>,
45        source: Arc<dyn TableSource>,
46    ) -> Self {
47        self.rel_sources.insert(rel_type.into(), source);
48        self
49    }
50}
51
52impl Default for InMemoryCatalog {
53    fn default() -> Self {
54        Self::new()
55    }
56}
57
58impl GraphSourceCatalog for InMemoryCatalog {
59    fn node_source(&self, label: &str) -> Option<Arc<dyn TableSource>> {
60        self.node_sources.get(label).cloned()
61    }
62
63    fn relationship_source(&self, rel_type: &str) -> Option<Arc<dyn TableSource>> {
64        self.rel_sources.get(rel_type).cloned()
65    }
66}
67
68/// A trivial logical table source with a fixed schema.
69pub struct SimpleTableSource {
70    schema: SchemaRef,
71}
72
73impl SimpleTableSource {
74    pub fn new(schema: SchemaRef) -> Self {
75        Self { schema }
76    }
77    pub fn empty() -> Self {
78        Self {
79            schema: Arc::new(Schema::empty()),
80        }
81    }
82}
83
84impl TableSource for SimpleTableSource {
85    fn as_any(&self) -> &dyn Any {
86        self
87    }
88    fn schema(&self) -> SchemaRef {
89        self.schema.clone()
90    }
91}
lance_graph/source_catalog.rs

lance_graph/
source_catalog.rs
