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.get(entity.id).is_some_and(|old| {
854 SegmentKeyLookup::key(old) != SegmentKeyLookup::key(entity)
855 || SegmentKeyLookup::queue_dedup(old) != SegmentKeyLookup::queue_dedup(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 visit_queue_dedup(
1766 &self,
1767 queue: &str,
1768 key: Option<&str>,
1769 now: u64,
1770 visit: &mut impl FnMut(&UnifiedEntity) -> bool,
1771 ) -> bool {
1772 let index = self
1773 .key_lookup
1774 .as_ref()
1775 .expect("segment owns its key index");
1776 let mut inspect = |id| {
1777 let Some(entity) = self.get(id) else {
1778 return true;
1779 };
1780 let Some((candidate_queue, candidate_key, expires)) =
1781 SegmentKeyLookup::queue_dedup(entity)
1782 else {
1783 return true;
1784 };
1785 if candidate_queue == queue
1786 && key.map_or(expires <= now, |key| {
1787 candidate_key == Some(key) && expires > now
1788 })
1789 {
1790 return visit(entity);
1791 }
1792 true
1793 };
1794 if let Some(key) = key {
1795 for id in index.queue_candidates(queue, key) {
1796 if !inspect(*id) {
1797 return false;
1798 }
1799 }
1800 } else {
1801 for id in index.queue_expired(queue, now) {
1802 if !inspect(id) {
1803 return false;
1804 }
1805 }
1806 }
1807 for id in index.mutated_candidates() {
1808 if !inspect(*id) {
1809 return false;
1810 }
1811 }
1812 true
1813 }
1814
1815 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1817 if let Some(id) = entity_id_from_probe_key(key) {
1818 return self.has_live_entity(id);
1819 }
1820
1821 let Ok(pk_str) = std::str::from_utf8(key) else {
1822 return false;
1823 };
1824 let pk_key = (self.collection.clone(), pk_str.to_string());
1825 self.pk_index
1826 .get(&pk_key)
1827 .is_some_and(|id| self.has_live_entity(*id))
1828 }
1829
1830 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1832 if let Some(index) = self.key_lookup.as_mut() {
1833 index.remove(entity);
1834 }
1835 if let Some(index) = self.graph_read_index.get_mut() {
1836 index.remove(entity);
1837 self.graph_read_index_bytes
1838 .store(index.memory_bytes(), Ordering::Relaxed);
1839 }
1840 let kind_key = entity.kind.storage_type().to_string();
1842 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1843 set.remove(&entity.id);
1844 }
1845
1846 if let EntityData::Row(row) = &entity.data {
1848 if let Some(first_col) = row.columns.first() {
1849 let pk_str = format!("{:?}", first_col);
1850 self.pk_index
1851 .remove(&(entity.kind.collection().to_string(), pk_str));
1852 }
1853 }
1854
1855 self.cross_ref_forward.remove(&entity.id);
1857 }
1859
1860 fn reindex_for_update(
1870 &mut self,
1871 old: &UpdateIndexSnapshot,
1872 new: &UnifiedEntity,
1873 modified_columns: Option<&[String]>,
1874 ) {
1875 let pk_changed = match &new.data {
1880 EntityData::Row(new_row) => {
1881 if let Some(cols) = modified_columns {
1882 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1885 new_row
1886 .schema
1887 .as_deref()
1888 .and_then(|schema| schema.first().map(|name| name.as_str()))
1889 });
1890 match pk_col_name {
1891 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1892 None => old.pk_value.as_ref() != new_row.columns.first(),
1894 }
1895 } else {
1896 old.pk_value.as_ref() != new_row.columns.first()
1897 }
1898 }
1899 _ => false,
1901 };
1902
1903 if pk_changed {
1904 if let Some((collection, pk_str)) = &old.pk_index_key {
1906 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1907 }
1908 if let EntityData::Row(row) = &new.data {
1910 if let Some(first_col) = row.columns.first() {
1911 let pk_str = format!("{:?}", first_col);
1912 self.pk_index
1913 .insert((new.kind.collection().to_string(), pk_str), new.id);
1914 }
1915 }
1916 }
1917
1918 let new_refs = new.cross_refs();
1920 if old.cross_refs.as_slice() != new_refs {
1921 self.cross_ref_forward.remove(&new.id);
1923 for cross_ref in &old.cross_refs {
1925 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1926 rev.retain(|(src, _)| *src != new.id);
1927 }
1928 }
1929 for cross_ref in new_refs {
1931 self.cross_ref_forward
1932 .entry(cross_ref.source)
1933 .or_default()
1934 .push((cross_ref.target, cross_ref.ref_type));
1935 self.cross_ref_reverse
1936 .entry(cross_ref.target)
1937 .or_default()
1938 .push((cross_ref.source, cross_ref.ref_type));
1939 }
1940 }
1941 }
1942
1943 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1945 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1946 }
1947
1948 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1950 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1951 }
1952
1953 pub fn age_secs(&self) -> u64 {
1955 let now = current_unix_secs();
1956 now.saturating_sub(self.created_at)
1957 }
1958
1959 pub fn idle_secs(&self) -> u64 {
1961 let now = current_unix_secs();
1962 now.saturating_sub(self.last_write_at)
1963 }
1964
1965 pub fn bulk_insert(
1974 &mut self,
1975 entities: Vec<UnifiedEntity>,
1976 ) -> Result<Vec<EntityId>, SegmentError> {
1977 if !self.state.is_writable() {
1978 return Err(SegmentError::NotWritable);
1979 }
1980
1981 for entity in &entities {
1982 self.mark_checkpoint_entity_dirty(entity.id);
1983 }
1984
1985 if let Some(index) = self.key_lookup.as_mut() {
1986 for entity in &entities {
1987 index.insert(entity);
1988 }
1989 }
1990
1991 let indexes = self.unique_key_lookup.get_mut();
1992 for index in indexes.iter_mut() {
1993 for entity in &entities {
1994 index.insert(entity);
1995 }
1996 }
1997 self.unique_key_lookup_bytes.store(
1998 indexes
1999 .iter()
2000 .map(|index| index.memory_bytes() as u64)
2001 .sum(),
2002 Ordering::Relaxed,
2003 );
2004
2005 let n = entities.len();
2006
2007 if entities.is_empty() {
2008 return Ok(Vec::new());
2009 }
2010
2011 for run in
2014 entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
2015 {
2016 let kind = run[0].kind.storage_type();
2017 self.observe_graph_kind(&run[0].kind);
2018 if matches!(
2019 run[0].kind,
2020 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
2021 ) {
2022 if let Some(index) = self.graph_read_index.get_mut() {
2023 for entity in run {
2024 index.insert(entity);
2025 }
2026 self.graph_read_index_bytes
2027 .store(index.memory_bytes(), Ordering::Relaxed);
2028 }
2029 }
2030 let kind_set = match self.kind_index.get_mut(kind) {
2031 Some(ids) => ids,
2032 None => self.kind_index.entry(kind.to_string()).or_default(),
2033 };
2034 kind_set.reserve(run.len());
2035 kind_set.extend(run.iter().map(|entity| entity.id));
2036 }
2037
2038 let now = current_unix_secs();
2039
2040 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
2041
2042 let mut ids = Vec::with_capacity(n);
2043
2044 if self.flat_entities.is_empty() && self.entities.is_empty() {
2046 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
2048 self.use_flat = true;
2049 }
2050
2051 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
2061 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
2062 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
2066
2067 let mut batch_bytes: usize = 0;
2072
2073 if self.use_flat {
2074 self.flat_entities.reserve(n);
2075 for (i, mut entity) in entities.into_iter().enumerate() {
2076 entity.sequence_id = base_seq + i as u64;
2077 let id = entity.id;
2078 ids.push(id);
2079 batch_bytes += Self::estimate_entity_size(&entity);
2080 if let EntityData::Row(row) = &entity.data {
2081 if row.schema.is_some() && !row.columns.is_empty() {
2082 if columnar_zone_updates.is_empty() {
2083 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
2084 columnar_schema = row.schema.clone();
2085 }
2086 for (ci, val) in row.columns.iter().enumerate() {
2087 if !matches!(val, Value::Null) {
2088 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
2089 bucket.push(val.clone());
2090 }
2091 }
2092 }
2093 } else {
2094 for (col, val) in row.iter_fields() {
2095 if !matches!(val, Value::Null) {
2096 named_zone_updates.push((col.to_string(), val.clone()));
2097 }
2098 }
2099 }
2100 }
2101 let expected = self.base_entity_id + self.flat_entities.len() as u64;
2117 if id.raw() == expected {
2118 self.flat_entities.push(entity);
2119 } else {
2120 self.hash_scan_ids.push(id);
2121 self.entities.insert(id, entity);
2122 }
2123 }
2124 } else {
2125 self.entities.reserve(n);
2127 self.hash_scan_ids.reserve(n);
2128 let mut pairs = Vec::with_capacity(n);
2129 for (i, mut entity) in entities.into_iter().enumerate() {
2130 entity.sequence_id = base_seq + i as u64;
2131 let id = entity.id;
2132 ids.push(id);
2133 batch_bytes += Self::estimate_entity_size(&entity);
2134 if let EntityData::Row(row) = &entity.data {
2135 for (col, val) in row.iter_fields() {
2136 if !matches!(val, Value::Null) {
2137 named_zone_updates.push((col.to_string(), val.clone()));
2138 }
2139 }
2140 }
2141 self.hash_scan_ids.push(id);
2142 pairs.push((id, entity));
2143 }
2144 self.entities.extend(pairs);
2145 }
2146
2147 if !columnar_zone_updates.is_empty() {
2151 let schema = columnar_schema.as_ref();
2152 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
2153 if values.is_empty() {
2154 continue;
2155 }
2156 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
2157 continue;
2158 };
2159 let mut iter = values.into_iter();
2165 if let Some(first) = iter.next() {
2166 let previous = self
2167 .col_zones
2168 .get(col_name)
2169 .map_or(0, ColZone::payload_bytes);
2170 let zone = self
2171 .col_zones
2172 .entry(col_name.clone())
2173 .and_modify(|z| z.update(&first))
2174 .or_insert_with(|| ColZone::new(first));
2175 for v in iter {
2176 zone.update(&v);
2177 }
2178 self.zone_payload_bytes = self
2179 .zone_payload_bytes
2180 .saturating_sub(previous)
2181 .saturating_add(zone.payload_bytes());
2182 }
2183 }
2184 }
2185 for (col, val) in named_zone_updates {
2186 self.update_column_zone(&col, &val);
2187 }
2188
2189 self.add_memory(batch_bytes);
2190 self.last_write_at = now;
2191
2192 if self.use_flat {
2194 self.published_flat_len
2195 .store(self.flat_entities.len(), Ordering::Release);
2196 }
2197
2198 Ok(ids)
2199 }
2200
2201 pub fn force_delete(&mut self, id: EntityId) -> bool {
2204 if self.use_flat {
2205 let raw = id.raw();
2206 if raw >= self.base_entity_id {
2207 let idx = (raw - self.base_entity_id) as usize;
2208 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2209 self.tombstone_flat_entity(idx, id);
2210 return true;
2211 }
2212 }
2213 return self.remove_hashmap_entity(id);
2214 }
2215
2216 self.remove_hashmap_entity(id)
2217 }
2218
2219 pub fn force_update_with_metadata(
2222 &mut self,
2223 entity: &UnifiedEntity,
2224 modified_columns: &[String],
2225 metadata: Option<&Metadata>,
2226 ) -> Result<(), SegmentError> {
2227 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
2228 }
2229}
2230
2231impl UnifiedSegment for GrowingSegment {
2232 fn id(&self) -> SegmentId {
2233 self.id
2234 }
2235
2236 fn state(&self) -> SegmentState {
2237 self.state
2238 }
2239
2240 fn collection(&self) -> &str {
2241 &self.collection
2242 }
2243
2244 fn stats(&self) -> SegmentStats {
2245 let mut stats = SegmentStats {
2246 entity_count: self.live_entity_count(),
2247 deleted_count: self.deleted.len(),
2248 memory_bytes: self.resident_bytes() as usize,
2249 ..Default::default()
2250 };
2251
2252 for entity in self.entities.values() {
2253 match &entity.kind {
2254 EntityKind::TableRow { .. } => stats.row_count += 1,
2255 EntityKind::GraphNode(_) => stats.node_count += 1,
2256 EntityKind::GraphEdge(_) => stats.edge_count += 1,
2257 EntityKind::Vector { .. } => stats.vector_count += 1,
2258 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
2259 EntityKind::QueueMessage { .. } => stats.row_count += 1,
2260 }
2261 stats.cross_ref_count += entity.cross_refs().len();
2262 }
2263
2264 stats
2265 }
2266
2267 fn entity_count(&self) -> usize {
2268 self.live_entity_count()
2269 }
2270
2271 fn contains(&self, id: EntityId) -> bool {
2272 self.has_live_entity(id)
2273 }
2274
2275 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2276 if self.deleted.contains_key(&id) {
2277 return None;
2278 }
2279 if self.use_flat {
2280 let raw = id.raw();
2281 if raw >= self.base_entity_id {
2282 let idx = (raw - self.base_entity_id) as usize;
2283 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
2284 return Some(entity);
2285 }
2286 }
2287 self.entities.get(&id)
2293 } else {
2294 self.entities.get(&id)
2295 }
2296 }
2297
2298 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
2299 self.mark_checkpoint_entity_dirty(id);
2300 self.invalidate_unique_key_lookup();
2301 self.invalidate_graph_read_index();
2302 if self.deleted.contains_key(&id) || !self.state.is_writable() {
2303 return None;
2304 }
2305 if let Some(mut index) = self.key_lookup.take() {
2306 if let Some(entity) = self.get(id) {
2307 index.mark_mutated(entity);
2308 }
2309 self.key_lookup = Some(index);
2310 }
2311 self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
2314 if self.use_flat {
2315 let raw = id.raw();
2316 if raw >= self.base_entity_id {
2317 let idx = (raw - self.base_entity_id) as usize;
2318 if self
2319 .flat_entities
2320 .get(idx)
2321 .map(|e| e.id == id)
2322 .unwrap_or(false)
2323 {
2324 return self.flat_entities.get_mut(idx);
2325 }
2326 }
2327 self.entities.get_mut(&id)
2330 } else {
2331 self.entities.get_mut(&id)
2332 }
2333 }
2334
2335 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2336 if !self.state.is_writable() {
2337 return Err(SegmentError::NotWritable);
2338 }
2339
2340 if self.entities.contains_key(&entity.id) {
2341 return Err(SegmentError::AlreadyExists(entity.id));
2342 }
2343
2344 entity.sequence_id = self.next_sequence();
2346
2347 let size = Self::estimate_entity_size(&entity);
2349 self.add_memory(size);
2350
2351 self.index_entity(&entity);
2353
2354 self.update_col_zones_from_entity(&entity);
2356
2357 let id = entity.id;
2359 self.entities.insert(id, entity);
2360
2361 self.last_write_at = current_unix_secs();
2363
2364 Ok(id)
2365 }
2366
2367 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2368 if !self.state.is_writable() {
2369 return Err(SegmentError::NotWritable);
2370 }
2371
2372 self.apply_update_with_metadata(&entity, None)?;
2373 self.last_write_at = current_unix_secs();
2374
2375 Ok(())
2376 }
2377
2378 fn update_hot(
2379 &mut self,
2380 entity: UnifiedEntity,
2381 modified_columns: &[String],
2382 ) -> Result<(), SegmentError> {
2383 if !self.state.is_writable() {
2384 return Err(SegmentError::NotWritable);
2385 }
2386
2387 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
2388 self.last_write_at = current_unix_secs();
2389 Ok(())
2390 }
2391
2392 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2393 if !self.state.is_writable() {
2394 return Err(SegmentError::NotWritable);
2395 }
2396
2397 if self.use_flat {
2399 let raw = id.raw();
2400 if raw >= self.base_entity_id {
2401 let idx = (raw - self.base_entity_id) as usize;
2402 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2403 self.tombstone_flat_entity(idx, id);
2404 return Ok(true);
2405 }
2406 }
2407 return Ok(self.remove_hashmap_entity(id));
2408 }
2409
2410 Ok(self.remove_hashmap_entity(id))
2411 }
2412
2413 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2414 if !self.has_live_entity(id) {
2415 return None;
2416 }
2417 Some(self.metadata.get_all(id))
2418 }
2419
2420 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2421 self.mark_checkpoint_entity_dirty(id);
2422 if !self.state.is_writable() {
2423 return Err(SegmentError::NotWritable);
2424 }
2425
2426 if !self.has_live_entity(id) {
2427 return Err(SegmentError::NotFound(id));
2428 }
2429
2430 self.metadata.set_all(id, &metadata);
2431 Ok(())
2432 }
2433
2434 fn seal(&mut self) -> Result<(), SegmentError> {
2435 if self.state != SegmentState::Growing {
2436 return Err(SegmentError::InvalidState(self.state));
2437 }
2438
2439 self.state = SegmentState::Sealing;
2440
2441 self.rebuild_sealed_col_zones();
2446
2447 self.state = SegmentState::Sealed;
2448 Ok(())
2449 }
2450
2451 fn should_seal(&self, config: &SegmentConfig) -> bool {
2452 if self.entities.len() >= config.max_entities {
2454 return true;
2455 }
2456
2457 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
2459 return true;
2460 }
2461
2462 if self.age_secs() >= config.max_age_secs {
2464 return true;
2465 }
2466
2467 false
2468 }
2469
2470 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2471 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2472 Box::new(self.flat_entities.iter().chain(self.entities.values()))
2476 } else {
2477 Box::new(self.entities.values())
2478 };
2479 if self.deleted.is_empty() {
2480 base
2481 } else {
2482 Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
2483 }
2484 }
2485
2486 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2487 let ids = self.kind_index.get(kind_filter);
2490 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2493 Box::new(self.flat_entities.iter())
2494 } else {
2495 Box::new(std::iter::empty())
2496 };
2497 Box::new(flat.chain(self.entities.values()).filter(move |e| {
2498 if self.deleted.contains_key(&e.id) {
2499 return false;
2500 }
2501 if let Some(ids) = ids {
2502 ids.contains(&e.id)
2503 } else {
2504 false
2505 }
2506 }))
2507 }
2508
2509 fn filter_metadata(
2510 &self,
2511 filters: &[(String, super::metadata::MetadataFilter)],
2512 ) -> Vec<EntityId> {
2513 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
2516 Box::new(self.flat_entities.iter().map(|e| e.id))
2517 } else {
2518 Box::new(std::iter::empty())
2519 };
2520 flat_ids
2521 .chain(self.entities.keys().copied())
2522 .filter(|id| {
2523 if self.deleted.contains_key(id) {
2524 return false;
2525 }
2526 let metadata = self.metadata.get_all(*id);
2527 metadata.matches_all(filters)
2528 })
2529 .collect()
2530 }
2531}
2532
2533fn build_minmax_multi_intervals(
2534 entries: &[(CanonicalKey, Value)],
2535 max_intervals: usize,
2536) -> Vec<ColZone> {
2537 if entries.is_empty() {
2538 return Vec::new();
2539 }
2540 if entries.len() == 1 || max_intervals <= 1 {
2541 return vec![ColZone::with_bounds(
2542 entries[0].1.clone(),
2543 entries[entries.len() - 1].1.clone(),
2544 )];
2545 }
2546
2547 let mut split_points = if entries.len() <= max_intervals {
2548 (1..entries.len()).collect::<Vec<_>>()
2549 } else {
2550 let target_splits = max_intervals - 1;
2551 let mut selected = select_gap_split_points(entries, target_splits);
2552 if selected.len() < target_splits {
2553 for bucket in 1..max_intervals {
2554 let idx = bucket * entries.len() / max_intervals;
2555 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
2556 continue;
2557 }
2558 selected.push(idx);
2559 if selected.len() >= target_splits {
2560 break;
2561 }
2562 }
2563 }
2564 selected.sort_unstable();
2565 selected.dedup();
2566 selected
2567 };
2568
2569 split_points.push(entries.len());
2570
2571 let mut out = Vec::with_capacity(split_points.len());
2572 let mut start = 0usize;
2573 for end in split_points {
2574 if end <= start {
2575 continue;
2576 }
2577 out.push(ColZone::with_bounds(
2578 entries[start].1.clone(),
2579 entries[end - 1].1.clone(),
2580 ));
2581 start = end;
2582 }
2583
2584 if out.is_empty() {
2585 out.push(ColZone::with_bounds(
2586 entries[0].1.clone(),
2587 entries[entries.len() - 1].1.clone(),
2588 ));
2589 }
2590
2591 out
2592}
2593
2594fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
2595 let mut gaps = Vec::new();
2596 for idx in 1..entries.len() {
2597 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
2598 if score > 0.0 {
2599 gaps.push((score, idx));
2600 }
2601 }
2602 }
2603 gaps.sort_by(|left, right| {
2604 right
2605 .0
2606 .partial_cmp(&left.0)
2607 .unwrap_or(std::cmp::Ordering::Equal)
2608 .then_with(|| left.1.cmp(&right.1))
2609 });
2610 gaps.into_iter()
2611 .take(max_splits)
2612 .map(|(_, idx)| idx)
2613 .collect()
2614}
2615
2616fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
2617 if left.family() != right.family() {
2618 return None;
2619 }
2620 match (left, right) {
2621 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
2622 Some(r.saturating_sub(*l) as f64)
2623 }
2624 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
2625 Some(r.saturating_sub(*l) as f64)
2626 }
2627 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
2628 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
2629 }
2630 _ => None,
2631 }
2632}
2633
2634impl UnifiedSegment for std::sync::Arc<parking_lot::RwLock<GrowingSegment>> {
2637 fn id(&self) -> SegmentId {
2638 self.read().id()
2639 }
2640
2641 fn state(&self) -> SegmentState {
2642 self.read().state()
2643 }
2644
2645 fn collection(&self) -> &str {
2646 "unknown"
2648 }
2649
2650 fn stats(&self) -> SegmentStats {
2651 self.read().stats()
2652 }
2653
2654 fn entity_count(&self) -> usize {
2655 self.read().entity_count()
2656 }
2657
2658 fn contains(&self, id: EntityId) -> bool {
2659 self.read().contains(id)
2660 }
2661
2662 fn get(&self, _id: EntityId) -> Option<&UnifiedEntity> {
2663 None
2666 }
2667
2668 fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
2669 None
2670 }
2671
2672 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2673 self.write().insert(entity)
2674 }
2675
2676 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2677 self.write().update(entity)
2678 }
2679
2680 fn update_hot(
2681 &mut self,
2682 entity: UnifiedEntity,
2683 modified_columns: &[String],
2684 ) -> Result<(), SegmentError> {
2685 self.write().update_hot(entity, modified_columns)
2686 }
2687
2688 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2689 self.write().delete(id)
2690 }
2691
2692 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2693 self.read().get_metadata(id)
2694 }
2695
2696 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2697 self.write().set_metadata(id, metadata)
2698 }
2699
2700 fn seal(&mut self) -> Result<(), SegmentError> {
2701 self.write().seal()
2702 }
2703
2704 fn should_seal(&self, config: &SegmentConfig) -> bool {
2705 self.read().should_seal(config)
2706 }
2707
2708 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2709 Box::new(std::iter::empty())
2711 }
2712
2713 fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2714 Box::new(std::iter::empty())
2715 }
2716
2717 fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2718 self.read().filter_metadata(filters)
2719 }
2720}
2721
2722#[cfg(test)]
2723mod tests {
2724 use super::*;
2725 use crate::storage::entity::RowData;
2726 use crate::storage::metadata::MetadataValue;
2727 use reddb_types::Value;
2728
2729 fn keyed_row(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
2730 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
2731 entity.data = EntityData::Row(RowData::with_names(
2732 vec![Value::text(key), Value::Integer(id as i64)],
2733 vec!["key".into(), "value".into()],
2734 ));
2735 entity.set_xmin(xmin);
2736 entity.set_xmax(xmax);
2737 entity
2738 }
2739
2740 fn keyed_probe(segment: &GrowingSegment, key: &str) -> Option<EntityId> {
2741 segment
2742 .find_key_entity(key, true, &mut |_| true)
2743 .map(|entity| entity.id)
2744 }
2745
2746 #[test]
2747 fn bulk_checkpoint_journal_preserves_prior_changes_and_ignores_rejected_batches() {
2748 let mut segment = GrowingSegment::new(1, "keys");
2749 segment.insert(keyed_row(1, "one", 1, 0)).expect("row");
2750 segment
2751 .bulk_insert(vec![keyed_row(2, "two", 1, 0), keyed_row(3, "three", 1, 0)])
2752 .expect("batch");
2753 assert_eq!(
2754 segment.checkpoint_mutations_since(0),
2755 Some(vec![EntityId::new(1), EntityId::new(2), EntityId::new(3)])
2756 );
2757 let revision = segment.checkpoint_revision();
2758 segment.bulk_insert(Vec::new()).expect("empty batch");
2759 assert_eq!(segment.checkpoint_revision(), revision);
2760 segment.seal().expect("seal segment");
2761 assert!(matches!(
2762 segment.bulk_insert(vec![keyed_row(4, "four", 1, 0)]),
2763 Err(SegmentError::NotWritable)
2764 ));
2765 assert_eq!(segment.checkpoint_revision(), revision);
2766 assert_eq!(segment.checkpoint_mutations_since(revision), Some(vec![]));
2767 }
2768
2769 #[test]
2770 fn checkpoint_journal_retains_bounded_row_versions_through_repeated_updates_and_gc() {
2771 const ROWS: u64 = 1700;
2772 let mut segment = GrowingSegment::new(1, "keys");
2773 for id in 1..=ROWS {
2774 segment
2775 .insert(keyed_row(id, "export", 1, 0))
2776 .expect("seed export");
2777 }
2778 let baseline = segment.checkpoint_revision();
2779 for id in 1..=ROWS {
2783 segment
2784 .update(keyed_row(id, "export", 1, 2))
2785 .expect("supersede");
2786 let delivered = ROWS + id;
2787 segment
2788 .insert(keyed_row(delivered, "export", 2, 0))
2789 .expect("delivered");
2790 segment
2791 .update(keyed_row(delivered, "export", 2, 3))
2792 .expect("acknowledge");
2793 segment
2794 .insert(keyed_row(2 * ROWS + id, "export", 3, 0))
2795 .expect("final version");
2796 assert_eq!(
2797 segment
2798 .delete_batch(&[EntityId::new(id), EntityId::new(delivered)])
2799 .expect("GC")
2800 .len(),
2801 2
2802 );
2803 }
2804 let changed = segment
2805 .checkpoint_mutations_since(baseline)
2806 .expect("all changed IDs still fit the bounded checkpoint journal");
2807 let actual: std::collections::BTreeSet<_> = changed.into_iter().collect();
2808 let expected: std::collections::BTreeSet<_> = (1..=3 * ROWS).map(EntityId::new).collect();
2809 assert_eq!(actual, expected, "deleted versions must remain tracked too");
2810 }
2811
2812 #[test]
2813 fn checkpoint_journal_preserves_cold_changes_during_hot_row_churn() {
2814 let mut segment = GrowingSegment::new(1, "keys");
2815 segment
2816 .insert(keyed_row(1, "cold", 1, 0))
2817 .expect("cold row");
2818 segment.insert(keyed_row(2, "hot", 1, 0)).expect("hot row");
2819 for xid in 2..=(5 * CHECKPOINT_MUTATIONS_MAX) as u64 {
2820 segment
2821 .update(keyed_row(2, "hot", xid, 0))
2822 .expect("hot edit");
2823 }
2824 let changed: std::collections::BTreeSet<_> = segment
2825 .checkpoint_mutations_since(0)
2826 .expect("two physical IDs never exhaust the distinct-ID bound")
2827 .into_iter()
2828 .collect();
2829 assert_eq!(
2830 changed,
2831 [EntityId::new(1), EntityId::new(2)].into_iter().collect()
2832 );
2833 assert!(
2834 segment.checkpoint_tracking_bytes() <= 1024 * 1024,
2835 "retained proof memory must stay bounded under unlimited repeated edits"
2836 );
2837 }
2838
2839 #[test]
2840 fn checkpoint_journal_complete_suffix_matches_uncompressed_mutation_history() {
2841 let mut segment = GrowingSegment::new(1, "keys");
2842 let mut history = Vec::new();
2843 let mut seed = 0x7b3d_5961_a902_4c85_u64;
2844 for n in 0..3 * CHECKPOINT_MUTATIONS_MAX {
2845 seed ^= seed << 13;
2846 seed ^= seed >> 7;
2847 seed ^= seed << 17;
2848 let raw = if n < CHECKPOINT_MUTATIONS_MAX + 17 {
2849 n as u64 + 1
2850 } else {
2851 seed % (CHECKPOINT_MUTATIONS_MAX as u64 + 17) + 1
2852 };
2853 let id = EntityId::new(raw);
2854 segment.mark_checkpoint_entity_dirty(id);
2855 history.push(id);
2856 }
2857 let now = segment.checkpoint_revision();
2858 assert!(
2859 segment.checkpoint_mutations_since(0).is_none(),
2860 "an evicted ID must still invalidate an older baseline"
2861 );
2862 let mut complete = 0;
2863 for baseline in [0, 1, now / 3, now / 2, now - 2048, now - 128, now] {
2864 if let Some(ids) = segment.checkpoint_mutations_since(baseline) {
2865 let actual: std::collections::BTreeSet<_> = ids.into_iter().collect();
2866 let expected: std::collections::BTreeSet<_> =
2867 history[baseline as usize..].iter().copied().collect();
2868 assert_eq!(actual, expected, "baseline {baseline}");
2869 complete += 1;
2870 }
2871 }
2872 assert!(
2873 complete >= 3,
2874 "recent baselines must retain a complete proof"
2875 );
2876 assert!(segment.checkpoint_mutations_since(now + 1).is_none());
2877 assert!(segment.checkpoint_tracking_bytes() <= 1024 * 1024);
2878 }
2879
2880 #[test]
2881 fn bulk_checkpoint_journal_exhaustion_keeps_a_bounded_complete_suffix() {
2882 let mut segment = GrowingSegment::new(1, "keys");
2883 segment
2884 .bulk_insert(
2885 (1..=CHECKPOINT_MUTATIONS_MAX as u64 + 1)
2886 .map(|id| keyed_row(id, "key", 1, 0))
2887 .collect(),
2888 )
2889 .expect("large batch");
2890 assert!(segment.checkpoint_mutations_since(0).is_none());
2891 assert_eq!(
2892 segment
2893 .checkpoint_mutations_since(1)
2894 .expect("complete suffix")
2895 .len(),
2896 CHECKPOINT_MUTATIONS_MAX
2897 );
2898 let revision = segment.checkpoint_revision();
2899 let id = EntityId::new(CHECKPOINT_MUTATIONS_MAX as u64 + 2);
2900 segment
2901 .bulk_insert(vec![keyed_row(id.raw(), "next", 1, 0)])
2902 .expect("later batch");
2903 assert_eq!(segment.checkpoint_mutations_since(revision), Some(vec![id]));
2904 assert_eq!(segment.checkpoint_mutations.len(), CHECKPOINT_MUTATIONS_MAX);
2905 }
2906
2907 fn dedup_row(id: u64, queue: &str, key: &str, expires: u64) -> UnifiedEntity {
2908 let mut entity = keyed_row(id, "unused", 1, 0);
2909 entity.data = EntityData::Row(RowData::with_names(
2910 vec![
2911 Value::text("queue_push_dedup"),
2912 Value::text(queue),
2913 Value::text(key),
2914 Value::UnsignedInteger(expires),
2915 ],
2916 vec![
2917 "kind".into(),
2918 "queue".into(),
2919 "dedup_key".into(),
2920 "expires_at_ns".into(),
2921 ],
2922 ));
2923 entity
2924 }
2925
2926 #[test]
2927 fn queue_dedup_candidates_exclude_unrelated_history_and_track_mutations() {
2928 let mut segment = GrowingSegment::new(1, "meta");
2929 segment
2930 .bulk_insert(
2931 (1..=4096)
2932 .map(|id| dedup_row(id, "other", "same", 1))
2933 .collect(),
2934 )
2935 .expect("unrelated history");
2936 segment
2937 .insert(dedup_row(5000, "wanted", "same", 100))
2938 .expect("live key");
2939 segment
2940 .insert(dedup_row(5001, "wanted", "old", 50))
2941 .expect("expired key");
2942 let candidates = |segment: &GrowingSegment, key, now| {
2943 let mut ids = Vec::new();
2944 segment.visit_queue_dedup("wanted", key, now, &mut |entity| {
2945 ids.push(entity.id.raw());
2946 true
2947 });
2948 ids
2949 };
2950 assert_eq!(candidates(&segment, Some("same"), 50), vec![5000]);
2951 assert!(candidates(&segment, Some("absent"), 50).is_empty());
2952 assert_eq!(candidates(&segment, None, 50), vec![5001]);
2953 segment
2954 .update(dedup_row(5000, "wanted", "changed", 40))
2955 .expect("replace indexed fields");
2956 assert!(candidates(&segment, Some("same"), 50).is_empty());
2957 let mut expired = candidates(&segment, None, 50);
2958 expired.sort_unstable();
2959 assert_eq!(expired, vec![5000, 5001]);
2960 *segment.get_mut(EntityId::new(5000)).expect("raw access") =
2961 dedup_row(5000, "wanted", "raw", 200);
2962 assert_eq!(candidates(&segment, Some("raw"), 100), vec![5000]);
2963 assert_eq!(candidates(&segment, None, 100), vec![5001]);
2964 segment.force_delete(EntityId::new(5001));
2965 assert!(candidates(&segment, None, 100).is_empty());
2966 }
2967
2968 #[test]
2969 fn key_lookup_visits_only_requested_key_versions() {
2970 let mut segment = GrowingSegment::new(1, "keys");
2971 segment
2972 .bulk_insert(
2973 (1..=4096)
2974 .map(|id| keyed_row(id, &format!("other-{id}"), 1, 0))
2975 .collect(),
2976 )
2977 .expect("unrelated rows");
2978 segment
2979 .insert(keyed_row(5000, "wanted", 2, 8))
2980 .expect("old version");
2981 segment
2982 .insert(keyed_row(5001, "wanted", 8, 0))
2983 .expect("new version");
2984 let mut visits = 0;
2985 let result = segment.find_key_entity("wanted", true, &mut |entity| {
2986 visits += 1;
2987 entity.xmin < 8
2988 });
2989 assert_eq!(result.map(|entity| entity.id), Some(EntityId::new(5000)));
2990 assert_eq!(
2991 visits, 2,
2992 "unrelated rows must never reach visibility evaluation"
2993 );
2994 assert_eq!(keyed_probe(&segment, "wanted"), Some(EntityId::new(5001)));
2995 let mut misses = 0;
2996 assert!(segment
2997 .find_key_entity("absent", true, &mut |_| {
2998 misses += 1;
2999 true
3000 })
3001 .is_none());
3002 assert_eq!(misses, 0);
3003
3004 let mut collision = SegmentKeyLookup::default();
3006 collision.insert(&keyed_row(1, "wanted", 1, 0));
3007 assert!(segment
3008 .find_key_candidate(&collision, "wanted", true, &mut |_| true)
3009 .is_none());
3010 }
3011
3012 #[test]
3013 fn raw_mutation_does_not_disable_unrelated_key_candidates() {
3014 let mut segment = GrowingSegment::new(1, "keys");
3015 segment
3016 .bulk_insert(
3017 (1..=4096)
3018 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
3019 .collect(),
3020 )
3021 .expect("rows");
3022 for _ in 0..3 {
3023 *segment.get_mut(EntityId::new(1)).expect("mutable") = keyed_row(1, "renamed", 2, 0);
3024 }
3025 let index = segment.key_lookup.as_ref().expect("index retained");
3026 assert_eq!(index.candidates("key-4096"), &[EntityId::new(4096)]);
3027 assert_eq!(index.mutated_candidates().count(), 1);
3028 assert!(index.candidates("key-1").is_empty());
3029 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(1)));
3030 assert_eq!(keyed_probe(&segment, "key-4096"), Some(EntityId::new(4096)));
3031 segment.delete(EntityId::new(1)).expect("delete raw row");
3032 assert_eq!(
3033 segment
3034 .key_lookup
3035 .as_ref()
3036 .unwrap()
3037 .mutated_candidates()
3038 .count(),
3039 0
3040 );
3041 assert_eq!(keyed_probe(&segment, "renamed"), None);
3042 }
3043
3044 #[test]
3045 fn key_lookup_tracks_mutations_sealing_and_adoption() {
3046 let mut segment = GrowingSegment::new(1, "keys");
3047 segment
3048 .bulk_insert(vec![keyed_row(1, "flat", 1, 0)])
3049 .expect("flat row");
3050 segment
3051 .insert(keyed_row(3, "hash", 1, 0))
3052 .expect("hash fallback");
3053 segment
3054 .bulk_insert(vec![keyed_row(5, "gap", 1, 0)])
3055 .expect("bulk gap");
3056 assert!(segment.key_lookup.as_ref().expect("index").memory_bytes() > 0);
3057 assert_eq!(keyed_probe(&segment, "gap"), Some(EntityId::new(5)));
3058 segment
3059 .update(keyed_row(3, "renamed", 2, 0))
3060 .expect("rename key");
3061 assert!(
3062 segment.key_lookup.is_some(),
3063 "key replacements maintain the index"
3064 );
3065 assert_eq!(keyed_probe(&segment, "hash"), None);
3066 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(3)));
3067
3068 *segment.get_mut(EntityId::new(3)).expect("mutable row") = keyed_row(3, "raw", 3, 0);
3069 assert!(
3070 segment.key_lookup.is_some(),
3071 "raw mutation preserves unaffected key candidates"
3072 );
3073 assert_eq!(keyed_probe(&segment, "renamed"), None);
3074 assert_eq!(keyed_probe(&segment, "raw"), Some(EntityId::new(3)));
3075 segment.seal().expect("seal");
3076 segment
3077 .force_update_with_metadata(&keyed_row(5, "sealed", 4, 0), &["key".into()], None)
3078 .expect("sealed rename");
3079 assert_eq!(keyed_probe(&segment, "gap"), None);
3080 assert_eq!(keyed_probe(&segment, "sealed"), Some(EntityId::new(5)));
3081 assert!(segment.force_delete(EntityId::new(1)));
3082 assert!(segment.force_delete(EntityId::new(3)));
3083 assert_eq!(keyed_probe(&segment, "flat"), None);
3084 assert_eq!(keyed_probe(&segment, "raw"), None);
3085
3086 let mut adopted = GrowingSegment::new(2, "keys");
3087 segment.for_each_fast(|entity| {
3088 adopted.adopt_entity(entity.clone(), None);
3089 true
3090 });
3091 assert_eq!(keyed_probe(&adopted, "sealed"), Some(EntityId::new(5)));
3092 assert!(adopted.evict_entity(EntityId::new(5)));
3093 assert_eq!(keyed_probe(&adopted, "sealed"), None);
3094 }
3095
3096 #[test]
3097 #[ignore = "diagnostic microbenchmark; no timing assertions"]
3098 fn key_lookup_scaling_diagnostic() {
3099 use std::{hint::black_box, time::Instant};
3100 for count in [100, 10_000, 100_000] {
3101 let mut segment = GrowingSegment::new(1, "keys");
3102 segment
3103 .bulk_insert(
3104 (1..=count)
3105 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
3106 .collect(),
3107 )
3108 .expect("rows");
3109 let key = format!("key-{count}");
3110 let started = Instant::now();
3111 for _ in 0..1000 {
3112 assert_eq!(
3113 keyed_probe(black_box(&segment), black_box(&key)),
3114 Some(EntityId::new(count))
3115 );
3116 }
3117 let indexed = started.elapsed();
3118 let started = Instant::now();
3119 for _ in 0..1000 {
3120 let mut found = None;
3121 segment.for_each_fast(|entity| {
3122 if SegmentKeyLookup::key(entity) == Some(black_box(key.as_str())) {
3123 found = Some(entity.id);
3124 return false;
3125 }
3126 true
3127 });
3128 assert_eq!(black_box(found), Some(EntityId::new(count)));
3129 }
3130 eprintln!(
3131 "rows={count} reads=1000 indexed={indexed:?} scan={:?}",
3132 started.elapsed()
3133 );
3134 }
3135 }
3136
3137 fn pruning_node(id: u64) -> UnifiedEntity {
3138 UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
3139 }
3140
3141 fn pruning_edge(id: u64) -> UnifiedEntity {
3142 UnifiedEntity::graph_edge(
3143 EntityId::new(id),
3144 "mixed",
3145 "a",
3146 "b",
3147 1.0,
3148 Default::default(),
3149 )
3150 }
3151
3152 #[test]
3153 fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
3154 let mut segment = GrowingSegment::new(1, "mixed");
3155 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
3156 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
3157 segment
3158 .bulk_insert(vec![
3159 UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
3160 pruning_node(2),
3161 pruning_edge(4),
3162 ])
3163 .expect("mixed bulk with gap");
3164 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3165 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3166 segment.seal().expect("seal");
3167 assert!(segment.force_delete(EntityId::new(2)));
3168 assert!(segment.force_delete(EntityId::new(4)));
3169 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3171 let mut rebuilt = GrowingSegment::new(2, "mixed");
3172 for item in segment.iter() {
3173 rebuilt.adopt_entity(item.clone(), None);
3174 }
3175 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
3176 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
3177 rebuilt.adopt_entity(pruning_node(5), None);
3178 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
3179 rebuilt.insert(pruning_edge(6)).expect("single edge");
3180 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
3181 }
3182
3183 #[test]
3184 fn graph_presence_covers_replacement_and_unrestricted_mutation() {
3185 for flat in [false, true] {
3186 for operation in 0..4 {
3187 let mut segment = GrowingSegment::new(1, "mixed");
3188 let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
3189 if flat {
3190 segment.bulk_insert(vec![vector]).expect("flat vector");
3191 } else {
3192 segment.insert(vector).expect("hashmap vector");
3193 }
3194 match operation {
3195 0 => segment.update(pruning_node(1)).expect("replace with node"),
3196 1 => segment
3197 .update_hot(pruning_node(1), &[])
3198 .expect("hot replacement"),
3199 2 => {
3200 segment.seal().expect("seal vector");
3201 segment
3202 .force_update_with_metadata(&pruning_node(1), &[], None)
3203 .expect("sealed replacement");
3204 }
3205 _ => {
3206 *segment.get_mut(EntityId::new(1)).expect("mutable vector") =
3207 pruning_node(1)
3208 }
3209 }
3210 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3211 assert!(matches!(
3212 segment.get(EntityId::new(1)).expect("node").kind,
3213 EntityKind::GraphNode(_)
3214 ));
3215 if operation == 3 {
3216 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3217 }
3218 segment
3219 .force_update_with_metadata(&pruning_edge(1), &[], None)
3220 .expect("replace with edge");
3221 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
3222 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
3223 }
3224 }
3225 }
3226
3227 fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
3228 UnifiedEntity::new(
3229 EntityId::new(id),
3230 EntityKind::TableRow {
3231 table: "keys".into(),
3232 row_id: id,
3233 },
3234 EntityData::Row(RowData::with_names(
3235 vec![Value::text("payload"), key],
3236 vec!["payload".into(), "primary".into()],
3237 )),
3238 )
3239 }
3240
3241 fn primary_key_probe(
3242 segment: &GrowingSegment,
3243 key: Value,
3244 exclude: Option<EntityId>,
3245 ) -> Option<EntityId> {
3246 segment.find_unique_key_conflict(
3247 &[vec!["primary".into()]],
3248 0,
3249 &[format!("{key:?}")],
3250 exclude,
3251 &|_| true,
3252 )
3253 }
3254
3255 #[test]
3256 fn unique_key_lookup_follows_segment_mutations_and_sealing() {
3257 let mut segment = GrowingSegment::new(1, "keys");
3258 segment
3259 .bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
3260 .expect("bulk row");
3261 let payload_bytes = segment.memory_bytes();
3262 assert_eq!(
3263 primary_key_probe(&segment, Value::Integer(10), None),
3264 Some(EntityId::new(1))
3265 );
3266 assert!(segment.memory_bytes() > payload_bytes);
3267 segment
3268 .insert(primary_key_row(3, Value::Integer(30)))
3269 .expect("hashmap fallback row");
3270 segment
3271 .bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
3272 .expect("bulk gap row");
3273 assert_eq!(
3274 primary_key_probe(&segment, Value::Integer(40), None),
3275 Some(EntityId::new(4))
3276 );
3277 assert_eq!(
3278 primary_key_probe(&segment, Value::Integer(30), None),
3279 Some(EntityId::new(3))
3280 );
3281 segment
3282 .update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
3283 .expect("replace flat key");
3284 assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
3285 assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
3286 assert_eq!(
3287 primary_key_probe(&segment, Value::Integer(11), None),
3288 Some(EntityId::new(1))
3289 );
3290 segment
3291 .delete_batch(&[EntityId::new(1), EntityId::new(3)])
3292 .expect("delete flat and map");
3293 assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
3294 assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
3295 segment.seal().expect("seal");
3296 assert_eq!(
3297 primary_key_probe(&segment, Value::Integer(40), None),
3298 Some(EntityId::new(4))
3299 );
3300 segment
3301 .force_update_with_metadata(
3302 &primary_key_row(4, Value::Integer(41)),
3303 &["primary".into()],
3304 None,
3305 )
3306 .expect("sealed update");
3307 assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
3308 assert_eq!(
3309 primary_key_probe(&segment, Value::Integer(41), None),
3310 Some(EntityId::new(4))
3311 );
3312 assert!(segment.force_delete(EntityId::new(4)));
3313 assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
3314 }
3315
3316 #[test]
3317 fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
3318 let mut segment = GrowingSegment::new(1, "keys");
3319 segment
3320 .bulk_insert(vec![
3321 primary_key_row(1, Value::Integer(10)),
3322 primary_key_row(2, Value::Integer(20)),
3323 ])
3324 .expect("initial rows");
3325 let definitions = vec![
3326 vec!["primary".into()],
3327 vec!["payload".into()],
3328 vec!["primary".into(), "payload".into()],
3329 ];
3330 let payload = format!("{:?}", Value::text("payload"));
3331 assert_eq!(
3332 segment.find_unique_key_conflict(
3333 &definitions,
3334 1,
3335 std::slice::from_ref(&payload),
3336 None,
3337 &|entity| entity.id != EntityId::new(1)
3338 ),
3339 Some(EntityId::new(2))
3340 );
3341 assert_eq!(segment.unique_key_lookup.read().len(), 3);
3342 segment
3343 .insert(primary_key_row(3, Value::Integer(30)))
3344 .expect("append");
3345 assert_eq!(
3346 segment.find_unique_key_conflict(
3347 &definitions,
3348 2,
3349 &[format!("{:?}", Value::Integer(30)), payload.clone()],
3350 None,
3351 &|_| true
3352 ),
3353 Some(EntityId::new(3))
3354 );
3355 let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
3356 assert_eq!(
3357 segment.find_unique_key_conflict(
3358 &definitions,
3359 0,
3360 &[format!("{:?}", Value::Integer(20))],
3361 None,
3362 &|_| true
3363 ),
3364 Some(EntityId::new(2))
3365 );
3366 assert_eq!(segment.unique_key_lookup.read().len(), 3);
3367 assert_eq!(
3368 segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
3369 bytes
3370 );
3371 let reduced = vec![vec!["payload".into()]];
3372 assert_eq!(
3373 segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
3374 Some(EntityId::new(1))
3375 );
3376 assert_eq!(segment.unique_key_lookup.read().len(), 1);
3377 assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
3378 }
3379
3380 #[test]
3381 fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
3382 let mut segment = GrowingSegment::new(1, "keys");
3383 let values = [
3384 Value::Float(f64::NAN),
3385 Value::Float(-0.0),
3386 Value::Float(0.0),
3387 Value::Integer(1),
3388 Value::UnsignedInteger(1),
3389 Value::text("1"),
3390 ];
3391 for (offset, value) in values.iter().enumerate() {
3392 let id = offset as u64 + 1;
3393 segment
3394 .insert(primary_key_row(id, value.clone()))
3395 .expect("distinct physical row");
3396 assert_eq!(
3397 primary_key_probe(&segment, value.clone(), None),
3398 Some(EntityId::new(id))
3399 );
3400 assert_eq!(
3401 primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
3402 None
3403 );
3404 }
3405 segment
3406 .insert(primary_key_row(20, Value::Integer(1)))
3407 .expect("duplicate physical key");
3408 assert_eq!(
3409 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
3410 Some(EntityId::new(20))
3411 );
3412 segment
3415 .entities
3416 .insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
3417 assert_eq!(
3418 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
3419 None
3420 );
3421 let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
3423 *entity = primary_key_row(4, Value::Integer(100));
3424 assert_eq!(
3425 primary_key_probe(&segment, Value::Integer(100), None),
3426 Some(EntityId::new(4))
3427 );
3428 }
3429
3430 #[test]
3431 fn test_growing_segment_basic() {
3432 let mut segment = GrowingSegment::new(1, "test");
3433
3434 let entity = UnifiedEntity::table_row(
3435 EntityId::new(1),
3436 "users",
3437 1,
3438 vec![Value::text("Alice".to_string())],
3439 );
3440
3441 let id = segment.insert(entity).unwrap();
3442 assert_eq!(id, EntityId::new(1));
3443 assert!(segment.contains(id));
3444
3445 let stats = segment.stats();
3446 assert_eq!(stats.entity_count, 1);
3447 assert_eq!(stats.row_count, 1);
3448 }
3449
3450 #[test]
3451 fn test_segment_metadata() {
3452 let mut segment = GrowingSegment::new(1, "test");
3453
3454 let entity = UnifiedEntity::table_row(
3455 EntityId::new(1),
3456 "users",
3457 1,
3458 vec![Value::text("Alice".to_string())],
3459 );
3460 segment.insert(entity).unwrap();
3461
3462 let mut meta = Metadata::new();
3463 meta.set("role", MetadataValue::String("admin".to_string()));
3464 meta.set("level", MetadataValue::Int(5));
3465
3466 segment.set_metadata(EntityId::new(1), meta).unwrap();
3467
3468 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
3469 assert_eq!(
3470 retrieved.get("role"),
3471 Some(&MetadataValue::String("admin".to_string()))
3472 );
3473 }
3474
3475 #[test]
3476 fn test_segment_seal() {
3477 let mut segment = GrowingSegment::new(1, "test");
3478
3479 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
3480 segment.insert(entity).unwrap();
3481
3482 assert!(segment.state().is_writable());
3484
3485 segment.seal().unwrap();
3487 assert_eq!(segment.state(), SegmentState::Sealed);
3488
3489 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
3491 assert!(segment.insert(entity2).is_err());
3492 }
3493
3494 #[test]
3495 fn test_should_seal() {
3496 let mut segment = GrowingSegment::new(1, "test");
3497
3498 let config = SegmentConfig {
3499 max_entities: 2,
3500 ..Default::default()
3501 };
3502
3503 assert!(!segment.should_seal(&config));
3504
3505 segment
3506 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
3507 .unwrap();
3508 assert!(!segment.should_seal(&config));
3509
3510 segment
3511 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
3512 .unwrap();
3513 assert!(segment.should_seal(&config));
3514 }
3515
3516 #[test]
3517 fn test_cross_references() {
3518 let mut segment = GrowingSegment::new(1, "test");
3519
3520 let mut entity1 = UnifiedEntity::table_row(
3521 EntityId::new(1),
3522 "hosts",
3523 1,
3524 vec![Value::text("192.168.1.1".to_string())],
3525 );
3526 entity1.add_cross_ref(CrossRef::new(
3527 EntityId::new(1),
3528 EntityId::new(2),
3529 "nodes",
3530 RefType::RowToNode,
3531 ));
3532 segment.insert(entity1).unwrap();
3533
3534 let refs_from = segment.get_references_from(EntityId::new(1));
3535 assert_eq!(refs_from.len(), 1);
3536 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
3537
3538 let refs_to = segment.get_references_to(EntityId::new(2));
3539 assert_eq!(refs_to.len(), 1);
3540 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
3541 }
3542
3543 #[test]
3544 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
3545 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
3546 zone.update(&Value::Email("delta@example.com".to_string()));
3547
3548 let probe = Value::Email("alpha@example.com".to_string());
3549 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
3550
3551 let in_range = Value::Email("charlie@example.com".to_string());
3552 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
3553 }
3554
3555 #[test]
3556 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
3557 let mut segment = GrowingSegment::new(1, "test");
3558
3559 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
3560 let entity = UnifiedEntity::new(
3561 EntityId::new(row_id),
3562 EntityKind::TableRow {
3563 table: "users".into(),
3564 row_id,
3565 },
3566 EntityData::Row(RowData::with_names(
3567 vec![Value::Integer(age)],
3568 vec!["age".to_string()],
3569 )),
3570 );
3571 segment.insert(entity).unwrap();
3572 }
3573
3574 segment.seal().unwrap();
3575
3576 let miss = Value::Integer(500);
3577 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
3578
3579 let hit = Value::Integer(1000);
3580 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
3581 }
3582}