use std::collections::HashMap;
use std::sync::Arc;
use arrow_schema::{Schema, SchemaRef};
use datafusion::logical_expr::TableSource;
pub trait GraphSourceCatalog: Send + Sync {
fn node_source(&self, label: &str) -> Option<Arc<dyn TableSource>>;
fn relationship_source(&self, rel_type: &str) -> Option<Arc<dyn TableSource>>;
}
pub struct InMemoryCatalog {
node_sources: HashMap<String, Arc<dyn TableSource>>,
rel_sources: HashMap<String, Arc<dyn TableSource>>,
}
impl InMemoryCatalog {
pub fn new() -> Self {
Self {
node_sources: HashMap::new(),
rel_sources: HashMap::new(),
}
}
pub fn with_node_source(
mut self,
label: impl Into<String>,
source: Arc<dyn TableSource>,
) -> Self {
self.node_sources
.insert(label.into().to_lowercase(), source);
self
}
pub fn with_relationship_source(
mut self,
rel_type: impl Into<String>,
source: Arc<dyn TableSource>,
) -> Self {
self.rel_sources
.insert(rel_type.into().to_lowercase(), source);
self
}
}
impl Default for InMemoryCatalog {
fn default() -> Self {
Self::new()
}
}
impl GraphSourceCatalog for InMemoryCatalog {
fn node_source(&self, label: &str) -> Option<Arc<dyn TableSource>> {
self.node_sources.get(&label.to_lowercase()).cloned()
}
fn relationship_source(&self, rel_type: &str) -> Option<Arc<dyn TableSource>> {
self.rel_sources.get(&rel_type.to_lowercase()).cloned()
}
}
pub struct SimpleTableSource {
schema: SchemaRef,
}
impl SimpleTableSource {
pub fn new(schema: SchemaRef) -> Self {
Self { schema }
}
pub fn empty() -> Self {
Self {
schema: Arc::new(Schema::empty()),
}
}
}
impl TableSource for SimpleTableSource {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
}