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

lance_csr_catalog/
source_catalog.rs
