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_unique_key::SegmentUniqueKey;
39use crate::storage::query::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
160fn current_unix_secs() -> u64 {
161 SystemTime::now()
162 .duration_since(UNIX_EPOCH)
163 .unwrap_or_default()
164 .as_secs()
165}
166
167fn entity_id_from_probe_key(key: &[u8]) -> Option<EntityId> {
168 if let Ok(text) = std::str::from_utf8(key) {
169 if let Ok(raw) = text.parse::<u64>() {
170 return Some(EntityId::new(raw));
171 }
172 }
173
174 let bytes: [u8; 8] = key.try_into().ok()?;
175 Some(EntityId::new(u64::from_le_bytes(bytes)))
176}
177
178const SEALED_MULTI_ZONE_MAX_INTERVALS: usize = 4;
179
180const TOMBSTONE_ENTRY_BYTES: u64 = 24;
188
189#[derive(Debug, Clone)]
190struct UpdateIndexSnapshot {
191 pk_column_name: Option<String>,
192 pk_value: Option<Value>,
193 pk_index_key: Option<(String, String)>,
194 cross_refs: Vec<CrossRef>,
195}
196
197impl UpdateIndexSnapshot {
198 fn from_entity(entity: &UnifiedEntity) -> Self {
199 let (pk_column_name, pk_value) = match &entity.data {
200 EntityData::Row(row) => (
201 row.schema
202 .as_deref()
203 .and_then(|schema| schema.first().cloned()),
204 row.columns.first().cloned(),
205 ),
206 _ => (None, None),
207 };
208 let pk_index_key = pk_value
209 .as_ref()
210 .map(|value| (entity.kind.collection().to_string(), format!("{:?}", value)));
211 Self {
212 pk_column_name,
213 pk_value,
214 pk_index_key,
215 cross_refs: entity.cross_refs().to_vec(),
216 }
217 }
218}
219
220pub trait UnifiedSegment: Send + Sync {
222 fn id(&self) -> SegmentId;
224
225 fn state(&self) -> SegmentState;
227
228 fn collection(&self) -> &str;
230
231 fn stats(&self) -> SegmentStats;
233
234 fn entity_count(&self) -> usize;
236
237 fn contains(&self, id: EntityId) -> bool;
239
240 fn get(&self, id: EntityId) -> Option<&UnifiedEntity>;
242
243 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity>;
245
246 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError>;
248
249 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError>;
251
252 fn update_hot(
256 &mut self,
257 entity: UnifiedEntity,
258 modified_columns: &[String],
259 ) -> Result<(), SegmentError> {
260 let _ = modified_columns;
261 self.update(entity)
262 }
263
264 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError>;
266
267 fn get_metadata(&self, id: EntityId) -> Option<Metadata>;
269
270 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError>;
272
273 fn seal(&mut self) -> Result<(), SegmentError>;
275
276 fn should_seal(&self, config: &SegmentConfig) -> bool;
278
279 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
281
282 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
284
285 fn filter_metadata(
287 &self,
288 filters: &[(String, super::metadata::MetadataFilter)],
289 ) -> Vec<EntityId>;
290}
291
292#[derive(Debug, Clone)]
299pub struct ColZone {
300 pub min: Value,
301 pub max: Value,
302 min_key: Option<CanonicalKey>,
303 max_key: Option<CanonicalKey>,
304}
305
306impl ColZone {
307 fn payload_bytes(&self) -> usize {
308 super::memory_size::value_heap_bytes(&self.min)
309 .saturating_add(super::memory_size::value_heap_bytes(&self.max))
310 .saturating_add(
311 self.min_key
312 .as_ref()
313 .map_or(0, super::memory_size::canonical_heap_bytes),
314 )
315 .saturating_add(
316 self.max_key
317 .as_ref()
318 .map_or(0, super::memory_size::canonical_heap_bytes),
319 )
320 }
321
322 fn new(v: Value) -> Self {
323 Self {
324 min_key: value_to_canonical_key(&v),
325 max_key: value_to_canonical_key(&v),
326 min: v.clone(),
327 max: v,
328 }
329 }
330
331 fn with_bounds(min: Value, max: Value) -> Self {
332 Self {
333 min_key: value_to_canonical_key(&min),
334 max_key: value_to_canonical_key(&max),
335 min,
336 max,
337 }
338 }
339
340 fn update(&mut self, v: &Value) {
341 if compare_zone_values(v, None, &self.min, self.min_key.as_ref())
342 .map(|o| o == std::cmp::Ordering::Less)
343 .unwrap_or(false)
344 {
345 self.min = v.clone();
346 self.min_key = value_to_canonical_key(v);
347 }
348 if compare_zone_values(v, None, &self.max, self.max_key.as_ref())
349 .map(|o| o == std::cmp::Ordering::Greater)
350 .unwrap_or(false)
351 {
352 self.max = v.clone();
353 self.max_key = value_to_canonical_key(v);
354 }
355 }
356}
357
358#[derive(Debug, Clone, Default)]
359pub struct MultiColZone {
360 pub intervals: Vec<ColZone>,
361}
362
363impl MultiColZone {
364 fn can_skip(&self, pred: &ZoneColPred<'_>) -> bool {
365 !self.intervals.is_empty() && self.intervals.iter().all(|zone| pred.can_skip(zone))
366 }
367}
368
369fn compare_zone_values(
370 left: &Value,
371 left_key: Option<&CanonicalKey>,
372 right: &Value,
373 right_key: Option<&CanonicalKey>,
374) -> Option<std::cmp::Ordering> {
375 partial_compare_values(left, right).or_else(|| {
376 let left_key = left_key.cloned().or_else(|| value_to_canonical_key(left))?;
377 let right_key = right_key
378 .cloned()
379 .or_else(|| value_to_canonical_key(right))?;
380 (left_key.family() == right_key.family()).then(|| left_key.cmp(&right_key))
381 })
382}
383
384#[derive(Debug, Clone, Copy, PartialEq, Eq)]
386pub enum ZoneColPredKind {
387 Eq,
388 Gt,
389 Gte,
390 Lt,
391 Lte,
392}
393
394#[derive(Debug, Clone)]
396pub enum ZoneColPred<'a> {
397 Eq(&'a Value),
398 Gt(&'a Value),
399 Gte(&'a Value),
400 Lt(&'a Value),
401 Lte(&'a Value),
402}
403
404impl<'a> ZoneColPred<'a> {
405 pub fn can_skip(&self, zone: &ColZone) -> bool {
407 match self {
408 ZoneColPred::Eq(val) => {
410 compare_zone_values(val, None, &zone.min, zone.min_key.as_ref())
411 .map(|o| o == std::cmp::Ordering::Less)
412 .unwrap_or(false)
413 || compare_zone_values(val, None, &zone.max, zone.max_key.as_ref())
414 .map(|o| o == std::cmp::Ordering::Greater)
415 .unwrap_or(false)
416 }
417 ZoneColPred::Gt(val) => {
419 compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
420 .map(|o| o != std::cmp::Ordering::Greater)
421 .unwrap_or(false)
422 }
423 ZoneColPred::Gte(val) => {
425 compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
426 .map(|o| o == std::cmp::Ordering::Less)
427 .unwrap_or(false)
428 }
429 ZoneColPred::Lt(val) => {
431 compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
432 .map(|o| o != std::cmp::Ordering::Less)
433 .unwrap_or(false)
434 }
435 ZoneColPred::Lte(val) => {
437 compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
438 .map(|o| o == std::cmp::Ordering::Greater)
439 .unwrap_or(false)
440 }
441 }
442 }
443}
444
445#[derive(Clone, Copy)]
447pub(crate) enum GraphEntityKind {
448 Node,
449 Edge,
450}
451
452impl GraphEntityKind {
453 fn mask(self) -> u8 {
454 match self {
455 Self::Node => 1,
456 Self::Edge => 2,
457 }
458 }
459
460 pub(crate) fn matches(self, kind: &EntityKind) -> bool {
461 matches!(
462 (self, kind),
463 (Self::Node, EntityKind::GraphNode(_)) | (Self::Edge, EntityKind::GraphEdge(_))
464 )
465 }
466}
467
468pub(crate) const SCAN_BATCH_SIZE: usize = 256;
470
471pub(crate) struct SegmentScanCursor {
473 flat_end: usize,
474 hash_end: usize,
475 deleted_end: usize,
476 position: usize,
477}
478
479impl SegmentScanCursor {
480 pub(crate) fn remaining(&self) -> usize {
481 self.flat_end + self.hash_end - self.position
482 }
483
484 pub(crate) fn finished(&self) -> bool {
485 self.remaining() == 0
486 }
487}
488
489pub(crate) struct GraphIndexCursor {
492 key_index: usize,
493 after: Option<EntityId>,
494 upper: EntityId,
495}
496
497impl GraphIndexCursor {
498 pub(crate) fn finished(&self, keys: &[EntityId]) -> bool {
499 self.key_index == keys.len()
500 }
501}
502
503pub struct GrowingSegment {
505 id: SegmentId,
507 collection: String,
509 state: SegmentState,
511 created_at: u64,
513 last_write_at: u64,
515
516 entities: HashMap<EntityId, UnifiedEntity>,
518 hash_scan_ids: Vec<EntityId>,
521 flat_entities: Vec<UnifiedEntity>,
524 base_entity_id: u64,
526 use_flat: bool,
528 deleted: HashMap<EntityId, usize>,
532 metadata: MetadataStorage,
534
535 pk_index: BTreeMap<(String, String), EntityId>,
537 unique_key_lookup: parking_lot::RwLock<Vec<SegmentUniqueKey>>,
540 unique_key_lookup_bytes: AtomicU64,
541 kind_index: HashMap<String, HashSet<EntityId>>,
543 graph_kind_mask: u8,
546 graph_read_index: parking_lot::RwLock<Option<SegmentGraphIndex>>,
548 graph_read_index_bytes: AtomicU64,
549 cross_ref_forward: HashMap<EntityId, Vec<(EntityId, RefType)>>,
551 cross_ref_reverse: HashMap<EntityId, Vec<(EntityId, RefType)>>,
553
554 col_zones: HashMap<String, ColZone>,
556 zone_payload_bytes: usize,
557 sealed_col_zones: HashMap<String, MultiColZone>,
559
560 sequence: AtomicU64,
562 checkpoint_revision: u64,
564 memory_bytes: AtomicU64,
569 dead_entity_bytes: u64,
572 dead_resident_count: usize,
574
575 pub(crate) published_flat_len: AtomicUsize,
582}
583
584impl GrowingSegment {
585 pub(crate) 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(crate) 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(crate) 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 kind_index: HashMap::new(),
724 graph_read_index: parking_lot::RwLock::new(Some(SegmentGraphIndex::default())),
725 graph_read_index_bytes: AtomicU64::new(0),
726 graph_kind_mask: 0,
727 cross_ref_forward: HashMap::new(),
728 cross_ref_reverse: HashMap::new(),
729 col_zones: HashMap::new(),
730 zone_payload_bytes: 0,
731 sealed_col_zones: HashMap::new(),
732 sequence: AtomicU64::new(0),
733 checkpoint_revision: 0,
734 memory_bytes: AtomicU64::new(0),
735 dead_entity_bytes: 0,
736 dead_resident_count: 0,
737 published_flat_len: AtomicUsize::new(0),
738 }
739 }
740
741 fn next_sequence(&self) -> u64 {
743 self.sequence.fetch_add(1, Ordering::SeqCst)
744 }
745
746 pub(crate) fn checkpoint_revision(&self) -> u64 {
747 self.checkpoint_revision
748 }
749
750 fn mark_checkpoint_dirty(&mut self) {
751 self.checkpoint_revision = self
752 .checkpoint_revision
753 .checked_add(1)
754 .expect("segment mutation counter must not overflow");
755 }
756
757 fn has_live_entity(&self, id: EntityId) -> bool {
758 if self.deleted.contains_key(&id) {
759 return false;
760 }
761 if self.use_flat {
762 let raw = id.raw();
763 if raw >= self.base_entity_id {
764 let idx = (raw - self.base_entity_id) as usize;
765 if self
766 .flat_entities
767 .get(idx)
768 .is_some_and(|entity| entity.id == id)
769 {
770 return true;
771 }
772 }
773 self.entities.contains_key(&id)
776 } else {
777 self.entities.contains_key(&id)
778 }
779 }
780
781 fn update_existing_entity_in_place(
782 &mut self,
783 entity: &UnifiedEntity,
784 ) -> Result<UpdateIndexSnapshot, SegmentError> {
785 self.mark_checkpoint_dirty();
786 self.invalidate_unique_key_lookup();
787 if (self.graph_kind_mask != 0
788 || matches!(
789 entity.kind,
790 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
791 ))
792 && self
793 .get(entity.id)
794 .is_some_and(|old| !SegmentGraphIndex::same_keys(old, entity))
795 {
796 self.invalidate_graph_read_index();
797 }
798 self.observe_graph_kind(&entity.kind);
801 if self.use_flat {
802 let raw = entity.id.raw();
803 if raw >= self.base_entity_id {
804 let idx = (raw - self.base_entity_id) as usize;
805 if let Some(slot) = self
806 .flat_entities
807 .get_mut(idx)
808 .filter(|slot| slot.id == entity.id)
809 {
810 let snapshot = UpdateIndexSnapshot::from_entity(slot);
811 let old_bytes = Self::estimate_entity_size(slot);
812 slot.clone_from(entity);
813 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
814 return Ok(snapshot);
815 }
816 }
817 let Some(slot) = self.entities.get_mut(&entity.id) else {
818 return Err(SegmentError::NotFound(entity.id));
819 };
820 let snapshot = UpdateIndexSnapshot::from_entity(slot);
821 let old_bytes = Self::estimate_entity_size(slot);
822 slot.clone_from(entity);
823 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
824 Ok(snapshot)
825 } else {
826 let Some(slot) = self.entities.get_mut(&entity.id) else {
827 return Err(SegmentError::NotFound(entity.id));
828 };
829 let snapshot = UpdateIndexSnapshot::from_entity(slot);
830 let old_bytes = Self::estimate_entity_size(slot);
831 slot.clone_from(entity);
832 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
833 Ok(snapshot)
834 }
835 }
836
837 fn apply_hot_update_with_metadata(
838 &mut self,
839 entity: &UnifiedEntity,
840 modified_columns: &[String],
841 metadata: Option<&Metadata>,
842 ) -> Result<(), SegmentError> {
843 let old = self.update_existing_entity_in_place(entity)?;
844 self.reindex_for_update(&old, entity, Some(modified_columns));
845 self.update_col_zones_from_entity(entity);
846 if let Some(metadata) = metadata {
847 self.metadata.set_all(entity.id, metadata);
848 }
849 Ok(())
850 }
851
852 fn apply_update_with_metadata(
853 &mut self,
854 entity: &UnifiedEntity,
855 metadata: Option<&Metadata>,
856 ) -> Result<(), SegmentError> {
857 let old = self.update_existing_entity_in_place(entity)?;
858 self.reindex_for_update(&old, entity, None);
859 self.update_col_zones_from_entity(entity);
860 if let Some(metadata) = metadata {
861 self.metadata.set_all(entity.id, metadata);
862 }
863 Ok(())
864 }
865
866 pub fn update_hot_batch_with_metadata<'a, I>(&mut self, items: I) -> Result<(), SegmentError>
867 where
868 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
869 {
870 if !self.state.is_writable() {
871 return Err(SegmentError::NotWritable);
872 }
873
874 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
875 items.into_iter().collect();
876 if items.is_empty() {
877 return Ok(());
878 }
879
880 for (entity, _, _) in &items {
881 if !self.has_live_entity(entity.id) {
882 return Err(SegmentError::NotFound(entity.id));
883 }
884 }
885
886 for (entity, modified_columns, metadata) in items {
887 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)?;
888 }
889
890 self.last_write_at = current_unix_secs();
891 Ok(())
892 }
893
894 fn record_tombstone(&mut self, id: EntityId) -> bool {
896 self.mark_checkpoint_dirty();
897 let ordinal = self.deleted.len();
898 match self.deleted.entry(id) {
899 std::collections::hash_map::Entry::Vacant(entry) => {
900 entry.insert(ordinal);
901 true
902 }
903 std::collections::hash_map::Entry::Occupied(_) => false,
904 }
905 }
906
907 fn tombstone_flat_entity(&mut self, idx: usize, id: EntityId) -> bool {
913 self.invalidate_unique_key_lookup();
914 if !self.record_tombstone(id) {
915 return false;
916 }
917 if let Some(index) = self.graph_read_index.get_mut() {
918 index.remove(&self.flat_entities[idx]);
919 self.graph_read_index_bytes
920 .store(index.memory_bytes(), Ordering::Relaxed);
921 }
922 let size = Self::estimate_entity_size(&self.flat_entities[idx]) as u64;
923 self.dead_entity_bytes += size;
924 self.dead_resident_count += 1;
925 self.metadata.remove_all(id);
926 true
927 }
928
929 fn remove_hashmap_entity(&mut self, id: EntityId) -> bool {
934 self.invalidate_unique_key_lookup();
935 let Some(entity) = self.entities.remove(&id) else {
936 return false;
937 };
938 self.release_memory(Self::estimate_entity_size(&entity));
939 self.unindex_entity(&entity);
940 self.metadata.remove_all(id);
941 self.record_tombstone(id);
942 true
943 }
944
945 pub fn delete_batch(&mut self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
946 if !self.state.is_writable() {
947 return Err(SegmentError::NotWritable);
948 }
949 if ids.is_empty() {
950 return Ok(Vec::new());
951 }
952
953 let mut deleted_ids = Vec::with_capacity(ids.len());
954
955 if self.use_flat {
956 for &id in ids {
957 let raw = id.raw();
958 if raw < self.base_entity_id {
959 if self.remove_hashmap_entity(id) {
960 deleted_ids.push(id);
961 }
962 continue;
963 }
964 let idx = (raw - self.base_entity_id) as usize;
965 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
966 if self.tombstone_flat_entity(idx, id) {
967 deleted_ids.push(id);
968 }
969 } else if self.remove_hashmap_entity(id) {
970 deleted_ids.push(id);
971 }
972 }
973 } else {
974 for &id in ids {
975 if self.remove_hashmap_entity(id) {
976 deleted_ids.push(id);
977 }
978 }
979 }
980
981 if !deleted_ids.is_empty() {
982 self.last_write_at = current_unix_secs();
983 }
984
985 Ok(deleted_ids)
986 }
987
988 fn add_memory(&self, bytes: usize) {
990 self.memory_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
991 }
992
993 fn replace_entity_memory(&self, previous: usize, current: usize) {
994 if current >= previous {
995 self.add_memory(current - previous);
996 } else {
997 self.release_memory(previous - current);
998 }
999 }
1000
1001 pub fn memory_bytes(&self) -> u64 {
1006 self.memory_bytes.load(Ordering::Relaxed)
1007 + self.zone_payload_bytes as u64
1008 + self.unique_key_lookup_bytes.load(Ordering::Relaxed)
1009 + self.graph_read_index_bytes.load(Ordering::Relaxed)
1010 + (self.hash_scan_ids.capacity() * std::mem::size_of::<EntityId>()) as u64
1011 }
1012
1013 fn release_memory(&self, bytes: usize) {
1016 let mut current = self.memory_bytes.load(Ordering::Relaxed);
1017 loop {
1018 let next = current.saturating_sub(bytes as u64);
1019 match self.memory_bytes.compare_exchange_weak(
1020 current,
1021 next,
1022 Ordering::Relaxed,
1023 Ordering::Relaxed,
1024 ) {
1025 Ok(_) => break,
1026 Err(observed) => current = observed,
1027 }
1028 }
1029 }
1030
1031 pub(crate) fn resident_entity_count(&self) -> usize {
1033 self.flat_entities.len() + self.entities.len()
1034 }
1035
1036 pub(crate) fn live_entity_count(&self) -> usize {
1038 self.resident_entity_count()
1039 .saturating_sub(self.dead_resident_count)
1040 }
1041
1042 pub(crate) fn tombstone_count(&self) -> usize {
1045 self.deleted.len()
1046 }
1047
1048 pub(crate) fn resident_bytes(&self) -> u64 {
1051 self.memory_bytes() + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1052 }
1053
1054 pub(crate) fn reclaimable_bytes(&self) -> u64 {
1057 self.dead_entity_bytes + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1058 }
1059
1060 pub(crate) fn live_entity_ids(&self) -> Vec<EntityId> {
1066 let mut ids = Vec::with_capacity(self.live_entity_count());
1067 self.for_each_fast(|entity| {
1068 ids.push(entity.id);
1069 true
1070 });
1071 ids
1072 }
1073
1074 pub(crate) fn adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
1082 let id = entity.id;
1083
1084 let next = entity.sequence_id.saturating_add(1);
1087 if self.sequence.load(Ordering::Relaxed) < next {
1088 self.sequence.store(next, Ordering::Relaxed);
1089 }
1090
1091 self.add_memory(Self::estimate_entity_size(&entity));
1092 self.index_entity(&entity);
1093 self.update_col_zones_from_entity(&entity);
1094 self.entities.insert(id, entity);
1095
1096 if let Some(metadata) = metadata {
1097 if !metadata.is_empty() {
1098 self.metadata.set_all(id, &metadata);
1099 }
1100 }
1101 }
1102
1103 pub(crate) fn evict_entity(&mut self, id: EntityId) -> bool {
1109 self.mark_checkpoint_dirty();
1110 self.invalidate_unique_key_lookup();
1111 let Some(entity) = self.entities.remove(&id) else {
1112 return false;
1113 };
1114 self.release_memory(Self::estimate_entity_size(&entity));
1115 self.unindex_entity(&entity);
1116 for cross_ref in entity.cross_refs() {
1117 if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1118 reverse.retain(|(source, _)| *source != id);
1119 }
1120 }
1121 self.metadata.remove_all(id);
1122 true
1123 }
1124
1125 fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
1127 super::memory_size::entity_bytes(entity)
1128 }
1129
1130 fn update_column_zone(&mut self, column: &str, value: &Value) {
1138 if matches!(value, Value::Null) {
1139 return;
1140 }
1141 match self.col_zones.entry(column.to_string()) {
1142 std::collections::hash_map::Entry::Occupied(mut entry) => {
1143 let before = entry.get().payload_bytes();
1144 entry.get_mut().update(value);
1145 self.zone_payload_bytes = self
1146 .zone_payload_bytes
1147 .saturating_sub(before)
1148 .saturating_add(entry.get().payload_bytes());
1149 }
1150 std::collections::hash_map::Entry::Vacant(entry) => {
1151 let zone = ColZone::new(value.clone());
1152 self.zone_payload_bytes =
1153 self.zone_payload_bytes.saturating_add(zone.payload_bytes());
1154 entry.insert(zone);
1155 }
1156 }
1157 }
1158
1159 fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
1160 if let EntityData::Row(row) = &entity.data {
1161 if let Some(named) = &row.named {
1162 for (column, value) in named {
1163 self.update_column_zone(column, value);
1164 }
1165 } else if let Some(schema) = &row.schema {
1166 for (column, value) in schema.iter().zip(&row.columns) {
1167 self.update_column_zone(column, value);
1168 }
1169 }
1170 }
1171 }
1172
1173 fn rebuild_sealed_col_zones(&mut self) {
1174 let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
1175 let mut family_by_col: HashMap<String, reddb_types::CanonicalKeyFamily> = HashMap::new();
1176 let mut mixed_family_cols = HashSet::new();
1177 let mut unsupported_cols = HashSet::new();
1178
1179 let mut observe_row = |row: &super::entity::RowData| {
1180 for (col, value) in row.iter_fields() {
1181 if matches!(value, Value::Null) {
1182 continue;
1183 }
1184 let Some(key) = value_to_canonical_key(value) else {
1185 unsupported_cols.insert(col.to_string());
1186 continue;
1187 };
1188 match family_by_col.get(col).copied() {
1189 Some(existing) if existing != key.family() => {
1190 mixed_family_cols.insert(col.to_string());
1191 }
1192 None => {
1193 family_by_col.insert(col.to_string(), key.family());
1194 }
1195 _ => {}
1196 }
1197 values_by_col
1198 .entry(col.to_string())
1199 .or_default()
1200 .push((key, value.clone()));
1201 }
1202 };
1203
1204 if self.use_flat {
1205 for entity in &self.flat_entities {
1206 if self.deleted.contains_key(&entity.id) {
1207 continue;
1208 }
1209 if let EntityData::Row(row) = &entity.data {
1210 observe_row(row);
1211 }
1212 }
1213 } else {
1214 for entity in self.entities.values() {
1215 if self.deleted.contains_key(&entity.id) {
1216 continue;
1217 }
1218 if let EntityData::Row(row) = &entity.data {
1219 observe_row(row);
1220 }
1221 }
1222 }
1223
1224 let mut sealed_col_zones = HashMap::new();
1225 for (col, mut entries) in values_by_col {
1226 if mixed_family_cols.contains(&col)
1227 || unsupported_cols.contains(&col)
1228 || entries.is_empty()
1229 {
1230 continue;
1231 }
1232 entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
1233 entries.dedup_by(|left, right| left.0 == right.0);
1234
1235 let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
1236 if intervals.len() > 1 {
1237 sealed_col_zones.insert(col, MultiColZone { intervals });
1238 }
1239 }
1240
1241 let previous = self
1242 .sealed_col_zones
1243 .values()
1244 .flat_map(|zones| &zones.intervals)
1245 .map(ColZone::payload_bytes)
1246 .fold(0, usize::saturating_add);
1247 let current = sealed_col_zones
1248 .values()
1249 .flat_map(|zones| &zones.intervals)
1250 .map(ColZone::payload_bytes)
1251 .fold(0, usize::saturating_add);
1252 self.zone_payload_bytes = self
1253 .zone_payload_bytes
1254 .saturating_sub(previous)
1255 .saturating_add(current);
1256 self.sealed_col_zones = sealed_col_zones;
1257 }
1258
1259 pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
1263 if preds.is_empty() {
1264 return false;
1265 }
1266 for (col, pred) in preds {
1267 if let Some(zone) = self.sealed_col_zones.get(*col) {
1268 if zone.can_skip(pred) {
1269 return true;
1270 }
1271 continue;
1272 }
1273 if let Some(zone) = self.col_zones.get(*col) {
1274 if pred.can_skip(zone) {
1275 return true; }
1277 }
1278 }
1279 false
1280 }
1281
1282 fn index_entity(&mut self, entity: &UnifiedEntity) {
1284 self.mark_checkpoint_dirty();
1285 self.hash_scan_ids.push(entity.id);
1286 self.observe_graph_kind(&entity.kind);
1287 if let Some(index) = self.graph_read_index.get_mut() {
1288 index.insert(entity);
1289 self.graph_read_index_bytes
1290 .store(index.memory_bytes(), Ordering::Relaxed);
1291 }
1292 let indexes = self.unique_key_lookup.get_mut();
1293 for index in indexes.iter_mut() {
1294 index.insert(entity);
1295 }
1296 self.unique_key_lookup_bytes.store(
1297 indexes
1298 .iter()
1299 .map(|index| index.memory_bytes() as u64)
1300 .sum(),
1301 Ordering::Relaxed,
1302 );
1303 let kind_key = entity.kind.storage_type().to_string();
1305 self.kind_index
1306 .entry(kind_key)
1307 .or_default()
1308 .insert(entity.id);
1309
1310 if let EntityData::Row(row) = &entity.data {
1312 if let Some(first_col) = row.columns.first() {
1313 let pk_str = format!("{:?}", first_col);
1314 self.pk_index
1315 .insert((entity.kind.collection().to_string(), pk_str), entity.id);
1316 }
1317 }
1318
1319 for cross_ref in entity.cross_refs() {
1321 self.cross_ref_forward
1322 .entry(cross_ref.source)
1323 .or_default()
1324 .push((cross_ref.target, cross_ref.ref_type));
1325
1326 self.cross_ref_reverse
1327 .entry(cross_ref.target)
1328 .or_default()
1329 .push((cross_ref.source, cross_ref.ref_type));
1330 }
1331 }
1332
1333 fn invalidate_graph_read_index(&mut self) {
1334 *self.graph_read_index.get_mut() = None;
1335 self.graph_read_index_bytes.store(0, Ordering::Relaxed);
1336 }
1337
1338 #[cfg(test)]
1341 pub(crate) fn visit_graph_index_candidates(
1342 &self,
1343 kind: GraphEntityKind,
1344 keys: &[EntityId],
1345 before_work: &mut impl FnMut() -> bool,
1346 visit: &mut impl FnMut(&UnifiedEntity) -> bool,
1347 ) -> bool {
1348 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1349 return false;
1350 };
1351 let index = cached.as_ref().expect("complete graph index");
1352 for key in keys {
1353 if !before_work() {
1354 return false;
1355 }
1356 for id in index.candidates(kind, *key) {
1357 if !before_work() {
1358 return false;
1359 }
1360 if let Some(entity) = self.get(id) {
1361 if !visit(entity) {
1362 return false;
1363 }
1364 }
1365 }
1366 }
1367 true
1368 }
1369
1370 fn graph_read_index_for_query(
1371 &self,
1372 before_work: &mut impl FnMut() -> bool,
1373 ) -> Option<parking_lot::RwLockReadGuard<'_, Option<SegmentGraphIndex>>> {
1374 let mut cached = self.graph_read_index.read();
1375 if cached.is_none() {
1376 drop(cached);
1377 let mut writer = self.graph_read_index.write();
1378 if writer.is_none() {
1379 let mut index = SegmentGraphIndex::default();
1380 if !self.for_each_fast(|entity| {
1383 if !before_work() {
1384 return false;
1385 }
1386 index.insert(entity);
1387 true
1388 }) {
1389 return None;
1390 }
1391 self.graph_read_index_bytes
1392 .store(index.memory_bytes(), Ordering::Relaxed);
1393 *writer = Some(index);
1394 }
1395 cached = parking_lot::RwLockWriteGuard::downgrade(writer);
1396 }
1397 Some(cached)
1398 }
1399
1400 pub(crate) fn graph_index_cursor(
1401 &self,
1402 kind: GraphEntityKind,
1403 keys: &[EntityId],
1404 before_work: &mut impl FnMut() -> bool,
1405 ) -> Option<GraphIndexCursor> {
1406 let cached = self.graph_read_index_for_query(before_work)?;
1407 let index = cached.as_ref().expect("complete graph index");
1408 let mut upper = EntityId::new(0);
1409 for key in keys {
1410 if !before_work() {
1411 return None;
1412 }
1413 if let Some(last) = index.candidates(kind, *key).next_back() {
1414 upper = upper.max(last);
1415 }
1416 }
1417 Some(GraphIndexCursor {
1418 key_index: 0,
1419 after: None,
1420 upper,
1421 })
1422 }
1423
1424 pub(crate) fn graph_index_batch(
1427 &self,
1428 kind: GraphEntityKind,
1429 keys: &[EntityId],
1430 cursor: &mut GraphIndexCursor,
1431 before_work: &mut impl FnMut() -> bool,
1432 visit: &mut impl FnMut(EntityId),
1433 ) -> bool {
1434 self.graph_index_batch_with_limit(kind, keys, cursor, before_work, visit, SCAN_BATCH_SIZE)
1435 }
1436
1437 pub(crate) fn graph_index_batch_with_limit(
1438 &self,
1439 kind: GraphEntityKind,
1440 keys: &[EntityId],
1441 cursor: &mut GraphIndexCursor,
1442 before_work: &mut impl FnMut() -> bool,
1443 visit: &mut impl FnMut(EntityId),
1444 limit: usize,
1445 ) -> bool {
1446 assert!(
1447 limit > 0 && limit <= SCAN_BATCH_SIZE,
1448 "bounded graph index batch"
1449 );
1450 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1451 return false;
1452 };
1453 let index = cached.as_ref().expect("complete graph index");
1454 let mut count = 0;
1455 while !cursor.finished(keys) && count < limit {
1456 let key = keys[cursor.key_index];
1457 if !before_work() {
1458 return false;
1459 }
1460 for id in index
1461 .candidates_after(kind, key, cursor.after, cursor.upper)
1462 .take(limit - count)
1463 {
1464 if !before_work() {
1465 return false;
1466 }
1467 cursor.after = Some(id);
1468 count += 1;
1469 visit(id);
1470 }
1471 if count == limit && cursor.after != Some(cursor.upper) {
1472 break;
1473 }
1474 cursor.key_index += 1;
1475 cursor.after = None;
1476 }
1477 true
1478 }
1479
1480 pub(crate) fn find_unique_key_conflict(
1485 &self,
1486 columns: &[Vec<String>],
1487 key_index: usize,
1488 signatures: &[String],
1489 exclude_id: Option<EntityId>,
1490 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1491 ) -> Option<EntityId> {
1492 assert!(
1493 key_index < columns.len(),
1494 "constraint index belongs to current schema"
1495 );
1496 let cached = self.unique_key_lookup.read();
1497 if cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1498 return self.find_unique_key_candidate(
1499 &cached[key_index],
1500 signatures,
1501 exclude_id,
1502 reserves_key,
1503 );
1504 }
1505 drop(cached);
1506 let mut cached = self.unique_key_lookup.write();
1507 if !cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1508 let mut indexes: Vec<_> = columns
1509 .iter()
1510 .map(|columns| SegmentUniqueKey::new(columns))
1511 .collect();
1512 self.for_each_fast(|entity| {
1513 for index in &mut indexes {
1514 index.insert(entity);
1515 }
1516 true
1517 });
1518 self.unique_key_lookup_bytes.store(
1519 indexes
1520 .iter()
1521 .map(|index| index.memory_bytes() as u64)
1522 .sum(),
1523 Ordering::Relaxed,
1524 );
1525 *cached = indexes;
1526 }
1527 self.find_unique_key_candidate(&cached[key_index], signatures, exclude_id, reserves_key)
1528 }
1529
1530 fn find_unique_key_candidate(
1531 &self,
1532 index: &SegmentUniqueKey,
1533 signatures: &[String],
1534 exclude_id: Option<EntityId>,
1535 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1536 ) -> Option<EntityId> {
1537 index.candidates(signatures).iter().copied().find(|id| {
1538 if Some(*id) == exclude_id {
1539 return false;
1540 }
1541 let Some(
1542 entity @ UnifiedEntity {
1543 data: EntityData::Row(row),
1544 ..
1545 },
1546 ) = self.get(*id)
1547 else {
1548 return false;
1549 };
1550 index
1551 .columns
1552 .iter()
1553 .zip(signatures)
1554 .all(|(column, signature)| {
1555 row.get_field(column)
1556 .is_some_and(|value| format!("{value:?}") == *signature)
1557 })
1558 && reserves_key(entity)
1559 })
1560 }
1561
1562 fn invalidate_unique_key_lookup(&mut self) {
1563 *self.unique_key_lookup.get_mut() = Vec::new();
1564 self.unique_key_lookup_bytes.store(0, Ordering::Relaxed);
1565 }
1566
1567 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1569 if let Some(id) = entity_id_from_probe_key(key) {
1570 return self.has_live_entity(id);
1571 }
1572
1573 let Ok(pk_str) = std::str::from_utf8(key) else {
1574 return false;
1575 };
1576 let pk_key = (self.collection.clone(), pk_str.to_string());
1577 self.pk_index
1578 .get(&pk_key)
1579 .is_some_and(|id| self.has_live_entity(*id))
1580 }
1581
1582 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1584 if let Some(index) = self.graph_read_index.get_mut() {
1585 index.remove(entity);
1586 self.graph_read_index_bytes
1587 .store(index.memory_bytes(), Ordering::Relaxed);
1588 }
1589 let kind_key = entity.kind.storage_type().to_string();
1591 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1592 set.remove(&entity.id);
1593 }
1594
1595 if let EntityData::Row(row) = &entity.data {
1597 if let Some(first_col) = row.columns.first() {
1598 let pk_str = format!("{:?}", first_col);
1599 self.pk_index
1600 .remove(&(entity.kind.collection().to_string(), pk_str));
1601 }
1602 }
1603
1604 self.cross_ref_forward.remove(&entity.id);
1606 }
1608
1609 fn reindex_for_update(
1619 &mut self,
1620 old: &UpdateIndexSnapshot,
1621 new: &UnifiedEntity,
1622 modified_columns: Option<&[String]>,
1623 ) {
1624 let pk_changed = match &new.data {
1629 EntityData::Row(new_row) => {
1630 if let Some(cols) = modified_columns {
1631 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1634 new_row
1635 .schema
1636 .as_deref()
1637 .and_then(|schema| schema.first().map(|name| name.as_str()))
1638 });
1639 match pk_col_name {
1640 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1641 None => old.pk_value.as_ref() != new_row.columns.first(),
1643 }
1644 } else {
1645 old.pk_value.as_ref() != new_row.columns.first()
1646 }
1647 }
1648 _ => false,
1650 };
1651
1652 if pk_changed {
1653 if let Some((collection, pk_str)) = &old.pk_index_key {
1655 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1656 }
1657 if let EntityData::Row(row) = &new.data {
1659 if let Some(first_col) = row.columns.first() {
1660 let pk_str = format!("{:?}", first_col);
1661 self.pk_index
1662 .insert((new.kind.collection().to_string(), pk_str), new.id);
1663 }
1664 }
1665 }
1666
1667 let new_refs = new.cross_refs();
1669 if old.cross_refs.as_slice() != new_refs {
1670 self.cross_ref_forward.remove(&new.id);
1672 for cross_ref in &old.cross_refs {
1674 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1675 rev.retain(|(src, _)| *src != new.id);
1676 }
1677 }
1678 for cross_ref in new_refs {
1680 self.cross_ref_forward
1681 .entry(cross_ref.source)
1682 .or_default()
1683 .push((cross_ref.target, cross_ref.ref_type));
1684 self.cross_ref_reverse
1685 .entry(cross_ref.target)
1686 .or_default()
1687 .push((cross_ref.source, cross_ref.ref_type));
1688 }
1689 }
1690 }
1691
1692 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1694 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1695 }
1696
1697 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1699 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1700 }
1701
1702 pub fn age_secs(&self) -> u64 {
1704 let now = current_unix_secs();
1705 now.saturating_sub(self.created_at)
1706 }
1707
1708 pub fn idle_secs(&self) -> u64 {
1710 let now = current_unix_secs();
1711 now.saturating_sub(self.last_write_at)
1712 }
1713
1714 pub fn bulk_insert(
1723 &mut self,
1724 entities: Vec<UnifiedEntity>,
1725 ) -> Result<Vec<EntityId>, SegmentError> {
1726 self.mark_checkpoint_dirty();
1727 if !self.state.is_writable() {
1728 return Err(SegmentError::NotWritable);
1729 }
1730
1731 let indexes = self.unique_key_lookup.get_mut();
1732 for index in indexes.iter_mut() {
1733 for entity in &entities {
1734 index.insert(entity);
1735 }
1736 }
1737 self.unique_key_lookup_bytes.store(
1738 indexes
1739 .iter()
1740 .map(|index| index.memory_bytes() as u64)
1741 .sum(),
1742 Ordering::Relaxed,
1743 );
1744
1745 let n = entities.len();
1746
1747 if entities.is_empty() {
1748 return Ok(Vec::new());
1749 }
1750
1751 for run in
1754 entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
1755 {
1756 let kind = run[0].kind.storage_type();
1757 self.observe_graph_kind(&run[0].kind);
1758 if matches!(
1759 run[0].kind,
1760 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
1761 ) {
1762 if let Some(index) = self.graph_read_index.get_mut() {
1763 for entity in run {
1764 index.insert(entity);
1765 }
1766 self.graph_read_index_bytes
1767 .store(index.memory_bytes(), Ordering::Relaxed);
1768 }
1769 }
1770 let kind_set = match self.kind_index.get_mut(kind) {
1771 Some(ids) => ids,
1772 None => self.kind_index.entry(kind.to_string()).or_default(),
1773 };
1774 kind_set.reserve(run.len());
1775 kind_set.extend(run.iter().map(|entity| entity.id));
1776 }
1777
1778 let now = current_unix_secs();
1779
1780 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
1781
1782 let mut ids = Vec::with_capacity(n);
1783
1784 if self.flat_entities.is_empty() && self.entities.is_empty() {
1786 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
1788 self.use_flat = true;
1789 }
1790
1791 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
1801 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
1802 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
1806
1807 let mut batch_bytes: usize = 0;
1812
1813 if self.use_flat {
1814 self.flat_entities.reserve(n);
1815 for (i, mut entity) in entities.into_iter().enumerate() {
1816 entity.sequence_id = base_seq + i as u64;
1817 let id = entity.id;
1818 ids.push(id);
1819 batch_bytes += Self::estimate_entity_size(&entity);
1820 if let EntityData::Row(row) = &entity.data {
1821 if row.schema.is_some() && !row.columns.is_empty() {
1822 if columnar_zone_updates.is_empty() {
1823 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
1824 columnar_schema = row.schema.clone();
1825 }
1826 for (ci, val) in row.columns.iter().enumerate() {
1827 if !matches!(val, Value::Null) {
1828 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
1829 bucket.push(val.clone());
1830 }
1831 }
1832 }
1833 } else {
1834 for (col, val) in row.iter_fields() {
1835 if !matches!(val, Value::Null) {
1836 named_zone_updates.push((col.to_string(), val.clone()));
1837 }
1838 }
1839 }
1840 }
1841 let expected = self.base_entity_id + self.flat_entities.len() as u64;
1857 if id.raw() == expected {
1858 self.flat_entities.push(entity);
1859 } else {
1860 self.hash_scan_ids.push(id);
1861 self.entities.insert(id, entity);
1862 }
1863 }
1864 } else {
1865 self.entities.reserve(n);
1867 self.hash_scan_ids.reserve(n);
1868 let mut pairs = Vec::with_capacity(n);
1869 for (i, mut entity) in entities.into_iter().enumerate() {
1870 entity.sequence_id = base_seq + i as u64;
1871 let id = entity.id;
1872 ids.push(id);
1873 batch_bytes += Self::estimate_entity_size(&entity);
1874 if let EntityData::Row(row) = &entity.data {
1875 for (col, val) in row.iter_fields() {
1876 if !matches!(val, Value::Null) {
1877 named_zone_updates.push((col.to_string(), val.clone()));
1878 }
1879 }
1880 }
1881 self.hash_scan_ids.push(id);
1882 pairs.push((id, entity));
1883 }
1884 self.entities.extend(pairs);
1885 }
1886
1887 if !columnar_zone_updates.is_empty() {
1891 let schema = columnar_schema.as_ref();
1892 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
1893 if values.is_empty() {
1894 continue;
1895 }
1896 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
1897 continue;
1898 };
1899 let mut iter = values.into_iter();
1905 if let Some(first) = iter.next() {
1906 let previous = self
1907 .col_zones
1908 .get(col_name)
1909 .map_or(0, ColZone::payload_bytes);
1910 let zone = self
1911 .col_zones
1912 .entry(col_name.clone())
1913 .and_modify(|z| z.update(&first))
1914 .or_insert_with(|| ColZone::new(first));
1915 for v in iter {
1916 zone.update(&v);
1917 }
1918 self.zone_payload_bytes = self
1919 .zone_payload_bytes
1920 .saturating_sub(previous)
1921 .saturating_add(zone.payload_bytes());
1922 }
1923 }
1924 }
1925 for (col, val) in named_zone_updates {
1926 self.update_column_zone(&col, &val);
1927 }
1928
1929 self.add_memory(batch_bytes);
1930 self.last_write_at = now;
1931
1932 if self.use_flat {
1934 self.published_flat_len
1935 .store(self.flat_entities.len(), Ordering::Release);
1936 }
1937
1938 Ok(ids)
1939 }
1940
1941 pub(crate) fn force_delete(&mut self, id: EntityId) -> bool {
1944 if self.use_flat {
1945 let raw = id.raw();
1946 if raw >= self.base_entity_id {
1947 let idx = (raw - self.base_entity_id) as usize;
1948 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
1949 self.tombstone_flat_entity(idx, id);
1950 return true;
1951 }
1952 }
1953 return self.remove_hashmap_entity(id);
1954 }
1955
1956 self.remove_hashmap_entity(id)
1957 }
1958
1959 pub(crate) fn force_update_with_metadata(
1962 &mut self,
1963 entity: &UnifiedEntity,
1964 modified_columns: &[String],
1965 metadata: Option<&Metadata>,
1966 ) -> Result<(), SegmentError> {
1967 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
1968 }
1969}
1970
1971impl UnifiedSegment for GrowingSegment {
1972 fn id(&self) -> SegmentId {
1973 self.id
1974 }
1975
1976 fn state(&self) -> SegmentState {
1977 self.state
1978 }
1979
1980 fn collection(&self) -> &str {
1981 &self.collection
1982 }
1983
1984 fn stats(&self) -> SegmentStats {
1985 let mut stats = SegmentStats {
1986 entity_count: self.live_entity_count(),
1987 deleted_count: self.deleted.len(),
1988 memory_bytes: self.resident_bytes() as usize,
1989 ..Default::default()
1990 };
1991
1992 for entity in self.entities.values() {
1993 match &entity.kind {
1994 EntityKind::TableRow { .. } => stats.row_count += 1,
1995 EntityKind::GraphNode(_) => stats.node_count += 1,
1996 EntityKind::GraphEdge(_) => stats.edge_count += 1,
1997 EntityKind::Vector { .. } => stats.vector_count += 1,
1998 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
1999 EntityKind::QueueMessage { .. } => stats.row_count += 1,
2000 }
2001 stats.cross_ref_count += entity.cross_refs().len();
2002 }
2003
2004 stats
2005 }
2006
2007 fn entity_count(&self) -> usize {
2008 self.live_entity_count()
2009 }
2010
2011 fn contains(&self, id: EntityId) -> bool {
2012 self.has_live_entity(id)
2013 }
2014
2015 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2016 if self.deleted.contains_key(&id) {
2017 return None;
2018 }
2019 if self.use_flat {
2020 let raw = id.raw();
2021 if raw >= self.base_entity_id {
2022 let idx = (raw - self.base_entity_id) as usize;
2023 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
2024 return Some(entity);
2025 }
2026 }
2027 self.entities.get(&id)
2033 } else {
2034 self.entities.get(&id)
2035 }
2036 }
2037
2038 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
2039 self.mark_checkpoint_dirty();
2040 self.invalidate_unique_key_lookup();
2041 self.invalidate_graph_read_index();
2042 if self.deleted.contains_key(&id) || !self.state.is_writable() {
2043 return None;
2044 }
2045 self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
2048 if self.use_flat {
2049 let raw = id.raw();
2050 if raw >= self.base_entity_id {
2051 let idx = (raw - self.base_entity_id) as usize;
2052 if self
2053 .flat_entities
2054 .get(idx)
2055 .map(|e| e.id == id)
2056 .unwrap_or(false)
2057 {
2058 return self.flat_entities.get_mut(idx);
2059 }
2060 }
2061 self.entities.get_mut(&id)
2064 } else {
2065 self.entities.get_mut(&id)
2066 }
2067 }
2068
2069 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2070 if !self.state.is_writable() {
2071 return Err(SegmentError::NotWritable);
2072 }
2073
2074 if self.entities.contains_key(&entity.id) {
2075 return Err(SegmentError::AlreadyExists(entity.id));
2076 }
2077
2078 entity.sequence_id = self.next_sequence();
2080
2081 let size = Self::estimate_entity_size(&entity);
2083 self.add_memory(size);
2084
2085 self.index_entity(&entity);
2087
2088 self.update_col_zones_from_entity(&entity);
2090
2091 let id = entity.id;
2093 self.entities.insert(id, entity);
2094
2095 self.last_write_at = current_unix_secs();
2097
2098 Ok(id)
2099 }
2100
2101 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2102 if !self.state.is_writable() {
2103 return Err(SegmentError::NotWritable);
2104 }
2105
2106 self.apply_update_with_metadata(&entity, None)?;
2107 self.last_write_at = current_unix_secs();
2108
2109 Ok(())
2110 }
2111
2112 fn update_hot(
2113 &mut self,
2114 entity: UnifiedEntity,
2115 modified_columns: &[String],
2116 ) -> Result<(), SegmentError> {
2117 if !self.state.is_writable() {
2118 return Err(SegmentError::NotWritable);
2119 }
2120
2121 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
2122 self.last_write_at = current_unix_secs();
2123 Ok(())
2124 }
2125
2126 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2127 if !self.state.is_writable() {
2128 return Err(SegmentError::NotWritable);
2129 }
2130
2131 if self.use_flat {
2133 let raw = id.raw();
2134 if raw >= self.base_entity_id {
2135 let idx = (raw - self.base_entity_id) as usize;
2136 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2137 self.tombstone_flat_entity(idx, id);
2138 return Ok(true);
2139 }
2140 }
2141 return Ok(self.remove_hashmap_entity(id));
2142 }
2143
2144 Ok(self.remove_hashmap_entity(id))
2145 }
2146
2147 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2148 if !self.has_live_entity(id) {
2149 return None;
2150 }
2151 Some(self.metadata.get_all(id))
2152 }
2153
2154 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2155 self.mark_checkpoint_dirty();
2156 if !self.state.is_writable() {
2157 return Err(SegmentError::NotWritable);
2158 }
2159
2160 if !self.has_live_entity(id) {
2161 return Err(SegmentError::NotFound(id));
2162 }
2163
2164 self.metadata.set_all(id, &metadata);
2165 Ok(())
2166 }
2167
2168 fn seal(&mut self) -> Result<(), SegmentError> {
2169 if self.state != SegmentState::Growing {
2170 return Err(SegmentError::InvalidState(self.state));
2171 }
2172
2173 self.state = SegmentState::Sealing;
2174
2175 self.rebuild_sealed_col_zones();
2180
2181 self.state = SegmentState::Sealed;
2182 Ok(())
2183 }
2184
2185 fn should_seal(&self, config: &SegmentConfig) -> bool {
2186 if self.entities.len() >= config.max_entities {
2188 return true;
2189 }
2190
2191 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
2193 return true;
2194 }
2195
2196 if self.age_secs() >= config.max_age_secs {
2198 return true;
2199 }
2200
2201 false
2202 }
2203
2204 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2205 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2206 Box::new(self.flat_entities.iter().chain(self.entities.values()))
2210 } else {
2211 Box::new(self.entities.values())
2212 };
2213 if self.deleted.is_empty() {
2214 base
2215 } else {
2216 Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
2217 }
2218 }
2219
2220 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2221 let ids = self.kind_index.get(kind_filter);
2224 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2227 Box::new(self.flat_entities.iter())
2228 } else {
2229 Box::new(std::iter::empty())
2230 };
2231 Box::new(flat.chain(self.entities.values()).filter(move |e| {
2232 if self.deleted.contains_key(&e.id) {
2233 return false;
2234 }
2235 if let Some(ids) = ids {
2236 ids.contains(&e.id)
2237 } else {
2238 false
2239 }
2240 }))
2241 }
2242
2243 fn filter_metadata(
2244 &self,
2245 filters: &[(String, super::metadata::MetadataFilter)],
2246 ) -> Vec<EntityId> {
2247 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
2250 Box::new(self.flat_entities.iter().map(|e| e.id))
2251 } else {
2252 Box::new(std::iter::empty())
2253 };
2254 flat_ids
2255 .chain(self.entities.keys().copied())
2256 .filter(|id| {
2257 if self.deleted.contains_key(id) {
2258 return false;
2259 }
2260 let metadata = self.metadata.get_all(*id);
2261 metadata.matches_all(filters)
2262 })
2263 .collect()
2264 }
2265}
2266
2267fn build_minmax_multi_intervals(
2268 entries: &[(CanonicalKey, Value)],
2269 max_intervals: usize,
2270) -> Vec<ColZone> {
2271 if entries.is_empty() {
2272 return Vec::new();
2273 }
2274 if entries.len() == 1 || max_intervals <= 1 {
2275 return vec![ColZone::with_bounds(
2276 entries[0].1.clone(),
2277 entries[entries.len() - 1].1.clone(),
2278 )];
2279 }
2280
2281 let mut split_points = if entries.len() <= max_intervals {
2282 (1..entries.len()).collect::<Vec<_>>()
2283 } else {
2284 let target_splits = max_intervals - 1;
2285 let mut selected = select_gap_split_points(entries, target_splits);
2286 if selected.len() < target_splits {
2287 for bucket in 1..max_intervals {
2288 let idx = bucket * entries.len() / max_intervals;
2289 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
2290 continue;
2291 }
2292 selected.push(idx);
2293 if selected.len() >= target_splits {
2294 break;
2295 }
2296 }
2297 }
2298 selected.sort_unstable();
2299 selected.dedup();
2300 selected
2301 };
2302
2303 split_points.push(entries.len());
2304
2305 let mut out = Vec::with_capacity(split_points.len());
2306 let mut start = 0usize;
2307 for end in split_points {
2308 if end <= start {
2309 continue;
2310 }
2311 out.push(ColZone::with_bounds(
2312 entries[start].1.clone(),
2313 entries[end - 1].1.clone(),
2314 ));
2315 start = end;
2316 }
2317
2318 if out.is_empty() {
2319 out.push(ColZone::with_bounds(
2320 entries[0].1.clone(),
2321 entries[entries.len() - 1].1.clone(),
2322 ));
2323 }
2324
2325 out
2326}
2327
2328fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
2329 let mut gaps = Vec::new();
2330 for idx in 1..entries.len() {
2331 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
2332 if score > 0.0 {
2333 gaps.push((score, idx));
2334 }
2335 }
2336 }
2337 gaps.sort_by(|left, right| {
2338 right
2339 .0
2340 .partial_cmp(&left.0)
2341 .unwrap_or(std::cmp::Ordering::Equal)
2342 .then_with(|| left.1.cmp(&right.1))
2343 });
2344 gaps.into_iter()
2345 .take(max_splits)
2346 .map(|(_, idx)| idx)
2347 .collect()
2348}
2349
2350fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
2351 if left.family() != right.family() {
2352 return None;
2353 }
2354 match (left, right) {
2355 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
2356 Some(r.saturating_sub(*l) as f64)
2357 }
2358 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
2359 Some(r.saturating_sub(*l) as f64)
2360 }
2361 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
2362 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
2363 }
2364 _ => None,
2365 }
2366}
2367
2368#[cfg(test)]
2369mod tests {
2370 use super::*;
2371 use crate::storage::unified::entity::RowData;
2372 use crate::storage::unified::MetadataValue;
2373 use reddb_types::Value;
2374
2375 fn pruning_node(id: u64) -> UnifiedEntity {
2376 UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
2377 }
2378
2379 fn pruning_edge(id: u64) -> UnifiedEntity {
2380 UnifiedEntity::graph_edge(
2381 EntityId::new(id),
2382 "mixed",
2383 "a",
2384 "b",
2385 1.0,
2386 Default::default(),
2387 )
2388 }
2389
2390 #[test]
2391 fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
2392 let mut segment = GrowingSegment::new(1, "mixed");
2393 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
2394 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
2395 segment
2396 .bulk_insert(vec![
2397 UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
2398 pruning_node(2),
2399 pruning_edge(4),
2400 ])
2401 .expect("mixed bulk with gap");
2402 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2403 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2404 segment.seal().expect("seal");
2405 assert!(segment.force_delete(EntityId::new(2)));
2406 assert!(segment.force_delete(EntityId::new(4)));
2407 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2409 let mut rebuilt = GrowingSegment::new(2, "mixed");
2410 for item in segment.iter() {
2411 rebuilt.adopt_entity(item.clone(), None);
2412 }
2413 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2414 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2415 rebuilt.adopt_entity(pruning_node(5), None);
2416 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2417 rebuilt.insert(pruning_edge(6)).expect("single edge");
2418 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2419 }
2420
2421 #[test]
2422 fn graph_presence_covers_replacement_and_unrestricted_mutation() {
2423 for flat in [false, true] {
2424 for operation in 0..4 {
2425 let mut segment = GrowingSegment::new(1, "mixed");
2426 let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
2427 if flat {
2428 segment.bulk_insert(vec![vector]).expect("flat vector");
2429 } else {
2430 segment.insert(vector).expect("hashmap vector");
2431 }
2432 match operation {
2433 0 => segment.update(pruning_node(1)).expect("replace with node"),
2434 1 => segment
2435 .update_hot(pruning_node(1), &[])
2436 .expect("hot replacement"),
2437 2 => {
2438 segment.seal().expect("seal vector");
2439 segment
2440 .force_update_with_metadata(&pruning_node(1), &[], None)
2441 .expect("sealed replacement");
2442 }
2443 _ => {
2444 *segment.get_mut(EntityId::new(1)).expect("mutable vector") =
2445 pruning_node(1)
2446 }
2447 }
2448 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2449 assert!(matches!(
2450 segment.get(EntityId::new(1)).expect("node").kind,
2451 EntityKind::GraphNode(_)
2452 ));
2453 if operation == 3 {
2454 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2455 }
2456 segment
2457 .force_update_with_metadata(&pruning_edge(1), &[], None)
2458 .expect("replace with edge");
2459 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2460 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2461 }
2462 }
2463 }
2464
2465 fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
2466 UnifiedEntity::new(
2467 EntityId::new(id),
2468 EntityKind::TableRow {
2469 table: "keys".into(),
2470 row_id: id,
2471 },
2472 EntityData::Row(RowData::with_names(
2473 vec![Value::text("payload"), key],
2474 vec!["payload".into(), "primary".into()],
2475 )),
2476 )
2477 }
2478
2479 fn primary_key_probe(
2480 segment: &GrowingSegment,
2481 key: Value,
2482 exclude: Option<EntityId>,
2483 ) -> Option<EntityId> {
2484 segment.find_unique_key_conflict(
2485 &[vec!["primary".into()]],
2486 0,
2487 &[format!("{key:?}")],
2488 exclude,
2489 &|_| true,
2490 )
2491 }
2492
2493 #[test]
2494 fn unique_key_lookup_follows_segment_mutations_and_sealing() {
2495 let mut segment = GrowingSegment::new(1, "keys");
2496 segment
2497 .bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
2498 .expect("bulk row");
2499 let payload_bytes = segment.memory_bytes();
2500 assert_eq!(
2501 primary_key_probe(&segment, Value::Integer(10), None),
2502 Some(EntityId::new(1))
2503 );
2504 assert!(segment.memory_bytes() > payload_bytes);
2505 segment
2506 .insert(primary_key_row(3, Value::Integer(30)))
2507 .expect("hashmap fallback row");
2508 segment
2509 .bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
2510 .expect("bulk gap row");
2511 assert_eq!(
2512 primary_key_probe(&segment, Value::Integer(40), None),
2513 Some(EntityId::new(4))
2514 );
2515 assert_eq!(
2516 primary_key_probe(&segment, Value::Integer(30), None),
2517 Some(EntityId::new(3))
2518 );
2519 segment
2520 .update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
2521 .expect("replace flat key");
2522 assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
2523 assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
2524 assert_eq!(
2525 primary_key_probe(&segment, Value::Integer(11), None),
2526 Some(EntityId::new(1))
2527 );
2528 segment
2529 .delete_batch(&[EntityId::new(1), EntityId::new(3)])
2530 .expect("delete flat and map");
2531 assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
2532 assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
2533 segment.seal().expect("seal");
2534 assert_eq!(
2535 primary_key_probe(&segment, Value::Integer(40), None),
2536 Some(EntityId::new(4))
2537 );
2538 segment
2539 .force_update_with_metadata(
2540 &primary_key_row(4, Value::Integer(41)),
2541 &["primary".into()],
2542 None,
2543 )
2544 .expect("sealed update");
2545 assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
2546 assert_eq!(
2547 primary_key_probe(&segment, Value::Integer(41), None),
2548 Some(EntityId::new(4))
2549 );
2550 assert!(segment.force_delete(EntityId::new(4)));
2551 assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
2552 }
2553
2554 #[test]
2555 fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
2556 let mut segment = GrowingSegment::new(1, "keys");
2557 segment
2558 .bulk_insert(vec![
2559 primary_key_row(1, Value::Integer(10)),
2560 primary_key_row(2, Value::Integer(20)),
2561 ])
2562 .expect("initial rows");
2563 let definitions = vec![
2564 vec!["primary".into()],
2565 vec!["payload".into()],
2566 vec!["primary".into(), "payload".into()],
2567 ];
2568 let payload = format!("{:?}", Value::text("payload"));
2569 assert_eq!(
2570 segment.find_unique_key_conflict(
2571 &definitions,
2572 1,
2573 &[payload.clone()],
2574 None,
2575 &|entity| entity.id != EntityId::new(1)
2576 ),
2577 Some(EntityId::new(2))
2578 );
2579 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2580 segment
2581 .insert(primary_key_row(3, Value::Integer(30)))
2582 .expect("append");
2583 assert_eq!(
2584 segment.find_unique_key_conflict(
2585 &definitions,
2586 2,
2587 &[format!("{:?}", Value::Integer(30)), payload.clone()],
2588 None,
2589 &|_| true
2590 ),
2591 Some(EntityId::new(3))
2592 );
2593 let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
2594 assert_eq!(
2595 segment.find_unique_key_conflict(
2596 &definitions,
2597 0,
2598 &[format!("{:?}", Value::Integer(20))],
2599 None,
2600 &|_| true
2601 ),
2602 Some(EntityId::new(2))
2603 );
2604 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2605 assert_eq!(
2606 segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
2607 bytes
2608 );
2609 let reduced = vec![vec!["payload".into()]];
2610 assert_eq!(
2611 segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
2612 Some(EntityId::new(1))
2613 );
2614 assert_eq!(segment.unique_key_lookup.read().len(), 1);
2615 assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
2616 }
2617
2618 #[test]
2619 fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
2620 let mut segment = GrowingSegment::new(1, "keys");
2621 let values = [
2622 Value::Float(f64::NAN),
2623 Value::Float(-0.0),
2624 Value::Float(0.0),
2625 Value::Integer(1),
2626 Value::UnsignedInteger(1),
2627 Value::text("1"),
2628 ];
2629 for (offset, value) in values.iter().enumerate() {
2630 let id = offset as u64 + 1;
2631 segment
2632 .insert(primary_key_row(id, value.clone()))
2633 .expect("distinct physical row");
2634 assert_eq!(
2635 primary_key_probe(&segment, value.clone(), None),
2636 Some(EntityId::new(id))
2637 );
2638 assert_eq!(
2639 primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
2640 None
2641 );
2642 }
2643 segment
2644 .insert(primary_key_row(20, Value::Integer(1)))
2645 .expect("duplicate physical key");
2646 assert_eq!(
2647 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
2648 Some(EntityId::new(20))
2649 );
2650 segment
2653 .entities
2654 .insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
2655 assert_eq!(
2656 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
2657 None
2658 );
2659 let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
2661 *entity = primary_key_row(4, Value::Integer(100));
2662 assert_eq!(
2663 primary_key_probe(&segment, Value::Integer(100), None),
2664 Some(EntityId::new(4))
2665 );
2666 }
2667
2668 #[test]
2669 fn test_growing_segment_basic() {
2670 let mut segment = GrowingSegment::new(1, "test");
2671
2672 let entity = UnifiedEntity::table_row(
2673 EntityId::new(1),
2674 "users",
2675 1,
2676 vec![Value::text("Alice".to_string())],
2677 );
2678
2679 let id = segment.insert(entity).unwrap();
2680 assert_eq!(id, EntityId::new(1));
2681 assert!(segment.contains(id));
2682
2683 let stats = segment.stats();
2684 assert_eq!(stats.entity_count, 1);
2685 assert_eq!(stats.row_count, 1);
2686 }
2687
2688 #[test]
2689 fn test_segment_metadata() {
2690 let mut segment = GrowingSegment::new(1, "test");
2691
2692 let entity = UnifiedEntity::table_row(
2693 EntityId::new(1),
2694 "users",
2695 1,
2696 vec![Value::text("Alice".to_string())],
2697 );
2698 segment.insert(entity).unwrap();
2699
2700 let mut meta = Metadata::new();
2701 meta.set("role", MetadataValue::String("admin".to_string()));
2702 meta.set("level", MetadataValue::Int(5));
2703
2704 segment.set_metadata(EntityId::new(1), meta).unwrap();
2705
2706 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
2707 assert_eq!(
2708 retrieved.get("role"),
2709 Some(&MetadataValue::String("admin".to_string()))
2710 );
2711 }
2712
2713 #[test]
2714 fn test_segment_seal() {
2715 let mut segment = GrowingSegment::new(1, "test");
2716
2717 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
2718 segment.insert(entity).unwrap();
2719
2720 assert!(segment.state().is_writable());
2722
2723 segment.seal().unwrap();
2725 assert_eq!(segment.state(), SegmentState::Sealed);
2726
2727 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
2729 assert!(segment.insert(entity2).is_err());
2730 }
2731
2732 #[test]
2733 fn test_should_seal() {
2734 let mut segment = GrowingSegment::new(1, "test");
2735
2736 let config = SegmentConfig {
2737 max_entities: 2,
2738 ..Default::default()
2739 };
2740
2741 assert!(!segment.should_seal(&config));
2742
2743 segment
2744 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
2745 .unwrap();
2746 assert!(!segment.should_seal(&config));
2747
2748 segment
2749 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
2750 .unwrap();
2751 assert!(segment.should_seal(&config));
2752 }
2753
2754 #[test]
2755 fn test_cross_references() {
2756 let mut segment = GrowingSegment::new(1, "test");
2757
2758 let mut entity1 = UnifiedEntity::table_row(
2759 EntityId::new(1),
2760 "hosts",
2761 1,
2762 vec![Value::text("192.168.1.1".to_string())],
2763 );
2764 entity1.add_cross_ref(CrossRef::new(
2765 EntityId::new(1),
2766 EntityId::new(2),
2767 "nodes",
2768 RefType::RowToNode,
2769 ));
2770 segment.insert(entity1).unwrap();
2771
2772 let refs_from = segment.get_references_from(EntityId::new(1));
2773 assert_eq!(refs_from.len(), 1);
2774 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
2775
2776 let refs_to = segment.get_references_to(EntityId::new(2));
2777 assert_eq!(refs_to.len(), 1);
2778 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
2779 }
2780
2781 #[test]
2782 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
2783 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
2784 zone.update(&Value::Email("delta@example.com".to_string()));
2785
2786 let probe = Value::Email("alpha@example.com".to_string());
2787 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
2788
2789 let in_range = Value::Email("charlie@example.com".to_string());
2790 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
2791 }
2792
2793 #[test]
2794 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
2795 let mut segment = GrowingSegment::new(1, "test");
2796
2797 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
2798 let entity = UnifiedEntity::new(
2799 EntityId::new(row_id),
2800 EntityKind::TableRow {
2801 table: "users".into(),
2802 row_id,
2803 },
2804 EntityData::Row(RowData::with_names(
2805 vec![Value::Integer(age)],
2806 vec!["age".to_string()],
2807 )),
2808 );
2809 segment.insert(entity).unwrap();
2810 }
2811
2812 segment.seal().unwrap();
2813
2814 let miss = Value::Integer(500);
2815 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
2816
2817 let hit = Value::Integer(1000);
2818 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
2819 }
2820}