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 adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
1146 let id = entity.id;
1147
1148 let next = entity.sequence_id.saturating_add(1);
1151 if self.sequence.load(Ordering::Relaxed) < next {
1152 self.sequence.store(next, Ordering::Relaxed);
1153 }
1154
1155 self.add_memory(Self::estimate_entity_size(&entity));
1156 self.index_entity(&entity);
1157 self.update_col_zones_from_entity(&entity);
1158 self.entities.insert(id, entity);
1159
1160 if let Some(metadata) = metadata {
1161 if !metadata.is_empty() {
1162 self.metadata.set_all(id, &metadata);
1163 }
1164 }
1165 }
1166
1167 pub(crate) fn evict_entity(&mut self, id: EntityId) -> bool {
1173 self.mark_checkpoint_entity_dirty(id);
1174 self.invalidate_unique_key_lookup();
1175 let Some(entity) = self.entities.remove(&id) else {
1176 return false;
1177 };
1178 self.release_memory(Self::estimate_entity_size(&entity));
1179 self.unindex_entity(&entity);
1180 for cross_ref in entity.cross_refs() {
1181 if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1182 reverse.retain(|(source, _)| *source != id);
1183 }
1184 }
1185 self.metadata.remove_all(id);
1186 true
1187 }
1188
1189 fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
1191 super::memory_size::entity_bytes(entity)
1192 }
1193
1194 fn update_column_zone(&mut self, column: &str, value: &Value) {
1202 if matches!(value, Value::Null) {
1203 return;
1204 }
1205 match self.col_zones.entry(column.to_string()) {
1206 std::collections::hash_map::Entry::Occupied(mut entry) => {
1207 let before = entry.get().payload_bytes();
1208 entry.get_mut().update(value);
1209 self.zone_payload_bytes = self
1210 .zone_payload_bytes
1211 .saturating_sub(before)
1212 .saturating_add(entry.get().payload_bytes());
1213 }
1214 std::collections::hash_map::Entry::Vacant(entry) => {
1215 let zone = ColZone::new(value.clone());
1216 self.zone_payload_bytes =
1217 self.zone_payload_bytes.saturating_add(zone.payload_bytes());
1218 entry.insert(zone);
1219 }
1220 }
1221 }
1222
1223 fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
1224 if let EntityData::Row(row) = &entity.data {
1225 if let Some(named) = &row.named {
1226 for (column, value) in named {
1227 self.update_column_zone(column, value);
1228 }
1229 } else if let Some(schema) = &row.schema {
1230 for (column, value) in schema.iter().zip(&row.columns) {
1231 self.update_column_zone(column, value);
1232 }
1233 }
1234 }
1235 }
1236
1237 fn rebuild_sealed_col_zones(&mut self) {
1238 let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
1239 let mut family_by_col: HashMap<String, reddb_types::CanonicalKeyFamily> = HashMap::new();
1240 let mut mixed_family_cols = HashSet::new();
1241 let mut unsupported_cols = HashSet::new();
1242
1243 let mut observe_row = |row: &super::entity::RowData| {
1244 for (col, value) in row.iter_fields() {
1245 if matches!(value, Value::Null) {
1246 continue;
1247 }
1248 let Some(key) = value_to_canonical_key(value) else {
1249 unsupported_cols.insert(col.to_string());
1250 continue;
1251 };
1252 match family_by_col.get(col).copied() {
1253 Some(existing) if existing != key.family() => {
1254 mixed_family_cols.insert(col.to_string());
1255 }
1256 None => {
1257 family_by_col.insert(col.to_string(), key.family());
1258 }
1259 _ => {}
1260 }
1261 values_by_col
1262 .entry(col.to_string())
1263 .or_default()
1264 .push((key, value.clone()));
1265 }
1266 };
1267
1268 if self.use_flat {
1269 for entity in &self.flat_entities {
1270 if self.deleted.contains_key(&entity.id) {
1271 continue;
1272 }
1273 if let EntityData::Row(row) = &entity.data {
1274 observe_row(row);
1275 }
1276 }
1277 } else {
1278 for entity in self.entities.values() {
1279 if self.deleted.contains_key(&entity.id) {
1280 continue;
1281 }
1282 if let EntityData::Row(row) = &entity.data {
1283 observe_row(row);
1284 }
1285 }
1286 }
1287
1288 let mut sealed_col_zones = HashMap::new();
1289 for (col, mut entries) in values_by_col {
1290 if mixed_family_cols.contains(&col)
1291 || unsupported_cols.contains(&col)
1292 || entries.is_empty()
1293 {
1294 continue;
1295 }
1296 entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
1297 entries.dedup_by(|left, right| left.0 == right.0);
1298
1299 let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
1300 if intervals.len() > 1 {
1301 sealed_col_zones.insert(col, MultiColZone { intervals });
1302 }
1303 }
1304
1305 let previous = self
1306 .sealed_col_zones
1307 .values()
1308 .flat_map(|zones| &zones.intervals)
1309 .map(ColZone::payload_bytes)
1310 .fold(0, usize::saturating_add);
1311 let current = sealed_col_zones
1312 .values()
1313 .flat_map(|zones| &zones.intervals)
1314 .map(ColZone::payload_bytes)
1315 .fold(0, usize::saturating_add);
1316 self.zone_payload_bytes = self
1317 .zone_payload_bytes
1318 .saturating_sub(previous)
1319 .saturating_add(current);
1320 self.sealed_col_zones = sealed_col_zones;
1321 }
1322
1323 pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
1327 if preds.is_empty() {
1328 return false;
1329 }
1330 for (col, pred) in preds {
1331 if let Some(zone) = self.sealed_col_zones.get(*col) {
1332 if zone.can_skip(pred) {
1333 return true;
1334 }
1335 continue;
1336 }
1337 if let Some(zone) = self.col_zones.get(*col) {
1338 if pred.can_skip(zone) {
1339 return true; }
1341 }
1342 }
1343 false
1344 }
1345
1346 fn index_entity(&mut self, entity: &UnifiedEntity) {
1348 self.mark_checkpoint_entity_dirty(entity.id);
1349 if let Some(index) = self.key_lookup.as_mut() {
1350 index.insert(entity);
1351 }
1352 self.hash_scan_ids.push(entity.id);
1353 self.observe_graph_kind(&entity.kind);
1354 if let Some(index) = self.graph_read_index.get_mut() {
1355 index.insert(entity);
1356 self.graph_read_index_bytes
1357 .store(index.memory_bytes(), Ordering::Relaxed);
1358 }
1359 let indexes = self.unique_key_lookup.get_mut();
1360 for index in indexes.iter_mut() {
1361 index.insert(entity);
1362 }
1363 self.unique_key_lookup_bytes.store(
1364 indexes
1365 .iter()
1366 .map(|index| index.memory_bytes() as u64)
1367 .sum(),
1368 Ordering::Relaxed,
1369 );
1370 let kind_key = entity.kind.storage_type().to_string();
1372 self.kind_index
1373 .entry(kind_key)
1374 .or_default()
1375 .insert(entity.id);
1376
1377 if let EntityData::Row(row) = &entity.data {
1379 if let Some(first_col) = row.columns.first() {
1380 let pk_str = format!("{:?}", first_col);
1381 self.pk_index
1382 .insert((entity.kind.collection().to_string(), pk_str), entity.id);
1383 }
1384 }
1385
1386 for cross_ref in entity.cross_refs() {
1388 self.cross_ref_forward
1389 .entry(cross_ref.source)
1390 .or_default()
1391 .push((cross_ref.target, cross_ref.ref_type));
1392
1393 self.cross_ref_reverse
1394 .entry(cross_ref.target)
1395 .or_default()
1396 .push((cross_ref.source, cross_ref.ref_type));
1397 }
1398 }
1399
1400 fn invalidate_graph_read_index(&mut self) {
1401 *self.graph_read_index.get_mut() = None;
1402 self.graph_read_index_bytes.store(0, Ordering::Relaxed);
1403 }
1404
1405 #[cfg(test)]
1408 pub(crate) fn visit_graph_index_candidates(
1409 &self,
1410 kind: GraphEntityKind,
1411 keys: &[EntityId],
1412 before_work: &mut impl FnMut() -> bool,
1413 visit: &mut impl FnMut(&UnifiedEntity) -> bool,
1414 ) -> bool {
1415 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1416 return false;
1417 };
1418 let index = cached.as_ref().expect("complete graph index");
1419 for key in keys {
1420 if !before_work() {
1421 return false;
1422 }
1423 for id in index.candidates(kind, *key) {
1424 if !before_work() {
1425 return false;
1426 }
1427 if let Some(entity) = self.get(id) {
1428 if !visit(entity) {
1429 return false;
1430 }
1431 }
1432 }
1433 }
1434 true
1435 }
1436
1437 fn graph_read_index_for_query(
1438 &self,
1439 before_work: &mut impl FnMut() -> bool,
1440 ) -> Option<parking_lot::RwLockReadGuard<'_, Option<SegmentGraphIndex>>> {
1441 let mut cached = self.graph_read_index.read();
1442 if cached.is_none() {
1443 drop(cached);
1444 let mut writer = self.graph_read_index.write();
1445 if writer.is_none() {
1446 let mut index = SegmentGraphIndex::default();
1447 if !self.for_each_fast(|entity| {
1450 if !before_work() {
1451 return false;
1452 }
1453 index.insert(entity);
1454 true
1455 }) {
1456 return None;
1457 }
1458 self.graph_read_index_bytes
1459 .store(index.memory_bytes(), Ordering::Relaxed);
1460 *writer = Some(index);
1461 }
1462 cached = parking_lot::RwLockWriteGuard::downgrade(writer);
1463 }
1464 Some(cached)
1465 }
1466
1467 pub(crate) fn graph_index_cursor(
1468 &self,
1469 kind: GraphEntityKind,
1470 keys: &[EntityId],
1471 before_work: &mut impl FnMut() -> bool,
1472 ) -> Option<GraphIndexCursor> {
1473 let cached = self.graph_read_index_for_query(before_work)?;
1474 let index = cached.as_ref().expect("complete graph index");
1475 let mut upper = EntityId::new(0);
1476 for key in keys {
1477 if !before_work() {
1478 return None;
1479 }
1480 if let Some(last) = index.candidates(kind, *key).next_back() {
1481 upper = upper.max(last);
1482 }
1483 }
1484 Some(GraphIndexCursor {
1485 key_index: 0,
1486 after: None,
1487 upper,
1488 })
1489 }
1490
1491 pub(crate) fn graph_index_batch(
1494 &self,
1495 kind: GraphEntityKind,
1496 keys: &[EntityId],
1497 cursor: &mut GraphIndexCursor,
1498 before_work: &mut impl FnMut() -> bool,
1499 visit: &mut impl FnMut(EntityId),
1500 ) -> bool {
1501 self.graph_index_batch_with_limit(kind, keys, cursor, before_work, visit, SCAN_BATCH_SIZE)
1502 }
1503
1504 pub(crate) fn graph_index_batch_with_limit(
1505 &self,
1506 kind: GraphEntityKind,
1507 keys: &[EntityId],
1508 cursor: &mut GraphIndexCursor,
1509 before_work: &mut impl FnMut() -> bool,
1510 visit: &mut impl FnMut(EntityId),
1511 limit: usize,
1512 ) -> bool {
1513 assert!(
1514 limit > 0 && limit <= SCAN_BATCH_SIZE,
1515 "bounded graph index batch"
1516 );
1517 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1518 return false;
1519 };
1520 let index = cached.as_ref().expect("complete graph index");
1521 let mut count = 0;
1522 while !cursor.finished(keys) && count < limit {
1523 let key = keys[cursor.key_index];
1524 if !before_work() {
1525 return false;
1526 }
1527 for id in index
1528 .candidates_after(kind, key, cursor.after, cursor.upper)
1529 .take(limit - count)
1530 {
1531 if !before_work() {
1532 return false;
1533 }
1534 cursor.after = Some(id);
1535 count += 1;
1536 visit(id);
1537 }
1538 if count == limit && cursor.after != Some(cursor.upper) {
1539 break;
1540 }
1541 cursor.key_index += 1;
1542 cursor.after = None;
1543 }
1544 true
1545 }
1546
1547 pub(crate) fn find_unique_key_conflict(
1552 &self,
1553 columns: &[Vec<String>],
1554 key_index: usize,
1555 signatures: &[String],
1556 exclude_id: Option<EntityId>,
1557 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1558 ) -> Option<EntityId> {
1559 assert!(
1560 key_index < columns.len(),
1561 "constraint index belongs to current schema"
1562 );
1563 let cached = self.unique_key_lookup.read();
1564 if cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1565 return self.find_unique_key_candidate(
1566 &cached[key_index],
1567 signatures,
1568 exclude_id,
1569 reserves_key,
1570 );
1571 }
1572 drop(cached);
1573 let mut cached = self.unique_key_lookup.write();
1574 if !cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1575 let mut indexes: Vec<_> = columns
1576 .iter()
1577 .map(|columns| SegmentUniqueKey::new(columns))
1578 .collect();
1579 self.for_each_fast(|entity| {
1580 for index in &mut indexes {
1581 index.insert(entity);
1582 }
1583 true
1584 });
1585 self.unique_key_lookup_bytes.store(
1586 indexes
1587 .iter()
1588 .map(|index| index.memory_bytes() as u64)
1589 .sum(),
1590 Ordering::Relaxed,
1591 );
1592 *cached = indexes;
1593 }
1594 self.find_unique_key_candidate(&cached[key_index], signatures, exclude_id, reserves_key)
1595 }
1596
1597 fn find_unique_key_candidate(
1598 &self,
1599 index: &SegmentUniqueKey,
1600 signatures: &[String],
1601 exclude_id: Option<EntityId>,
1602 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1603 ) -> Option<EntityId> {
1604 index.candidates(signatures).iter().copied().find(|id| {
1605 if Some(*id) == exclude_id {
1606 return false;
1607 }
1608 let Some(
1609 entity @ UnifiedEntity {
1610 data: EntityData::Row(row),
1611 ..
1612 },
1613 ) = self.get(*id)
1614 else {
1615 return false;
1616 };
1617 index
1618 .columns
1619 .iter()
1620 .zip(signatures)
1621 .all(|(column, signature)| {
1622 row.get_field(column)
1623 .is_some_and(|value| format!("{value:?}") == *signature)
1624 })
1625 && reserves_key(entity)
1626 })
1627 }
1628
1629 fn invalidate_unique_key_lookup(&mut self) {
1630 *self.unique_key_lookup.get_mut() = Vec::new();
1631 self.unique_key_lookup_bytes.store(0, Ordering::Relaxed);
1632 }
1633
1634 pub(crate) fn find_key_entity(
1637 &self,
1638 key: &str,
1639 newest: bool,
1640 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1641 ) -> Option<&UnifiedEntity> {
1642 let index = self
1645 .key_lookup
1646 .as_ref()
1647 .expect("segment owns its key index");
1648 self.find_key_candidate(index, key, newest, visible)
1649 }
1650
1651 fn find_key_candidate(
1652 &self,
1653 index: &SegmentKeyLookup,
1654 key: &str,
1655 newest: bool,
1656 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1657 ) -> Option<&UnifiedEntity> {
1658 self.select_key_entity(
1659 index
1660 .candidates(key)
1661 .iter()
1662 .chain(index.mutated_candidates())
1663 .filter_map(|id| self.get(*id)),
1664 key,
1665 newest,
1666 visible,
1667 )
1668 }
1669
1670 fn select_key_entity<'a>(
1671 &self,
1672 entities: impl Iterator<Item = &'a UnifiedEntity>,
1673 key: &str,
1674 newest: bool,
1675 visible: &mut impl FnMut(&UnifiedEntity) -> bool,
1676 ) -> Option<&'a UnifiedEntity> {
1677 let mut best: Option<&UnifiedEntity> = None;
1678 for entity in entities {
1679 if SegmentKeyLookup::key(entity) != Some(key) || !visible(entity) {
1681 continue;
1682 }
1683 if !newest {
1684 return Some(entity);
1685 }
1686 if best.is_none_or(|current| entity.xmin >= current.xmin) {
1687 best = Some(entity);
1688 }
1689 }
1690 best
1691 }
1692
1693 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1695 if let Some(id) = entity_id_from_probe_key(key) {
1696 return self.has_live_entity(id);
1697 }
1698
1699 let Ok(pk_str) = std::str::from_utf8(key) else {
1700 return false;
1701 };
1702 let pk_key = (self.collection.clone(), pk_str.to_string());
1703 self.pk_index
1704 .get(&pk_key)
1705 .is_some_and(|id| self.has_live_entity(*id))
1706 }
1707
1708 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1710 if let Some(index) = self.key_lookup.as_mut() {
1711 index.remove(entity);
1712 }
1713 if let Some(index) = self.graph_read_index.get_mut() {
1714 index.remove(entity);
1715 self.graph_read_index_bytes
1716 .store(index.memory_bytes(), Ordering::Relaxed);
1717 }
1718 let kind_key = entity.kind.storage_type().to_string();
1720 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1721 set.remove(&entity.id);
1722 }
1723
1724 if let EntityData::Row(row) = &entity.data {
1726 if let Some(first_col) = row.columns.first() {
1727 let pk_str = format!("{:?}", first_col);
1728 self.pk_index
1729 .remove(&(entity.kind.collection().to_string(), pk_str));
1730 }
1731 }
1732
1733 self.cross_ref_forward.remove(&entity.id);
1735 }
1737
1738 fn reindex_for_update(
1748 &mut self,
1749 old: &UpdateIndexSnapshot,
1750 new: &UnifiedEntity,
1751 modified_columns: Option<&[String]>,
1752 ) {
1753 let pk_changed = match &new.data {
1758 EntityData::Row(new_row) => {
1759 if let Some(cols) = modified_columns {
1760 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1763 new_row
1764 .schema
1765 .as_deref()
1766 .and_then(|schema| schema.first().map(|name| name.as_str()))
1767 });
1768 match pk_col_name {
1769 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1770 None => old.pk_value.as_ref() != new_row.columns.first(),
1772 }
1773 } else {
1774 old.pk_value.as_ref() != new_row.columns.first()
1775 }
1776 }
1777 _ => false,
1779 };
1780
1781 if pk_changed {
1782 if let Some((collection, pk_str)) = &old.pk_index_key {
1784 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1785 }
1786 if let EntityData::Row(row) = &new.data {
1788 if let Some(first_col) = row.columns.first() {
1789 let pk_str = format!("{:?}", first_col);
1790 self.pk_index
1791 .insert((new.kind.collection().to_string(), pk_str), new.id);
1792 }
1793 }
1794 }
1795
1796 let new_refs = new.cross_refs();
1798 if old.cross_refs.as_slice() != new_refs {
1799 self.cross_ref_forward.remove(&new.id);
1801 for cross_ref in &old.cross_refs {
1803 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1804 rev.retain(|(src, _)| *src != new.id);
1805 }
1806 }
1807 for cross_ref in new_refs {
1809 self.cross_ref_forward
1810 .entry(cross_ref.source)
1811 .or_default()
1812 .push((cross_ref.target, cross_ref.ref_type));
1813 self.cross_ref_reverse
1814 .entry(cross_ref.target)
1815 .or_default()
1816 .push((cross_ref.source, cross_ref.ref_type));
1817 }
1818 }
1819 }
1820
1821 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1823 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1824 }
1825
1826 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1828 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1829 }
1830
1831 pub fn age_secs(&self) -> u64 {
1833 let now = current_unix_secs();
1834 now.saturating_sub(self.created_at)
1835 }
1836
1837 pub fn idle_secs(&self) -> u64 {
1839 let now = current_unix_secs();
1840 now.saturating_sub(self.last_write_at)
1841 }
1842
1843 pub fn bulk_insert(
1852 &mut self,
1853 entities: Vec<UnifiedEntity>,
1854 ) -> Result<Vec<EntityId>, SegmentError> {
1855 self.mark_checkpoint_dirty();
1856 if !self.state.is_writable() {
1857 return Err(SegmentError::NotWritable);
1858 }
1859
1860 if let Some(index) = self.key_lookup.as_mut() {
1861 for entity in &entities {
1862 index.insert(entity);
1863 }
1864 }
1865
1866 let indexes = self.unique_key_lookup.get_mut();
1867 for index in indexes.iter_mut() {
1868 for entity in &entities {
1869 index.insert(entity);
1870 }
1871 }
1872 self.unique_key_lookup_bytes.store(
1873 indexes
1874 .iter()
1875 .map(|index| index.memory_bytes() as u64)
1876 .sum(),
1877 Ordering::Relaxed,
1878 );
1879
1880 let n = entities.len();
1881
1882 if entities.is_empty() {
1883 return Ok(Vec::new());
1884 }
1885
1886 for run in
1889 entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
1890 {
1891 let kind = run[0].kind.storage_type();
1892 self.observe_graph_kind(&run[0].kind);
1893 if matches!(
1894 run[0].kind,
1895 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
1896 ) {
1897 if let Some(index) = self.graph_read_index.get_mut() {
1898 for entity in run {
1899 index.insert(entity);
1900 }
1901 self.graph_read_index_bytes
1902 .store(index.memory_bytes(), Ordering::Relaxed);
1903 }
1904 }
1905 let kind_set = match self.kind_index.get_mut(kind) {
1906 Some(ids) => ids,
1907 None => self.kind_index.entry(kind.to_string()).or_default(),
1908 };
1909 kind_set.reserve(run.len());
1910 kind_set.extend(run.iter().map(|entity| entity.id));
1911 }
1912
1913 let now = current_unix_secs();
1914
1915 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
1916
1917 let mut ids = Vec::with_capacity(n);
1918
1919 if self.flat_entities.is_empty() && self.entities.is_empty() {
1921 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
1923 self.use_flat = true;
1924 }
1925
1926 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
1936 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
1937 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
1941
1942 let mut batch_bytes: usize = 0;
1947
1948 if self.use_flat {
1949 self.flat_entities.reserve(n);
1950 for (i, mut entity) in entities.into_iter().enumerate() {
1951 entity.sequence_id = base_seq + i as u64;
1952 let id = entity.id;
1953 ids.push(id);
1954 batch_bytes += Self::estimate_entity_size(&entity);
1955 if let EntityData::Row(row) = &entity.data {
1956 if row.schema.is_some() && !row.columns.is_empty() {
1957 if columnar_zone_updates.is_empty() {
1958 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
1959 columnar_schema = row.schema.clone();
1960 }
1961 for (ci, val) in row.columns.iter().enumerate() {
1962 if !matches!(val, Value::Null) {
1963 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
1964 bucket.push(val.clone());
1965 }
1966 }
1967 }
1968 } else {
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 }
1976 let expected = self.base_entity_id + self.flat_entities.len() as u64;
1992 if id.raw() == expected {
1993 self.flat_entities.push(entity);
1994 } else {
1995 self.hash_scan_ids.push(id);
1996 self.entities.insert(id, entity);
1997 }
1998 }
1999 } else {
2000 self.entities.reserve(n);
2002 self.hash_scan_ids.reserve(n);
2003 let mut pairs = Vec::with_capacity(n);
2004 for (i, mut entity) in entities.into_iter().enumerate() {
2005 entity.sequence_id = base_seq + i as u64;
2006 let id = entity.id;
2007 ids.push(id);
2008 batch_bytes += Self::estimate_entity_size(&entity);
2009 if let EntityData::Row(row) = &entity.data {
2010 for (col, val) in row.iter_fields() {
2011 if !matches!(val, Value::Null) {
2012 named_zone_updates.push((col.to_string(), val.clone()));
2013 }
2014 }
2015 }
2016 self.hash_scan_ids.push(id);
2017 pairs.push((id, entity));
2018 }
2019 self.entities.extend(pairs);
2020 }
2021
2022 if !columnar_zone_updates.is_empty() {
2026 let schema = columnar_schema.as_ref();
2027 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
2028 if values.is_empty() {
2029 continue;
2030 }
2031 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
2032 continue;
2033 };
2034 let mut iter = values.into_iter();
2040 if let Some(first) = iter.next() {
2041 let previous = self
2042 .col_zones
2043 .get(col_name)
2044 .map_or(0, ColZone::payload_bytes);
2045 let zone = self
2046 .col_zones
2047 .entry(col_name.clone())
2048 .and_modify(|z| z.update(&first))
2049 .or_insert_with(|| ColZone::new(first));
2050 for v in iter {
2051 zone.update(&v);
2052 }
2053 self.zone_payload_bytes = self
2054 .zone_payload_bytes
2055 .saturating_sub(previous)
2056 .saturating_add(zone.payload_bytes());
2057 }
2058 }
2059 }
2060 for (col, val) in named_zone_updates {
2061 self.update_column_zone(&col, &val);
2062 }
2063
2064 self.add_memory(batch_bytes);
2065 self.last_write_at = now;
2066
2067 if self.use_flat {
2069 self.published_flat_len
2070 .store(self.flat_entities.len(), Ordering::Release);
2071 }
2072
2073 Ok(ids)
2074 }
2075
2076 pub(crate) fn force_delete(&mut self, id: EntityId) -> bool {
2079 if self.use_flat {
2080 let raw = id.raw();
2081 if raw >= self.base_entity_id {
2082 let idx = (raw - self.base_entity_id) as usize;
2083 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2084 self.tombstone_flat_entity(idx, id);
2085 return true;
2086 }
2087 }
2088 return self.remove_hashmap_entity(id);
2089 }
2090
2091 self.remove_hashmap_entity(id)
2092 }
2093
2094 pub(crate) fn force_update_with_metadata(
2097 &mut self,
2098 entity: &UnifiedEntity,
2099 modified_columns: &[String],
2100 metadata: Option<&Metadata>,
2101 ) -> Result<(), SegmentError> {
2102 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
2103 }
2104}
2105
2106impl UnifiedSegment for GrowingSegment {
2107 fn id(&self) -> SegmentId {
2108 self.id
2109 }
2110
2111 fn state(&self) -> SegmentState {
2112 self.state
2113 }
2114
2115 fn collection(&self) -> &str {
2116 &self.collection
2117 }
2118
2119 fn stats(&self) -> SegmentStats {
2120 let mut stats = SegmentStats {
2121 entity_count: self.live_entity_count(),
2122 deleted_count: self.deleted.len(),
2123 memory_bytes: self.resident_bytes() as usize,
2124 ..Default::default()
2125 };
2126
2127 for entity in self.entities.values() {
2128 match &entity.kind {
2129 EntityKind::TableRow { .. } => stats.row_count += 1,
2130 EntityKind::GraphNode(_) => stats.node_count += 1,
2131 EntityKind::GraphEdge(_) => stats.edge_count += 1,
2132 EntityKind::Vector { .. } => stats.vector_count += 1,
2133 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
2134 EntityKind::QueueMessage { .. } => stats.row_count += 1,
2135 }
2136 stats.cross_ref_count += entity.cross_refs().len();
2137 }
2138
2139 stats
2140 }
2141
2142 fn entity_count(&self) -> usize {
2143 self.live_entity_count()
2144 }
2145
2146 fn contains(&self, id: EntityId) -> bool {
2147 self.has_live_entity(id)
2148 }
2149
2150 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2151 if self.deleted.contains_key(&id) {
2152 return None;
2153 }
2154 if self.use_flat {
2155 let raw = id.raw();
2156 if raw >= self.base_entity_id {
2157 let idx = (raw - self.base_entity_id) as usize;
2158 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
2159 return Some(entity);
2160 }
2161 }
2162 self.entities.get(&id)
2168 } else {
2169 self.entities.get(&id)
2170 }
2171 }
2172
2173 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
2174 self.mark_checkpoint_entity_dirty(id);
2175 self.invalidate_unique_key_lookup();
2176 self.invalidate_graph_read_index();
2177 if self.deleted.contains_key(&id) || !self.state.is_writable() {
2178 return None;
2179 }
2180 if let Some(mut index) = self.key_lookup.take() {
2181 if let Some(entity) = self.get(id) {
2182 index.mark_mutated(entity);
2183 }
2184 self.key_lookup = Some(index);
2185 }
2186 self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
2189 if self.use_flat {
2190 let raw = id.raw();
2191 if raw >= self.base_entity_id {
2192 let idx = (raw - self.base_entity_id) as usize;
2193 if self
2194 .flat_entities
2195 .get(idx)
2196 .map(|e| e.id == id)
2197 .unwrap_or(false)
2198 {
2199 return self.flat_entities.get_mut(idx);
2200 }
2201 }
2202 self.entities.get_mut(&id)
2205 } else {
2206 self.entities.get_mut(&id)
2207 }
2208 }
2209
2210 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2211 if !self.state.is_writable() {
2212 return Err(SegmentError::NotWritable);
2213 }
2214
2215 if self.entities.contains_key(&entity.id) {
2216 return Err(SegmentError::AlreadyExists(entity.id));
2217 }
2218
2219 entity.sequence_id = self.next_sequence();
2221
2222 let size = Self::estimate_entity_size(&entity);
2224 self.add_memory(size);
2225
2226 self.index_entity(&entity);
2228
2229 self.update_col_zones_from_entity(&entity);
2231
2232 let id = entity.id;
2234 self.entities.insert(id, entity);
2235
2236 self.last_write_at = current_unix_secs();
2238
2239 Ok(id)
2240 }
2241
2242 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2243 if !self.state.is_writable() {
2244 return Err(SegmentError::NotWritable);
2245 }
2246
2247 self.apply_update_with_metadata(&entity, None)?;
2248 self.last_write_at = current_unix_secs();
2249
2250 Ok(())
2251 }
2252
2253 fn update_hot(
2254 &mut self,
2255 entity: UnifiedEntity,
2256 modified_columns: &[String],
2257 ) -> Result<(), SegmentError> {
2258 if !self.state.is_writable() {
2259 return Err(SegmentError::NotWritable);
2260 }
2261
2262 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
2263 self.last_write_at = current_unix_secs();
2264 Ok(())
2265 }
2266
2267 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2268 if !self.state.is_writable() {
2269 return Err(SegmentError::NotWritable);
2270 }
2271
2272 if self.use_flat {
2274 let raw = id.raw();
2275 if raw >= self.base_entity_id {
2276 let idx = (raw - self.base_entity_id) as usize;
2277 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2278 self.tombstone_flat_entity(idx, id);
2279 return Ok(true);
2280 }
2281 }
2282 return Ok(self.remove_hashmap_entity(id));
2283 }
2284
2285 Ok(self.remove_hashmap_entity(id))
2286 }
2287
2288 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2289 if !self.has_live_entity(id) {
2290 return None;
2291 }
2292 Some(self.metadata.get_all(id))
2293 }
2294
2295 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2296 self.mark_checkpoint_entity_dirty(id);
2297 if !self.state.is_writable() {
2298 return Err(SegmentError::NotWritable);
2299 }
2300
2301 if !self.has_live_entity(id) {
2302 return Err(SegmentError::NotFound(id));
2303 }
2304
2305 self.metadata.set_all(id, &metadata);
2306 Ok(())
2307 }
2308
2309 fn seal(&mut self) -> Result<(), SegmentError> {
2310 if self.state != SegmentState::Growing {
2311 return Err(SegmentError::InvalidState(self.state));
2312 }
2313
2314 self.state = SegmentState::Sealing;
2315
2316 self.rebuild_sealed_col_zones();
2321
2322 self.state = SegmentState::Sealed;
2323 Ok(())
2324 }
2325
2326 fn should_seal(&self, config: &SegmentConfig) -> bool {
2327 if self.entities.len() >= config.max_entities {
2329 return true;
2330 }
2331
2332 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
2334 return true;
2335 }
2336
2337 if self.age_secs() >= config.max_age_secs {
2339 return true;
2340 }
2341
2342 false
2343 }
2344
2345 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2346 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2347 Box::new(self.flat_entities.iter().chain(self.entities.values()))
2351 } else {
2352 Box::new(self.entities.values())
2353 };
2354 if self.deleted.is_empty() {
2355 base
2356 } else {
2357 Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
2358 }
2359 }
2360
2361 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2362 let ids = self.kind_index.get(kind_filter);
2365 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2368 Box::new(self.flat_entities.iter())
2369 } else {
2370 Box::new(std::iter::empty())
2371 };
2372 Box::new(flat.chain(self.entities.values()).filter(move |e| {
2373 if self.deleted.contains_key(&e.id) {
2374 return false;
2375 }
2376 if let Some(ids) = ids {
2377 ids.contains(&e.id)
2378 } else {
2379 false
2380 }
2381 }))
2382 }
2383
2384 fn filter_metadata(
2385 &self,
2386 filters: &[(String, super::metadata::MetadataFilter)],
2387 ) -> Vec<EntityId> {
2388 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
2391 Box::new(self.flat_entities.iter().map(|e| e.id))
2392 } else {
2393 Box::new(std::iter::empty())
2394 };
2395 flat_ids
2396 .chain(self.entities.keys().copied())
2397 .filter(|id| {
2398 if self.deleted.contains_key(id) {
2399 return false;
2400 }
2401 let metadata = self.metadata.get_all(*id);
2402 metadata.matches_all(filters)
2403 })
2404 .collect()
2405 }
2406}
2407
2408fn build_minmax_multi_intervals(
2409 entries: &[(CanonicalKey, Value)],
2410 max_intervals: usize,
2411) -> Vec<ColZone> {
2412 if entries.is_empty() {
2413 return Vec::new();
2414 }
2415 if entries.len() == 1 || max_intervals <= 1 {
2416 return vec![ColZone::with_bounds(
2417 entries[0].1.clone(),
2418 entries[entries.len() - 1].1.clone(),
2419 )];
2420 }
2421
2422 let mut split_points = if entries.len() <= max_intervals {
2423 (1..entries.len()).collect::<Vec<_>>()
2424 } else {
2425 let target_splits = max_intervals - 1;
2426 let mut selected = select_gap_split_points(entries, target_splits);
2427 if selected.len() < target_splits {
2428 for bucket in 1..max_intervals {
2429 let idx = bucket * entries.len() / max_intervals;
2430 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
2431 continue;
2432 }
2433 selected.push(idx);
2434 if selected.len() >= target_splits {
2435 break;
2436 }
2437 }
2438 }
2439 selected.sort_unstable();
2440 selected.dedup();
2441 selected
2442 };
2443
2444 split_points.push(entries.len());
2445
2446 let mut out = Vec::with_capacity(split_points.len());
2447 let mut start = 0usize;
2448 for end in split_points {
2449 if end <= start {
2450 continue;
2451 }
2452 out.push(ColZone::with_bounds(
2453 entries[start].1.clone(),
2454 entries[end - 1].1.clone(),
2455 ));
2456 start = end;
2457 }
2458
2459 if out.is_empty() {
2460 out.push(ColZone::with_bounds(
2461 entries[0].1.clone(),
2462 entries[entries.len() - 1].1.clone(),
2463 ));
2464 }
2465
2466 out
2467}
2468
2469fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
2470 let mut gaps = Vec::new();
2471 for idx in 1..entries.len() {
2472 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
2473 if score > 0.0 {
2474 gaps.push((score, idx));
2475 }
2476 }
2477 }
2478 gaps.sort_by(|left, right| {
2479 right
2480 .0
2481 .partial_cmp(&left.0)
2482 .unwrap_or(std::cmp::Ordering::Equal)
2483 .then_with(|| left.1.cmp(&right.1))
2484 });
2485 gaps.into_iter()
2486 .take(max_splits)
2487 .map(|(_, idx)| idx)
2488 .collect()
2489}
2490
2491fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
2492 if left.family() != right.family() {
2493 return None;
2494 }
2495 match (left, right) {
2496 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
2497 Some(r.saturating_sub(*l) as f64)
2498 }
2499 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
2500 Some(r.saturating_sub(*l) as f64)
2501 }
2502 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
2503 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
2504 }
2505 _ => None,
2506 }
2507}
2508
2509#[cfg(test)]
2510mod tests {
2511 use super::*;
2512 use crate::storage::unified::entity::RowData;
2513 use crate::storage::unified::MetadataValue;
2514 use reddb_types::Value;
2515
2516 fn keyed_row(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
2517 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
2518 entity.data = EntityData::Row(RowData::with_names(
2519 vec![Value::text(key), Value::Integer(id as i64)],
2520 vec!["key".into(), "value".into()],
2521 ));
2522 entity.set_xmin(xmin);
2523 entity.set_xmax(xmax);
2524 entity
2525 }
2526
2527 fn keyed_probe(segment: &GrowingSegment, key: &str) -> Option<EntityId> {
2528 segment
2529 .find_key_entity(key, true, &mut |_| true)
2530 .map(|entity| entity.id)
2531 }
2532
2533 #[test]
2534 fn key_lookup_visits_only_requested_key_versions() {
2535 let mut segment = GrowingSegment::new(1, "keys");
2536 segment
2537 .bulk_insert(
2538 (1..=4096)
2539 .map(|id| keyed_row(id, &format!("other-{id}"), 1, 0))
2540 .collect(),
2541 )
2542 .expect("unrelated rows");
2543 segment
2544 .insert(keyed_row(5000, "wanted", 2, 8))
2545 .expect("old version");
2546 segment
2547 .insert(keyed_row(5001, "wanted", 8, 0))
2548 .expect("new version");
2549 let mut visits = 0;
2550 let result = segment.find_key_entity("wanted", true, &mut |entity| {
2551 visits += 1;
2552 entity.xmin < 8
2553 });
2554 assert_eq!(result.map(|entity| entity.id), Some(EntityId::new(5000)));
2555 assert_eq!(
2556 visits, 2,
2557 "unrelated rows must never reach visibility evaluation"
2558 );
2559 assert_eq!(keyed_probe(&segment, "wanted"), Some(EntityId::new(5001)));
2560 let mut misses = 0;
2561 assert!(segment
2562 .find_key_entity("absent", true, &mut |_| {
2563 misses += 1;
2564 true
2565 })
2566 .is_none());
2567 assert_eq!(misses, 0);
2568
2569 let mut collision = SegmentKeyLookup::default();
2571 collision.insert(&keyed_row(1, "wanted", 1, 0));
2572 assert!(segment
2573 .find_key_candidate(&collision, "wanted", true, &mut |_| true)
2574 .is_none());
2575 }
2576
2577 #[test]
2578 fn raw_mutation_does_not_disable_unrelated_key_candidates() {
2579 let mut segment = GrowingSegment::new(1, "keys");
2580 segment
2581 .bulk_insert(
2582 (1..=4096)
2583 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2584 .collect(),
2585 )
2586 .expect("rows");
2587 for _ in 0..3 {
2588 *segment.get_mut(EntityId::new(1)).expect("mutable") = keyed_row(1, "renamed", 2, 0);
2589 }
2590 let index = segment.key_lookup.as_ref().expect("index retained");
2591 assert_eq!(index.candidates("key-4096"), &[EntityId::new(4096)]);
2592 assert_eq!(index.mutated_candidates().count(), 1);
2593 assert!(index.candidates("key-1").is_empty());
2594 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(1)));
2595 assert_eq!(keyed_probe(&segment, "key-4096"), Some(EntityId::new(4096)));
2596 segment.delete(EntityId::new(1)).expect("delete raw row");
2597 assert_eq!(
2598 segment
2599 .key_lookup
2600 .as_ref()
2601 .unwrap()
2602 .mutated_candidates()
2603 .count(),
2604 0
2605 );
2606 assert_eq!(keyed_probe(&segment, "renamed"), None);
2607 }
2608
2609 #[test]
2610 fn key_lookup_tracks_mutations_sealing_and_adoption() {
2611 let mut segment = GrowingSegment::new(1, "keys");
2612 segment
2613 .bulk_insert(vec![keyed_row(1, "flat", 1, 0)])
2614 .expect("flat row");
2615 segment
2616 .insert(keyed_row(3, "hash", 1, 0))
2617 .expect("hash fallback");
2618 segment
2619 .bulk_insert(vec![keyed_row(5, "gap", 1, 0)])
2620 .expect("bulk gap");
2621 assert!(segment.key_lookup.as_ref().expect("index").memory_bytes() > 0);
2622 assert_eq!(keyed_probe(&segment, "gap"), Some(EntityId::new(5)));
2623 segment
2624 .update(keyed_row(3, "renamed", 2, 0))
2625 .expect("rename key");
2626 assert!(
2627 segment.key_lookup.is_some(),
2628 "key replacements maintain the index"
2629 );
2630 assert_eq!(keyed_probe(&segment, "hash"), None);
2631 assert_eq!(keyed_probe(&segment, "renamed"), Some(EntityId::new(3)));
2632
2633 *segment.get_mut(EntityId::new(3)).expect("mutable row") = keyed_row(3, "raw", 3, 0);
2634 assert!(
2635 segment.key_lookup.is_some(),
2636 "raw mutation preserves unaffected key candidates"
2637 );
2638 assert_eq!(keyed_probe(&segment, "renamed"), None);
2639 assert_eq!(keyed_probe(&segment, "raw"), Some(EntityId::new(3)));
2640 segment.seal().expect("seal");
2641 segment
2642 .force_update_with_metadata(&keyed_row(5, "sealed", 4, 0), &["key".into()], None)
2643 .expect("sealed rename");
2644 assert_eq!(keyed_probe(&segment, "gap"), None);
2645 assert_eq!(keyed_probe(&segment, "sealed"), Some(EntityId::new(5)));
2646 assert!(segment.force_delete(EntityId::new(1)));
2647 assert!(segment.force_delete(EntityId::new(3)));
2648 assert_eq!(keyed_probe(&segment, "flat"), None);
2649 assert_eq!(keyed_probe(&segment, "raw"), None);
2650
2651 let mut adopted = GrowingSegment::new(2, "keys");
2652 segment.for_each_fast(|entity| {
2653 adopted.adopt_entity(entity.clone(), None);
2654 true
2655 });
2656 assert_eq!(keyed_probe(&adopted, "sealed"), Some(EntityId::new(5)));
2657 assert!(adopted.evict_entity(EntityId::new(5)));
2658 assert_eq!(keyed_probe(&adopted, "sealed"), None);
2659 }
2660
2661 #[test]
2662 #[ignore = "diagnostic microbenchmark; no timing assertions"]
2663 fn key_lookup_scaling_diagnostic() {
2664 use std::{hint::black_box, time::Instant};
2665 for count in [100, 10_000, 100_000] {
2666 let mut segment = GrowingSegment::new(1, "keys");
2667 segment
2668 .bulk_insert(
2669 (1..=count)
2670 .map(|id| keyed_row(id, &format!("key-{id}"), 1, 0))
2671 .collect(),
2672 )
2673 .expect("rows");
2674 let key = format!("key-{count}");
2675 let started = Instant::now();
2676 for _ in 0..1000 {
2677 assert_eq!(
2678 keyed_probe(black_box(&segment), black_box(&key)),
2679 Some(EntityId::new(count))
2680 );
2681 }
2682 let indexed = started.elapsed();
2683 let started = Instant::now();
2684 for _ in 0..1000 {
2685 let mut found = None;
2686 segment.for_each_fast(|entity| {
2687 if SegmentKeyLookup::key(entity) == Some(black_box(key.as_str())) {
2688 found = Some(entity.id);
2689 return false;
2690 }
2691 true
2692 });
2693 assert_eq!(black_box(found), Some(EntityId::new(count)));
2694 }
2695 eprintln!(
2696 "rows={count} reads=1000 indexed={indexed:?} scan={:?}",
2697 started.elapsed()
2698 );
2699 }
2700 }
2701
2702 fn pruning_node(id: u64) -> UnifiedEntity {
2703 UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
2704 }
2705
2706 fn pruning_edge(id: u64) -> UnifiedEntity {
2707 UnifiedEntity::graph_edge(
2708 EntityId::new(id),
2709 "mixed",
2710 "a",
2711 "b",
2712 1.0,
2713 Default::default(),
2714 )
2715 }
2716
2717 #[test]
2718 fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
2719 let mut segment = GrowingSegment::new(1, "mixed");
2720 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
2721 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
2722 segment
2723 .bulk_insert(vec![
2724 UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
2725 pruning_node(2),
2726 pruning_edge(4),
2727 ])
2728 .expect("mixed bulk with gap");
2729 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2730 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2731 segment.seal().expect("seal");
2732 assert!(segment.force_delete(EntityId::new(2)));
2733 assert!(segment.force_delete(EntityId::new(4)));
2734 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2736 let mut rebuilt = GrowingSegment::new(2, "mixed");
2737 for item in segment.iter() {
2738 rebuilt.adopt_entity(item.clone(), None);
2739 }
2740 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2741 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2742 rebuilt.adopt_entity(pruning_node(5), None);
2743 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2744 rebuilt.insert(pruning_edge(6)).expect("single edge");
2745 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2746 }
2747
2748 #[test]
2749 fn graph_presence_covers_replacement_and_unrestricted_mutation() {
2750 for flat in [false, true] {
2751 for operation in 0..4 {
2752 let mut segment = GrowingSegment::new(1, "mixed");
2753 let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
2754 if flat {
2755 segment.bulk_insert(vec![vector]).expect("flat vector");
2756 } else {
2757 segment.insert(vector).expect("hashmap vector");
2758 }
2759 match operation {
2760 0 => segment.update(pruning_node(1)).expect("replace with node"),
2761 1 => segment
2762 .update_hot(pruning_node(1), &[])
2763 .expect("hot replacement"),
2764 2 => {
2765 segment.seal().expect("seal vector");
2766 segment
2767 .force_update_with_metadata(&pruning_node(1), &[], None)
2768 .expect("sealed replacement");
2769 }
2770 _ => {
2771 *segment.get_mut(EntityId::new(1)).expect("mutable vector") =
2772 pruning_node(1)
2773 }
2774 }
2775 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2776 assert!(matches!(
2777 segment.get(EntityId::new(1)).expect("node").kind,
2778 EntityKind::GraphNode(_)
2779 ));
2780 if operation == 3 {
2781 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2782 }
2783 segment
2784 .force_update_with_metadata(&pruning_edge(1), &[], None)
2785 .expect("replace with edge");
2786 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2787 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2788 }
2789 }
2790 }
2791
2792 fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
2793 UnifiedEntity::new(
2794 EntityId::new(id),
2795 EntityKind::TableRow {
2796 table: "keys".into(),
2797 row_id: id,
2798 },
2799 EntityData::Row(RowData::with_names(
2800 vec![Value::text("payload"), key],
2801 vec!["payload".into(), "primary".into()],
2802 )),
2803 )
2804 }
2805
2806 fn primary_key_probe(
2807 segment: &GrowingSegment,
2808 key: Value,
2809 exclude: Option<EntityId>,
2810 ) -> Option<EntityId> {
2811 segment.find_unique_key_conflict(
2812 &[vec!["primary".into()]],
2813 0,
2814 &[format!("{key:?}")],
2815 exclude,
2816 &|_| true,
2817 )
2818 }
2819
2820 #[test]
2821 fn unique_key_lookup_follows_segment_mutations_and_sealing() {
2822 let mut segment = GrowingSegment::new(1, "keys");
2823 segment
2824 .bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
2825 .expect("bulk row");
2826 let payload_bytes = segment.memory_bytes();
2827 assert_eq!(
2828 primary_key_probe(&segment, Value::Integer(10), None),
2829 Some(EntityId::new(1))
2830 );
2831 assert!(segment.memory_bytes() > payload_bytes);
2832 segment
2833 .insert(primary_key_row(3, Value::Integer(30)))
2834 .expect("hashmap fallback row");
2835 segment
2836 .bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
2837 .expect("bulk gap row");
2838 assert_eq!(
2839 primary_key_probe(&segment, Value::Integer(40), None),
2840 Some(EntityId::new(4))
2841 );
2842 assert_eq!(
2843 primary_key_probe(&segment, Value::Integer(30), None),
2844 Some(EntityId::new(3))
2845 );
2846 segment
2847 .update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
2848 .expect("replace flat key");
2849 assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
2850 assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
2851 assert_eq!(
2852 primary_key_probe(&segment, Value::Integer(11), None),
2853 Some(EntityId::new(1))
2854 );
2855 segment
2856 .delete_batch(&[EntityId::new(1), EntityId::new(3)])
2857 .expect("delete flat and map");
2858 assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
2859 assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
2860 segment.seal().expect("seal");
2861 assert_eq!(
2862 primary_key_probe(&segment, Value::Integer(40), None),
2863 Some(EntityId::new(4))
2864 );
2865 segment
2866 .force_update_with_metadata(
2867 &primary_key_row(4, Value::Integer(41)),
2868 &["primary".into()],
2869 None,
2870 )
2871 .expect("sealed update");
2872 assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
2873 assert_eq!(
2874 primary_key_probe(&segment, Value::Integer(41), None),
2875 Some(EntityId::new(4))
2876 );
2877 assert!(segment.force_delete(EntityId::new(4)));
2878 assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
2879 }
2880
2881 #[test]
2882 fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
2883 let mut segment = GrowingSegment::new(1, "keys");
2884 segment
2885 .bulk_insert(vec![
2886 primary_key_row(1, Value::Integer(10)),
2887 primary_key_row(2, Value::Integer(20)),
2888 ])
2889 .expect("initial rows");
2890 let definitions = vec![
2891 vec!["primary".into()],
2892 vec!["payload".into()],
2893 vec!["primary".into(), "payload".into()],
2894 ];
2895 let payload = format!("{:?}", Value::text("payload"));
2896 assert_eq!(
2897 segment.find_unique_key_conflict(
2898 &definitions,
2899 1,
2900 &[payload.clone()],
2901 None,
2902 &|entity| entity.id != EntityId::new(1)
2903 ),
2904 Some(EntityId::new(2))
2905 );
2906 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2907 segment
2908 .insert(primary_key_row(3, Value::Integer(30)))
2909 .expect("append");
2910 assert_eq!(
2911 segment.find_unique_key_conflict(
2912 &definitions,
2913 2,
2914 &[format!("{:?}", Value::Integer(30)), payload.clone()],
2915 None,
2916 &|_| true
2917 ),
2918 Some(EntityId::new(3))
2919 );
2920 let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
2921 assert_eq!(
2922 segment.find_unique_key_conflict(
2923 &definitions,
2924 0,
2925 &[format!("{:?}", Value::Integer(20))],
2926 None,
2927 &|_| true
2928 ),
2929 Some(EntityId::new(2))
2930 );
2931 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2932 assert_eq!(
2933 segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
2934 bytes
2935 );
2936 let reduced = vec![vec!["payload".into()]];
2937 assert_eq!(
2938 segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
2939 Some(EntityId::new(1))
2940 );
2941 assert_eq!(segment.unique_key_lookup.read().len(), 1);
2942 assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
2943 }
2944
2945 #[test]
2946 fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
2947 let mut segment = GrowingSegment::new(1, "keys");
2948 let values = [
2949 Value::Float(f64::NAN),
2950 Value::Float(-0.0),
2951 Value::Float(0.0),
2952 Value::Integer(1),
2953 Value::UnsignedInteger(1),
2954 Value::text("1"),
2955 ];
2956 for (offset, value) in values.iter().enumerate() {
2957 let id = offset as u64 + 1;
2958 segment
2959 .insert(primary_key_row(id, value.clone()))
2960 .expect("distinct physical row");
2961 assert_eq!(
2962 primary_key_probe(&segment, value.clone(), None),
2963 Some(EntityId::new(id))
2964 );
2965 assert_eq!(
2966 primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
2967 None
2968 );
2969 }
2970 segment
2971 .insert(primary_key_row(20, Value::Integer(1)))
2972 .expect("duplicate physical key");
2973 assert_eq!(
2974 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
2975 Some(EntityId::new(20))
2976 );
2977 segment
2980 .entities
2981 .insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
2982 assert_eq!(
2983 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
2984 None
2985 );
2986 let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
2988 *entity = primary_key_row(4, Value::Integer(100));
2989 assert_eq!(
2990 primary_key_probe(&segment, Value::Integer(100), None),
2991 Some(EntityId::new(4))
2992 );
2993 }
2994
2995 #[test]
2996 fn test_growing_segment_basic() {
2997 let mut segment = GrowingSegment::new(1, "test");
2998
2999 let entity = UnifiedEntity::table_row(
3000 EntityId::new(1),
3001 "users",
3002 1,
3003 vec![Value::text("Alice".to_string())],
3004 );
3005
3006 let id = segment.insert(entity).unwrap();
3007 assert_eq!(id, EntityId::new(1));
3008 assert!(segment.contains(id));
3009
3010 let stats = segment.stats();
3011 assert_eq!(stats.entity_count, 1);
3012 assert_eq!(stats.row_count, 1);
3013 }
3014
3015 #[test]
3016 fn test_segment_metadata() {
3017 let mut segment = GrowingSegment::new(1, "test");
3018
3019 let entity = UnifiedEntity::table_row(
3020 EntityId::new(1),
3021 "users",
3022 1,
3023 vec![Value::text("Alice".to_string())],
3024 );
3025 segment.insert(entity).unwrap();
3026
3027 let mut meta = Metadata::new();
3028 meta.set("role", MetadataValue::String("admin".to_string()));
3029 meta.set("level", MetadataValue::Int(5));
3030
3031 segment.set_metadata(EntityId::new(1), meta).unwrap();
3032
3033 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
3034 assert_eq!(
3035 retrieved.get("role"),
3036 Some(&MetadataValue::String("admin".to_string()))
3037 );
3038 }
3039
3040 #[test]
3041 fn test_segment_seal() {
3042 let mut segment = GrowingSegment::new(1, "test");
3043
3044 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
3045 segment.insert(entity).unwrap();
3046
3047 assert!(segment.state().is_writable());
3049
3050 segment.seal().unwrap();
3052 assert_eq!(segment.state(), SegmentState::Sealed);
3053
3054 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
3056 assert!(segment.insert(entity2).is_err());
3057 }
3058
3059 #[test]
3060 fn test_should_seal() {
3061 let mut segment = GrowingSegment::new(1, "test");
3062
3063 let config = SegmentConfig {
3064 max_entities: 2,
3065 ..Default::default()
3066 };
3067
3068 assert!(!segment.should_seal(&config));
3069
3070 segment
3071 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
3072 .unwrap();
3073 assert!(!segment.should_seal(&config));
3074
3075 segment
3076 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
3077 .unwrap();
3078 assert!(segment.should_seal(&config));
3079 }
3080
3081 #[test]
3082 fn test_cross_references() {
3083 let mut segment = GrowingSegment::new(1, "test");
3084
3085 let mut entity1 = UnifiedEntity::table_row(
3086 EntityId::new(1),
3087 "hosts",
3088 1,
3089 vec![Value::text("192.168.1.1".to_string())],
3090 );
3091 entity1.add_cross_ref(CrossRef::new(
3092 EntityId::new(1),
3093 EntityId::new(2),
3094 "nodes",
3095 RefType::RowToNode,
3096 ));
3097 segment.insert(entity1).unwrap();
3098
3099 let refs_from = segment.get_references_from(EntityId::new(1));
3100 assert_eq!(refs_from.len(), 1);
3101 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
3102
3103 let refs_to = segment.get_references_to(EntityId::new(2));
3104 assert_eq!(refs_to.len(), 1);
3105 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
3106 }
3107
3108 #[test]
3109 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
3110 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
3111 zone.update(&Value::Email("delta@example.com".to_string()));
3112
3113 let probe = Value::Email("alpha@example.com".to_string());
3114 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
3115
3116 let in_range = Value::Email("charlie@example.com".to_string());
3117 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
3118 }
3119
3120 #[test]
3121 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
3122 let mut segment = GrowingSegment::new(1, "test");
3123
3124 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
3125 let entity = UnifiedEntity::new(
3126 EntityId::new(row_id),
3127 EntityKind::TableRow {
3128 table: "users".into(),
3129 row_id,
3130 },
3131 EntityData::Row(RowData::with_names(
3132 vec![Value::Integer(age)],
3133 vec!["age".to_string()],
3134 )),
3135 );
3136 segment.insert(entity).unwrap();
3137 }
3138
3139 segment.seal().unwrap();
3140
3141 let miss = Value::Integer(500);
3142 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
3143
3144 let hit = Value::Integer(1000);
3145 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
3146 }
3147}