1use std::collections::HashMap;
7use std::sync::Arc;
8
9use reddb_types::Value;
10
11pub const FIRST_USER_ENTITY_ID: u64 = 1024;
23
24pub use super::entity_id::EntityId;
25
26#[derive(Debug, Clone, PartialEq, Eq, Hash)]
28pub enum EntityKind {
29 TableRow { table: Arc<str>, row_id: u64 },
31 GraphNode(Box<GraphNodeKind>),
33 GraphEdge(Box<GraphEdgeKind>),
35 Vector { collection: String },
37 TimeSeriesPoint(Box<TimeSeriesPointKind>),
39 QueueMessage { queue: String, position: u64 },
41}
42
43#[derive(Debug, Clone, PartialEq, Eq, Hash)]
44pub struct GraphNodeKind {
45 pub label: String,
46 pub node_type: String,
47}
48
49#[derive(Debug, Clone, PartialEq, Eq, Hash)]
50pub struct GraphEdgeKind {
51 pub label: String,
52 pub from_node: String,
53 pub to_node: String,
54 pub weight: u32,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq, Hash)]
58pub struct TimeSeriesPointKind {
59 pub series: String,
60 pub metric: String,
61}
62
63impl EntityKind {
64 pub fn storage_type(&self) -> &'static str {
66 match self {
67 Self::TableRow { .. } => "table",
68 Self::GraphNode(_) => "graph_node",
69 Self::GraphEdge(_) => "graph_edge",
70 Self::Vector { .. } => "vector",
71 Self::TimeSeriesPoint(_) => "timeseries",
72 Self::QueueMessage { .. } => "queue",
73 }
74 }
75
76 pub fn collection(&self) -> &str {
78 match self {
79 Self::TableRow { table, .. } => table,
80 Self::GraphNode(n) => &n.label,
81 Self::GraphEdge(e) => &e.label,
82 Self::Vector { collection } => collection,
83 Self::TimeSeriesPoint(ts) => &ts.series,
84 Self::QueueMessage { queue, .. } => queue,
85 }
86 }
87}
88
89#[derive(Debug, Clone, PartialEq)]
91pub enum EntityData {
92 Row(RowData),
94 Node(NodeData),
96 Edge(EdgeData),
98 Vector(VectorData),
100 TimeSeries(TimeSeriesData),
102 QueueMessage(QueueMessageData),
104}
105
106impl EntityData {
107 pub fn is_row(&self) -> bool {
109 matches!(self, Self::Row(_))
110 }
111
112 pub fn is_node(&self) -> bool {
114 matches!(self, Self::Node(_))
115 }
116
117 pub fn is_edge(&self) -> bool {
119 matches!(self, Self::Edge(_))
120 }
121
122 pub fn is_vector(&self) -> bool {
124 matches!(self, Self::Vector(_))
125 }
126
127 pub fn as_row(&self) -> Option<&RowData> {
129 match self {
130 Self::Row(r) => Some(r),
131 _ => None,
132 }
133 }
134
135 pub fn as_node(&self) -> Option<&NodeData> {
137 match self {
138 Self::Node(n) => Some(n),
139 _ => None,
140 }
141 }
142
143 pub fn as_edge(&self) -> Option<&EdgeData> {
145 match self {
146 Self::Edge(e) => Some(e),
147 _ => None,
148 }
149 }
150
151 pub fn as_vector(&self) -> Option<&VectorData> {
153 match self {
154 Self::Vector(v) => Some(v),
155 _ => None,
156 }
157 }
158}
159
160#[derive(Debug, Clone, PartialEq)]
162pub struct RowData {
163 pub columns: Vec<Value>,
165 pub named: Option<HashMap<String, Value>>,
167 pub schema: Option<std::sync::Arc<Vec<String>>>,
171}
172
173impl RowData {
174 pub fn new(columns: Vec<Value>) -> Self {
176 Self {
177 columns,
178 named: None,
179 schema: None,
180 }
181 }
182
183 pub fn with_names(columns: Vec<Value>, names: Vec<String>) -> Self {
185 let named: HashMap<String, Value> =
186 names.into_iter().zip(columns.iter().cloned()).collect();
187 Self {
188 columns,
189 named: Some(named),
190 schema: None,
191 }
192 }
193
194 #[cfg(feature = "full")]
197 pub(crate) fn compact_with_schema(&mut self, schema: &std::sync::Arc<Vec<String>>) -> bool {
198 let Some(named) = &self.named else {
199 return false;
200 };
201 if named.len() != schema.len() || schema.iter().any(|name| !named.contains_key(name)) {
202 return false;
203 }
204 if !self.columns.is_empty()
207 && (self.columns.len() != schema.len()
208 || schema
209 .iter()
210 .zip(&self.columns)
211 .any(|(name, value)| named.get(name) != Some(value)))
212 {
213 return false;
214 }
215 let mut named = self.named.take().expect("validated named row");
216 self.columns = schema
217 .iter()
218 .map(|name| named.remove(name).expect("validated schema field"))
219 .collect();
220 self.schema = Some(std::sync::Arc::clone(schema));
221 true
222 }
223
224 pub fn get_field(&self, name: &str) -> Option<&Value> {
226 if let Some(ref named) = self.named {
228 return named.get(name);
229 }
230 if let Some(ref schema) = self.schema {
232 if let Some(idx) = schema.iter().position(|s| s == name) {
233 return self.columns.get(idx);
234 }
235 }
236 None
237 }
238
239 pub fn iter_fields(&self) -> Box<dyn Iterator<Item = (&str, &Value)> + '_> {
241 if let Some(ref named) = self.named {
242 Box::new(named.iter().map(|(k, v)| (k.as_str(), v)))
243 } else if let Some(ref schema) = self.schema {
244 Box::new(
245 schema
246 .iter()
247 .zip(self.columns.iter())
248 .map(|(k, v)| (k.as_str(), v)),
249 )
250 } else {
251 Box::new(std::iter::empty())
252 }
253 }
254
255 pub fn get(&self, index: usize) -> Option<&Value> {
257 self.columns.get(index)
258 }
259
260 pub fn get_by_name(&self, name: &str) -> Option<&Value> {
262 self.named.as_ref()?.get(name)
263 }
264
265 pub fn len(&self) -> usize {
267 self.columns.len()
268 }
269
270 pub fn is_empty(&self) -> bool {
272 self.columns.is_empty()
273 }
274}
275
276#[derive(Debug, Clone, PartialEq)]
278pub struct NodeData {
279 pub properties: HashMap<String, Value>,
281}
282
283impl NodeData {
284 pub fn new() -> Self {
286 Self {
287 properties: HashMap::new(),
288 }
289 }
290
291 pub fn with_properties(properties: HashMap<String, Value>) -> Self {
293 Self { properties }
294 }
295
296 pub fn set(&mut self, key: impl Into<String>, value: Value) {
298 self.properties.insert(key.into(), value);
299 }
300
301 pub fn get(&self, key: &str) -> Option<&Value> {
303 self.properties.get(key)
304 }
305
306 pub fn has(&self, key: &str) -> bool {
308 self.properties.contains_key(key)
309 }
310}
311
312impl Default for NodeData {
313 fn default() -> Self {
314 Self::new()
315 }
316}
317
318#[derive(Debug, Clone, PartialEq)]
320pub struct EdgeData {
321 pub properties: HashMap<String, Value>,
323 pub weight: f32,
325}
326
327impl EdgeData {
328 pub fn new(weight: f32) -> Self {
330 Self {
331 properties: HashMap::new(),
332 weight,
333 }
334 }
335
336 pub fn with_properties(weight: f32, properties: HashMap<String, Value>) -> Self {
338 Self { properties, weight }
339 }
340
341 pub fn set(&mut self, key: impl Into<String>, value: Value) {
343 self.properties.insert(key.into(), value);
344 }
345
346 pub fn get(&self, key: &str) -> Option<&Value> {
348 self.properties.get(key)
349 }
350}
351
352impl Default for EdgeData {
353 fn default() -> Self {
354 Self::new(1.0)
355 }
356}
357
358#[derive(Debug, Clone, PartialEq)]
360pub struct VectorData {
361 pub dense: Vec<f32>,
363 pub sparse: Option<SparseVector>,
365 pub content: Option<String>,
367}
368
369impl VectorData {
370 pub fn new(dense: Vec<f32>) -> Self {
372 Self {
373 dense,
374 sparse: None,
375 content: None,
376 }
377 }
378
379 pub fn with_sparse(dense: Vec<f32>, sparse: SparseVector) -> Self {
381 Self {
382 dense,
383 sparse: Some(sparse),
384 content: None,
385 }
386 }
387
388 pub fn with_content(mut self, content: impl Into<String>) -> Self {
390 self.content = Some(content.into());
391 self
392 }
393
394 pub fn dimension(&self) -> usize {
396 self.dense.len()
397 }
398
399 pub fn is_hybrid(&self) -> bool {
401 self.sparse.is_some()
402 }
403}
404
405#[derive(Debug, Clone, PartialEq)]
407pub struct TimeSeriesData {
408 pub metric: String,
410 pub series_id: Option<u64>,
412 pub timestamp_ns: u64,
414 pub value: f64,
416 pub tags: std::collections::HashMap<String, String>,
418 pub fields: std::collections::HashMap<String, Value>,
420}
421
422#[derive(Debug, Clone, PartialEq)]
424pub struct QueueMessageData {
425 pub payload: Value,
427 pub priority: Option<i32>,
429 pub enqueued_at_ns: u64,
431 pub attempts: u32,
433 pub max_attempts: u32,
435 pub acked: bool,
437}
438
439#[derive(Debug, Clone, PartialEq)]
441pub struct SparseVector {
442 pub indices: Vec<u32>,
444 pub values: Vec<f32>,
446 pub dimension: usize,
448}
449
450impl SparseVector {
451 pub fn new(indices: Vec<u32>, values: Vec<f32>, dimension: usize) -> Self {
453 debug_assert_eq!(indices.len(), values.len());
454 Self {
455 indices,
456 values,
457 dimension,
458 }
459 }
460
461 pub fn nnz(&self) -> usize {
463 self.indices.len()
464 }
465
466 pub fn sparsity(&self) -> f32 {
468 if self.dimension == 0 {
469 1.0
470 } else {
471 1.0 - (self.nnz() as f32 / self.dimension as f32)
472 }
473 }
474
475 pub fn get(&self, index: u32) -> f32 {
477 self.indices
478 .iter()
479 .position(|&i| i == index)
480 .map(|pos| self.values[pos])
481 .unwrap_or(0.0)
482 }
483}
484
485#[derive(Debug, Clone, PartialEq)]
487pub struct EmbeddingSlot {
488 pub name: String,
490 pub vector: Vec<f32>,
492 pub model: String,
494 pub dimension: usize,
496 pub generated_at: u64,
498}
499
500use crate::clock::current_unix_secs;
501
502impl EmbeddingSlot {
503 pub fn new(name: impl Into<String>, vector: Vec<f32>, model: impl Into<String>) -> Self {
505 let dimension = vector.len();
506 Self {
507 name: name.into(),
508 vector,
509 model: model.into(),
510 dimension,
511 generated_at: current_unix_secs(),
512 }
513 }
514}
515
516#[derive(Debug, Clone)]
518pub struct UnifiedEntity {
519 pub id: EntityId,
521 logical_id: Option<EntityId>,
525 pub kind: EntityKind,
527 pub created_at: u64,
529 pub updated_at: u64,
531 pub data: EntityData,
533 pub sequence_id: u64,
535 pub field_bloom: u64,
546 pub xmin: u64,
555 pub xmax: u64,
561 aux: Option<Box<EntityAux>>,
564}
565
566#[derive(Debug, Clone, Default)]
568pub struct EntityAux {
569 pub embeddings: Vec<EmbeddingSlot>,
571 pub cross_refs: Vec<CrossRef>,
573 pub branch: Option<BranchVersion>,
576}
577
578#[derive(Debug, Clone, Default, PartialEq, Eq)]
579pub struct BranchVersion {
580 pub origin: u64,
581 pub deletions: Vec<(u64, u64)>,
583}
584
585impl UnifiedEntity {
586 pub fn canonicalize_system_collection_names(&mut self) {
590 match &mut self.kind {
591 EntityKind::TableRow { table, .. } => {
592 if let Some(canonical) = super::system_collection_names::canonical_name(table) {
593 *table = Arc::from(canonical);
594 }
595 }
596 EntityKind::Vector { collection } => {
597 if let Some(canonical) = super::system_collection_names::canonical_name(collection)
598 {
599 *collection = canonical.to_owned();
600 }
601 }
602 EntityKind::TimeSeriesPoint(point) => {
603 if let Some(canonical) =
604 super::system_collection_names::canonical_name(&point.series)
605 {
606 point.series = canonical.to_owned();
607 }
608 }
609 EntityKind::QueueMessage { queue, .. } => {
610 if let Some(canonical) = super::system_collection_names::canonical_name(queue) {
611 *queue = canonical.to_owned();
612 }
613 }
614 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_) => {}
615 }
616 if let Some(aux) = &mut self.aux {
617 for reference in &mut aux.cross_refs {
618 if let Some(canonical) =
619 super::system_collection_names::canonical_name(&reference.target_collection)
620 {
621 reference.target_collection = canonical.to_owned();
622 }
623 }
624 }
625 }
626
627 pub fn embeddings(&self) -> &[EmbeddingSlot] {
629 self.aux
630 .as_ref()
631 .map(|a| a.embeddings.as_slice())
632 .unwrap_or(&[])
633 }
634
635 pub fn cross_refs(&self) -> &[CrossRef] {
637 self.aux
638 .as_ref()
639 .map(|a| a.cross_refs.as_slice())
640 .unwrap_or(&[])
641 }
642
643 pub fn embeddings_mut(&mut self) -> &mut Vec<EmbeddingSlot> {
645 &mut self.aux.get_or_insert_with(Default::default).embeddings
646 }
647
648 pub fn cross_refs_mut(&mut self) -> &mut Vec<CrossRef> {
650 &mut self.aux.get_or_insert_with(Default::default).cross_refs
651 }
652
653 pub fn has_aux(&self) -> bool {
655 self.aux.is_some()
656 }
657}
658
659#[inline]
665pub fn field_name_bloom(name: &str) -> u64 {
666 let b = name.as_bytes();
667 if b.is_empty() {
668 return 0;
669 }
670 1u64 << (b[b.len() / 2] & 63)
671}
672
673pub fn compute_entity_field_bloom(data: &EntityData) -> u64 {
677 match data {
678 EntityData::Row(row) => {
679 if row.schema.is_some() {
680 return 0;
683 }
684 if let Some(named) = &row.named {
685 let mut bloom = named.keys().fold(0u64, |acc, k| acc | field_name_bloom(k));
686 if let Some(Value::Json(bytes)) = named.get("body") {
694 if let Ok(names) = reddb_types::document_body_codec::field_names(bytes) {
695 for name in names {
696 bloom |= field_name_bloom(&name);
697 }
698 }
699 }
700 bloom
701 } else {
702 0
703 }
704 }
705 EntityData::Node(node) => node
706 .properties
707 .keys()
708 .fold(0u64, |acc, k| acc | field_name_bloom(k)),
709 EntityData::Edge(edge) => edge
710 .properties
711 .keys()
712 .fold(0u64, |acc, k| acc | field_name_bloom(k)),
713 _ => 0,
715 }
716}
717
718impl UnifiedEntity {
719 pub fn new(id: EntityId, kind: EntityKind, data: EntityData) -> Self {
721 let now = current_unix_secs();
722 let field_bloom = compute_entity_field_bloom(&data);
723
724 Self {
725 id,
726 logical_id: None,
727 kind,
728 created_at: now,
729 updated_at: now,
730 data,
731 sequence_id: 0,
732 field_bloom,
733 xmin: 0,
736 xmax: 0,
737 aux: None,
738 }
739 }
740
741 #[inline]
752 pub fn is_visible(&self, snapshot_xid: u64) -> bool {
753 if self.xmin != 0 && self.xmin > snapshot_xid {
754 return false;
755 }
756 if self.xmax != 0 && self.xmax <= snapshot_xid {
757 return false;
758 }
759 true
760 }
761
762 #[inline]
765 pub fn set_xmin(&mut self, xid: u64) {
766 self.xmin = xid;
767 }
768
769 #[inline]
773 pub fn set_xmax(&mut self, xid: u64) {
774 self.xmax = xid;
775 }
776
777 pub fn branch_version(&self) -> Option<&BranchVersion> {
778 self.aux.as_deref()?.branch.as_ref()
779 }
780
781 pub fn set_branch_origin(&mut self, epoch: u64) {
782 self.aux.get_or_insert_with(Default::default).branch = Some(BranchVersion {
783 origin: epoch,
784 deletions: Vec::new(),
785 });
786 }
787
788 pub fn clear_branch_version(&mut self) {
789 if let Some(aux) = self.aux.as_deref_mut() {
790 aux.branch = None;
791 }
792 }
793
794 pub fn branch_xmax(&self, epoch: u64) -> u64 {
795 match self.branch_version() {
796 Some(branch) if branch.origin != epoch => branch
797 .deletions
798 .binary_search_by_key(&epoch, |(owner, _)| *owner)
799 .map_or(0, |index| branch.deletions[index].1),
800 _ => self.xmax,
801 }
802 }
803
804 pub fn set_branch_xmax(&mut self, epoch: u64, xid: u64) {
807 let Some(branch) = self.aux.as_deref_mut().and_then(|aux| aux.branch.as_mut()) else {
808 self.xmax = xid;
809 return;
810 };
811 if branch.origin == epoch {
812 self.xmax = xid;
813 return;
814 }
815 match branch
816 .deletions
817 .binary_search_by_key(&epoch, |(owner, _)| *owner)
818 {
819 Ok(index) if xid == 0 => {
820 branch.deletions.remove(index);
821 }
822 Ok(index) => branch.deletions[index].1 = xid,
823 Err(index) if xid != 0 => branch.deletions.insert(index, (epoch, xid)),
824 Err(_) => {}
825 }
826 }
827
828 pub fn retain_branch_deletions(&mut self, mut keep: impl FnMut(u64) -> bool) {
829 if let Some(branch) = self.aux.as_deref_mut().and_then(|aux| aux.branch.as_mut()) {
830 branch.deletions.retain(|(_, xid)| keep(*xid));
831 }
832 }
833
834 pub fn max_transaction_id(&self) -> u64 {
835 self.branch_version()
836 .map_or(self.xmin.max(self.xmax), |branch| {
837 branch
838 .deletions
839 .iter()
840 .fold(self.xmin.max(self.xmax), |high, (_, xid)| high.max(*xid))
841 })
842 }
843
844 #[inline]
847 pub fn logical_id(&self) -> EntityId {
848 self.logical_id.unwrap_or(self.id)
849 }
850
851 #[inline]
853 pub fn has_explicit_logical_id(&self) -> bool {
854 self.logical_id.is_some()
855 }
856
857 #[inline]
859 pub fn set_logical_id(&mut self, logical_id: EntityId) {
860 self.logical_id = Some(logical_id);
861 }
862
863 #[inline]
875 pub fn ensure_table_logical_id(&mut self) {
876 if self.logical_id.is_none()
877 && matches!(
878 self.kind,
879 EntityKind::TableRow { .. } | EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
880 )
881 {
882 self.logical_id = Some(self.id);
883 }
884 }
885
886 pub fn table_row(
888 id: EntityId,
889 table: impl Into<Arc<str>>,
890 row_id: u64,
891 columns: Vec<Value>,
892 ) -> Self {
893 Self::new(
894 id,
895 EntityKind::TableRow {
896 table: table.into(),
897 row_id,
898 },
899 EntityData::Row(RowData::new(columns)),
900 )
901 }
902
903 pub fn graph_node(
905 id: EntityId,
906 label: impl Into<String>,
907 node_type: impl Into<String>,
908 properties: HashMap<String, Value>,
909 ) -> Self {
910 Self::new(
911 id,
912 EntityKind::GraphNode(Box::new(GraphNodeKind {
913 label: label.into(),
914 node_type: node_type.into(),
915 })),
916 EntityData::Node(NodeData::with_properties(properties)),
917 )
918 }
919
920 pub fn graph_edge(
922 id: EntityId,
923 label: impl Into<String>,
924 from: impl Into<String>,
925 to: impl Into<String>,
926 weight: f32,
927 properties: HashMap<String, Value>,
928 ) -> Self {
929 Self::new(
930 id,
931 EntityKind::GraphEdge(Box::new(GraphEdgeKind {
932 label: label.into(),
933 from_node: from.into(),
934 to_node: to.into(),
935 weight: (weight * 1000.0) as u32,
936 })),
937 EntityData::Edge(EdgeData::with_properties(weight, properties)),
938 )
939 }
940
941 pub fn vector(id: EntityId, collection: impl Into<String>, vector: Vec<f32>) -> Self {
943 Self::new(
944 id,
945 EntityKind::Vector {
946 collection: collection.into(),
947 },
948 EntityData::Vector(VectorData::new(vector)),
949 )
950 }
951
952 pub fn add_embedding(&mut self, slot: EmbeddingSlot) {
954 self.embeddings_mut().push(slot);
955 self.touch();
956 }
957
958 pub fn add_cross_ref(&mut self, cross_ref: CrossRef) {
960 self.cross_refs_mut().push(cross_ref);
961 self.touch();
962 }
963
964 pub fn get_embedding(&self, name: &str) -> Option<&EmbeddingSlot> {
966 self.embeddings().iter().find(|e| e.name == name)
967 }
968
969 fn touch(&mut self) {
971 self.updated_at = current_unix_secs();
972 }
973
974 pub fn is_stale(&self, max_age_secs: u64) -> bool {
976 let now = current_unix_secs();
977 now.saturating_sub(self.updated_at) > max_age_secs
978 }
979}
980
981#[derive(Debug, Clone, PartialEq)]
983pub struct CrossRef {
984 pub source: EntityId,
986 pub target: EntityId,
988 pub target_collection: String,
990 pub ref_type: RefType,
992 pub weight: f32,
994 pub created_at: u64,
996}
997
998impl CrossRef {
999 pub fn new(
1001 source: EntityId,
1002 target: EntityId,
1003 target_collection: impl Into<String>,
1004 ref_type: RefType,
1005 ) -> Self {
1006 Self {
1007 source,
1008 target,
1009 target_collection: target_collection.into(),
1010 ref_type,
1011 weight: 1.0,
1012 created_at: current_unix_secs(),
1013 }
1014 }
1015
1016 pub fn with_weight(
1018 source: EntityId,
1019 target: EntityId,
1020 target_collection: impl Into<String>,
1021 ref_type: RefType,
1022 weight: f32,
1023 ) -> Self {
1024 let mut cr = Self::new(source, target, target_collection, ref_type);
1025 cr.weight = weight;
1026 cr
1027 }
1028}
1029
1030#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1032pub enum RefType {
1033 RowToNode, RowToEdge, NodeToRow, RowToVector, VectorToRow, NodeToVector, EdgeToVector, VectorToNode, SimilarTo, RelatedTo, DerivesFrom, Mentions, Contains, DependsOn, }
1055
1056impl RefType {
1057 pub fn inverse(&self) -> Option<Self> {
1059 match self {
1060 Self::RowToNode => Some(Self::NodeToRow),
1061 Self::NodeToRow => Some(Self::RowToNode),
1062 Self::RowToVector => Some(Self::VectorToRow),
1063 Self::VectorToRow => Some(Self::RowToVector),
1064 Self::NodeToVector => Some(Self::VectorToNode),
1065 Self::VectorToNode => Some(Self::NodeToVector),
1066 Self::SimilarTo => Some(Self::SimilarTo), Self::RelatedTo => Some(Self::RelatedTo), _ => None, }
1070 }
1071
1072 pub fn is_symmetric(&self) -> bool {
1074 matches!(self, Self::SimilarTo | Self::RelatedTo)
1075 }
1076
1077 pub fn to_byte(&self) -> u8 {
1079 match self {
1080 Self::RowToNode => 0,
1081 Self::RowToEdge => 1,
1082 Self::NodeToRow => 2,
1083 Self::RowToVector => 3,
1084 Self::VectorToRow => 4,
1085 Self::NodeToVector => 5,
1086 Self::EdgeToVector => 6,
1087 Self::VectorToNode => 7,
1088 Self::SimilarTo => 8,
1089 Self::RelatedTo => 9,
1090 Self::DerivesFrom => 10,
1091 Self::Mentions => 11,
1092 Self::Contains => 12,
1093 Self::DependsOn => 13,
1094 }
1095 }
1096
1097 pub fn from_byte(byte: u8) -> Self {
1099 match byte {
1100 0 => Self::RowToNode,
1101 1 => Self::RowToEdge,
1102 2 => Self::NodeToRow,
1103 3 => Self::RowToVector,
1104 4 => Self::VectorToRow,
1105 5 => Self::NodeToVector,
1106 6 => Self::EdgeToVector,
1107 7 => Self::VectorToNode,
1108 8 => Self::SimilarTo,
1109 9 => Self::RelatedTo,
1110 10 => Self::DerivesFrom,
1111 11 => Self::Mentions,
1112 12 => Self::Contains,
1113 13 => Self::DependsOn,
1114 _ => Self::RelatedTo, }
1116 }
1117}
1118
1119impl From<Vec<Value>> for RowData {
1121 fn from(columns: Vec<Value>) -> Self {
1122 RowData::new(columns)
1123 }
1124}
1125
1126impl From<HashMap<String, Value>> for NodeData {
1128 fn from(properties: HashMap<String, Value>) -> Self {
1129 NodeData::with_properties(properties)
1130 }
1131}
1132
1133impl From<Vec<f32>> for VectorData {
1135 fn from(dense: Vec<f32>) -> Self {
1136 VectorData::new(dense)
1137 }
1138}
1139
1140impl From<(Vec<f32>, SparseVector)> for VectorData {
1142 fn from((dense, sparse): (Vec<f32>, SparseVector)) -> Self {
1143 VectorData::with_sparse(dense, sparse)
1144 }
1145}
1146
1147impl UnifiedEntity {
1149 pub fn from_properties(
1151 id: EntityId,
1152 label: impl Into<String>,
1153 node_type: impl Into<String>,
1154 properties: impl IntoIterator<Item = (impl Into<String>, Value)>,
1155 ) -> Self {
1156 let props: HashMap<String, Value> =
1157 properties.into_iter().map(|(k, v)| (k.into(), v)).collect();
1158 Self::graph_node(id, label, node_type, props)
1159 }
1160
1161 pub fn into_row(self) -> Option<RowData> {
1163 match self.data {
1164 EntityData::Row(r) => Some(r),
1165 _ => None,
1166 }
1167 }
1168
1169 pub fn into_node(self) -> Option<NodeData> {
1171 match self.data {
1172 EntityData::Node(n) => Some(n),
1173 _ => None,
1174 }
1175 }
1176
1177 pub fn into_edge(self) -> Option<EdgeData> {
1179 match self.data {
1180 EntityData::Edge(e) => Some(e),
1181 _ => None,
1182 }
1183 }
1184
1185 pub fn into_vector(self) -> Option<VectorData> {
1187 match self.data {
1188 EntityData::Vector(v) => Some(v),
1189 _ => None,
1190 }
1191 }
1192}
1193
1194#[cfg(test)]
1195mod tests {
1196 use super::*;
1197
1198 #[test]
1199 fn test_entity_creation() {
1200 let id = EntityId::new(1);
1201 let entity = UnifiedEntity::table_row(
1202 id,
1203 "users",
1204 100,
1205 vec![Value::text("alice".to_string()), Value::Integer(25)],
1206 );
1207
1208 assert!(entity.data.is_row());
1209 assert_eq!(entity.kind.storage_type(), "table");
1210 assert_eq!(entity.kind.collection(), "users");
1211 }
1212
1213 #[test]
1214 fn test_cross_refs() {
1215 let id1 = EntityId::new(1);
1216 let id2 = EntityId::new(2);
1217
1218 let cross_ref = CrossRef::new(id1, id2, "nodes", RefType::RowToNode);
1219 assert_eq!(cross_ref.source, id1);
1220 assert_eq!(cross_ref.target, id2);
1221 assert_eq!(cross_ref.ref_type.inverse(), Some(RefType::NodeToRow));
1222 }
1223
1224 #[test]
1225 fn test_sparse_vector() {
1226 let sparse = SparseVector::new(vec![0, 5, 10], vec![1.0, 2.0, 3.0], 100);
1227
1228 assert_eq!(sparse.nnz(), 3);
1229 assert_eq!(sparse.get(5), 2.0);
1230 assert_eq!(sparse.get(3), 0.0);
1231 assert!(sparse.sparsity() > 0.9);
1232 }
1233
1234 #[test]
1235 fn test_embedding_slots() {
1236 let mut entity = UnifiedEntity::table_row(
1237 EntityId::new(1),
1238 "documents",
1239 1,
1240 vec![Value::text("Hello world".to_string())],
1241 );
1242
1243 entity.add_embedding(EmbeddingSlot::new(
1244 "content",
1245 vec![0.1, 0.2, 0.3],
1246 "text-embedding-3-small",
1247 ));
1248
1249 assert_eq!(entity.embeddings().len(), 1);
1250 assert!(entity.get_embedding("content").is_some());
1251 assert!(entity.get_embedding("summary").is_none());
1252 }
1253}