Skip to main content

vecgraph_core/
keys.rs

1use crate::{EdgeId, NodeId, SearchKind};
2
3pub enum StorageKey {
4    Node(NodeId),
5    NodeVector {
6        kind: String,
7        namespace: Option<String>,
8        node_id: NodeId,
9    },
10
11    Edge(EdgeId),
12    EdgeVector {
13        kind: String,
14        namespace: Option<String>,
15        source_node_id: NodeId,
16        target_node_id: NodeId,
17    },
18
19    EdgesForNode {
20        node_id: NodeId,
21        edge_id: EdgeId,
22    },
23
24    EdgesTargetingNode {
25        node_id: NodeId,
26        edge_id: EdgeId,
27    },
28
29    NameMapping {
30        kind: String,
31        name: String,
32    },
33}
34
35impl StorageKey {
36    pub fn partition(&self) -> &'static str {
37        match self {
38            StorageKey::Node(_) => "nodes",
39            StorageKey::NodeVector { .. } => "node_vectors",
40            StorageKey::Edge(_) => "edges",
41            StorageKey::EdgeVector { .. } => "edge_vectors",
42            StorageKey::EdgesForNode { .. } => "edges_for_node",
43            StorageKey::EdgesTargetingNode { .. } => "edges_targeting_node",
44            StorageKey::NameMapping { .. } => "names",
45        }
46    }
47
48    pub fn key(&self) -> String {
49        match self {
50            StorageKey::Node(id) => format!("node:{}", id.as_str()),
51            StorageKey::NodeVector {
52                kind,
53                namespace,
54                node_id,
55            } => {
56                let namespace = namespace.as_deref().unwrap_or("_");
57                format!("vec:{}:{}:{}", kind, namespace, node_id.as_str())
58            }
59            StorageKey::Edge(id) => format!("edge:{}", id.as_str()),
60            StorageKey::EdgeVector {
61                kind,
62                namespace,
63                source_node_id,
64                target_node_id,
65            } => {
66                let namespace = namespace.as_deref().unwrap_or("_");
67                format!(
68                    "vec:{}:{}:{}:{}",
69                    kind,
70                    namespace,
71                    source_node_id.as_str(),
72                    target_node_id.as_str()
73                )
74            }
75            StorageKey::EdgesForNode { node_id, edge_id } => {
76                format!("efn:{}:{}", node_id.as_str(), edge_id.as_str())
77            }
78            StorageKey::EdgesTargetingNode { node_id, edge_id } => {
79                format!("etn:{}:{}", node_id.as_str(), edge_id.as_str())
80            }
81            StorageKey::NameMapping { kind, name } => {
82                format!("name:{}:{}", kind, name)
83            }
84        }
85    }
86
87    pub fn edges_for_node_prefix(node_id: &NodeId) -> (&'static str, String) {
88        ("edges_for_node", format!("efn:{}:", node_id.as_str()))
89    }
90
91    pub fn edges_targeting_node_prefix(node_id: &NodeId) -> (&'static str, String) {
92        ("edges_targeting_node", format!("etn:{}:", node_id.as_str()))
93    }
94}
95
96pub struct VectorScanQuery {
97    pub kind: String,
98    pub namespace: Option<String>,
99    pub search_kind: SearchKind,
100}
101
102impl VectorScanQuery {
103    pub fn new(kind: impl Into<String>, search_kind: impl Into<SearchKind>) -> Self {
104        Self {
105            kind: kind.into(),
106            namespace: None,
107            search_kind: search_kind.into(),
108        }
109    }
110
111    pub fn with_namespace(mut self, namespace: impl Into<String>) -> Self {
112        self.namespace = Some(namespace.into());
113        self
114    }
115
116    pub fn partitions(&self) -> Vec<String> {
117        let mut partitions = Vec::new();
118        match self.search_kind {
119            SearchKind::Edge => {
120                partitions.push("edge_vectors".to_string());
121            }
122            SearchKind::Node => {
123                partitions.push("node_vectors".to_string());
124            }
125            SearchKind::All => {
126                partitions.push("edge_vectors".to_string());
127                partitions.push("node_vectors".to_string());
128            }
129        }
130        partitions
131    }
132
133    pub fn scan_prefix(&self) -> String {
134        match &self.namespace {
135            Some(ns) => format!("vec:{}:{}:", self.kind, ns),
136            None => format!("vec:{}:", self.kind),
137        }
138    }
139
140    pub fn search_kind_from_partition(partition: &str) -> Option<SearchKind> {
141        match partition {
142            "node_vectors" => Some(SearchKind::Node),
143            "edge_vectors" => Some(SearchKind::Edge),
144            _ => None,
145        }
146    }
147
148    pub fn node_id_from_key(key: &str) -> Option<&str> {
149        let mut parts = key.splitn(5, ':');
150        let _prefix = parts.next()?; // "vec"
151        let _kind = parts.next()?; // kind or node kind
152        let _namespace = parts.next()?; // namespace or "_"
153        parts.next() // source_node_id (for edges) or node_id (for nodes)
154    }
155}