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)]
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)]
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 pub fn get_field(&self, name: &str) -> Option<&Value> {
196 if let Some(ref named) = self.named {
198 return named.get(name);
199 }
200 if let Some(ref schema) = self.schema {
202 if let Some(idx) = schema.iter().position(|s| s == name) {
203 return self.columns.get(idx);
204 }
205 }
206 None
207 }
208
209 pub fn iter_fields(&self) -> Box<dyn Iterator<Item = (&str, &Value)> + '_> {
211 if let Some(ref named) = self.named {
212 Box::new(named.iter().map(|(k, v)| (k.as_str(), v)))
213 } else if let Some(ref schema) = self.schema {
214 Box::new(
215 schema
216 .iter()
217 .zip(self.columns.iter())
218 .map(|(k, v)| (k.as_str(), v)),
219 )
220 } else {
221 Box::new(std::iter::empty())
222 }
223 }
224
225 pub fn get(&self, index: usize) -> Option<&Value> {
227 self.columns.get(index)
228 }
229
230 pub fn get_by_name(&self, name: &str) -> Option<&Value> {
232 self.named.as_ref()?.get(name)
233 }
234
235 pub fn len(&self) -> usize {
237 self.columns.len()
238 }
239
240 pub fn is_empty(&self) -> bool {
242 self.columns.is_empty()
243 }
244}
245
246#[derive(Debug, Clone)]
248pub struct NodeData {
249 pub properties: HashMap<String, Value>,
251}
252
253impl NodeData {
254 pub fn new() -> Self {
256 Self {
257 properties: HashMap::new(),
258 }
259 }
260
261 pub fn with_properties(properties: HashMap<String, Value>) -> Self {
263 Self { properties }
264 }
265
266 pub fn set(&mut self, key: impl Into<String>, value: Value) {
268 self.properties.insert(key.into(), value);
269 }
270
271 pub fn get(&self, key: &str) -> Option<&Value> {
273 self.properties.get(key)
274 }
275
276 pub fn has(&self, key: &str) -> bool {
278 self.properties.contains_key(key)
279 }
280}
281
282impl Default for NodeData {
283 fn default() -> Self {
284 Self::new()
285 }
286}
287
288#[derive(Debug, Clone)]
290pub struct EdgeData {
291 pub properties: HashMap<String, Value>,
293 pub weight: f32,
295}
296
297impl EdgeData {
298 pub fn new(weight: f32) -> Self {
300 Self {
301 properties: HashMap::new(),
302 weight,
303 }
304 }
305
306 pub fn with_properties(weight: f32, properties: HashMap<String, Value>) -> Self {
308 Self { properties, weight }
309 }
310
311 pub fn set(&mut self, key: impl Into<String>, value: Value) {
313 self.properties.insert(key.into(), value);
314 }
315
316 pub fn get(&self, key: &str) -> Option<&Value> {
318 self.properties.get(key)
319 }
320}
321
322impl Default for EdgeData {
323 fn default() -> Self {
324 Self::new(1.0)
325 }
326}
327
328#[derive(Debug, Clone)]
330pub struct VectorData {
331 pub dense: Vec<f32>,
333 pub sparse: Option<SparseVector>,
335 pub content: Option<String>,
337}
338
339impl VectorData {
340 pub fn new(dense: Vec<f32>) -> Self {
342 Self {
343 dense,
344 sparse: None,
345 content: None,
346 }
347 }
348
349 pub fn with_sparse(dense: Vec<f32>, sparse: SparseVector) -> Self {
351 Self {
352 dense,
353 sparse: Some(sparse),
354 content: None,
355 }
356 }
357
358 pub fn with_content(mut self, content: impl Into<String>) -> Self {
360 self.content = Some(content.into());
361 self
362 }
363
364 pub fn dimension(&self) -> usize {
366 self.dense.len()
367 }
368
369 pub fn is_hybrid(&self) -> bool {
371 self.sparse.is_some()
372 }
373}
374
375#[derive(Debug, Clone)]
377pub struct TimeSeriesData {
378 pub metric: String,
380 pub series_id: Option<u64>,
382 pub timestamp_ns: u64,
384 pub value: f64,
386 pub tags: std::collections::HashMap<String, String>,
388 pub fields: std::collections::HashMap<String, Value>,
390}
391
392#[derive(Debug, Clone)]
394pub struct QueueMessageData {
395 pub payload: Value,
397 pub priority: Option<i32>,
399 pub enqueued_at_ns: u64,
401 pub attempts: u32,
403 pub max_attempts: u32,
405 pub acked: bool,
407}
408
409#[derive(Debug, Clone)]
411pub struct SparseVector {
412 pub indices: Vec<u32>,
414 pub values: Vec<f32>,
416 pub dimension: usize,
418}
419
420impl SparseVector {
421 pub fn new(indices: Vec<u32>, values: Vec<f32>, dimension: usize) -> Self {
423 debug_assert_eq!(indices.len(), values.len());
424 Self {
425 indices,
426 values,
427 dimension,
428 }
429 }
430
431 pub fn nnz(&self) -> usize {
433 self.indices.len()
434 }
435
436 pub fn sparsity(&self) -> f32 {
438 if self.dimension == 0 {
439 1.0
440 } else {
441 1.0 - (self.nnz() as f32 / self.dimension as f32)
442 }
443 }
444
445 pub fn get(&self, index: u32) -> f32 {
447 self.indices
448 .iter()
449 .position(|&i| i == index)
450 .map(|pos| self.values[pos])
451 .unwrap_or(0.0)
452 }
453}
454
455#[derive(Debug, Clone)]
457pub struct EmbeddingSlot {
458 pub name: String,
460 pub vector: Vec<f32>,
462 pub model: String,
464 pub dimension: usize,
466 pub generated_at: u64,
468}
469
470use crate::clock::current_unix_secs;
471
472impl EmbeddingSlot {
473 pub fn new(name: impl Into<String>, vector: Vec<f32>, model: impl Into<String>) -> Self {
475 let dimension = vector.len();
476 Self {
477 name: name.into(),
478 vector,
479 model: model.into(),
480 dimension,
481 generated_at: current_unix_secs(),
482 }
483 }
484}
485
486#[derive(Debug, Clone)]
488pub struct UnifiedEntity {
489 pub id: EntityId,
491 logical_id: Option<EntityId>,
495 pub kind: EntityKind,
497 pub created_at: u64,
499 pub updated_at: u64,
501 pub data: EntityData,
503 pub sequence_id: u64,
505 pub field_bloom: u64,
516 pub xmin: u64,
525 pub xmax: u64,
531 aux: Option<Box<EntityAux>>,
534}
535
536#[derive(Debug, Clone, Default)]
538pub struct EntityAux {
539 pub embeddings: Vec<EmbeddingSlot>,
541 pub cross_refs: Vec<CrossRef>,
543}
544
545impl UnifiedEntity {
546 pub fn embeddings(&self) -> &[EmbeddingSlot] {
548 self.aux
549 .as_ref()
550 .map(|a| a.embeddings.as_slice())
551 .unwrap_or(&[])
552 }
553
554 pub fn cross_refs(&self) -> &[CrossRef] {
556 self.aux
557 .as_ref()
558 .map(|a| a.cross_refs.as_slice())
559 .unwrap_or(&[])
560 }
561
562 pub fn embeddings_mut(&mut self) -> &mut Vec<EmbeddingSlot> {
564 &mut self.aux.get_or_insert_with(Default::default).embeddings
565 }
566
567 pub fn cross_refs_mut(&mut self) -> &mut Vec<CrossRef> {
569 &mut self.aux.get_or_insert_with(Default::default).cross_refs
570 }
571
572 pub fn has_aux(&self) -> bool {
574 self.aux.is_some()
575 }
576}
577
578#[inline]
584pub fn field_name_bloom(name: &str) -> u64 {
585 let b = name.as_bytes();
586 if b.is_empty() {
587 return 0;
588 }
589 1u64 << (b[b.len() / 2] & 63)
590}
591
592pub fn compute_entity_field_bloom(data: &EntityData) -> u64 {
596 match data {
597 EntityData::Row(row) => {
598 if row.schema.is_some() {
599 return 0;
602 }
603 if let Some(named) = &row.named {
604 let mut bloom = named.keys().fold(0u64, |acc, k| acc | field_name_bloom(k));
605 if let Some(Value::Json(bytes)) = named.get("body") {
613 if let Ok(names) = reddb_types::document_body_codec::field_names(bytes) {
614 for name in names {
615 bloom |= field_name_bloom(&name);
616 }
617 }
618 }
619 bloom
620 } else {
621 0
622 }
623 }
624 EntityData::Node(node) => node
625 .properties
626 .keys()
627 .fold(0u64, |acc, k| acc | field_name_bloom(k)),
628 EntityData::Edge(edge) => edge
629 .properties
630 .keys()
631 .fold(0u64, |acc, k| acc | field_name_bloom(k)),
632 _ => 0,
634 }
635}
636
637impl UnifiedEntity {
638 pub fn new(id: EntityId, kind: EntityKind, data: EntityData) -> Self {
640 let now = current_unix_secs();
641 let field_bloom = compute_entity_field_bloom(&data);
642
643 Self {
644 id,
645 logical_id: None,
646 kind,
647 created_at: now,
648 updated_at: now,
649 data,
650 sequence_id: 0,
651 field_bloom,
652 xmin: 0,
655 xmax: 0,
656 aux: None,
657 }
658 }
659
660 #[inline]
671 pub fn is_visible(&self, snapshot_xid: u64) -> bool {
672 if self.xmin != 0 && self.xmin > snapshot_xid {
673 return false;
674 }
675 if self.xmax != 0 && self.xmax <= snapshot_xid {
676 return false;
677 }
678 true
679 }
680
681 #[inline]
684 pub fn set_xmin(&mut self, xid: u64) {
685 self.xmin = xid;
686 }
687
688 #[inline]
692 pub fn set_xmax(&mut self, xid: u64) {
693 self.xmax = xid;
694 }
695
696 #[inline]
699 pub fn logical_id(&self) -> EntityId {
700 self.logical_id.unwrap_or(self.id)
701 }
702
703 #[inline]
705 pub fn has_explicit_logical_id(&self) -> bool {
706 self.logical_id.is_some()
707 }
708
709 #[inline]
711 pub fn set_logical_id(&mut self, logical_id: EntityId) {
712 self.logical_id = Some(logical_id);
713 }
714
715 #[inline]
727 pub fn ensure_table_logical_id(&mut self) {
728 if self.logical_id.is_none()
729 && matches!(
730 self.kind,
731 EntityKind::TableRow { .. } | EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
732 )
733 {
734 self.logical_id = Some(self.id);
735 }
736 }
737
738 pub fn table_row(
740 id: EntityId,
741 table: impl Into<Arc<str>>,
742 row_id: u64,
743 columns: Vec<Value>,
744 ) -> Self {
745 Self::new(
746 id,
747 EntityKind::TableRow {
748 table: table.into(),
749 row_id,
750 },
751 EntityData::Row(RowData::new(columns)),
752 )
753 }
754
755 pub fn graph_node(
757 id: EntityId,
758 label: impl Into<String>,
759 node_type: impl Into<String>,
760 properties: HashMap<String, Value>,
761 ) -> Self {
762 Self::new(
763 id,
764 EntityKind::GraphNode(Box::new(GraphNodeKind {
765 label: label.into(),
766 node_type: node_type.into(),
767 })),
768 EntityData::Node(NodeData::with_properties(properties)),
769 )
770 }
771
772 pub fn graph_edge(
774 id: EntityId,
775 label: impl Into<String>,
776 from: impl Into<String>,
777 to: impl Into<String>,
778 weight: f32,
779 properties: HashMap<String, Value>,
780 ) -> Self {
781 Self::new(
782 id,
783 EntityKind::GraphEdge(Box::new(GraphEdgeKind {
784 label: label.into(),
785 from_node: from.into(),
786 to_node: to.into(),
787 weight: (weight * 1000.0) as u32,
788 })),
789 EntityData::Edge(EdgeData::with_properties(weight, properties)),
790 )
791 }
792
793 pub fn vector(id: EntityId, collection: impl Into<String>, vector: Vec<f32>) -> Self {
795 Self::new(
796 id,
797 EntityKind::Vector {
798 collection: collection.into(),
799 },
800 EntityData::Vector(VectorData::new(vector)),
801 )
802 }
803
804 pub fn add_embedding(&mut self, slot: EmbeddingSlot) {
806 self.embeddings_mut().push(slot);
807 self.touch();
808 }
809
810 pub fn add_cross_ref(&mut self, cross_ref: CrossRef) {
812 self.cross_refs_mut().push(cross_ref);
813 self.touch();
814 }
815
816 pub fn get_embedding(&self, name: &str) -> Option<&EmbeddingSlot> {
818 self.embeddings().iter().find(|e| e.name == name)
819 }
820
821 fn touch(&mut self) {
823 self.updated_at = current_unix_secs();
824 }
825
826 pub fn is_stale(&self, max_age_secs: u64) -> bool {
828 let now = current_unix_secs();
829 now.saturating_sub(self.updated_at) > max_age_secs
830 }
831}
832
833#[derive(Debug, Clone, PartialEq)]
835pub struct CrossRef {
836 pub source: EntityId,
838 pub target: EntityId,
840 pub target_collection: String,
842 pub ref_type: RefType,
844 pub weight: f32,
846 pub created_at: u64,
848}
849
850impl CrossRef {
851 pub fn new(
853 source: EntityId,
854 target: EntityId,
855 target_collection: impl Into<String>,
856 ref_type: RefType,
857 ) -> Self {
858 Self {
859 source,
860 target,
861 target_collection: target_collection.into(),
862 ref_type,
863 weight: 1.0,
864 created_at: current_unix_secs(),
865 }
866 }
867
868 pub fn with_weight(
870 source: EntityId,
871 target: EntityId,
872 target_collection: impl Into<String>,
873 ref_type: RefType,
874 weight: f32,
875 ) -> Self {
876 let mut cr = Self::new(source, target, target_collection, ref_type);
877 cr.weight = weight;
878 cr
879 }
880}
881
882#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
884pub enum RefType {
885 RowToNode, RowToEdge, NodeToRow, RowToVector, VectorToRow, NodeToVector, EdgeToVector, VectorToNode, SimilarTo, RelatedTo, DerivesFrom, Mentions, Contains, DependsOn, }
907
908impl RefType {
909 pub fn inverse(&self) -> Option<Self> {
911 match self {
912 Self::RowToNode => Some(Self::NodeToRow),
913 Self::NodeToRow => Some(Self::RowToNode),
914 Self::RowToVector => Some(Self::VectorToRow),
915 Self::VectorToRow => Some(Self::RowToVector),
916 Self::NodeToVector => Some(Self::VectorToNode),
917 Self::VectorToNode => Some(Self::NodeToVector),
918 Self::SimilarTo => Some(Self::SimilarTo), Self::RelatedTo => Some(Self::RelatedTo), _ => None, }
922 }
923
924 pub fn is_symmetric(&self) -> bool {
926 matches!(self, Self::SimilarTo | Self::RelatedTo)
927 }
928
929 pub fn to_byte(&self) -> u8 {
931 match self {
932 Self::RowToNode => 0,
933 Self::RowToEdge => 1,
934 Self::NodeToRow => 2,
935 Self::RowToVector => 3,
936 Self::VectorToRow => 4,
937 Self::NodeToVector => 5,
938 Self::EdgeToVector => 6,
939 Self::VectorToNode => 7,
940 Self::SimilarTo => 8,
941 Self::RelatedTo => 9,
942 Self::DerivesFrom => 10,
943 Self::Mentions => 11,
944 Self::Contains => 12,
945 Self::DependsOn => 13,
946 }
947 }
948
949 pub fn from_byte(byte: u8) -> Self {
951 match byte {
952 0 => Self::RowToNode,
953 1 => Self::RowToEdge,
954 2 => Self::NodeToRow,
955 3 => Self::RowToVector,
956 4 => Self::VectorToRow,
957 5 => Self::NodeToVector,
958 6 => Self::EdgeToVector,
959 7 => Self::VectorToNode,
960 8 => Self::SimilarTo,
961 9 => Self::RelatedTo,
962 10 => Self::DerivesFrom,
963 11 => Self::Mentions,
964 12 => Self::Contains,
965 13 => Self::DependsOn,
966 _ => Self::RelatedTo, }
968 }
969}
970
971impl From<Vec<Value>> for RowData {
973 fn from(columns: Vec<Value>) -> Self {
974 RowData::new(columns)
975 }
976}
977
978impl From<HashMap<String, Value>> for NodeData {
980 fn from(properties: HashMap<String, Value>) -> Self {
981 NodeData::with_properties(properties)
982 }
983}
984
985impl From<Vec<f32>> for VectorData {
987 fn from(dense: Vec<f32>) -> Self {
988 VectorData::new(dense)
989 }
990}
991
992impl From<(Vec<f32>, SparseVector)> for VectorData {
994 fn from((dense, sparse): (Vec<f32>, SparseVector)) -> Self {
995 VectorData::with_sparse(dense, sparse)
996 }
997}
998
999impl UnifiedEntity {
1001 pub fn from_properties(
1003 id: EntityId,
1004 label: impl Into<String>,
1005 node_type: impl Into<String>,
1006 properties: impl IntoIterator<Item = (impl Into<String>, Value)>,
1007 ) -> Self {
1008 let props: HashMap<String, Value> =
1009 properties.into_iter().map(|(k, v)| (k.into(), v)).collect();
1010 Self::graph_node(id, label, node_type, props)
1011 }
1012
1013 pub fn into_row(self) -> Option<RowData> {
1015 match self.data {
1016 EntityData::Row(r) => Some(r),
1017 _ => None,
1018 }
1019 }
1020
1021 pub fn into_node(self) -> Option<NodeData> {
1023 match self.data {
1024 EntityData::Node(n) => Some(n),
1025 _ => None,
1026 }
1027 }
1028
1029 pub fn into_edge(self) -> Option<EdgeData> {
1031 match self.data {
1032 EntityData::Edge(e) => Some(e),
1033 _ => None,
1034 }
1035 }
1036
1037 pub fn into_vector(self) -> Option<VectorData> {
1039 match self.data {
1040 EntityData::Vector(v) => Some(v),
1041 _ => None,
1042 }
1043 }
1044}
1045
1046#[cfg(test)]
1047mod tests {
1048 use super::*;
1049
1050 #[test]
1051 fn test_entity_creation() {
1052 let id = EntityId::new(1);
1053 let entity = UnifiedEntity::table_row(
1054 id,
1055 "users",
1056 100,
1057 vec![Value::text("alice".to_string()), Value::Integer(25)],
1058 );
1059
1060 assert!(entity.data.is_row());
1061 assert_eq!(entity.kind.storage_type(), "table");
1062 assert_eq!(entity.kind.collection(), "users");
1063 }
1064
1065 #[test]
1066 fn test_cross_refs() {
1067 let id1 = EntityId::new(1);
1068 let id2 = EntityId::new(2);
1069
1070 let cross_ref = CrossRef::new(id1, id2, "nodes", RefType::RowToNode);
1071 assert_eq!(cross_ref.source, id1);
1072 assert_eq!(cross_ref.target, id2);
1073 assert_eq!(cross_ref.ref_type.inverse(), Some(RefType::NodeToRow));
1074 }
1075
1076 #[test]
1077 fn test_sparse_vector() {
1078 let sparse = SparseVector::new(vec![0, 5, 10], vec![1.0, 2.0, 3.0], 100);
1079
1080 assert_eq!(sparse.nnz(), 3);
1081 assert_eq!(sparse.get(5), 2.0);
1082 assert_eq!(sparse.get(3), 0.0);
1083 assert!(sparse.sparsity() > 0.9);
1084 }
1085
1086 #[test]
1087 fn test_embedding_slots() {
1088 let mut entity = UnifiedEntity::table_row(
1089 EntityId::new(1),
1090 "documents",
1091 1,
1092 vec![Value::text("Hello world".to_string())],
1093 );
1094
1095 entity.add_embedding(EmbeddingSlot::new(
1096 "content",
1097 vec![0.1, 0.2, 0.3],
1098 "text-embedding-3-small",
1099 ));
1100
1101 assert_eq!(entity.embeddings().len(), 1);
1102 assert!(entity.get_embedding("content").is_some());
1103 assert!(entity.get_embedding("summary").is_none());
1104 }
1105}