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