1use std::collections::{BTreeMap, HashMap, HashSet};
32use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
33use std::time::{SystemTime, UNIX_EPOCH};
34
35use super::entity::{CrossRef, EntityData, EntityId, EntityKind, RefType, UnifiedEntity};
36use super::metadata::{Metadata, MetadataStorage};
37use super::segment_graph_index::SegmentGraphIndex;
38use super::segment_key_lookup::SegmentKeyLookup;
39use super::segment_unique_key::SegmentUniqueKey;
40use crate::storage::query::value_compare::partial_compare_values;
41use reddb_types::{value_to_canonical_key, CanonicalKey, Value};
42
43pub type SegmentId = u64;
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum SegmentState {
49 Growing,
51 Sealing,
53 Sealed,
55 Flushed,
57 Archived,
59}
60
61impl SegmentState {
62 pub fn is_writable(&self) -> bool {
64 matches!(self, Self::Growing)
65 }
66
67 pub fn is_queryable(&self) -> bool {
69 !matches!(self, Self::Sealing)
70 }
71
72 pub fn is_immutable(&self) -> bool {
74 matches!(self, Self::Sealed | Self::Flushed | Self::Archived)
75 }
76}
77
78#[derive(Debug, Clone)]
80pub struct SegmentConfig {
81 pub max_entities: usize,
83 pub max_bytes: usize,
85 pub max_age_secs: u64,
87 pub build_vector_index: bool,
89 pub build_graph_index: bool,
91 pub compression_level: u8,
93}
94
95impl Default for SegmentConfig {
96 fn default() -> Self {
97 Self {
98 max_entities: 100_000,
99 max_bytes: 256 * 1024 * 1024, max_age_secs: 3600, build_vector_index: true,
102 build_graph_index: true,
103 compression_level: 6,
104 }
105 }
106}
107
108#[derive(Debug, Clone, Default)]
110pub struct SegmentStats {
111 pub entity_count: usize,
113 pub deleted_count: usize,
115 pub memory_bytes: usize,
117 pub vector_count: usize,
119 pub node_count: usize,
121 pub edge_count: usize,
123 pub row_count: usize,
125 pub cross_ref_count: usize,
127}
128
129#[derive(Debug, Clone)]
131pub enum SegmentError {
132 NotWritable,
134 NotFound(EntityId),
136 AlreadyExists(EntityId),
138 Full,
140 InvalidState(SegmentState),
142 Internal(String),
144}
145
146impl std::fmt::Display for SegmentError {
147 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
148 match self {
149 Self::NotWritable => write!(f, "segment is not writable"),
150 Self::NotFound(id) => write!(f, "entity not found: {}", id),
151 Self::AlreadyExists(id) => write!(f, "entity already exists: {}", id),
152 Self::Full => write!(f, "segment is full"),
153 Self::InvalidState(state) => write!(f, "invalid operation for state: {:?}", state),
154 Self::Internal(msg) => write!(f, "internal error: {}", msg),
155 }
156 }
157}
158
159impl std::error::Error for SegmentError {}
160
161fn current_unix_secs() -> u64 {
162 SystemTime::now()
163 .duration_since(UNIX_EPOCH)
164 .unwrap_or_default()
165 .as_secs()
166}
167
168fn entity_id_from_probe_key(key: &[u8]) -> Option<EntityId> {
169 if let Ok(text) = std::str::from_utf8(key) {
170 if let Ok(raw) = text.parse::<u64>() {
171 return Some(EntityId::new(raw));
172 }
173 }
174
175 let bytes: [u8; 8] = key.try_into().ok()?;
176 Some(EntityId::new(u64::from_le_bytes(bytes)))
177}
178
179const SEALED_MULTI_ZONE_MAX_INTERVALS: usize = 4;
180
181const TOMBSTONE_ENTRY_BYTES: u64 = 24;
189
190#[derive(Debug, Clone)]
191struct UpdateIndexSnapshot {
192 pk_column_name: Option<String>,
193 pk_value: Option<Value>,
194 pk_index_key: Option<(String, String)>,
195 cross_refs: Vec<CrossRef>,
196}
197
198impl UpdateIndexSnapshot {
199 fn from_entity(entity: &UnifiedEntity) -> Self {
200 let (pk_column_name, pk_value) = match &entity.data {
201 EntityData::Row(row) => (
202 row.schema
203 .as_deref()
204 .and_then(|schema| schema.first().cloned()),
205 row.columns.first().cloned(),
206 ),
207 _ => (None, None),
208 };
209 let pk_index_key = pk_value
210 .as_ref()
211 .map(|value| (entity.kind.collection().to_string(), format!("{:?}", value)));
212 Self {
213 pk_column_name,
214 pk_value,
215 pk_index_key,
216 cross_refs: entity.cross_refs().to_vec(),
217 }
218 }
219}
220
221pub trait UnifiedSegment: Send + Sync {
223 fn id(&self) -> SegmentId;
225
226 fn state(&self) -> SegmentState;
228
229 fn collection(&self) -> &str;
231
232 fn stats(&self) -> SegmentStats;
234
235 fn entity_count(&self) -> usize;
237
238 fn contains(&self, id: EntityId) -> bool;
240
241 fn get(&self, id: EntityId) -> Option<&UnifiedEntity>;
243
244 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity>;
246
247 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError>;
249
250 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError>;
252
253 fn update_hot(
257 &mut self,
258 entity: UnifiedEntity,
259 modified_columns: &[String],
260 ) -> Result<(), SegmentError> {
261 let _ = modified_columns;
262 self.update(entity)
263 }
264
265 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError>;
267
268 fn get_metadata(&self, id: EntityId) -> Option<Metadata>;
270
271 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError>;
273
274 fn seal(&mut self) -> Result<(), SegmentError>;
276
277 fn should_seal(&self, config: &SegmentConfig) -> bool;
279
280 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
282
283 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
285
286 fn filter_metadata(
288 &self,
289 filters: &[(String, super::metadata::MetadataFilter)],
290 ) -> Vec<EntityId>;
291}
292
293#[derive(Debug, Clone)]
300pub struct ColZone {
301 pub min: Value,
302 pub max: Value,
303 min_key: Option<CanonicalKey>,
304 max_key: Option<CanonicalKey>,
305}
306
307impl ColZone {
308 fn payload_bytes(&self) -> usize {
309 super::memory_size::value_heap_bytes(&self.min)
310 .saturating_add(super::memory_size::value_heap_bytes(&self.max))
311 .saturating_add(
312 self.min_key
313 .as_ref()
314 .map_or(0, super::memory_size::canonical_heap_bytes),
315 )
316 .saturating_add(
317 self.max_key
318 .as_ref()
319 .map_or(0, super::memory_size::canonical_heap_bytes),
320 )
321 }
322
323 fn new(v: Value) -> Self {
324 Self {
325 min_key: value_to_canonical_key(&v),
326 max_key: value_to_canonical_key(&v),
327 min: v.clone(),
328 max: v,
329 }
330 }
331
332 fn with_bounds(min: Value, max: Value) -> Self {
333 Self {
334 min_key: value_to_canonical_key(&min),
335 max_key: value_to_canonical_key(&max),
336 min,
337 max,
338 }
339 }
340
341 fn update(&mut self, v: &Value) {
342 if compare_zone_values(v, None, &self.min, self.min_key.as_ref())
343 .map(|o| o == std::cmp::Ordering::Less)
344 .unwrap_or(false)
345 {
346 self.min = v.clone();
347 self.min_key = value_to_canonical_key(v);
348 }
349 if compare_zone_values(v, None, &self.max, self.max_key.as_ref())
350 .map(|o| o == std::cmp::Ordering::Greater)
351 .unwrap_or(false)
352 {
353 self.max = v.clone();
354 self.max_key = value_to_canonical_key(v);
355 }
356 }
357}
358
359#[derive(Debug, Clone, Default)]
360pub struct MultiColZone {
361 pub intervals: Vec<ColZone>,
362}
363
364impl MultiColZone {
365 fn can_skip(&self, pred: &ZoneColPred<'_>) -> bool {
366 !self.intervals.is_empty() && self.intervals.iter().all(|zone| pred.can_skip(zone))
367 }
368}
369
370fn compare_zone_values(
371 left: &Value,
372 left_key: Option<&CanonicalKey>,
373 right: &Value,
374 right_key: Option<&CanonicalKey>,
375) -> Option<std::cmp::Ordering> {
376 partial_compare_values(left, right).or_else(|| {
377 let left_key = left_key.cloned().or_else(|| value_to_canonical_key(left))?;
378 let right_key = right_key
379 .cloned()
380 .or_else(|| value_to_canonical_key(right))?;
381 (left_key.family() == right_key.family()).then(|| left_key.cmp(&right_key))
382 })
383}
384
385#[derive(Debug, Clone, Copy, PartialEq, Eq)]
387pub enum ZoneColPredKind {
388 Eq,
389 Gt,
390 Gte,
391 Lt,
392 Lte,
393}
394
395#[derive(Debug, Clone)]
397pub enum ZoneColPred<'a> {
398 Eq(&'a Value),
399 Gt(&'a Value),
400 Gte(&'a Value),
401 Lt(&'a Value),
402 Lte(&'a Value),
403}
404
405impl<'a> ZoneColPred<'a> {
406 pub fn can_skip(&self, zone: &ColZone) -> bool {
408 match self {
409 ZoneColPred::Eq(val) => {
411 compare_zone_values(val, None, &zone.min, zone.min_key.as_ref())
412 .map(|o| o == std::cmp::Ordering::Less)
413 .unwrap_or(false)
414 || compare_zone_values(val, None, &zone.max, zone.max_key.as_ref())
415 .map(|o| o == std::cmp::Ordering::Greater)
416 .unwrap_or(false)
417 }
418 ZoneColPred::Gt(val) => {
420 compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
421 .map(|o| o != std::cmp::Ordering::Greater)
422 .unwrap_or(false)
423 }
424 ZoneColPred::Gte(val) => {
426 compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
427 .map(|o| o == std::cmp::Ordering::Less)
428 .unwrap_or(false)
429 }
430 ZoneColPred::Lt(val) => {
432 compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
433 .map(|o| o != std::cmp::Ordering::Less)
434 .unwrap_or(false)
435 }
436 ZoneColPred::Lte(val) => {
438 compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
439 .map(|o| o == std::cmp::Ordering::Greater)
440 .unwrap_or(false)
441 }
442 }
443 }
444}
445
446#[derive(Clone, Copy)]
448pub(crate) enum GraphEntityKind {
449 Node,
450 Edge,
451}
452
453impl GraphEntityKind {
454 fn mask(self) -> u8 {
455 match self {
456 Self::Node => 1,
457 Self::Edge => 2,
458 }
459 }
460
461 pub(crate) fn matches(self, kind: &EntityKind) -> bool {
462 matches!(
463 (self, kind),
464 (Self::Node, EntityKind::GraphNode(_)) | (Self::Edge, EntityKind::GraphEdge(_))
465 )
466 }
467}
468
469pub(crate) const SCAN_BATCH_SIZE: usize = 256;
471
472pub(crate) struct SegmentScanCursor {
474 flat_end: usize,
475 hash_end: usize,
476 deleted_end: usize,
477 position: usize,
478}
479
480impl SegmentScanCursor {
481 pub(crate) fn remaining(&self) -> usize {
482 self.flat_end + self.hash_end - self.position
483 }
484
485 pub(crate) fn finished(&self) -> bool {
486 self.remaining() == 0
487 }
488}
489
490pub(crate) struct GraphIndexCursor {
493 key_index: usize,
494 after: Option<EntityId>,
495 upper: EntityId,
496}
497
498impl GraphIndexCursor {
499 pub(crate) 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 memory_bytes: AtomicU64,
571 dead_entity_bytes: u64,
574 dead_resident_count: usize,
576
577 pub(crate) published_flat_len: AtomicUsize,
584}
585
586impl GrowingSegment {
587 pub(crate) fn may_contain_graph_kind(&self, kind: GraphEntityKind) -> bool {
588 self.graph_kind_mask & kind.mask() != 0
589 }
590
591 fn observe_graph_kind(&mut self, kind: &EntityKind) {
592 self.graph_kind_mask |= match kind {
593 EntityKind::GraphNode(_) => GraphEntityKind::Node.mask(),
594 EntityKind::GraphEdge(_) => GraphEntityKind::Edge.mask(),
595 _ => 0,
596 };
597 }
598
599 pub(crate) fn scan_cursor(&self) -> SegmentScanCursor {
600 SegmentScanCursor {
601 flat_end: self.flat_entities.len(),
602 hash_end: self.hash_scan_ids.len(),
603 deleted_end: self.deleted.len(),
604 position: 0,
605 }
606 }
607
608 pub(crate) fn scan_batch(
612 &self,
613 cursor: &mut SegmentScanCursor,
614 before_work: &mut impl FnMut() -> bool,
615 visit: &mut impl FnMut(EntityId),
616 ) -> bool {
617 for _ in 0..SCAN_BATCH_SIZE {
618 if cursor.finished() {
619 break;
620 }
621 if !before_work() {
622 return false;
623 }
624 let id = if cursor.position < cursor.flat_end {
625 self.flat_entities[cursor.position].id
626 } else {
627 self.hash_scan_ids[cursor.position - cursor.flat_end]
628 };
629 cursor.position += 1;
630 let eligible = match self.deleted.get(&id) {
631 Some(&ordinal) => ordinal >= cursor.deleted_end,
632 None => self.get(id).is_some(),
635 };
636 if eligible {
637 visit(id);
638 }
639 }
640 true
641 }
642
643 #[inline]
646 pub fn for_each_fast<F>(&self, mut f: F) -> bool
647 where
648 F: FnMut(&UnifiedEntity) -> bool,
649 {
650 if self.use_flat {
651 if self.deleted.is_empty() {
653 for entity in &self.flat_entities {
654 if !f(entity) {
655 return false;
656 }
657 }
658 for entity in self.entities.values() {
661 if !f(entity) {
662 return false;
663 }
664 }
665 } else {
666 for entity in &self.flat_entities {
667 if self.deleted.contains_key(&entity.id) {
668 continue;
669 }
670 if !f(entity) {
671 return false;
672 }
673 }
674 for entity in self.entities.values() {
675 if self.deleted.contains_key(&entity.id) {
676 continue;
677 }
678 if !f(entity) {
679 return false;
680 }
681 }
682 }
683 } else {
684 if self.deleted.is_empty() {
686 for entity in self.entities.values() {
687 if !f(entity) {
688 return false;
689 }
690 }
691 } else {
692 for entity in self.entities.values() {
693 if self.deleted.contains_key(&entity.id) {
694 continue;
695 }
696 if !f(entity) {
697 return false;
698 }
699 }
700 }
701 }
702 true
703 }
704
705 pub fn new(id: SegmentId, collection: impl Into<String>) -> Self {
707 let now = current_unix_secs();
708
709 Self {
710 id,
711 collection: collection.into(),
712 state: SegmentState::Growing,
713 created_at: now,
714 last_write_at: now,
715 entities: HashMap::new(),
716 hash_scan_ids: Vec::new(),
717 flat_entities: Vec::new(),
718 base_entity_id: 0,
719 use_flat: false,
720 deleted: HashMap::new(),
721 metadata: MetadataStorage::new(),
722 pk_index: BTreeMap::new(),
723 unique_key_lookup: parking_lot::RwLock::new(Vec::new()),
724 unique_key_lookup_bytes: AtomicU64::new(0),
725 key_lookup: Some(SegmentKeyLookup::default()),
726 kind_index: HashMap::new(),
727 graph_read_index: parking_lot::RwLock::new(Some(SegmentGraphIndex::default())),
728 graph_read_index_bytes: AtomicU64::new(0),
729 graph_kind_mask: 0,
730 cross_ref_forward: HashMap::new(),
731 cross_ref_reverse: HashMap::new(),
732 col_zones: HashMap::new(),
733 zone_payload_bytes: 0,
734 sealed_col_zones: HashMap::new(),
735 sequence: AtomicU64::new(0),
736 checkpoint_revision: 0,
737 memory_bytes: AtomicU64::new(0),
738 dead_entity_bytes: 0,
739 dead_resident_count: 0,
740 published_flat_len: AtomicUsize::new(0),
741 }
742 }
743
744 fn next_sequence(&self) -> u64 {
746 self.sequence.fetch_add(1, Ordering::SeqCst)
747 }
748
749 pub(crate) fn checkpoint_revision(&self) -> u64 {
750 self.checkpoint_revision
751 }
752
753 fn mark_checkpoint_dirty(&mut self) {
754 self.checkpoint_revision = self
755 .checkpoint_revision
756 .checked_add(1)
757 .expect("segment mutation counter must not overflow");
758 }
759
760 fn has_live_entity(&self, id: EntityId) -> bool {
761 if self.deleted.contains_key(&id) {
762 return false;
763 }
764 if self.use_flat {
765 let raw = id.raw();
766 if raw >= self.base_entity_id {
767 let idx = (raw - self.base_entity_id) as usize;
768 if self
769 .flat_entities
770 .get(idx)
771 .is_some_and(|entity| entity.id == id)
772 {
773 return true;
774 }
775 }
776 self.entities.contains_key(&id)
779 } else {
780 self.entities.contains_key(&id)
781 }
782 }
783
784 fn update_existing_entity_in_place(
785 &mut self,
786 entity: &UnifiedEntity,
787 ) -> Result<UpdateIndexSnapshot, SegmentError> {
788 self.mark_checkpoint_dirty();
789 self.invalidate_unique_key_lookup();
790 if self
791 .get(entity.id)
792 .is_some_and(|old| SegmentKeyLookup::key(old) != SegmentKeyLookup::key(entity))
793 {
794 if let Some(mut index) = self.key_lookup.take() {
795 index.remove(
796 self.get(entity.id)
797 .expect("invariant: existing replacement row"),
798 );
799 index.insert(entity);
800 self.key_lookup = Some(index);
801 }
802 }
803 if (self.graph_kind_mask != 0
804 || matches!(
805 entity.kind,
806 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
807 ))
808 && self
809 .get(entity.id)
810 .is_some_and(|old| !SegmentGraphIndex::same_keys(old, entity))
811 {
812 self.invalidate_graph_read_index();
813 }
814 self.observe_graph_kind(&entity.kind);
817 if self.use_flat {
818 let raw = entity.id.raw();
819 if raw >= self.base_entity_id {
820 let idx = (raw - self.base_entity_id) as usize;
821 if let Some(slot) = self
822 .flat_entities
823 .get_mut(idx)
824 .filter(|slot| slot.id == entity.id)
825 {
826 let snapshot = UpdateIndexSnapshot::from_entity(slot);
827 let old_bytes = Self::estimate_entity_size(slot);
828 slot.clone_from(entity);
829 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
830 return Ok(snapshot);
831 }
832 }
833 let Some(slot) = self.entities.get_mut(&entity.id) else {
834 return Err(SegmentError::NotFound(entity.id));
835 };
836 let snapshot = UpdateIndexSnapshot::from_entity(slot);
837 let old_bytes = Self::estimate_entity_size(slot);
838 slot.clone_from(entity);
839 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
840 Ok(snapshot)
841 } else {
842 let Some(slot) = self.entities.get_mut(&entity.id) else {
843 return Err(SegmentError::NotFound(entity.id));
844 };
845 let snapshot = UpdateIndexSnapshot::from_entity(slot);
846 let old_bytes = Self::estimate_entity_size(slot);
847 slot.clone_from(entity);
848 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
849 Ok(snapshot)
850 }
851 }
852
853 fn apply_hot_update_with_metadata(
854 &mut self,
855 entity: &UnifiedEntity,
856 modified_columns: &[String],
857 metadata: Option<&Metadata>,
858 ) -> Result<(), SegmentError> {
859 let old = self.update_existing_entity_in_place(entity)?;
860 self.reindex_for_update(&old, entity, Some(modified_columns));
861 self.update_col_zones_from_entity(entity);
862 if let Some(metadata) = metadata {
863 self.metadata.set_all(entity.id, metadata);
864 }
865 Ok(())
866 }
867
868 fn apply_update_with_metadata(
869 &mut self,
870 entity: &UnifiedEntity,
871 metadata: Option<&Metadata>,
872 ) -> Result<(), SegmentError> {
873 let old = self.update_existing_entity_in_place(entity)?;
874 self.reindex_for_update(&old, entity, None);
875 self.update_col_zones_from_entity(entity);
876 if let Some(metadata) = metadata {
877 self.metadata.set_all(entity.id, metadata);
878 }
879 Ok(())
880 }
881
882 pub fn update_hot_batch_with_metadata<'a, I>(&mut self, items: I) -> Result<(), SegmentError>
883 where
884 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
885 {
886 if !self.state.is_writable() {
887 return Err(SegmentError::NotWritable);
888 }
889
890 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
891 items.into_iter().collect();
892 if items.is_empty() {
893 return Ok(());
894 }
895
896 for (entity, _, _) in &items {
897 if !self.has_live_entity(entity.id) {
898 return Err(SegmentError::NotFound(entity.id));
899 }
900 }
901
902 for (entity, modified_columns, metadata) in items {
903 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)?;
904 }
905
906 self.last_write_at = current_unix_secs();
907 Ok(())
908 }
909
910 fn record_tombstone(&mut self, id: EntityId) -> bool {
912 self.mark_checkpoint_dirty();
913 let ordinal = self.deleted.len();
914 match self.deleted.entry(id) {
915 std::collections::hash_map::Entry::Vacant(entry) => {
916 entry.insert(ordinal);
917 true
918 }
919 std::collections::hash_map::Entry::Occupied(_) => false,
920 }
921 }
922
923 fn tombstone_flat_entity(&mut self, idx: usize, id: EntityId) -> bool {
929 self.invalidate_unique_key_lookup();
930 if !self.record_tombstone(id) {
931 return false;
932 }
933 if let Some(index) = self.key_lookup.as_mut() {
934 index.remove(&self.flat_entities[idx]);
935 }
936 if let Some(index) = self.graph_read_index.get_mut() {
937 index.remove(&self.flat_entities[idx]);
938 self.graph_read_index_bytes
939 .store(index.memory_bytes(), Ordering::Relaxed);
940 }
941 let size = Self::estimate_entity_size(&self.flat_entities[idx]) as u64;
942 self.dead_entity_bytes += size;
943 self.dead_resident_count += 1;
944 self.metadata.remove_all(id);
945 true
946 }
947
948 fn remove_hashmap_entity(&mut self, id: EntityId) -> bool {
953 self.invalidate_unique_key_lookup();
954 let Some(entity) = self.entities.remove(&id) else {
955 return false;
956 };
957 self.release_memory(Self::estimate_entity_size(&entity));
958 self.unindex_entity(&entity);
959 self.metadata.remove_all(id);
960 self.record_tombstone(id);
961 true
962 }
963
964 pub fn delete_batch(&mut self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
965 if !self.state.is_writable() {
966 return Err(SegmentError::NotWritable);
967 }
968 if ids.is_empty() {
969 return Ok(Vec::new());
970 }
971
972 let mut deleted_ids = Vec::with_capacity(ids.len());
973
974 if self.use_flat {
975 for &id in ids {
976 let raw = id.raw();
977 if raw < self.base_entity_id {
978 if self.remove_hashmap_entity(id) {
979 deleted_ids.push(id);
980 }
981 continue;
982 }
983 let idx = (raw - self.base_entity_id) as usize;
984 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
985 if self.tombstone_flat_entity(idx, id) {
986 deleted_ids.push(id);
987 }
988 } else if self.remove_hashmap_entity(id) {
989 deleted_ids.push(id);
990 }
991 }
992 } else {
993 for &id in ids {
994 if self.remove_hashmap_entity(id) {
995 deleted_ids.push(id);
996 }
997 }
998 }
999
1000 if !deleted_ids.is_empty() {
1001 self.last_write_at = current_unix_secs();
1002 }
1003
1004 Ok(deleted_ids)
1005 }
1006
1007 fn add_memory(&self, bytes: usize) {
1009 self.memory_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
1010 }
1011
1012 fn replace_entity_memory(&self, previous: usize, current: usize) {
1013 if current >= previous {
1014 self.add_memory(current - previous);
1015 } else {
1016 self.release_memory(previous - current);
1017 }
1018 }
1019
1020 pub fn memory_bytes(&self) -> u64 {
1025 self.memory_bytes.load(Ordering::Relaxed)
1026 + self.zone_payload_bytes as u64
1027 + self.unique_key_lookup_bytes.load(Ordering::Relaxed)
1028 + self
1029 .key_lookup
1030 .as_ref()
1031 .map_or(0, SegmentKeyLookup::memory_bytes)
1032 + self.graph_read_index_bytes.load(Ordering::Relaxed)
1033 + (self.hash_scan_ids.capacity() * std::mem::size_of::<EntityId>()) as u64
1034 }
1035
1036 fn release_memory(&self, bytes: usize) {
1039 let mut current = self.memory_bytes.load(Ordering::Relaxed);
1040 loop {
1041 let next = current.saturating_sub(bytes as u64);
1042 match self.memory_bytes.compare_exchange_weak(
1043 current,
1044 next,
1045 Ordering::Relaxed,
1046 Ordering::Relaxed,
1047 ) {
1048 Ok(_) => break,
1049 Err(observed) => current = observed,
1050 }
1051 }
1052 }
1053
1054 pub(crate) fn resident_entity_count(&self) -> usize {
1056 self.flat_entities.len() + self.entities.len()
1057 }
1058
1059 pub(crate) fn live_entity_count(&self) -> usize {
1061 self.resident_entity_count()
1062 .saturating_sub(self.dead_resident_count)
1063 }
1064
1065 pub(crate) fn tombstone_count(&self) -> usize {
1068 self.deleted.len()
1069 }
1070
1071 pub(crate) fn resident_bytes(&self) -> u64 {
1074 self.memory_bytes() + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1075 }
1076
1077 pub(crate) fn reclaimable_bytes(&self) -> u64 {
1080 self.dead_entity_bytes + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1081 }
1082
1083 pub(crate) fn live_entity_ids(&self) -> Vec<EntityId> {
1089 let mut ids = Vec::with_capacity(self.live_entity_count());
1090 self.for_each_fast(|entity| {
1091 ids.push(entity.id);
1092 true
1093 });
1094 ids
1095 }
1096
1097 pub(crate) fn adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
1105 let id = entity.id;
1106
1107 let next = entity.sequence_id.saturating_add(1);
1110 if self.sequence.load(Ordering::Relaxed) < next {
1111 self.sequence.store(next, Ordering::Relaxed);
1112 }
1113
1114 self.add_memory(Self::estimate_entity_size(&entity));
1115 self.index_entity(&entity);
1116 self.update_col_zones_from_entity(&entity);
1117 self.entities.insert(id, entity);
1118
1119 if let Some(metadata) = metadata {
1120 if !metadata.is_empty() {
1121 self.metadata.set_all(id, &metadata);
1122 }
1123 }
1124 }
1125
1126 pub(crate) fn evict_entity(&mut self, id: EntityId) -> bool {
1132 self.mark_checkpoint_dirty();
1133 self.invalidate_unique_key_lookup();
1134 let Some(entity) = self.entities.remove(&id) else {
1135 return false;
1136 };
1137 self.release_memory(Self::estimate_entity_size(&entity));
1138 self.unindex_entity(&entity);
1139 for cross_ref in entity.cross_refs() {
1140 if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1141 reverse.retain(|(source, _)| *source != id);
1142 }
1143 }
1144 self.metadata.remove_all(id);
1145 true
1146 }
1147
1148 fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
1150 super::memory_size::entity_bytes(entity)
1151 }
1152
1153 fn update_column_zone(&mut self, column: &str, value: &Value) {
1161 if matches!(value, Value::Null) {
1162 return;
1163 }
1164 match self.col_zones.entry(column.to_string()) {
1165 std::collections::hash_map::Entry::Occupied(mut entry) => {
1166 let before = entry.get().payload_bytes();
1167 entry.get_mut().update(value);
1168 self.zone_payload_bytes = self
1169 .zone_payload_bytes
1170 .saturating_sub(before)
1171 .saturating_add(entry.get().payload_bytes());
1172 }
1173 std::collections::hash_map::Entry::Vacant(entry) => {
1174 let zone = ColZone::new(value.clone());
1175 self.zone_payload_bytes =
1176 self.zone_payload_bytes.saturating_add(zone.payload_bytes());
1177 entry.insert(zone);
1178 }
1179 }
1180 }
1181
1182 fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
1183 if let EntityData::Row(row) = &entity.data {
1184 if let Some(named) = &row.named {
1185 for (column, value) in named {
1186 self.update_column_zone(column, value);
1187 }
1188 } else if let Some(schema) = &row.schema {
1189 for (column, value) in schema.iter().zip(&row.columns) {
1190 self.update_column_zone(column, value);
1191 }
1192 }
1193 }
1194 }
1195
1196 fn rebuild_sealed_col_zones(&mut self) {
1197 let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
1198 let mut family_by_col: HashMap<String, reddb_types::CanonicalKeyFamily> = HashMap::new();
1199 let mut mixed_family_cols = HashSet::new();
1200 let mut unsupported_cols = HashSet::new();
1201
1202 let mut observe_row = |row: &super::entity::RowData| {
1203 for (col, value) in row.iter_fields() {
1204 if matches!(value, Value::Null) {
1205 continue;
1206 }
1207 let Some(key) = value_to_canonical_key(value) else {
1208 unsupported_cols.insert(col.to_string());
1209 continue;
1210 };
1211 match family_by_col.get(col).copied() {
1212 Some(existing) if existing != key.family() => {
1213 mixed_family_cols.insert(col.to_string());
1214 }
1215 None => {
1216 family_by_col.insert(col.to_string(), key.family());
1217 }
1218 _ => {}
1219 }
1220 values_by_col
1221 .entry(col.to_string())
1222 .or_default()
1223 .push((key, value.clone()));
1224 }
1225 };
1226
1227 if self.use_flat {
1228 for entity in &self.flat_entities {
1229 if self.deleted.contains_key(&entity.id) {
1230 continue;
1231 }
1232 if let EntityData::Row(row) = &entity.data {
1233 observe_row(row);
1234 }
1235 }
1236 } else {
1237 for entity in self.entities.values() {
1238 if self.deleted.contains_key(&entity.id) {
1239 continue;
1240 }
1241 if let EntityData::Row(row) = &entity.data {
1242 observe_row(row);
1243 }
1244 }
1245 }
1246
1247 let mut sealed_col_zones = HashMap::new();
1248 for (col, mut entries) in values_by_col {
1249 if mixed_family_cols.contains(&col)
1250 || unsupported_cols.contains(&col)
1251 || entries.is_empty()
1252 {
1253 continue;
1254 }
1255 entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
1256 entries.dedup_by(|left, right| left.0 == right.0);
1257
1258 let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
1259 if intervals.len() > 1 {
1260 sealed_col_zones.insert(col, MultiColZone { intervals });
1261 }
1262 }
1263
1264 let previous = self
1265 .sealed_col_zones
1266 .values()
1267 .flat_map(|zones| &zones.intervals)
1268 .map(ColZone::payload_bytes)
1269 .fold(0, usize::saturating_add);
1270 let current = sealed_col_zones
1271 .values()
1272 .flat_map(|zones| &zones.intervals)
1273 .map(ColZone::payload_bytes)
1274 .fold(0, usize::saturating_add);
1275 self.zone_payload_bytes = self
1276 .zone_payload_bytes
1277 .saturating_sub(previous)
1278 .saturating_add(current);
1279 self.sealed_col_zones = sealed_col_zones;
1280 }
1281
1282 pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
1286 if preds.is_empty() {
1287 return false;
1288 }
1289 for (col, pred) in preds {
1290 if let Some(zone) = self.sealed_col_zones.get(*col) {
1291 if zone.can_skip(pred) {
1292 return true;
1293 }
1294 continue;
1295 }
1296 if let Some(zone) = self.col_zones.get(*col) {
1297 if pred.can_skip(zone) {
1298 return true; }
1300 }
1301 }
1302 false
1303 }
1304
1305 fn index_entity(&mut self, entity: &UnifiedEntity) {
1307 self.mark_checkpoint_dirty();
1308 if let Some(index) = self.key_lookup.as_mut() {
1309 index.insert(entity);
1310 }
1311 self.hash_scan_ids.push(entity.id);
1312 self.observe_graph_kind(&entity.kind);
1313 if let Some(index) = self.graph_read_index.get_mut() {
1314 index.insert(entity);
1315 self.graph_read_index_bytes
1316 .store(index.memory_bytes(), Ordering::Relaxed);
1317 }
1318 let indexes = self.unique_key_lookup.get_mut();
1319 for index in indexes.iter_mut() {
1320 index.insert(entity);
1321 }
1322 self.unique_key_lookup_bytes.store(
1323 indexes
1324 .iter()
1325 .map(|index| index.memory_bytes() as u64)
1326 .sum(),
1327 Ordering::Relaxed,
1328 );
1329 let kind_key = entity.kind.storage_type().to_string();
1331 self.kind_index
1332 .entry(kind_key)
1333 .or_default()
1334 .insert(entity.id);
1335
1336 if let EntityData::Row(row) = &entity.data {
1338 if let Some(first_col) = row.columns.first() {
1339 let pk_str = format!("{:?}", first_col);
1340 self.pk_index
1341 .insert((entity.kind.collection().to_string(), pk_str), entity.id);
1342 }
1343 }
1344
1345 for cross_ref in entity.cross_refs() {
1347 self.cross_ref_forward
1348 .entry(cross_ref.source)
1349 .or_default()
1350 .push((cross_ref.target, cross_ref.ref_type));
1351
1352 self.cross_ref_reverse
1353 .entry(cross_ref.target)
1354 .or_default()
1355 .push((cross_ref.source, cross_ref.ref_type));
1356 }
1357 }
1358
1359 fn invalidate_graph_read_index(&mut self) {
1360 *self.graph_read_index.get_mut() = None;
1361 self.graph_read_index_bytes.store(0, Ordering::Relaxed);
1362 }
1363
1364 #[cfg(test)]
1367 pub(crate) fn visit_graph_index_candidates(
1368 &self,
1369 kind: GraphEntityKind,
1370 keys: &[EntityId],
1371 before_work: &mut impl FnMut() -> bool,
1372 visit: &mut impl FnMut(&UnifiedEntity) -> bool,
1373 ) -> bool {
1374 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1375 return false;
1376 };
1377 let index = cached.as_ref().expect("complete graph index");
1378 for key in keys {
1379 if !before_work() {
1380 return false;
1381 }
1382 for id in index.candidates(kind, *key) {
1383 if !before_work() {
1384 return false;
1385 }
1386 if let Some(entity) = self.get(id) {
1387 if !visit(entity) {
1388 return false;
1389 }
1390 }
1391 }
1392 }
1393 true
1394 }
1395
1396 fn graph_read_index_for_query(
1397 &self,
1398 before_work: &mut impl FnMut() -> bool,
1399 ) -> Option<parking_lot::RwLockReadGuard<'_, Option<SegmentGraphIndex>>> {
1400 let mut cached = self.graph_read_index.read();
1401 if cached.is_none() {
1402 drop(cached);
1403 let mut writer = self.graph_read_index.write();
1404 if writer.is_none() {
1405 let mut index = SegmentGraphIndex::default();
1406 if !self.for_each_fast(|entity| {
1409 if !before_work() {
1410 return false;
1411 }
1412 index.insert(entity);
1413 true
1414 }) {
1415 return None;
1416 }
1417 self.graph_read_index_bytes
1418 .store(index.memory_bytes(), Ordering::Relaxed);
1419 *writer = Some(index);
1420 }
1421 cached = parking_lot::RwLockWriteGuard::downgrade(writer);
1422 }
1423 Some(cached)
1424 }
1425
1426 pub(crate) fn graph_index_cursor(
1427 &self,
1428 kind: GraphEntityKind,
1429 keys: &[EntityId],
1430 before_work: &mut impl FnMut() -> bool,
1431 ) -> Option<GraphIndexCursor> {
1432 let cached = self.graph_read_index_for_query(before_work)?;
1433 let index = cached.as_ref().expect("complete graph index");
1434 let mut upper = EntityId::new(0);
1435 for key in keys {
1436 if !before_work() {
1437 return None;
1438 }
1439 if let Some(last) = index.candidates(kind, *key).next_back() {
1440 upper = upper.max(last);
1441 }
1442 }
1443 Some(GraphIndexCursor {
1444 key_index: 0,
1445 after: None,
1446 upper,
1447 })
1448 }
1449
1450 pub(crate) fn graph_index_batch(
1453 &self,
1454 kind: GraphEntityKind,
1455 keys: &[EntityId],
1456 cursor: &mut GraphIndexCursor,
1457 before_work: &mut impl FnMut() -> bool,
1458 visit: &mut impl FnMut(EntityId),
1459 ) -> bool {
1460 self.graph_index_batch_with_limit(kind, keys, cursor, before_work, visit, SCAN_BATCH_SIZE)
1461 }
1462
1463 pub(crate) fn graph_index_batch_with_limit(
1464 &self,
1465 kind: GraphEntityKind,
1466 keys: &[EntityId],
1467 cursor: &mut GraphIndexCursor,
1468 before_work: &mut impl FnMut() -> bool,
1469 visit: &mut impl FnMut(EntityId),
1470 limit: usize,
1471 ) -> bool {
1472 assert!(
1473 limit > 0 && limit <= SCAN_BATCH_SIZE,
1474 "bounded graph index batch"
1475 );
1476 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1477 return false;
1478 };
1479 let index = cached.as_ref().expect("complete graph index");
1480 let mut count = 0;
1481 while !cursor.finished(keys) && count < limit {
1482 let key = keys[cursor.key_index];
1483 if !before_work() {
1484 return false;
1485 }
1486 for id in index
1487 .candidates_after(kind, key, cursor.after, cursor.upper)
1488 .take(limit - count)
1489 {
1490 if !before_work() {
1491 return false;
1492 }
1493 cursor.after = Some(id);
1494 count += 1;
1495 visit(id);
1496 }
1497 if count == limit && cursor.after != Some(cursor.upper) {
1498 break;
1499 }
1500 cursor.key_index += 1;
1501 cursor.after = None;
1502 }
1503 true
1504 }
1505
1506 pub(crate) fn find_unique_key_conflict(
1511 &self,
1512 columns: &[Vec<String>],
1513 key_index: usize,
1514 signatures: &[String],
1515 exclude_id: Option<EntityId>,
1516 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1517 ) -> Option<EntityId> {
1518 assert!(
1519 key_index < columns.len(),
1520 "constraint index belongs to current schema"
1521 );
1522 let cached = self.unique_key_lookup.read();
1523 if cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1524 return self.find_unique_key_candidate(
1525 &cached[key_index],
1526 signatures,
1527 exclude_id,
1528 reserves_key,
1529 );
1530 }
1531 drop(cached);
1532 let mut cached = self.unique_key_lookup.write();
1533 if !cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1534 let mut indexes: Vec<_> = columns
1535 .iter()
1536 .map(|columns| SegmentUniqueKey::new(columns))
1537 .collect();
1538 self.for_each_fast(|entity| {
1539 for index in &mut indexes {
1540 index.insert(entity);
1541 }
1542 true
1543 });
1544 self.unique_key_lookup_bytes.store(
1545 indexes
1546 .iter()
1547 .map(|index| index.memory_bytes() as u64)
1548 .sum(),
1549 Ordering::Relaxed,
1550 );
1551 *cached = indexes;
1552 }
1553 self.find_unique_key_candidate(&cached[key_index], signatures, exclude_id, reserves_key)
1554 }
1555
1556 fn find_unique_key_candidate(
1557 &self,
1558 index: &SegmentUniqueKey,
1559 signatures: &[String],
1560 exclude_id: Option<EntityId>,
1561 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1562 ) -> Option<EntityId> {
1563 index.candidates(signatures).iter().copied().find(|id| {
1564 if Some(*id) == exclude_id {
1565 return false;
1566 }
1567 let Some(
1568 entity @ UnifiedEntity {
1569 data: EntityData::Row(row),
1570 ..
1571 },
1572 ) = self.get(*id)
1573 else {
1574 return false;
1575 };
1576 index
1577 .columns
1578 .iter()
1579 .zip(signatures)
1580 .all(|(column, signature)| {
1581 row.get_field(column)
1582 .is_some_and(|value| format!("{value:?}") == *signature)
1583 })
1584 && reserves_key(entity)
1585 })
1586 }
1587
1588 fn invalidate_unique_key_lookup(&mut self) {
1589 *self.unique_key_lookup.get_mut() = Vec::new();
1590 self.unique_key_lookup_bytes.store(0, Ordering::Relaxed);
1591 }
1592
1593 pub(crate) fn find_key_entity(
1596 &self,
1597 key: &str,
1598 newest: bool,
1599 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1600 ) -> Option<&UnifiedEntity> {
1601 let index = self
1604 .key_lookup
1605 .as_ref()
1606 .expect("segment owns its key index");
1607 self.find_key_candidate(index, key, newest, visible)
1608 }
1609
1610 fn find_key_candidate(
1611 &self,
1612 index: &SegmentKeyLookup,
1613 key: &str,
1614 newest: bool,
1615 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1616 ) -> Option<&UnifiedEntity> {
1617 self.select_key_entity(
1618 index
1619 .candidates(key)
1620 .iter()
1621 .chain(index.mutated_candidates())
1622 .filter_map(|id| self.get(*id)),
1623 key,
1624 newest,
1625 visible,
1626 )
1627 }
1628
1629 fn select_key_entity<'a>(
1630 &self,
1631 entities: impl Iterator<Item = &'a UnifiedEntity>,
1632 key: &str,
1633 newest: bool,
1634 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1635 ) -> Option<&'a UnifiedEntity> {
1636 let mut best: Option<&UnifiedEntity> = None;
1637 for entity in entities {
1638 if SegmentKeyLookup::key(entity) != Some(key) || !visible(entity) {
1640 continue;
1641 }
1642 if !newest {
1643 return Some(entity);
1644 }
1645 if best.is_none_or(|current| entity.xmin >= current.xmin) {
1646 best = Some(entity);
1647 }
1648 }
1649 best
1650 }
1651
1652 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1654 if let Some(id) = entity_id_from_probe_key(key) {
1655 return self.has_live_entity(id);
1656 }
1657
1658 let Ok(pk_str) = std::str::from_utf8(key) else {
1659 return false;
1660 };
1661 let pk_key = (self.collection.clone(), pk_str.to_string());
1662 self.pk_index
1663 .get(&pk_key)
1664 .is_some_and(|id| self.has_live_entity(*id))
1665 }
1666
1667 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1669 if let Some(index) = self.key_lookup.as_mut() {
1670 index.remove(entity);
1671 }
1672 if let Some(index) = self.graph_read_index.get_mut() {
1673 index.remove(entity);
1674 self.graph_read_index_bytes
1675 .store(index.memory_bytes(), Ordering::Relaxed);
1676 }
1677 let kind_key = entity.kind.storage_type().to_string();
1679 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1680 set.remove(&entity.id);
1681 }
1682
1683 if let EntityData::Row(row) = &entity.data {
1685 if let Some(first_col) = row.columns.first() {
1686 let pk_str = format!("{:?}", first_col);
1687 self.pk_index
1688 .remove(&(entity.kind.collection().to_string(), pk_str));
1689 }
1690 }
1691
1692 self.cross_ref_forward.remove(&entity.id);
1694 }
1696
1697 fn reindex_for_update(
1707 &mut self,
1708 old: &UpdateIndexSnapshot,
1709 new: &UnifiedEntity,
1710 modified_columns: Option<&[String]>,
1711 ) {
1712 let pk_changed = match &new.data {
1717 EntityData::Row(new_row) => {
1718 if let Some(cols) = modified_columns {
1719 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1722 new_row
1723 .schema
1724 .as_deref()
1725 .and_then(|schema| schema.first().map(|name| name.as_str()))
1726 });
1727 match pk_col_name {
1728 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1729 None => old.pk_value.as_ref() != new_row.columns.first(),
1731 }
1732 } else {
1733 old.pk_value.as_ref() != new_row.columns.first()
1734 }
1735 }
1736 _ => false,
1738 };
1739
1740 if pk_changed {
1741 if let Some((collection, pk_str)) = &old.pk_index_key {
1743 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1744 }
1745 if let EntityData::Row(row) = &new.data {
1747 if let Some(first_col) = row.columns.first() {
1748 let pk_str = format!("{:?}", first_col);
1749 self.pk_index
1750 .insert((new.kind.collection().to_string(), pk_str), new.id);
1751 }
1752 }
1753 }
1754
1755 let new_refs = new.cross_refs();
1757 if old.cross_refs.as_slice() != new_refs {
1758 self.cross_ref_forward.remove(&new.id);
1760 for cross_ref in &old.cross_refs {
1762 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1763 rev.retain(|(src, _)| *src != new.id);
1764 }
1765 }
1766 for cross_ref in new_refs {
1768 self.cross_ref_forward
1769 .entry(cross_ref.source)
1770 .or_default()
1771 .push((cross_ref.target, cross_ref.ref_type));
1772 self.cross_ref_reverse
1773 .entry(cross_ref.target)
1774 .or_default()
1775 .push((cross_ref.source, cross_ref.ref_type));
1776 }
1777 }
1778 }
1779
1780 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1782 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1783 }
1784
1785 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1787 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1788 }
1789
1790 pub fn age_secs(&self) -> u64 {
1792 let now = current_unix_secs();
1793 now.saturating_sub(self.created_at)
1794 }
1795
1796 pub fn idle_secs(&self) -> u64 {
1798 let now = current_unix_secs();
1799 now.saturating_sub(self.last_write_at)
1800 }
1801
1802 pub fn bulk_insert(
1811 &mut self,
1812 entities: Vec<UnifiedEntity>,
1813 ) -> Result<Vec<EntityId>, SegmentError> {
1814 self.mark_checkpoint_dirty();
1815 if !self.state.is_writable() {
1816 return Err(SegmentError::NotWritable);
1817 }
1818
1819 if let Some(index) = self.key_lookup.as_mut() {
1820 for entity in &entities {
1821 index.insert(entity);
1822 }
1823 }
1824
1825 let indexes = self.unique_key_lookup.get_mut();
1826 for index in indexes.iter_mut() {
1827 for entity in &entities {
1828 index.insert(entity);
1829 }
1830 }
1831 self.unique_key_lookup_bytes.store(
1832 indexes
1833 .iter()
1834 .map(|index| index.memory_bytes() as u64)
1835 .sum(),
1836 Ordering::Relaxed,
1837 );
1838
1839 let n = entities.len();
1840
1841 if entities.is_empty() {
1842 return Ok(Vec::new());
1843 }
1844
1845 for run in
1848 entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
1849 {
1850 let kind = run[0].kind.storage_type();
1851 self.observe_graph_kind(&run[0].kind);
1852 if matches!(
1853 run[0].kind,
1854 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
1855 ) {
1856 if let Some(index) = self.graph_read_index.get_mut() {
1857 for entity in run {
1858 index.insert(entity);
1859 }
1860 self.graph_read_index_bytes
1861 .store(index.memory_bytes(), Ordering::Relaxed);
1862 }
1863 }
1864 let kind_set = match self.kind_index.get_mut(kind) {
1865 Some(ids) => ids,
1866 None => self.kind_index.entry(kind.to_string()).or_default(),
1867 };
1868 kind_set.reserve(run.len());
1869 kind_set.extend(run.iter().map(|entity| entity.id));
1870 }
1871
1872 let now = current_unix_secs();
1873
1874 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
1875
1876 let mut ids = Vec::with_capacity(n);
1877
1878 if self.flat_entities.is_empty() && self.entities.is_empty() {
1880 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
1882 self.use_flat = true;
1883 }
1884
1885 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
1895 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
1896 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
1900
1901 let mut batch_bytes: usize = 0;
1906
1907 if self.use_flat {
1908 self.flat_entities.reserve(n);
1909 for (i, mut entity) in entities.into_iter().enumerate() {
1910 entity.sequence_id = base_seq + i as u64;
1911 let id = entity.id;
1912 ids.push(id);
1913 batch_bytes += Self::estimate_entity_size(&entity);
1914 if let EntityData::Row(row) = &entity.data {
1915 if row.schema.is_some() && !row.columns.is_empty() {
1916 if columnar_zone_updates.is_empty() {
1917 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
1918 columnar_schema = row.schema.clone();
1919 }
1920 for (ci, val) in row.columns.iter().enumerate() {
1921 if !matches!(val, Value::Null) {
1922 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
1923 bucket.push(val.clone());
1924 }
1925 }
1926 }
1927 } else {
1928 for (col, val) in row.iter_fields() {
1929 if !matches!(val, Value::Null) {
1930 named_zone_updates.push((col.to_string(), val.clone()));
1931 }
1932 }
1933 }
1934 }
1935 let expected = self.base_entity_id + self.flat_entities.len() as u64;
1951 if id.raw() == expected {
1952 self.flat_entities.push(entity);
1953 } else {
1954 self.hash_scan_ids.push(id);
1955 self.entities.insert(id, entity);
1956 }
1957 }
1958 } else {
1959 self.entities.reserve(n);
1961 self.hash_scan_ids.reserve(n);
1962 let mut pairs = Vec::with_capacity(n);
1963 for (i, mut entity) in entities.into_iter().enumerate() {
1964 entity.sequence_id = base_seq + i as u64;
1965 let id = entity.id;
1966 ids.push(id);
1967 batch_bytes += Self::estimate_entity_size(&entity);
1968 if let EntityData::Row(row) = &entity.data {
1969 for (col, val) in row.iter_fields() {
1970 if !matches!(val, Value::Null) {
1971 named_zone_updates.push((col.to_string(), val.clone()));
1972 }
1973 }
1974 }
1975 self.hash_scan_ids.push(id);
1976 pairs.push((id, entity));
1977 }
1978 self.entities.extend(pairs);
1979 }
1980
1981 if !columnar_zone_updates.is_empty() {
1985 let schema = columnar_schema.as_ref();
1986 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
1987 if values.is_empty() {
1988 continue;
1989 }
1990 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
1991 continue;
1992 };
1993 let mut iter = values.into_iter();
1999 if let Some(first) = iter.next() {
2000 let previous = self
2001 .col_zones
2002 .get(col_name)
2003 .map_or(0, ColZone::payload_bytes);
2004 let zone = self
2005 .col_zones
2006 .entry(col_name.clone())
2007 .and_modify(|z| z.update(&first))
2008 .or_insert_with(|| ColZone::new(first));
2009 for v in iter {
2010 zone.update(&v);
2011 }
2012 self.zone_payload_bytes = self
2013 .zone_payload_bytes
2014 .saturating_sub(previous)
2015 .saturating_add(zone.payload_bytes());
2016 }
2017 }
2018 }
2019 for (col, val) in named_zone_updates {
2020 self.update_column_zone(&col, &val);
2021 }
2022
2023 self.add_memory(batch_bytes);
2024 self.last_write_at = now;
2025
2026 if self.use_flat {
2028 self.published_flat_len
2029 .store(self.flat_entities.len(), Ordering::Release);
2030 }
2031
2032 Ok(ids)
2033 }
2034
2035 pub(crate) fn force_delete(&mut self, id: EntityId) -> bool {
2038 if self.use_flat {
2039 let raw = id.raw();
2040 if raw >= self.base_entity_id {
2041 let idx = (raw - self.base_entity_id) as usize;
2042 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2043 self.tombstone_flat_entity(idx, id);
2044 return true;
2045 }
2046 }
2047 return self.remove_hashmap_entity(id);
2048 }
2049
2050 self.remove_hashmap_entity(id)
2051 }
2052
2053 pub(crate) fn force_update_with_metadata(
2056 &mut self,
2057 entity: &UnifiedEntity,
2058 modified_columns: &[String],
2059 metadata: Option<&Metadata>,
2060 ) -> Result<(), SegmentError> {
2061 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
2062 }
2063}
2064
2065impl UnifiedSegment for GrowingSegment {
2066 fn id(&self) -> SegmentId {
2067 self.id
2068 }
2069
2070 fn state(&self) -> SegmentState {
2071 self.state
2072 }
2073
2074 fn collection(&self) -> &str {
2075 &self.collection
2076 }
2077
2078 fn stats(&self) -> SegmentStats {
2079 let mut stats = SegmentStats {
2080 entity_count: self.live_entity_count(),
2081 deleted_count: self.deleted.len(),
2082 memory_bytes: self.resident_bytes() as usize,
2083 ..Default::default()
2084 };
2085
2086 for entity in self.entities.values() {
2087 match &entity.kind {
2088 EntityKind::TableRow { .. } => stats.row_count += 1,
2089 EntityKind::GraphNode(_) => stats.node_count += 1,
2090 EntityKind::GraphEdge(_) => stats.edge_count += 1,
2091 EntityKind::Vector { .. } => stats.vector_count += 1,
2092 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
2093 EntityKind::QueueMessage { .. } => stats.row_count += 1,
2094 }
2095 stats.cross_ref_count += entity.cross_refs().len();
2096 }
2097
2098 stats
2099 }
2100
2101 fn entity_count(&self) -> usize {
2102 self.live_entity_count()
2103 }
2104
2105 fn contains(&self, id: EntityId) -> bool {
2106 self.has_live_entity(id)
2107 }
2108
2109 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2110 if self.deleted.contains_key(&id) {
2111 return None;
2112 }
2113 if self.use_flat {
2114 let raw = id.raw();
2115 if raw >= self.base_entity_id {
2116 let idx = (raw - self.base_entity_id) as usize;
2117 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
2118 return Some(entity);
2119 }
2120 }
2121 self.entities.get(&id)
2127 } else {
2128 self.entities.get(&id)
2129 }
2130 }
2131
2132 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
2133 self.mark_checkpoint_dirty();
2134 self.invalidate_unique_key_lookup();
2135 self.invalidate_graph_read_index();
2136 if self.deleted.contains_key(&id) || !self.state.is_writable() {
2137 return None;
2138 }
2139 if let Some(mut index) = self.key_lookup.take() {
2140 if let Some(entity) = self.get(id) {
2141 index.mark_mutated(entity);
2142 }
2143 self.key_lookup = Some(index);
2144 }
2145 self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
2148 if self.use_flat {
2149 let raw = id.raw();
2150 if raw >= self.base_entity_id {
2151 let idx = (raw - self.base_entity_id) as usize;
2152 if self
2153 .flat_entities
2154 .get(idx)
2155 .map(|e| e.id == id)
2156 .unwrap_or(false)
2157 {
2158 return self.flat_entities.get_mut(idx);
2159 }
2160 }
2161 self.entities.get_mut(&id)
2164 } else {
2165 self.entities.get_mut(&id)
2166 }
2167 }
2168
2169 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2170 if !self.state.is_writable() {
2171 return Err(SegmentError::NotWritable);
2172 }
2173
2174 if self.entities.contains_key(&entity.id) {
2175 return Err(SegmentError::AlreadyExists(entity.id));
2176 }
2177
2178 entity.sequence_id = self.next_sequence();
2180
2181 let size = Self::estimate_entity_size(&entity);
2183 self.add_memory(size);
2184
2185 self.index_entity(&entity);
2187
2188 self.update_col_zones_from_entity(&entity);
2190
2191 let id = entity.id;
2193 self.entities.insert(id, entity);
2194
2195 self.last_write_at = current_unix_secs();
2197
2198 Ok(id)
2199 }
2200
2201 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2202 if !self.state.is_writable() {
2203 return Err(SegmentError::NotWritable);
2204 }
2205
2206 self.apply_update_with_metadata(&entity, None)?;
2207 self.last_write_at = current_unix_secs();
2208
2209 Ok(())
2210 }
2211
2212 fn update_hot(
2213 &mut self,
2214 entity: UnifiedEntity,
2215 modified_columns: &[String],
2216 ) -> Result<(), SegmentError> {
2217 if !self.state.is_writable() {
2218 return Err(SegmentError::NotWritable);
2219 }
2220
2221 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
2222 self.last_write_at = current_unix_secs();
2223 Ok(())
2224 }
2225
2226 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2227 if !self.state.is_writable() {
2228 return Err(SegmentError::NotWritable);
2229 }
2230
2231 if self.use_flat {
2233 let raw = id.raw();
2234 if raw >= self.base_entity_id {
2235 let idx = (raw - self.base_entity_id) as usize;
2236 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2237 self.tombstone_flat_entity(idx, id);
2238 return Ok(true);
2239 }
2240 }
2241 return Ok(self.remove_hashmap_entity(id));
2242 }
2243
2244 Ok(self.remove_hashmap_entity(id))
2245 }
2246
2247 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2248 if !self.has_live_entity(id) {
2249 return None;
2250 }
2251 Some(self.metadata.get_all(id))
2252 }
2253
2254 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2255 self.mark_checkpoint_dirty();
2256 if !self.state.is_writable() {
2257 return Err(SegmentError::NotWritable);
2258 }
2259
2260 if !self.has_live_entity(id) {
2261 return Err(SegmentError::NotFound(id));
2262 }
2263
2264 self.metadata.set_all(id, &metadata);
2265 Ok(())
2266 }
2267
2268 fn seal(&mut self) -> Result<(), SegmentError> {
2269 if self.state != SegmentState::Growing {
2270 return Err(SegmentError::InvalidState(self.state));
2271 }
2272
2273 self.state = SegmentState::Sealing;
2274
2275 self.rebuild_sealed_col_zones();
2280
2281 self.state = SegmentState::Sealed;
2282 Ok(())
2283 }
2284
2285 fn should_seal(&self, config: &SegmentConfig) -> bool {
2286 if self.entities.len() >= config.max_entities {
2288 return true;
2289 }
2290
2291 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
2293 return true;
2294 }
2295
2296 if self.age_secs() >= config.max_age_secs {
2298 return true;
2299 }
2300
2301 false
2302 }
2303
2304 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2305 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2306 Box::new(self.flat_entities.iter().chain(self.entities.values()))
2310 } else {
2311 Box::new(self.entities.values())
2312 };
2313 if self.deleted.is_empty() {
2314 base
2315 } else {
2316 Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
2317 }
2318 }
2319
2320 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2321 let ids = self.kind_index.get(kind_filter);
2324 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2327 Box::new(self.flat_entities.iter())
2328 } else {
2329 Box::new(std::iter::empty())
2330 };
2331 Box::new(flat.chain(self.entities.values()).filter(move |e| {
2332 if self.deleted.contains_key(&e.id) {
2333 return false;
2334 }
2335 if let Some(ids) = ids {
2336 ids.contains(&e.id)
2337 } else {
2338 false
2339 }
2340 }))
2341 }
2342
2343 fn filter_metadata(
2344 &self,
2345 filters: &[(String, super::metadata::MetadataFilter)],
2346 ) -> Vec<EntityId> {
2347 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
2350 Box::new(self.flat_entities.iter().map(|e| e.id))
2351 } else {
2352 Box::new(std::iter::empty())
2353 };
2354 flat_ids
2355 .chain(self.entities.keys().copied())
2356 .filter(|id| {
2357 if self.deleted.contains_key(id) {
2358 return false;
2359 }
2360 let metadata = self.metadata.get_all(*id);
2361 metadata.matches_all(filters)
2362 })
2363 .collect()
2364 }
2365}
2366
2367fn build_minmax_multi_intervals(
2368 entries: &[(CanonicalKey, Value)],
2369 max_intervals: usize,
2370) -> Vec<ColZone> {
2371 if entries.is_empty() {
2372 return Vec::new();
2373 }
2374 if entries.len() == 1 || max_intervals <= 1 {
2375 return vec![ColZone::with_bounds(
2376 entries[0].1.clone(),
2377 entries[entries.len() - 1].1.clone(),
2378 )];
2379 }
2380
2381 let mut split_points = if entries.len() <= max_intervals {
2382 (1..entries.len()).collect::<Vec<_>>()
2383 } else {
2384 let target_splits = max_intervals - 1;
2385 let mut selected = select_gap_split_points(entries, target_splits);
2386 if selected.len() < target_splits {
2387 for bucket in 1..max_intervals {
2388 let idx = bucket * entries.len() / max_intervals;
2389 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
2390 continue;
2391 }
2392 selected.push(idx);
2393 if selected.len() >= target_splits {
2394 break;
2395 }
2396 }
2397 }
2398 selected.sort_unstable();
2399 selected.dedup();
2400 selected
2401 };
2402
2403 split_points.push(entries.len());
2404
2405 let mut out = Vec::with_capacity(split_points.len());
2406 let mut start = 0usize;
2407 for end in split_points {
2408 if end <= start {
2409 continue;
2410 }
2411 out.push(ColZone::with_bounds(
2412 entries[start].1.clone(),
2413 entries[end - 1].1.clone(),
2414 ));
2415 start = end;
2416 }
2417
2418 if out.is_empty() {
2419 out.push(ColZone::with_bounds(
2420 entries[0].1.clone(),
2421 entries[entries.len() - 1].1.clone(),
2422 ));
2423 }
2424
2425 out
2426}
2427
2428fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
2429 let mut gaps = Vec::new();
2430 for idx in 1..entries.len() {
2431 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
2432 if score > 0.0 {
2433 gaps.push((score, idx));
2434 }
2435 }
2436 }
2437 gaps.sort_by(|left, right| {
2438 right
2439 .0
2440 .partial_cmp(&left.0)
2441 .unwrap_or(std::cmp::Ordering::Equal)
2442 .then_with(|| left.1.cmp(&right.1))
2443 });
2444 gaps.into_iter()
2445 .take(max_splits)
2446 .map(|(_, idx)| idx)
2447 .collect()
2448}
2449
2450fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
2451 if left.family() != right.family() {
2452 return None;
2453 }
2454 match (left, right) {
2455 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
2456 Some(r.saturating_sub(*l) as f64)
2457 }
2458 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
2459 Some(r.saturating_sub(*l) as f64)
2460 }
2461 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
2462 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
2463 }
2464 _ => None,
2465 }
2466}
2467
2468#[cfg(test)]
2469mod tests {
2470 use super::*;
2471 use crate::storage::unified::entity::RowData;
2472 use crate::storage::unified::MetadataValue;
2473 use reddb_types::Value;
2474
2475 fn keyed_row(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
2476 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
2477 entity.data = EntityData::Row(RowData::with_names(
2478 vec![Value::text(key), Value::Integer(id as i64)],
2479 vec!["key".into(), "value".into()],
2480 ));
2481 entity.set_xmin(xmin);
2482 entity.set_xmax(xmax);
2483 entity
2484 }
2485
2486 fn keyed_probe(segment: &GrowingSegment, key: &str) -> Option<EntityId> {
2487 segment
2488 .find_key_entity(key, true, &mut |_| true)
2489 .map(|entity| entity.id)
2490 }
2491
2492 #[test]
2493 fn key_lookup_visits_only_requested_key_versions() {
2494 let mut segment = GrowingSegment::new(1, "keys");
2495 segment
2496 .bulk_insert(
2497 (1..=4096)
2498 .map(|id| keyed_row(id, &format!("other-{id}"), 1, 0))
2499 .collect(),
2500 )
2501 .expect("unrelated rows");
2502 segment
2503 .insert(keyed_row(5000, "wanted", 2, 8))
2504 .expect("old version");
2505 segment
2506 .insert(keyed_row(5001, "wanted", 8, 0))
2507 .expect("new version");
2508 let mut visits = 0;
2509 let result = segment.find_key_entity("wanted", true, &mut |entity| {
2510 visits += 1;
2511 entity.xmin < 8
2512 });
2513 assert_eq!(result.map(|entity| entity.id), Some(EntityId::new(5000)));
2514 assert_eq!(
2515 visits, 2,
2516 "unrelated rows must never reach visibility evaluation"
2517 );
2518 assert_eq!(keyed_probe(&segment, "wanted"), Some(EntityId::new(5001)));
2519 let mut misses = 0;
2520 assert!(segment
2521 .find_key_entity("absent", true, &mut |_| {
2522 misses += 1;
2523 true
2524 })
2525 .is_none());
2526 assert_eq!(misses, 0);
2527
2528 let mut collision = SegmentKeyLookup::default();
2530 collision.insert(&keyed_row(1, "wanted", 1, 0));
2531 assert!(segment
2532 .find_key_candidate(&collision, "wanted", true, &mut |_| true)
2533 .is_none());
2534 }
2535
2536 #[test]
2537 fn raw_mutation_does_not_disable_unrelated_key_candidates() {
2538 let mut segment = GrowingSegment::new(1, "keys");
2539 segment
2540 .bulk_insert(
2541 (1..=4096)
2542 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2543 .collect(),
2544 )
2545 .expect("rows");
2546 for _ in 0..3 {
2547 *segment.get_mut(EntityId::new(1)).expect("mutable") = keyed_row(1, "renamed", 2, 0);
2548 }
2549 let index = segment.key_lookup.as_ref().expect("index retained");
2550 assert_eq!(index.candidates("key-4096"), &[EntityId::new(4096)]);
2551 assert_eq!(index.mutated_candidates().count(), 1);
2552 assert!(index.candidates("key-1").is_empty());
2553 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(1)));
2554 assert_eq!(keyed_probe(&segment, "key-4096"), Some(EntityId::new(4096)));
2555 segment.delete(EntityId::new(1)).expect("delete raw row");
2556 assert_eq!(
2557 segment
2558 .key_lookup
2559 .as_ref()
2560 .unwrap()
2561 .mutated_candidates()
2562 .count(),
2563 0
2564 );
2565 assert_eq!(keyed_probe(&segment, "renamed"), None);
2566 }
2567
2568 #[test]
2569 fn key_lookup_tracks_mutations_sealing_and_adoption() {
2570 let mut segment = GrowingSegment::new(1, "keys");
2571 segment
2572 .bulk_insert(vec![keyed_row(1, "flat", 1, 0)])
2573 .expect("flat row");
2574 segment
2575 .insert(keyed_row(3, "hash", 1, 0))
2576 .expect("hash fallback");
2577 segment
2578 .bulk_insert(vec![keyed_row(5, "gap", 1, 0)])
2579 .expect("bulk gap");
2580 assert!(segment.key_lookup.as_ref().expect("index").memory_bytes() > 0);
2581 assert_eq!(keyed_probe(&segment, "gap"), Some(EntityId::new(5)));
2582 segment
2583 .update(keyed_row(3, "renamed", 2, 0))
2584 .expect("rename key");
2585 assert!(
2586 segment.key_lookup.is_some(),
2587 "key replacements maintain the index"
2588 );
2589 assert_eq!(keyed_probe(&segment, "hash"), None);
2590 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(3)));
2591
2592 *segment.get_mut(EntityId::new(3)).expect("mutable row") = keyed_row(3, "raw", 3, 0);
2593 assert!(
2594 segment.key_lookup.is_some(),
2595 "raw mutation preserves unaffected key candidates"
2596 );
2597 assert_eq!(keyed_probe(&segment, "renamed"), None);
2598 assert_eq!(keyed_probe(&segment, "raw"), Some(EntityId::new(3)));
2599 segment.seal().expect("seal");
2600 segment
2601 .force_update_with_metadata(&keyed_row(5, "sealed", 4, 0), &["key".into()], None)
2602 .expect("sealed rename");
2603 assert_eq!(keyed_probe(&segment, "gap"), None);
2604 assert_eq!(keyed_probe(&segment, "sealed"), Some(EntityId::new(5)));
2605 assert!(segment.force_delete(EntityId::new(1)));
2606 assert!(segment.force_delete(EntityId::new(3)));
2607 assert_eq!(keyed_probe(&segment, "flat"), None);
2608 assert_eq!(keyed_probe(&segment, "raw"), None);
2609
2610 let mut adopted = GrowingSegment::new(2, "keys");
2611 segment.for_each_fast(|entity| {
2612 adopted.adopt_entity(entity.clone(), None);
2613 true
2614 });
2615 assert_eq!(keyed_probe(&adopted, "sealed"), Some(EntityId::new(5)));
2616 assert!(adopted.evict_entity(EntityId::new(5)));
2617 assert_eq!(keyed_probe(&adopted, "sealed"), None);
2618 }
2619
2620 #[test]
2621 #[ignore = "diagnostic microbenchmark; no timing assertions"]
2622 fn key_lookup_scaling_diagnostic() {
2623 use std::{hint::black_box, time::Instant};
2624 for count in [100, 10_000, 100_000] {
2625 let mut segment = GrowingSegment::new(1, "keys");
2626 segment
2627 .bulk_insert(
2628 (1..=count)
2629 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2630 .collect(),
2631 )
2632 .expect("rows");
2633 let key = format!("key-{count}");
2634 let started = Instant::now();
2635 for _ in 0..1000 {
2636 assert_eq!(
2637 keyed_probe(black_box(&segment), black_box(&key)),
2638 Some(EntityId::new(count))
2639 );
2640 }
2641 let indexed = started.elapsed();
2642 let started = Instant::now();
2643 for _ in 0..1000 {
2644 let mut found = None;
2645 segment.for_each_fast(|entity| {
2646 if SegmentKeyLookup::key(entity) == Some(black_box(key.as_str())) {
2647 found = Some(entity.id);
2648 return false;
2649 }
2650 true
2651 });
2652 assert_eq!(black_box(found), Some(EntityId::new(count)));
2653 }
2654 eprintln!(
2655 "rows={count} reads=1000 indexed={indexed:?} scan={:?}",
2656 started.elapsed()
2657 );
2658 }
2659 }
2660
2661 fn pruning_node(id: u64) -> UnifiedEntity {
2662 UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
2663 }
2664
2665 fn pruning_edge(id: u64) -> UnifiedEntity {
2666 UnifiedEntity::graph_edge(
2667 EntityId::new(id),
2668 "mixed",
2669 "a",
2670 "b",
2671 1.0,
2672 Default::default(),
2673 )
2674 }
2675
2676 #[test]
2677 fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
2678 let mut segment = GrowingSegment::new(1, "mixed");
2679 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
2680 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
2681 segment
2682 .bulk_insert(vec![
2683 UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
2684 pruning_node(2),
2685 pruning_edge(4),
2686 ])
2687 .expect("mixed bulk with gap");
2688 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2689 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2690 segment.seal().expect("seal");
2691 assert!(segment.force_delete(EntityId::new(2)));
2692 assert!(segment.force_delete(EntityId::new(4)));
2693 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2695 let mut rebuilt = GrowingSegment::new(2, "mixed");
2696 for item in segment.iter() {
2697 rebuilt.adopt_entity(item.clone(), None);
2698 }
2699 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2700 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2701 rebuilt.adopt_entity(pruning_node(5), None);
2702 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2703 rebuilt.insert(pruning_edge(6)).expect("single edge");
2704 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2705 }
2706
2707 #[test]
2708 fn graph_presence_covers_replacement_and_unrestricted_mutation() {
2709 for flat in [false, true] {
2710 for operation in 0..4 {
2711 let mut segment = GrowingSegment::new(1, "mixed");
2712 let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
2713 if flat {
2714 segment.bulk_insert(vec![vector]).expect("flat vector");
2715 } else {
2716 segment.insert(vector).expect("hashmap vector");
2717 }
2718 match operation {
2719 0 => segment.update(pruning_node(1)).expect("replace with node"),
2720 1 => segment
2721 .update_hot(pruning_node(1), &[])
2722 .expect("hot replacement"),
2723 2 => {
2724 segment.seal().expect("seal vector");
2725 segment
2726 .force_update_with_metadata(&pruning_node(1), &[], None)
2727 .expect("sealed replacement");
2728 }
2729 _ => {
2730 *segment.get_mut(EntityId::new(1)).expect("mutable vector") =
2731 pruning_node(1)
2732 }
2733 }
2734 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2735 assert!(matches!(
2736 segment.get(EntityId::new(1)).expect("node").kind,
2737 EntityKind::GraphNode(_)
2738 ));
2739 if operation == 3 {
2740 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2741 }
2742 segment
2743 .force_update_with_metadata(&pruning_edge(1), &[], None)
2744 .expect("replace with edge");
2745 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2746 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2747 }
2748 }
2749 }
2750
2751 fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
2752 UnifiedEntity::new(
2753 EntityId::new(id),
2754 EntityKind::TableRow {
2755 table: "keys".into(),
2756 row_id: id,
2757 },
2758 EntityData::Row(RowData::with_names(
2759 vec![Value::text("payload"), key],
2760 vec!["payload".into(), "primary".into()],
2761 )),
2762 )
2763 }
2764
2765 fn primary_key_probe(
2766 segment: &GrowingSegment,
2767 key: Value,
2768 exclude: Option<EntityId>,
2769 ) -> Option<EntityId> {
2770 segment.find_unique_key_conflict(
2771 &[vec!["primary".into()]],
2772 0,
2773 &[format!("{key:?}")],
2774 exclude,
2775 &|_| true,
2776 )
2777 }
2778
2779 #[test]
2780 fn unique_key_lookup_follows_segment_mutations_and_sealing() {
2781 let mut segment = GrowingSegment::new(1, "keys");
2782 segment
2783 .bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
2784 .expect("bulk row");
2785 let payload_bytes = segment.memory_bytes();
2786 assert_eq!(
2787 primary_key_probe(&segment, Value::Integer(10), None),
2788 Some(EntityId::new(1))
2789 );
2790 assert!(segment.memory_bytes() > payload_bytes);
2791 segment
2792 .insert(primary_key_row(3, Value::Integer(30)))
2793 .expect("hashmap fallback row");
2794 segment
2795 .bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
2796 .expect("bulk gap row");
2797 assert_eq!(
2798 primary_key_probe(&segment, Value::Integer(40), None),
2799 Some(EntityId::new(4))
2800 );
2801 assert_eq!(
2802 primary_key_probe(&segment, Value::Integer(30), None),
2803 Some(EntityId::new(3))
2804 );
2805 segment
2806 .update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
2807 .expect("replace flat key");
2808 assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
2809 assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
2810 assert_eq!(
2811 primary_key_probe(&segment, Value::Integer(11), None),
2812 Some(EntityId::new(1))
2813 );
2814 segment
2815 .delete_batch(&[EntityId::new(1), EntityId::new(3)])
2816 .expect("delete flat and map");
2817 assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
2818 assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
2819 segment.seal().expect("seal");
2820 assert_eq!(
2821 primary_key_probe(&segment, Value::Integer(40), None),
2822 Some(EntityId::new(4))
2823 );
2824 segment
2825 .force_update_with_metadata(
2826 &primary_key_row(4, Value::Integer(41)),
2827 &["primary".into()],
2828 None,
2829 )
2830 .expect("sealed update");
2831 assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
2832 assert_eq!(
2833 primary_key_probe(&segment, Value::Integer(41), None),
2834 Some(EntityId::new(4))
2835 );
2836 assert!(segment.force_delete(EntityId::new(4)));
2837 assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
2838 }
2839
2840 #[test]
2841 fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
2842 let mut segment = GrowingSegment::new(1, "keys");
2843 segment
2844 .bulk_insert(vec![
2845 primary_key_row(1, Value::Integer(10)),
2846 primary_key_row(2, Value::Integer(20)),
2847 ])
2848 .expect("initial rows");
2849 let definitions = vec![
2850 vec!["primary".into()],
2851 vec!["payload".into()],
2852 vec!["primary".into(), "payload".into()],
2853 ];
2854 let payload = format!("{:?}", Value::text("payload"));
2855 assert_eq!(
2856 segment.find_unique_key_conflict(
2857 &definitions,
2858 1,
2859 &[payload.clone()],
2860 None,
2861 &|entity| entity.id != EntityId::new(1)
2862 ),
2863 Some(EntityId::new(2))
2864 );
2865 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2866 segment
2867 .insert(primary_key_row(3, Value::Integer(30)))
2868 .expect("append");
2869 assert_eq!(
2870 segment.find_unique_key_conflict(
2871 &definitions,
2872 2,
2873 &[format!("{:?}", Value::Integer(30)), payload.clone()],
2874 None,
2875 &|_| true
2876 ),
2877 Some(EntityId::new(3))
2878 );
2879 let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
2880 assert_eq!(
2881 segment.find_unique_key_conflict(
2882 &definitions,
2883 0,
2884 &[format!("{:?}", Value::Integer(20))],
2885 None,
2886 &|_| true
2887 ),
2888 Some(EntityId::new(2))
2889 );
2890 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2891 assert_eq!(
2892 segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
2893 bytes
2894 );
2895 let reduced = vec![vec!["payload".into()]];
2896 assert_eq!(
2897 segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
2898 Some(EntityId::new(1))
2899 );
2900 assert_eq!(segment.unique_key_lookup.read().len(), 1);
2901 assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
2902 }
2903
2904 #[test]
2905 fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
2906 let mut segment = GrowingSegment::new(1, "keys");
2907 let values = [
2908 Value::Float(f64::NAN),
2909 Value::Float(-0.0),
2910 Value::Float(0.0),
2911 Value::Integer(1),
2912 Value::UnsignedInteger(1),
2913 Value::text("1"),
2914 ];
2915 for (offset, value) in values.iter().enumerate() {
2916 let id = offset as u64 + 1;
2917 segment
2918 .insert(primary_key_row(id, value.clone()))
2919 .expect("distinct physical row");
2920 assert_eq!(
2921 primary_key_probe(&segment, value.clone(), None),
2922 Some(EntityId::new(id))
2923 );
2924 assert_eq!(
2925 primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
2926 None
2927 );
2928 }
2929 segment
2930 .insert(primary_key_row(20, Value::Integer(1)))
2931 .expect("duplicate physical key");
2932 assert_eq!(
2933 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
2934 Some(EntityId::new(20))
2935 );
2936 segment
2939 .entities
2940 .insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
2941 assert_eq!(
2942 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
2943 None
2944 );
2945 let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
2947 *entity = primary_key_row(4, Value::Integer(100));
2948 assert_eq!(
2949 primary_key_probe(&segment, Value::Integer(100), None),
2950 Some(EntityId::new(4))
2951 );
2952 }
2953
2954 #[test]
2955 fn test_growing_segment_basic() {
2956 let mut segment = GrowingSegment::new(1, "test");
2957
2958 let entity = UnifiedEntity::table_row(
2959 EntityId::new(1),
2960 "users",
2961 1,
2962 vec![Value::text("Alice".to_string())],
2963 );
2964
2965 let id = segment.insert(entity).unwrap();
2966 assert_eq!(id, EntityId::new(1));
2967 assert!(segment.contains(id));
2968
2969 let stats = segment.stats();
2970 assert_eq!(stats.entity_count, 1);
2971 assert_eq!(stats.row_count, 1);
2972 }
2973
2974 #[test]
2975 fn test_segment_metadata() {
2976 let mut segment = GrowingSegment::new(1, "test");
2977
2978 let entity = UnifiedEntity::table_row(
2979 EntityId::new(1),
2980 "users",
2981 1,
2982 vec![Value::text("Alice".to_string())],
2983 );
2984 segment.insert(entity).unwrap();
2985
2986 let mut meta = Metadata::new();
2987 meta.set("role", MetadataValue::String("admin".to_string()));
2988 meta.set("level", MetadataValue::Int(5));
2989
2990 segment.set_metadata(EntityId::new(1), meta).unwrap();
2991
2992 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
2993 assert_eq!(
2994 retrieved.get("role"),
2995 Some(&MetadataValue::String("admin".to_string()))
2996 );
2997 }
2998
2999 #[test]
3000 fn test_segment_seal() {
3001 let mut segment = GrowingSegment::new(1, "test");
3002
3003 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
3004 segment.insert(entity).unwrap();
3005
3006 assert!(segment.state().is_writable());
3008
3009 segment.seal().unwrap();
3011 assert_eq!(segment.state(), SegmentState::Sealed);
3012
3013 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
3015 assert!(segment.insert(entity2).is_err());
3016 }
3017
3018 #[test]
3019 fn test_should_seal() {
3020 let mut segment = GrowingSegment::new(1, "test");
3021
3022 let config = SegmentConfig {
3023 max_entities: 2,
3024 ..Default::default()
3025 };
3026
3027 assert!(!segment.should_seal(&config));
3028
3029 segment
3030 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
3031 .unwrap();
3032 assert!(!segment.should_seal(&config));
3033
3034 segment
3035 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
3036 .unwrap();
3037 assert!(segment.should_seal(&config));
3038 }
3039
3040 #[test]
3041 fn test_cross_references() {
3042 let mut segment = GrowingSegment::new(1, "test");
3043
3044 let mut entity1 = UnifiedEntity::table_row(
3045 EntityId::new(1),
3046 "hosts",
3047 1,
3048 vec![Value::text("192.168.1.1".to_string())],
3049 );
3050 entity1.add_cross_ref(CrossRef::new(
3051 EntityId::new(1),
3052 EntityId::new(2),
3053 "nodes",
3054 RefType::RowToNode,
3055 ));
3056 segment.insert(entity1).unwrap();
3057
3058 let refs_from = segment.get_references_from(EntityId::new(1));
3059 assert_eq!(refs_from.len(), 1);
3060 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
3061
3062 let refs_to = segment.get_references_to(EntityId::new(2));
3063 assert_eq!(refs_to.len(), 1);
3064 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
3065 }
3066
3067 #[test]
3068 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
3069 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
3070 zone.update(&Value::Email("delta@example.com".to_string()));
3071
3072 let probe = Value::Email("alpha@example.com".to_string());
3073 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
3074
3075 let in_range = Value::Email("charlie@example.com".to_string());
3076 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
3077 }
3078
3079 #[test]
3080 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
3081 let mut segment = GrowingSegment::new(1, "test");
3082
3083 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
3084 let entity = UnifiedEntity::new(
3085 EntityId::new(row_id),
3086 EntityKind::TableRow {
3087 table: "users".into(),
3088 row_id,
3089 },
3090 EntityData::Row(RowData::with_names(
3091 vec![Value::Integer(age)],
3092 vec!["age".to_string()],
3093 )),
3094 );
3095 segment.insert(entity).unwrap();
3096 }
3097
3098 segment.seal().unwrap();
3099
3100 let miss = Value::Integer(500);
3101 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
3102
3103 let hit = Value::Integer(1000);
3104 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
3105 }
3106}