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(crate) fn is_tombstoned(&self, id: EntityId) -> bool {
1118 self.deleted.contains_key(&id)
1119 }
1120
1121 pub fn tombstone_count(&self) -> usize {
1124 self.deleted.len()
1125 }
1126
1127 pub fn resident_bytes(&self) -> u64 {
1130 self.memory_bytes() + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1131 }
1132
1133 pub fn reclaimable_bytes(&self) -> u64 {
1136 self.dead_entity_bytes + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1137 }
1138
1139 pub fn live_entity_ids(&self) -> Vec<EntityId> {
1145 let mut ids = Vec::with_capacity(self.live_entity_count());
1146 self.for_each_fast(|entity| {
1147 ids.push(entity.id);
1148 true
1149 });
1150 ids
1151 }
1152
1153 pub fn insert_preserving_sequence(
1157 &mut self,
1158 entity: UnifiedEntity,
1159 ) -> Result<EntityId, SegmentError> {
1160 if !self.state.is_writable() {
1161 return Err(SegmentError::NotWritable);
1162 }
1163 if self.entities.contains_key(&entity.id) {
1164 return Err(SegmentError::AlreadyExists(entity.id));
1165 }
1166 let id = entity.id;
1167 self.adopt_entity(entity, None);
1168 self.last_write_at = current_unix_secs();
1169 Ok(id)
1170 }
1171
1172 pub fn adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
1180 let id = entity.id;
1181
1182 let next = entity.sequence_id.saturating_add(1);
1185 if self.sequence.load(Ordering::Relaxed) < next {
1186 self.sequence.store(next, Ordering::Relaxed);
1187 }
1188
1189 self.add_memory(Self::estimate_entity_size(&entity));
1190 self.index_entity(&entity);
1191 self.update_col_zones_from_entity(&entity);
1192 self.entities.insert(id, entity);
1193
1194 if let Some(metadata) = metadata {
1195 if !metadata.is_empty() {
1196 self.metadata.set_all(id, &metadata);
1197 }
1198 }
1199 }
1200
1201 pub fn evict_entity(&mut self, id: EntityId) -> bool {
1207 self.mark_checkpoint_entity_dirty(id);
1208 self.invalidate_unique_key_lookup();
1209 let Some(entity) = self.entities.remove(&id) else {
1210 return false;
1211 };
1212 self.release_memory(Self::estimate_entity_size(&entity));
1213 self.unindex_entity(&entity);
1214 for cross_ref in entity.cross_refs() {
1215 if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1216 reverse.retain(|(source, _)| *source != id);
1217 }
1218 }
1219 self.metadata.remove_all(id);
1220 true
1221 }
1222
1223 fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
1225 super::memory_size::entity_bytes(entity)
1226 }
1227
1228 fn update_column_zone(&mut self, column: &str, value: &Value) {
1236 if matches!(value, Value::Null) {
1237 return;
1238 }
1239 match self.col_zones.entry(column.to_string()) {
1240 std::collections::hash_map::Entry::Occupied(mut entry) => {
1241 let before = entry.get().payload_bytes();
1242 entry.get_mut().update(value);
1243 self.zone_payload_bytes = self
1244 .zone_payload_bytes
1245 .saturating_sub(before)
1246 .saturating_add(entry.get().payload_bytes());
1247 }
1248 std::collections::hash_map::Entry::Vacant(entry) => {
1249 let zone = ColZone::new(value.clone());
1250 self.zone_payload_bytes =
1251 self.zone_payload_bytes.saturating_add(zone.payload_bytes());
1252 entry.insert(zone);
1253 }
1254 }
1255 }
1256
1257 fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
1258 if let EntityData::Row(row) = &entity.data {
1259 if let Some(named) = &row.named {
1260 for (column, value) in named {
1261 self.update_column_zone(column, value);
1262 }
1263 } else if let Some(schema) = &row.schema {
1264 for (column, value) in schema.iter().zip(&row.columns) {
1265 self.update_column_zone(column, value);
1266 }
1267 }
1268 }
1269 }
1270
1271 fn rebuild_sealed_col_zones(&mut self) {
1272 let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
1273 let mut family_by_col: HashMap<String, reddb_types::CanonicalKeyFamily> = HashMap::new();
1274 let mut mixed_family_cols = HashSet::new();
1275 let mut unsupported_cols = HashSet::new();
1276
1277 let mut observe_row = |row: &super::entity::RowData| {
1278 for (col, value) in row.iter_fields() {
1279 if matches!(value, Value::Null) {
1280 continue;
1281 }
1282 let Some(key) = value_to_canonical_key(value) else {
1283 unsupported_cols.insert(col.to_string());
1284 continue;
1285 };
1286 match family_by_col.get(col).copied() {
1287 Some(existing) if existing != key.family() => {
1288 mixed_family_cols.insert(col.to_string());
1289 }
1290 None => {
1291 family_by_col.insert(col.to_string(), key.family());
1292 }
1293 _ => {}
1294 }
1295 values_by_col
1296 .entry(col.to_string())
1297 .or_default()
1298 .push((key, value.clone()));
1299 }
1300 };
1301
1302 if self.use_flat {
1303 for entity in &self.flat_entities {
1304 if self.deleted.contains_key(&entity.id) {
1305 continue;
1306 }
1307 if let EntityData::Row(row) = &entity.data {
1308 observe_row(row);
1309 }
1310 }
1311 } else {
1312 for entity in self.entities.values() {
1313 if self.deleted.contains_key(&entity.id) {
1314 continue;
1315 }
1316 if let EntityData::Row(row) = &entity.data {
1317 observe_row(row);
1318 }
1319 }
1320 }
1321
1322 let mut sealed_col_zones = HashMap::new();
1323 for (col, mut entries) in values_by_col {
1324 if mixed_family_cols.contains(&col)
1325 || unsupported_cols.contains(&col)
1326 || entries.is_empty()
1327 {
1328 continue;
1329 }
1330 entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
1331 entries.dedup_by(|left, right| left.0 == right.0);
1332
1333 let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
1334 if intervals.len() > 1 {
1335 sealed_col_zones.insert(col, MultiColZone { intervals });
1336 }
1337 }
1338
1339 let previous = self
1340 .sealed_col_zones
1341 .values()
1342 .flat_map(|zones| &zones.intervals)
1343 .map(ColZone::payload_bytes)
1344 .fold(0, usize::saturating_add);
1345 let current = sealed_col_zones
1346 .values()
1347 .flat_map(|zones| &zones.intervals)
1348 .map(ColZone::payload_bytes)
1349 .fold(0, usize::saturating_add);
1350 self.zone_payload_bytes = self
1351 .zone_payload_bytes
1352 .saturating_sub(previous)
1353 .saturating_add(current);
1354 self.sealed_col_zones = sealed_col_zones;
1355 }
1356
1357 pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
1361 if preds.is_empty() {
1362 return false;
1363 }
1364 for (col, pred) in preds {
1365 if let Some(zone) = self.sealed_col_zones.get(*col) {
1366 if zone.can_skip(pred) {
1367 return true;
1368 }
1369 continue;
1370 }
1371 if let Some(zone) = self.col_zones.get(*col) {
1372 if pred.can_skip(zone) {
1373 return true; }
1375 }
1376 }
1377 false
1378 }
1379
1380 fn index_entity(&mut self, entity: &UnifiedEntity) {
1382 self.mark_checkpoint_entity_dirty(entity.id);
1383 if let Some(index) = self.key_lookup.as_mut() {
1384 index.insert(entity);
1385 }
1386 self.hash_scan_ids.push(entity.id);
1387 self.observe_graph_kind(&entity.kind);
1388 if let Some(index) = self.graph_read_index.get_mut() {
1389 index.insert(entity);
1390 self.graph_read_index_bytes
1391 .store(index.memory_bytes(), Ordering::Relaxed);
1392 }
1393 let indexes = self.unique_key_lookup.get_mut();
1394 for index in indexes.iter_mut() {
1395 index.insert(entity);
1396 }
1397 self.unique_key_lookup_bytes.store(
1398 indexes
1399 .iter()
1400 .map(|index| index.memory_bytes() as u64)
1401 .sum(),
1402 Ordering::Relaxed,
1403 );
1404 let kind_key = entity.kind.storage_type().to_string();
1406 self.kind_index
1407 .entry(kind_key)
1408 .or_default()
1409 .insert(entity.id);
1410
1411 if let EntityData::Row(row) = &entity.data {
1413 if let Some(first_col) = row.columns.first() {
1414 let pk_str = format!("{:?}", first_col);
1415 self.pk_index
1416 .insert((entity.kind.collection().to_string(), pk_str), entity.id);
1417 }
1418 }
1419
1420 for cross_ref in entity.cross_refs() {
1422 self.cross_ref_forward
1423 .entry(cross_ref.source)
1424 .or_default()
1425 .push((cross_ref.target, cross_ref.ref_type));
1426
1427 self.cross_ref_reverse
1428 .entry(cross_ref.target)
1429 .or_default()
1430 .push((cross_ref.source, cross_ref.ref_type));
1431 }
1432 }
1433
1434 fn invalidate_graph_read_index(&mut self) {
1435 *self.graph_read_index.get_mut() = None;
1436 self.graph_read_index_bytes.store(0, Ordering::Relaxed);
1437 }
1438
1439 #[cfg(test)]
1442 pub fn visit_graph_index_candidates(
1443 &self,
1444 kind: GraphEntityKind,
1445 keys: &[EntityId],
1446 before_work: &mut impl FnMut() -> bool,
1447 visit: &mut impl FnMut(&UnifiedEntity) -> bool,
1448 ) -> bool {
1449 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1450 return false;
1451 };
1452 let index = cached.as_ref().expect("complete graph index");
1453 for key in keys {
1454 if !before_work() {
1455 return false;
1456 }
1457 for id in index.candidates(kind, *key) {
1458 if !before_work() {
1459 return false;
1460 }
1461 if let Some(entity) = self.get(id) {
1462 if !visit(entity) {
1463 return false;
1464 }
1465 }
1466 }
1467 }
1468 true
1469 }
1470
1471 fn graph_read_index_for_query(
1472 &self,
1473 before_work: &mut impl FnMut() -> bool,
1474 ) -> Option<parking_lot::RwLockReadGuard<'_, Option<SegmentGraphIndex>>> {
1475 let mut cached = self.graph_read_index.read();
1476 if cached.is_none() {
1477 drop(cached);
1478 let mut writer = self.graph_read_index.write();
1479 if writer.is_none() {
1480 let mut index = SegmentGraphIndex::default();
1481 if !self.for_each_fast(|entity| {
1484 if !before_work() {
1485 return false;
1486 }
1487 index.insert(entity);
1488 true
1489 }) {
1490 return None;
1491 }
1492 self.graph_read_index_bytes
1493 .store(index.memory_bytes(), Ordering::Relaxed);
1494 *writer = Some(index);
1495 }
1496 cached = parking_lot::RwLockWriteGuard::downgrade(writer);
1497 }
1498 Some(cached)
1499 }
1500
1501 pub fn graph_index_cursor(
1502 &self,
1503 kind: GraphEntityKind,
1504 keys: &[EntityId],
1505 before_work: &mut impl FnMut() -> bool,
1506 ) -> Option<GraphIndexCursor> {
1507 let cached = self.graph_read_index_for_query(before_work)?;
1508 let index = cached.as_ref().expect("complete graph index");
1509 let mut upper = EntityId::new(0);
1510 for key in keys {
1511 if !before_work() {
1512 return None;
1513 }
1514 if let Some(last) = index.candidates(kind, *key).next_back() {
1515 upper = upper.max(last);
1516 }
1517 }
1518 Some(GraphIndexCursor {
1519 key_index: 0,
1520 after: None,
1521 upper,
1522 })
1523 }
1524
1525 pub fn graph_index_batch(
1528 &self,
1529 kind: GraphEntityKind,
1530 keys: &[EntityId],
1531 cursor: &mut GraphIndexCursor,
1532 before_work: &mut impl FnMut() -> bool,
1533 visit: &mut impl FnMut(EntityId),
1534 ) -> bool {
1535 self.graph_index_batch_with_limit(kind, keys, cursor, before_work, visit, SCAN_BATCH_SIZE)
1536 }
1537
1538 pub fn graph_index_batch_with_limit(
1539 &self,
1540 kind: GraphEntityKind,
1541 keys: &[EntityId],
1542 cursor: &mut GraphIndexCursor,
1543 before_work: &mut impl FnMut() -> bool,
1544 visit: &mut impl FnMut(EntityId),
1545 limit: usize,
1546 ) -> bool {
1547 assert!(
1548 limit > 0 && limit <= SCAN_BATCH_SIZE,
1549 "bounded graph index batch"
1550 );
1551 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1552 return false;
1553 };
1554 let index = cached.as_ref().expect("complete graph index");
1555 let mut count = 0;
1556 while !cursor.finished(keys) && count < limit {
1557 let key = keys[cursor.key_index];
1558 if !before_work() {
1559 return false;
1560 }
1561 for id in index
1562 .candidates_after(kind, key, cursor.after, cursor.upper)
1563 .take(limit - count)
1564 {
1565 if !before_work() {
1566 return false;
1567 }
1568 cursor.after = Some(id);
1569 count += 1;
1570 visit(id);
1571 }
1572 if count == limit && cursor.after != Some(cursor.upper) {
1573 break;
1574 }
1575 cursor.key_index += 1;
1576 cursor.after = None;
1577 }
1578 true
1579 }
1580
1581 pub fn find_unique_key_conflict(
1586 &self,
1587 columns: &[Vec<String>],
1588 key_index: usize,
1589 signatures: &[String],
1590 exclude_id: Option<EntityId>,
1591 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1592 ) -> Option<EntityId> {
1593 assert!(
1594 key_index < columns.len(),
1595 "constraint index belongs to current schema"
1596 );
1597 let cached = self.unique_key_lookup.read();
1598 if cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1599 return self.find_unique_key_candidate(
1600 &cached[key_index],
1601 signatures,
1602 exclude_id,
1603 reserves_key,
1604 );
1605 }
1606 drop(cached);
1607 let mut cached = self.unique_key_lookup.write();
1608 if !cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1609 let mut indexes: Vec<_> = columns
1610 .iter()
1611 .map(|columns| SegmentUniqueKey::new(columns))
1612 .collect();
1613 self.for_each_fast(|entity| {
1614 for index in &mut indexes {
1615 index.insert(entity);
1616 }
1617 true
1618 });
1619 self.unique_key_lookup_bytes.store(
1620 indexes
1621 .iter()
1622 .map(|index| index.memory_bytes() as u64)
1623 .sum(),
1624 Ordering::Relaxed,
1625 );
1626 *cached = indexes;
1627 }
1628 self.find_unique_key_candidate(&cached[key_index], signatures, exclude_id, reserves_key)
1629 }
1630
1631 fn find_unique_key_candidate(
1632 &self,
1633 index: &SegmentUniqueKey,
1634 signatures: &[String],
1635 exclude_id: Option<EntityId>,
1636 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1637 ) -> Option<EntityId> {
1638 index.candidates(signatures).iter().copied().find(|id| {
1639 if Some(*id) == exclude_id {
1640 return false;
1641 }
1642 let Some(
1643 entity @ UnifiedEntity {
1644 data: EntityData::Row(row),
1645 ..
1646 },
1647 ) = self.get(*id)
1648 else {
1649 return false;
1650 };
1651 index
1652 .columns
1653 .iter()
1654 .zip(signatures)
1655 .all(|(column, signature)| {
1656 row.get_field(column)
1657 .is_some_and(|value| format!("{value:?}") == *signature)
1658 })
1659 && reserves_key(entity)
1660 })
1661 }
1662
1663 fn invalidate_unique_key_lookup(&mut self) {
1664 *self.unique_key_lookup.get_mut() = Vec::new();
1665 self.unique_key_lookup_bytes.store(0, Ordering::Relaxed);
1666 }
1667
1668 pub fn find_key_entity(
1671 &self,
1672 key: &str,
1673 newest: bool,
1674 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1675 ) -> Option<&UnifiedEntity> {
1676 let index = self
1679 .key_lookup
1680 .as_ref()
1681 .expect("segment owns its key index");
1682 self.find_key_candidate(index, key, newest, visible)
1683 }
1684
1685 fn find_key_candidate(
1686 &self,
1687 index: &SegmentKeyLookup,
1688 key: &str,
1689 newest: bool,
1690 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1691 ) -> Option<&UnifiedEntity> {
1692 self.select_key_entity(
1693 index
1694 .candidates(key)
1695 .iter()
1696 .chain(index.mutated_candidates())
1697 .filter_map(|id| self.get(*id)),
1698 key,
1699 newest,
1700 visible,
1701 )
1702 }
1703
1704 fn select_key_entity<'a>(
1705 &self,
1706 entities: impl Iterator<Item = &'a UnifiedEntity>,
1707 key: &str,
1708 newest: bool,
1709 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1710 ) -> Option<&'a UnifiedEntity> {
1711 let mut best: Option<&UnifiedEntity> = None;
1712 for entity in entities {
1713 if SegmentKeyLookup::key(entity) != Some(key) || !visible(entity) {
1715 continue;
1716 }
1717 if !newest {
1718 return Some(entity);
1719 }
1720 if best.is_none_or(|current| entity.xmin >= current.xmin) {
1721 best = Some(entity);
1722 }
1723 }
1724 best
1725 }
1726
1727 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1729 if let Some(id) = entity_id_from_probe_key(key) {
1730 return self.has_live_entity(id);
1731 }
1732
1733 let Ok(pk_str) = std::str::from_utf8(key) else {
1734 return false;
1735 };
1736 let pk_key = (self.collection.clone(), pk_str.to_string());
1737 self.pk_index
1738 .get(&pk_key)
1739 .is_some_and(|id| self.has_live_entity(*id))
1740 }
1741
1742 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1744 if let Some(index) = self.key_lookup.as_mut() {
1745 index.remove(entity);
1746 }
1747 if let Some(index) = self.graph_read_index.get_mut() {
1748 index.remove(entity);
1749 self.graph_read_index_bytes
1750 .store(index.memory_bytes(), Ordering::Relaxed);
1751 }
1752 let kind_key = entity.kind.storage_type().to_string();
1754 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1755 set.remove(&entity.id);
1756 }
1757
1758 if let EntityData::Row(row) = &entity.data {
1760 if let Some(first_col) = row.columns.first() {
1761 let pk_str = format!("{:?}", first_col);
1762 self.pk_index
1763 .remove(&(entity.kind.collection().to_string(), pk_str));
1764 }
1765 }
1766
1767 self.cross_ref_forward.remove(&entity.id);
1769 }
1771
1772 fn reindex_for_update(
1782 &mut self,
1783 old: &UpdateIndexSnapshot,
1784 new: &UnifiedEntity,
1785 modified_columns: Option<&[String]>,
1786 ) {
1787 let pk_changed = match &new.data {
1792 EntityData::Row(new_row) => {
1793 if let Some(cols) = modified_columns {
1794 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1797 new_row
1798 .schema
1799 .as_deref()
1800 .and_then(|schema| schema.first().map(|name| name.as_str()))
1801 });
1802 match pk_col_name {
1803 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1804 None => old.pk_value.as_ref() != new_row.columns.first(),
1806 }
1807 } else {
1808 old.pk_value.as_ref() != new_row.columns.first()
1809 }
1810 }
1811 _ => false,
1813 };
1814
1815 if pk_changed {
1816 if let Some((collection, pk_str)) = &old.pk_index_key {
1818 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1819 }
1820 if let EntityData::Row(row) = &new.data {
1822 if let Some(first_col) = row.columns.first() {
1823 let pk_str = format!("{:?}", first_col);
1824 self.pk_index
1825 .insert((new.kind.collection().to_string(), pk_str), new.id);
1826 }
1827 }
1828 }
1829
1830 let new_refs = new.cross_refs();
1832 if old.cross_refs.as_slice() != new_refs {
1833 self.cross_ref_forward.remove(&new.id);
1835 for cross_ref in &old.cross_refs {
1837 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1838 rev.retain(|(src, _)| *src != new.id);
1839 }
1840 }
1841 for cross_ref in new_refs {
1843 self.cross_ref_forward
1844 .entry(cross_ref.source)
1845 .or_default()
1846 .push((cross_ref.target, cross_ref.ref_type));
1847 self.cross_ref_reverse
1848 .entry(cross_ref.target)
1849 .or_default()
1850 .push((cross_ref.source, cross_ref.ref_type));
1851 }
1852 }
1853 }
1854
1855 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1857 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1858 }
1859
1860 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1862 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1863 }
1864
1865 pub fn age_secs(&self) -> u64 {
1867 let now = current_unix_secs();
1868 now.saturating_sub(self.created_at)
1869 }
1870
1871 pub fn idle_secs(&self) -> u64 {
1873 let now = current_unix_secs();
1874 now.saturating_sub(self.last_write_at)
1875 }
1876
1877 pub fn bulk_insert(
1886 &mut self,
1887 entities: Vec<UnifiedEntity>,
1888 ) -> Result<Vec<EntityId>, SegmentError> {
1889 if !self.state.is_writable() {
1890 return Err(SegmentError::NotWritable);
1891 }
1892
1893 for entity in &entities {
1894 self.mark_checkpoint_entity_dirty(entity.id);
1895 }
1896
1897 if let Some(index) = self.key_lookup.as_mut() {
1898 for entity in &entities {
1899 index.insert(entity);
1900 }
1901 }
1902
1903 let indexes = self.unique_key_lookup.get_mut();
1904 for index in indexes.iter_mut() {
1905 for entity in &entities {
1906 index.insert(entity);
1907 }
1908 }
1909 self.unique_key_lookup_bytes.store(
1910 indexes
1911 .iter()
1912 .map(|index| index.memory_bytes() as u64)
1913 .sum(),
1914 Ordering::Relaxed,
1915 );
1916
1917 let n = entities.len();
1918
1919 if entities.is_empty() {
1920 return Ok(Vec::new());
1921 }
1922
1923 for run in
1926 entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
1927 {
1928 let kind = run[0].kind.storage_type();
1929 self.observe_graph_kind(&run[0].kind);
1930 if matches!(
1931 run[0].kind,
1932 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
1933 ) {
1934 if let Some(index) = self.graph_read_index.get_mut() {
1935 for entity in run {
1936 index.insert(entity);
1937 }
1938 self.graph_read_index_bytes
1939 .store(index.memory_bytes(), Ordering::Relaxed);
1940 }
1941 }
1942 let kind_set = match self.kind_index.get_mut(kind) {
1943 Some(ids) => ids,
1944 None => self.kind_index.entry(kind.to_string()).or_default(),
1945 };
1946 kind_set.reserve(run.len());
1947 kind_set.extend(run.iter().map(|entity| entity.id));
1948 }
1949
1950 let now = current_unix_secs();
1951
1952 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
1953
1954 let mut ids = Vec::with_capacity(n);
1955
1956 if self.flat_entities.is_empty() && self.entities.is_empty() {
1958 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
1960 self.use_flat = true;
1961 }
1962
1963 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
1973 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
1974 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
1978
1979 let mut batch_bytes: usize = 0;
1984
1985 if self.use_flat {
1986 self.flat_entities.reserve(n);
1987 for (i, mut entity) in entities.into_iter().enumerate() {
1988 entity.sequence_id = base_seq + i as u64;
1989 let id = entity.id;
1990 ids.push(id);
1991 batch_bytes += Self::estimate_entity_size(&entity);
1992 if let EntityData::Row(row) = &entity.data {
1993 if row.schema.is_some() && !row.columns.is_empty() {
1994 if columnar_zone_updates.is_empty() {
1995 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
1996 columnar_schema = row.schema.clone();
1997 }
1998 for (ci, val) in row.columns.iter().enumerate() {
1999 if !matches!(val, Value::Null) {
2000 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
2001 bucket.push(val.clone());
2002 }
2003 }
2004 }
2005 } else {
2006 for (col, val) in row.iter_fields() {
2007 if !matches!(val, Value::Null) {
2008 named_zone_updates.push((col.to_string(), val.clone()));
2009 }
2010 }
2011 }
2012 }
2013 let expected = self.base_entity_id + self.flat_entities.len() as u64;
2029 if id.raw() == expected {
2030 self.flat_entities.push(entity);
2031 } else {
2032 self.hash_scan_ids.push(id);
2033 self.entities.insert(id, entity);
2034 }
2035 }
2036 } else {
2037 self.entities.reserve(n);
2039 self.hash_scan_ids.reserve(n);
2040 let mut pairs = Vec::with_capacity(n);
2041 for (i, mut entity) in entities.into_iter().enumerate() {
2042 entity.sequence_id = base_seq + i as u64;
2043 let id = entity.id;
2044 ids.push(id);
2045 batch_bytes += Self::estimate_entity_size(&entity);
2046 if let EntityData::Row(row) = &entity.data {
2047 for (col, val) in row.iter_fields() {
2048 if !matches!(val, Value::Null) {
2049 named_zone_updates.push((col.to_string(), val.clone()));
2050 }
2051 }
2052 }
2053 self.hash_scan_ids.push(id);
2054 pairs.push((id, entity));
2055 }
2056 self.entities.extend(pairs);
2057 }
2058
2059 if !columnar_zone_updates.is_empty() {
2063 let schema = columnar_schema.as_ref();
2064 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
2065 if values.is_empty() {
2066 continue;
2067 }
2068 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
2069 continue;
2070 };
2071 let mut iter = values.into_iter();
2077 if let Some(first) = iter.next() {
2078 let previous = self
2079 .col_zones
2080 .get(col_name)
2081 .map_or(0, ColZone::payload_bytes);
2082 let zone = self
2083 .col_zones
2084 .entry(col_name.clone())
2085 .and_modify(|z| z.update(&first))
2086 .or_insert_with(|| ColZone::new(first));
2087 for v in iter {
2088 zone.update(&v);
2089 }
2090 self.zone_payload_bytes = self
2091 .zone_payload_bytes
2092 .saturating_sub(previous)
2093 .saturating_add(zone.payload_bytes());
2094 }
2095 }
2096 }
2097 for (col, val) in named_zone_updates {
2098 self.update_column_zone(&col, &val);
2099 }
2100
2101 self.add_memory(batch_bytes);
2102 self.last_write_at = now;
2103
2104 if self.use_flat {
2106 self.published_flat_len
2107 .store(self.flat_entities.len(), Ordering::Release);
2108 }
2109
2110 Ok(ids)
2111 }
2112
2113 pub fn force_delete(&mut self, id: EntityId) -> bool {
2116 if self.use_flat {
2117 let raw = id.raw();
2118 if raw >= self.base_entity_id {
2119 let idx = (raw - self.base_entity_id) as usize;
2120 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2121 self.tombstone_flat_entity(idx, id);
2122 return true;
2123 }
2124 }
2125 return self.remove_hashmap_entity(id);
2126 }
2127
2128 self.remove_hashmap_entity(id)
2129 }
2130
2131 pub fn force_update_with_metadata(
2134 &mut self,
2135 entity: &UnifiedEntity,
2136 modified_columns: &[String],
2137 metadata: Option<&Metadata>,
2138 ) -> Result<(), SegmentError> {
2139 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
2140 }
2141}
2142
2143impl UnifiedSegment for GrowingSegment {
2144 fn id(&self) -> SegmentId {
2145 self.id
2146 }
2147
2148 fn state(&self) -> SegmentState {
2149 self.state
2150 }
2151
2152 fn collection(&self) -> &str {
2153 &self.collection
2154 }
2155
2156 fn stats(&self) -> SegmentStats {
2157 let mut stats = SegmentStats {
2158 entity_count: self.live_entity_count(),
2159 deleted_count: self.deleted.len(),
2160 memory_bytes: self.resident_bytes() as usize,
2161 ..Default::default()
2162 };
2163
2164 for entity in self.entities.values() {
2165 match &entity.kind {
2166 EntityKind::TableRow { .. } => stats.row_count += 1,
2167 EntityKind::GraphNode(_) => stats.node_count += 1,
2168 EntityKind::GraphEdge(_) => stats.edge_count += 1,
2169 EntityKind::Vector { .. } => stats.vector_count += 1,
2170 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
2171 EntityKind::QueueMessage { .. } => stats.row_count += 1,
2172 }
2173 stats.cross_ref_count += entity.cross_refs().len();
2174 }
2175
2176 stats
2177 }
2178
2179 fn entity_count(&self) -> usize {
2180 self.live_entity_count()
2181 }
2182
2183 fn contains(&self, id: EntityId) -> bool {
2184 self.has_live_entity(id)
2185 }
2186
2187 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2188 if self.deleted.contains_key(&id) {
2189 return None;
2190 }
2191 if self.use_flat {
2192 let raw = id.raw();
2193 if raw >= self.base_entity_id {
2194 let idx = (raw - self.base_entity_id) as usize;
2195 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
2196 return Some(entity);
2197 }
2198 }
2199 self.entities.get(&id)
2205 } else {
2206 self.entities.get(&id)
2207 }
2208 }
2209
2210 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
2211 self.mark_checkpoint_entity_dirty(id);
2212 self.invalidate_unique_key_lookup();
2213 self.invalidate_graph_read_index();
2214 if self.deleted.contains_key(&id) || !self.state.is_writable() {
2215 return None;
2216 }
2217 if let Some(mut index) = self.key_lookup.take() {
2218 if let Some(entity) = self.get(id) {
2219 index.mark_mutated(entity);
2220 }
2221 self.key_lookup = Some(index);
2222 }
2223 self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
2226 if self.use_flat {
2227 let raw = id.raw();
2228 if raw >= self.base_entity_id {
2229 let idx = (raw - self.base_entity_id) as usize;
2230 if self
2231 .flat_entities
2232 .get(idx)
2233 .map(|e| e.id == id)
2234 .unwrap_or(false)
2235 {
2236 return self.flat_entities.get_mut(idx);
2237 }
2238 }
2239 self.entities.get_mut(&id)
2242 } else {
2243 self.entities.get_mut(&id)
2244 }
2245 }
2246
2247 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2248 if !self.state.is_writable() {
2249 return Err(SegmentError::NotWritable);
2250 }
2251
2252 if self.entities.contains_key(&entity.id) {
2253 return Err(SegmentError::AlreadyExists(entity.id));
2254 }
2255
2256 entity.sequence_id = self.next_sequence();
2258
2259 let size = Self::estimate_entity_size(&entity);
2261 self.add_memory(size);
2262
2263 self.index_entity(&entity);
2265
2266 self.update_col_zones_from_entity(&entity);
2268
2269 let id = entity.id;
2271 self.entities.insert(id, entity);
2272
2273 self.last_write_at = current_unix_secs();
2275
2276 Ok(id)
2277 }
2278
2279 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2280 if !self.state.is_writable() {
2281 return Err(SegmentError::NotWritable);
2282 }
2283
2284 self.apply_update_with_metadata(&entity, None)?;
2285 self.last_write_at = current_unix_secs();
2286
2287 Ok(())
2288 }
2289
2290 fn update_hot(
2291 &mut self,
2292 entity: UnifiedEntity,
2293 modified_columns: &[String],
2294 ) -> Result<(), SegmentError> {
2295 if !self.state.is_writable() {
2296 return Err(SegmentError::NotWritable);
2297 }
2298
2299 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
2300 self.last_write_at = current_unix_secs();
2301 Ok(())
2302 }
2303
2304 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2305 if !self.state.is_writable() {
2306 return Err(SegmentError::NotWritable);
2307 }
2308
2309 if self.use_flat {
2311 let raw = id.raw();
2312 if raw >= self.base_entity_id {
2313 let idx = (raw - self.base_entity_id) as usize;
2314 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2315 self.tombstone_flat_entity(idx, id);
2316 return Ok(true);
2317 }
2318 }
2319 return Ok(self.remove_hashmap_entity(id));
2320 }
2321
2322 Ok(self.remove_hashmap_entity(id))
2323 }
2324
2325 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2326 if !self.has_live_entity(id) {
2327 return None;
2328 }
2329 Some(self.metadata.get_all(id))
2330 }
2331
2332 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2333 self.mark_checkpoint_entity_dirty(id);
2334 if !self.state.is_writable() {
2335 return Err(SegmentError::NotWritable);
2336 }
2337
2338 if !self.has_live_entity(id) {
2339 return Err(SegmentError::NotFound(id));
2340 }
2341
2342 self.metadata.set_all(id, &metadata);
2343 Ok(())
2344 }
2345
2346 fn seal(&mut self) -> Result<(), SegmentError> {
2347 if self.state != SegmentState::Growing {
2348 return Err(SegmentError::InvalidState(self.state));
2349 }
2350
2351 self.state = SegmentState::Sealing;
2352
2353 self.rebuild_sealed_col_zones();
2358
2359 self.state = SegmentState::Sealed;
2360 Ok(())
2361 }
2362
2363 fn should_seal(&self, config: &SegmentConfig) -> bool {
2364 if self.entities.len() >= config.max_entities {
2366 return true;
2367 }
2368
2369 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
2371 return true;
2372 }
2373
2374 if self.age_secs() >= config.max_age_secs {
2376 return true;
2377 }
2378
2379 false
2380 }
2381
2382 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2383 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2384 Box::new(self.flat_entities.iter().chain(self.entities.values()))
2388 } else {
2389 Box::new(self.entities.values())
2390 };
2391 if self.deleted.is_empty() {
2392 base
2393 } else {
2394 Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
2395 }
2396 }
2397
2398 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2399 let ids = self.kind_index.get(kind_filter);
2402 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2405 Box::new(self.flat_entities.iter())
2406 } else {
2407 Box::new(std::iter::empty())
2408 };
2409 Box::new(flat.chain(self.entities.values()).filter(move |e| {
2410 if self.deleted.contains_key(&e.id) {
2411 return false;
2412 }
2413 if let Some(ids) = ids {
2414 ids.contains(&e.id)
2415 } else {
2416 false
2417 }
2418 }))
2419 }
2420
2421 fn filter_metadata(
2422 &self,
2423 filters: &[(String, super::metadata::MetadataFilter)],
2424 ) -> Vec<EntityId> {
2425 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
2428 Box::new(self.flat_entities.iter().map(|e| e.id))
2429 } else {
2430 Box::new(std::iter::empty())
2431 };
2432 flat_ids
2433 .chain(self.entities.keys().copied())
2434 .filter(|id| {
2435 if self.deleted.contains_key(id) {
2436 return false;
2437 }
2438 let metadata = self.metadata.get_all(*id);
2439 metadata.matches_all(filters)
2440 })
2441 .collect()
2442 }
2443}
2444
2445fn build_minmax_multi_intervals(
2446 entries: &[(CanonicalKey, Value)],
2447 max_intervals: usize,
2448) -> Vec<ColZone> {
2449 if entries.is_empty() {
2450 return Vec::new();
2451 }
2452 if entries.len() == 1 || max_intervals <= 1 {
2453 return vec![ColZone::with_bounds(
2454 entries[0].1.clone(),
2455 entries[entries.len() - 1].1.clone(),
2456 )];
2457 }
2458
2459 let mut split_points = if entries.len() <= max_intervals {
2460 (1..entries.len()).collect::<Vec<_>>()
2461 } else {
2462 let target_splits = max_intervals - 1;
2463 let mut selected = select_gap_split_points(entries, target_splits);
2464 if selected.len() < target_splits {
2465 for bucket in 1..max_intervals {
2466 let idx = bucket * entries.len() / max_intervals;
2467 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
2468 continue;
2469 }
2470 selected.push(idx);
2471 if selected.len() >= target_splits {
2472 break;
2473 }
2474 }
2475 }
2476 selected.sort_unstable();
2477 selected.dedup();
2478 selected
2479 };
2480
2481 split_points.push(entries.len());
2482
2483 let mut out = Vec::with_capacity(split_points.len());
2484 let mut start = 0usize;
2485 for end in split_points {
2486 if end <= start {
2487 continue;
2488 }
2489 out.push(ColZone::with_bounds(
2490 entries[start].1.clone(),
2491 entries[end - 1].1.clone(),
2492 ));
2493 start = end;
2494 }
2495
2496 if out.is_empty() {
2497 out.push(ColZone::with_bounds(
2498 entries[0].1.clone(),
2499 entries[entries.len() - 1].1.clone(),
2500 ));
2501 }
2502
2503 out
2504}
2505
2506fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
2507 let mut gaps = Vec::new();
2508 for idx in 1..entries.len() {
2509 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
2510 if score > 0.0 {
2511 gaps.push((score, idx));
2512 }
2513 }
2514 }
2515 gaps.sort_by(|left, right| {
2516 right
2517 .0
2518 .partial_cmp(&left.0)
2519 .unwrap_or(std::cmp::Ordering::Equal)
2520 .then_with(|| left.1.cmp(&right.1))
2521 });
2522 gaps.into_iter()
2523 .take(max_splits)
2524 .map(|(_, idx)| idx)
2525 .collect()
2526}
2527
2528fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
2529 if left.family() != right.family() {
2530 return None;
2531 }
2532 match (left, right) {
2533 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
2534 Some(r.saturating_sub(*l) as f64)
2535 }
2536 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
2537 Some(r.saturating_sub(*l) as f64)
2538 }
2539 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
2540 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
2541 }
2542 _ => None,
2543 }
2544}
2545
2546impl UnifiedSegment for std::sync::Arc<parking_lot::RwLock<GrowingSegment>> {
2549 fn id(&self) -> SegmentId {
2550 self.read().id()
2551 }
2552
2553 fn state(&self) -> SegmentState {
2554 self.read().state()
2555 }
2556
2557 fn collection(&self) -> &str {
2558 "unknown"
2560 }
2561
2562 fn stats(&self) -> SegmentStats {
2563 self.read().stats()
2564 }
2565
2566 fn entity_count(&self) -> usize {
2567 self.read().entity_count()
2568 }
2569
2570 fn contains(&self, id: EntityId) -> bool {
2571 self.read().contains(id)
2572 }
2573
2574 fn get(&self, _id: EntityId) -> Option<&UnifiedEntity> {
2575 None
2578 }
2579
2580 fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
2581 None
2582 }
2583
2584 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2585 self.write().insert(entity)
2586 }
2587
2588 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2589 self.write().update(entity)
2590 }
2591
2592 fn update_hot(
2593 &mut self,
2594 entity: UnifiedEntity,
2595 modified_columns: &[String],
2596 ) -> Result<(), SegmentError> {
2597 self.write().update_hot(entity, modified_columns)
2598 }
2599
2600 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2601 self.write().delete(id)
2602 }
2603
2604 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2605 self.read().get_metadata(id)
2606 }
2607
2608 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2609 self.write().set_metadata(id, metadata)
2610 }
2611
2612 fn seal(&mut self) -> Result<(), SegmentError> {
2613 self.write().seal()
2614 }
2615
2616 fn should_seal(&self, config: &SegmentConfig) -> bool {
2617 self.read().should_seal(config)
2618 }
2619
2620 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2621 Box::new(std::iter::empty())
2623 }
2624
2625 fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2626 Box::new(std::iter::empty())
2627 }
2628
2629 fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2630 self.read().filter_metadata(filters)
2631 }
2632}
2633
2634#[cfg(test)]
2635mod tests {
2636 use super::*;
2637 use crate::storage::entity::RowData;
2638 use crate::storage::metadata::MetadataValue;
2639 use reddb_types::Value;
2640
2641 fn keyed_row(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
2642 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
2643 entity.data = EntityData::Row(RowData::with_names(
2644 vec![Value::text(key), Value::Integer(id as i64)],
2645 vec!["key".into(), "value".into()],
2646 ));
2647 entity.set_xmin(xmin);
2648 entity.set_xmax(xmax);
2649 entity
2650 }
2651
2652 fn keyed_probe(segment: &GrowingSegment, key: &str) -> Option<EntityId> {
2653 segment
2654 .find_key_entity(key, true, &mut |_| true)
2655 .map(|entity| entity.id)
2656 }
2657
2658 #[test]
2659 fn bulk_checkpoint_journal_preserves_prior_changes_and_ignores_rejected_batches() {
2660 let mut segment = GrowingSegment::new(1, "keys");
2661 segment.insert(keyed_row(1, "one", 1, 0)).expect("row");
2662 segment
2663 .bulk_insert(vec![keyed_row(2, "two", 1, 0), keyed_row(3, "three", 1, 0)])
2664 .expect("batch");
2665 assert_eq!(
2666 segment.checkpoint_mutations_since(0),
2667 Some(vec![EntityId::new(1), EntityId::new(2), EntityId::new(3)])
2668 );
2669 let revision = segment.checkpoint_revision();
2670 segment.bulk_insert(Vec::new()).expect("empty batch");
2671 assert_eq!(segment.checkpoint_revision(), revision);
2672 segment.seal().expect("seal segment");
2673 assert!(matches!(
2674 segment.bulk_insert(vec![keyed_row(4, "four", 1, 0)]),
2675 Err(SegmentError::NotWritable)
2676 ));
2677 assert_eq!(segment.checkpoint_revision(), revision);
2678 assert_eq!(segment.checkpoint_mutations_since(revision), Some(vec![]));
2679 }
2680
2681 #[test]
2682 fn checkpoint_journal_retains_bounded_row_versions_through_repeated_updates_and_gc() {
2683 const ROWS: u64 = 1700;
2684 let mut segment = GrowingSegment::new(1, "keys");
2685 for id in 1..=ROWS {
2686 segment
2687 .insert(keyed_row(id, "export", 1, 0))
2688 .expect("seed export");
2689 }
2690 let baseline = segment.checkpoint_revision();
2691 for id in 1..=ROWS {
2695 segment
2696 .update(keyed_row(id, "export", 1, 2))
2697 .expect("supersede");
2698 let delivered = ROWS + id;
2699 segment
2700 .insert(keyed_row(delivered, "export", 2, 0))
2701 .expect("delivered");
2702 segment
2703 .update(keyed_row(delivered, "export", 2, 3))
2704 .expect("acknowledge");
2705 segment
2706 .insert(keyed_row(2 * ROWS + id, "export", 3, 0))
2707 .expect("final version");
2708 assert_eq!(
2709 segment
2710 .delete_batch(&[EntityId::new(id), EntityId::new(delivered)])
2711 .expect("GC")
2712 .len(),
2713 2
2714 );
2715 }
2716 let changed = segment
2717 .checkpoint_mutations_since(baseline)
2718 .expect("all changed IDs still fit the bounded checkpoint journal");
2719 let actual: std::collections::BTreeSet<_> = changed.into_iter().collect();
2720 let expected: std::collections::BTreeSet<_> = (1..=3 * ROWS).map(EntityId::new).collect();
2721 assert_eq!(actual, expected, "deleted versions must remain tracked too");
2722 }
2723
2724 #[test]
2725 fn checkpoint_journal_preserves_cold_changes_during_hot_row_churn() {
2726 let mut segment = GrowingSegment::new(1, "keys");
2727 segment
2728 .insert(keyed_row(1, "cold", 1, 0))
2729 .expect("cold row");
2730 segment.insert(keyed_row(2, "hot", 1, 0)).expect("hot row");
2731 for xid in 2..=(5 * CHECKPOINT_MUTATIONS_MAX) as u64 {
2732 segment
2733 .update(keyed_row(2, "hot", xid, 0))
2734 .expect("hot edit");
2735 }
2736 let changed: std::collections::BTreeSet<_> = segment
2737 .checkpoint_mutations_since(0)
2738 .expect("two physical IDs never exhaust the distinct-ID bound")
2739 .into_iter()
2740 .collect();
2741 assert_eq!(
2742 changed,
2743 [EntityId::new(1), EntityId::new(2)].into_iter().collect()
2744 );
2745 assert!(
2746 segment.checkpoint_tracking_bytes() <= 1024 * 1024,
2747 "retained proof memory must stay bounded under unlimited repeated edits"
2748 );
2749 }
2750
2751 #[test]
2752 fn checkpoint_journal_complete_suffix_matches_uncompressed_mutation_history() {
2753 let mut segment = GrowingSegment::new(1, "keys");
2754 let mut history = Vec::new();
2755 let mut seed = 0x7b3d_5961_a902_4c85_u64;
2756 for n in 0..3 * CHECKPOINT_MUTATIONS_MAX {
2757 seed ^= seed << 13;
2758 seed ^= seed >> 7;
2759 seed ^= seed << 17;
2760 let raw = if n < CHECKPOINT_MUTATIONS_MAX + 17 {
2761 n as u64 + 1
2762 } else {
2763 seed % (CHECKPOINT_MUTATIONS_MAX as u64 + 17) + 1
2764 };
2765 let id = EntityId::new(raw);
2766 segment.mark_checkpoint_entity_dirty(id);
2767 history.push(id);
2768 }
2769 let now = segment.checkpoint_revision();
2770 assert!(
2771 segment.checkpoint_mutations_since(0).is_none(),
2772 "an evicted ID must still invalidate an older baseline"
2773 );
2774 let mut complete = 0;
2775 for baseline in [0, 1, now / 3, now / 2, now - 2048, now - 128, now] {
2776 if let Some(ids) = segment.checkpoint_mutations_since(baseline) {
2777 let actual: std::collections::BTreeSet<_> = ids.into_iter().collect();
2778 let expected: std::collections::BTreeSet<_> =
2779 history[baseline as usize..].iter().copied().collect();
2780 assert_eq!(actual, expected, "baseline {baseline}");
2781 complete += 1;
2782 }
2783 }
2784 assert!(
2785 complete >= 3,
2786 "recent baselines must retain a complete proof"
2787 );
2788 assert!(segment.checkpoint_mutations_since(now + 1).is_none());
2789 assert!(segment.checkpoint_tracking_bytes() <= 1024 * 1024);
2790 }
2791
2792 #[test]
2793 fn bulk_checkpoint_journal_exhaustion_keeps_a_bounded_complete_suffix() {
2794 let mut segment = GrowingSegment::new(1, "keys");
2795 segment
2796 .bulk_insert(
2797 (1..=CHECKPOINT_MUTATIONS_MAX as u64 + 1)
2798 .map(|id| keyed_row(id, "key", 1, 0))
2799 .collect(),
2800 )
2801 .expect("large batch");
2802 assert!(segment.checkpoint_mutations_since(0).is_none());
2803 assert_eq!(
2804 segment
2805 .checkpoint_mutations_since(1)
2806 .expect("complete suffix")
2807 .len(),
2808 CHECKPOINT_MUTATIONS_MAX
2809 );
2810 let revision = segment.checkpoint_revision();
2811 let id = EntityId::new(CHECKPOINT_MUTATIONS_MAX as u64 + 2);
2812 segment
2813 .bulk_insert(vec![keyed_row(id.raw(), "next", 1, 0)])
2814 .expect("later batch");
2815 assert_eq!(segment.checkpoint_mutations_since(revision), Some(vec![id]));
2816 assert_eq!(segment.checkpoint_mutations.len(), CHECKPOINT_MUTATIONS_MAX);
2817 }
2818
2819 #[test]
2820 fn key_lookup_visits_only_requested_key_versions() {
2821 let mut segment = GrowingSegment::new(1, "keys");
2822 segment
2823 .bulk_insert(
2824 (1..=4096)
2825 .map(|id| keyed_row(id, &format!("other-{id}"), 1, 0))
2826 .collect(),
2827 )
2828 .expect("unrelated rows");
2829 segment
2830 .insert(keyed_row(5000, "wanted", 2, 8))
2831 .expect("old version");
2832 segment
2833 .insert(keyed_row(5001, "wanted", 8, 0))
2834 .expect("new version");
2835 let mut visits = 0;
2836 let result = segment.find_key_entity("wanted", true, &mut |entity| {
2837 visits += 1;
2838 entity.xmin < 8
2839 });
2840 assert_eq!(result.map(|entity| entity.id), Some(EntityId::new(5000)));
2841 assert_eq!(
2842 visits, 2,
2843 "unrelated rows must never reach visibility evaluation"
2844 );
2845 assert_eq!(keyed_probe(&segment, "wanted"), Some(EntityId::new(5001)));
2846 let mut misses = 0;
2847 assert!(segment
2848 .find_key_entity("absent", true, &mut |_| {
2849 misses += 1;
2850 true
2851 })
2852 .is_none());
2853 assert_eq!(misses, 0);
2854
2855 let mut collision = SegmentKeyLookup::default();
2857 collision.insert(&keyed_row(1, "wanted", 1, 0));
2858 assert!(segment
2859 .find_key_candidate(&collision, "wanted", true, &mut |_| true)
2860 .is_none());
2861 }
2862
2863 #[test]
2864 fn raw_mutation_does_not_disable_unrelated_key_candidates() {
2865 let mut segment = GrowingSegment::new(1, "keys");
2866 segment
2867 .bulk_insert(
2868 (1..=4096)
2869 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2870 .collect(),
2871 )
2872 .expect("rows");
2873 for _ in 0..3 {
2874 *segment.get_mut(EntityId::new(1)).expect("mutable") = keyed_row(1, "renamed", 2, 0);
2875 }
2876 let index = segment.key_lookup.as_ref().expect("index retained");
2877 assert_eq!(index.candidates("key-4096"), &[EntityId::new(4096)]);
2878 assert_eq!(index.mutated_candidates().count(), 1);
2879 assert!(index.candidates("key-1").is_empty());
2880 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(1)));
2881 assert_eq!(keyed_probe(&segment, "key-4096"), Some(EntityId::new(4096)));
2882 segment.delete(EntityId::new(1)).expect("delete raw row");
2883 assert_eq!(
2884 segment
2885 .key_lookup
2886 .as_ref()
2887 .unwrap()
2888 .mutated_candidates()
2889 .count(),
2890 0
2891 );
2892 assert_eq!(keyed_probe(&segment, "renamed"), None);
2893 }
2894
2895 #[test]
2896 fn key_lookup_tracks_mutations_sealing_and_adoption() {
2897 let mut segment = GrowingSegment::new(1, "keys");
2898 segment
2899 .bulk_insert(vec![keyed_row(1, "flat", 1, 0)])
2900 .expect("flat row");
2901 segment
2902 .insert(keyed_row(3, "hash", 1, 0))
2903 .expect("hash fallback");
2904 segment
2905 .bulk_insert(vec![keyed_row(5, "gap", 1, 0)])
2906 .expect("bulk gap");
2907 assert!(segment.key_lookup.as_ref().expect("index").memory_bytes() > 0);
2908 assert_eq!(keyed_probe(&segment, "gap"), Some(EntityId::new(5)));
2909 segment
2910 .update(keyed_row(3, "renamed", 2, 0))
2911 .expect("rename key");
2912 assert!(
2913 segment.key_lookup.is_some(),
2914 "key replacements maintain the index"
2915 );
2916 assert_eq!(keyed_probe(&segment, "hash"), None);
2917 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(3)));
2918
2919 *segment.get_mut(EntityId::new(3)).expect("mutable row") = keyed_row(3, "raw", 3, 0);
2920 assert!(
2921 segment.key_lookup.is_some(),
2922 "raw mutation preserves unaffected key candidates"
2923 );
2924 assert_eq!(keyed_probe(&segment, "renamed"), None);
2925 assert_eq!(keyed_probe(&segment, "raw"), Some(EntityId::new(3)));
2926 segment.seal().expect("seal");
2927 segment
2928 .force_update_with_metadata(&keyed_row(5, "sealed", 4, 0), &["key".into()], None)
2929 .expect("sealed rename");
2930 assert_eq!(keyed_probe(&segment, "gap"), None);
2931 assert_eq!(keyed_probe(&segment, "sealed"), Some(EntityId::new(5)));
2932 assert!(segment.force_delete(EntityId::new(1)));
2933 assert!(segment.force_delete(EntityId::new(3)));
2934 assert_eq!(keyed_probe(&segment, "flat"), None);
2935 assert_eq!(keyed_probe(&segment, "raw"), None);
2936
2937 let mut adopted = GrowingSegment::new(2, "keys");
2938 segment.for_each_fast(|entity| {
2939 adopted.adopt_entity(entity.clone(), None);
2940 true
2941 });
2942 assert_eq!(keyed_probe(&adopted, "sealed"), Some(EntityId::new(5)));
2943 assert!(adopted.evict_entity(EntityId::new(5)));
2944 assert_eq!(keyed_probe(&adopted, "sealed"), None);
2945 }
2946
2947 #[test]
2948 #[ignore = "diagnostic microbenchmark; no timing assertions"]
2949 fn key_lookup_scaling_diagnostic() {
2950 use std::{hint::black_box, time::Instant};
2951 for count in [100, 10_000, 100_000] {
2952 let mut segment = GrowingSegment::new(1, "keys");
2953 segment
2954 .bulk_insert(
2955 (1..=count)
2956 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2957 .collect(),
2958 )
2959 .expect("rows");
2960 let key = format!("key-{count}");
2961 let started = Instant::now();
2962 for _ in 0..1000 {
2963 assert_eq!(
2964 keyed_probe(black_box(&segment), black_box(&key)),
2965 Some(EntityId::new(count))
2966 );
2967 }
2968 let indexed = started.elapsed();
2969 let started = Instant::now();
2970 for _ in 0..1000 {
2971 let mut found = None;
2972 segment.for_each_fast(|entity| {
2973 if SegmentKeyLookup::key(entity) == Some(black_box(key.as_str())) {
2974 found = Some(entity.id);
2975 return false;
2976 }
2977 true
2978 });
2979 assert_eq!(black_box(found), Some(EntityId::new(count)));
2980 }
2981 eprintln!(
2982 "rows={count} reads=1000 indexed={indexed:?} scan={:?}",
2983 started.elapsed()
2984 );
2985 }
2986 }
2987
2988 fn pruning_node(id: u64) -> UnifiedEntity {
2989 UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
2990 }
2991
2992 fn pruning_edge(id: u64) -> UnifiedEntity {
2993 UnifiedEntity::graph_edge(
2994 EntityId::new(id),
2995 "mixed",
2996 "a",
2997 "b",
2998 1.0,
2999 Default::default(),
3000 )
3001 }
3002
3003 #[test]
3004 fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
3005 let mut segment = GrowingSegment::new(1, "mixed");
3006 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
3007 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
3008 segment
3009 .bulk_insert(vec![
3010 UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
3011 pruning_node(2),
3012 pruning_edge(4),
3013 ])
3014 .expect("mixed bulk with gap");
3015 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3016 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3017 segment.seal().expect("seal");
3018 assert!(segment.force_delete(EntityId::new(2)));
3019 assert!(segment.force_delete(EntityId::new(4)));
3020 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3022 let mut rebuilt = GrowingSegment::new(2, "mixed");
3023 for item in segment.iter() {
3024 rebuilt.adopt_entity(item.clone(), None);
3025 }
3026 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
3027 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
3028 rebuilt.adopt_entity(pruning_node(5), None);
3029 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
3030 rebuilt.insert(pruning_edge(6)).expect("single edge");
3031 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
3032 }
3033
3034 #[test]
3035 fn graph_presence_covers_replacement_and_unrestricted_mutation() {
3036 for flat in [false, true] {
3037 for operation in 0..4 {
3038 let mut segment = GrowingSegment::new(1, "mixed");
3039 let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
3040 if flat {
3041 segment.bulk_insert(vec![vector]).expect("flat vector");
3042 } else {
3043 segment.insert(vector).expect("hashmap vector");
3044 }
3045 match operation {
3046 0 => segment.update(pruning_node(1)).expect("replace with node"),
3047 1 => segment
3048 .update_hot(pruning_node(1), &[])
3049 .expect("hot replacement"),
3050 2 => {
3051 segment.seal().expect("seal vector");
3052 segment
3053 .force_update_with_metadata(&pruning_node(1), &[], None)
3054 .expect("sealed replacement");
3055 }
3056 _ => {
3057 *segment.get_mut(EntityId::new(1)).expect("mutable vector") =
3058 pruning_node(1)
3059 }
3060 }
3061 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3062 assert!(matches!(
3063 segment.get(EntityId::new(1)).expect("node").kind,
3064 EntityKind::GraphNode(_)
3065 ));
3066 if operation == 3 {
3067 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3068 }
3069 segment
3070 .force_update_with_metadata(&pruning_edge(1), &[], None)
3071 .expect("replace with edge");
3072 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3073 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3074 }
3075 }
3076 }
3077
3078 fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
3079 UnifiedEntity::new(
3080 EntityId::new(id),
3081 EntityKind::TableRow {
3082 table: "keys".into(),
3083 row_id: id,
3084 },
3085 EntityData::Row(RowData::with_names(
3086 vec![Value::text("payload"), key],
3087 vec!["payload".into(), "primary".into()],
3088 )),
3089 )
3090 }
3091
3092 fn primary_key_probe(
3093 segment: &GrowingSegment,
3094 key: Value,
3095 exclude: Option<EntityId>,
3096 ) -> Option<EntityId> {
3097 segment.find_unique_key_conflict(
3098 &[vec!["primary".into()]],
3099 0,
3100 &[format!("{key:?}")],
3101 exclude,
3102 &|_| true,
3103 )
3104 }
3105
3106 #[test]
3107 fn unique_key_lookup_follows_segment_mutations_and_sealing() {
3108 let mut segment = GrowingSegment::new(1, "keys");
3109 segment
3110 .bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
3111 .expect("bulk row");
3112 let payload_bytes = segment.memory_bytes();
3113 assert_eq!(
3114 primary_key_probe(&segment, Value::Integer(10), None),
3115 Some(EntityId::new(1))
3116 );
3117 assert!(segment.memory_bytes() > payload_bytes);
3118 segment
3119 .insert(primary_key_row(3, Value::Integer(30)))
3120 .expect("hashmap fallback row");
3121 segment
3122 .bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
3123 .expect("bulk gap row");
3124 assert_eq!(
3125 primary_key_probe(&segment, Value::Integer(40), None),
3126 Some(EntityId::new(4))
3127 );
3128 assert_eq!(
3129 primary_key_probe(&segment, Value::Integer(30), None),
3130 Some(EntityId::new(3))
3131 );
3132 segment
3133 .update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
3134 .expect("replace flat key");
3135 assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
3136 assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
3137 assert_eq!(
3138 primary_key_probe(&segment, Value::Integer(11), None),
3139 Some(EntityId::new(1))
3140 );
3141 segment
3142 .delete_batch(&[EntityId::new(1), EntityId::new(3)])
3143 .expect("delete flat and map");
3144 assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
3145 assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
3146 segment.seal().expect("seal");
3147 assert_eq!(
3148 primary_key_probe(&segment, Value::Integer(40), None),
3149 Some(EntityId::new(4))
3150 );
3151 segment
3152 .force_update_with_metadata(
3153 &primary_key_row(4, Value::Integer(41)),
3154 &["primary".into()],
3155 None,
3156 )
3157 .expect("sealed update");
3158 assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
3159 assert_eq!(
3160 primary_key_probe(&segment, Value::Integer(41), None),
3161 Some(EntityId::new(4))
3162 );
3163 assert!(segment.force_delete(EntityId::new(4)));
3164 assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
3165 }
3166
3167 #[test]
3168 fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
3169 let mut segment = GrowingSegment::new(1, "keys");
3170 segment
3171 .bulk_insert(vec![
3172 primary_key_row(1, Value::Integer(10)),
3173 primary_key_row(2, Value::Integer(20)),
3174 ])
3175 .expect("initial rows");
3176 let definitions = vec![
3177 vec!["primary".into()],
3178 vec!["payload".into()],
3179 vec!["primary".into(), "payload".into()],
3180 ];
3181 let payload = format!("{:?}", Value::text("payload"));
3182 assert_eq!(
3183 segment.find_unique_key_conflict(
3184 &definitions,
3185 1,
3186 std::slice::from_ref(&payload),
3187 None,
3188 &|entity| entity.id != EntityId::new(1)
3189 ),
3190 Some(EntityId::new(2))
3191 );
3192 assert_eq!(segment.unique_key_lookup.read().len(), 3);
3193 segment
3194 .insert(primary_key_row(3, Value::Integer(30)))
3195 .expect("append");
3196 assert_eq!(
3197 segment.find_unique_key_conflict(
3198 &definitions,
3199 2,
3200 &[format!("{:?}", Value::Integer(30)), payload.clone()],
3201 None,
3202 &|_| true
3203 ),
3204 Some(EntityId::new(3))
3205 );
3206 let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
3207 assert_eq!(
3208 segment.find_unique_key_conflict(
3209 &definitions,
3210 0,
3211 &[format!("{:?}", Value::Integer(20))],
3212 None,
3213 &|_| true
3214 ),
3215 Some(EntityId::new(2))
3216 );
3217 assert_eq!(segment.unique_key_lookup.read().len(), 3);
3218 assert_eq!(
3219 segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
3220 bytes
3221 );
3222 let reduced = vec![vec!["payload".into()]];
3223 assert_eq!(
3224 segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
3225 Some(EntityId::new(1))
3226 );
3227 assert_eq!(segment.unique_key_lookup.read().len(), 1);
3228 assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
3229 }
3230
3231 #[test]
3232 fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
3233 let mut segment = GrowingSegment::new(1, "keys");
3234 let values = [
3235 Value::Float(f64::NAN),
3236 Value::Float(-0.0),
3237 Value::Float(0.0),
3238 Value::Integer(1),
3239 Value::UnsignedInteger(1),
3240 Value::text("1"),
3241 ];
3242 for (offset, value) in values.iter().enumerate() {
3243 let id = offset as u64 + 1;
3244 segment
3245 .insert(primary_key_row(id, value.clone()))
3246 .expect("distinct physical row");
3247 assert_eq!(
3248 primary_key_probe(&segment, value.clone(), None),
3249 Some(EntityId::new(id))
3250 );
3251 assert_eq!(
3252 primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
3253 None
3254 );
3255 }
3256 segment
3257 .insert(primary_key_row(20, Value::Integer(1)))
3258 .expect("duplicate physical key");
3259 assert_eq!(
3260 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
3261 Some(EntityId::new(20))
3262 );
3263 segment
3266 .entities
3267 .insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
3268 assert_eq!(
3269 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
3270 None
3271 );
3272 let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
3274 *entity = primary_key_row(4, Value::Integer(100));
3275 assert_eq!(
3276 primary_key_probe(&segment, Value::Integer(100), None),
3277 Some(EntityId::new(4))
3278 );
3279 }
3280
3281 #[test]
3282 fn test_growing_segment_basic() {
3283 let mut segment = GrowingSegment::new(1, "test");
3284
3285 let entity = UnifiedEntity::table_row(
3286 EntityId::new(1),
3287 "users",
3288 1,
3289 vec![Value::text("Alice".to_string())],
3290 );
3291
3292 let id = segment.insert(entity).unwrap();
3293 assert_eq!(id, EntityId::new(1));
3294 assert!(segment.contains(id));
3295
3296 let stats = segment.stats();
3297 assert_eq!(stats.entity_count, 1);
3298 assert_eq!(stats.row_count, 1);
3299 }
3300
3301 #[test]
3302 fn test_segment_metadata() {
3303 let mut segment = GrowingSegment::new(1, "test");
3304
3305 let entity = UnifiedEntity::table_row(
3306 EntityId::new(1),
3307 "users",
3308 1,
3309 vec![Value::text("Alice".to_string())],
3310 );
3311 segment.insert(entity).unwrap();
3312
3313 let mut meta = Metadata::new();
3314 meta.set("role", MetadataValue::String("admin".to_string()));
3315 meta.set("level", MetadataValue::Int(5));
3316
3317 segment.set_metadata(EntityId::new(1), meta).unwrap();
3318
3319 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
3320 assert_eq!(
3321 retrieved.get("role"),
3322 Some(&MetadataValue::String("admin".to_string()))
3323 );
3324 }
3325
3326 #[test]
3327 fn test_segment_seal() {
3328 let mut segment = GrowingSegment::new(1, "test");
3329
3330 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
3331 segment.insert(entity).unwrap();
3332
3333 assert!(segment.state().is_writable());
3335
3336 segment.seal().unwrap();
3338 assert_eq!(segment.state(), SegmentState::Sealed);
3339
3340 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
3342 assert!(segment.insert(entity2).is_err());
3343 }
3344
3345 #[test]
3346 fn test_should_seal() {
3347 let mut segment = GrowingSegment::new(1, "test");
3348
3349 let config = SegmentConfig {
3350 max_entities: 2,
3351 ..Default::default()
3352 };
3353
3354 assert!(!segment.should_seal(&config));
3355
3356 segment
3357 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
3358 .unwrap();
3359 assert!(!segment.should_seal(&config));
3360
3361 segment
3362 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
3363 .unwrap();
3364 assert!(segment.should_seal(&config));
3365 }
3366
3367 #[test]
3368 fn test_cross_references() {
3369 let mut segment = GrowingSegment::new(1, "test");
3370
3371 let mut entity1 = UnifiedEntity::table_row(
3372 EntityId::new(1),
3373 "hosts",
3374 1,
3375 vec![Value::text("192.168.1.1".to_string())],
3376 );
3377 entity1.add_cross_ref(CrossRef::new(
3378 EntityId::new(1),
3379 EntityId::new(2),
3380 "nodes",
3381 RefType::RowToNode,
3382 ));
3383 segment.insert(entity1).unwrap();
3384
3385 let refs_from = segment.get_references_from(EntityId::new(1));
3386 assert_eq!(refs_from.len(), 1);
3387 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
3388
3389 let refs_to = segment.get_references_to(EntityId::new(2));
3390 assert_eq!(refs_to.len(), 1);
3391 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
3392 }
3393
3394 #[test]
3395 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
3396 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
3397 zone.update(&Value::Email("delta@example.com".to_string()));
3398
3399 let probe = Value::Email("alpha@example.com".to_string());
3400 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
3401
3402 let in_range = Value::Email("charlie@example.com".to_string());
3403 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
3404 }
3405
3406 #[test]
3407 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
3408 let mut segment = GrowingSegment::new(1, "test");
3409
3410 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
3411 let entity = UnifiedEntity::new(
3412 EntityId::new(row_id),
3413 EntityKind::TableRow {
3414 table: "users".into(),
3415 row_id,
3416 },
3417 EntityData::Row(RowData::with_names(
3418 vec![Value::Integer(age)],
3419 vec!["age".to_string()],
3420 )),
3421 );
3422 segment.insert(entity).unwrap();
3423 }
3424
3425 segment.seal().unwrap();
3426
3427 let miss = Value::Integer(500);
3428 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
3429
3430 let hit = Value::Integer(1000);
3431 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
3432 }
3433}