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