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) const CHECKPOINT_MUTATIONS_MAX: usize = 8192;
474
475pub(crate) struct SegmentScanCursor {
477 flat_end: usize,
478 hash_end: usize,
479 deleted_end: usize,
480 position: usize,
481}
482
483impl SegmentScanCursor {
484 pub(crate) fn remaining(&self) -> usize {
485 self.flat_end + self.hash_end - self.position
486 }
487
488 pub(crate) fn finished(&self) -> bool {
489 self.remaining() == 0
490 }
491}
492
493pub(crate) struct GraphIndexCursor {
496 key_index: usize,
497 after: Option<EntityId>,
498 upper: EntityId,
499}
500
501impl GraphIndexCursor {
502 pub(crate) fn finished(&self, keys: &[EntityId]) -> bool {
503 self.key_index == keys.len()
504 }
505}
506
507pub struct GrowingSegment {
509 id: SegmentId,
511 collection: String,
513 state: SegmentState,
515 created_at: u64,
517 last_write_at: u64,
519
520 entities: HashMap<EntityId, UnifiedEntity>,
522 hash_scan_ids: Vec<EntityId>,
525 flat_entities: Vec<UnifiedEntity>,
528 base_entity_id: u64,
530 use_flat: bool,
532 deleted: HashMap<EntityId, usize>,
536 metadata: MetadataStorage,
538
539 pk_index: BTreeMap<(String, String), EntityId>,
541 unique_key_lookup: parking_lot::RwLock<Vec<SegmentUniqueKey>>,
544 unique_key_lookup_bytes: AtomicU64,
545 key_lookup: Option<SegmentKeyLookup>,
546 kind_index: HashMap<String, HashSet<EntityId>>,
548 graph_kind_mask: u8,
551 graph_read_index: parking_lot::RwLock<Option<SegmentGraphIndex>>,
553 graph_read_index_bytes: AtomicU64,
554 cross_ref_forward: HashMap<EntityId, Vec<(EntityId, RefType)>>,
556 cross_ref_reverse: HashMap<EntityId, Vec<(EntityId, RefType)>>,
558
559 col_zones: HashMap<String, ColZone>,
561 zone_payload_bytes: usize,
562 sealed_col_zones: HashMap<String, MultiColZone>,
564
565 sequence: AtomicU64,
567 checkpoint_revision: u64,
569 checkpoint_mutations: std::collections::VecDeque<(u64, EntityId)>,
570 memory_bytes: AtomicU64,
575 dead_entity_bytes: u64,
578 dead_resident_count: usize,
580
581 pub(crate) published_flat_len: AtomicUsize,
588}
589
590impl GrowingSegment {
591 pub(crate) fn may_contain_graph_kind(&self, kind: GraphEntityKind) -> bool {
592 self.graph_kind_mask & kind.mask() != 0
593 }
594
595 fn observe_graph_kind(&mut self, kind: &EntityKind) {
596 self.graph_kind_mask |= match kind {
597 EntityKind::GraphNode(_) => GraphEntityKind::Node.mask(),
598 EntityKind::GraphEdge(_) => GraphEntityKind::Edge.mask(),
599 _ => 0,
600 };
601 }
602
603 pub(crate) fn scan_cursor(&self) -> SegmentScanCursor {
604 SegmentScanCursor {
605 flat_end: self.flat_entities.len(),
606 hash_end: self.hash_scan_ids.len(),
607 deleted_end: self.deleted.len(),
608 position: 0,
609 }
610 }
611
612 pub(crate) fn scan_batch(
616 &self,
617 cursor: &mut SegmentScanCursor,
618 before_work: &mut impl FnMut() -> bool,
619 visit: &mut impl FnMut(EntityId),
620 ) -> bool {
621 for _ in 0..SCAN_BATCH_SIZE {
622 if cursor.finished() {
623 break;
624 }
625 if !before_work() {
626 return false;
627 }
628 let id = if cursor.position < cursor.flat_end {
629 self.flat_entities[cursor.position].id
630 } else {
631 self.hash_scan_ids[cursor.position - cursor.flat_end]
632 };
633 cursor.position += 1;
634 let eligible = match self.deleted.get(&id) {
635 Some(&ordinal) => ordinal >= cursor.deleted_end,
636 None => self.get(id).is_some(),
639 };
640 if eligible {
641 visit(id);
642 }
643 }
644 true
645 }
646
647 #[inline]
650 pub fn for_each_fast<F>(&self, mut f: F) -> bool
651 where
652 F: FnMut(&UnifiedEntity) -> bool,
653 {
654 if self.use_flat {
655 if self.deleted.is_empty() {
657 for entity in &self.flat_entities {
658 if !f(entity) {
659 return false;
660 }
661 }
662 for entity in self.entities.values() {
665 if !f(entity) {
666 return false;
667 }
668 }
669 } else {
670 for entity in &self.flat_entities {
671 if self.deleted.contains_key(&entity.id) {
672 continue;
673 }
674 if !f(entity) {
675 return false;
676 }
677 }
678 for entity in self.entities.values() {
679 if self.deleted.contains_key(&entity.id) {
680 continue;
681 }
682 if !f(entity) {
683 return false;
684 }
685 }
686 }
687 } else {
688 if self.deleted.is_empty() {
690 for entity in self.entities.values() {
691 if !f(entity) {
692 return false;
693 }
694 }
695 } else {
696 for entity in self.entities.values() {
697 if self.deleted.contains_key(&entity.id) {
698 continue;
699 }
700 if !f(entity) {
701 return false;
702 }
703 }
704 }
705 }
706 true
707 }
708
709 pub fn new(id: SegmentId, collection: impl Into<String>) -> Self {
711 let now = current_unix_secs();
712
713 Self {
714 id,
715 collection: collection.into(),
716 state: SegmentState::Growing,
717 created_at: now,
718 last_write_at: now,
719 entities: HashMap::new(),
720 hash_scan_ids: Vec::new(),
721 flat_entities: Vec::new(),
722 base_entity_id: 0,
723 use_flat: false,
724 deleted: HashMap::new(),
725 metadata: MetadataStorage::new(),
726 pk_index: BTreeMap::new(),
727 unique_key_lookup: parking_lot::RwLock::new(Vec::new()),
728 unique_key_lookup_bytes: AtomicU64::new(0),
729 key_lookup: Some(SegmentKeyLookup::default()),
730 kind_index: HashMap::new(),
731 graph_read_index: parking_lot::RwLock::new(Some(SegmentGraphIndex::default())),
732 graph_read_index_bytes: AtomicU64::new(0),
733 graph_kind_mask: 0,
734 cross_ref_forward: HashMap::new(),
735 cross_ref_reverse: HashMap::new(),
736 col_zones: HashMap::new(),
737 zone_payload_bytes: 0,
738 sealed_col_zones: HashMap::new(),
739 sequence: AtomicU64::new(0),
740 checkpoint_revision: 0,
741 checkpoint_mutations: std::collections::VecDeque::new(),
742 memory_bytes: AtomicU64::new(0),
743 dead_entity_bytes: 0,
744 dead_resident_count: 0,
745 published_flat_len: AtomicUsize::new(0),
746 }
747 }
748
749 fn next_sequence(&self) -> u64 {
751 self.sequence.fetch_add(1, Ordering::SeqCst)
752 }
753
754 pub(crate) fn checkpoint_revision(&self) -> u64 {
755 self.checkpoint_revision
756 }
757
758 fn mark_checkpoint_dirty(&mut self) {
759 self.checkpoint_mutations.clear();
760 self.checkpoint_revision = self
761 .checkpoint_revision
762 .checked_add(1)
763 .expect("segment mutation counter must not overflow");
764 }
765
766 fn mark_checkpoint_entity_dirty(&mut self, id: EntityId) {
767 self.checkpoint_revision = self
768 .checkpoint_revision
769 .checked_add(1)
770 .expect("segment mutation counter must not overflow");
771 if self.checkpoint_mutations.len() == CHECKPOINT_MUTATIONS_MAX {
772 self.checkpoint_mutations.pop_front();
773 }
774 self.checkpoint_mutations
775 .push_back((self.checkpoint_revision, id));
776 }
777
778 pub(crate) fn checkpoint_mutations_since(&self, revision: u64) -> Option<Vec<EntityId>> {
779 if revision == self.checkpoint_revision {
780 return Some(Vec::new());
781 }
782 let mut expected = revision.checked_add(1)?;
783 let mut ids = Vec::new();
784 for &(sequence, id) in &self.checkpoint_mutations {
785 if sequence <= revision {
786 continue;
787 }
788 if sequence != expected {
789 return None;
790 }
791 ids.push(id);
792 expected = expected.checked_add(1)?;
793 }
794 (expected == self.checkpoint_revision.checked_add(1)?).then_some(ids)
795 }
796
797 pub(crate) fn checkpoint_tracking_bytes(&self) -> usize {
798 self.checkpoint_mutations.capacity() * std::mem::size_of::<(u64, EntityId)>()
799 }
800
801 fn has_live_entity(&self, id: EntityId) -> bool {
802 if self.deleted.contains_key(&id) {
803 return false;
804 }
805 if self.use_flat {
806 let raw = id.raw();
807 if raw >= self.base_entity_id {
808 let idx = (raw - self.base_entity_id) as usize;
809 if self
810 .flat_entities
811 .get(idx)
812 .is_some_and(|entity| entity.id == id)
813 {
814 return true;
815 }
816 }
817 self.entities.contains_key(&id)
820 } else {
821 self.entities.contains_key(&id)
822 }
823 }
824
825 fn update_existing_entity_in_place(
826 &mut self,
827 entity: &UnifiedEntity,
828 ) -> Result<UpdateIndexSnapshot, SegmentError> {
829 self.mark_checkpoint_entity_dirty(entity.id);
830 self.invalidate_unique_key_lookup();
831 if self
832 .get(entity.id)
833 .is_some_and(|old| SegmentKeyLookup::key(old) != SegmentKeyLookup::key(entity))
834 {
835 if let Some(mut index) = self.key_lookup.take() {
836 index.remove(
837 self.get(entity.id)
838 .expect("invariant: existing replacement row"),
839 );
840 index.insert(entity);
841 self.key_lookup = Some(index);
842 }
843 }
844 if (self.graph_kind_mask != 0
845 || matches!(
846 entity.kind,
847 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
848 ))
849 && self
850 .get(entity.id)
851 .is_some_and(|old| !SegmentGraphIndex::same_keys(old, entity))
852 {
853 self.invalidate_graph_read_index();
854 }
855 self.observe_graph_kind(&entity.kind);
858 if self.use_flat {
859 let raw = entity.id.raw();
860 if raw >= self.base_entity_id {
861 let idx = (raw - self.base_entity_id) as usize;
862 if let Some(slot) = self
863 .flat_entities
864 .get_mut(idx)
865 .filter(|slot| slot.id == entity.id)
866 {
867 let snapshot = UpdateIndexSnapshot::from_entity(slot);
868 let old_bytes = Self::estimate_entity_size(slot);
869 slot.clone_from(entity);
870 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
871 return Ok(snapshot);
872 }
873 }
874 let Some(slot) = self.entities.get_mut(&entity.id) else {
875 return Err(SegmentError::NotFound(entity.id));
876 };
877 let snapshot = UpdateIndexSnapshot::from_entity(slot);
878 let old_bytes = Self::estimate_entity_size(slot);
879 slot.clone_from(entity);
880 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
881 Ok(snapshot)
882 } else {
883 let Some(slot) = self.entities.get_mut(&entity.id) else {
884 return Err(SegmentError::NotFound(entity.id));
885 };
886 let snapshot = UpdateIndexSnapshot::from_entity(slot);
887 let old_bytes = Self::estimate_entity_size(slot);
888 slot.clone_from(entity);
889 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
890 Ok(snapshot)
891 }
892 }
893
894 fn apply_hot_update_with_metadata(
895 &mut self,
896 entity: &UnifiedEntity,
897 modified_columns: &[String],
898 metadata: Option<&Metadata>,
899 ) -> Result<(), SegmentError> {
900 let old = self.update_existing_entity_in_place(entity)?;
901 self.reindex_for_update(&old, entity, Some(modified_columns));
902 self.update_col_zones_from_entity(entity);
903 if let Some(metadata) = metadata {
904 self.metadata.set_all(entity.id, metadata);
905 }
906 Ok(())
907 }
908
909 fn apply_update_with_metadata(
910 &mut self,
911 entity: &UnifiedEntity,
912 metadata: Option<&Metadata>,
913 ) -> Result<(), SegmentError> {
914 let old = self.update_existing_entity_in_place(entity)?;
915 self.reindex_for_update(&old, entity, None);
916 self.update_col_zones_from_entity(entity);
917 if let Some(metadata) = metadata {
918 self.metadata.set_all(entity.id, metadata);
919 }
920 Ok(())
921 }
922
923 pub fn update_hot_batch_with_metadata<'a, I>(&mut self, items: I) -> Result<(), SegmentError>
924 where
925 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
926 {
927 if !self.state.is_writable() {
928 return Err(SegmentError::NotWritable);
929 }
930
931 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
932 items.into_iter().collect();
933 if items.is_empty() {
934 return Ok(());
935 }
936
937 for (entity, _, _) in &items {
938 if !self.has_live_entity(entity.id) {
939 return Err(SegmentError::NotFound(entity.id));
940 }
941 }
942
943 for (entity, modified_columns, metadata) in items {
944 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)?;
945 }
946
947 self.last_write_at = current_unix_secs();
948 Ok(())
949 }
950
951 fn record_tombstone(&mut self, id: EntityId) -> bool {
953 self.mark_checkpoint_entity_dirty(id);
954 let ordinal = self.deleted.len();
955 match self.deleted.entry(id) {
956 std::collections::hash_map::Entry::Vacant(entry) => {
957 entry.insert(ordinal);
958 true
959 }
960 std::collections::hash_map::Entry::Occupied(_) => false,
961 }
962 }
963
964 fn tombstone_flat_entity(&mut self, idx: usize, id: EntityId) -> bool {
970 self.invalidate_unique_key_lookup();
971 if !self.record_tombstone(id) {
972 return false;
973 }
974 if let Some(index) = self.key_lookup.as_mut() {
975 index.remove(&self.flat_entities[idx]);
976 }
977 if let Some(index) = self.graph_read_index.get_mut() {
978 index.remove(&self.flat_entities[idx]);
979 self.graph_read_index_bytes
980 .store(index.memory_bytes(), Ordering::Relaxed);
981 }
982 let size = Self::estimate_entity_size(&self.flat_entities[idx]) as u64;
983 self.dead_entity_bytes += size;
984 self.dead_resident_count += 1;
985 self.metadata.remove_all(id);
986 true
987 }
988
989 fn remove_hashmap_entity(&mut self, id: EntityId) -> bool {
994 self.invalidate_unique_key_lookup();
995 let Some(entity) = self.entities.remove(&id) else {
996 return false;
997 };
998 self.release_memory(Self::estimate_entity_size(&entity));
999 self.unindex_entity(&entity);
1000 self.metadata.remove_all(id);
1001 self.record_tombstone(id);
1002 true
1003 }
1004
1005 pub fn delete_batch(&mut self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1006 if !self.state.is_writable() {
1007 return Err(SegmentError::NotWritable);
1008 }
1009 if ids.is_empty() {
1010 return Ok(Vec::new());
1011 }
1012
1013 let mut deleted_ids = Vec::with_capacity(ids.len());
1014
1015 if self.use_flat {
1016 for &id in ids {
1017 let raw = id.raw();
1018 if raw < self.base_entity_id {
1019 if self.remove_hashmap_entity(id) {
1020 deleted_ids.push(id);
1021 }
1022 continue;
1023 }
1024 let idx = (raw - self.base_entity_id) as usize;
1025 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
1026 if self.tombstone_flat_entity(idx, id) {
1027 deleted_ids.push(id);
1028 }
1029 } else if self.remove_hashmap_entity(id) {
1030 deleted_ids.push(id);
1031 }
1032 }
1033 } else {
1034 for &id in ids {
1035 if self.remove_hashmap_entity(id) {
1036 deleted_ids.push(id);
1037 }
1038 }
1039 }
1040
1041 if !deleted_ids.is_empty() {
1042 self.last_write_at = current_unix_secs();
1043 }
1044
1045 Ok(deleted_ids)
1046 }
1047
1048 fn add_memory(&self, bytes: usize) {
1050 self.memory_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
1051 }
1052
1053 fn replace_entity_memory(&self, previous: usize, current: usize) {
1054 if current >= previous {
1055 self.add_memory(current - previous);
1056 } else {
1057 self.release_memory(previous - current);
1058 }
1059 }
1060
1061 pub fn memory_bytes(&self) -> u64 {
1066 self.memory_bytes.load(Ordering::Relaxed)
1067 + self.zone_payload_bytes as u64
1068 + self.unique_key_lookup_bytes.load(Ordering::Relaxed)
1069 + self
1070 .key_lookup
1071 .as_ref()
1072 .map_or(0, SegmentKeyLookup::memory_bytes)
1073 + self.graph_read_index_bytes.load(Ordering::Relaxed)
1074 + (self.hash_scan_ids.capacity() * std::mem::size_of::<EntityId>()) as u64
1075 }
1076
1077 fn release_memory(&self, bytes: usize) {
1080 let mut current = self.memory_bytes.load(Ordering::Relaxed);
1081 loop {
1082 let next = current.saturating_sub(bytes as u64);
1083 match self.memory_bytes.compare_exchange_weak(
1084 current,
1085 next,
1086 Ordering::Relaxed,
1087 Ordering::Relaxed,
1088 ) {
1089 Ok(_) => break,
1090 Err(observed) => current = observed,
1091 }
1092 }
1093 }
1094
1095 pub(crate) fn resident_entity_count(&self) -> usize {
1097 self.flat_entities.len() + self.entities.len()
1098 }
1099
1100 pub(crate) fn live_entity_count(&self) -> usize {
1102 self.resident_entity_count()
1103 .saturating_sub(self.dead_resident_count)
1104 }
1105
1106 pub(crate) fn tombstone_count(&self) -> usize {
1109 self.deleted.len()
1110 }
1111
1112 pub(crate) fn resident_bytes(&self) -> u64 {
1115 self.memory_bytes() + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1116 }
1117
1118 pub(crate) fn reclaimable_bytes(&self) -> u64 {
1121 self.dead_entity_bytes + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1122 }
1123
1124 pub(crate) fn live_entity_ids(&self) -> Vec<EntityId> {
1130 let mut ids = Vec::with_capacity(self.live_entity_count());
1131 self.for_each_fast(|entity| {
1132 ids.push(entity.id);
1133 true
1134 });
1135 ids
1136 }
1137
1138 pub(crate) fn insert_preserving_sequence(
1142 &mut self,
1143 entity: UnifiedEntity,
1144 ) -> Result<EntityId, SegmentError> {
1145 if !self.state.is_writable() {
1146 return Err(SegmentError::NotWritable);
1147 }
1148 if self.entities.contains_key(&entity.id) {
1149 return Err(SegmentError::AlreadyExists(entity.id));
1150 }
1151 let id = entity.id;
1152 self.adopt_entity(entity, None);
1153 self.last_write_at = current_unix_secs();
1154 Ok(id)
1155 }
1156
1157 pub(crate) fn adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
1165 let id = entity.id;
1166
1167 let next = entity.sequence_id.saturating_add(1);
1170 if self.sequence.load(Ordering::Relaxed) < next {
1171 self.sequence.store(next, Ordering::Relaxed);
1172 }
1173
1174 self.add_memory(Self::estimate_entity_size(&entity));
1175 self.index_entity(&entity);
1176 self.update_col_zones_from_entity(&entity);
1177 self.entities.insert(id, entity);
1178
1179 if let Some(metadata) = metadata {
1180 if !metadata.is_empty() {
1181 self.metadata.set_all(id, &metadata);
1182 }
1183 }
1184 }
1185
1186 pub(crate) fn evict_entity(&mut self, id: EntityId) -> bool {
1192 self.mark_checkpoint_entity_dirty(id);
1193 self.invalidate_unique_key_lookup();
1194 let Some(entity) = self.entities.remove(&id) else {
1195 return false;
1196 };
1197 self.release_memory(Self::estimate_entity_size(&entity));
1198 self.unindex_entity(&entity);
1199 for cross_ref in entity.cross_refs() {
1200 if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1201 reverse.retain(|(source, _)| *source != id);
1202 }
1203 }
1204 self.metadata.remove_all(id);
1205 true
1206 }
1207
1208 fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
1210 super::memory_size::entity_bytes(entity)
1211 }
1212
1213 fn update_column_zone(&mut self, column: &str, value: &Value) {
1221 if matches!(value, Value::Null) {
1222 return;
1223 }
1224 match self.col_zones.entry(column.to_string()) {
1225 std::collections::hash_map::Entry::Occupied(mut entry) => {
1226 let before = entry.get().payload_bytes();
1227 entry.get_mut().update(value);
1228 self.zone_payload_bytes = self
1229 .zone_payload_bytes
1230 .saturating_sub(before)
1231 .saturating_add(entry.get().payload_bytes());
1232 }
1233 std::collections::hash_map::Entry::Vacant(entry) => {
1234 let zone = ColZone::new(value.clone());
1235 self.zone_payload_bytes =
1236 self.zone_payload_bytes.saturating_add(zone.payload_bytes());
1237 entry.insert(zone);
1238 }
1239 }
1240 }
1241
1242 fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
1243 if let EntityData::Row(row) = &entity.data {
1244 if let Some(named) = &row.named {
1245 for (column, value) in named {
1246 self.update_column_zone(column, value);
1247 }
1248 } else if let Some(schema) = &row.schema {
1249 for (column, value) in schema.iter().zip(&row.columns) {
1250 self.update_column_zone(column, value);
1251 }
1252 }
1253 }
1254 }
1255
1256 fn rebuild_sealed_col_zones(&mut self) {
1257 let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
1258 let mut family_by_col: HashMap<String, reddb_types::CanonicalKeyFamily> = HashMap::new();
1259 let mut mixed_family_cols = HashSet::new();
1260 let mut unsupported_cols = HashSet::new();
1261
1262 let mut observe_row = |row: &super::entity::RowData| {
1263 for (col, value) in row.iter_fields() {
1264 if matches!(value, Value::Null) {
1265 continue;
1266 }
1267 let Some(key) = value_to_canonical_key(value) else {
1268 unsupported_cols.insert(col.to_string());
1269 continue;
1270 };
1271 match family_by_col.get(col).copied() {
1272 Some(existing) if existing != key.family() => {
1273 mixed_family_cols.insert(col.to_string());
1274 }
1275 None => {
1276 family_by_col.insert(col.to_string(), key.family());
1277 }
1278 _ => {}
1279 }
1280 values_by_col
1281 .entry(col.to_string())
1282 .or_default()
1283 .push((key, value.clone()));
1284 }
1285 };
1286
1287 if self.use_flat {
1288 for entity in &self.flat_entities {
1289 if self.deleted.contains_key(&entity.id) {
1290 continue;
1291 }
1292 if let EntityData::Row(row) = &entity.data {
1293 observe_row(row);
1294 }
1295 }
1296 } else {
1297 for entity in self.entities.values() {
1298 if self.deleted.contains_key(&entity.id) {
1299 continue;
1300 }
1301 if let EntityData::Row(row) = &entity.data {
1302 observe_row(row);
1303 }
1304 }
1305 }
1306
1307 let mut sealed_col_zones = HashMap::new();
1308 for (col, mut entries) in values_by_col {
1309 if mixed_family_cols.contains(&col)
1310 || unsupported_cols.contains(&col)
1311 || entries.is_empty()
1312 {
1313 continue;
1314 }
1315 entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
1316 entries.dedup_by(|left, right| left.0 == right.0);
1317
1318 let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
1319 if intervals.len() > 1 {
1320 sealed_col_zones.insert(col, MultiColZone { intervals });
1321 }
1322 }
1323
1324 let previous = self
1325 .sealed_col_zones
1326 .values()
1327 .flat_map(|zones| &zones.intervals)
1328 .map(ColZone::payload_bytes)
1329 .fold(0, usize::saturating_add);
1330 let current = sealed_col_zones
1331 .values()
1332 .flat_map(|zones| &zones.intervals)
1333 .map(ColZone::payload_bytes)
1334 .fold(0, usize::saturating_add);
1335 self.zone_payload_bytes = self
1336 .zone_payload_bytes
1337 .saturating_sub(previous)
1338 .saturating_add(current);
1339 self.sealed_col_zones = sealed_col_zones;
1340 }
1341
1342 pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
1346 if preds.is_empty() {
1347 return false;
1348 }
1349 for (col, pred) in preds {
1350 if let Some(zone) = self.sealed_col_zones.get(*col) {
1351 if zone.can_skip(pred) {
1352 return true;
1353 }
1354 continue;
1355 }
1356 if let Some(zone) = self.col_zones.get(*col) {
1357 if pred.can_skip(zone) {
1358 return true; }
1360 }
1361 }
1362 false
1363 }
1364
1365 fn index_entity(&mut self, entity: &UnifiedEntity) {
1367 self.mark_checkpoint_entity_dirty(entity.id);
1368 if let Some(index) = self.key_lookup.as_mut() {
1369 index.insert(entity);
1370 }
1371 self.hash_scan_ids.push(entity.id);
1372 self.observe_graph_kind(&entity.kind);
1373 if let Some(index) = self.graph_read_index.get_mut() {
1374 index.insert(entity);
1375 self.graph_read_index_bytes
1376 .store(index.memory_bytes(), Ordering::Relaxed);
1377 }
1378 let indexes = self.unique_key_lookup.get_mut();
1379 for index in indexes.iter_mut() {
1380 index.insert(entity);
1381 }
1382 self.unique_key_lookup_bytes.store(
1383 indexes
1384 .iter()
1385 .map(|index| index.memory_bytes() as u64)
1386 .sum(),
1387 Ordering::Relaxed,
1388 );
1389 let kind_key = entity.kind.storage_type().to_string();
1391 self.kind_index
1392 .entry(kind_key)
1393 .or_default()
1394 .insert(entity.id);
1395
1396 if let EntityData::Row(row) = &entity.data {
1398 if let Some(first_col) = row.columns.first() {
1399 let pk_str = format!("{:?}", first_col);
1400 self.pk_index
1401 .insert((entity.kind.collection().to_string(), pk_str), entity.id);
1402 }
1403 }
1404
1405 for cross_ref in entity.cross_refs() {
1407 self.cross_ref_forward
1408 .entry(cross_ref.source)
1409 .or_default()
1410 .push((cross_ref.target, cross_ref.ref_type));
1411
1412 self.cross_ref_reverse
1413 .entry(cross_ref.target)
1414 .or_default()
1415 .push((cross_ref.source, cross_ref.ref_type));
1416 }
1417 }
1418
1419 fn invalidate_graph_read_index(&mut self) {
1420 *self.graph_read_index.get_mut() = None;
1421 self.graph_read_index_bytes.store(0, Ordering::Relaxed);
1422 }
1423
1424 #[cfg(test)]
1427 pub(crate) fn visit_graph_index_candidates(
1428 &self,
1429 kind: GraphEntityKind,
1430 keys: &[EntityId],
1431 before_work: &mut impl FnMut() -> bool,
1432 visit: &mut impl FnMut(&UnifiedEntity) -> bool,
1433 ) -> bool {
1434 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1435 return false;
1436 };
1437 let index = cached.as_ref().expect("complete graph index");
1438 for key in keys {
1439 if !before_work() {
1440 return false;
1441 }
1442 for id in index.candidates(kind, *key) {
1443 if !before_work() {
1444 return false;
1445 }
1446 if let Some(entity) = self.get(id) {
1447 if !visit(entity) {
1448 return false;
1449 }
1450 }
1451 }
1452 }
1453 true
1454 }
1455
1456 fn graph_read_index_for_query(
1457 &self,
1458 before_work: &mut impl FnMut() -> bool,
1459 ) -> Option<parking_lot::RwLockReadGuard<'_, Option<SegmentGraphIndex>>> {
1460 let mut cached = self.graph_read_index.read();
1461 if cached.is_none() {
1462 drop(cached);
1463 let mut writer = self.graph_read_index.write();
1464 if writer.is_none() {
1465 let mut index = SegmentGraphIndex::default();
1466 if !self.for_each_fast(|entity| {
1469 if !before_work() {
1470 return false;
1471 }
1472 index.insert(entity);
1473 true
1474 }) {
1475 return None;
1476 }
1477 self.graph_read_index_bytes
1478 .store(index.memory_bytes(), Ordering::Relaxed);
1479 *writer = Some(index);
1480 }
1481 cached = parking_lot::RwLockWriteGuard::downgrade(writer);
1482 }
1483 Some(cached)
1484 }
1485
1486 pub(crate) fn graph_index_cursor(
1487 &self,
1488 kind: GraphEntityKind,
1489 keys: &[EntityId],
1490 before_work: &mut impl FnMut() -> bool,
1491 ) -> Option<GraphIndexCursor> {
1492 let cached = self.graph_read_index_for_query(before_work)?;
1493 let index = cached.as_ref().expect("complete graph index");
1494 let mut upper = EntityId::new(0);
1495 for key in keys {
1496 if !before_work() {
1497 return None;
1498 }
1499 if let Some(last) = index.candidates(kind, *key).next_back() {
1500 upper = upper.max(last);
1501 }
1502 }
1503 Some(GraphIndexCursor {
1504 key_index: 0,
1505 after: None,
1506 upper,
1507 })
1508 }
1509
1510 pub(crate) fn graph_index_batch(
1513 &self,
1514 kind: GraphEntityKind,
1515 keys: &[EntityId],
1516 cursor: &mut GraphIndexCursor,
1517 before_work: &mut impl FnMut() -> bool,
1518 visit: &mut impl FnMut(EntityId),
1519 ) -> bool {
1520 self.graph_index_batch_with_limit(kind, keys, cursor, before_work, visit, SCAN_BATCH_SIZE)
1521 }
1522
1523 pub(crate) fn graph_index_batch_with_limit(
1524 &self,
1525 kind: GraphEntityKind,
1526 keys: &[EntityId],
1527 cursor: &mut GraphIndexCursor,
1528 before_work: &mut impl FnMut() -> bool,
1529 visit: &mut impl FnMut(EntityId),
1530 limit: usize,
1531 ) -> bool {
1532 assert!(
1533 limit > 0 && limit <= SCAN_BATCH_SIZE,
1534 "bounded graph index batch"
1535 );
1536 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1537 return false;
1538 };
1539 let index = cached.as_ref().expect("complete graph index");
1540 let mut count = 0;
1541 while !cursor.finished(keys) && count < limit {
1542 let key = keys[cursor.key_index];
1543 if !before_work() {
1544 return false;
1545 }
1546 for id in index
1547 .candidates_after(kind, key, cursor.after, cursor.upper)
1548 .take(limit - count)
1549 {
1550 if !before_work() {
1551 return false;
1552 }
1553 cursor.after = Some(id);
1554 count += 1;
1555 visit(id);
1556 }
1557 if count == limit && cursor.after != Some(cursor.upper) {
1558 break;
1559 }
1560 cursor.key_index += 1;
1561 cursor.after = None;
1562 }
1563 true
1564 }
1565
1566 pub(crate) fn find_unique_key_conflict(
1571 &self,
1572 columns: &[Vec<String>],
1573 key_index: usize,
1574 signatures: &[String],
1575 exclude_id: Option<EntityId>,
1576 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1577 ) -> Option<EntityId> {
1578 assert!(
1579 key_index < columns.len(),
1580 "constraint index belongs to current schema"
1581 );
1582 let cached = self.unique_key_lookup.read();
1583 if cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1584 return self.find_unique_key_candidate(
1585 &cached[key_index],
1586 signatures,
1587 exclude_id,
1588 reserves_key,
1589 );
1590 }
1591 drop(cached);
1592 let mut cached = self.unique_key_lookup.write();
1593 if !cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1594 let mut indexes: Vec<_> = columns
1595 .iter()
1596 .map(|columns| SegmentUniqueKey::new(columns))
1597 .collect();
1598 self.for_each_fast(|entity| {
1599 for index in &mut indexes {
1600 index.insert(entity);
1601 }
1602 true
1603 });
1604 self.unique_key_lookup_bytes.store(
1605 indexes
1606 .iter()
1607 .map(|index| index.memory_bytes() as u64)
1608 .sum(),
1609 Ordering::Relaxed,
1610 );
1611 *cached = indexes;
1612 }
1613 self.find_unique_key_candidate(&cached[key_index], signatures, exclude_id, reserves_key)
1614 }
1615
1616 fn find_unique_key_candidate(
1617 &self,
1618 index: &SegmentUniqueKey,
1619 signatures: &[String],
1620 exclude_id: Option<EntityId>,
1621 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1622 ) -> Option<EntityId> {
1623 index.candidates(signatures).iter().copied().find(|id| {
1624 if Some(*id) == exclude_id {
1625 return false;
1626 }
1627 let Some(
1628 entity @ UnifiedEntity {
1629 data: EntityData::Row(row),
1630 ..
1631 },
1632 ) = self.get(*id)
1633 else {
1634 return false;
1635 };
1636 index
1637 .columns
1638 .iter()
1639 .zip(signatures)
1640 .all(|(column, signature)| {
1641 row.get_field(column)
1642 .is_some_and(|value| format!("{value:?}") == *signature)
1643 })
1644 && reserves_key(entity)
1645 })
1646 }
1647
1648 fn invalidate_unique_key_lookup(&mut self) {
1649 *self.unique_key_lookup.get_mut() = Vec::new();
1650 self.unique_key_lookup_bytes.store(0, Ordering::Relaxed);
1651 }
1652
1653 pub(crate) fn find_key_entity(
1656 &self,
1657 key: &str,
1658 newest: bool,
1659 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1660 ) -> Option<&UnifiedEntity> {
1661 let index = self
1664 .key_lookup
1665 .as_ref()
1666 .expect("segment owns its key index");
1667 self.find_key_candidate(index, key, newest, visible)
1668 }
1669
1670 fn find_key_candidate(
1671 &self,
1672 index: &SegmentKeyLookup,
1673 key: &str,
1674 newest: bool,
1675 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1676 ) -> Option<&UnifiedEntity> {
1677 self.select_key_entity(
1678 index
1679 .candidates(key)
1680 .iter()
1681 .chain(index.mutated_candidates())
1682 .filter_map(|id| self.get(*id)),
1683 key,
1684 newest,
1685 visible,
1686 )
1687 }
1688
1689 fn select_key_entity<'a>(
1690 &self,
1691 entities: impl Iterator<Item = &'a UnifiedEntity>,
1692 key: &str,
1693 newest: bool,
1694 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1695 ) -> Option<&'a UnifiedEntity> {
1696 let mut best: Option<&UnifiedEntity> = None;
1697 for entity in entities {
1698 if SegmentKeyLookup::key(entity) != Some(key) || !visible(entity) {
1700 continue;
1701 }
1702 if !newest {
1703 return Some(entity);
1704 }
1705 if best.is_none_or(|current| entity.xmin >= current.xmin) {
1706 best = Some(entity);
1707 }
1708 }
1709 best
1710 }
1711
1712 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1714 if let Some(id) = entity_id_from_probe_key(key) {
1715 return self.has_live_entity(id);
1716 }
1717
1718 let Ok(pk_str) = std::str::from_utf8(key) else {
1719 return false;
1720 };
1721 let pk_key = (self.collection.clone(), pk_str.to_string());
1722 self.pk_index
1723 .get(&pk_key)
1724 .is_some_and(|id| self.has_live_entity(*id))
1725 }
1726
1727 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1729 if let Some(index) = self.key_lookup.as_mut() {
1730 index.remove(entity);
1731 }
1732 if let Some(index) = self.graph_read_index.get_mut() {
1733 index.remove(entity);
1734 self.graph_read_index_bytes
1735 .store(index.memory_bytes(), Ordering::Relaxed);
1736 }
1737 let kind_key = entity.kind.storage_type().to_string();
1739 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1740 set.remove(&entity.id);
1741 }
1742
1743 if let EntityData::Row(row) = &entity.data {
1745 if let Some(first_col) = row.columns.first() {
1746 let pk_str = format!("{:?}", first_col);
1747 self.pk_index
1748 .remove(&(entity.kind.collection().to_string(), pk_str));
1749 }
1750 }
1751
1752 self.cross_ref_forward.remove(&entity.id);
1754 }
1756
1757 fn reindex_for_update(
1767 &mut self,
1768 old: &UpdateIndexSnapshot,
1769 new: &UnifiedEntity,
1770 modified_columns: Option<&[String]>,
1771 ) {
1772 let pk_changed = match &new.data {
1777 EntityData::Row(new_row) => {
1778 if let Some(cols) = modified_columns {
1779 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1782 new_row
1783 .schema
1784 .as_deref()
1785 .and_then(|schema| schema.first().map(|name| name.as_str()))
1786 });
1787 match pk_col_name {
1788 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1789 None => old.pk_value.as_ref() != new_row.columns.first(),
1791 }
1792 } else {
1793 old.pk_value.as_ref() != new_row.columns.first()
1794 }
1795 }
1796 _ => false,
1798 };
1799
1800 if pk_changed {
1801 if let Some((collection, pk_str)) = &old.pk_index_key {
1803 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1804 }
1805 if let EntityData::Row(row) = &new.data {
1807 if let Some(first_col) = row.columns.first() {
1808 let pk_str = format!("{:?}", first_col);
1809 self.pk_index
1810 .insert((new.kind.collection().to_string(), pk_str), new.id);
1811 }
1812 }
1813 }
1814
1815 let new_refs = new.cross_refs();
1817 if old.cross_refs.as_slice() != new_refs {
1818 self.cross_ref_forward.remove(&new.id);
1820 for cross_ref in &old.cross_refs {
1822 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1823 rev.retain(|(src, _)| *src != new.id);
1824 }
1825 }
1826 for cross_ref in new_refs {
1828 self.cross_ref_forward
1829 .entry(cross_ref.source)
1830 .or_default()
1831 .push((cross_ref.target, cross_ref.ref_type));
1832 self.cross_ref_reverse
1833 .entry(cross_ref.target)
1834 .or_default()
1835 .push((cross_ref.source, cross_ref.ref_type));
1836 }
1837 }
1838 }
1839
1840 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1842 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1843 }
1844
1845 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1847 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1848 }
1849
1850 pub fn age_secs(&self) -> u64 {
1852 let now = current_unix_secs();
1853 now.saturating_sub(self.created_at)
1854 }
1855
1856 pub fn idle_secs(&self) -> u64 {
1858 let now = current_unix_secs();
1859 now.saturating_sub(self.last_write_at)
1860 }
1861
1862 pub fn bulk_insert(
1871 &mut self,
1872 entities: Vec<UnifiedEntity>,
1873 ) -> Result<Vec<EntityId>, SegmentError> {
1874 self.mark_checkpoint_dirty();
1875 if !self.state.is_writable() {
1876 return Err(SegmentError::NotWritable);
1877 }
1878
1879 if let Some(index) = self.key_lookup.as_mut() {
1880 for entity in &entities {
1881 index.insert(entity);
1882 }
1883 }
1884
1885 let indexes = self.unique_key_lookup.get_mut();
1886 for index in indexes.iter_mut() {
1887 for entity in &entities {
1888 index.insert(entity);
1889 }
1890 }
1891 self.unique_key_lookup_bytes.store(
1892 indexes
1893 .iter()
1894 .map(|index| index.memory_bytes() as u64)
1895 .sum(),
1896 Ordering::Relaxed,
1897 );
1898
1899 let n = entities.len();
1900
1901 if entities.is_empty() {
1902 return Ok(Vec::new());
1903 }
1904
1905 for run in
1908 entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
1909 {
1910 let kind = run[0].kind.storage_type();
1911 self.observe_graph_kind(&run[0].kind);
1912 if matches!(
1913 run[0].kind,
1914 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
1915 ) {
1916 if let Some(index) = self.graph_read_index.get_mut() {
1917 for entity in run {
1918 index.insert(entity);
1919 }
1920 self.graph_read_index_bytes
1921 .store(index.memory_bytes(), Ordering::Relaxed);
1922 }
1923 }
1924 let kind_set = match self.kind_index.get_mut(kind) {
1925 Some(ids) => ids,
1926 None => self.kind_index.entry(kind.to_string()).or_default(),
1927 };
1928 kind_set.reserve(run.len());
1929 kind_set.extend(run.iter().map(|entity| entity.id));
1930 }
1931
1932 let now = current_unix_secs();
1933
1934 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
1935
1936 let mut ids = Vec::with_capacity(n);
1937
1938 if self.flat_entities.is_empty() && self.entities.is_empty() {
1940 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
1942 self.use_flat = true;
1943 }
1944
1945 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
1955 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
1956 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
1960
1961 let mut batch_bytes: usize = 0;
1966
1967 if self.use_flat {
1968 self.flat_entities.reserve(n);
1969 for (i, mut entity) in entities.into_iter().enumerate() {
1970 entity.sequence_id = base_seq + i as u64;
1971 let id = entity.id;
1972 ids.push(id);
1973 batch_bytes += Self::estimate_entity_size(&entity);
1974 if let EntityData::Row(row) = &entity.data {
1975 if row.schema.is_some() && !row.columns.is_empty() {
1976 if columnar_zone_updates.is_empty() {
1977 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
1978 columnar_schema = row.schema.clone();
1979 }
1980 for (ci, val) in row.columns.iter().enumerate() {
1981 if !matches!(val, Value::Null) {
1982 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
1983 bucket.push(val.clone());
1984 }
1985 }
1986 }
1987 } else {
1988 for (col, val) in row.iter_fields() {
1989 if !matches!(val, Value::Null) {
1990 named_zone_updates.push((col.to_string(), val.clone()));
1991 }
1992 }
1993 }
1994 }
1995 let expected = self.base_entity_id + self.flat_entities.len() as u64;
2011 if id.raw() == expected {
2012 self.flat_entities.push(entity);
2013 } else {
2014 self.hash_scan_ids.push(id);
2015 self.entities.insert(id, entity);
2016 }
2017 }
2018 } else {
2019 self.entities.reserve(n);
2021 self.hash_scan_ids.reserve(n);
2022 let mut pairs = Vec::with_capacity(n);
2023 for (i, mut entity) in entities.into_iter().enumerate() {
2024 entity.sequence_id = base_seq + i as u64;
2025 let id = entity.id;
2026 ids.push(id);
2027 batch_bytes += Self::estimate_entity_size(&entity);
2028 if let EntityData::Row(row) = &entity.data {
2029 for (col, val) in row.iter_fields() {
2030 if !matches!(val, Value::Null) {
2031 named_zone_updates.push((col.to_string(), val.clone()));
2032 }
2033 }
2034 }
2035 self.hash_scan_ids.push(id);
2036 pairs.push((id, entity));
2037 }
2038 self.entities.extend(pairs);
2039 }
2040
2041 if !columnar_zone_updates.is_empty() {
2045 let schema = columnar_schema.as_ref();
2046 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
2047 if values.is_empty() {
2048 continue;
2049 }
2050 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
2051 continue;
2052 };
2053 let mut iter = values.into_iter();
2059 if let Some(first) = iter.next() {
2060 let previous = self
2061 .col_zones
2062 .get(col_name)
2063 .map_or(0, ColZone::payload_bytes);
2064 let zone = self
2065 .col_zones
2066 .entry(col_name.clone())
2067 .and_modify(|z| z.update(&first))
2068 .or_insert_with(|| ColZone::new(first));
2069 for v in iter {
2070 zone.update(&v);
2071 }
2072 self.zone_payload_bytes = self
2073 .zone_payload_bytes
2074 .saturating_sub(previous)
2075 .saturating_add(zone.payload_bytes());
2076 }
2077 }
2078 }
2079 for (col, val) in named_zone_updates {
2080 self.update_column_zone(&col, &val);
2081 }
2082
2083 self.add_memory(batch_bytes);
2084 self.last_write_at = now;
2085
2086 if self.use_flat {
2088 self.published_flat_len
2089 .store(self.flat_entities.len(), Ordering::Release);
2090 }
2091
2092 Ok(ids)
2093 }
2094
2095 pub(crate) fn force_delete(&mut self, id: EntityId) -> bool {
2098 if self.use_flat {
2099 let raw = id.raw();
2100 if raw >= self.base_entity_id {
2101 let idx = (raw - self.base_entity_id) as usize;
2102 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2103 self.tombstone_flat_entity(idx, id);
2104 return true;
2105 }
2106 }
2107 return self.remove_hashmap_entity(id);
2108 }
2109
2110 self.remove_hashmap_entity(id)
2111 }
2112
2113 pub(crate) fn force_update_with_metadata(
2116 &mut self,
2117 entity: &UnifiedEntity,
2118 modified_columns: &[String],
2119 metadata: Option<&Metadata>,
2120 ) -> Result<(), SegmentError> {
2121 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
2122 }
2123}
2124
2125impl UnifiedSegment for GrowingSegment {
2126 fn id(&self) -> SegmentId {
2127 self.id
2128 }
2129
2130 fn state(&self) -> SegmentState {
2131 self.state
2132 }
2133
2134 fn collection(&self) -> &str {
2135 &self.collection
2136 }
2137
2138 fn stats(&self) -> SegmentStats {
2139 let mut stats = SegmentStats {
2140 entity_count: self.live_entity_count(),
2141 deleted_count: self.deleted.len(),
2142 memory_bytes: self.resident_bytes() as usize,
2143 ..Default::default()
2144 };
2145
2146 for entity in self.entities.values() {
2147 match &entity.kind {
2148 EntityKind::TableRow { .. } => stats.row_count += 1,
2149 EntityKind::GraphNode(_) => stats.node_count += 1,
2150 EntityKind::GraphEdge(_) => stats.edge_count += 1,
2151 EntityKind::Vector { .. } => stats.vector_count += 1,
2152 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
2153 EntityKind::QueueMessage { .. } => stats.row_count += 1,
2154 }
2155 stats.cross_ref_count += entity.cross_refs().len();
2156 }
2157
2158 stats
2159 }
2160
2161 fn entity_count(&self) -> usize {
2162 self.live_entity_count()
2163 }
2164
2165 fn contains(&self, id: EntityId) -> bool {
2166 self.has_live_entity(id)
2167 }
2168
2169 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2170 if self.deleted.contains_key(&id) {
2171 return None;
2172 }
2173 if self.use_flat {
2174 let raw = id.raw();
2175 if raw >= self.base_entity_id {
2176 let idx = (raw - self.base_entity_id) as usize;
2177 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
2178 return Some(entity);
2179 }
2180 }
2181 self.entities.get(&id)
2187 } else {
2188 self.entities.get(&id)
2189 }
2190 }
2191
2192 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
2193 self.mark_checkpoint_entity_dirty(id);
2194 self.invalidate_unique_key_lookup();
2195 self.invalidate_graph_read_index();
2196 if self.deleted.contains_key(&id) || !self.state.is_writable() {
2197 return None;
2198 }
2199 if let Some(mut index) = self.key_lookup.take() {
2200 if let Some(entity) = self.get(id) {
2201 index.mark_mutated(entity);
2202 }
2203 self.key_lookup = Some(index);
2204 }
2205 self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
2208 if self.use_flat {
2209 let raw = id.raw();
2210 if raw >= self.base_entity_id {
2211 let idx = (raw - self.base_entity_id) as usize;
2212 if self
2213 .flat_entities
2214 .get(idx)
2215 .map(|e| e.id == id)
2216 .unwrap_or(false)
2217 {
2218 return self.flat_entities.get_mut(idx);
2219 }
2220 }
2221 self.entities.get_mut(&id)
2224 } else {
2225 self.entities.get_mut(&id)
2226 }
2227 }
2228
2229 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2230 if !self.state.is_writable() {
2231 return Err(SegmentError::NotWritable);
2232 }
2233
2234 if self.entities.contains_key(&entity.id) {
2235 return Err(SegmentError::AlreadyExists(entity.id));
2236 }
2237
2238 entity.sequence_id = self.next_sequence();
2240
2241 let size = Self::estimate_entity_size(&entity);
2243 self.add_memory(size);
2244
2245 self.index_entity(&entity);
2247
2248 self.update_col_zones_from_entity(&entity);
2250
2251 let id = entity.id;
2253 self.entities.insert(id, entity);
2254
2255 self.last_write_at = current_unix_secs();
2257
2258 Ok(id)
2259 }
2260
2261 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2262 if !self.state.is_writable() {
2263 return Err(SegmentError::NotWritable);
2264 }
2265
2266 self.apply_update_with_metadata(&entity, None)?;
2267 self.last_write_at = current_unix_secs();
2268
2269 Ok(())
2270 }
2271
2272 fn update_hot(
2273 &mut self,
2274 entity: UnifiedEntity,
2275 modified_columns: &[String],
2276 ) -> Result<(), SegmentError> {
2277 if !self.state.is_writable() {
2278 return Err(SegmentError::NotWritable);
2279 }
2280
2281 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
2282 self.last_write_at = current_unix_secs();
2283 Ok(())
2284 }
2285
2286 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2287 if !self.state.is_writable() {
2288 return Err(SegmentError::NotWritable);
2289 }
2290
2291 if self.use_flat {
2293 let raw = id.raw();
2294 if raw >= self.base_entity_id {
2295 let idx = (raw - self.base_entity_id) as usize;
2296 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2297 self.tombstone_flat_entity(idx, id);
2298 return Ok(true);
2299 }
2300 }
2301 return Ok(self.remove_hashmap_entity(id));
2302 }
2303
2304 Ok(self.remove_hashmap_entity(id))
2305 }
2306
2307 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2308 if !self.has_live_entity(id) {
2309 return None;
2310 }
2311 Some(self.metadata.get_all(id))
2312 }
2313
2314 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2315 self.mark_checkpoint_entity_dirty(id);
2316 if !self.state.is_writable() {
2317 return Err(SegmentError::NotWritable);
2318 }
2319
2320 if !self.has_live_entity(id) {
2321 return Err(SegmentError::NotFound(id));
2322 }
2323
2324 self.metadata.set_all(id, &metadata);
2325 Ok(())
2326 }
2327
2328 fn seal(&mut self) -> Result<(), SegmentError> {
2329 if self.state != SegmentState::Growing {
2330 return Err(SegmentError::InvalidState(self.state));
2331 }
2332
2333 self.state = SegmentState::Sealing;
2334
2335 self.rebuild_sealed_col_zones();
2340
2341 self.state = SegmentState::Sealed;
2342 Ok(())
2343 }
2344
2345 fn should_seal(&self, config: &SegmentConfig) -> bool {
2346 if self.entities.len() >= config.max_entities {
2348 return true;
2349 }
2350
2351 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
2353 return true;
2354 }
2355
2356 if self.age_secs() >= config.max_age_secs {
2358 return true;
2359 }
2360
2361 false
2362 }
2363
2364 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2365 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2366 Box::new(self.flat_entities.iter().chain(self.entities.values()))
2370 } else {
2371 Box::new(self.entities.values())
2372 };
2373 if self.deleted.is_empty() {
2374 base
2375 } else {
2376 Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
2377 }
2378 }
2379
2380 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2381 let ids = self.kind_index.get(kind_filter);
2384 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2387 Box::new(self.flat_entities.iter())
2388 } else {
2389 Box::new(std::iter::empty())
2390 };
2391 Box::new(flat.chain(self.entities.values()).filter(move |e| {
2392 if self.deleted.contains_key(&e.id) {
2393 return false;
2394 }
2395 if let Some(ids) = ids {
2396 ids.contains(&e.id)
2397 } else {
2398 false
2399 }
2400 }))
2401 }
2402
2403 fn filter_metadata(
2404 &self,
2405 filters: &[(String, super::metadata::MetadataFilter)],
2406 ) -> Vec<EntityId> {
2407 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
2410 Box::new(self.flat_entities.iter().map(|e| e.id))
2411 } else {
2412 Box::new(std::iter::empty())
2413 };
2414 flat_ids
2415 .chain(self.entities.keys().copied())
2416 .filter(|id| {
2417 if self.deleted.contains_key(id) {
2418 return false;
2419 }
2420 let metadata = self.metadata.get_all(*id);
2421 metadata.matches_all(filters)
2422 })
2423 .collect()
2424 }
2425}
2426
2427fn build_minmax_multi_intervals(
2428 entries: &[(CanonicalKey, Value)],
2429 max_intervals: usize,
2430) -> Vec<ColZone> {
2431 if entries.is_empty() {
2432 return Vec::new();
2433 }
2434 if entries.len() == 1 || max_intervals <= 1 {
2435 return vec![ColZone::with_bounds(
2436 entries[0].1.clone(),
2437 entries[entries.len() - 1].1.clone(),
2438 )];
2439 }
2440
2441 let mut split_points = if entries.len() <= max_intervals {
2442 (1..entries.len()).collect::<Vec<_>>()
2443 } else {
2444 let target_splits = max_intervals - 1;
2445 let mut selected = select_gap_split_points(entries, target_splits);
2446 if selected.len() < target_splits {
2447 for bucket in 1..max_intervals {
2448 let idx = bucket * entries.len() / max_intervals;
2449 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
2450 continue;
2451 }
2452 selected.push(idx);
2453 if selected.len() >= target_splits {
2454 break;
2455 }
2456 }
2457 }
2458 selected.sort_unstable();
2459 selected.dedup();
2460 selected
2461 };
2462
2463 split_points.push(entries.len());
2464
2465 let mut out = Vec::with_capacity(split_points.len());
2466 let mut start = 0usize;
2467 for end in split_points {
2468 if end <= start {
2469 continue;
2470 }
2471 out.push(ColZone::with_bounds(
2472 entries[start].1.clone(),
2473 entries[end - 1].1.clone(),
2474 ));
2475 start = end;
2476 }
2477
2478 if out.is_empty() {
2479 out.push(ColZone::with_bounds(
2480 entries[0].1.clone(),
2481 entries[entries.len() - 1].1.clone(),
2482 ));
2483 }
2484
2485 out
2486}
2487
2488fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
2489 let mut gaps = Vec::new();
2490 for idx in 1..entries.len() {
2491 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
2492 if score > 0.0 {
2493 gaps.push((score, idx));
2494 }
2495 }
2496 }
2497 gaps.sort_by(|left, right| {
2498 right
2499 .0
2500 .partial_cmp(&left.0)
2501 .unwrap_or(std::cmp::Ordering::Equal)
2502 .then_with(|| left.1.cmp(&right.1))
2503 });
2504 gaps.into_iter()
2505 .take(max_splits)
2506 .map(|(_, idx)| idx)
2507 .collect()
2508}
2509
2510fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
2511 if left.family() != right.family() {
2512 return None;
2513 }
2514 match (left, right) {
2515 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
2516 Some(r.saturating_sub(*l) as f64)
2517 }
2518 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
2519 Some(r.saturating_sub(*l) as f64)
2520 }
2521 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
2522 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
2523 }
2524 _ => None,
2525 }
2526}
2527
2528#[cfg(test)]
2529mod tests {
2530 use super::*;
2531 use crate::storage::unified::entity::RowData;
2532 use crate::storage::unified::MetadataValue;
2533 use reddb_types::Value;
2534
2535 fn keyed_row(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
2536 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
2537 entity.data = EntityData::Row(RowData::with_names(
2538 vec![Value::text(key), Value::Integer(id as i64)],
2539 vec!["key".into(), "value".into()],
2540 ));
2541 entity.set_xmin(xmin);
2542 entity.set_xmax(xmax);
2543 entity
2544 }
2545
2546 fn keyed_probe(segment: &GrowingSegment, key: &str) -> Option<EntityId> {
2547 segment
2548 .find_key_entity(key, true, &mut |_| true)
2549 .map(|entity| entity.id)
2550 }
2551
2552 #[test]
2553 fn key_lookup_visits_only_requested_key_versions() {
2554 let mut segment = GrowingSegment::new(1, "keys");
2555 segment
2556 .bulk_insert(
2557 (1..=4096)
2558 .map(|id| keyed_row(id, &format!("other-{id}"), 1, 0))
2559 .collect(),
2560 )
2561 .expect("unrelated rows");
2562 segment
2563 .insert(keyed_row(5000, "wanted", 2, 8))
2564 .expect("old version");
2565 segment
2566 .insert(keyed_row(5001, "wanted", 8, 0))
2567 .expect("new version");
2568 let mut visits = 0;
2569 let result = segment.find_key_entity("wanted", true, &mut |entity| {
2570 visits += 1;
2571 entity.xmin < 8
2572 });
2573 assert_eq!(result.map(|entity| entity.id), Some(EntityId::new(5000)));
2574 assert_eq!(
2575 visits, 2,
2576 "unrelated rows must never reach visibility evaluation"
2577 );
2578 assert_eq!(keyed_probe(&segment, "wanted"), Some(EntityId::new(5001)));
2579 let mut misses = 0;
2580 assert!(segment
2581 .find_key_entity("absent", true, &mut |_| {
2582 misses += 1;
2583 true
2584 })
2585 .is_none());
2586 assert_eq!(misses, 0);
2587
2588 let mut collision = SegmentKeyLookup::default();
2590 collision.insert(&keyed_row(1, "wanted", 1, 0));
2591 assert!(segment
2592 .find_key_candidate(&collision, "wanted", true, &mut |_| true)
2593 .is_none());
2594 }
2595
2596 #[test]
2597 fn raw_mutation_does_not_disable_unrelated_key_candidates() {
2598 let mut segment = GrowingSegment::new(1, "keys");
2599 segment
2600 .bulk_insert(
2601 (1..=4096)
2602 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2603 .collect(),
2604 )
2605 .expect("rows");
2606 for _ in 0..3 {
2607 *segment.get_mut(EntityId::new(1)).expect("mutable") = keyed_row(1, "renamed", 2, 0);
2608 }
2609 let index = segment.key_lookup.as_ref().expect("index retained");
2610 assert_eq!(index.candidates("key-4096"), &[EntityId::new(4096)]);
2611 assert_eq!(index.mutated_candidates().count(), 1);
2612 assert!(index.candidates("key-1").is_empty());
2613 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(1)));
2614 assert_eq!(keyed_probe(&segment, "key-4096"), Some(EntityId::new(4096)));
2615 segment.delete(EntityId::new(1)).expect("delete raw row");
2616 assert_eq!(
2617 segment
2618 .key_lookup
2619 .as_ref()
2620 .unwrap()
2621 .mutated_candidates()
2622 .count(),
2623 0
2624 );
2625 assert_eq!(keyed_probe(&segment, "renamed"), None);
2626 }
2627
2628 #[test]
2629 fn key_lookup_tracks_mutations_sealing_and_adoption() {
2630 let mut segment = GrowingSegment::new(1, "keys");
2631 segment
2632 .bulk_insert(vec![keyed_row(1, "flat", 1, 0)])
2633 .expect("flat row");
2634 segment
2635 .insert(keyed_row(3, "hash", 1, 0))
2636 .expect("hash fallback");
2637 segment
2638 .bulk_insert(vec![keyed_row(5, "gap", 1, 0)])
2639 .expect("bulk gap");
2640 assert!(segment.key_lookup.as_ref().expect("index").memory_bytes() > 0);
2641 assert_eq!(keyed_probe(&segment, "gap"), Some(EntityId::new(5)));
2642 segment
2643 .update(keyed_row(3, "renamed", 2, 0))
2644 .expect("rename key");
2645 assert!(
2646 segment.key_lookup.is_some(),
2647 "key replacements maintain the index"
2648 );
2649 assert_eq!(keyed_probe(&segment, "hash"), None);
2650 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(3)));
2651
2652 *segment.get_mut(EntityId::new(3)).expect("mutable row") = keyed_row(3, "raw", 3, 0);
2653 assert!(
2654 segment.key_lookup.is_some(),
2655 "raw mutation preserves unaffected key candidates"
2656 );
2657 assert_eq!(keyed_probe(&segment, "renamed"), None);
2658 assert_eq!(keyed_probe(&segment, "raw"), Some(EntityId::new(3)));
2659 segment.seal().expect("seal");
2660 segment
2661 .force_update_with_metadata(&keyed_row(5, "sealed", 4, 0), &["key".into()], None)
2662 .expect("sealed rename");
2663 assert_eq!(keyed_probe(&segment, "gap"), None);
2664 assert_eq!(keyed_probe(&segment, "sealed"), Some(EntityId::new(5)));
2665 assert!(segment.force_delete(EntityId::new(1)));
2666 assert!(segment.force_delete(EntityId::new(3)));
2667 assert_eq!(keyed_probe(&segment, "flat"), None);
2668 assert_eq!(keyed_probe(&segment, "raw"), None);
2669
2670 let mut adopted = GrowingSegment::new(2, "keys");
2671 segment.for_each_fast(|entity| {
2672 adopted.adopt_entity(entity.clone(), None);
2673 true
2674 });
2675 assert_eq!(keyed_probe(&adopted, "sealed"), Some(EntityId::new(5)));
2676 assert!(adopted.evict_entity(EntityId::new(5)));
2677 assert_eq!(keyed_probe(&adopted, "sealed"), None);
2678 }
2679
2680 #[test]
2681 #[ignore = "diagnostic microbenchmark; no timing assertions"]
2682 fn key_lookup_scaling_diagnostic() {
2683 use std::{hint::black_box, time::Instant};
2684 for count in [100, 10_000, 100_000] {
2685 let mut segment = GrowingSegment::new(1, "keys");
2686 segment
2687 .bulk_insert(
2688 (1..=count)
2689 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2690 .collect(),
2691 )
2692 .expect("rows");
2693 let key = format!("key-{count}");
2694 let started = Instant::now();
2695 for _ in 0..1000 {
2696 assert_eq!(
2697 keyed_probe(black_box(&segment), black_box(&key)),
2698 Some(EntityId::new(count))
2699 );
2700 }
2701 let indexed = started.elapsed();
2702 let started = Instant::now();
2703 for _ in 0..1000 {
2704 let mut found = None;
2705 segment.for_each_fast(|entity| {
2706 if SegmentKeyLookup::key(entity) == Some(black_box(key.as_str())) {
2707 found = Some(entity.id);
2708 return false;
2709 }
2710 true
2711 });
2712 assert_eq!(black_box(found), Some(EntityId::new(count)));
2713 }
2714 eprintln!(
2715 "rows={count} reads=1000 indexed={indexed:?} scan={:?}",
2716 started.elapsed()
2717 );
2718 }
2719 }
2720
2721 fn pruning_node(id: u64) -> UnifiedEntity {
2722 UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
2723 }
2724
2725 fn pruning_edge(id: u64) -> UnifiedEntity {
2726 UnifiedEntity::graph_edge(
2727 EntityId::new(id),
2728 "mixed",
2729 "a",
2730 "b",
2731 1.0,
2732 Default::default(),
2733 )
2734 }
2735
2736 #[test]
2737 fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
2738 let mut segment = GrowingSegment::new(1, "mixed");
2739 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
2740 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
2741 segment
2742 .bulk_insert(vec![
2743 UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
2744 pruning_node(2),
2745 pruning_edge(4),
2746 ])
2747 .expect("mixed bulk with gap");
2748 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2749 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2750 segment.seal().expect("seal");
2751 assert!(segment.force_delete(EntityId::new(2)));
2752 assert!(segment.force_delete(EntityId::new(4)));
2753 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2755 let mut rebuilt = GrowingSegment::new(2, "mixed");
2756 for item in segment.iter() {
2757 rebuilt.adopt_entity(item.clone(), None);
2758 }
2759 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2760 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2761 rebuilt.adopt_entity(pruning_node(5), None);
2762 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2763 rebuilt.insert(pruning_edge(6)).expect("single edge");
2764 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2765 }
2766
2767 #[test]
2768 fn graph_presence_covers_replacement_and_unrestricted_mutation() {
2769 for flat in [false, true] {
2770 for operation in 0..4 {
2771 let mut segment = GrowingSegment::new(1, "mixed");
2772 let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
2773 if flat {
2774 segment.bulk_insert(vec![vector]).expect("flat vector");
2775 } else {
2776 segment.insert(vector).expect("hashmap vector");
2777 }
2778 match operation {
2779 0 => segment.update(pruning_node(1)).expect("replace with node"),
2780 1 => segment
2781 .update_hot(pruning_node(1), &[])
2782 .expect("hot replacement"),
2783 2 => {
2784 segment.seal().expect("seal vector");
2785 segment
2786 .force_update_with_metadata(&pruning_node(1), &[], None)
2787 .expect("sealed replacement");
2788 }
2789 _ => {
2790 *segment.get_mut(EntityId::new(1)).expect("mutable vector") =
2791 pruning_node(1)
2792 }
2793 }
2794 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2795 assert!(matches!(
2796 segment.get(EntityId::new(1)).expect("node").kind,
2797 EntityKind::GraphNode(_)
2798 ));
2799 if operation == 3 {
2800 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2801 }
2802 segment
2803 .force_update_with_metadata(&pruning_edge(1), &[], None)
2804 .expect("replace with edge");
2805 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2806 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2807 }
2808 }
2809 }
2810
2811 fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
2812 UnifiedEntity::new(
2813 EntityId::new(id),
2814 EntityKind::TableRow {
2815 table: "keys".into(),
2816 row_id: id,
2817 },
2818 EntityData::Row(RowData::with_names(
2819 vec![Value::text("payload"), key],
2820 vec!["payload".into(), "primary".into()],
2821 )),
2822 )
2823 }
2824
2825 fn primary_key_probe(
2826 segment: &GrowingSegment,
2827 key: Value,
2828 exclude: Option<EntityId>,
2829 ) -> Option<EntityId> {
2830 segment.find_unique_key_conflict(
2831 &[vec!["primary".into()]],
2832 0,
2833 &[format!("{key:?}")],
2834 exclude,
2835 &|_| true,
2836 )
2837 }
2838
2839 #[test]
2840 fn unique_key_lookup_follows_segment_mutations_and_sealing() {
2841 let mut segment = GrowingSegment::new(1, "keys");
2842 segment
2843 .bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
2844 .expect("bulk row");
2845 let payload_bytes = segment.memory_bytes();
2846 assert_eq!(
2847 primary_key_probe(&segment, Value::Integer(10), None),
2848 Some(EntityId::new(1))
2849 );
2850 assert!(segment.memory_bytes() > payload_bytes);
2851 segment
2852 .insert(primary_key_row(3, Value::Integer(30)))
2853 .expect("hashmap fallback row");
2854 segment
2855 .bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
2856 .expect("bulk gap row");
2857 assert_eq!(
2858 primary_key_probe(&segment, Value::Integer(40), None),
2859 Some(EntityId::new(4))
2860 );
2861 assert_eq!(
2862 primary_key_probe(&segment, Value::Integer(30), None),
2863 Some(EntityId::new(3))
2864 );
2865 segment
2866 .update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
2867 .expect("replace flat key");
2868 assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
2869 assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
2870 assert_eq!(
2871 primary_key_probe(&segment, Value::Integer(11), None),
2872 Some(EntityId::new(1))
2873 );
2874 segment
2875 .delete_batch(&[EntityId::new(1), EntityId::new(3)])
2876 .expect("delete flat and map");
2877 assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
2878 assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
2879 segment.seal().expect("seal");
2880 assert_eq!(
2881 primary_key_probe(&segment, Value::Integer(40), None),
2882 Some(EntityId::new(4))
2883 );
2884 segment
2885 .force_update_with_metadata(
2886 &primary_key_row(4, Value::Integer(41)),
2887 &["primary".into()],
2888 None,
2889 )
2890 .expect("sealed update");
2891 assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
2892 assert_eq!(
2893 primary_key_probe(&segment, Value::Integer(41), None),
2894 Some(EntityId::new(4))
2895 );
2896 assert!(segment.force_delete(EntityId::new(4)));
2897 assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
2898 }
2899
2900 #[test]
2901 fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
2902 let mut segment = GrowingSegment::new(1, "keys");
2903 segment
2904 .bulk_insert(vec![
2905 primary_key_row(1, Value::Integer(10)),
2906 primary_key_row(2, Value::Integer(20)),
2907 ])
2908 .expect("initial rows");
2909 let definitions = vec![
2910 vec!["primary".into()],
2911 vec!["payload".into()],
2912 vec!["primary".into(), "payload".into()],
2913 ];
2914 let payload = format!("{:?}", Value::text("payload"));
2915 assert_eq!(
2916 segment.find_unique_key_conflict(
2917 &definitions,
2918 1,
2919 &[payload.clone()],
2920 None,
2921 &|entity| entity.id != EntityId::new(1)
2922 ),
2923 Some(EntityId::new(2))
2924 );
2925 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2926 segment
2927 .insert(primary_key_row(3, Value::Integer(30)))
2928 .expect("append");
2929 assert_eq!(
2930 segment.find_unique_key_conflict(
2931 &definitions,
2932 2,
2933 &[format!("{:?}", Value::Integer(30)), payload.clone()],
2934 None,
2935 &|_| true
2936 ),
2937 Some(EntityId::new(3))
2938 );
2939 let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
2940 assert_eq!(
2941 segment.find_unique_key_conflict(
2942 &definitions,
2943 0,
2944 &[format!("{:?}", Value::Integer(20))],
2945 None,
2946 &|_| true
2947 ),
2948 Some(EntityId::new(2))
2949 );
2950 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2951 assert_eq!(
2952 segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
2953 bytes
2954 );
2955 let reduced = vec![vec!["payload".into()]];
2956 assert_eq!(
2957 segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
2958 Some(EntityId::new(1))
2959 );
2960 assert_eq!(segment.unique_key_lookup.read().len(), 1);
2961 assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
2962 }
2963
2964 #[test]
2965 fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
2966 let mut segment = GrowingSegment::new(1, "keys");
2967 let values = [
2968 Value::Float(f64::NAN),
2969 Value::Float(-0.0),
2970 Value::Float(0.0),
2971 Value::Integer(1),
2972 Value::UnsignedInteger(1),
2973 Value::text("1"),
2974 ];
2975 for (offset, value) in values.iter().enumerate() {
2976 let id = offset as u64 + 1;
2977 segment
2978 .insert(primary_key_row(id, value.clone()))
2979 .expect("distinct physical row");
2980 assert_eq!(
2981 primary_key_probe(&segment, value.clone(), None),
2982 Some(EntityId::new(id))
2983 );
2984 assert_eq!(
2985 primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
2986 None
2987 );
2988 }
2989 segment
2990 .insert(primary_key_row(20, Value::Integer(1)))
2991 .expect("duplicate physical key");
2992 assert_eq!(
2993 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
2994 Some(EntityId::new(20))
2995 );
2996 segment
2999 .entities
3000 .insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
3001 assert_eq!(
3002 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
3003 None
3004 );
3005 let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
3007 *entity = primary_key_row(4, Value::Integer(100));
3008 assert_eq!(
3009 primary_key_probe(&segment, Value::Integer(100), None),
3010 Some(EntityId::new(4))
3011 );
3012 }
3013
3014 #[test]
3015 fn test_growing_segment_basic() {
3016 let mut segment = GrowingSegment::new(1, "test");
3017
3018 let entity = UnifiedEntity::table_row(
3019 EntityId::new(1),
3020 "users",
3021 1,
3022 vec![Value::text("Alice".to_string())],
3023 );
3024
3025 let id = segment.insert(entity).unwrap();
3026 assert_eq!(id, EntityId::new(1));
3027 assert!(segment.contains(id));
3028
3029 let stats = segment.stats();
3030 assert_eq!(stats.entity_count, 1);
3031 assert_eq!(stats.row_count, 1);
3032 }
3033
3034 #[test]
3035 fn test_segment_metadata() {
3036 let mut segment = GrowingSegment::new(1, "test");
3037
3038 let entity = UnifiedEntity::table_row(
3039 EntityId::new(1),
3040 "users",
3041 1,
3042 vec![Value::text("Alice".to_string())],
3043 );
3044 segment.insert(entity).unwrap();
3045
3046 let mut meta = Metadata::new();
3047 meta.set("role", MetadataValue::String("admin".to_string()));
3048 meta.set("level", MetadataValue::Int(5));
3049
3050 segment.set_metadata(EntityId::new(1), meta).unwrap();
3051
3052 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
3053 assert_eq!(
3054 retrieved.get("role"),
3055 Some(&MetadataValue::String("admin".to_string()))
3056 );
3057 }
3058
3059 #[test]
3060 fn test_segment_seal() {
3061 let mut segment = GrowingSegment::new(1, "test");
3062
3063 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
3064 segment.insert(entity).unwrap();
3065
3066 assert!(segment.state().is_writable());
3068
3069 segment.seal().unwrap();
3071 assert_eq!(segment.state(), SegmentState::Sealed);
3072
3073 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
3075 assert!(segment.insert(entity2).is_err());
3076 }
3077
3078 #[test]
3079 fn test_should_seal() {
3080 let mut segment = GrowingSegment::new(1, "test");
3081
3082 let config = SegmentConfig {
3083 max_entities: 2,
3084 ..Default::default()
3085 };
3086
3087 assert!(!segment.should_seal(&config));
3088
3089 segment
3090 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
3091 .unwrap();
3092 assert!(!segment.should_seal(&config));
3093
3094 segment
3095 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
3096 .unwrap();
3097 assert!(segment.should_seal(&config));
3098 }
3099
3100 #[test]
3101 fn test_cross_references() {
3102 let mut segment = GrowingSegment::new(1, "test");
3103
3104 let mut entity1 = UnifiedEntity::table_row(
3105 EntityId::new(1),
3106 "hosts",
3107 1,
3108 vec![Value::text("192.168.1.1".to_string())],
3109 );
3110 entity1.add_cross_ref(CrossRef::new(
3111 EntityId::new(1),
3112 EntityId::new(2),
3113 "nodes",
3114 RefType::RowToNode,
3115 ));
3116 segment.insert(entity1).unwrap();
3117
3118 let refs_from = segment.get_references_from(EntityId::new(1));
3119 assert_eq!(refs_from.len(), 1);
3120 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
3121
3122 let refs_to = segment.get_references_to(EntityId::new(2));
3123 assert_eq!(refs_to.len(), 1);
3124 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
3125 }
3126
3127 #[test]
3128 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
3129 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
3130 zone.update(&Value::Email("delta@example.com".to_string()));
3131
3132 let probe = Value::Email("alpha@example.com".to_string());
3133 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
3134
3135 let in_range = Value::Email("charlie@example.com".to_string());
3136 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
3137 }
3138
3139 #[test]
3140 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
3141 let mut segment = GrowingSegment::new(1, "test");
3142
3143 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
3144 let entity = UnifiedEntity::new(
3145 EntityId::new(row_id),
3146 EntityKind::TableRow {
3147 table: "users".into(),
3148 row_id,
3149 },
3150 EntityData::Row(RowData::with_names(
3151 vec![Value::Integer(age)],
3152 vec!["age".to_string()],
3153 )),
3154 );
3155 segment.insert(entity).unwrap();
3156 }
3157
3158 segment.seal().unwrap();
3159
3160 let miss = Value::Integer(500);
3161 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
3162
3163 let hit = Value::Integer(1000);
3164 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
3165 }
3166}