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