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 crate::storage::query::value_compare::partial_compare_values;
38use crate::storage::schema::{value_to_canonical_key, CanonicalKey, Value};
39
40pub type SegmentId = u64;
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum SegmentState {
46 Growing,
48 Sealing,
50 Sealed,
52 Flushed,
54 Archived,
56}
57
58impl SegmentState {
59 pub fn is_writable(&self) -> bool {
61 matches!(self, Self::Growing)
62 }
63
64 pub fn is_queryable(&self) -> bool {
66 !matches!(self, Self::Sealing)
67 }
68
69 pub fn is_immutable(&self) -> bool {
71 matches!(self, Self::Sealed | Self::Flushed | Self::Archived)
72 }
73}
74
75#[derive(Debug, Clone)]
77pub struct SegmentConfig {
78 pub max_entities: usize,
80 pub max_bytes: usize,
82 pub max_age_secs: u64,
84 pub build_vector_index: bool,
86 pub build_graph_index: bool,
88 pub compression_level: u8,
90}
91
92impl Default for SegmentConfig {
93 fn default() -> Self {
94 Self {
95 max_entities: 100_000,
96 max_bytes: 256 * 1024 * 1024, max_age_secs: 3600, build_vector_index: true,
99 build_graph_index: true,
100 compression_level: 6,
101 }
102 }
103}
104
105#[derive(Debug, Clone, Default)]
107pub struct SegmentStats {
108 pub entity_count: usize,
110 pub deleted_count: usize,
112 pub memory_bytes: usize,
114 pub vector_count: usize,
116 pub node_count: usize,
118 pub edge_count: usize,
120 pub row_count: usize,
122 pub cross_ref_count: usize,
124}
125
126#[derive(Debug, Clone)]
128pub enum SegmentError {
129 NotWritable,
131 NotFound(EntityId),
133 AlreadyExists(EntityId),
135 Full,
137 InvalidState(SegmentState),
139 Internal(String),
141}
142
143impl std::fmt::Display for SegmentError {
144 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145 match self {
146 Self::NotWritable => write!(f, "segment is not writable"),
147 Self::NotFound(id) => write!(f, "entity not found: {}", id),
148 Self::AlreadyExists(id) => write!(f, "entity already exists: {}", id),
149 Self::Full => write!(f, "segment is full"),
150 Self::InvalidState(state) => write!(f, "invalid operation for state: {:?}", state),
151 Self::Internal(msg) => write!(f, "internal error: {}", msg),
152 }
153 }
154}
155
156impl std::error::Error for SegmentError {}
157
158fn current_unix_secs() -> u64 {
159 SystemTime::now()
160 .duration_since(UNIX_EPOCH)
161 .unwrap_or_default()
162 .as_secs()
163}
164
165fn entity_id_from_probe_key(key: &[u8]) -> Option<EntityId> {
166 if let Ok(text) = std::str::from_utf8(key) {
167 if let Ok(raw) = text.parse::<u64>() {
168 return Some(EntityId::new(raw));
169 }
170 }
171
172 let bytes: [u8; 8] = key.try_into().ok()?;
173 Some(EntityId::new(u64::from_le_bytes(bytes)))
174}
175
176const SEALED_MULTI_ZONE_MAX_INTERVALS: usize = 4;
177
178const TOMBSTONE_ENTRY_BYTES: u64 = 16;
186
187#[derive(Debug, Clone)]
188struct UpdateIndexSnapshot {
189 pk_column_name: Option<String>,
190 pk_value: Option<Value>,
191 pk_index_key: Option<(String, String)>,
192 cross_refs: Vec<CrossRef>,
193}
194
195impl UpdateIndexSnapshot {
196 fn from_entity(entity: &UnifiedEntity) -> Self {
197 let (pk_column_name, pk_value) = match &entity.data {
198 EntityData::Row(row) => (
199 row.schema
200 .as_deref()
201 .and_then(|schema| schema.first().cloned()),
202 row.columns.first().cloned(),
203 ),
204 _ => (None, None),
205 };
206 let pk_index_key = pk_value
207 .as_ref()
208 .map(|value| (entity.kind.collection().to_string(), format!("{:?}", value)));
209 Self {
210 pk_column_name,
211 pk_value,
212 pk_index_key,
213 cross_refs: entity.cross_refs().to_vec(),
214 }
215 }
216}
217
218pub trait UnifiedSegment: Send + Sync {
220 fn id(&self) -> SegmentId;
222
223 fn state(&self) -> SegmentState;
225
226 fn collection(&self) -> &str;
228
229 fn stats(&self) -> SegmentStats;
231
232 fn entity_count(&self) -> usize;
234
235 fn contains(&self, id: EntityId) -> bool;
237
238 fn get(&self, id: EntityId) -> Option<&UnifiedEntity>;
240
241 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity>;
243
244 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError>;
246
247 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError>;
249
250 fn update_hot(
254 &mut self,
255 entity: UnifiedEntity,
256 modified_columns: &[String],
257 ) -> Result<(), SegmentError> {
258 let _ = modified_columns;
259 self.update(entity)
260 }
261
262 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError>;
264
265 fn get_metadata(&self, id: EntityId) -> Option<Metadata>;
267
268 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError>;
270
271 fn seal(&mut self) -> Result<(), SegmentError>;
273
274 fn should_seal(&self, config: &SegmentConfig) -> bool;
276
277 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
279
280 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
282
283 fn filter_metadata(
285 &self,
286 filters: &[(String, super::metadata::MetadataFilter)],
287 ) -> Vec<EntityId>;
288}
289
290#[derive(Debug, Clone)]
297pub struct ColZone {
298 pub min: Value,
299 pub max: Value,
300 min_key: Option<CanonicalKey>,
301 max_key: Option<CanonicalKey>,
302}
303
304impl ColZone {
305 fn new(v: Value) -> Self {
306 Self {
307 min_key: value_to_canonical_key(&v),
308 max_key: value_to_canonical_key(&v),
309 min: v.clone(),
310 max: v,
311 }
312 }
313
314 fn with_bounds(min: Value, max: Value) -> Self {
315 Self {
316 min_key: value_to_canonical_key(&min),
317 max_key: value_to_canonical_key(&max),
318 min,
319 max,
320 }
321 }
322
323 fn update(&mut self, v: &Value) {
324 if compare_zone_values(v, None, &self.min, self.min_key.as_ref())
325 .map(|o| o == std::cmp::Ordering::Less)
326 .unwrap_or(false)
327 {
328 self.min = v.clone();
329 self.min_key = value_to_canonical_key(v);
330 }
331 if compare_zone_values(v, None, &self.max, self.max_key.as_ref())
332 .map(|o| o == std::cmp::Ordering::Greater)
333 .unwrap_or(false)
334 {
335 self.max = v.clone();
336 self.max_key = value_to_canonical_key(v);
337 }
338 }
339}
340
341#[derive(Debug, Clone, Default)]
342pub struct MultiColZone {
343 pub intervals: Vec<ColZone>,
344}
345
346impl MultiColZone {
347 fn can_skip(&self, pred: &ZoneColPred<'_>) -> bool {
348 !self.intervals.is_empty() && self.intervals.iter().all(|zone| pred.can_skip(zone))
349 }
350}
351
352fn compare_zone_values(
353 left: &Value,
354 left_key: Option<&CanonicalKey>,
355 right: &Value,
356 right_key: Option<&CanonicalKey>,
357) -> Option<std::cmp::Ordering> {
358 partial_compare_values(left, right).or_else(|| {
359 let left_key = left_key.cloned().or_else(|| value_to_canonical_key(left))?;
360 let right_key = right_key
361 .cloned()
362 .or_else(|| value_to_canonical_key(right))?;
363 (left_key.family() == right_key.family()).then(|| left_key.cmp(&right_key))
364 })
365}
366
367#[derive(Debug, Clone, Copy, PartialEq, Eq)]
369pub enum ZoneColPredKind {
370 Eq,
371 Gt,
372 Gte,
373 Lt,
374 Lte,
375}
376
377#[derive(Debug, Clone)]
379pub enum ZoneColPred<'a> {
380 Eq(&'a Value),
381 Gt(&'a Value),
382 Gte(&'a Value),
383 Lt(&'a Value),
384 Lte(&'a Value),
385}
386
387impl<'a> ZoneColPred<'a> {
388 pub fn can_skip(&self, zone: &ColZone) -> bool {
390 match self {
391 ZoneColPred::Eq(val) => {
393 compare_zone_values(val, None, &zone.min, zone.min_key.as_ref())
394 .map(|o| o == std::cmp::Ordering::Less)
395 .unwrap_or(false)
396 || compare_zone_values(val, None, &zone.max, zone.max_key.as_ref())
397 .map(|o| o == std::cmp::Ordering::Greater)
398 .unwrap_or(false)
399 }
400 ZoneColPred::Gt(val) => {
402 compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
403 .map(|o| o != std::cmp::Ordering::Greater)
404 .unwrap_or(false)
405 }
406 ZoneColPred::Gte(val) => {
408 compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
409 .map(|o| o == std::cmp::Ordering::Less)
410 .unwrap_or(false)
411 }
412 ZoneColPred::Lt(val) => {
414 compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
415 .map(|o| o != std::cmp::Ordering::Less)
416 .unwrap_or(false)
417 }
418 ZoneColPred::Lte(val) => {
420 compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
421 .map(|o| o == std::cmp::Ordering::Greater)
422 .unwrap_or(false)
423 }
424 }
425 }
426}
427
428pub struct GrowingSegment {
430 id: SegmentId,
432 collection: String,
434 state: SegmentState,
436 created_at: u64,
438 last_write_at: u64,
440
441 entities: HashMap<EntityId, UnifiedEntity>,
443 flat_entities: Vec<UnifiedEntity>,
446 base_entity_id: u64,
448 use_flat: bool,
450 deleted: HashSet<EntityId>,
452 metadata: MetadataStorage,
454
455 pk_index: BTreeMap<(String, String), EntityId>,
457 kind_index: HashMap<String, HashSet<EntityId>>,
459 cross_ref_forward: HashMap<EntityId, Vec<(EntityId, RefType)>>,
461 cross_ref_reverse: HashMap<EntityId, Vec<(EntityId, RefType)>>,
463
464 col_zones: HashMap<String, ColZone>,
466 sealed_col_zones: HashMap<String, MultiColZone>,
468
469 sequence: AtomicU64,
471 memory_bytes: AtomicU64,
476 dead_entity_bytes: u64,
479 dead_resident_count: usize,
481
482 pub(crate) published_flat_len: AtomicUsize,
489}
490
491impl GrowingSegment {
492 #[inline]
495 pub fn for_each_fast<F>(&self, mut f: F) -> bool
496 where
497 F: FnMut(&UnifiedEntity) -> bool,
498 {
499 if self.use_flat {
500 if self.deleted.is_empty() {
502 for entity in &self.flat_entities {
503 if !f(entity) {
504 return false;
505 }
506 }
507 for entity in self.entities.values() {
510 if !f(entity) {
511 return false;
512 }
513 }
514 } else {
515 for entity in &self.flat_entities {
516 if self.deleted.contains(&entity.id) {
517 continue;
518 }
519 if !f(entity) {
520 return false;
521 }
522 }
523 for entity in self.entities.values() {
524 if self.deleted.contains(&entity.id) {
525 continue;
526 }
527 if !f(entity) {
528 return false;
529 }
530 }
531 }
532 } else {
533 if self.deleted.is_empty() {
535 for entity in self.entities.values() {
536 if !f(entity) {
537 return false;
538 }
539 }
540 } else {
541 for entity in self.entities.values() {
542 if self.deleted.contains(&entity.id) {
543 continue;
544 }
545 if !f(entity) {
546 return false;
547 }
548 }
549 }
550 }
551 true
552 }
553
554 pub fn new(id: SegmentId, collection: impl Into<String>) -> Self {
556 let now = current_unix_secs();
557
558 Self {
559 id,
560 collection: collection.into(),
561 state: SegmentState::Growing,
562 created_at: now,
563 last_write_at: now,
564 entities: HashMap::new(),
565 flat_entities: Vec::new(),
566 base_entity_id: 0,
567 use_flat: false,
568 deleted: HashSet::new(),
569 metadata: MetadataStorage::new(),
570 pk_index: BTreeMap::new(),
571 kind_index: HashMap::new(),
572 cross_ref_forward: HashMap::new(),
573 cross_ref_reverse: HashMap::new(),
574 col_zones: HashMap::new(),
575 sealed_col_zones: HashMap::new(),
576 sequence: AtomicU64::new(0),
577 memory_bytes: AtomicU64::new(0),
578 dead_entity_bytes: 0,
579 dead_resident_count: 0,
580 published_flat_len: AtomicUsize::new(0),
581 }
582 }
583
584 fn next_sequence(&self) -> u64 {
586 self.sequence.fetch_add(1, Ordering::SeqCst)
587 }
588
589 fn has_live_entity(&self, id: EntityId) -> bool {
590 if self.deleted.contains(&id) {
591 return false;
592 }
593 if self.use_flat {
594 let raw = id.raw();
595 if raw >= self.base_entity_id {
596 let idx = (raw - self.base_entity_id) as usize;
597 if self
598 .flat_entities
599 .get(idx)
600 .is_some_and(|entity| entity.id == id)
601 {
602 return true;
603 }
604 }
605 self.entities.contains_key(&id)
608 } else {
609 self.entities.contains_key(&id)
610 }
611 }
612
613 fn update_existing_entity_in_place(
614 &mut self,
615 entity: &UnifiedEntity,
616 ) -> Result<UpdateIndexSnapshot, SegmentError> {
617 if self.use_flat {
618 let raw = entity.id.raw();
619 if raw < self.base_entity_id {
620 return Err(SegmentError::NotFound(entity.id));
621 }
622 let idx = (raw - self.base_entity_id) as usize;
623 let Some(slot) = self.flat_entities.get_mut(idx) else {
624 return Err(SegmentError::NotFound(entity.id));
625 };
626 if slot.id != entity.id {
627 return Err(SegmentError::NotFound(entity.id));
628 }
629 let snapshot = UpdateIndexSnapshot::from_entity(slot);
630 slot.clone_from(entity);
631 Ok(snapshot)
632 } else {
633 let Some(slot) = self.entities.get_mut(&entity.id) else {
634 return Err(SegmentError::NotFound(entity.id));
635 };
636 let snapshot = UpdateIndexSnapshot::from_entity(slot);
637 slot.clone_from(entity);
638 Ok(snapshot)
639 }
640 }
641
642 fn apply_hot_update_with_metadata(
643 &mut self,
644 entity: &UnifiedEntity,
645 modified_columns: &[String],
646 metadata: Option<&Metadata>,
647 ) -> Result<(), SegmentError> {
648 let old = self.update_existing_entity_in_place(entity)?;
649 self.reindex_for_update(&old, entity, Some(modified_columns));
650 self.update_col_zones_from_entity(entity);
651 if let Some(metadata) = metadata {
652 self.metadata.set_all(entity.id, metadata);
653 }
654 Ok(())
655 }
656
657 fn apply_update_with_metadata(
658 &mut self,
659 entity: &UnifiedEntity,
660 metadata: Option<&Metadata>,
661 ) -> Result<(), SegmentError> {
662 let old = self.update_existing_entity_in_place(entity)?;
663 self.reindex_for_update(&old, entity, None);
664 self.update_col_zones_from_entity(entity);
665 if let Some(metadata) = metadata {
666 self.metadata.set_all(entity.id, metadata);
667 }
668 Ok(())
669 }
670
671 pub fn update_hot_batch_with_metadata<'a, I>(&mut self, items: I) -> Result<(), SegmentError>
672 where
673 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
674 {
675 if !self.state.is_writable() {
676 return Err(SegmentError::NotWritable);
677 }
678
679 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
680 items.into_iter().collect();
681 if items.is_empty() {
682 return Ok(());
683 }
684
685 for (entity, _, _) in &items {
686 if !self.has_live_entity(entity.id) {
687 return Err(SegmentError::NotFound(entity.id));
688 }
689 }
690
691 for (entity, modified_columns, metadata) in items {
692 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)?;
693 }
694
695 self.last_write_at = current_unix_secs();
696 Ok(())
697 }
698
699 fn tombstone_flat_entity(&mut self, idx: usize, id: EntityId) -> bool {
705 if !self.deleted.insert(id) {
706 return false;
707 }
708 let size = Self::estimate_entity_size(&self.flat_entities[idx]) as u64;
709 self.dead_entity_bytes += size;
710 self.dead_resident_count += 1;
711 self.metadata.remove_all(id);
712 true
713 }
714
715 fn remove_hashmap_entity(&mut self, id: EntityId) -> bool {
720 let Some(entity) = self.entities.remove(&id) else {
721 return false;
722 };
723 self.release_memory(Self::estimate_entity_size(&entity));
724 self.unindex_entity(&entity);
725 self.metadata.remove_all(id);
726 self.deleted.insert(id);
727 true
728 }
729
730 pub fn delete_batch(&mut self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
731 if !self.state.is_writable() {
732 return Err(SegmentError::NotWritable);
733 }
734 if ids.is_empty() {
735 return Ok(Vec::new());
736 }
737
738 let mut deleted_ids = Vec::with_capacity(ids.len());
739
740 if self.use_flat {
741 for &id in ids {
742 let raw = id.raw();
743 if raw < self.base_entity_id {
744 if self.remove_hashmap_entity(id) {
745 deleted_ids.push(id);
746 }
747 continue;
748 }
749 let idx = (raw - self.base_entity_id) as usize;
750 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
751 if self.tombstone_flat_entity(idx, id) {
752 deleted_ids.push(id);
753 }
754 } else if self.remove_hashmap_entity(id) {
755 deleted_ids.push(id);
756 }
757 }
758 } else {
759 for &id in ids {
760 if self.remove_hashmap_entity(id) {
761 deleted_ids.push(id);
762 }
763 }
764 }
765
766 if !deleted_ids.is_empty() {
767 self.last_write_at = current_unix_secs();
768 }
769
770 Ok(deleted_ids)
771 }
772
773 fn add_memory(&self, bytes: usize) {
775 self.memory_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
776 }
777
778 pub fn memory_bytes(&self) -> u64 {
783 self.memory_bytes.load(Ordering::Relaxed)
784 }
785
786 fn release_memory(&self, bytes: usize) {
789 let mut current = self.memory_bytes.load(Ordering::Relaxed);
790 loop {
791 let next = current.saturating_sub(bytes as u64);
792 match self.memory_bytes.compare_exchange_weak(
793 current,
794 next,
795 Ordering::Relaxed,
796 Ordering::Relaxed,
797 ) {
798 Ok(_) => break,
799 Err(observed) => current = observed,
800 }
801 }
802 }
803
804 pub(crate) fn resident_entity_count(&self) -> usize {
806 self.flat_entities.len() + self.entities.len()
807 }
808
809 pub(crate) fn live_entity_count(&self) -> usize {
811 self.resident_entity_count()
812 .saturating_sub(self.dead_resident_count)
813 }
814
815 pub(crate) fn tombstone_count(&self) -> usize {
818 self.deleted.len()
819 }
820
821 pub(crate) fn resident_bytes(&self) -> u64 {
824 self.memory_bytes.load(Ordering::Relaxed)
825 + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
826 }
827
828 pub(crate) fn reclaimable_bytes(&self) -> u64 {
831 self.dead_entity_bytes + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
832 }
833
834 pub(crate) fn live_entity_ids(&self) -> Vec<EntityId> {
840 let mut ids = Vec::with_capacity(self.live_entity_count());
841 self.for_each_fast(|entity| {
842 ids.push(entity.id);
843 true
844 });
845 ids
846 }
847
848 pub(crate) fn adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
856 let id = entity.id;
857
858 let next = entity.sequence_id.saturating_add(1);
861 if self.sequence.load(Ordering::Relaxed) < next {
862 self.sequence.store(next, Ordering::Relaxed);
863 }
864
865 self.add_memory(Self::estimate_entity_size(&entity));
866 self.index_entity(&entity);
867 self.update_col_zones_from_entity(&entity);
868 self.entities.insert(id, entity);
869
870 if let Some(metadata) = metadata {
871 if !metadata.is_empty() {
872 self.metadata.set_all(id, &metadata);
873 }
874 }
875 }
876
877 pub(crate) fn evict_entity(&mut self, id: EntityId) -> bool {
883 let Some(entity) = self.entities.remove(&id) else {
884 return false;
885 };
886 self.release_memory(Self::estimate_entity_size(&entity));
887 self.unindex_entity(&entity);
888 for cross_ref in entity.cross_refs() {
889 if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
890 reverse.retain(|(source, _)| *source != id);
891 }
892 }
893 self.metadata.remove_all(id);
894 true
895 }
896
897 fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
899 let mut size = std::mem::size_of::<UnifiedEntity>();
900
901 size += match &entity.data {
903 EntityData::Row(row) => row.columns.len() * 64, EntityData::Node(node) => node.properties.len() * 128,
905 EntityData::Edge(edge) => edge.properties.len() * 128,
906 EntityData::Vector(vec) => {
907 vec.dense.len() * 4 + vec.sparse.as_ref().map_or(0, |s| s.indices.len() * 8)
908 }
909 EntityData::TimeSeries(_) => 64,
910 EntityData::QueueMessage(_) => 128,
911 };
912
913 for emb in entity.embeddings() {
915 size += emb.vector.len() * 4 + emb.name.len() + emb.model.len();
916 }
917
918 size += std::mem::size_of_val(entity.cross_refs());
920
921 size
922 }
923
924 fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
932 if let EntityData::Row(row) = &entity.data {
933 if let Some(named) = &row.named {
934 for (col, val) in named {
936 if matches!(val, Value::Null) {
937 continue;
938 }
939 self.col_zones
940 .entry(col.clone())
941 .and_modify(|z| z.update(val))
942 .or_insert_with(|| ColZone::new(val.clone()));
943 }
944 } else if let Some(schema) = &row.schema {
945 for (col, val) in schema.iter().zip(row.columns.iter()) {
948 if matches!(val, Value::Null) {
949 continue;
950 }
951 self.col_zones
952 .entry(col.clone())
953 .and_modify(|z| z.update(val))
954 .or_insert_with(|| ColZone::new(val.clone()));
955 }
956 }
957 }
958 }
959
960 fn rebuild_sealed_col_zones(&mut self) {
961 let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
962 let mut family_by_col: HashMap<String, crate::storage::schema::CanonicalKeyFamily> =
963 HashMap::new();
964 let mut mixed_family_cols = HashSet::new();
965 let mut unsupported_cols = HashSet::new();
966
967 let mut observe_row = |row: &super::entity::RowData| {
968 for (col, value) in row.iter_fields() {
969 if matches!(value, Value::Null) {
970 continue;
971 }
972 let Some(key) = value_to_canonical_key(value) else {
973 unsupported_cols.insert(col.to_string());
974 continue;
975 };
976 match family_by_col.get(col).copied() {
977 Some(existing) if existing != key.family() => {
978 mixed_family_cols.insert(col.to_string());
979 }
980 None => {
981 family_by_col.insert(col.to_string(), key.family());
982 }
983 _ => {}
984 }
985 values_by_col
986 .entry(col.to_string())
987 .or_default()
988 .push((key, value.clone()));
989 }
990 };
991
992 if self.use_flat {
993 for entity in &self.flat_entities {
994 if self.deleted.contains(&entity.id) {
995 continue;
996 }
997 if let EntityData::Row(row) = &entity.data {
998 observe_row(row);
999 }
1000 }
1001 } else {
1002 for entity in self.entities.values() {
1003 if self.deleted.contains(&entity.id) {
1004 continue;
1005 }
1006 if let EntityData::Row(row) = &entity.data {
1007 observe_row(row);
1008 }
1009 }
1010 }
1011
1012 let mut sealed_col_zones = HashMap::new();
1013 for (col, mut entries) in values_by_col {
1014 if mixed_family_cols.contains(&col)
1015 || unsupported_cols.contains(&col)
1016 || entries.is_empty()
1017 {
1018 continue;
1019 }
1020 entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
1021 entries.dedup_by(|left, right| left.0 == right.0);
1022
1023 let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
1024 if intervals.len() > 1 {
1025 sealed_col_zones.insert(col, MultiColZone { intervals });
1026 }
1027 }
1028
1029 self.sealed_col_zones = sealed_col_zones;
1030 }
1031
1032 pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
1036 if preds.is_empty() {
1037 return false;
1038 }
1039 for (col, pred) in preds {
1040 if let Some(zone) = self.sealed_col_zones.get(*col) {
1041 if zone.can_skip(pred) {
1042 return true;
1043 }
1044 continue;
1045 }
1046 if let Some(zone) = self.col_zones.get(*col) {
1047 if pred.can_skip(zone) {
1048 return true; }
1050 }
1051 }
1052 false
1053 }
1054
1055 fn index_entity(&mut self, entity: &UnifiedEntity) {
1057 let kind_key = entity.kind.storage_type().to_string();
1059 self.kind_index
1060 .entry(kind_key)
1061 .or_default()
1062 .insert(entity.id);
1063
1064 if let EntityData::Row(row) = &entity.data {
1066 if let Some(first_col) = row.columns.first() {
1067 let pk_str = format!("{:?}", first_col);
1068 self.pk_index
1069 .insert((entity.kind.collection().to_string(), pk_str), entity.id);
1070 }
1071 }
1072
1073 for cross_ref in entity.cross_refs() {
1075 self.cross_ref_forward
1076 .entry(cross_ref.source)
1077 .or_default()
1078 .push((cross_ref.target, cross_ref.ref_type));
1079
1080 self.cross_ref_reverse
1081 .entry(cross_ref.target)
1082 .or_default()
1083 .push((cross_ref.source, cross_ref.ref_type));
1084 }
1085 }
1086
1087 pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
1089 if let Some(id) = entity_id_from_probe_key(key) {
1090 return self.has_live_entity(id);
1091 }
1092
1093 let Ok(pk_str) = std::str::from_utf8(key) else {
1094 return false;
1095 };
1096 let pk_key = (self.collection.clone(), pk_str.to_string());
1097 self.pk_index
1098 .get(&pk_key)
1099 .is_some_and(|id| self.has_live_entity(*id))
1100 }
1101
1102 fn unindex_entity(&mut self, entity: &UnifiedEntity) {
1104 let kind_key = entity.kind.storage_type().to_string();
1106 if let Some(set) = self.kind_index.get_mut(&kind_key) {
1107 set.remove(&entity.id);
1108 }
1109
1110 if let EntityData::Row(row) = &entity.data {
1112 if let Some(first_col) = row.columns.first() {
1113 let pk_str = format!("{:?}", first_col);
1114 self.pk_index
1115 .remove(&(entity.kind.collection().to_string(), pk_str));
1116 }
1117 }
1118
1119 self.cross_ref_forward.remove(&entity.id);
1121 }
1123
1124 fn reindex_for_update(
1134 &mut self,
1135 old: &UpdateIndexSnapshot,
1136 new: &UnifiedEntity,
1137 modified_columns: Option<&[String]>,
1138 ) {
1139 let pk_changed = match &new.data {
1144 EntityData::Row(new_row) => {
1145 if let Some(cols) = modified_columns {
1146 let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
1149 new_row
1150 .schema
1151 .as_deref()
1152 .and_then(|schema| schema.first().map(|name| name.as_str()))
1153 });
1154 match pk_col_name {
1155 Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
1156 None => old.pk_value.as_ref() != new_row.columns.first(),
1158 }
1159 } else {
1160 old.pk_value.as_ref() != new_row.columns.first()
1161 }
1162 }
1163 _ => false,
1165 };
1166
1167 if pk_changed {
1168 if let Some((collection, pk_str)) = &old.pk_index_key {
1170 self.pk_index.remove(&(collection.clone(), pk_str.clone()));
1171 }
1172 if let EntityData::Row(row) = &new.data {
1174 if let Some(first_col) = row.columns.first() {
1175 let pk_str = format!("{:?}", first_col);
1176 self.pk_index
1177 .insert((new.kind.collection().to_string(), pk_str), new.id);
1178 }
1179 }
1180 }
1181
1182 let new_refs = new.cross_refs();
1184 if old.cross_refs.as_slice() != new_refs {
1185 self.cross_ref_forward.remove(&new.id);
1187 for cross_ref in &old.cross_refs {
1189 if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
1190 rev.retain(|(src, _)| *src != new.id);
1191 }
1192 }
1193 for cross_ref in new_refs {
1195 self.cross_ref_forward
1196 .entry(cross_ref.source)
1197 .or_default()
1198 .push((cross_ref.target, cross_ref.ref_type));
1199 self.cross_ref_reverse
1200 .entry(cross_ref.target)
1201 .or_default()
1202 .push((cross_ref.source, cross_ref.ref_type));
1203 }
1204 }
1205 }
1206
1207 pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1209 self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
1210 }
1211
1212 pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
1214 self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
1215 }
1216
1217 pub fn age_secs(&self) -> u64 {
1219 let now = current_unix_secs();
1220 now.saturating_sub(self.created_at)
1221 }
1222
1223 pub fn idle_secs(&self) -> u64 {
1225 let now = current_unix_secs();
1226 now.saturating_sub(self.last_write_at)
1227 }
1228
1229 pub fn bulk_insert(
1238 &mut self,
1239 entities: Vec<UnifiedEntity>,
1240 ) -> Result<Vec<EntityId>, SegmentError> {
1241 if !self.state.is_writable() {
1242 return Err(SegmentError::NotWritable);
1243 }
1244
1245 let n = entities.len();
1246
1247 let kind_key = if let Some(first) = entities.first() {
1249 first.kind.storage_type().to_string()
1250 } else {
1251 return Ok(Vec::new());
1252 };
1253
1254 let kind_set = self.kind_index.entry(kind_key).or_default();
1255 kind_set.reserve(n);
1256
1257 let now = current_unix_secs();
1258
1259 let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
1260
1261 let mut ids = Vec::with_capacity(n);
1262
1263 if self.flat_entities.is_empty() && self.entities.is_empty() {
1265 self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
1267 self.use_flat = true;
1268 }
1269
1270 let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
1280 let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
1281 let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
1285
1286 let mut batch_bytes: usize = 0;
1291
1292 if self.use_flat {
1293 self.flat_entities.reserve(n);
1294 for (i, mut entity) in entities.into_iter().enumerate() {
1295 entity.sequence_id = base_seq + i as u64;
1296 let id = entity.id;
1297 kind_set.insert(id);
1298 ids.push(id);
1299 batch_bytes += Self::estimate_entity_size(&entity);
1300 if let EntityData::Row(row) = &entity.data {
1301 if row.schema.is_some() && !row.columns.is_empty() {
1302 if columnar_zone_updates.is_empty() {
1303 columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
1304 columnar_schema = row.schema.clone();
1305 }
1306 for (ci, val) in row.columns.iter().enumerate() {
1307 if !matches!(val, Value::Null) {
1308 if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
1309 bucket.push(val.clone());
1310 }
1311 }
1312 }
1313 } else {
1314 for (col, val) in row.iter_fields() {
1315 if !matches!(val, Value::Null) {
1316 named_zone_updates.push((col.to_string(), val.clone()));
1317 }
1318 }
1319 }
1320 }
1321 let expected = self.base_entity_id + self.flat_entities.len() as u64;
1337 if id.raw() == expected {
1338 self.flat_entities.push(entity);
1339 } else {
1340 self.entities.insert(id, entity);
1341 }
1342 }
1343 } else {
1344 self.entities.reserve(n);
1346 let mut pairs = Vec::with_capacity(n);
1347 for (i, mut entity) in entities.into_iter().enumerate() {
1348 entity.sequence_id = base_seq + i as u64;
1349 let id = entity.id;
1350 kind_set.insert(id);
1351 ids.push(id);
1352 batch_bytes += Self::estimate_entity_size(&entity);
1353 if let EntityData::Row(row) = &entity.data {
1354 for (col, val) in row.iter_fields() {
1355 if !matches!(val, Value::Null) {
1356 named_zone_updates.push((col.to_string(), val.clone()));
1357 }
1358 }
1359 }
1360 pairs.push((id, entity));
1361 }
1362 self.entities.extend(pairs);
1363 }
1364
1365 let _ = kind_set;
1369 if !columnar_zone_updates.is_empty() {
1370 let schema = columnar_schema.as_ref();
1371 for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
1372 if values.is_empty() {
1373 continue;
1374 }
1375 let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
1376 continue;
1377 };
1378 let mut iter = values.into_iter();
1384 if let Some(first) = iter.next() {
1385 let zone = self
1386 .col_zones
1387 .entry(col_name.clone())
1388 .and_modify(|z| z.update(&first))
1389 .or_insert_with(|| ColZone::new(first));
1390 for v in iter {
1391 zone.update(&v);
1392 }
1393 }
1394 }
1395 }
1396 for (col, val) in named_zone_updates {
1397 self.col_zones
1398 .entry(col)
1399 .and_modify(|z| z.update(&val))
1400 .or_insert_with(|| ColZone::new(val));
1401 }
1402
1403 self.add_memory(batch_bytes);
1404 self.last_write_at = now;
1405
1406 if self.use_flat {
1408 self.published_flat_len
1409 .store(self.flat_entities.len(), Ordering::Release);
1410 }
1411
1412 Ok(ids)
1413 }
1414
1415 pub(crate) fn force_delete(&mut self, id: EntityId) -> bool {
1418 if self.use_flat {
1419 let raw = id.raw();
1420 if raw >= self.base_entity_id {
1421 let idx = (raw - self.base_entity_id) as usize;
1422 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
1423 self.tombstone_flat_entity(idx, id);
1424 return true;
1425 }
1426 }
1427 return self.remove_hashmap_entity(id);
1428 }
1429
1430 self.remove_hashmap_entity(id)
1431 }
1432
1433 pub(crate) fn force_update_with_metadata(
1436 &mut self,
1437 entity: &UnifiedEntity,
1438 modified_columns: &[String],
1439 metadata: Option<&Metadata>,
1440 ) -> Result<(), SegmentError> {
1441 self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
1442 }
1443}
1444
1445impl UnifiedSegment for GrowingSegment {
1446 fn id(&self) -> SegmentId {
1447 self.id
1448 }
1449
1450 fn state(&self) -> SegmentState {
1451 self.state
1452 }
1453
1454 fn collection(&self) -> &str {
1455 &self.collection
1456 }
1457
1458 fn stats(&self) -> SegmentStats {
1459 let mut stats = SegmentStats {
1460 entity_count: self.live_entity_count(),
1461 deleted_count: self.deleted.len(),
1462 memory_bytes: self.resident_bytes() as usize,
1463 ..Default::default()
1464 };
1465
1466 for entity in self.entities.values() {
1467 match &entity.kind {
1468 EntityKind::TableRow { .. } => stats.row_count += 1,
1469 EntityKind::GraphNode(_) => stats.node_count += 1,
1470 EntityKind::GraphEdge(_) => stats.edge_count += 1,
1471 EntityKind::Vector { .. } => stats.vector_count += 1,
1472 EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
1473 EntityKind::QueueMessage { .. } => stats.row_count += 1,
1474 }
1475 stats.cross_ref_count += entity.cross_refs().len();
1476 }
1477
1478 stats
1479 }
1480
1481 fn entity_count(&self) -> usize {
1482 self.live_entity_count()
1483 }
1484
1485 fn contains(&self, id: EntityId) -> bool {
1486 self.has_live_entity(id)
1487 }
1488
1489 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
1490 if self.deleted.contains(&id) {
1491 return None;
1492 }
1493 if self.use_flat {
1494 let raw = id.raw();
1495 if raw >= self.base_entity_id {
1496 let idx = (raw - self.base_entity_id) as usize;
1497 if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
1498 return Some(entity);
1499 }
1500 }
1501 self.entities.get(&id)
1507 } else {
1508 self.entities.get(&id)
1509 }
1510 }
1511
1512 fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
1513 if self.deleted.contains(&id) || !self.state.is_writable() {
1514 return None;
1515 }
1516 if self.use_flat {
1517 let raw = id.raw();
1518 if raw >= self.base_entity_id {
1519 let idx = (raw - self.base_entity_id) as usize;
1520 if self
1521 .flat_entities
1522 .get(idx)
1523 .map(|e| e.id == id)
1524 .unwrap_or(false)
1525 {
1526 return self.flat_entities.get_mut(idx);
1527 }
1528 }
1529 self.entities.get_mut(&id)
1532 } else {
1533 self.entities.get_mut(&id)
1534 }
1535 }
1536
1537 fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
1538 if !self.state.is_writable() {
1539 return Err(SegmentError::NotWritable);
1540 }
1541
1542 if self.entities.contains_key(&entity.id) {
1543 return Err(SegmentError::AlreadyExists(entity.id));
1544 }
1545
1546 entity.sequence_id = self.next_sequence();
1548
1549 let size = Self::estimate_entity_size(&entity);
1551 self.add_memory(size);
1552
1553 self.index_entity(&entity);
1555
1556 self.update_col_zones_from_entity(&entity);
1558
1559 let id = entity.id;
1561 self.entities.insert(id, entity);
1562
1563 self.last_write_at = current_unix_secs();
1565
1566 Ok(id)
1567 }
1568
1569 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
1570 if !self.state.is_writable() {
1571 return Err(SegmentError::NotWritable);
1572 }
1573
1574 self.apply_update_with_metadata(&entity, None)?;
1575 self.last_write_at = current_unix_secs();
1576
1577 Ok(())
1578 }
1579
1580 fn update_hot(
1581 &mut self,
1582 entity: UnifiedEntity,
1583 modified_columns: &[String],
1584 ) -> Result<(), SegmentError> {
1585 if !self.state.is_writable() {
1586 return Err(SegmentError::NotWritable);
1587 }
1588
1589 self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
1590 self.last_write_at = current_unix_secs();
1591 Ok(())
1592 }
1593
1594 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
1595 if !self.state.is_writable() {
1596 return Err(SegmentError::NotWritable);
1597 }
1598
1599 if self.use_flat {
1601 let raw = id.raw();
1602 if raw >= self.base_entity_id {
1603 let idx = (raw - self.base_entity_id) as usize;
1604 if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
1605 self.tombstone_flat_entity(idx, id);
1606 return Ok(true);
1607 }
1608 }
1609 return Ok(self.remove_hashmap_entity(id));
1610 }
1611
1612 Ok(self.remove_hashmap_entity(id))
1613 }
1614
1615 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1616 if !self.has_live_entity(id) {
1617 return None;
1618 }
1619 Some(self.metadata.get_all(id))
1620 }
1621
1622 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1623 if !self.state.is_writable() {
1624 return Err(SegmentError::NotWritable);
1625 }
1626
1627 if !self.has_live_entity(id) {
1628 return Err(SegmentError::NotFound(id));
1629 }
1630
1631 self.metadata.set_all(id, &metadata);
1632 Ok(())
1633 }
1634
1635 fn seal(&mut self) -> Result<(), SegmentError> {
1636 if self.state != SegmentState::Growing {
1637 return Err(SegmentError::InvalidState(self.state));
1638 }
1639
1640 self.state = SegmentState::Sealing;
1641
1642 self.rebuild_sealed_col_zones();
1647
1648 self.state = SegmentState::Sealed;
1649 Ok(())
1650 }
1651
1652 fn should_seal(&self, config: &SegmentConfig) -> bool {
1653 if self.entities.len() >= config.max_entities {
1655 return true;
1656 }
1657
1658 if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
1660 return true;
1661 }
1662
1663 if self.age_secs() >= config.max_age_secs {
1665 return true;
1666 }
1667
1668 false
1669 }
1670
1671 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
1672 let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
1673 Box::new(self.flat_entities.iter().chain(self.entities.values()))
1677 } else {
1678 Box::new(self.entities.values())
1679 };
1680 if self.deleted.is_empty() {
1681 base
1682 } else {
1683 Box::new(base.filter(|e| !self.deleted.contains(&e.id)))
1684 }
1685 }
1686
1687 fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
1688 let ids = self.kind_index.get(kind_filter).cloned();
1689 let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
1692 Box::new(self.flat_entities.iter())
1693 } else {
1694 Box::new(std::iter::empty())
1695 };
1696 Box::new(flat.chain(self.entities.values()).filter(move |e| {
1697 if self.deleted.contains(&e.id) {
1698 return false;
1699 }
1700 if let Some(ref ids) = ids {
1701 ids.contains(&e.id)
1702 } else {
1703 false
1704 }
1705 }))
1706 }
1707
1708 fn filter_metadata(
1709 &self,
1710 filters: &[(String, super::metadata::MetadataFilter)],
1711 ) -> Vec<EntityId> {
1712 let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
1715 Box::new(self.flat_entities.iter().map(|e| e.id))
1716 } else {
1717 Box::new(std::iter::empty())
1718 };
1719 flat_ids
1720 .chain(self.entities.keys().copied())
1721 .filter(|id| {
1722 if self.deleted.contains(id) {
1723 return false;
1724 }
1725 let metadata = self.metadata.get_all(*id);
1726 metadata.matches_all(filters)
1727 })
1728 .collect()
1729 }
1730}
1731
1732fn build_minmax_multi_intervals(
1733 entries: &[(CanonicalKey, Value)],
1734 max_intervals: usize,
1735) -> Vec<ColZone> {
1736 if entries.is_empty() {
1737 return Vec::new();
1738 }
1739 if entries.len() == 1 || max_intervals <= 1 {
1740 return vec![ColZone::with_bounds(
1741 entries[0].1.clone(),
1742 entries[entries.len() - 1].1.clone(),
1743 )];
1744 }
1745
1746 let mut split_points = if entries.len() <= max_intervals {
1747 (1..entries.len()).collect::<Vec<_>>()
1748 } else {
1749 let target_splits = max_intervals - 1;
1750 let mut selected = select_gap_split_points(entries, target_splits);
1751 if selected.len() < target_splits {
1752 for bucket in 1..max_intervals {
1753 let idx = bucket * entries.len() / max_intervals;
1754 if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
1755 continue;
1756 }
1757 selected.push(idx);
1758 if selected.len() >= target_splits {
1759 break;
1760 }
1761 }
1762 }
1763 selected.sort_unstable();
1764 selected.dedup();
1765 selected
1766 };
1767
1768 split_points.push(entries.len());
1769
1770 let mut out = Vec::with_capacity(split_points.len());
1771 let mut start = 0usize;
1772 for end in split_points {
1773 if end <= start {
1774 continue;
1775 }
1776 out.push(ColZone::with_bounds(
1777 entries[start].1.clone(),
1778 entries[end - 1].1.clone(),
1779 ));
1780 start = end;
1781 }
1782
1783 if out.is_empty() {
1784 out.push(ColZone::with_bounds(
1785 entries[0].1.clone(),
1786 entries[entries.len() - 1].1.clone(),
1787 ));
1788 }
1789
1790 out
1791}
1792
1793fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
1794 let mut gaps = Vec::new();
1795 for idx in 1..entries.len() {
1796 if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
1797 if score > 0.0 {
1798 gaps.push((score, idx));
1799 }
1800 }
1801 }
1802 gaps.sort_by(|left, right| {
1803 right
1804 .0
1805 .partial_cmp(&left.0)
1806 .unwrap_or(std::cmp::Ordering::Equal)
1807 .then_with(|| left.1.cmp(&right.1))
1808 });
1809 gaps.into_iter()
1810 .take(max_splits)
1811 .map(|(_, idx)| idx)
1812 .collect()
1813}
1814
1815fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
1816 if left.family() != right.family() {
1817 return None;
1818 }
1819 match (left, right) {
1820 (CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
1821 Some(r.saturating_sub(*l) as f64)
1822 }
1823 (CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
1824 Some(r.saturating_sub(*l) as f64)
1825 }
1826 (CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
1827 Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
1828 }
1829 _ => None,
1830 }
1831}
1832
1833#[cfg(test)]
1834mod tests {
1835 use super::*;
1836 use crate::storage::schema::Value;
1837 use crate::storage::unified::entity::RowData;
1838 use crate::storage::unified::MetadataValue;
1839
1840 #[test]
1841 fn test_growing_segment_basic() {
1842 let mut segment = GrowingSegment::new(1, "test");
1843
1844 let entity = UnifiedEntity::table_row(
1845 EntityId::new(1),
1846 "users",
1847 1,
1848 vec![Value::text("Alice".to_string())],
1849 );
1850
1851 let id = segment.insert(entity).unwrap();
1852 assert_eq!(id, EntityId::new(1));
1853 assert!(segment.contains(id));
1854
1855 let stats = segment.stats();
1856 assert_eq!(stats.entity_count, 1);
1857 assert_eq!(stats.row_count, 1);
1858 }
1859
1860 #[test]
1861 fn test_segment_metadata() {
1862 let mut segment = GrowingSegment::new(1, "test");
1863
1864 let entity = UnifiedEntity::table_row(
1865 EntityId::new(1),
1866 "users",
1867 1,
1868 vec![Value::text("Alice".to_string())],
1869 );
1870 segment.insert(entity).unwrap();
1871
1872 let mut meta = Metadata::new();
1873 meta.set("role", MetadataValue::String("admin".to_string()));
1874 meta.set("level", MetadataValue::Int(5));
1875
1876 segment.set_metadata(EntityId::new(1), meta).unwrap();
1877
1878 let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
1879 assert_eq!(
1880 retrieved.get("role"),
1881 Some(&MetadataValue::String("admin".to_string()))
1882 );
1883 }
1884
1885 #[test]
1886 fn test_segment_seal() {
1887 let mut segment = GrowingSegment::new(1, "test");
1888
1889 let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
1890 segment.insert(entity).unwrap();
1891
1892 assert!(segment.state().is_writable());
1894
1895 segment.seal().unwrap();
1897 assert_eq!(segment.state(), SegmentState::Sealed);
1898
1899 let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
1901 assert!(segment.insert(entity2).is_err());
1902 }
1903
1904 #[test]
1905 fn test_should_seal() {
1906 let mut segment = GrowingSegment::new(1, "test");
1907
1908 let config = SegmentConfig {
1909 max_entities: 2,
1910 ..Default::default()
1911 };
1912
1913 assert!(!segment.should_seal(&config));
1914
1915 segment
1916 .insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
1917 .unwrap();
1918 assert!(!segment.should_seal(&config));
1919
1920 segment
1921 .insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
1922 .unwrap();
1923 assert!(segment.should_seal(&config));
1924 }
1925
1926 #[test]
1927 fn test_cross_references() {
1928 let mut segment = GrowingSegment::new(1, "test");
1929
1930 let mut entity1 = UnifiedEntity::table_row(
1931 EntityId::new(1),
1932 "hosts",
1933 1,
1934 vec![Value::text("192.168.1.1".to_string())],
1935 );
1936 entity1.add_cross_ref(CrossRef::new(
1937 EntityId::new(1),
1938 EntityId::new(2),
1939 "nodes",
1940 RefType::RowToNode,
1941 ));
1942 segment.insert(entity1).unwrap();
1943
1944 let refs_from = segment.get_references_from(EntityId::new(1));
1945 assert_eq!(refs_from.len(), 1);
1946 assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
1947
1948 let refs_to = segment.get_references_to(EntityId::new(2));
1949 assert_eq!(refs_to.len(), 1);
1950 assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
1951 }
1952
1953 #[test]
1954 fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
1955 let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
1956 zone.update(&Value::Email("delta@example.com".to_string()));
1957
1958 let probe = Value::Email("alpha@example.com".to_string());
1959 assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
1960
1961 let in_range = Value::Email("charlie@example.com".to_string());
1962 assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
1963 }
1964
1965 #[test]
1966 fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
1967 let mut segment = GrowingSegment::new(1, "test");
1968
1969 for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
1970 let entity = UnifiedEntity::new(
1971 EntityId::new(row_id),
1972 EntityKind::TableRow {
1973 table: "users".into(),
1974 row_id,
1975 },
1976 EntityData::Row(RowData::with_names(
1977 vec![Value::Integer(age)],
1978 vec!["age".to_string()],
1979 )),
1980 );
1981 segment.insert(entity).unwrap();
1982 }
1983
1984 segment.seal().unwrap();
1985
1986 let miss = Value::Integer(500);
1987 assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
1988
1989 let hit = Value::Integer(1000);
1990 assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
1991 }
1992}