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 memory_bytes: AtomicU64,
567 dead_entity_bytes: u64,
570 dead_resident_count: usize,
572
573 pub(crate) published_flat_len: AtomicUsize,
580}
581
582impl GrowingSegment {
583 pub(crate) fn may_contain_graph_kind(&self, kind: GraphEntityKind) -> bool {
584 self.graph_kind_mask & kind.mask() != 0
585 }
586
587 fn observe_graph_kind(&mut self, kind: &EntityKind) {
588 self.graph_kind_mask |= match kind {
589 EntityKind::GraphNode(_) => GraphEntityKind::Node.mask(),
590 EntityKind::GraphEdge(_) => GraphEntityKind::Edge.mask(),
591 _ => 0,
592 };
593 }
594
595 pub(crate) fn scan_cursor(&self) -> SegmentScanCursor {
596 SegmentScanCursor {
597 flat_end: self.flat_entities.len(),
598 hash_end: self.hash_scan_ids.len(),
599 deleted_end: self.deleted.len(),
600 position: 0,
601 }
602 }
603
604 pub(crate) fn scan_batch(
608 &self,
609 cursor: &mut SegmentScanCursor,
610 before_work: &mut impl FnMut() -> bool,
611 visit: &mut impl FnMut(EntityId),
612 ) -> bool {
613 for _ in 0..SCAN_BATCH_SIZE {
614 if cursor.finished() {
615 break;
616 }
617 if !before_work() {
618 return false;
619 }
620 let id = if cursor.position < cursor.flat_end {
621 self.flat_entities[cursor.position].id
622 } else {
623 self.hash_scan_ids[cursor.position - cursor.flat_end]
624 };
625 cursor.position += 1;
626 let eligible = match self.deleted.get(&id) {
627 Some(&ordinal) => ordinal >= cursor.deleted_end,
628 None => self.get(id).is_some(),
631 };
632 if eligible {
633 visit(id);
634 }
635 }
636 true
637 }
638
639 #[inline]
642 pub fn for_each_fast<F>(&self, mut f: F) -> bool
643 where
644 F: FnMut(&UnifiedEntity) -> bool,
645 {
646 if self.use_flat {
647 if self.deleted.is_empty() {
649 for entity in &self.flat_entities {
650 if !f(entity) {
651 return false;
652 }
653 }
654 for entity in self.entities.values() {
657 if !f(entity) {
658 return false;
659 }
660 }
661 } else {
662 for entity in &self.flat_entities {
663 if self.deleted.contains_key(&entity.id) {
664 continue;
665 }
666 if !f(entity) {
667 return false;
668 }
669 }
670 for entity in self.entities.values() {
671 if self.deleted.contains_key(&entity.id) {
672 continue;
673 }
674 if !f(entity) {
675 return false;
676 }
677 }
678 }
679 } else {
680 if self.deleted.is_empty() {
682 for entity in self.entities.values() {
683 if !f(entity) {
684 return false;
685 }
686 }
687 } else {
688 for entity in self.entities.values() {
689 if self.deleted.contains_key(&entity.id) {
690 continue;
691 }
692 if !f(entity) {
693 return false;
694 }
695 }
696 }
697 }
698 true
699 }
700
701 pub fn new(id: SegmentId, collection: impl Into<String>) -> Self {
703 let now = current_unix_secs();
704
705 Self {
706 id,
707 collection: collection.into(),
708 state: SegmentState::Growing,
709 created_at: now,
710 last_write_at: now,
711 entities: HashMap::new(),
712 hash_scan_ids: Vec::new(),
713 flat_entities: Vec::new(),
714 base_entity_id: 0,
715 use_flat: false,
716 deleted: HashMap::new(),
717 metadata: MetadataStorage::new(),
718 pk_index: BTreeMap::new(),
719 unique_key_lookup: parking_lot::RwLock::new(Vec::new()),
720 unique_key_lookup_bytes: AtomicU64::new(0),
721 kind_index: HashMap::new(),
722 graph_read_index: parking_lot::RwLock::new(Some(SegmentGraphIndex::default())),
723 graph_read_index_bytes: AtomicU64::new(0),
724 graph_kind_mask: 0,
725 cross_ref_forward: HashMap::new(),
726 cross_ref_reverse: HashMap::new(),
727 col_zones: HashMap::new(),
728 zone_payload_bytes: 0,
729 sealed_col_zones: HashMap::new(),
730 sequence: AtomicU64::new(0),
731 memory_bytes: AtomicU64::new(0),
732 dead_entity_bytes: 0,
733 dead_resident_count: 0,
734 published_flat_len: AtomicUsize::new(0),
735 }
736 }
737
738 fn next_sequence(&self) -> u64 {
740 self.sequence.fetch_add(1, Ordering::SeqCst)
741 }
742
743 fn has_live_entity(&self, id: EntityId) -> bool {
744 if self.deleted.contains_key(&id) {
745 return false;
746 }
747 if self.use_flat {
748 let raw = id.raw();
749 if raw >= self.base_entity_id {
750 let idx = (raw - self.base_entity_id) as usize;
751 if self
752 .flat_entities
753 .get(idx)
754 .is_some_and(|entity| entity.id == id)
755 {
756 return true;
757 }
758 }
759 self.entities.contains_key(&id)
762 } else {
763 self.entities.contains_key(&id)
764 }
765 }
766
767 fn update_existing_entity_in_place(
768 &mut self,
769 entity: &UnifiedEntity,
770 ) -> Result<UpdateIndexSnapshot, SegmentError> {
771 self.invalidate_unique_key_lookup();
772 if (self.graph_kind_mask != 0
773 || matches!(
774 entity.kind,
775 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
776 ))
777 && self
778 .get(entity.id)
779 .is_some_and(|old| !SegmentGraphIndex::same_keys(old, entity))
780 {
781 self.invalidate_graph_read_index();
782 }
783 self.observe_graph_kind(&entity.kind);
786 if self.use_flat {
787 let raw = entity.id.raw();
788 if raw >= self.base_entity_id {
789 let idx = (raw - self.base_entity_id) as usize;
790 if let Some(slot) = self
791 .flat_entities
792 .get_mut(idx)
793 .filter(|slot| slot.id == entity.id)
794 {
795 let snapshot = UpdateIndexSnapshot::from_entity(slot);
796 let old_bytes = Self::estimate_entity_size(slot);
797 slot.clone_from(entity);
798 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
799 return Ok(snapshot);
800 }
801 }
802 let Some(slot) = self.entities.get_mut(&entity.id) else {
803 return Err(SegmentError::NotFound(entity.id));
804 };
805 let snapshot = UpdateIndexSnapshot::from_entity(slot);
806 let old_bytes = Self::estimate_entity_size(slot);
807 slot.clone_from(entity);
808 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
809 Ok(snapshot)
810 } else {
811 let Some(slot) = self.entities.get_mut(&entity.id) else {
812 return Err(SegmentError::NotFound(entity.id));
813 };
814 let snapshot = UpdateIndexSnapshot::from_entity(slot);
815 let old_bytes = Self::estimate_entity_size(slot);
816 slot.clone_from(entity);
817 self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
818 Ok(snapshot)
819 }
820 }
821
822 fn apply_hot_update_with_metadata(
823 &mut self,
824 entity: &UnifiedEntity,
825 modified_columns: &[String],
826 metadata: Option<&Metadata>,
827 ) -> Result<(), SegmentError> {
828 let old = self.update_existing_entity_in_place(entity)?;
829 self.reindex_for_update(&old, entity, Some(modified_columns));
830 self.update_col_zones_from_entity(entity);
831 if let Some(metadata) = metadata {
832 self.metadata.set_all(entity.id, metadata);
833 }
834 Ok(())
835 }
836
837 fn apply_update_with_metadata(
838 &mut self,
839 entity: &UnifiedEntity,
840 metadata: Option<&Metadata>,
841 ) -> Result<(), SegmentError> {
842 let old = self.update_existing_entity_in_place(entity)?;
843 self.reindex_for_update(&old, entity, None);
844 self.update_col_zones_from_entity(entity);
845 if let Some(metadata) = metadata {
846 self.metadata.set_all(entity.id, metadata);
847 }
848 Ok(())
849 }
850
851 pub fn update_hot_batch_with_metadata<'a, I>(&mut self, items: I) -> Result<(), SegmentError>
852 where
853 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
854 {
855 if !self.state.is_writable() {
856 return Err(SegmentError::NotWritable);
857 }
858
859 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
860 items.into_iter().collect();
861 if items.is_empty() {
862 return Ok(());
863 }
864
865 for (entity, _, _) in &items {
866 if !self.has_live_entity(entity.id) {
867 return Err(SegmentError::NotFound(entity.id));
868 }
869 }
870
871 for (entity, modified_columns, metadata) in items {
872 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)?;
873 }
874
875 self.last_write_at = current_unix_secs();
876 Ok(())
877 }
878
879 fn record_tombstone(&mut self, id: EntityId) -> bool {
881 let ordinal = self.deleted.len();
882 match self.deleted.entry(id) {
883 std::collections::hash_map::Entry::Vacant(entry) => {
884 entry.insert(ordinal);
885 true
886 }
887 std::collections::hash_map::Entry::Occupied(_) => false,
888 }
889 }
890
891 fn tombstone_flat_entity(&mut self, idx: usize, id: EntityId) -> bool {
897 self.invalidate_unique_key_lookup();
898 if !self.record_tombstone(id) {
899 return false;
900 }
901 if let Some(index) = self.graph_read_index.get_mut() {
902 index.remove(&self.flat_entities[idx]);
903 self.graph_read_index_bytes
904 .store(index.memory_bytes(), Ordering::Relaxed);
905 }
906 let size = Self::estimate_entity_size(&self.flat_entities[idx]) as u64;
907 self.dead_entity_bytes += size;
908 self.dead_resident_count += 1;
909 self.metadata.remove_all(id);
910 true
911 }
912
913 fn remove_hashmap_entity(&mut self, id: EntityId) -> bool {
918 self.invalidate_unique_key_lookup();
919 let Some(entity) = self.entities.remove(&id) else {
920 return false;
921 };
922 self.release_memory(Self::estimate_entity_size(&entity));
923 self.unindex_entity(&entity);
924 self.metadata.remove_all(id);
925 self.record_tombstone(id);
926 true
927 }
928
929 pub fn delete_batch(&mut self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
930 if !self.state.is_writable() {
931 return Err(SegmentError::NotWritable);
932 }
933 if ids.is_empty() {
934 return Ok(Vec::new());
935 }
936
937 let mut deleted_ids = Vec::with_capacity(ids.len());
938
939 if self.use_flat {
940 for &id in ids {
941 let raw = id.raw();
942 if raw < self.base_entity_id {
943 if self.remove_hashmap_entity(id) {
944 deleted_ids.push(id);
945 }
946 continue;
947 }
948 let idx = (raw - self.base_entity_id) as usize;
949 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
950 if self.tombstone_flat_entity(idx, id) {
951 deleted_ids.push(id);
952 }
953 } else if self.remove_hashmap_entity(id) {
954 deleted_ids.push(id);
955 }
956 }
957 } else {
958 for &id in ids {
959 if self.remove_hashmap_entity(id) {
960 deleted_ids.push(id);
961 }
962 }
963 }
964
965 if !deleted_ids.is_empty() {
966 self.last_write_at = current_unix_secs();
967 }
968
969 Ok(deleted_ids)
970 }
971
972 fn add_memory(&self, bytes: usize) {
974 self.memory_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
975 }
976
977 fn replace_entity_memory(&self, previous: usize, current: usize) {
978 if current >= previous {
979 self.add_memory(current - previous);
980 } else {
981 self.release_memory(previous - current);
982 }
983 }
984
985 pub fn memory_bytes(&self) -> u64 {
990 self.memory_bytes.load(Ordering::Relaxed)
991 + self.zone_payload_bytes as u64
992 + self.unique_key_lookup_bytes.load(Ordering::Relaxed)
993 + self.graph_read_index_bytes.load(Ordering::Relaxed)
994 + (self.hash_scan_ids.capacity() * std::mem::size_of::<EntityId>()) as u64
995 }
996
997 fn release_memory(&self, bytes: usize) {
1000 let mut current = self.memory_bytes.load(Ordering::Relaxed);
1001 loop {
1002 let next = current.saturating_sub(bytes as u64);
1003 match self.memory_bytes.compare_exchange_weak(
1004 current,
1005 next,
1006 Ordering::Relaxed,
1007 Ordering::Relaxed,
1008 ) {
1009 Ok(_) => break,
1010 Err(observed) => current = observed,
1011 }
1012 }
1013 }
1014
1015 pub(crate) fn resident_entity_count(&self) -> usize {
1017 self.flat_entities.len() + self.entities.len()
1018 }
1019
1020 pub(crate) fn live_entity_count(&self) -> usize {
1022 self.resident_entity_count()
1023 .saturating_sub(self.dead_resident_count)
1024 }
1025
1026 pub(crate) fn tombstone_count(&self) -> usize {
1029 self.deleted.len()
1030 }
1031
1032 pub(crate) fn resident_bytes(&self) -> u64 {
1035 self.memory_bytes() + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1036 }
1037
1038 pub(crate) fn reclaimable_bytes(&self) -> u64 {
1041 self.dead_entity_bytes + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
1042 }
1043
1044 pub(crate) fn live_entity_ids(&self) -> Vec<EntityId> {
1050 let mut ids = Vec::with_capacity(self.live_entity_count());
1051 self.for_each_fast(|entity| {
1052 ids.push(entity.id);
1053 true
1054 });
1055 ids
1056 }
1057
1058 pub(crate) fn adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
1066 let id = entity.id;
1067
1068 let next = entity.sequence_id.saturating_add(1);
1071 if self.sequence.load(Ordering::Relaxed) < next {
1072 self.sequence.store(next, Ordering::Relaxed);
1073 }
1074
1075 self.add_memory(Self::estimate_entity_size(&entity));
1076 self.index_entity(&entity);
1077 self.update_col_zones_from_entity(&entity);
1078 self.entities.insert(id, entity);
1079
1080 if let Some(metadata) = metadata {
1081 if !metadata.is_empty() {
1082 self.metadata.set_all(id, &metadata);
1083 }
1084 }
1085 }
1086
1087 pub(crate) fn evict_entity(&mut self, id: EntityId) -> bool {
1093 self.invalidate_unique_key_lookup();
1094 let Some(entity) = self.entities.remove(&id) else {
1095 return false;
1096 };
1097 self.release_memory(Self::estimate_entity_size(&entity));
1098 self.unindex_entity(&entity);
1099 for cross_ref in entity.cross_refs() {
1100 if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1101 reverse.retain(|(source, _)| *source != id);
1102 }
1103 }
1104 self.metadata.remove_all(id);
1105 true
1106 }
1107
1108 fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
1110 super::memory_size::entity_bytes(entity)
1111 }
1112
1113 fn update_column_zone(&mut self, column: &str, value: &Value) {
1121 if matches!(value, Value::Null) {
1122 return;
1123 }
1124 match self.col_zones.entry(column.to_string()) {
1125 std::collections::hash_map::Entry::Occupied(mut entry) => {
1126 let before = entry.get().payload_bytes();
1127 entry.get_mut().update(value);
1128 self.zone_payload_bytes = self
1129 .zone_payload_bytes
1130 .saturating_sub(before)
1131 .saturating_add(entry.get().payload_bytes());
1132 }
1133 std::collections::hash_map::Entry::Vacant(entry) => {
1134 let zone = ColZone::new(value.clone());
1135 self.zone_payload_bytes =
1136 self.zone_payload_bytes.saturating_add(zone.payload_bytes());
1137 entry.insert(zone);
1138 }
1139 }
1140 }
1141
1142 fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
1143 if let EntityData::Row(row) = &entity.data {
1144 if let Some(named) = &row.named {
1145 for (column, value) in named {
1146 self.update_column_zone(column, value);
1147 }
1148 } else if let Some(schema) = &row.schema {
1149 for (column, value) in schema.iter().zip(&row.columns) {
1150 self.update_column_zone(column, value);
1151 }
1152 }
1153 }
1154 }
1155
1156 fn rebuild_sealed_col_zones(&mut self) {
1157 let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
1158 let mut family_by_col: HashMap<String, reddb_types::CanonicalKeyFamily> = HashMap::new();
1159 let mut mixed_family_cols = HashSet::new();
1160 let mut unsupported_cols = HashSet::new();
1161
1162 let mut observe_row = |row: &super::entity::RowData| {
1163 for (col, value) in row.iter_fields() {
1164 if matches!(value, Value::Null) {
1165 continue;
1166 }
1167 let Some(key) = value_to_canonical_key(value) else {
1168 unsupported_cols.insert(col.to_string());
1169 continue;
1170 };
1171 match family_by_col.get(col).copied() {
1172 Some(existing) if existing != key.family() => {
1173 mixed_family_cols.insert(col.to_string());
1174 }
1175 None => {
1176 family_by_col.insert(col.to_string(), key.family());
1177 }
1178 _ => {}
1179 }
1180 values_by_col
1181 .entry(col.to_string())
1182 .or_default()
1183 .push((key, value.clone()));
1184 }
1185 };
1186
1187 if self.use_flat {
1188 for entity in &self.flat_entities {
1189 if self.deleted.contains_key(&entity.id) {
1190 continue;
1191 }
1192 if let EntityData::Row(row) = &entity.data {
1193 observe_row(row);
1194 }
1195 }
1196 } else {
1197 for entity in self.entities.values() {
1198 if self.deleted.contains_key(&entity.id) {
1199 continue;
1200 }
1201 if let EntityData::Row(row) = &entity.data {
1202 observe_row(row);
1203 }
1204 }
1205 }
1206
1207 let mut sealed_col_zones = HashMap::new();
1208 for (col, mut entries) in values_by_col {
1209 if mixed_family_cols.contains(&col)
1210 || unsupported_cols.contains(&col)
1211 || entries.is_empty()
1212 {
1213 continue;
1214 }
1215 entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
1216 entries.dedup_by(|left, right| left.0 == right.0);
1217
1218 let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
1219 if intervals.len() > 1 {
1220 sealed_col_zones.insert(col, MultiColZone { intervals });
1221 }
1222 }
1223
1224 let previous = self
1225 .sealed_col_zones
1226 .values()
1227 .flat_map(|zones| &zones.intervals)
1228 .map(ColZone::payload_bytes)
1229 .fold(0, usize::saturating_add);
1230 let current = sealed_col_zones
1231 .values()
1232 .flat_map(|zones| &zones.intervals)
1233 .map(ColZone::payload_bytes)
1234 .fold(0, usize::saturating_add);
1235 self.zone_payload_bytes = self
1236 .zone_payload_bytes
1237 .saturating_sub(previous)
1238 .saturating_add(current);
1239 self.sealed_col_zones = sealed_col_zones;
1240 }
1241
1242 pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
1246 if preds.is_empty() {
1247 return false;
1248 }
1249 for (col, pred) in preds {
1250 if let Some(zone) = self.sealed_col_zones.get(*col) {
1251 if zone.can_skip(pred) {
1252 return true;
1253 }
1254 continue;
1255 }
1256 if let Some(zone) = self.col_zones.get(*col) {
1257 if pred.can_skip(zone) {
1258 return true; }
1260 }
1261 }
1262 false
1263 }
1264
1265 fn index_entity(&mut self, entity: &UnifiedEntity) {
1267 self.hash_scan_ids.push(entity.id);
1268 self.observe_graph_kind(&entity.kind);
1269 if let Some(index) = self.graph_read_index.get_mut() {
1270 index.insert(entity);
1271 self.graph_read_index_bytes
1272 .store(index.memory_bytes(), Ordering::Relaxed);
1273 }
1274 let indexes = self.unique_key_lookup.get_mut();
1275 for index in indexes.iter_mut() {
1276 index.insert(entity);
1277 }
1278 self.unique_key_lookup_bytes.store(
1279 indexes
1280 .iter()
1281 .map(|index| index.memory_bytes() as u64)
1282 .sum(),
1283 Ordering::Relaxed,
1284 );
1285 let kind_key = entity.kind.storage_type().to_string();
1287 self.kind_index
1288 .entry(kind_key)
1289 .or_default()
1290 .insert(entity.id);
1291
1292 if let EntityData::Row(row) = &entity.data {
1294 if let Some(first_col) = row.columns.first() {
1295 let pk_str = format!("{:?}", first_col);
1296 self.pk_index
1297 .insert((entity.kind.collection().to_string(), pk_str), entity.id);
1298 }
1299 }
1300
1301 for cross_ref in entity.cross_refs() {
1303 self.cross_ref_forward
1304 .entry(cross_ref.source)
1305 .or_default()
1306 .push((cross_ref.target, cross_ref.ref_type));
1307
1308 self.cross_ref_reverse
1309 .entry(cross_ref.target)
1310 .or_default()
1311 .push((cross_ref.source, cross_ref.ref_type));
1312 }
1313 }
1314
1315 fn invalidate_graph_read_index(&mut self) {
1316 *self.graph_read_index.get_mut() = None;
1317 self.graph_read_index_bytes.store(0, Ordering::Relaxed);
1318 }
1319
1320 #[cfg(test)]
1323 pub(crate) fn visit_graph_index_candidates(
1324 &self,
1325 kind: GraphEntityKind,
1326 keys: &[EntityId],
1327 before_work: &mut impl FnMut() -> bool,
1328 visit: &mut impl FnMut(&UnifiedEntity) -> bool,
1329 ) -> bool {
1330 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1331 return false;
1332 };
1333 let index = cached.as_ref().expect("complete graph index");
1334 for key in keys {
1335 if !before_work() {
1336 return false;
1337 }
1338 for id in index.candidates(kind, *key) {
1339 if !before_work() {
1340 return false;
1341 }
1342 if let Some(entity) = self.get(id) {
1343 if !visit(entity) {
1344 return false;
1345 }
1346 }
1347 }
1348 }
1349 true
1350 }
1351
1352 fn graph_read_index_for_query(
1353 &self,
1354 before_work: &mut impl FnMut() -> bool,
1355 ) -> Option<parking_lot::RwLockReadGuard<'_, Option<SegmentGraphIndex>>> {
1356 let mut cached = self.graph_read_index.read();
1357 if cached.is_none() {
1358 drop(cached);
1359 let mut writer = self.graph_read_index.write();
1360 if writer.is_none() {
1361 let mut index = SegmentGraphIndex::default();
1362 if !self.for_each_fast(|entity| {
1365 if !before_work() {
1366 return false;
1367 }
1368 index.insert(entity);
1369 true
1370 }) {
1371 return None;
1372 }
1373 self.graph_read_index_bytes
1374 .store(index.memory_bytes(), Ordering::Relaxed);
1375 *writer = Some(index);
1376 }
1377 cached = parking_lot::RwLockWriteGuard::downgrade(writer);
1378 }
1379 Some(cached)
1380 }
1381
1382 pub(crate) fn graph_index_cursor(
1383 &self,
1384 kind: GraphEntityKind,
1385 keys: &[EntityId],
1386 before_work: &mut impl FnMut() -> bool,
1387 ) -> Option<GraphIndexCursor> {
1388 let cached = self.graph_read_index_for_query(before_work)?;
1389 let index = cached.as_ref().expect("complete graph index");
1390 let mut upper = EntityId::new(0);
1391 for key in keys {
1392 if !before_work() {
1393 return None;
1394 }
1395 if let Some(last) = index.candidates(kind, *key).next_back() {
1396 upper = upper.max(last);
1397 }
1398 }
1399 Some(GraphIndexCursor {
1400 key_index: 0,
1401 after: None,
1402 upper,
1403 })
1404 }
1405
1406 pub(crate) fn graph_index_batch(
1409 &self,
1410 kind: GraphEntityKind,
1411 keys: &[EntityId],
1412 cursor: &mut GraphIndexCursor,
1413 before_work: &mut impl FnMut() -> bool,
1414 visit: &mut impl FnMut(EntityId),
1415 ) -> bool {
1416 self.graph_index_batch_with_limit(kind, keys, cursor, before_work, visit, SCAN_BATCH_SIZE)
1417 }
1418
1419 pub(crate) fn graph_index_batch_with_limit(
1420 &self,
1421 kind: GraphEntityKind,
1422 keys: &[EntityId],
1423 cursor: &mut GraphIndexCursor,
1424 before_work: &mut impl FnMut() -> bool,
1425 visit: &mut impl FnMut(EntityId),
1426 limit: usize,
1427 ) -> bool {
1428 assert!(
1429 limit > 0 && limit <= SCAN_BATCH_SIZE,
1430 "bounded graph index batch"
1431 );
1432 let Some(cached) = self.graph_read_index_for_query(before_work) else {
1433 return false;
1434 };
1435 let index = cached.as_ref().expect("complete graph index");
1436 let mut count = 0;
1437 while !cursor.finished(keys) && count < limit {
1438 let key = keys[cursor.key_index];
1439 if !before_work() {
1440 return false;
1441 }
1442 for id in index
1443 .candidates_after(kind, key, cursor.after, cursor.upper)
1444 .take(limit - count)
1445 {
1446 if !before_work() {
1447 return false;
1448 }
1449 cursor.after = Some(id);
1450 count += 1;
1451 visit(id);
1452 }
1453 if count == limit && cursor.after != Some(cursor.upper) {
1454 break;
1455 }
1456 cursor.key_index += 1;
1457 cursor.after = None;
1458 }
1459 true
1460 }
1461
1462 pub(crate) fn find_unique_key_conflict(
1467 &self,
1468 columns: &[Vec<String>],
1469 key_index: usize,
1470 signatures: &[String],
1471 exclude_id: Option<EntityId>,
1472 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1473 ) -> Option<EntityId> {
1474 assert!(
1475 key_index < columns.len(),
1476 "constraint index belongs to current schema"
1477 );
1478 let cached = self.unique_key_lookup.read();
1479 if cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1480 return self.find_unique_key_candidate(
1481 &cached[key_index],
1482 signatures,
1483 exclude_id,
1484 reserves_key,
1485 );
1486 }
1487 drop(cached);
1488 let mut cached = self.unique_key_lookup.write();
1489 if !cached.iter().map(|index| &index.columns).eq(columns.iter()) {
1490 let mut indexes: Vec<_> = columns
1491 .iter()
1492 .map(|columns| SegmentUniqueKey::new(columns))
1493 .collect();
1494 self.for_each_fast(|entity| {
1495 for index in &mut indexes {
1496 index.insert(entity);
1497 }
1498 true
1499 });
1500 self.unique_key_lookup_bytes.store(
1501 indexes
1502 .iter()
1503 .map(|index| index.memory_bytes() as u64)
1504 .sum(),
1505 Ordering::Relaxed,
1506 );
1507 *cached = indexes;
1508 }
1509 self.find_unique_key_candidate(&cached[key_index], signatures, exclude_id, reserves_key)
1510 }
1511
1512 fn find_unique_key_candidate(
1513 &self,
1514 index: &SegmentUniqueKey,
1515 signatures: &[String],
1516 exclude_id: Option<EntityId>,
1517 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1518 ) -> Option<EntityId> {
1519 index.candidates(signatures).iter().copied().find(|id| {
1520 if Some(*id) == exclude_id {
1521 return false;
1522 }
1523 let Some(
1524 entity @ UnifiedEntity {
1525 data: EntityData::Row(row),
1526 ..
1527 },
1528 ) = self.get(*id)
1529 else {
1530 return false;
1531 };
1532 index
1533 .columns
1534 .iter()
1535 .zip(signatures)
1536 .all(|(column, signature)| {
1537 row.get_field(column)
1538 .is_some_and(|value| format!("{value:?}") == *signature)
1539 })
1540 && reserves_key(entity)
1541 })
1542 }
1543
1544 fn invalidate_unique_key_lookup(&mut self) {
1545 *self.unique_key_lookup.get_mut() = Vec::new();
1546 self.unique_key_lookup_bytes.store(0, Ordering::Relaxed);
1547 }
1548
1549 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1551 if let Some(id) = entity_id_from_probe_key(key) {
1552 return self.has_live_entity(id);
1553 }
1554
1555 let Ok(pk_str) = std::str::from_utf8(key) else {
1556 return false;
1557 };
1558 let pk_key = (self.collection.clone(), pk_str.to_string());
1559 self.pk_index
1560 .get(&pk_key)
1561 .is_some_and(|id| self.has_live_entity(*id))
1562 }
1563
1564 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1566 if let Some(index) = self.graph_read_index.get_mut() {
1567 index.remove(entity);
1568 self.graph_read_index_bytes
1569 .store(index.memory_bytes(), Ordering::Relaxed);
1570 }
1571 let kind_key = entity.kind.storage_type().to_string();
1573 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1574 set.remove(&entity.id);
1575 }
1576
1577 if let EntityData::Row(row) = &entity.data {
1579 if let Some(first_col) = row.columns.first() {
1580 let pk_str = format!("{:?}", first_col);
1581 self.pk_index
1582 .remove(&(entity.kind.collection().to_string(), pk_str));
1583 }
1584 }
1585
1586 self.cross_ref_forward.remove(&entity.id);
1588 }
1590
1591 fn reindex_for_update(
1601 &mut self,
1602 old: &UpdateIndexSnapshot,
1603 new: &UnifiedEntity,
1604 modified_columns: Option<&[String]>,
1605 ) {
1606 let pk_changed = match &new.data {
1611 EntityData::Row(new_row) => {
1612 if let Some(cols) = modified_columns {
1613 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1616 new_row
1617 .schema
1618 .as_deref()
1619 .and_then(|schema| schema.first().map(|name| name.as_str()))
1620 });
1621 match pk_col_name {
1622 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1623 None => old.pk_value.as_ref() != new_row.columns.first(),
1625 }
1626 } else {
1627 old.pk_value.as_ref() != new_row.columns.first()
1628 }
1629 }
1630 _ => false,
1632 };
1633
1634 if pk_changed {
1635 if let Some((collection, pk_str)) = &old.pk_index_key {
1637 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1638 }
1639 if let EntityData::Row(row) = &new.data {
1641 if let Some(first_col) = row.columns.first() {
1642 let pk_str = format!("{:?}", first_col);
1643 self.pk_index
1644 .insert((new.kind.collection().to_string(), pk_str), new.id);
1645 }
1646 }
1647 }
1648
1649 let new_refs = new.cross_refs();
1651 if old.cross_refs.as_slice() != new_refs {
1652 self.cross_ref_forward.remove(&new.id);
1654 for cross_ref in &old.cross_refs {
1656 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1657 rev.retain(|(src, _)| *src != new.id);
1658 }
1659 }
1660 for cross_ref in new_refs {
1662 self.cross_ref_forward
1663 .entry(cross_ref.source)
1664 .or_default()
1665 .push((cross_ref.target, cross_ref.ref_type));
1666 self.cross_ref_reverse
1667 .entry(cross_ref.target)
1668 .or_default()
1669 .push((cross_ref.source, cross_ref.ref_type));
1670 }
1671 }
1672 }
1673
1674 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1676 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1677 }
1678
1679 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1681 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1682 }
1683
1684 pub fn age_secs(&self) -> u64 {
1686 let now = current_unix_secs();
1687 now.saturating_sub(self.created_at)
1688 }
1689
1690 pub fn idle_secs(&self) -> u64 {
1692 let now = current_unix_secs();
1693 now.saturating_sub(self.last_write_at)
1694 }
1695
1696 pub fn bulk_insert(
1705 &mut self,
1706 entities: Vec<UnifiedEntity>,
1707 ) -> Result<Vec<EntityId>, SegmentError> {
1708 if !self.state.is_writable() {
1709 return Err(SegmentError::NotWritable);
1710 }
1711
1712 let indexes = self.unique_key_lookup.get_mut();
1713 for index in indexes.iter_mut() {
1714 for entity in &entities {
1715 index.insert(entity);
1716 }
1717 }
1718 self.unique_key_lookup_bytes.store(
1719 indexes
1720 .iter()
1721 .map(|index| index.memory_bytes() as u64)
1722 .sum(),
1723 Ordering::Relaxed,
1724 );
1725
1726 let n = entities.len();
1727
1728 if entities.is_empty() {
1729 return Ok(Vec::new());
1730 }
1731
1732 for run in
1735 entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
1736 {
1737 let kind = run[0].kind.storage_type();
1738 self.observe_graph_kind(&run[0].kind);
1739 if matches!(
1740 run[0].kind,
1741 EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
1742 ) {
1743 if let Some(index) = self.graph_read_index.get_mut() {
1744 for entity in run {
1745 index.insert(entity);
1746 }
1747 self.graph_read_index_bytes
1748 .store(index.memory_bytes(), Ordering::Relaxed);
1749 }
1750 }
1751 let kind_set = match self.kind_index.get_mut(kind) {
1752 Some(ids) => ids,
1753 None => self.kind_index.entry(kind.to_string()).or_default(),
1754 };
1755 kind_set.reserve(run.len());
1756 kind_set.extend(run.iter().map(|entity| entity.id));
1757 }
1758
1759 let now = current_unix_secs();
1760
1761 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
1762
1763 let mut ids = Vec::with_capacity(n);
1764
1765 if self.flat_entities.is_empty() && self.entities.is_empty() {
1767 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
1769 self.use_flat = true;
1770 }
1771
1772 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
1782 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
1783 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
1787
1788 let mut batch_bytes: usize = 0;
1793
1794 if self.use_flat {
1795 self.flat_entities.reserve(n);
1796 for (i, mut entity) in entities.into_iter().enumerate() {
1797 entity.sequence_id = base_seq + i as u64;
1798 let id = entity.id;
1799 ids.push(id);
1800 batch_bytes += Self::estimate_entity_size(&entity);
1801 if let EntityData::Row(row) = &entity.data {
1802 if row.schema.is_some() && !row.columns.is_empty() {
1803 if columnar_zone_updates.is_empty() {
1804 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
1805 columnar_schema = row.schema.clone();
1806 }
1807 for (ci, val) in row.columns.iter().enumerate() {
1808 if !matches!(val, Value::Null) {
1809 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
1810 bucket.push(val.clone());
1811 }
1812 }
1813 }
1814 } else {
1815 for (col, val) in row.iter_fields() {
1816 if !matches!(val, Value::Null) {
1817 named_zone_updates.push((col.to_string(), val.clone()));
1818 }
1819 }
1820 }
1821 }
1822 let expected = self.base_entity_id + self.flat_entities.len() as u64;
1838 if id.raw() == expected {
1839 self.flat_entities.push(entity);
1840 } else {
1841 self.hash_scan_ids.push(id);
1842 self.entities.insert(id, entity);
1843 }
1844 }
1845 } else {
1846 self.entities.reserve(n);
1848 self.hash_scan_ids.reserve(n);
1849 let mut pairs = Vec::with_capacity(n);
1850 for (i, mut entity) in entities.into_iter().enumerate() {
1851 entity.sequence_id = base_seq + i as u64;
1852 let id = entity.id;
1853 ids.push(id);
1854 batch_bytes += Self::estimate_entity_size(&entity);
1855 if let EntityData::Row(row) = &entity.data {
1856 for (col, val) in row.iter_fields() {
1857 if !matches!(val, Value::Null) {
1858 named_zone_updates.push((col.to_string(), val.clone()));
1859 }
1860 }
1861 }
1862 self.hash_scan_ids.push(id);
1863 pairs.push((id, entity));
1864 }
1865 self.entities.extend(pairs);
1866 }
1867
1868 if !columnar_zone_updates.is_empty() {
1872 let schema = columnar_schema.as_ref();
1873 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
1874 if values.is_empty() {
1875 continue;
1876 }
1877 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
1878 continue;
1879 };
1880 let mut iter = values.into_iter();
1886 if let Some(first) = iter.next() {
1887 let previous = self
1888 .col_zones
1889 .get(col_name)
1890 .map_or(0, ColZone::payload_bytes);
1891 let zone = self
1892 .col_zones
1893 .entry(col_name.clone())
1894 .and_modify(|z| z.update(&first))
1895 .or_insert_with(|| ColZone::new(first));
1896 for v in iter {
1897 zone.update(&v);
1898 }
1899 self.zone_payload_bytes = self
1900 .zone_payload_bytes
1901 .saturating_sub(previous)
1902 .saturating_add(zone.payload_bytes());
1903 }
1904 }
1905 }
1906 for (col, val) in named_zone_updates {
1907 self.update_column_zone(&col, &val);
1908 }
1909
1910 self.add_memory(batch_bytes);
1911 self.last_write_at = now;
1912
1913 if self.use_flat {
1915 self.published_flat_len
1916 .store(self.flat_entities.len(), Ordering::Release);
1917 }
1918
1919 Ok(ids)
1920 }
1921
1922 pub(crate) fn force_delete(&mut self, id: EntityId) -> bool {
1925 if self.use_flat {
1926 let raw = id.raw();
1927 if raw >= self.base_entity_id {
1928 let idx = (raw - self.base_entity_id) as usize;
1929 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
1930 self.tombstone_flat_entity(idx, id);
1931 return true;
1932 }
1933 }
1934 return self.remove_hashmap_entity(id);
1935 }
1936
1937 self.remove_hashmap_entity(id)
1938 }
1939
1940 pub(crate) fn force_update_with_metadata(
1943 &mut self,
1944 entity: &UnifiedEntity,
1945 modified_columns: &[String],
1946 metadata: Option<&Metadata>,
1947 ) -> Result<(), SegmentError> {
1948 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
1949 }
1950}
1951
1952impl UnifiedSegment for GrowingSegment {
1953 fn id(&self) -> SegmentId {
1954 self.id
1955 }
1956
1957 fn state(&self) -> SegmentState {
1958 self.state
1959 }
1960
1961 fn collection(&self) -> &str {
1962 &self.collection
1963 }
1964
1965 fn stats(&self) -> SegmentStats {
1966 let mut stats = SegmentStats {
1967 entity_count: self.live_entity_count(),
1968 deleted_count: self.deleted.len(),
1969 memory_bytes: self.resident_bytes() as usize,
1970 ..Default::default()
1971 };
1972
1973 for entity in self.entities.values() {
1974 match &entity.kind {
1975 EntityKind::TableRow { .. } => stats.row_count += 1,
1976 EntityKind::GraphNode(_) => stats.node_count += 1,
1977 EntityKind::GraphEdge(_) => stats.edge_count += 1,
1978 EntityKind::Vector { .. } => stats.vector_count += 1,
1979 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
1980 EntityKind::QueueMessage { .. } => stats.row_count += 1,
1981 }
1982 stats.cross_ref_count += entity.cross_refs().len();
1983 }
1984
1985 stats
1986 }
1987
1988 fn entity_count(&self) -> usize {
1989 self.live_entity_count()
1990 }
1991
1992 fn contains(&self, id: EntityId) -> bool {
1993 self.has_live_entity(id)
1994 }
1995
1996 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
1997 if self.deleted.contains_key(&id) {
1998 return None;
1999 }
2000 if self.use_flat {
2001 let raw = id.raw();
2002 if raw >= self.base_entity_id {
2003 let idx = (raw - self.base_entity_id) as usize;
2004 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
2005 return Some(entity);
2006 }
2007 }
2008 self.entities.get(&id)
2014 } else {
2015 self.entities.get(&id)
2016 }
2017 }
2018
2019 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
2020 self.invalidate_unique_key_lookup();
2021 self.invalidate_graph_read_index();
2022 if self.deleted.contains_key(&id) || !self.state.is_writable() {
2023 return None;
2024 }
2025 self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
2028 if self.use_flat {
2029 let raw = id.raw();
2030 if raw >= self.base_entity_id {
2031 let idx = (raw - self.base_entity_id) as usize;
2032 if self
2033 .flat_entities
2034 .get(idx)
2035 .map(|e| e.id == id)
2036 .unwrap_or(false)
2037 {
2038 return self.flat_entities.get_mut(idx);
2039 }
2040 }
2041 self.entities.get_mut(&id)
2044 } else {
2045 self.entities.get_mut(&id)
2046 }
2047 }
2048
2049 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2050 if !self.state.is_writable() {
2051 return Err(SegmentError::NotWritable);
2052 }
2053
2054 if self.entities.contains_key(&entity.id) {
2055 return Err(SegmentError::AlreadyExists(entity.id));
2056 }
2057
2058 entity.sequence_id = self.next_sequence();
2060
2061 let size = Self::estimate_entity_size(&entity);
2063 self.add_memory(size);
2064
2065 self.index_entity(&entity);
2067
2068 self.update_col_zones_from_entity(&entity);
2070
2071 let id = entity.id;
2073 self.entities.insert(id, entity);
2074
2075 self.last_write_at = current_unix_secs();
2077
2078 Ok(id)
2079 }
2080
2081 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2082 if !self.state.is_writable() {
2083 return Err(SegmentError::NotWritable);
2084 }
2085
2086 self.apply_update_with_metadata(&entity, None)?;
2087 self.last_write_at = current_unix_secs();
2088
2089 Ok(())
2090 }
2091
2092 fn update_hot(
2093 &mut self,
2094 entity: UnifiedEntity,
2095 modified_columns: &[String],
2096 ) -> Result<(), SegmentError> {
2097 if !self.state.is_writable() {
2098 return Err(SegmentError::NotWritable);
2099 }
2100
2101 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
2102 self.last_write_at = current_unix_secs();
2103 Ok(())
2104 }
2105
2106 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2107 if !self.state.is_writable() {
2108 return Err(SegmentError::NotWritable);
2109 }
2110
2111 if self.use_flat {
2113 let raw = id.raw();
2114 if raw >= self.base_entity_id {
2115 let idx = (raw - self.base_entity_id) as usize;
2116 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
2117 self.tombstone_flat_entity(idx, id);
2118 return Ok(true);
2119 }
2120 }
2121 return Ok(self.remove_hashmap_entity(id));
2122 }
2123
2124 Ok(self.remove_hashmap_entity(id))
2125 }
2126
2127 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2128 if !self.has_live_entity(id) {
2129 return None;
2130 }
2131 Some(self.metadata.get_all(id))
2132 }
2133
2134 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2135 if !self.state.is_writable() {
2136 return Err(SegmentError::NotWritable);
2137 }
2138
2139 if !self.has_live_entity(id) {
2140 return Err(SegmentError::NotFound(id));
2141 }
2142
2143 self.metadata.set_all(id, &metadata);
2144 Ok(())
2145 }
2146
2147 fn seal(&mut self) -> Result<(), SegmentError> {
2148 if self.state != SegmentState::Growing {
2149 return Err(SegmentError::InvalidState(self.state));
2150 }
2151
2152 self.state = SegmentState::Sealing;
2153
2154 self.rebuild_sealed_col_zones();
2159
2160 self.state = SegmentState::Sealed;
2161 Ok(())
2162 }
2163
2164 fn should_seal(&self, config: &SegmentConfig) -> bool {
2165 if self.entities.len() >= config.max_entities {
2167 return true;
2168 }
2169
2170 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
2172 return true;
2173 }
2174
2175 if self.age_secs() >= config.max_age_secs {
2177 return true;
2178 }
2179
2180 false
2181 }
2182
2183 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2184 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2185 Box::new(self.flat_entities.iter().chain(self.entities.values()))
2189 } else {
2190 Box::new(self.entities.values())
2191 };
2192 if self.deleted.is_empty() {
2193 base
2194 } else {
2195 Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
2196 }
2197 }
2198
2199 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2200 let ids = self.kind_index.get(kind_filter);
2203 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
2206 Box::new(self.flat_entities.iter())
2207 } else {
2208 Box::new(std::iter::empty())
2209 };
2210 Box::new(flat.chain(self.entities.values()).filter(move |e| {
2211 if self.deleted.contains_key(&e.id) {
2212 return false;
2213 }
2214 if let Some(ids) = ids {
2215 ids.contains(&e.id)
2216 } else {
2217 false
2218 }
2219 }))
2220 }
2221
2222 fn filter_metadata(
2223 &self,
2224 filters: &[(String, super::metadata::MetadataFilter)],
2225 ) -> Vec<EntityId> {
2226 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
2229 Box::new(self.flat_entities.iter().map(|e| e.id))
2230 } else {
2231 Box::new(std::iter::empty())
2232 };
2233 flat_ids
2234 .chain(self.entities.keys().copied())
2235 .filter(|id| {
2236 if self.deleted.contains_key(id) {
2237 return false;
2238 }
2239 let metadata = self.metadata.get_all(*id);
2240 metadata.matches_all(filters)
2241 })
2242 .collect()
2243 }
2244}
2245
2246fn build_minmax_multi_intervals(
2247 entries: &[(CanonicalKey, Value)],
2248 max_intervals: usize,
2249) -> Vec<ColZone> {
2250 if entries.is_empty() {
2251 return Vec::new();
2252 }
2253 if entries.len() == 1 || max_intervals <= 1 {
2254 return vec![ColZone::with_bounds(
2255 entries[0].1.clone(),
2256 entries[entries.len() - 1].1.clone(),
2257 )];
2258 }
2259
2260 let mut split_points = if entries.len() <= max_intervals {
2261 (1..entries.len()).collect::<Vec<_>>()
2262 } else {
2263 let target_splits = max_intervals - 1;
2264 let mut selected = select_gap_split_points(entries, target_splits);
2265 if selected.len() < target_splits {
2266 for bucket in 1..max_intervals {
2267 let idx = bucket * entries.len() / max_intervals;
2268 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
2269 continue;
2270 }
2271 selected.push(idx);
2272 if selected.len() >= target_splits {
2273 break;
2274 }
2275 }
2276 }
2277 selected.sort_unstable();
2278 selected.dedup();
2279 selected
2280 };
2281
2282 split_points.push(entries.len());
2283
2284 let mut out = Vec::with_capacity(split_points.len());
2285 let mut start = 0usize;
2286 for end in split_points {
2287 if end <= start {
2288 continue;
2289 }
2290 out.push(ColZone::with_bounds(
2291 entries[start].1.clone(),
2292 entries[end - 1].1.clone(),
2293 ));
2294 start = end;
2295 }
2296
2297 if out.is_empty() {
2298 out.push(ColZone::with_bounds(
2299 entries[0].1.clone(),
2300 entries[entries.len() - 1].1.clone(),
2301 ));
2302 }
2303
2304 out
2305}
2306
2307fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
2308 let mut gaps = Vec::new();
2309 for idx in 1..entries.len() {
2310 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
2311 if score > 0.0 {
2312 gaps.push((score, idx));
2313 }
2314 }
2315 }
2316 gaps.sort_by(|left, right| {
2317 right
2318 .0
2319 .partial_cmp(&left.0)
2320 .unwrap_or(std::cmp::Ordering::Equal)
2321 .then_with(|| left.1.cmp(&right.1))
2322 });
2323 gaps.into_iter()
2324 .take(max_splits)
2325 .map(|(_, idx)| idx)
2326 .collect()
2327}
2328
2329fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
2330 if left.family() != right.family() {
2331 return None;
2332 }
2333 match (left, right) {
2334 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
2335 Some(r.saturating_sub(*l) as f64)
2336 }
2337 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
2338 Some(r.saturating_sub(*l) as f64)
2339 }
2340 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
2341 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
2342 }
2343 _ => None,
2344 }
2345}
2346
2347#[cfg(test)]
2348mod tests {
2349 use super::*;
2350 use crate::storage::unified::entity::RowData;
2351 use crate::storage::unified::MetadataValue;
2352 use reddb_types::Value;
2353
2354 fn pruning_node(id: u64) -> UnifiedEntity {
2355 UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
2356 }
2357
2358 fn pruning_edge(id: u64) -> UnifiedEntity {
2359 UnifiedEntity::graph_edge(
2360 EntityId::new(id),
2361 "mixed",
2362 "a",
2363 "b",
2364 1.0,
2365 Default::default(),
2366 )
2367 }
2368
2369 #[test]
2370 fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
2371 let mut segment = GrowingSegment::new(1, "mixed");
2372 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
2373 assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
2374 segment
2375 .bulk_insert(vec![
2376 UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
2377 pruning_node(2),
2378 pruning_edge(4),
2379 ])
2380 .expect("mixed bulk with gap");
2381 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2382 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2383 segment.seal().expect("seal");
2384 assert!(segment.force_delete(EntityId::new(2)));
2385 assert!(segment.force_delete(EntityId::new(4)));
2386 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2388 let mut rebuilt = GrowingSegment::new(2, "mixed");
2389 for item in segment.iter() {
2390 rebuilt.adopt_entity(item.clone(), None);
2391 }
2392 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2393 assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2394 rebuilt.adopt_entity(pruning_node(5), None);
2395 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
2396 rebuilt.insert(pruning_edge(6)).expect("single edge");
2397 assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
2398 }
2399
2400 #[test]
2401 fn graph_presence_covers_replacement_and_unrestricted_mutation() {
2402 for flat in [false, true] {
2403 for operation in 0..4 {
2404 let mut segment = GrowingSegment::new(1, "mixed");
2405 let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
2406 if flat {
2407 segment.bulk_insert(vec![vector]).expect("flat vector");
2408 } else {
2409 segment.insert(vector).expect("hashmap vector");
2410 }
2411 match operation {
2412 0 => segment.update(pruning_node(1)).expect("replace with node"),
2413 1 => segment
2414 .update_hot(pruning_node(1), &[])
2415 .expect("hot replacement"),
2416 2 => {
2417 segment.seal().expect("seal vector");
2418 segment
2419 .force_update_with_metadata(&pruning_node(1), &[], None)
2420 .expect("sealed replacement");
2421 }
2422 _ => {
2423 *segment.get_mut(EntityId::new(1)).expect("mutable vector") =
2424 pruning_node(1)
2425 }
2426 }
2427 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2428 assert!(matches!(
2429 segment.get(EntityId::new(1)).expect("node").kind,
2430 EntityKind::GraphNode(_)
2431 ));
2432 if operation == 3 {
2433 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2434 }
2435 segment
2436 .force_update_with_metadata(&pruning_edge(1), &[], None)
2437 .expect("replace with edge");
2438 assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
2439 assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
2440 }
2441 }
2442 }
2443
2444 fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
2445 UnifiedEntity::new(
2446 EntityId::new(id),
2447 EntityKind::TableRow {
2448 table: "keys".into(),
2449 row_id: id,
2450 },
2451 EntityData::Row(RowData::with_names(
2452 vec![Value::text("payload"), key],
2453 vec!["payload".into(), "primary".into()],
2454 )),
2455 )
2456 }
2457
2458 fn primary_key_probe(
2459 segment: &GrowingSegment,
2460 key: Value,
2461 exclude: Option<EntityId>,
2462 ) -> Option<EntityId> {
2463 segment.find_unique_key_conflict(
2464 &[vec!["primary".into()]],
2465 0,
2466 &[format!("{key:?}")],
2467 exclude,
2468 &|_| true,
2469 )
2470 }
2471
2472 #[test]
2473 fn unique_key_lookup_follows_segment_mutations_and_sealing() {
2474 let mut segment = GrowingSegment::new(1, "keys");
2475 segment
2476 .bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
2477 .expect("bulk row");
2478 let payload_bytes = segment.memory_bytes();
2479 assert_eq!(
2480 primary_key_probe(&segment, Value::Integer(10), None),
2481 Some(EntityId::new(1))
2482 );
2483 assert!(segment.memory_bytes() > payload_bytes);
2484 segment
2485 .insert(primary_key_row(3, Value::Integer(30)))
2486 .expect("hashmap fallback row");
2487 segment
2488 .bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
2489 .expect("bulk gap row");
2490 assert_eq!(
2491 primary_key_probe(&segment, Value::Integer(40), None),
2492 Some(EntityId::new(4))
2493 );
2494 assert_eq!(
2495 primary_key_probe(&segment, Value::Integer(30), None),
2496 Some(EntityId::new(3))
2497 );
2498 segment
2499 .update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
2500 .expect("replace flat key");
2501 assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
2502 assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
2503 assert_eq!(
2504 primary_key_probe(&segment, Value::Integer(11), None),
2505 Some(EntityId::new(1))
2506 );
2507 segment
2508 .delete_batch(&[EntityId::new(1), EntityId::new(3)])
2509 .expect("delete flat and map");
2510 assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
2511 assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
2512 segment.seal().expect("seal");
2513 assert_eq!(
2514 primary_key_probe(&segment, Value::Integer(40), None),
2515 Some(EntityId::new(4))
2516 );
2517 segment
2518 .force_update_with_metadata(
2519 &primary_key_row(4, Value::Integer(41)),
2520 &["primary".into()],
2521 None,
2522 )
2523 .expect("sealed update");
2524 assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
2525 assert_eq!(
2526 primary_key_probe(&segment, Value::Integer(41), None),
2527 Some(EntityId::new(4))
2528 );
2529 assert!(segment.force_delete(EntityId::new(4)));
2530 assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
2531 }
2532
2533 #[test]
2534 fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
2535 let mut segment = GrowingSegment::new(1, "keys");
2536 segment
2537 .bulk_insert(vec![
2538 primary_key_row(1, Value::Integer(10)),
2539 primary_key_row(2, Value::Integer(20)),
2540 ])
2541 .expect("initial rows");
2542 let definitions = vec![
2543 vec!["primary".into()],
2544 vec!["payload".into()],
2545 vec!["primary".into(), "payload".into()],
2546 ];
2547 let payload = format!("{:?}", Value::text("payload"));
2548 assert_eq!(
2549 segment.find_unique_key_conflict(
2550 &definitions,
2551 1,
2552 &[payload.clone()],
2553 None,
2554 &|entity| entity.id != EntityId::new(1)
2555 ),
2556 Some(EntityId::new(2))
2557 );
2558 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2559 segment
2560 .insert(primary_key_row(3, Value::Integer(30)))
2561 .expect("append");
2562 assert_eq!(
2563 segment.find_unique_key_conflict(
2564 &definitions,
2565 2,
2566 &[format!("{:?}", Value::Integer(30)), payload.clone()],
2567 None,
2568 &|_| true
2569 ),
2570 Some(EntityId::new(3))
2571 );
2572 let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
2573 assert_eq!(
2574 segment.find_unique_key_conflict(
2575 &definitions,
2576 0,
2577 &[format!("{:?}", Value::Integer(20))],
2578 None,
2579 &|_| true
2580 ),
2581 Some(EntityId::new(2))
2582 );
2583 assert_eq!(segment.unique_key_lookup.read().len(), 3);
2584 assert_eq!(
2585 segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
2586 bytes
2587 );
2588 let reduced = vec![vec!["payload".into()]];
2589 assert_eq!(
2590 segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
2591 Some(EntityId::new(1))
2592 );
2593 assert_eq!(segment.unique_key_lookup.read().len(), 1);
2594 assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
2595 }
2596
2597 #[test]
2598 fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
2599 let mut segment = GrowingSegment::new(1, "keys");
2600 let values = [
2601 Value::Float(f64::NAN),
2602 Value::Float(-0.0),
2603 Value::Float(0.0),
2604 Value::Integer(1),
2605 Value::UnsignedInteger(1),
2606 Value::text("1"),
2607 ];
2608 for (offset, value) in values.iter().enumerate() {
2609 let id = offset as u64 + 1;
2610 segment
2611 .insert(primary_key_row(id, value.clone()))
2612 .expect("distinct physical row");
2613 assert_eq!(
2614 primary_key_probe(&segment, value.clone(), None),
2615 Some(EntityId::new(id))
2616 );
2617 assert_eq!(
2618 primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
2619 None
2620 );
2621 }
2622 segment
2623 .insert(primary_key_row(20, Value::Integer(1)))
2624 .expect("duplicate physical key");
2625 assert_eq!(
2626 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
2627 Some(EntityId::new(20))
2628 );
2629 segment
2632 .entities
2633 .insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
2634 assert_eq!(
2635 primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
2636 None
2637 );
2638 let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
2640 *entity = primary_key_row(4, Value::Integer(100));
2641 assert_eq!(
2642 primary_key_probe(&segment, Value::Integer(100), None),
2643 Some(EntityId::new(4))
2644 );
2645 }
2646
2647 #[test]
2648 fn test_growing_segment_basic() {
2649 let mut segment = GrowingSegment::new(1, "test");
2650
2651 let entity = UnifiedEntity::table_row(
2652 EntityId::new(1),
2653 "users",
2654 1,
2655 vec![Value::text("Alice".to_string())],
2656 );
2657
2658 let id = segment.insert(entity).unwrap();
2659 assert_eq!(id, EntityId::new(1));
2660 assert!(segment.contains(id));
2661
2662 let stats = segment.stats();
2663 assert_eq!(stats.entity_count, 1);
2664 assert_eq!(stats.row_count, 1);
2665 }
2666
2667 #[test]
2668 fn test_segment_metadata() {
2669 let mut segment = GrowingSegment::new(1, "test");
2670
2671 let entity = UnifiedEntity::table_row(
2672 EntityId::new(1),
2673 "users",
2674 1,
2675 vec![Value::text("Alice".to_string())],
2676 );
2677 segment.insert(entity).unwrap();
2678
2679 let mut meta = Metadata::new();
2680 meta.set("role", MetadataValue::String("admin".to_string()));
2681 meta.set("level", MetadataValue::Int(5));
2682
2683 segment.set_metadata(EntityId::new(1), meta).unwrap();
2684
2685 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
2686 assert_eq!(
2687 retrieved.get("role"),
2688 Some(&MetadataValue::String("admin".to_string()))
2689 );
2690 }
2691
2692 #[test]
2693 fn test_segment_seal() {
2694 let mut segment = GrowingSegment::new(1, "test");
2695
2696 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
2697 segment.insert(entity).unwrap();
2698
2699 assert!(segment.state().is_writable());
2701
2702 segment.seal().unwrap();
2704 assert_eq!(segment.state(), SegmentState::Sealed);
2705
2706 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
2708 assert!(segment.insert(entity2).is_err());
2709 }
2710
2711 #[test]
2712 fn test_should_seal() {
2713 let mut segment = GrowingSegment::new(1, "test");
2714
2715 let config = SegmentConfig {
2716 max_entities: 2,
2717 ..Default::default()
2718 };
2719
2720 assert!(!segment.should_seal(&config));
2721
2722 segment
2723 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
2724 .unwrap();
2725 assert!(!segment.should_seal(&config));
2726
2727 segment
2728 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
2729 .unwrap();
2730 assert!(segment.should_seal(&config));
2731 }
2732
2733 #[test]
2734 fn test_cross_references() {
2735 let mut segment = GrowingSegment::new(1, "test");
2736
2737 let mut entity1 = UnifiedEntity::table_row(
2738 EntityId::new(1),
2739 "hosts",
2740 1,
2741 vec![Value::text("192.168.1.1".to_string())],
2742 );
2743 entity1.add_cross_ref(CrossRef::new(
2744 EntityId::new(1),
2745 EntityId::new(2),
2746 "nodes",
2747 RefType::RowToNode,
2748 ));
2749 segment.insert(entity1).unwrap();
2750
2751 let refs_from = segment.get_references_from(EntityId::new(1));
2752 assert_eq!(refs_from.len(), 1);
2753 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
2754
2755 let refs_to = segment.get_references_to(EntityId::new(2));
2756 assert_eq!(refs_to.len(), 1);
2757 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
2758 }
2759
2760 #[test]
2761 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
2762 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
2763 zone.update(&Value::Email("delta@example.com".to_string()));
2764
2765 let probe = Value::Email("alpha@example.com".to_string());
2766 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
2767
2768 let in_range = Value::Email("charlie@example.com".to_string());
2769 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
2770 }
2771
2772 #[test]
2773 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
2774 let mut segment = GrowingSegment::new(1, "test");
2775
2776 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
2777 let entity = UnifiedEntity::new(
2778 EntityId::new(row_id),
2779 EntityKind::TableRow {
2780 table: "users".into(),
2781 row_id,
2782 },
2783 EntityData::Row(RowData::with_names(
2784 vec![Value::Integer(age)],
2785 vec!["age".to_string()],
2786 )),
2787 );
2788 segment.insert(entity).unwrap();
2789 }
2790
2791 segment.seal().unwrap();
2792
2793 let miss = Value::Integer(500);
2794 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
2795
2796 let hit = Value::Integer(1000);
2797 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
2798 }
2799}