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()?; let _kind = parts.next()?; let _namespace = parts.next()?; parts.next() }
155}