1use std::collections::{BTreeMap, HashMap, HashSet};
32use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
33
34use super::entity::{CrossRef, EntityData, EntityId, EntityKind, RefType, UnifiedEntity};
35use super::metadata::{Metadata, MetadataFilter, MetadataStorage};
36use super::segment_graph_index::SegmentGraphIndex;
37use super::segment_key_lookup::SegmentKeyLookup;
38use super::segment_unique_key::SegmentUniqueKey;
39use reddb_types::value_compare::partial_compare_values;
40use reddb_types::{value_to_canonical_key, CanonicalKey, Value};
41
42pub type SegmentId = u64;
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum SegmentState {
48 Growing,
50 Sealing,
52 Sealed,
54 Flushed,
56 Archived,
58}
59
60impl SegmentState {
61 pub fn is_writable(&self) -> bool {
63 matches!(self, Self::Growing)
64 }
65
66 pub fn is_queryable(&self) -> bool {
68 !matches!(self, Self::Sealing)
69 }
70
71 pub fn is_immutable(&self) -> bool {
73 matches!(self, Self::Sealed | Self::Flushed | Self::Archived)
74 }
75}
76
77#[derive(Debug, Clone)]
79pub struct SegmentConfig {
80 pub max_entities: usize,
82 pub max_bytes: usize,
84 pub max_age_secs: u64,
86 pub build_vector_index: bool,
88 pub build_graph_index: bool,
90 pub compression_level: u8,
92}
93
94impl Default for SegmentConfig {
95 fn default() -> Self {
96 Self {
97 max_entities: 100_000,
98 max_bytes: 256 * 1024 * 1024, max_age_secs: 3600, build_vector_index: true,
101 build_graph_index: true,
102 compression_level: 6,
103 }
104 }
105}
106
107#[derive(Debug, Clone, Default)]
109pub struct SegmentStats {
110 pub entity_count: usize,
112 pub deleted_count: usize,
114 pub memory_bytes: usize,
116 pub vector_count: usize,
118 pub node_count: usize,
120 pub edge_count: usize,
122 pub row_count: usize,
124 pub cross_ref_count: usize,
126}
127
128#[derive(Debug, Clone)]
130pub enum SegmentError {
131 NotWritable,
133 NotFound(EntityId),
135 AlreadyExists(EntityId),
137 Full,
139 InvalidState(SegmentState),
141 Internal(String),
143}
144
145impl std::fmt::Display for SegmentError {
146 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147 match self {
148 Self::NotWritable => write!(f, "segment is not writable"),
149 Self::NotFound(id) => write!(f, "entity not found: {}", id),
150 Self::AlreadyExists(id) => write!(f, "entity already exists: {}", id),
151 Self::Full => write!(f, "segment is full"),
152 Self::InvalidState(state) => write!(f, "invalid operation for state: {:?}", state),
153 Self::Internal(msg) => write!(f, "internal error: {}", msg),
154 }
155 }
156}
157
158impl std::error::Error for SegmentError {}
159
160use crate::clock::current_unix_secs;
161
162fn entity_id_from_probe_key(key: &[u8]) -> Option<EntityId> {
163 if let Ok(text) = std::str::from_utf8(key) {
164 if let Ok(raw) = text.parse::<u64>() {
165 return Some(EntityId::new(raw));
166 }
167 }
168
169 let bytes: [u8; 8] = key.try_into().ok()?;
170 Some(EntityId::new(u64::from_le_bytes(bytes)))
171}
172
173const SEALED_MULTI_ZONE_MAX_INTERVALS: usize = 4;
174
175const TOMBSTONE_ENTRY_BYTES: u64 = 24;
183
184#[derive(Debug, Clone)]
185struct UpdateIndexSnapshot {
186 pk_column_name: Option<String>,
187 pk_value: Option<Value>,
188 pk_index_key: Option<(String, String)>,
189 cross_refs: Vec<CrossRef>,
190}
191
192impl UpdateIndexSnapshot {
193 fn from_entity(entity: &UnifiedEntity) -> Self {
194 let (pk_column_name, pk_value) = match &entity.data {
195 EntityData::Row(row) => (
196 row.schema
197 .as_deref()
198 .and_then(|schema| schema.first().cloned()),
199 row.columns.first().cloned(),
200 ),
201 _ => (None, None),
202 };
203 let pk_index_key = pk_value
204 .as_ref()
205 .map(|value| (entity.kind.collection().to_string(), format!("{:?}", value)));
206 Self {
207 pk_column_name,
208 pk_value,
209 pk_index_key,
210 cross_refs: entity.cross_refs().to_vec(),
211 }
212 }
213}
214
215pub trait UnifiedSegment: Send + Sync {
217 fn id(&self) -> SegmentId;
219
220 fn state(&self) -> SegmentState;
222
223 fn collection(&self) -> &str;
225
226 fn stats(&self) -> SegmentStats;
228
229 fn entity_count(&self) -> usize;
231
232 fn contains(&self, id: EntityId) -> bool;
234
235 fn get(&self, id: EntityId) -> Option<&UnifiedEntity>;
237
238 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity>;
240
241 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError>;
243
244 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError>;
246
247 fn update_hot(
251 &mut self,
252 entity: UnifiedEntity,
253 modified_columns: &[String],
254 ) -> Result<(), SegmentError> {
255 let _ = modified_columns;
256 self.update(entity)
257 }
258
259 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError>;
261
262 fn get_metadata(&self, id: EntityId) -> Option<Metadata>;
264
265 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError>;
267
268 fn seal(&mut self) -> Result<(), SegmentError>;
270
271 fn should_seal(&self, config: &SegmentConfig) -> bool;
273
274 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
276
277 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
279
280 fn filter_metadata(
282 &self,
283 filters: &[(String, super::metadata::MetadataFilter)],
284 ) -> Vec<EntityId>;
285}
286
287#[derive(Debug, Clone)]
294pub struct ColZone {
295 pub min: Value,
296 pub max: Value,
297 min_key: Option<CanonicalKey>,
298 max_key: Option<CanonicalKey>,
299}
300
301impl ColZone {
302 fn payload_bytes(&self) -> usize {
303 super::memory_size::value_heap_bytes(&self.min)
304 .saturating_add(super::memory_size::value_heap_bytes(&self.max))
305 .saturating_add(
306 self.min_key
307 .as_ref()
308 .map_or(0, super::memory_size::canonical_heap_bytes),
309 )
310 .saturating_add(
311 self.max_key
312 .as_ref()
313 .map_or(0, super::memory_size::canonical_heap_bytes),
314 )
315 }
316
317 fn new(v: Value) -> Self {
318 Self {
319 min_key: value_to_canonical_key(&v),
320 max_key: value_to_canonical_key(&v),
321 min: v.clone(),
322 max: v,
323 }
324 }
325
326 fn with_bounds(min: Value, max: Value) -> Self {
327 Self {
328 min_key: value_to_canonical_key(&min),
329 max_key: value_to_canonical_key(&max),
330 min,
331 max,
332 }
333 }
334
335 fn update(&mut self, v: &Value) {
336 if compare_zone_values(v, None, &self.min, self.min_key.as_ref())
337 .map(|o| o == std::cmp::Ordering::Less)
338 .unwrap_or(false)
339 {
340 self.min = v.clone();
341 self.min_key = value_to_canonical_key(v);
342 }
343 if compare_zone_values(v, None, &self.max, self.max_key.as_ref())
344 .map(|o| o == std::cmp::Ordering::Greater)
345 .unwrap_or(false)
346 {
347 self.max = v.clone();
348 self.max_key = value_to_canonical_key(v);
349 }
350 }
351}
352
353#[derive(Debug, Clone, Default)]
354pub struct MultiColZone {
355 pub intervals: Vec<ColZone>,
356}
357
358impl MultiColZone {
359 fn can_skip(&self, pred: &ZoneColPred<'_>) -> bool {
360 !self.intervals.is_empty() && self.intervals.iter().all(|zone| pred.can_skip(zone))
361 }
362}
363
364fn compare_zone_values(
365 left: &Value,
366 left_key: Option<&CanonicalKey>,
367 right: &Value,
368 right_key: Option<&CanonicalKey>,
369) -> Option<std::cmp::Ordering> {
370 partial_compare_values(left, right).or_else(|| {
371 let left_key = left_key.cloned().or_else(|| value_to_canonical_key(left))?;
372 let right_key = right_key
373 .cloned()
374 .or_else(|| value_to_canonical_key(right))?;
375 (left_key.family() == right_key.family()).then(|| left_key.cmp(&right_key))
376 })
377}
378
379#[derive(Debug, Clone, Copy, PartialEq, Eq)]
381pub enum ZoneColPredKind {
382 Eq,
383 Gt,
384 Gte,
385 Lt,
386 Lte,
387}
388
389#[derive(Debug, Clone)]
391pub enum ZoneColPred<'a> {
392 Eq(&'a Value),
393 Gt(&'a Value),
394 Gte(&'a Value),
395 Lt(&'a Value),
396 Lte(&'a Value),
397}
398
399impl<'a> ZoneColPred<'a> {
400 pub fn can_skip(&self, zone: &ColZone) -> bool {
402 match self {
403 ZoneColPred::Eq(val) => {
405 compare_zone_values(val, None, &zone.min, zone.min_key.as_ref())
406 .map(|o| o == std::cmp::Ordering::Less)
407 .unwrap_or(false)
408 || compare_zone_values(val, None, &zone.max, zone.max_key.as_ref())
409 .map(|o| o == std::cmp::Ordering::Greater)
410 .unwrap_or(false)
411 }
412 ZoneColPred::Gt(val) => {
414 compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
415 .map(|o| o != std::cmp::Ordering::Greater)
416 .unwrap_or(false)
417 }
418 ZoneColPred::Gte(val) => {
420 compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
421 .map(|o| o == std::cmp::Ordering::Less)
422 .unwrap_or(false)
423 }
424 ZoneColPred::Lt(val) => {
426 compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
427 .map(|o| o != std::cmp::Ordering::Less)
428 .unwrap_or(false)
429 }
430 ZoneColPred::Lte(val) => {
432 compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
433 .map(|o| o == std::cmp::Ordering::Greater)
434 .unwrap_or(false)
435 }
436 }
437 }
438}
439
440#[derive(Clone, Copy)]
442pub enum GraphEntityKind {
443 Node,
444 Edge,
445}
446
447impl GraphEntityKind {
448 fn mask(self) -> u8 {
449 match self {
450 Self::Node => 1,
451 Self::Edge => 2,
452 }
453 }
454
455 pub fn matches(self, kind: &EntityKind) -> bool {
456 matches!(
457 (self, kind),
458 (Self::Node, EntityKind::GraphNode(_)) | (Self::Edge, EntityKind::GraphEdge(_))
459 )
460 }
461}
462
463pub const SCAN_BATCH_SIZE: usize = 256;
465
466pub const CHECKPOINT_MUTATIONS_MAX: usize = 8192;
468
469pub struct SegmentScanCursor {
471 flat_end: usize,
472 hash_end: usize,
473 deleted_end: usize,
474 position: usize,
475}
476
477impl SegmentScanCursor {
478 pub fn remaining(&self) -> usize {
479 self.flat_end + self.hash_end - self.position
480 }
481
482 pub fn finished(&self) -> bool {
483 self.remaining() == 0
484 }
485}
486
487pub struct GraphIndexCursor {
490 key_index: usize,
491 after: Option<EntityId>,
492 upper: EntityId,
493}
494
495impl GraphIndexCursor {
496 pub fn finished(&self, keys: &[EntityId]) -> bool {
497 self.key_index == keys.len()
498 }
499}
500
501pub struct GrowingSegment {
503 id: SegmentId,
505 collection: String,
507 state: SegmentState,
509 created_at: u64,
511 last_write_at: u64,
513
514 entities: HashMap<EntityId, UnifiedEntity>,
516 hash_scan_ids: Vec<EntityId>,
519 flat_entities: Vec<UnifiedEntity>,
522 base_entity_id: u64,
524 use_flat: bool,
526 deleted: HashMap<EntityId, usize>,
530 metadata: MetadataStorage,
532
533 pk_index: BTreeMap<(String, String), EntityId>,
535 unique_key_lookup: parking_lot::RwLock<Vec<SegmentUniqueKey>>,
538 unique_key_lookup_bytes: AtomicU64,
539 key_lookup: Option<SegmentKeyLookup>,
540 kind_index: HashMap<String, HashSet<EntityId>>,
542 graph_kind_mask: u8,
545 graph_read_index: parking_lot::RwLock<Option<SegmentGraphIndex>>,
547 graph_read_index_bytes: AtomicU64,
548 cross_ref_forward: HashMap<EntityId, Vec<(EntityId, RefType)>>,
550 cross_ref_reverse: HashMap<EntityId, Vec<(EntityId, RefType)>>,
552
553 col_zones: HashMap<String, ColZone>,
555 zone_payload_bytes: usize,
556 sealed_col_zones: HashMap<String, MultiColZone>,
558
559 sequence: AtomicU64,
561 checkpoint_revision: u64,
563 checkpoint_mutations: std::collections::VecDeque<(u64, EntityId)>,
564 checkpoint_latest_mutations: HashMap<EntityId, u64>,
567 checkpoint_mutation_floor: u64,
568 memory_bytes: AtomicU64,
573 dead_entity_bytes: u64,
576 dead_resident_count: usize,
578
579 published_flat_len: AtomicUsize,
586}
587
588impl GrowingSegment {
589 pub fn may_contain_graph_kind(&self, kind: GraphEntityKind) -> bool {
590 self.graph_kind_mask & kind.mask() != 0
591 }
592
593 fn observe_graph_kind(&mut self, kind: &EntityKind) {
594 self.graph_kind_mask |= match kind {
595 EntityKind::GraphNode(_) => GraphEntityKind::Node.mask(),
596 EntityKind::GraphEdge(_) => GraphEntityKind::Edge.mask(),
597 _ => 0,
598 };
599 }
600
601 pub fn scan_cursor(&self) -> SegmentScanCursor {
602 SegmentScanCursor {
603 flat_end: self.flat_entities.len(),
604 hash_end: self.hash_scan_ids.len(),
605 deleted_end: self.deleted.len(),
606 position: 0,
607 }
608 }
609
610 pub fn scan_batch(
614 &self,
615 cursor: &mut SegmentScanCursor,
616 before_work: &mut impl FnMut() -> bool,
617 visit: &mut impl FnMut(EntityId),
618 ) -> bool {
619 for _ in 0..SCAN_BATCH_SIZE {
620 if cursor.finished() {
621 break;
622 }
623 if !before_work() {
624 return false;
625 }
626 let id = if cursor.position < cursor.flat_end {
627 self.flat_entities[cursor.position].id
628 } else {
629 self.hash_scan_ids[cursor.position - cursor.flat_end]
630 };
631 cursor.position += 1;
632 let eligible = match self.deleted.get(&id) {
633 Some(&ordinal) => ordinal >= cursor.deleted_end,
634 None => self.get(id).is_some(),
637 };
638 if eligible {
639 visit(id);
640 }
641 }
642 true
643 }
644
645 #[inline]
648 pub fn for_each_fast<F>(&self, mut f: F) -> bool
649 where
650 F: FnMut(&UnifiedEntity) -> bool,
651 {
652 if self.use_flat {
653 if self.deleted.is_empty() {
655 for entity in &self.flat_entities {
656 if !f(entity) {
657 return false;
658 }
659 }
660 for entity in self.entities.values() {
663 if !f(entity) {
664 return false;
665 }
666 }
667 } else {
668 for entity in &self.flat_entities {
669 if self.deleted.contains_key(&entity.id) {
670 continue;
671 }
672 if !f(entity) {
673 return false;
674 }
675 }
676 for entity in self.entities.values() {
677 if self.deleted.contains_key(&entity.id) {
678 continue;
679 }
680 if !f(entity) {
681 return false;
682 }
683 }
684 }
685 } else {
686 if self.deleted.is_empty() {
688 for entity in self.entities.values() {
689 if !f(entity) {
690 return false;
691 }
692 }
693 } else {
694 for entity in self.entities.values() {
695 if self.deleted.contains_key(&entity.id) {
696 continue;
697 }
698 if !f(entity) {
699 return false;
700 }
701 }
702 }
703 }
704 true
705 }
706
707 pub fn new(id: SegmentId, collection: impl Into<String>) -> Self {
709 let now = current_unix_secs();
710
711 Self {
712 id,
713 collection: collection.into(),
714 state: SegmentState::Growing,
715 created_at: now,
716 last_write_at: now,
717 entities: HashMap::new(),
718 hash_scan_ids: Vec::new(),
719 flat_entities: Vec::new(),
720 base_entity_id: 0,
721 use_flat: false,
722 deleted: HashMap::new(),
723 metadata: MetadataStorage::new(),
724 pk_index: BTreeMap::new(),
725 unique_key_lookup: parking_lot::RwLock::new(Vec::new()),
726 unique_key_lookup_bytes: AtomicU64::new(0),
727 key_lookup: Some(SegmentKeyLookup::default()),
728 kind_index: HashMap::new(),
729 graph_read_index: parking_lot::RwLock::new(Some(SegmentGraphIndex::default())),
730 graph_read_index_bytes: AtomicU64::new(0),
731 graph_kind_mask: 0,
732 cross_ref_forward: HashMap::new(),
733 cross_ref_reverse: HashMap::new(),
734 col_zones: HashMap::new(),
735 zone_payload_bytes: 0,
736 sealed_col_zones: HashMap::new(),
737 sequence: AtomicU64::new(0),
738 checkpoint_revision: 0,
739 checkpoint_mutations: std::collections::VecDeque::new(),
740 checkpoint_latest_mutations: HashMap::new(),
741 checkpoint_mutation_floor: 0,
742 memory_bytes: AtomicU64::new(0),
743 dead_entity_bytes: 0,
744 dead_resident_count: 0,
745 published_flat_len: AtomicUsize::new(0),
746 }
747 }
748
749 fn next_sequence(&self) -> u64 {
751 self.sequence.fetch_add(1, Ordering::SeqCst)
752 }
753
754 pub fn checkpoint_revision(&self) -> u64 {
755 self.checkpoint_revision
756 }
757
758 fn mark_checkpoint_entity_dirty(&mut self, id: EntityId) {
759 self.checkpoint_revision = self
760 .checkpoint_revision
761 .checked_add(1)
762 .expect("segment mutation counter must not overflow");
763 if !self.checkpoint_latest_mutations.contains_key(&id)
764 && self.checkpoint_latest_mutations.len() == CHECKPOINT_MUTATIONS_MAX
765 {
766 while let Some((revision, previous)) = self.checkpoint_mutations.pop_front() {
769 if self.checkpoint_latest_mutations.get(&previous) == Some(&revision) {
770 self.checkpoint_latest_mutations.remove(&previous);
771 self.checkpoint_mutation_floor = self.checkpoint_mutation_floor.max(revision);
772 break;
773 }
774 }
775 }
776 self.checkpoint_latest_mutations
777 .insert(id, self.checkpoint_revision);
778 if self.checkpoint_mutations.len() == 2 * CHECKPOINT_MUTATIONS_MAX {
779 self.checkpoint_mutations.retain(|(revision, previous)| {
782 self.checkpoint_latest_mutations.get(previous) == Some(revision)
783 });
784 }
785 self.checkpoint_mutations
786 .push_back((self.checkpoint_revision, id));
787 }
788
789 pub fn checkpoint_mutations_since(&self, revision: u64) -> Option<Vec<EntityId>> {
790 if revision == self.checkpoint_revision {
791 return Some(Vec::new());
792 }
793 if revision < self.checkpoint_mutation_floor || revision > self.checkpoint_revision {
794 return None;
795 }
796 let mut ids = Vec::new();
797 for &(sequence, id) in &self.checkpoint_mutations {
798 if sequence > revision && self.checkpoint_latest_mutations.get(&id) == Some(&sequence) {
799 ids.push(id);
800 }
801 }
802 Some(ids)
803 }
804
805 pub fn checkpoint_tracking_bytes(&self) -> usize {
806 self.checkpoint_mutations.capacity() * std::mem::size_of::<(u64, EntityId)>()
807 + self.checkpoint_latest_mutations.capacity()
809 * (std::mem::size_of::<(EntityId, u64)>() + 16)
810 }
811
812 fn has_live_entity(&self, id: EntityId) -> bool {
813 if self.deleted.contains_key(&id) {
814 return false;
815 }
816 if self.use_flat {
817 let raw = id.raw();
818 if raw >= self.base_entity_id {
819 let idx = (raw - self.base_entity_id) as usize;
820 if self
821 .flat_entities
822 .get(idx)
823 .is_some_and(|entity| entity.id == id)
824 {
825 return true;
826 }
827 }
828 self.entities.contains_key(&id)
831 } else {
832 self.entities.contains_key(&id)
833 }
834 }
835
836 fn update_existing_entity_in_place(
837 &mut self,
838 entity: &UnifiedEntity,
839 ) -> Result<UpdateIndexSnapshot, SegmentError> {
840 self.mark_checkpoint_entity_dirty(entity.id);
841 self.invalidate_unique_key_lookup();
842 if self
843 .get(entity.id)
844 .is_some_and(|old| SegmentKeyLookup::key(old) != SegmentKeyLookup::key(entity))
845 {
846 if let Some(mut index) = self.key_lookup.take() {
847 index.remove(
848 self.get(entity.id)
849 .expect("invariant: existing replacement row"),
850 );
851 index.insert(entity);
852 self.key_lookup = Some(index);
853 }
854 }
855 if (self.graph_kind_mask != 0
856 || matches!(
857 entity.kind,
858 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
859 ))
860 && self
861 .get(entity.id)
862 .is_some_and(|old| !SegmentGraphIndex::same_keys(old, entity))
863 {
864 self.invalidate_graph_read_index();
865 }
866 self.observe_graph_kind(&entity.kind);
869 if self.use_flat {
870 let raw = entity.id.raw();
871 if raw >= self.base_entity_id {
872 let idx = (raw - self.base_entity_id) as usize;
873 if let Some(slot) = self
874 .flat_entities
875 .get_mut(idx)
876 .filter(|slot| slot.id == entity.id)
877 {
878 let snapshot = UpdateIndexSnapshot::from_entity(slot);
879 let old_bytes = Self::estimate_entity_size(slot);
880 slot.clone_from(entity);
881 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
882 return Ok(snapshot);
883 }
884 }
885 let Some(slot) = self.entities.get_mut(&entity.id) else {
886 return Err(SegmentError::NotFound(entity.id));
887 };
888 let snapshot = UpdateIndexSnapshot::from_entity(slot);
889 let old_bytes = Self::estimate_entity_size(slot);
890 slot.clone_from(entity);
891 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
892 Ok(snapshot)
893 } else {
894 let Some(slot) = self.entities.get_mut(&entity.id) else {
895 return Err(SegmentError::NotFound(entity.id));
896 };
897 let snapshot = UpdateIndexSnapshot::from_entity(slot);
898 let old_bytes = Self::estimate_entity_size(slot);
899 slot.clone_from(entity);
900 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
901 Ok(snapshot)
902 }
903 }
904
905 fn apply_hot_update_with_metadata(
906 &mut self,
907 entity: &UnifiedEntity,
908 modified_columns: &[String],
909 metadata: Option<&Metadata>,
910 ) -> Result<(), SegmentError> {
911 let old = self.update_existing_entity_in_place(entity)?;
912 self.reindex_for_update(&old, entity, Some(modified_columns));
913 self.update_col_zones_from_entity(entity);
914 if let Some(metadata) = metadata {
915 self.metadata.set_all(entity.id, metadata);
916 }
917 Ok(())
918 }
919
920 fn apply_update_with_metadata(
921 &mut self,
922 entity: &UnifiedEntity,
923 metadata: Option<&Metadata>,
924 ) -> Result<(), SegmentError> {
925 let old = self.update_existing_entity_in_place(entity)?;
926 self.reindex_for_update(&old, entity, None);
927 self.update_col_zones_from_entity(entity);
928 if let Some(metadata) = metadata {
929 self.metadata.set_all(entity.id, metadata);
930 }
931 Ok(())
932 }
933
934 pub fn update_hot_batch_with_metadata<'a, I>(&mut self, items: I) -> Result<(), SegmentError>
935 where
936 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
937 {
938 if !self.state.is_writable() {
939 return Err(SegmentError::NotWritable);
940 }
941
942 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
943 items.into_iter().collect();
944 if items.is_empty() {
945 return Ok(());
946 }
947
948 for (entity, _, _) in &items {
949 if !self.has_live_entity(entity.id) {
950 return Err(SegmentError::NotFound(entity.id));
951 }
952 }
953
954 for (entity, modified_columns, metadata) in items {
955 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)?;
956 }
957
958 self.last_write_at = current_unix_secs();
959 Ok(())
960 }
961
962 fn record_tombstone(&mut self, id: EntityId) -> bool {
964 self.mark_checkpoint_entity_dirty(id);
965 let ordinal = self.deleted.len();
966 match self.deleted.entry(id) {
967 std::collections::hash_map::Entry::Vacant(entry) => {
968 entry.insert(ordinal);
969 true
970 }
971 std::collections::hash_map::Entry::Occupied(_) => false,
972 }
973 }
974
975 fn tombstone_flat_entity(&mut self, idx: usize, id: EntityId) -> bool {
981 self.invalidate_unique_key_lookup();
982 if !self.record_tombstone(id) {
983 return false;
984 }
985 if let Some(index) = self.key_lookup.as_mut() {
986 index.remove(&self.flat_entities[idx]);
987 }
988 if let Some(index) = self.graph_read_index.get_mut() {
989 index.remove(&self.flat_entities[idx]);
990 self.graph_read_index_bytes
991 .store(index.memory_bytes(), Ordering::Relaxed);
992 }
993 let size = Self::estimate_entity_size(&self.flat_entities[idx]) as u64;
994 self.dead_entity_bytes += size;
995 self.dead_resident_count += 1;
996 self.metadata.remove_all(id);
997 true
998 }
999
1000 fn remove_hashmap_entity(&mut self, id: EntityId) -> bool {
1005 self.invalidate_unique_key_lookup();
1006 let Some(entity) = self.entities.remove(&id) else {
1007 return false;
1008 };
1009 self.release_memory(Self::estimate_entity_size(&entity));
1010 self.unindex_entity(&entity);
1011 self.metadata.remove_all(id);
1012 self.record_tombstone(id);
1013 true
1014 }
1015
1016 pub fn delete_batch(&mut self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1017 if !self.state.is_writable() {
1018 return Err(SegmentError::NotWritable);
1019 }
1020 if ids.is_empty() {
1021 return Ok(Vec::new());
1022 }
1023
1024 let mut deleted_ids = Vec::with_capacity(ids.len());
1025
1026 if self.use_flat {
1027 for &id in ids {
1028 let raw = id.raw();
1029 if raw < self.base_entity_id {
1030 if self.remove_hashmap_entity(id) {
1031 deleted_ids.push(id);
1032 }
1033 continue;
1034 }
1035 let idx = (raw - self.base_entity_id) as usize;
1036 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
1037 if self.tombstone_flat_entity(idx, id) {
1038 deleted_ids.push(id);
1039 }
1040 } else if self.remove_hashmap_entity(id) {
1041 deleted_ids.push(id);
1042 }
1043 }
1044 } else {
1045 for &id in ids {
1046 if self.remove_hashmap_entity(id) {
1047 deleted_ids.push(id);
1048 }
1049 }
1050 }
1051
1052 if !deleted_ids.is_empty() {
1053 self.last_write_at = current_unix_secs();
1054 }
1055
1056 Ok(deleted_ids)
1057 }
1058
1059 fn add_memory(&self, bytes: usize) {
1061 self.memory_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
1062 }
1063
1064 fn replace_entity_memory(&self, previous: usize, current: usize) {
1065 if current >= previous {
1066 self.add_memory(current - previous);
1067 } else {
1068 self.release_memory(previous - current);
1069 }
1070 }
1071
1072 pub fn memory_bytes(&self) -> u64 {
1077 self.memory_bytes.load(Ordering::Relaxed)
1078 + self.zone_payload_bytes as u64
1079 + self.unique_key_lookup_bytes.load(Ordering::Relaxed)
1080 + self
1081 .key_lookup
1082 .as_ref()
1083 .map_or(0, SegmentKeyLookup::memory_bytes)
1084 + self.graph_read_index_bytes.load(Ordering::Relaxed)
1085 + (self.hash_scan_ids.capacity() * std::mem::size_of::<EntityId>()) as u64
1086 }
1087
1088 fn release_memory(&self, bytes: usize) {
1091 let mut current = self.memory_bytes.load(Ordering::Relaxed);
1092 loop {
1093 let next = current.saturating_sub(bytes as u64);
1094 match self.memory_bytes.compare_exchange_weak(
1095 current,
1096 next,
1097 Ordering::Relaxed,
1098 Ordering::Relaxed,
1099 ) {
1100 Ok(_) => break,
1101 Err(observed) => current = observed,
1102 }
1103 }
1104 }
1105
1106 pub fn resident_entity_count(&self) -> usize {
1108 self.flat_entities.len() + self.entities.len()
1109 }
1110
1111 pub fn live_entity_count(&self) -> usize {
1113 self.resident_entity_count()
1114 .saturating_sub(self.dead_resident_count)
1115 }
1116
1117 pub fn tombstone_count(&self) -> usize {
1120 self.deleted.len()
1121 }
1122
1123 pub fn resident_bytes(&self) -> u64 {
1126 self.memory_bytes() + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1127 }
1128
1129 pub fn reclaimable_bytes(&self) -> u64 {
1132 self.dead_entity_bytes + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1133 }
1134
1135 pub fn live_entity_ids(&self) -> Vec<EntityId> {
1141 let mut ids = Vec::with_capacity(self.live_entity_count());
1142 self.for_each_fast(|entity| {
1143 ids.push(entity.id);
1144 true
1145 });
1146 ids
1147 }
1148
1149 pub fn insert_preserving_sequence(
1153 &mut self,
1154 entity: UnifiedEntity,
1155 ) -> Result<EntityId, SegmentError> {
1156 if !self.state.is_writable() {
1157 return Err(SegmentError::NotWritable);
1158 }
1159 if self.entities.contains_key(&entity.id) {
1160 return Err(SegmentError::AlreadyExists(entity.id));
1161 }
1162 let id = entity.id;
1163 self.adopt_entity(entity, None);
1164 self.last_write_at = current_unix_secs();
1165 Ok(id)
1166 }
1167
1168 pub fn adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
1176 let id = entity.id;
1177
1178 let next = entity.sequence_id.saturating_add(1);
1181 if self.sequence.load(Ordering::Relaxed) < next {
1182 self.sequence.store(next, Ordering::Relaxed);
1183 }
1184
1185 self.add_memory(Self::estimate_entity_size(&entity));
1186 self.index_entity(&entity);
1187 self.update_col_zones_from_entity(&entity);
1188 self.entities.insert(id, entity);
1189
1190 if let Some(metadata) = metadata {
1191 if !metadata.is_empty() {
1192 self.metadata.set_all(id, &metadata);
1193 }
1194 }
1195 }
1196
1197 pub fn evict_entity(&mut self, id: EntityId) -> bool {
1203 self.mark_checkpoint_entity_dirty(id);
1204 self.invalidate_unique_key_lookup();
1205 let Some(entity) = self.entities.remove(&id) else {
1206 return false;
1207 };
1208 self.release_memory(Self::estimate_entity_size(&entity));
1209 self.unindex_entity(&entity);
1210 for cross_ref in entity.cross_refs() {
1211 if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1212 reverse.retain(|(source, _)| *source != id);
1213 }
1214 }
1215 self.metadata.remove_all(id);
1216 true
1217 }
1218
1219 fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
1221 super::memory_size::entity_bytes(entity)
1222 }
1223
1224 fn update_column_zone(&mut self, column: &str, value: &Value) {
1232 if matches!(value, Value::Null) {
1233 return;
1234 }
1235 match self.col_zones.entry(column.to_string()) {
1236 std::collections::hash_map::Entry::Occupied(mut entry) => {
1237 let before = entry.get().payload_bytes();
1238 entry.get_mut().update(value);
1239 self.zone_payload_bytes = self
1240 .zone_payload_bytes
1241 .saturating_sub(before)
1242 .saturating_add(entry.get().payload_bytes());
1243 }
1244 std::collections::hash_map::Entry::Vacant(entry) => {
1245 let zone = ColZone::new(value.clone());
1246 self.zone_payload_bytes =
1247 self.zone_payload_bytes.saturating_add(zone.payload_bytes());
1248 entry.insert(zone);
1249 }
1250 }
1251 }
1252
1253 fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
1254 if let EntityData::Row(row) = &entity.data {
1255 if let Some(named) = &row.named {
1256 for (column, value) in named {
1257 self.update_column_zone(column, value);
1258 }
1259 } else if let Some(schema) = &row.schema {
1260 for (column, value) in schema.iter().zip(&row.columns) {
1261 self.update_column_zone(column, value);
1262 }
1263 }
1264 }
1265 }
1266
1267 fn rebuild_sealed_col_zones(&mut self) {
1268 let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
1269 let mut family_by_col: HashMap<String, reddb_types::CanonicalKeyFamily> = HashMap::new();
1270 let mut mixed_family_cols = HashSet::new();
1271 let mut unsupported_cols = HashSet::new();
1272
1273 let mut observe_row = |row: &super::entity::RowData| {
1274 for (col, value) in row.iter_fields() {
1275 if matches!(value, Value::Null) {
1276 continue;
1277 }
1278 let Some(key) = value_to_canonical_key(value) else {
1279 unsupported_cols.insert(col.to_string());
1280 continue;
1281 };
1282 match family_by_col.get(col).copied() {
1283 Some(existing) if existing != key.family() => {
1284 mixed_family_cols.insert(col.to_string());
1285 }
1286 None => {
1287 family_by_col.insert(col.to_string(), key.family());
1288 }
1289 _ => {}
1290 }
1291 values_by_col
1292 .entry(col.to_string())
1293 .or_default()
1294 .push((key, value.clone()));
1295 }
1296 };
1297
1298 if self.use_flat {
1299 for entity in &self.flat_entities {
1300 if self.deleted.contains_key(&entity.id) {
1301 continue;
1302 }
1303 if let EntityData::Row(row) = &entity.data {
1304 observe_row(row);
1305 }
1306 }
1307 } else {
1308 for entity in self.entities.values() {
1309 if self.deleted.contains_key(&entity.id) {
1310 continue;
1311 }
1312 if let EntityData::Row(row) = &entity.data {
1313 observe_row(row);
1314 }
1315 }
1316 }
1317
1318 let mut sealed_col_zones = HashMap::new();
1319 for (col, mut entries) in values_by_col {
1320 if mixed_family_cols.contains(&col)
1321 || unsupported_cols.contains(&col)
1322 || entries.is_empty()
1323 {
1324 continue;
1325 }
1326 entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
1327 entries.dedup_by(|left, right| left.0 == right.0);
1328
1329 let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
1330 if intervals.len() > 1 {
1331 sealed_col_zones.insert(col, MultiColZone { intervals });
1332 }
1333 }
1334
1335 let previous = self
1336 .sealed_col_zones
1337 .values()
1338 .flat_map(|zones| &zones.intervals)
1339 .map(ColZone::payload_bytes)
1340 .fold(0, usize::saturating_add);
1341 let current = sealed_col_zones
1342 .values()
1343 .flat_map(|zones| &zones.intervals)
1344 .map(ColZone::payload_bytes)
1345 .fold(0, usize::saturating_add);
1346 self.zone_payload_bytes = self
1347 .zone_payload_bytes
1348 .saturating_sub(previous)
1349 .saturating_add(current);
1350 self.sealed_col_zones = sealed_col_zones;
1351 }
1352
1353 pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
1357 if preds.is_empty() {
1358 return false;
1359 }
1360 for (col, pred) in preds {
1361 if let Some(zone) = self.sealed_col_zones.get(*col) {
1362 if zone.can_skip(pred) {
1363 return true;
1364 }
1365 continue;
1366 }
1367 if let Some(zone) = self.col_zones.get(*col) {
1368 if pred.can_skip(zone) {
1369 return true; }
1371 }
1372 }
1373 false
1374 }
1375
1376 fn index_entity(&mut self, entity: &UnifiedEntity) {
1378 self.mark_checkpoint_entity_dirty(entity.id);
1379 if let Some(index) = self.key_lookup.as_mut() {
1380 index.insert(entity);
1381 }
1382 self.hash_scan_ids.push(entity.id);
1383 self.observe_graph_kind(&entity.kind);
1384 if let Some(index) = self.graph_read_index.get_mut() {
1385 index.insert(entity);
1386 self.graph_read_index_bytes
1387 .store(index.memory_bytes(), Ordering::Relaxed);
1388 }
1389 let indexes = self.unique_key_lookup.get_mut();
1390 for index in indexes.iter_mut() {
1391 index.insert(entity);
1392 }
1393 self.unique_key_lookup_bytes.store(
1394 indexes
1395 .iter()
1396 .map(|index| index.memory_bytes() as u64)
1397 .sum(),
1398 Ordering::Relaxed,
1399 );
1400 let kind_key = entity.kind.storage_type().to_string();
1402 self.kind_index
1403 .entry(kind_key)
1404 .or_default()
1405 .insert(entity.id);
1406
1407 if let EntityData::Row(row) = &entity.data {
1409 if let Some(first_col) = row.columns.first() {
1410 let pk_str = format!("{:?}", first_col);
1411 self.pk_index
1412 .insert((entity.kind.collection().to_string(), pk_str), entity.id);
1413 }
1414 }
1415
1416 for cross_ref in entity.cross_refs() {
1418 self.cross_ref_forward
1419 .entry(cross_ref.source)
1420 .or_default()
1421 .push((cross_ref.target, cross_ref.ref_type));
1422
1423 self.cross_ref_reverse
1424 .entry(cross_ref.target)
1425 .or_default()
1426 .push((cross_ref.source, cross_ref.ref_type));
1427 }
1428 }
1429
1430 fn invalidate_graph_read_index(&mut self) {
1431 *self.graph_read_index.get_mut() = None;
1432 self.graph_read_index_bytes.store(0, Ordering::Relaxed);
1433 }
1434
1435 #[cfg(test)]
1438 pub fn visit_graph_index_candidates(
1439 &self,
1440 kind: GraphEntityKind,
1441 keys: &[EntityId],
1442 before_work: &mut impl FnMut() -> bool,
1443 visit: &mut impl FnMut(&UnifiedEntity) -> bool,
1444 ) -> bool {
1445 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1446 return false;
1447 };
1448 let index = cached.as_ref().expect("complete graph index");
1449 for key in keys {
1450 if !before_work() {
1451 return false;
1452 }
1453 for id in index.candidates(kind, *key) {
1454 if !before_work() {
1455 return false;
1456 }
1457 if let Some(entity) = self.get(id) {
1458 if !visit(entity) {
1459 return false;
1460 }
1461 }
1462 }
1463 }
1464 true
1465 }
1466
1467 fn graph_read_index_for_query(
1468 &self,
1469 before_work: &mut impl FnMut() -> bool,
1470 ) -> Option<parking_lot::RwLockReadGuard<'_, Option<SegmentGraphIndex>>> {
1471 let mut cached = self.graph_read_index.read();
1472 if cached.is_none() {
1473 drop(cached);
1474 let mut writer = self.graph_read_index.write();
1475 if writer.is_none() {
1476 let mut index = SegmentGraphIndex::default();
1477 if !self.for_each_fast(|entity| {
1480 if !before_work() {
1481 return false;
1482 }
1483 index.insert(entity);
1484 true
1485 }) {
1486 return None;
1487 }
1488 self.graph_read_index_bytes
1489 .store(index.memory_bytes(), Ordering::Relaxed);
1490 *writer = Some(index);
1491 }
1492 cached = parking_lot::RwLockWriteGuard::downgrade(writer);
1493 }
1494 Some(cached)
1495 }
1496
1497 pub fn graph_index_cursor(
1498 &self,
1499 kind: GraphEntityKind,
1500 keys: &[EntityId],
1501 before_work: &mut impl FnMut() -> bool,
1502 ) -> Option<GraphIndexCursor> {
1503 let cached = self.graph_read_index_for_query(before_work)?;
1504 let index = cached.as_ref().expect("complete graph index");
1505 let mut upper = EntityId::new(0);
1506 for key in keys {
1507 if !before_work() {
1508 return None;
1509 }
1510 if let Some(last) = index.candidates(kind, *key).next_back() {
1511 upper = upper.max(last);
1512 }
1513 }
1514 Some(GraphIndexCursor {
1515 key_index: 0,
1516 after: None,
1517 upper,
1518 })
1519 }
1520
1521 pub fn graph_index_batch(
1524 &self,
1525 kind: GraphEntityKind,
1526 keys: &[EntityId],
1527 cursor: &mut GraphIndexCursor,
1528 before_work: &mut impl FnMut() -> bool,
1529 visit: &mut impl FnMut(EntityId),
1530 ) -> bool {
1531 self.graph_index_batch_with_limit(kind, keys, cursor, before_work, visit, SCAN_BATCH_SIZE)
1532 }
1533
1534 pub fn graph_index_batch_with_limit(
1535 &self,
1536 kind: GraphEntityKind,
1537 keys: &[EntityId],
1538 cursor: &mut GraphIndexCursor,
1539 before_work: &mut impl FnMut() -> bool,
1540 visit: &mut impl FnMut(EntityId),
1541 limit: usize,
1542 ) -> bool {
1543 assert!(
1544 limit > 0 && limit <= SCAN_BATCH_SIZE,
1545 "bounded graph index batch"
1546 );
1547 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1548 return false;
1549 };
1550 let index = cached.as_ref().expect("complete graph index");
1551 let mut count = 0;
1552 while !cursor.finished(keys) && count < limit {
1553 let key = keys[cursor.key_index];
1554 if !before_work() {
1555 return false;
1556 }
1557 for id in index
1558 .candidates_after(kind, key, cursor.after, cursor.upper)
1559 .take(limit - count)
1560 {
1561 if !before_work() {
1562 return false;
1563 }
1564 cursor.after = Some(id);
1565 count += 1;
1566 visit(id);
1567 }
1568 if count == limit && cursor.after != Some(cursor.upper) {
1569 break;
1570 }
1571 cursor.key_index += 1;
1572 cursor.after = None;
1573 }
1574 true
1575 }
1576
1577 pub fn find_unique_key_conflict(
1582 &self,
1583 columns: &[Vec<String>],
1584 key_index: usize,
1585 signatures: &[String],
1586 exclude_id: Option<EntityId>,
1587 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1588 ) -> Option<EntityId> {
1589 assert!(
1590 key_index < columns.len(),
1591 "constraint index belongs to current schema"
1592 );
1593 let cached = self.unique_key_lookup.read();
1594 if cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1595 return self.find_unique_key_candidate(
1596 &cached[key_index],
1597 signatures,
1598 exclude_id,
1599 reserves_key,
1600 );
1601 }
1602 drop(cached);
1603 let mut cached = self.unique_key_lookup.write();
1604 if !cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1605 let mut indexes: Vec<_> = columns
1606 .iter()
1607 .map(|columns| SegmentUniqueKey::new(columns))
1608 .collect();
1609 self.for_each_fast(|entity| {
1610 for index in &mut indexes {
1611 index.insert(entity);
1612 }
1613 true
1614 });
1615 self.unique_key_lookup_bytes.store(
1616 indexes
1617 .iter()
1618 .map(|index| index.memory_bytes() as u64)
1619 .sum(),
1620 Ordering::Relaxed,
1621 );
1622 *cached = indexes;
1623 }
1624 self.find_unique_key_candidate(&cached[key_index], signatures, exclude_id, reserves_key)
1625 }
1626
1627 fn find_unique_key_candidate(
1628 &self,
1629 index: &SegmentUniqueKey,
1630 signatures: &[String],
1631 exclude_id: Option<EntityId>,
1632 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1633 ) -> Option<EntityId> {
1634 index.candidates(signatures).iter().copied().find(|id| {
1635 if Some(*id) == exclude_id {
1636 return false;
1637 }
1638 let Some(
1639 entity @ UnifiedEntity {
1640 data: EntityData::Row(row),
1641 ..
1642 },
1643 ) = self.get(*id)
1644 else {
1645 return false;
1646 };
1647 index
1648 .columns
1649 .iter()
1650 .zip(signatures)
1651 .all(|(column, signature)| {
1652 row.get_field(column)
1653 .is_some_and(|value| format!("{value:?}") == *signature)
1654 })
1655 && reserves_key(entity)
1656 })
1657 }
1658
1659 fn invalidate_unique_key_lookup(&mut self) {
1660 *self.unique_key_lookup.get_mut() = Vec::new();
1661 self.unique_key_lookup_bytes.store(0, Ordering::Relaxed);
1662 }
1663
1664 pub fn find_key_entity(
1667 &self,
1668 key: &str,
1669 newest: bool,
1670 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1671 ) -> Option<&UnifiedEntity> {
1672 let index = self
1675 .key_lookup
1676 .as_ref()
1677 .expect("segment owns its key index");
1678 self.find_key_candidate(index, key, newest, visible)
1679 }
1680
1681 fn find_key_candidate(
1682 &self,
1683 index: &SegmentKeyLookup,
1684 key: &str,
1685 newest: bool,
1686 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1687 ) -> Option<&UnifiedEntity> {
1688 self.select_key_entity(
1689 index
1690 .candidates(key)
1691 .iter()
1692 .chain(index.mutated_candidates())
1693 .filter_map(|id| self.get(*id)),
1694 key,
1695 newest,
1696 visible,
1697 )
1698 }
1699
1700 fn select_key_entity<'a>(
1701 &self,
1702 entities: impl Iterator<Item = &'a UnifiedEntity>,
1703 key: &str,
1704 newest: bool,
1705 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1706 ) -> Option<&'a UnifiedEntity> {
1707 let mut best: Option<&UnifiedEntity> = None;
1708 for entity in entities {
1709 if SegmentKeyLookup::key(entity) != Some(key) || !visible(entity) {
1711 continue;
1712 }
1713 if !newest {
1714 return Some(entity);
1715 }
1716 if best.is_none_or(|current| entity.xmin >= current.xmin) {
1717 best = Some(entity);
1718 }
1719 }
1720 best
1721 }
1722
1723 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1725 if let Some(id) = entity_id_from_probe_key(key) {
1726 return self.has_live_entity(id);
1727 }
1728
1729 let Ok(pk_str) = std::str::from_utf8(key) else {
1730 return false;
1731 };
1732 let pk_key = (self.collection.clone(), pk_str.to_string());
1733 self.pk_index
1734 .get(&pk_key)
1735 .is_some_and(|id| self.has_live_entity(*id))
1736 }
1737
1738 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1740 if let Some(index) = self.key_lookup.as_mut() {
1741 index.remove(entity);
1742 }
1743 if let Some(index) = self.graph_read_index.get_mut() {
1744 index.remove(entity);
1745 self.graph_read_index_bytes
1746 .store(index.memory_bytes(), Ordering::Relaxed);
1747 }
1748 let kind_key = entity.kind.storage_type().to_string();
1750 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1751 set.remove(&entity.id);
1752 }
1753
1754 if let EntityData::Row(row) = &entity.data {
1756 if let Some(first_col) = row.columns.first() {
1757 let pk_str = format!("{:?}", first_col);
1758 self.pk_index
1759 .remove(&(entity.kind.collection().to_string(), pk_str));
1760 }
1761 }
1762
1763 self.cross_ref_forward.remove(&entity.id);
1765 }
1767
1768 fn reindex_for_update(
1778 &mut self,
1779 old: &UpdateIndexSnapshot,
1780 new: &UnifiedEntity,
1781 modified_columns: Option<&[String]>,
1782 ) {
1783 let pk_changed = match &new.data {
1788 EntityData::Row(new_row) => {
1789 if let Some(cols) = modified_columns {
1790 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1793 new_row
1794 .schema
1795 .as_deref()
1796 .and_then(|schema| schema.first().map(|name| name.as_str()))
1797 });
1798 match pk_col_name {
1799 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1800 None => old.pk_value.as_ref() != new_row.columns.first(),
1802 }
1803 } else {
1804 old.pk_value.as_ref() != new_row.columns.first()
1805 }
1806 }
1807 _ => false,
1809 };
1810
1811 if pk_changed {
1812 if let Some((collection, pk_str)) = &old.pk_index_key {
1814 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1815 }
1816 if let EntityData::Row(row) = &new.data {
1818 if let Some(first_col) = row.columns.first() {
1819 let pk_str = format!("{:?}", first_col);
1820 self.pk_index
1821 .insert((new.kind.collection().to_string(), pk_str), new.id);
1822 }
1823 }
1824 }
1825
1826 let new_refs = new.cross_refs();
1828 if old.cross_refs.as_slice() != new_refs {
1829 self.cross_ref_forward.remove(&new.id);
1831 for cross_ref in &old.cross_refs {
1833 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1834 rev.retain(|(src, _)| *src != new.id);
1835 }
1836 }
1837 for cross_ref in new_refs {
1839 self.cross_ref_forward
1840 .entry(cross_ref.source)
1841 .or_default()
1842 .push((cross_ref.target, cross_ref.ref_type));
1843 self.cross_ref_reverse
1844 .entry(cross_ref.target)
1845 .or_default()
1846 .push((cross_ref.source, cross_ref.ref_type));
1847 }
1848 }
1849 }
1850
1851 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1853 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1854 }
1855
1856 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1858 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1859 }
1860
1861 pub fn age_secs(&self) -> u64 {
1863 let now = current_unix_secs();
1864 now.saturating_sub(self.created_at)
1865 }
1866
1867 pub fn idle_secs(&self) -> u64 {
1869 let now = current_unix_secs();
1870 now.saturating_sub(self.last_write_at)
1871 }
1872
1873 pub fn bulk_insert(
1882 &mut self,
1883 entities: Vec<UnifiedEntity>,
1884 ) -> Result<Vec<EntityId>, SegmentError> {
1885 if !self.state.is_writable() {
1886 return Err(SegmentError::NotWritable);
1887 }
1888
1889 for entity in &entities {
1890 self.mark_checkpoint_entity_dirty(entity.id);
1891 }
1892
1893 if let Some(index) = self.key_lookup.as_mut() {
1894 for entity in &entities {
1895 index.insert(entity);
1896 }
1897 }
1898
1899 let indexes = self.unique_key_lookup.get_mut();
1900 for index in indexes.iter_mut() {
1901 for entity in &entities {
1902 index.insert(entity);
1903 }
1904 }
1905 self.unique_key_lookup_bytes.store(
1906 indexes
1907 .iter()
1908 .map(|index| index.memory_bytes() as u64)
1909 .sum(),
1910 Ordering::Relaxed,
1911 );
1912
1913 let n = entities.len();
1914
1915 if entities.is_empty() {
1916 return Ok(Vec::new());
1917 }
1918
1919 for run in
1922 entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
1923 {
1924 let kind = run[0].kind.storage_type();
1925 self.observe_graph_kind(&run[0].kind);
1926 if matches!(
1927 run[0].kind,
1928 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
1929 ) {
1930 if let Some(index) = self.graph_read_index.get_mut() {
1931 for entity in run {
1932 index.insert(entity);
1933 }
1934 self.graph_read_index_bytes
1935 .store(index.memory_bytes(), Ordering::Relaxed);
1936 }
1937 }
1938 let kind_set = match self.kind_index.get_mut(kind) {
1939 Some(ids) => ids,
1940 None => self.kind_index.entry(kind.to_string()).or_default(),
1941 };
1942 kind_set.reserve(run.len());
1943 kind_set.extend(run.iter().map(|entity| entity.id));
1944 }
1945
1946 let now = current_unix_secs();
1947
1948 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
1949
1950 let mut ids = Vec::with_capacity(n);
1951
1952 if self.flat_entities.is_empty() && self.entities.is_empty() {
1954 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
1956 self.use_flat = true;
1957 }
1958
1959 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
1969 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
1970 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
1974
1975 let mut batch_bytes: usize = 0;
1980
1981 if self.use_flat {
1982 self.flat_entities.reserve(n);
1983 for (i, mut entity) in entities.into_iter().enumerate() {
1984 entity.sequence_id = base_seq + i as u64;
1985 let id = entity.id;
1986 ids.push(id);
1987 batch_bytes += Self::estimate_entity_size(&entity);
1988 if let EntityData::Row(row) = &entity.data {
1989 if row.schema.is_some() && !row.columns.is_empty() {
1990 if columnar_zone_updates.is_empty() {
1991 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
1992 columnar_schema = row.schema.clone();
1993 }
1994 for (ci, val) in row.columns.iter().enumerate() {
1995 if !matches!(val, Value::Null) {
1996 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
1997 bucket.push(val.clone());
1998 }
1999 }
2000 }
2001 } else {
2002 for (col, val) in row.iter_fields() {
2003 if !matches!(val, Value::Null) {
2004 named_zone_updates.push((col.to_string(), val.clone()));
2005 }
2006 }
2007 }
2008 }
2009 let expected = self.base_entity_id + self.flat_entities.len() as u64;
2025 if id.raw() == expected {
2026 self.flat_entities.push(entity);
2027 } else {
2028 self.hash_scan_ids.push(id);
2029 self.entities.insert(id, entity);
2030 }
2031 }
2032 } else {
2033 self.entities.reserve(n);
2035 self.hash_scan_ids.reserve(n);
2036 let mut pairs = Vec::with_capacity(n);
2037 for (i, mut entity) in entities.into_iter().enumerate() {
2038 entity.sequence_id = base_seq + i as u64;
2039 let id = entity.id;
2040 ids.push(id);
2041 batch_bytes += Self::estimate_entity_size(&entity);
2042 if let EntityData::Row(row) = &entity.data {
2043 for (col, val) in row.iter_fields() {
2044 if !matches!(val, Value::Null) {
2045 named_zone_updates.push((col.to_string(), val.clone()));
2046 }
2047 }
2048 }
2049 self.hash_scan_ids.push(id);
2050 pairs.push((id, entity));
2051 }
2052 self.entities.extend(pairs);
2053 }
2054
2055 if !columnar_zone_updates.is_empty() {
2059 let schema = columnar_schema.as_ref();
2060 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
2061 if values.is_empty() {
2062 continue;
2063 }
2064 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
2065 continue;
2066 };
2067 let mut iter = values.into_iter();
2073 if let Some(first) = iter.next() {
2074 let previous = self
2075 .col_zones
2076 .get(col_name)
2077 .map_or(0, ColZone::payload_bytes);
2078 let zone = self
2079 .col_zones
2080 .entry(col_name.clone())
2081 .and_modify(|z| z.update(&first))
2082 .or_insert_with(|| ColZone::new(first));
2083 for v in iter {
2084 zone.update(&v);
2085 }
2086 self.zone_payload_bytes = self
2087 .zone_payload_bytes
2088 .saturating_sub(previous)
2089 .saturating_add(zone.payload_bytes());
2090 }
2091 }
2092 }
2093 for (col, val) in named_zone_updates {
2094 self.update_column_zone(&col, &val);
2095 }
2096
2097 self.add_memory(batch_bytes);
2098 self.last_write_at = now;
2099
2100 if self.use_flat {
2102 self.published_flat_len
2103 .store(self.flat_entities.len(), Ordering::Release);
2104 }
2105
2106 Ok(ids)
2107 }
2108
2109 pub fn force_delete(&mut self, id: EntityId) -> bool {
2112 if self.use_flat {
2113 let raw = id.raw();
2114 if raw >= self.base_entity_id {
2115 let idx = (raw - self.base_entity_id) as usize;
2116 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2117 self.tombstone_flat_entity(idx, id);
2118 return true;
2119 }
2120 }
2121 return self.remove_hashmap_entity(id);
2122 }
2123
2124 self.remove_hashmap_entity(id)
2125 }
2126
2127 pub fn force_update_with_metadata(
2130 &mut self,
2131 entity: &UnifiedEntity,
2132 modified_columns: &[String],
2133 metadata: Option<&Metadata>,
2134 ) -> Result<(), SegmentError> {
2135 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
2136 }
2137}
2138
2139impl UnifiedSegment for GrowingSegment {
2140 fn id(&self) -> SegmentId {
2141 self.id
2142 }
2143
2144 fn state(&self) -> SegmentState {
2145 self.state
2146 }
2147
2148 fn collection(&self) -> &str {
2149 &self.collection
2150 }
2151
2152 fn stats(&self) -> SegmentStats {
2153 let mut stats = SegmentStats {
2154 entity_count: self.live_entity_count(),
2155 deleted_count: self.deleted.len(),
2156 memory_bytes: self.resident_bytes() as usize,
2157 ..Default::default()
2158 };
2159
2160 for entity in self.entities.values() {
2161 match &entity.kind {
2162 EntityKind::TableRow { .. } => stats.row_count += 1,
2163 EntityKind::GraphNode(_) => stats.node_count += 1,
2164 EntityKind::GraphEdge(_) => stats.edge_count += 1,
2165 EntityKind::Vector { .. } => stats.vector_count += 1,
2166 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
2167 EntityKind::QueueMessage { .. } => stats.row_count += 1,
2168 }
2169 stats.cross_ref_count += entity.cross_refs().len();
2170 }
2171
2172 stats
2173 }
2174
2175 fn entity_count(&self) -> usize {
2176 self.live_entity_count()
2177 }
2178
2179 fn contains(&self, id: EntityId) -> bool {
2180 self.has_live_entity(id)
2181 }
2182
2183 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2184 if self.deleted.contains_key(&id) {
2185 return None;
2186 }
2187 if self.use_flat {
2188 let raw = id.raw();
2189 if raw >= self.base_entity_id {
2190 let idx = (raw - self.base_entity_id) as usize;
2191 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
2192 return Some(entity);
2193 }
2194 }
2195 self.entities.get(&id)
2201 } else {
2202 self.entities.get(&id)
2203 }
2204 }
2205
2206 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
2207 self.mark_checkpoint_entity_dirty(id);
2208 self.invalidate_unique_key_lookup();
2209 self.invalidate_graph_read_index();
2210 if self.deleted.contains_key(&id) || !self.state.is_writable() {
2211 return None;
2212 }
2213 if let Some(mut index) = self.key_lookup.take() {
2214 if let Some(entity) = self.get(id) {
2215 index.mark_mutated(entity);
2216 }
2217 self.key_lookup = Some(index);
2218 }
2219 self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
2222 if self.use_flat {
2223 let raw = id.raw();
2224 if raw >= self.base_entity_id {
2225 let idx = (raw - self.base_entity_id) as usize;
2226 if self
2227 .flat_entities
2228 .get(idx)
2229 .map(|e| e.id == id)
2230 .unwrap_or(false)
2231 {
2232 return self.flat_entities.get_mut(idx);
2233 }
2234 }
2235 self.entities.get_mut(&id)
2238 } else {
2239 self.entities.get_mut(&id)
2240 }
2241 }
2242
2243 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2244 if !self.state.is_writable() {
2245 return Err(SegmentError::NotWritable);
2246 }
2247
2248 if self.entities.contains_key(&entity.id) {
2249 return Err(SegmentError::AlreadyExists(entity.id));
2250 }
2251
2252 entity.sequence_id = self.next_sequence();
2254
2255 let size = Self::estimate_entity_size(&entity);
2257 self.add_memory(size);
2258
2259 self.index_entity(&entity);
2261
2262 self.update_col_zones_from_entity(&entity);
2264
2265 let id = entity.id;
2267 self.entities.insert(id, entity);
2268
2269 self.last_write_at = current_unix_secs();
2271
2272 Ok(id)
2273 }
2274
2275 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2276 if !self.state.is_writable() {
2277 return Err(SegmentError::NotWritable);
2278 }
2279
2280 self.apply_update_with_metadata(&entity, None)?;
2281 self.last_write_at = current_unix_secs();
2282
2283 Ok(())
2284 }
2285
2286 fn update_hot(
2287 &mut self,
2288 entity: UnifiedEntity,
2289 modified_columns: &[String],
2290 ) -> Result<(), SegmentError> {
2291 if !self.state.is_writable() {
2292 return Err(SegmentError::NotWritable);
2293 }
2294
2295 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
2296 self.last_write_at = current_unix_secs();
2297 Ok(())
2298 }
2299
2300 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2301 if !self.state.is_writable() {
2302 return Err(SegmentError::NotWritable);
2303 }
2304
2305 if self.use_flat {
2307 let raw = id.raw();
2308 if raw >= self.base_entity_id {
2309 let idx = (raw - self.base_entity_id) as usize;
2310 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2311 self.tombstone_flat_entity(idx, id);
2312 return Ok(true);
2313 }
2314 }
2315 return Ok(self.remove_hashmap_entity(id));
2316 }
2317
2318 Ok(self.remove_hashmap_entity(id))
2319 }
2320
2321 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2322 if !self.has_live_entity(id) {
2323 return None;
2324 }
2325 Some(self.metadata.get_all(id))
2326 }
2327
2328 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2329 self.mark_checkpoint_entity_dirty(id);
2330 if !self.state.is_writable() {
2331 return Err(SegmentError::NotWritable);
2332 }
2333
2334 if !self.has_live_entity(id) {
2335 return Err(SegmentError::NotFound(id));
2336 }
2337
2338 self.metadata.set_all(id, &metadata);
2339 Ok(())
2340 }
2341
2342 fn seal(&mut self) -> Result<(), SegmentError> {
2343 if self.state != SegmentState::Growing {
2344 return Err(SegmentError::InvalidState(self.state));
2345 }
2346
2347 self.state = SegmentState::Sealing;
2348
2349 self.rebuild_sealed_col_zones();
2354
2355 self.state = SegmentState::Sealed;
2356 Ok(())
2357 }
2358
2359 fn should_seal(&self, config: &SegmentConfig) -> bool {
2360 if self.entities.len() >= config.max_entities {
2362 return true;
2363 }
2364
2365 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
2367 return true;
2368 }
2369
2370 if self.age_secs() >= config.max_age_secs {
2372 return true;
2373 }
2374
2375 false
2376 }
2377
2378 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2379 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2380 Box::new(self.flat_entities.iter().chain(self.entities.values()))
2384 } else {
2385 Box::new(self.entities.values())
2386 };
2387 if self.deleted.is_empty() {
2388 base
2389 } else {
2390 Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
2391 }
2392 }
2393
2394 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2395 let ids = self.kind_index.get(kind_filter);
2398 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2401 Box::new(self.flat_entities.iter())
2402 } else {
2403 Box::new(std::iter::empty())
2404 };
2405 Box::new(flat.chain(self.entities.values()).filter(move |e| {
2406 if self.deleted.contains_key(&e.id) {
2407 return false;
2408 }
2409 if let Some(ids) = ids {
2410 ids.contains(&e.id)
2411 } else {
2412 false
2413 }
2414 }))
2415 }
2416
2417 fn filter_metadata(
2418 &self,
2419 filters: &[(String, super::metadata::MetadataFilter)],
2420 ) -> Vec<EntityId> {
2421 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
2424 Box::new(self.flat_entities.iter().map(|e| e.id))
2425 } else {
2426 Box::new(std::iter::empty())
2427 };
2428 flat_ids
2429 .chain(self.entities.keys().copied())
2430 .filter(|id| {
2431 if self.deleted.contains_key(id) {
2432 return false;
2433 }
2434 let metadata = self.metadata.get_all(*id);
2435 metadata.matches_all(filters)
2436 })
2437 .collect()
2438 }
2439}
2440
2441fn build_minmax_multi_intervals(
2442 entries: &[(CanonicalKey, Value)],
2443 max_intervals: usize,
2444) -> Vec<ColZone> {
2445 if entries.is_empty() {
2446 return Vec::new();
2447 }
2448 if entries.len() == 1 || max_intervals <= 1 {
2449 return vec![ColZone::with_bounds(
2450 entries[0].1.clone(),
2451 entries[entries.len() - 1].1.clone(),
2452 )];
2453 }
2454
2455 let mut split_points = if entries.len() <= max_intervals {
2456 (1..entries.len()).collect::<Vec<_>>()
2457 } else {
2458 let target_splits = max_intervals - 1;
2459 let mut selected = select_gap_split_points(entries, target_splits);
2460 if selected.len() < target_splits {
2461 for bucket in 1..max_intervals {
2462 let idx = bucket * entries.len() / max_intervals;
2463 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
2464 continue;
2465 }
2466 selected.push(idx);
2467 if selected.len() >= target_splits {
2468 break;
2469 }
2470 }
2471 }
2472 selected.sort_unstable();
2473 selected.dedup();
2474 selected
2475 };
2476
2477 split_points.push(entries.len());
2478
2479 let mut out = Vec::with_capacity(split_points.len());
2480 let mut start = 0usize;
2481 for end in split_points {
2482 if end <= start {
2483 continue;
2484 }
2485 out.push(ColZone::with_bounds(
2486 entries[start].1.clone(),
2487 entries[end - 1].1.clone(),
2488 ));
2489 start = end;
2490 }
2491
2492 if out.is_empty() {
2493 out.push(ColZone::with_bounds(
2494 entries[0].1.clone(),
2495 entries[entries.len() - 1].1.clone(),
2496 ));
2497 }
2498
2499 out
2500}
2501
2502fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
2503 let mut gaps = Vec::new();
2504 for idx in 1..entries.len() {
2505 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
2506 if score > 0.0 {
2507 gaps.push((score, idx));
2508 }
2509 }
2510 }
2511 gaps.sort_by(|left, right| {
2512 right
2513 .0
2514 .partial_cmp(&left.0)
2515 .unwrap_or(std::cmp::Ordering::Equal)
2516 .then_with(|| left.1.cmp(&right.1))
2517 });
2518 gaps.into_iter()
2519 .take(max_splits)
2520 .map(|(_, idx)| idx)
2521 .collect()
2522}
2523
2524fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
2525 if left.family() != right.family() {
2526 return None;
2527 }
2528 match (left, right) {
2529 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
2530 Some(r.saturating_sub(*l) as f64)
2531 }
2532 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
2533 Some(r.saturating_sub(*l) as f64)
2534 }
2535 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
2536 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
2537 }
2538 _ => None,
2539 }
2540}
2541
2542impl UnifiedSegment for std::sync::Arc<parking_lot::RwLock<GrowingSegment>> {
2545 fn id(&self) -> SegmentId {
2546 self.read().id()
2547 }
2548
2549 fn state(&self) -> SegmentState {
2550 self.read().state()
2551 }
2552
2553 fn collection(&self) -> &str {
2554 "unknown"
2556 }
2557
2558 fn stats(&self) -> SegmentStats {
2559 self.read().stats()
2560 }
2561
2562 fn entity_count(&self) -> usize {
2563 self.read().entity_count()
2564 }
2565
2566 fn contains(&self, id: EntityId) -> bool {
2567 self.read().contains(id)
2568 }
2569
2570 fn get(&self, _id: EntityId) -> Option<&UnifiedEntity> {
2571 None
2574 }
2575
2576 fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
2577 None
2578 }
2579
2580 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2581 self.write().insert(entity)
2582 }
2583
2584 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2585 self.write().update(entity)
2586 }
2587
2588 fn update_hot(
2589 &mut self,
2590 entity: UnifiedEntity,
2591 modified_columns: &[String],
2592 ) -> Result<(), SegmentError> {
2593 self.write().update_hot(entity, modified_columns)
2594 }
2595
2596 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2597 self.write().delete(id)
2598 }
2599
2600 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2601 self.read().get_metadata(id)
2602 }
2603
2604 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2605 self.write().set_metadata(id, metadata)
2606 }
2607
2608 fn seal(&mut self) -> Result<(), SegmentError> {
2609 self.write().seal()
2610 }
2611
2612 fn should_seal(&self, config: &SegmentConfig) -> bool {
2613 self.read().should_seal(config)
2614 }
2615
2616 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2617 Box::new(std::iter::empty())
2619 }
2620
2621 fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2622 Box::new(std::iter::empty())
2623 }
2624
2625 fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2626 self.read().filter_metadata(filters)
2627 }
2628}
2629
2630#[cfg(test)]
2631mod tests {
2632 use super::*;
2633 use crate::storage::entity::RowData;
2634 use crate::storage::metadata::MetadataValue;
2635 use reddb_types::Value;
2636
2637 fn keyed_row(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
2638 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
2639 entity.data = EntityData::Row(RowData::with_names(
2640 vec![Value::text(key), Value::Integer(id as i64)],
2641 vec!["key".into(), "value".into()],
2642 ));
2643 entity.set_xmin(xmin);
2644 entity.set_xmax(xmax);
2645 entity
2646 }
2647
2648 fn keyed_probe(segment: &GrowingSegment, key: &str) -> Option<EntityId> {
2649 segment
2650 .find_key_entity(key, true, &mut |_| true)
2651 .map(|entity| entity.id)
2652 }
2653
2654 #[test]
2655 fn bulk_checkpoint_journal_preserves_prior_changes_and_ignores_rejected_batches() {
2656 let mut segment = GrowingSegment::new(1, "keys");
2657 segment.insert(keyed_row(1, "one", 1, 0)).expect("row");
2658 segment
2659 .bulk_insert(vec![keyed_row(2, "two", 1, 0), keyed_row(3, "three", 1, 0)])
2660 .expect("batch");
2661 assert_eq!(
2662 segment.checkpoint_mutations_since(0),
2663 Some(vec![EntityId::new(1), EntityId::new(2), EntityId::new(3)])
2664 );
2665 let revision = segment.checkpoint_revision();
2666 segment.bulk_insert(Vec::new()).expect("empty batch");
2667 assert_eq!(segment.checkpoint_revision(), revision);
2668 segment.seal().expect("seal segment");
2669 assert!(matches!(
2670 segment.bulk_insert(vec![keyed_row(4, "four", 1, 0)]),
2671 Err(SegmentError::NotWritable)
2672 ));
2673 assert_eq!(segment.checkpoint_revision(), revision);
2674 assert_eq!(segment.checkpoint_mutations_since(revision), Some(vec![]));
2675 }
2676
2677 #[test]
2678 fn checkpoint_journal_retains_bounded_row_versions_through_repeated_updates_and_gc() {
2679 const ROWS: u64 = 1700;
2680 let mut segment = GrowingSegment::new(1, "keys");
2681 for id in 1..=ROWS {
2682 segment
2683 .insert(keyed_row(id, "export", 1, 0))
2684 .expect("seed export");
2685 }
2686 let baseline = segment.checkpoint_revision();
2687 for id in 1..=ROWS {
2691 segment
2692 .update(keyed_row(id, "export", 1, 2))
2693 .expect("supersede");
2694 let delivered = ROWS + id;
2695 segment
2696 .insert(keyed_row(delivered, "export", 2, 0))
2697 .expect("delivered");
2698 segment
2699 .update(keyed_row(delivered, "export", 2, 3))
2700 .expect("acknowledge");
2701 segment
2702 .insert(keyed_row(2 * ROWS + id, "export", 3, 0))
2703 .expect("final version");
2704 assert_eq!(
2705 segment
2706 .delete_batch(&[EntityId::new(id), EntityId::new(delivered)])
2707 .expect("GC")
2708 .len(),
2709 2
2710 );
2711 }
2712 let changed = segment
2713 .checkpoint_mutations_since(baseline)
2714 .expect("all changed IDs still fit the bounded checkpoint journal");
2715 let actual: std::collections::BTreeSet<_> = changed.into_iter().collect();
2716 let expected: std::collections::BTreeSet<_> = (1..=3 * ROWS).map(EntityId::new).collect();
2717 assert_eq!(actual, expected, "deleted versions must remain tracked too");
2718 }
2719
2720 #[test]
2721 fn checkpoint_journal_preserves_cold_changes_during_hot_row_churn() {
2722 let mut segment = GrowingSegment::new(1, "keys");
2723 segment
2724 .insert(keyed_row(1, "cold", 1, 0))
2725 .expect("cold row");
2726 segment.insert(keyed_row(2, "hot", 1, 0)).expect("hot row");
2727 for xid in 2..=(5 * CHECKPOINT_MUTATIONS_MAX) as u64 {
2728 segment
2729 .update(keyed_row(2, "hot", xid, 0))
2730 .expect("hot edit");
2731 }
2732 let changed: std::collections::BTreeSet<_> = segment
2733 .checkpoint_mutations_since(0)
2734 .expect("two physical IDs never exhaust the distinct-ID bound")
2735 .into_iter()
2736 .collect();
2737 assert_eq!(
2738 changed,
2739 [EntityId::new(1), EntityId::new(2)].into_iter().collect()
2740 );
2741 assert!(
2742 segment.checkpoint_tracking_bytes() <= 1024 * 1024,
2743 "retained proof memory must stay bounded under unlimited repeated edits"
2744 );
2745 }
2746
2747 #[test]
2748 fn checkpoint_journal_complete_suffix_matches_uncompressed_mutation_history() {
2749 let mut segment = GrowingSegment::new(1, "keys");
2750 let mut history = Vec::new();
2751 let mut seed = 0x7b3d_5961_a902_4c85_u64;
2752 for n in 0..3 * CHECKPOINT_MUTATIONS_MAX {
2753 seed ^= seed << 13;
2754 seed ^= seed >> 7;
2755 seed ^= seed << 17;
2756 let raw = if n < CHECKPOINT_MUTATIONS_MAX + 17 {
2757 n as u64 + 1
2758 } else {
2759 seed % (CHECKPOINT_MUTATIONS_MAX as u64 + 17) + 1
2760 };
2761 let id = EntityId::new(raw);
2762 segment.mark_checkpoint_entity_dirty(id);
2763 history.push(id);
2764 }
2765 let now = segment.checkpoint_revision();
2766 assert!(
2767 segment.checkpoint_mutations_since(0).is_none(),
2768 "an evicted ID must still invalidate an older baseline"
2769 );
2770 let mut complete = 0;
2771 for baseline in [0, 1, now / 3, now / 2, now - 2048, now - 128, now] {
2772 if let Some(ids) = segment.checkpoint_mutations_since(baseline) {
2773 let actual: std::collections::BTreeSet<_> = ids.into_iter().collect();
2774 let expected: std::collections::BTreeSet<_> =
2775 history[baseline as usize..].iter().copied().collect();
2776 assert_eq!(actual, expected, "baseline {baseline}");
2777 complete += 1;
2778 }
2779 }
2780 assert!(
2781 complete >= 3,
2782 "recent baselines must retain a complete proof"
2783 );
2784 assert!(segment.checkpoint_mutations_since(now + 1).is_none());
2785 assert!(segment.checkpoint_tracking_bytes() <= 1024 * 1024);
2786 }
2787
2788 #[test]
2789 fn bulk_checkpoint_journal_exhaustion_keeps_a_bounded_complete_suffix() {
2790 let mut segment = GrowingSegment::new(1, "keys");
2791 segment
2792 .bulk_insert(
2793 (1..=CHECKPOINT_MUTATIONS_MAX as u64 + 1)
2794 .map(|id| keyed_row(id, "key", 1, 0))
2795 .collect(),
2796 )
2797 .expect("large batch");
2798 assert!(segment.checkpoint_mutations_since(0).is_none());
2799 assert_eq!(
2800 segment
2801 .checkpoint_mutations_since(1)
2802 .expect("complete suffix")
2803 .len(),
2804 CHECKPOINT_MUTATIONS_MAX
2805 );
2806 let revision = segment.checkpoint_revision();
2807 let id = EntityId::new(CHECKPOINT_MUTATIONS_MAX as u64 + 2);
2808 segment
2809 .bulk_insert(vec![keyed_row(id.raw(), "next", 1, 0)])
2810 .expect("later batch");
2811 assert_eq!(segment.checkpoint_mutations_since(revision), Some(vec![id]));
2812 assert_eq!(segment.checkpoint_mutations.len(), CHECKPOINT_MUTATIONS_MAX);
2813 }
2814
2815 #[test]
2816 fn key_lookup_visits_only_requested_key_versions() {
2817 let mut segment = GrowingSegment::new(1, "keys");
2818 segment
2819 .bulk_insert(
2820 (1..=4096)
2821 .map(|id| keyed_row(id, &format!("other-{id}"), 1, 0))
2822 .collect(),
2823 )
2824 .expect("unrelated rows");
2825 segment
2826 .insert(keyed_row(5000, "wanted", 2, 8))
2827 .expect("old version");
2828 segment
2829 .insert(keyed_row(5001, "wanted", 8, 0))
2830 .expect("new version");
2831 let mut visits = 0;
2832 let result = segment.find_key_entity("wanted", true, &mut |entity| {
2833 visits += 1;
2834 entity.xmin < 8
2835 });
2836 assert_eq!(result.map(|entity| entity.id), Some(EntityId::new(5000)));
2837 assert_eq!(
2838 visits, 2,
2839 "unrelated rows must never reach visibility evaluation"
2840 );
2841 assert_eq!(keyed_probe(&segment, "wanted"), Some(EntityId::new(5001)));
2842 let mut misses = 0;
2843 assert!(segment
2844 .find_key_entity("absent", true, &mut |_| {
2845 misses += 1;
2846 true
2847 })
2848 .is_none());
2849 assert_eq!(misses, 0);
2850
2851 let mut collision = SegmentKeyLookup::default();
2853 collision.insert(&keyed_row(1, "wanted", 1, 0));
2854 assert!(segment
2855 .find_key_candidate(&collision, "wanted", true, &mut |_| true)
2856 .is_none());
2857 }
2858
2859 #[test]
2860 fn raw_mutation_does_not_disable_unrelated_key_candidates() {
2861 let mut segment = GrowingSegment::new(1, "keys");
2862 segment
2863 .bulk_insert(
2864 (1..=4096)
2865 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2866 .collect(),
2867 )
2868 .expect("rows");
2869 for _ in 0..3 {
2870 *segment.get_mut(EntityId::new(1)).expect("mutable") = keyed_row(1, "renamed", 2, 0);
2871 }
2872 let index = segment.key_lookup.as_ref().expect("index retained");
2873 assert_eq!(index.candidates("key-4096"), &[EntityId::new(4096)]);
2874 assert_eq!(index.mutated_candidates().count(), 1);
2875 assert!(index.candidates("key-1").is_empty());
2876 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(1)));
2877 assert_eq!(keyed_probe(&segment, "key-4096"), Some(EntityId::new(4096)));
2878 segment.delete(EntityId::new(1)).expect("delete raw row");
2879 assert_eq!(
2880 segment
2881 .key_lookup
2882 .as_ref()
2883 .unwrap()
2884 .mutated_candidates()
2885 .count(),
2886 0
2887 );
2888 assert_eq!(keyed_probe(&segment, "renamed"), None);
2889 }
2890
2891 #[test]
2892 fn key_lookup_tracks_mutations_sealing_and_adoption() {
2893 let mut segment = GrowingSegment::new(1, "keys");
2894 segment
2895 .bulk_insert(vec![keyed_row(1, "flat", 1, 0)])
2896 .expect("flat row");
2897 segment
2898 .insert(keyed_row(3, "hash", 1, 0))
2899 .expect("hash fallback");
2900 segment
2901 .bulk_insert(vec![keyed_row(5, "gap", 1, 0)])
2902 .expect("bulk gap");
2903 assert!(segment.key_lookup.as_ref().expect("index").memory_bytes() > 0);
2904 assert_eq!(keyed_probe(&segment, "gap"), Some(EntityId::new(5)));
2905 segment
2906 .update(keyed_row(3, "renamed", 2, 0))
2907 .expect("rename key");
2908 assert!(
2909 segment.key_lookup.is_some(),
2910 "key replacements maintain the index"
2911 );
2912 assert_eq!(keyed_probe(&segment, "hash"), None);
2913 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(3)));
2914
2915 *segment.get_mut(EntityId::new(3)).expect("mutable row") = keyed_row(3, "raw", 3, 0);
2916 assert!(
2917 segment.key_lookup.is_some(),
2918 "raw mutation preserves unaffected key candidates"
2919 );
2920 assert_eq!(keyed_probe(&segment, "renamed"), None);
2921 assert_eq!(keyed_probe(&segment, "raw"), Some(EntityId::new(3)));
2922 segment.seal().expect("seal");
2923 segment
2924 .force_update_with_metadata(&keyed_row(5, "sealed", 4, 0), &["key".into()], None)
2925 .expect("sealed rename");
2926 assert_eq!(keyed_probe(&segment, "gap"), None);
2927 assert_eq!(keyed_probe(&segment, "sealed"), Some(EntityId::new(5)));
2928 assert!(segment.force_delete(EntityId::new(1)));
2929 assert!(segment.force_delete(EntityId::new(3)));
2930 assert_eq!(keyed_probe(&segment, "flat"), None);
2931 assert_eq!(keyed_probe(&segment, "raw"), None);
2932
2933 let mut adopted = GrowingSegment::new(2, "keys");
2934 segment.for_each_fast(|entity| {
2935 adopted.adopt_entity(entity.clone(), None);
2936 true
2937 });
2938 assert_eq!(keyed_probe(&adopted, "sealed"), Some(EntityId::new(5)));
2939 assert!(adopted.evict_entity(EntityId::new(5)));
2940 assert_eq!(keyed_probe(&adopted, "sealed"), None);
2941 }
2942
2943 #[test]
2944 #[ignore = "diagnostic microbenchmark; no timing assertions"]
2945 fn key_lookup_scaling_diagnostic() {
2946 use std::{hint::black_box, time::Instant};
2947 for count in [100, 10_000, 100_000] {
2948 let mut segment = GrowingSegment::new(1, "keys");
2949 segment
2950 .bulk_insert(
2951 (1..=count)
2952 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2953 .collect(),
2954 )
2955 .expect("rows");
2956 let key = format!("key-{count}");
2957 let started = Instant::now();
2958 for _ in 0..1000 {
2959 assert_eq!(
2960 keyed_probe(black_box(&segment), black_box(&key)),
2961 Some(EntityId::new(count))
2962 );
2963 }
2964 let indexed = started.elapsed();
2965 let started = Instant::now();
2966 for _ in 0..1000 {
2967 let mut found = None;
2968 segment.for_each_fast(|entity| {
2969 if SegmentKeyLookup::key(entity) == Some(black_box(key.as_str())) {
2970 found = Some(entity.id);
2971 return false;
2972 }
2973 true
2974 });
2975 assert_eq!(black_box(found), Some(EntityId::new(count)));
2976 }
2977 eprintln!(
2978 "rows={count} reads=1000 indexed={indexed:?} scan={:?}",
2979 started.elapsed()
2980 );
2981 }
2982 }
2983
2984 fn pruning_node(id: u64) -> UnifiedEntity {
2985 UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
2986 }
2987
2988 fn pruning_edge(id: u64) -> UnifiedEntity {
2989 UnifiedEntity::graph_edge(
2990 EntityId::new(id),
2991 "mixed",
2992 "a",
2993 "b",
2994 1.0,
2995 Default::default(),
2996 )
2997 }
2998
2999 #[test]
3000 fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
3001 let mut segment = GrowingSegment::new(1, "mixed");
3002 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
3003 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
3004 segment
3005 .bulk_insert(vec![
3006 UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
3007 pruning_node(2),
3008 pruning_edge(4),
3009 ])
3010 .expect("mixed bulk with gap");
3011 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3012 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3013 segment.seal().expect("seal");
3014 assert!(segment.force_delete(EntityId::new(2)));
3015 assert!(segment.force_delete(EntityId::new(4)));
3016 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3018 let mut rebuilt = GrowingSegment::new(2, "mixed");
3019 for item in segment.iter() {
3020 rebuilt.adopt_entity(item.clone(), None);
3021 }
3022 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
3023 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
3024 rebuilt.adopt_entity(pruning_node(5), None);
3025 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
3026 rebuilt.insert(pruning_edge(6)).expect("single edge");
3027 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
3028 }
3029
3030 #[test]
3031 fn graph_presence_covers_replacement_and_unrestricted_mutation() {
3032 for flat in [false, true] {
3033 for operation in 0..4 {
3034 let mut segment = GrowingSegment::new(1, "mixed");
3035 let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
3036 if flat {
3037 segment.bulk_insert(vec![vector]).expect("flat vector");
3038 } else {
3039 segment.insert(vector).expect("hashmap vector");
3040 }
3041 match operation {
3042 0 => segment.update(pruning_node(1)).expect("replace with node"),
3043 1 => segment
3044 .update_hot(pruning_node(1), &[])
3045 .expect("hot replacement"),
3046 2 => {
3047 segment.seal().expect("seal vector");
3048 segment
3049 .force_update_with_metadata(&pruning_node(1), &[], None)
3050 .expect("sealed replacement");
3051 }
3052 _ => {
3053 *segment.get_mut(EntityId::new(1)).expect("mutable vector") =
3054 pruning_node(1)
3055 }
3056 }
3057 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3058 assert!(matches!(
3059 segment.get(EntityId::new(1)).expect("node").kind,
3060 EntityKind::GraphNode(_)
3061 ));
3062 if operation == 3 {
3063 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3064 }
3065 segment
3066 .force_update_with_metadata(&pruning_edge(1), &[], None)
3067 .expect("replace with edge");
3068 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3069 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3070 }
3071 }
3072 }
3073
3074 fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
3075 UnifiedEntity::new(
3076 EntityId::new(id),
3077 EntityKind::TableRow {
3078 table: "keys".into(),
3079 row_id: id,
3080 },
3081 EntityData::Row(RowData::with_names(
3082 vec![Value::text("payload"), key],
3083 vec!["payload".into(), "primary".into()],
3084 )),
3085 )
3086 }
3087
3088 fn primary_key_probe(
3089 segment: &GrowingSegment,
3090 key: Value,
3091 exclude: Option<EntityId>,
3092 ) -> Option<EntityId> {
3093 segment.find_unique_key_conflict(
3094 &[vec!["primary".into()]],
3095 0,
3096 &[format!("{key:?}")],
3097 exclude,
3098 &|_| true,
3099 )
3100 }
3101
3102 #[test]
3103 fn unique_key_lookup_follows_segment_mutations_and_sealing() {
3104 let mut segment = GrowingSegment::new(1, "keys");
3105 segment
3106 .bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
3107 .expect("bulk row");
3108 let payload_bytes = segment.memory_bytes();
3109 assert_eq!(
3110 primary_key_probe(&segment, Value::Integer(10), None),
3111 Some(EntityId::new(1))
3112 );
3113 assert!(segment.memory_bytes() > payload_bytes);
3114 segment
3115 .insert(primary_key_row(3, Value::Integer(30)))
3116 .expect("hashmap fallback row");
3117 segment
3118 .bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
3119 .expect("bulk gap row");
3120 assert_eq!(
3121 primary_key_probe(&segment, Value::Integer(40), None),
3122 Some(EntityId::new(4))
3123 );
3124 assert_eq!(
3125 primary_key_probe(&segment, Value::Integer(30), None),
3126 Some(EntityId::new(3))
3127 );
3128 segment
3129 .update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
3130 .expect("replace flat key");
3131 assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
3132 assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
3133 assert_eq!(
3134 primary_key_probe(&segment, Value::Integer(11), None),
3135 Some(EntityId::new(1))
3136 );
3137 segment
3138 .delete_batch(&[EntityId::new(1), EntityId::new(3)])
3139 .expect("delete flat and map");
3140 assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
3141 assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
3142 segment.seal().expect("seal");
3143 assert_eq!(
3144 primary_key_probe(&segment, Value::Integer(40), None),
3145 Some(EntityId::new(4))
3146 );
3147 segment
3148 .force_update_with_metadata(
3149 &primary_key_row(4, Value::Integer(41)),
3150 &["primary".into()],
3151 None,
3152 )
3153 .expect("sealed update");
3154 assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
3155 assert_eq!(
3156 primary_key_probe(&segment, Value::Integer(41), None),
3157 Some(EntityId::new(4))
3158 );
3159 assert!(segment.force_delete(EntityId::new(4)));
3160 assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
3161 }
3162
3163 #[test]
3164 fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
3165 let mut segment = GrowingSegment::new(1, "keys");
3166 segment
3167 .bulk_insert(vec![
3168 primary_key_row(1, Value::Integer(10)),
3169 primary_key_row(2, Value::Integer(20)),
3170 ])
3171 .expect("initial rows");
3172 let definitions = vec![
3173 vec!["primary".into()],
3174 vec!["payload".into()],
3175 vec!["primary".into(), "payload".into()],
3176 ];
3177 let payload = format!("{:?}", Value::text("payload"));
3178 assert_eq!(
3179 segment.find_unique_key_conflict(
3180 &definitions,
3181 1,
3182 std::slice::from_ref(&payload),
3183 None,
3184 &|entity| entity.id != EntityId::new(1)
3185 ),
3186 Some(EntityId::new(2))
3187 );
3188 assert_eq!(segment.unique_key_lookup.read().len(), 3);
3189 segment
3190 .insert(primary_key_row(3, Value::Integer(30)))
3191 .expect("append");
3192 assert_eq!(
3193 segment.find_unique_key_conflict(
3194 &definitions,
3195 2,
3196 &[format!("{:?}", Value::Integer(30)), payload.clone()],
3197 None,
3198 &|_| true
3199 ),
3200 Some(EntityId::new(3))
3201 );
3202 let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
3203 assert_eq!(
3204 segment.find_unique_key_conflict(
3205 &definitions,
3206 0,
3207 &[format!("{:?}", Value::Integer(20))],
3208 None,
3209 &|_| true
3210 ),
3211 Some(EntityId::new(2))
3212 );
3213 assert_eq!(segment.unique_key_lookup.read().len(), 3);
3214 assert_eq!(
3215 segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
3216 bytes
3217 );
3218 let reduced = vec![vec!["payload".into()]];
3219 assert_eq!(
3220 segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
3221 Some(EntityId::new(1))
3222 );
3223 assert_eq!(segment.unique_key_lookup.read().len(), 1);
3224 assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
3225 }
3226
3227 #[test]
3228 fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
3229 let mut segment = GrowingSegment::new(1, "keys");
3230 let values = [
3231 Value::Float(f64::NAN),
3232 Value::Float(-0.0),
3233 Value::Float(0.0),
3234 Value::Integer(1),
3235 Value::UnsignedInteger(1),
3236 Value::text("1"),
3237 ];
3238 for (offset, value) in values.iter().enumerate() {
3239 let id = offset as u64 + 1;
3240 segment
3241 .insert(primary_key_row(id, value.clone()))
3242 .expect("distinct physical row");
3243 assert_eq!(
3244 primary_key_probe(&segment, value.clone(), None),
3245 Some(EntityId::new(id))
3246 );
3247 assert_eq!(
3248 primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
3249 None
3250 );
3251 }
3252 segment
3253 .insert(primary_key_row(20, Value::Integer(1)))
3254 .expect("duplicate physical key");
3255 assert_eq!(
3256 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
3257 Some(EntityId::new(20))
3258 );
3259 segment
3262 .entities
3263 .insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
3264 assert_eq!(
3265 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
3266 None
3267 );
3268 let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
3270 *entity = primary_key_row(4, Value::Integer(100));
3271 assert_eq!(
3272 primary_key_probe(&segment, Value::Integer(100), None),
3273 Some(EntityId::new(4))
3274 );
3275 }
3276
3277 #[test]
3278 fn test_growing_segment_basic() {
3279 let mut segment = GrowingSegment::new(1, "test");
3280
3281 let entity = UnifiedEntity::table_row(
3282 EntityId::new(1),
3283 "users",
3284 1,
3285 vec![Value::text("Alice".to_string())],
3286 );
3287
3288 let id = segment.insert(entity).unwrap();
3289 assert_eq!(id, EntityId::new(1));
3290 assert!(segment.contains(id));
3291
3292 let stats = segment.stats();
3293 assert_eq!(stats.entity_count, 1);
3294 assert_eq!(stats.row_count, 1);
3295 }
3296
3297 #[test]
3298 fn test_segment_metadata() {
3299 let mut segment = GrowingSegment::new(1, "test");
3300
3301 let entity = UnifiedEntity::table_row(
3302 EntityId::new(1),
3303 "users",
3304 1,
3305 vec![Value::text("Alice".to_string())],
3306 );
3307 segment.insert(entity).unwrap();
3308
3309 let mut meta = Metadata::new();
3310 meta.set("role", MetadataValue::String("admin".to_string()));
3311 meta.set("level", MetadataValue::Int(5));
3312
3313 segment.set_metadata(EntityId::new(1), meta).unwrap();
3314
3315 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
3316 assert_eq!(
3317 retrieved.get("role"),
3318 Some(&MetadataValue::String("admin".to_string()))
3319 );
3320 }
3321
3322 #[test]
3323 fn test_segment_seal() {
3324 let mut segment = GrowingSegment::new(1, "test");
3325
3326 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
3327 segment.insert(entity).unwrap();
3328
3329 assert!(segment.state().is_writable());
3331
3332 segment.seal().unwrap();
3334 assert_eq!(segment.state(), SegmentState::Sealed);
3335
3336 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
3338 assert!(segment.insert(entity2).is_err());
3339 }
3340
3341 #[test]
3342 fn test_should_seal() {
3343 let mut segment = GrowingSegment::new(1, "test");
3344
3345 let config = SegmentConfig {
3346 max_entities: 2,
3347 ..Default::default()
3348 };
3349
3350 assert!(!segment.should_seal(&config));
3351
3352 segment
3353 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
3354 .unwrap();
3355 assert!(!segment.should_seal(&config));
3356
3357 segment
3358 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
3359 .unwrap();
3360 assert!(segment.should_seal(&config));
3361 }
3362
3363 #[test]
3364 fn test_cross_references() {
3365 let mut segment = GrowingSegment::new(1, "test");
3366
3367 let mut entity1 = UnifiedEntity::table_row(
3368 EntityId::new(1),
3369 "hosts",
3370 1,
3371 vec![Value::text("192.168.1.1".to_string())],
3372 );
3373 entity1.add_cross_ref(CrossRef::new(
3374 EntityId::new(1),
3375 EntityId::new(2),
3376 "nodes",
3377 RefType::RowToNode,
3378 ));
3379 segment.insert(entity1).unwrap();
3380
3381 let refs_from = segment.get_references_from(EntityId::new(1));
3382 assert_eq!(refs_from.len(), 1);
3383 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
3384
3385 let refs_to = segment.get_references_to(EntityId::new(2));
3386 assert_eq!(refs_to.len(), 1);
3387 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
3388 }
3389
3390 #[test]
3391 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
3392 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
3393 zone.update(&Value::Email("delta@example.com".to_string()));
3394
3395 let probe = Value::Email("alpha@example.com".to_string());
3396 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
3397
3398 let in_range = Value::Email("charlie@example.com".to_string());
3399 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
3400 }
3401
3402 #[test]
3403 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
3404 let mut segment = GrowingSegment::new(1, "test");
3405
3406 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
3407 let entity = UnifiedEntity::new(
3408 EntityId::new(row_id),
3409 EntityKind::TableRow {
3410 table: "users".into(),
3411 row_id,
3412 },
3413 EntityData::Row(RowData::with_names(
3414 vec![Value::Integer(age)],
3415 vec!["age".to_string()],
3416 )),
3417 );
3418 segment.insert(entity).unwrap();
3419 }
3420
3421 segment.seal().unwrap();
3422
3423 let miss = Value::Integer(500);
3424 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
3425
3426 let hit = Value::Integer(1000);
3427 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
3428 }
3429}