use std::collections::{BTreeMap, HashMap, HashSet};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use super::entity::{CrossRef, EntityData, EntityId, EntityKind, RefType, UnifiedEntity};
use super::metadata::{Metadata, MetadataStorage};
use super::segment_graph_index::SegmentGraphIndex;
use super::segment_unique_key::SegmentUniqueKey;
use crate::storage::query::value_compare::partial_compare_values;
use reddb_types::{value_to_canonical_key, CanonicalKey, Value};
pub type SegmentId = u64;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SegmentState {
Growing,
Sealing,
Sealed,
Flushed,
Archived,
}
impl SegmentState {
pub fn is_writable(&self) -> bool {
matches!(self, Self::Growing)
}
pub fn is_queryable(&self) -> bool {
!matches!(self, Self::Sealing)
}
pub fn is_immutable(&self) -> bool {
matches!(self, Self::Sealed | Self::Flushed | Self::Archived)
}
}
#[derive(Debug, Clone)]
pub struct SegmentConfig {
pub max_entities: usize,
pub max_bytes: usize,
pub max_age_secs: u64,
pub build_vector_index: bool,
pub build_graph_index: bool,
pub compression_level: u8,
}
impl Default for SegmentConfig {
fn default() -> Self {
Self {
max_entities: 100_000,
max_bytes: 256 * 1024 * 1024, max_age_secs: 3600, build_vector_index: true,
build_graph_index: true,
compression_level: 6,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct SegmentStats {
pub entity_count: usize,
pub deleted_count: usize,
pub memory_bytes: usize,
pub vector_count: usize,
pub node_count: usize,
pub edge_count: usize,
pub row_count: usize,
pub cross_ref_count: usize,
}
#[derive(Debug, Clone)]
pub enum SegmentError {
NotWritable,
NotFound(EntityId),
AlreadyExists(EntityId),
Full,
InvalidState(SegmentState),
Internal(String),
}
impl std::fmt::Display for SegmentError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotWritable => write!(f, "segment is not writable"),
Self::NotFound(id) => write!(f, "entity not found: {}", id),
Self::AlreadyExists(id) => write!(f, "entity already exists: {}", id),
Self::Full => write!(f, "segment is full"),
Self::InvalidState(state) => write!(f, "invalid operation for state: {:?}", state),
Self::Internal(msg) => write!(f, "internal error: {}", msg),
}
}
}
impl std::error::Error for SegmentError {}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
fn entity_id_from_probe_key(key: &[u8]) -> Option<EntityId> {
if let Ok(text) = std::str::from_utf8(key) {
if let Ok(raw) = text.parse::<u64>() {
return Some(EntityId::new(raw));
}
}
let bytes: [u8; 8] = key.try_into().ok()?;
Some(EntityId::new(u64::from_le_bytes(bytes)))
}
const SEALED_MULTI_ZONE_MAX_INTERVALS: usize = 4;
const TOMBSTONE_ENTRY_BYTES: u64 = 24;
#[derive(Debug, Clone)]
struct UpdateIndexSnapshot {
pk_column_name: Option<String>,
pk_value: Option<Value>,
pk_index_key: Option<(String, String)>,
cross_refs: Vec<CrossRef>,
}
impl UpdateIndexSnapshot {
fn from_entity(entity: &UnifiedEntity) -> Self {
let (pk_column_name, pk_value) = match &entity.data {
EntityData::Row(row) => (
row.schema
.as_deref()
.and_then(|schema| schema.first().cloned()),
row.columns.first().cloned(),
),
_ => (None, None),
};
let pk_index_key = pk_value
.as_ref()
.map(|value| (entity.kind.collection().to_string(), format!("{:?}", value)));
Self {
pk_column_name,
pk_value,
pk_index_key,
cross_refs: entity.cross_refs().to_vec(),
}
}
}
pub trait UnifiedSegment: Send + Sync {
fn id(&self) -> SegmentId;
fn state(&self) -> SegmentState;
fn collection(&self) -> &str;
fn stats(&self) -> SegmentStats;
fn entity_count(&self) -> usize;
fn contains(&self, id: EntityId) -> bool;
fn get(&self, id: EntityId) -> Option<&UnifiedEntity>;
fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity>;
fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError>;
fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError>;
fn update_hot(
&mut self,
entity: UnifiedEntity,
modified_columns: &[String],
) -> Result<(), SegmentError> {
let _ = modified_columns;
self.update(entity)
}
fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError>;
fn get_metadata(&self, id: EntityId) -> Option<Metadata>;
fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError>;
fn seal(&mut self) -> Result<(), SegmentError>;
fn should_seal(&self, config: &SegmentConfig) -> bool;
fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_>;
fn filter_metadata(
&self,
filters: &[(String, super::metadata::MetadataFilter)],
) -> Vec<EntityId>;
}
#[derive(Debug, Clone)]
pub struct ColZone {
pub min: Value,
pub max: Value,
min_key: Option<CanonicalKey>,
max_key: Option<CanonicalKey>,
}
impl ColZone {
fn payload_bytes(&self) -> usize {
super::memory_size::value_heap_bytes(&self.min)
.saturating_add(super::memory_size::value_heap_bytes(&self.max))
.saturating_add(
self.min_key
.as_ref()
.map_or(0, super::memory_size::canonical_heap_bytes),
)
.saturating_add(
self.max_key
.as_ref()
.map_or(0, super::memory_size::canonical_heap_bytes),
)
}
fn new(v: Value) -> Self {
Self {
min_key: value_to_canonical_key(&v),
max_key: value_to_canonical_key(&v),
min: v.clone(),
max: v,
}
}
fn with_bounds(min: Value, max: Value) -> Self {
Self {
min_key: value_to_canonical_key(&min),
max_key: value_to_canonical_key(&max),
min,
max,
}
}
fn update(&mut self, v: &Value) {
if compare_zone_values(v, None, &self.min, self.min_key.as_ref())
.map(|o| o == std::cmp::Ordering::Less)
.unwrap_or(false)
{
self.min = v.clone();
self.min_key = value_to_canonical_key(v);
}
if compare_zone_values(v, None, &self.max, self.max_key.as_ref())
.map(|o| o == std::cmp::Ordering::Greater)
.unwrap_or(false)
{
self.max = v.clone();
self.max_key = value_to_canonical_key(v);
}
}
}
#[derive(Debug, Clone, Default)]
pub struct MultiColZone {
pub intervals: Vec<ColZone>,
}
impl MultiColZone {
fn can_skip(&self, pred: &ZoneColPred<'_>) -> bool {
!self.intervals.is_empty() && self.intervals.iter().all(|zone| pred.can_skip(zone))
}
}
fn compare_zone_values(
left: &Value,
left_key: Option<&CanonicalKey>,
right: &Value,
right_key: Option<&CanonicalKey>,
) -> Option<std::cmp::Ordering> {
partial_compare_values(left, right).or_else(|| {
let left_key = left_key.cloned().or_else(|| value_to_canonical_key(left))?;
let right_key = right_key
.cloned()
.or_else(|| value_to_canonical_key(right))?;
(left_key.family() == right_key.family()).then(|| left_key.cmp(&right_key))
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ZoneColPredKind {
Eq,
Gt,
Gte,
Lt,
Lte,
}
#[derive(Debug, Clone)]
pub enum ZoneColPred<'a> {
Eq(&'a Value),
Gt(&'a Value),
Gte(&'a Value),
Lt(&'a Value),
Lte(&'a Value),
}
impl<'a> ZoneColPred<'a> {
pub fn can_skip(&self, zone: &ColZone) -> bool {
match self {
ZoneColPred::Eq(val) => {
compare_zone_values(val, None, &zone.min, zone.min_key.as_ref())
.map(|o| o == std::cmp::Ordering::Less)
.unwrap_or(false)
|| compare_zone_values(val, None, &zone.max, zone.max_key.as_ref())
.map(|o| o == std::cmp::Ordering::Greater)
.unwrap_or(false)
}
ZoneColPred::Gt(val) => {
compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
.map(|o| o != std::cmp::Ordering::Greater)
.unwrap_or(false)
}
ZoneColPred::Gte(val) => {
compare_zone_values(&zone.max, zone.max_key.as_ref(), val, None)
.map(|o| o == std::cmp::Ordering::Less)
.unwrap_or(false)
}
ZoneColPred::Lt(val) => {
compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
.map(|o| o != std::cmp::Ordering::Less)
.unwrap_or(false)
}
ZoneColPred::Lte(val) => {
compare_zone_values(&zone.min, zone.min_key.as_ref(), val, None)
.map(|o| o == std::cmp::Ordering::Greater)
.unwrap_or(false)
}
}
}
}
#[derive(Clone, Copy)]
pub(crate) enum GraphEntityKind {
Node,
Edge,
}
impl GraphEntityKind {
fn mask(self) -> u8 {
match self {
Self::Node => 1,
Self::Edge => 2,
}
}
pub(crate) fn matches(self, kind: &EntityKind) -> bool {
matches!(
(self, kind),
(Self::Node, EntityKind::GraphNode(_)) | (Self::Edge, EntityKind::GraphEdge(_))
)
}
}
pub(crate) const SCAN_BATCH_SIZE: usize = 256;
pub(crate) struct SegmentScanCursor {
flat_end: usize,
hash_end: usize,
deleted_end: usize,
position: usize,
}
impl SegmentScanCursor {
pub(crate) fn remaining(&self) -> usize {
self.flat_end + self.hash_end - self.position
}
pub(crate) fn finished(&self) -> bool {
self.remaining() == 0
}
}
pub(crate) struct GraphIndexCursor {
key_index: usize,
after: Option<EntityId>,
upper: EntityId,
}
impl GraphIndexCursor {
pub(crate) fn finished(&self, keys: &[EntityId]) -> bool {
self.key_index == keys.len()
}
}
pub struct GrowingSegment {
id: SegmentId,
collection: String,
state: SegmentState,
created_at: u64,
last_write_at: u64,
entities: HashMap<EntityId, UnifiedEntity>,
hash_scan_ids: Vec<EntityId>,
flat_entities: Vec<UnifiedEntity>,
base_entity_id: u64,
use_flat: bool,
deleted: HashMap<EntityId, usize>,
metadata: MetadataStorage,
pk_index: BTreeMap<(String, String), EntityId>,
unique_key_lookup: parking_lot::RwLock<Vec<SegmentUniqueKey>>,
unique_key_lookup_bytes: AtomicU64,
kind_index: HashMap<String, HashSet<EntityId>>,
graph_kind_mask: u8,
graph_read_index: parking_lot::RwLock<Option<SegmentGraphIndex>>,
graph_read_index_bytes: AtomicU64,
cross_ref_forward: HashMap<EntityId, Vec<(EntityId, RefType)>>,
cross_ref_reverse: HashMap<EntityId, Vec<(EntityId, RefType)>>,
col_zones: HashMap<String, ColZone>,
zone_payload_bytes: usize,
sealed_col_zones: HashMap<String, MultiColZone>,
sequence: AtomicU64,
memory_bytes: AtomicU64,
dead_entity_bytes: u64,
dead_resident_count: usize,
pub(crate) published_flat_len: AtomicUsize,
}
impl GrowingSegment {
pub(crate) fn may_contain_graph_kind(&self, kind: GraphEntityKind) -> bool {
self.graph_kind_mask & kind.mask() != 0
}
fn observe_graph_kind(&mut self, kind: &EntityKind) {
self.graph_kind_mask |= match kind {
EntityKind::GraphNode(_) => GraphEntityKind::Node.mask(),
EntityKind::GraphEdge(_) => GraphEntityKind::Edge.mask(),
_ => 0,
};
}
pub(crate) fn scan_cursor(&self) -> SegmentScanCursor {
SegmentScanCursor {
flat_end: self.flat_entities.len(),
hash_end: self.hash_scan_ids.len(),
deleted_end: self.deleted.len(),
position: 0,
}
}
pub(crate) fn scan_batch(
&self,
cursor: &mut SegmentScanCursor,
before_work: &mut impl FnMut() -> bool,
visit: &mut impl FnMut(EntityId),
) -> bool {
for _ in 0..SCAN_BATCH_SIZE {
if cursor.finished() {
break;
}
if !before_work() {
return false;
}
let id = if cursor.position < cursor.flat_end {
self.flat_entities[cursor.position].id
} else {
self.hash_scan_ids[cursor.position - cursor.flat_end]
};
cursor.position += 1;
let eligible = match self.deleted.get(&id) {
Some(&ordinal) => ordinal >= cursor.deleted_end,
None => self.get(id).is_some(),
};
if eligible {
visit(id);
}
}
true
}
#[inline]
pub fn for_each_fast<F>(&self, mut f: F) -> bool
where
F: FnMut(&UnifiedEntity) -> bool,
{
if self.use_flat {
if self.deleted.is_empty() {
for entity in &self.flat_entities {
if !f(entity) {
return false;
}
}
for entity in self.entities.values() {
if !f(entity) {
return false;
}
}
} else {
for entity in &self.flat_entities {
if self.deleted.contains_key(&entity.id) {
continue;
}
if !f(entity) {
return false;
}
}
for entity in self.entities.values() {
if self.deleted.contains_key(&entity.id) {
continue;
}
if !f(entity) {
return false;
}
}
}
} else {
if self.deleted.is_empty() {
for entity in self.entities.values() {
if !f(entity) {
return false;
}
}
} else {
for entity in self.entities.values() {
if self.deleted.contains_key(&entity.id) {
continue;
}
if !f(entity) {
return false;
}
}
}
}
true
}
pub fn new(id: SegmentId, collection: impl Into<String>) -> Self {
let now = current_unix_secs();
Self {
id,
collection: collection.into(),
state: SegmentState::Growing,
created_at: now,
last_write_at: now,
entities: HashMap::new(),
hash_scan_ids: Vec::new(),
flat_entities: Vec::new(),
base_entity_id: 0,
use_flat: false,
deleted: HashMap::new(),
metadata: MetadataStorage::new(),
pk_index: BTreeMap::new(),
unique_key_lookup: parking_lot::RwLock::new(Vec::new()),
unique_key_lookup_bytes: AtomicU64::new(0),
kind_index: HashMap::new(),
graph_read_index: parking_lot::RwLock::new(Some(SegmentGraphIndex::default())),
graph_read_index_bytes: AtomicU64::new(0),
graph_kind_mask: 0,
cross_ref_forward: HashMap::new(),
cross_ref_reverse: HashMap::new(),
col_zones: HashMap::new(),
zone_payload_bytes: 0,
sealed_col_zones: HashMap::new(),
sequence: AtomicU64::new(0),
memory_bytes: AtomicU64::new(0),
dead_entity_bytes: 0,
dead_resident_count: 0,
published_flat_len: AtomicUsize::new(0),
}
}
fn next_sequence(&self) -> u64 {
self.sequence.fetch_add(1, Ordering::SeqCst)
}
fn has_live_entity(&self, id: EntityId) -> bool {
if self.deleted.contains_key(&id) {
return false;
}
if self.use_flat {
let raw = id.raw();
if raw >= self.base_entity_id {
let idx = (raw - self.base_entity_id) as usize;
if self
.flat_entities
.get(idx)
.is_some_and(|entity| entity.id == id)
{
return true;
}
}
self.entities.contains_key(&id)
} else {
self.entities.contains_key(&id)
}
}
fn update_existing_entity_in_place(
&mut self,
entity: &UnifiedEntity,
) -> Result<UpdateIndexSnapshot, SegmentError> {
self.invalidate_unique_key_lookup();
if (self.graph_kind_mask != 0
|| matches!(
entity.kind,
EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
))
&& self
.get(entity.id)
.is_some_and(|old| !SegmentGraphIndex::same_keys(old, entity))
{
self.invalidate_graph_read_index();
}
self.observe_graph_kind(&entity.kind);
if self.use_flat {
let raw = entity.id.raw();
if raw >= self.base_entity_id {
let idx = (raw - self.base_entity_id) as usize;
if let Some(slot) = self
.flat_entities
.get_mut(idx)
.filter(|slot| slot.id == entity.id)
{
let snapshot = UpdateIndexSnapshot::from_entity(slot);
let old_bytes = Self::estimate_entity_size(slot);
slot.clone_from(entity);
self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
return Ok(snapshot);
}
}
let Some(slot) = self.entities.get_mut(&entity.id) else {
return Err(SegmentError::NotFound(entity.id));
};
let snapshot = UpdateIndexSnapshot::from_entity(slot);
let old_bytes = Self::estimate_entity_size(slot);
slot.clone_from(entity);
self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
Ok(snapshot)
} else {
let Some(slot) = self.entities.get_mut(&entity.id) else {
return Err(SegmentError::NotFound(entity.id));
};
let snapshot = UpdateIndexSnapshot::from_entity(slot);
let old_bytes = Self::estimate_entity_size(slot);
slot.clone_from(entity);
self.replace_entity_memory(old_bytes, Self::estimate_entity_size(entity));
Ok(snapshot)
}
}
fn apply_hot_update_with_metadata(
&mut self,
entity: &UnifiedEntity,
modified_columns: &[String],
metadata: Option<&Metadata>,
) -> Result<(), SegmentError> {
let old = self.update_existing_entity_in_place(entity)?;
self.reindex_for_update(&old, entity, Some(modified_columns));
self.update_col_zones_from_entity(entity);
if let Some(metadata) = metadata {
self.metadata.set_all(entity.id, metadata);
}
Ok(())
}
fn apply_update_with_metadata(
&mut self,
entity: &UnifiedEntity,
metadata: Option<&Metadata>,
) -> Result<(), SegmentError> {
let old = self.update_existing_entity_in_place(entity)?;
self.reindex_for_update(&old, entity, None);
self.update_col_zones_from_entity(entity);
if let Some(metadata) = metadata {
self.metadata.set_all(entity.id, metadata);
}
Ok(())
}
pub fn update_hot_batch_with_metadata<'a, I>(&mut self, items: I) -> Result<(), SegmentError>
where
I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
{
if !self.state.is_writable() {
return Err(SegmentError::NotWritable);
}
let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
items.into_iter().collect();
if items.is_empty() {
return Ok(());
}
for (entity, _, _) in &items {
if !self.has_live_entity(entity.id) {
return Err(SegmentError::NotFound(entity.id));
}
}
for (entity, modified_columns, metadata) in items {
self.apply_hot_update_with_metadata(entity, modified_columns, metadata)?;
}
self.last_write_at = current_unix_secs();
Ok(())
}
fn record_tombstone(&mut self, id: EntityId) -> bool {
let ordinal = self.deleted.len();
match self.deleted.entry(id) {
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert(ordinal);
true
}
std::collections::hash_map::Entry::Occupied(_) => false,
}
}
fn tombstone_flat_entity(&mut self, idx: usize, id: EntityId) -> bool {
self.invalidate_unique_key_lookup();
if !self.record_tombstone(id) {
return false;
}
if let Some(index) = self.graph_read_index.get_mut() {
index.remove(&self.flat_entities[idx]);
self.graph_read_index_bytes
.store(index.memory_bytes(), Ordering::Relaxed);
}
let size = Self::estimate_entity_size(&self.flat_entities[idx]) as u64;
self.dead_entity_bytes += size;
self.dead_resident_count += 1;
self.metadata.remove_all(id);
true
}
fn remove_hashmap_entity(&mut self, id: EntityId) -> bool {
self.invalidate_unique_key_lookup();
let Some(entity) = self.entities.remove(&id) else {
return false;
};
self.release_memory(Self::estimate_entity_size(&entity));
self.unindex_entity(&entity);
self.metadata.remove_all(id);
self.record_tombstone(id);
true
}
pub fn delete_batch(&mut self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
if !self.state.is_writable() {
return Err(SegmentError::NotWritable);
}
if ids.is_empty() {
return Ok(Vec::new());
}
let mut deleted_ids = Vec::with_capacity(ids.len());
if self.use_flat {
for &id in ids {
let raw = id.raw();
if raw < self.base_entity_id {
if self.remove_hashmap_entity(id) {
deleted_ids.push(id);
}
continue;
}
let idx = (raw - self.base_entity_id) as usize;
if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
if self.tombstone_flat_entity(idx, id) {
deleted_ids.push(id);
}
} else if self.remove_hashmap_entity(id) {
deleted_ids.push(id);
}
}
} else {
for &id in ids {
if self.remove_hashmap_entity(id) {
deleted_ids.push(id);
}
}
}
if !deleted_ids.is_empty() {
self.last_write_at = current_unix_secs();
}
Ok(deleted_ids)
}
fn add_memory(&self, bytes: usize) {
self.memory_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
}
fn replace_entity_memory(&self, previous: usize, current: usize) {
if current >= previous {
self.add_memory(current - previous);
} else {
self.release_memory(previous - current);
}
}
pub fn memory_bytes(&self) -> u64 {
self.memory_bytes.load(Ordering::Relaxed)
+ self.zone_payload_bytes as u64
+ self.unique_key_lookup_bytes.load(Ordering::Relaxed)
+ self.graph_read_index_bytes.load(Ordering::Relaxed)
+ (self.hash_scan_ids.capacity() * std::mem::size_of::<EntityId>()) as u64
}
fn release_memory(&self, bytes: usize) {
let mut current = self.memory_bytes.load(Ordering::Relaxed);
loop {
let next = current.saturating_sub(bytes as u64);
match self.memory_bytes.compare_exchange_weak(
current,
next,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => break,
Err(observed) => current = observed,
}
}
}
pub(crate) fn resident_entity_count(&self) -> usize {
self.flat_entities.len() + self.entities.len()
}
pub(crate) fn live_entity_count(&self) -> usize {
self.resident_entity_count()
.saturating_sub(self.dead_resident_count)
}
pub(crate) fn tombstone_count(&self) -> usize {
self.deleted.len()
}
pub(crate) fn resident_bytes(&self) -> u64 {
self.memory_bytes() + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
}
pub(crate) fn reclaimable_bytes(&self) -> u64 {
self.dead_entity_bytes + self.deleted.len() as u64 * TOMBSTONE_ENTRY_BYTES
}
pub(crate) fn live_entity_ids(&self) -> Vec<EntityId> {
let mut ids = Vec::with_capacity(self.live_entity_count());
self.for_each_fast(|entity| {
ids.push(entity.id);
true
});
ids
}
pub(crate) fn adopt_entity(&mut self, entity: UnifiedEntity, metadata: Option<Metadata>) {
let id = entity.id;
let next = entity.sequence_id.saturating_add(1);
if self.sequence.load(Ordering::Relaxed) < next {
self.sequence.store(next, Ordering::Relaxed);
}
self.add_memory(Self::estimate_entity_size(&entity));
self.index_entity(&entity);
self.update_col_zones_from_entity(&entity);
self.entities.insert(id, entity);
if let Some(metadata) = metadata {
if !metadata.is_empty() {
self.metadata.set_all(id, &metadata);
}
}
}
pub(crate) fn evict_entity(&mut self, id: EntityId) -> bool {
self.invalidate_unique_key_lookup();
let Some(entity) = self.entities.remove(&id) else {
return false;
};
self.release_memory(Self::estimate_entity_size(&entity));
self.unindex_entity(&entity);
for cross_ref in entity.cross_refs() {
if let Some(reverse) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
reverse.retain(|(source, _)| *source != id);
}
}
self.metadata.remove_all(id);
true
}
fn estimate_entity_size(entity: &UnifiedEntity) -> usize {
super::memory_size::entity_bytes(entity)
}
fn update_column_zone(&mut self, column: &str, value: &Value) {
if matches!(value, Value::Null) {
return;
}
match self.col_zones.entry(column.to_string()) {
std::collections::hash_map::Entry::Occupied(mut entry) => {
let before = entry.get().payload_bytes();
entry.get_mut().update(value);
self.zone_payload_bytes = self
.zone_payload_bytes
.saturating_sub(before)
.saturating_add(entry.get().payload_bytes());
}
std::collections::hash_map::Entry::Vacant(entry) => {
let zone = ColZone::new(value.clone());
self.zone_payload_bytes =
self.zone_payload_bytes.saturating_add(zone.payload_bytes());
entry.insert(zone);
}
}
}
fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
if let EntityData::Row(row) = &entity.data {
if let Some(named) = &row.named {
for (column, value) in named {
self.update_column_zone(column, value);
}
} else if let Some(schema) = &row.schema {
for (column, value) in schema.iter().zip(&row.columns) {
self.update_column_zone(column, value);
}
}
}
}
fn rebuild_sealed_col_zones(&mut self) {
let mut values_by_col: HashMap<String, Vec<(CanonicalKey, Value)>> = HashMap::new();
let mut family_by_col: HashMap<String, reddb_types::CanonicalKeyFamily> = HashMap::new();
let mut mixed_family_cols = HashSet::new();
let mut unsupported_cols = HashSet::new();
let mut observe_row = |row: &super::entity::RowData| {
for (col, value) in row.iter_fields() {
if matches!(value, Value::Null) {
continue;
}
let Some(key) = value_to_canonical_key(value) else {
unsupported_cols.insert(col.to_string());
continue;
};
match family_by_col.get(col).copied() {
Some(existing) if existing != key.family() => {
mixed_family_cols.insert(col.to_string());
}
None => {
family_by_col.insert(col.to_string(), key.family());
}
_ => {}
}
values_by_col
.entry(col.to_string())
.or_default()
.push((key, value.clone()));
}
};
if self.use_flat {
for entity in &self.flat_entities {
if self.deleted.contains_key(&entity.id) {
continue;
}
if let EntityData::Row(row) = &entity.data {
observe_row(row);
}
}
} else {
for entity in self.entities.values() {
if self.deleted.contains_key(&entity.id) {
continue;
}
if let EntityData::Row(row) = &entity.data {
observe_row(row);
}
}
}
let mut sealed_col_zones = HashMap::new();
for (col, mut entries) in values_by_col {
if mixed_family_cols.contains(&col)
|| unsupported_cols.contains(&col)
|| entries.is_empty()
{
continue;
}
entries.sort_unstable_by(|left, right| left.0.cmp(&right.0));
entries.dedup_by(|left, right| left.0 == right.0);
let intervals = build_minmax_multi_intervals(&entries, SEALED_MULTI_ZONE_MAX_INTERVALS);
if intervals.len() > 1 {
sealed_col_zones.insert(col, MultiColZone { intervals });
}
}
let previous = self
.sealed_col_zones
.values()
.flat_map(|zones| &zones.intervals)
.map(ColZone::payload_bytes)
.fold(0, usize::saturating_add);
let current = sealed_col_zones
.values()
.flat_map(|zones| &zones.intervals)
.map(ColZone::payload_bytes)
.fold(0, usize::saturating_add);
self.zone_payload_bytes = self
.zone_payload_bytes
.saturating_sub(previous)
.saturating_add(current);
self.sealed_col_zones = sealed_col_zones;
}
pub fn can_skip_zone_preds(&self, preds: &[(&str, ZoneColPred<'_>)]) -> bool {
if preds.is_empty() {
return false;
}
for (col, pred) in preds {
if let Some(zone) = self.sealed_col_zones.get(*col) {
if zone.can_skip(pred) {
return true;
}
continue;
}
if let Some(zone) = self.col_zones.get(*col) {
if pred.can_skip(zone) {
return true; }
}
}
false
}
fn index_entity(&mut self, entity: &UnifiedEntity) {
self.hash_scan_ids.push(entity.id);
self.observe_graph_kind(&entity.kind);
if let Some(index) = self.graph_read_index.get_mut() {
index.insert(entity);
self.graph_read_index_bytes
.store(index.memory_bytes(), Ordering::Relaxed);
}
let indexes = self.unique_key_lookup.get_mut();
for index in indexes.iter_mut() {
index.insert(entity);
}
self.unique_key_lookup_bytes.store(
indexes
.iter()
.map(|index| index.memory_bytes() as u64)
.sum(),
Ordering::Relaxed,
);
let kind_key = entity.kind.storage_type().to_string();
self.kind_index
.entry(kind_key)
.or_default()
.insert(entity.id);
if let EntityData::Row(row) = &entity.data {
if let Some(first_col) = row.columns.first() {
let pk_str = format!("{:?}", first_col);
self.pk_index
.insert((entity.kind.collection().to_string(), pk_str), entity.id);
}
}
for cross_ref in entity.cross_refs() {
self.cross_ref_forward
.entry(cross_ref.source)
.or_default()
.push((cross_ref.target, cross_ref.ref_type));
self.cross_ref_reverse
.entry(cross_ref.target)
.or_default()
.push((cross_ref.source, cross_ref.ref_type));
}
}
fn invalidate_graph_read_index(&mut self) {
*self.graph_read_index.get_mut() = None;
self.graph_read_index_bytes.store(0, Ordering::Relaxed);
}
#[cfg(test)]
pub(crate) fn visit_graph_index_candidates(
&self,
kind: GraphEntityKind,
keys: &[EntityId],
before_work: &mut impl FnMut() -> bool,
visit: &mut impl FnMut(&UnifiedEntity) -> bool,
) -> bool {
let Some(cached) = self.graph_read_index_for_query(before_work) else {
return false;
};
let index = cached.as_ref().expect("complete graph index");
for key in keys {
if !before_work() {
return false;
}
for id in index.candidates(kind, *key) {
if !before_work() {
return false;
}
if let Some(entity) = self.get(id) {
if !visit(entity) {
return false;
}
}
}
}
true
}
fn graph_read_index_for_query(
&self,
before_work: &mut impl FnMut() -> bool,
) -> Option<parking_lot::RwLockReadGuard<'_, Option<SegmentGraphIndex>>> {
let mut cached = self.graph_read_index.read();
if cached.is_none() {
drop(cached);
let mut writer = self.graph_read_index.write();
if writer.is_none() {
let mut index = SegmentGraphIndex::default();
if !self.for_each_fast(|entity| {
if !before_work() {
return false;
}
index.insert(entity);
true
}) {
return None;
}
self.graph_read_index_bytes
.store(index.memory_bytes(), Ordering::Relaxed);
*writer = Some(index);
}
cached = parking_lot::RwLockWriteGuard::downgrade(writer);
}
Some(cached)
}
pub(crate) fn graph_index_cursor(
&self,
kind: GraphEntityKind,
keys: &[EntityId],
before_work: &mut impl FnMut() -> bool,
) -> Option<GraphIndexCursor> {
let cached = self.graph_read_index_for_query(before_work)?;
let index = cached.as_ref().expect("complete graph index");
let mut upper = EntityId::new(0);
for key in keys {
if !before_work() {
return None;
}
if let Some(last) = index.candidates(kind, *key).next_back() {
upper = upper.max(last);
}
}
Some(GraphIndexCursor {
key_index: 0,
after: None,
upper,
})
}
pub(crate) fn graph_index_batch(
&self,
kind: GraphEntityKind,
keys: &[EntityId],
cursor: &mut GraphIndexCursor,
before_work: &mut impl FnMut() -> bool,
visit: &mut impl FnMut(EntityId),
) -> bool {
self.graph_index_batch_with_limit(kind, keys, cursor, before_work, visit, SCAN_BATCH_SIZE)
}
pub(crate) fn graph_index_batch_with_limit(
&self,
kind: GraphEntityKind,
keys: &[EntityId],
cursor: &mut GraphIndexCursor,
before_work: &mut impl FnMut() -> bool,
visit: &mut impl FnMut(EntityId),
limit: usize,
) -> bool {
assert!(
limit > 0 && limit <= SCAN_BATCH_SIZE,
"bounded graph index batch"
);
let Some(cached) = self.graph_read_index_for_query(before_work) else {
return false;
};
let index = cached.as_ref().expect("complete graph index");
let mut count = 0;
while !cursor.finished(keys) && count < limit {
let key = keys[cursor.key_index];
if !before_work() {
return false;
}
for id in index
.candidates_after(kind, key, cursor.after, cursor.upper)
.take(limit - count)
{
if !before_work() {
return false;
}
cursor.after = Some(id);
count += 1;
visit(id);
}
if count == limit && cursor.after != Some(cursor.upper) {
break;
}
cursor.key_index += 1;
cursor.after = None;
}
true
}
pub(crate) fn find_unique_key_conflict(
&self,
columns: &[Vec<String>],
key_index: usize,
signatures: &[String],
exclude_id: Option<EntityId>,
reserves_key: &impl Fn(&UnifiedEntity) -> bool,
) -> Option<EntityId> {
assert!(
key_index < columns.len(),
"constraint index belongs to current schema"
);
let cached = self.unique_key_lookup.read();
if cached.iter().map(|index| &index.columns).eq(columns.iter()) {
return self.find_unique_key_candidate(
&cached[key_index],
signatures,
exclude_id,
reserves_key,
);
}
drop(cached);
let mut cached = self.unique_key_lookup.write();
if !cached.iter().map(|index| &index.columns).eq(columns.iter()) {
let mut indexes: Vec<_> = columns
.iter()
.map(|columns| SegmentUniqueKey::new(columns))
.collect();
self.for_each_fast(|entity| {
for index in &mut indexes {
index.insert(entity);
}
true
});
self.unique_key_lookup_bytes.store(
indexes
.iter()
.map(|index| index.memory_bytes() as u64)
.sum(),
Ordering::Relaxed,
);
*cached = indexes;
}
self.find_unique_key_candidate(&cached[key_index], signatures, exclude_id, reserves_key)
}
fn find_unique_key_candidate(
&self,
index: &SegmentUniqueKey,
signatures: &[String],
exclude_id: Option<EntityId>,
reserves_key: &impl Fn(&UnifiedEntity) -> bool,
) -> Option<EntityId> {
index.candidates(signatures).iter().copied().find(|id| {
if Some(*id) == exclude_id {
return false;
}
let Some(
entity @ UnifiedEntity {
data: EntityData::Row(row),
..
},
) = self.get(*id)
else {
return false;
};
index
.columns
.iter()
.zip(signatures)
.all(|(column, signature)| {
row.get_field(column)
.is_some_and(|value| format!("{value:?}") == *signature)
})
&& reserves_key(entity)
})
}
fn invalidate_unique_key_lookup(&mut self) {
*self.unique_key_lookup.get_mut() = Vec::new();
self.unique_key_lookup_bytes.store(0, Ordering::Relaxed);
}
pub fn may_contain_exact_key(&self, key: &[u8]) -> bool {
if let Some(id) = entity_id_from_probe_key(key) {
return self.has_live_entity(id);
}
let Ok(pk_str) = std::str::from_utf8(key) else {
return false;
};
let pk_key = (self.collection.clone(), pk_str.to_string());
self.pk_index
.get(&pk_key)
.is_some_and(|id| self.has_live_entity(*id))
}
fn unindex_entity(&mut self, entity: &UnifiedEntity) {
if let Some(index) = self.graph_read_index.get_mut() {
index.remove(entity);
self.graph_read_index_bytes
.store(index.memory_bytes(), Ordering::Relaxed);
}
let kind_key = entity.kind.storage_type().to_string();
if let Some(set) = self.kind_index.get_mut(&kind_key) {
set.remove(&entity.id);
}
if let EntityData::Row(row) = &entity.data {
if let Some(first_col) = row.columns.first() {
let pk_str = format!("{:?}", first_col);
self.pk_index
.remove(&(entity.kind.collection().to_string(), pk_str));
}
}
self.cross_ref_forward.remove(&entity.id);
}
fn reindex_for_update(
&mut self,
old: &UpdateIndexSnapshot,
new: &UnifiedEntity,
modified_columns: Option<&[String]>,
) {
let pk_changed = match &new.data {
EntityData::Row(new_row) => {
if let Some(cols) = modified_columns {
let pk_col_name = old.pk_column_name.as_deref().or_else(|| {
new_row
.schema
.as_deref()
.and_then(|schema| schema.first().map(|name| name.as_str()))
});
match pk_col_name {
Some(pk_name) => cols.iter().any(|c| c.eq_ignore_ascii_case(pk_name)),
None => old.pk_value.as_ref() != new_row.columns.first(),
}
} else {
old.pk_value.as_ref() != new_row.columns.first()
}
}
_ => false,
};
if pk_changed {
if let Some((collection, pk_str)) = &old.pk_index_key {
self.pk_index.remove(&(collection.clone(), pk_str.clone()));
}
if let EntityData::Row(row) = &new.data {
if let Some(first_col) = row.columns.first() {
let pk_str = format!("{:?}", first_col);
self.pk_index
.insert((new.kind.collection().to_string(), pk_str), new.id);
}
}
}
let new_refs = new.cross_refs();
if old.cross_refs.as_slice() != new_refs {
self.cross_ref_forward.remove(&new.id);
for cross_ref in &old.cross_refs {
if let Some(rev) = self.cross_ref_reverse.get_mut(&cross_ref.target) {
rev.retain(|(src, _)| *src != new.id);
}
}
for cross_ref in new_refs {
self.cross_ref_forward
.entry(cross_ref.source)
.or_default()
.push((cross_ref.target, cross_ref.ref_type));
self.cross_ref_reverse
.entry(cross_ref.target)
.or_default()
.push((cross_ref.source, cross_ref.ref_type));
}
}
}
pub fn get_references_to(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
self.cross_ref_reverse.get(&id).cloned().unwrap_or_default()
}
pub fn get_references_from(&self, id: EntityId) -> Vec<(EntityId, RefType)> {
self.cross_ref_forward.get(&id).cloned().unwrap_or_default()
}
pub fn age_secs(&self) -> u64 {
let now = current_unix_secs();
now.saturating_sub(self.created_at)
}
pub fn idle_secs(&self) -> u64 {
let now = current_unix_secs();
now.saturating_sub(self.last_write_at)
}
pub fn bulk_insert(
&mut self,
entities: Vec<UnifiedEntity>,
) -> Result<Vec<EntityId>, SegmentError> {
if !self.state.is_writable() {
return Err(SegmentError::NotWritable);
}
let indexes = self.unique_key_lookup.get_mut();
for index in indexes.iter_mut() {
for entity in &entities {
index.insert(entity);
}
}
self.unique_key_lookup_bytes.store(
indexes
.iter()
.map(|index| index.memory_bytes() as u64)
.sum(),
Ordering::Relaxed,
);
let n = entities.len();
if entities.is_empty() {
return Ok(Vec::new());
}
for run in
entities.chunk_by(|left, right| left.kind.storage_type() == right.kind.storage_type())
{
let kind = run[0].kind.storage_type();
self.observe_graph_kind(&run[0].kind);
if matches!(
run[0].kind,
EntityKind::GraphNode(_) | EntityKind::GraphEdge(_)
) {
if let Some(index) = self.graph_read_index.get_mut() {
for entity in run {
index.insert(entity);
}
self.graph_read_index_bytes
.store(index.memory_bytes(), Ordering::Relaxed);
}
}
let kind_set = match self.kind_index.get_mut(kind) {
Some(ids) => ids,
None => self.kind_index.entry(kind.to_string()).or_default(),
};
kind_set.reserve(run.len());
kind_set.extend(run.iter().map(|entity| entity.id));
}
let now = current_unix_secs();
let base_seq = self.sequence.fetch_add(n as u64, Ordering::Relaxed);
let mut ids = Vec::with_capacity(n);
if self.flat_entities.is_empty() && self.entities.is_empty() {
self.base_entity_id = entities.first().map(|e| e.id.raw()).unwrap_or(0);
self.use_flat = true;
}
let mut columnar_zone_updates: Vec<Vec<Value>> = Vec::new();
let mut columnar_schema: Option<std::sync::Arc<Vec<String>>> = None;
let mut named_zone_updates: Vec<(String, Value)> = Vec::new();
let mut batch_bytes: usize = 0;
if self.use_flat {
self.flat_entities.reserve(n);
for (i, mut entity) in entities.into_iter().enumerate() {
entity.sequence_id = base_seq + i as u64;
let id = entity.id;
ids.push(id);
batch_bytes += Self::estimate_entity_size(&entity);
if let EntityData::Row(row) = &entity.data {
if row.schema.is_some() && !row.columns.is_empty() {
if columnar_zone_updates.is_empty() {
columnar_zone_updates = vec![Vec::with_capacity(n); row.columns.len()];
columnar_schema = row.schema.clone();
}
for (ci, val) in row.columns.iter().enumerate() {
if !matches!(val, Value::Null) {
if let Some(bucket) = columnar_zone_updates.get_mut(ci) {
bucket.push(val.clone());
}
}
}
} else {
for (col, val) in row.iter_fields() {
if !matches!(val, Value::Null) {
named_zone_updates.push((col.to_string(), val.clone()));
}
}
}
}
let expected = self.base_entity_id + self.flat_entities.len() as u64;
if id.raw() == expected {
self.flat_entities.push(entity);
} else {
self.hash_scan_ids.push(id);
self.entities.insert(id, entity);
}
}
} else {
self.entities.reserve(n);
self.hash_scan_ids.reserve(n);
let mut pairs = Vec::with_capacity(n);
for (i, mut entity) in entities.into_iter().enumerate() {
entity.sequence_id = base_seq + i as u64;
let id = entity.id;
ids.push(id);
batch_bytes += Self::estimate_entity_size(&entity);
if let EntityData::Row(row) = &entity.data {
for (col, val) in row.iter_fields() {
if !matches!(val, Value::Null) {
named_zone_updates.push((col.to_string(), val.clone()));
}
}
}
self.hash_scan_ids.push(id);
pairs.push((id, entity));
}
self.entities.extend(pairs);
}
if !columnar_zone_updates.is_empty() {
let schema = columnar_schema.as_ref();
for (ci, values) in columnar_zone_updates.into_iter().enumerate() {
if values.is_empty() {
continue;
}
let Some(col_name) = schema.and_then(|s| s.get(ci)) else {
continue;
};
let mut iter = values.into_iter();
if let Some(first) = iter.next() {
let previous = self
.col_zones
.get(col_name)
.map_or(0, ColZone::payload_bytes);
let zone = self
.col_zones
.entry(col_name.clone())
.and_modify(|z| z.update(&first))
.or_insert_with(|| ColZone::new(first));
for v in iter {
zone.update(&v);
}
self.zone_payload_bytes = self
.zone_payload_bytes
.saturating_sub(previous)
.saturating_add(zone.payload_bytes());
}
}
}
for (col, val) in named_zone_updates {
self.update_column_zone(&col, &val);
}
self.add_memory(batch_bytes);
self.last_write_at = now;
if self.use_flat {
self.published_flat_len
.store(self.flat_entities.len(), Ordering::Release);
}
Ok(ids)
}
pub(crate) fn force_delete(&mut self, id: EntityId) -> bool {
if self.use_flat {
let raw = id.raw();
if raw >= self.base_entity_id {
let idx = (raw - self.base_entity_id) as usize;
if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
self.tombstone_flat_entity(idx, id);
return true;
}
}
return self.remove_hashmap_entity(id);
}
self.remove_hashmap_entity(id)
}
pub(crate) fn force_update_with_metadata(
&mut self,
entity: &UnifiedEntity,
modified_columns: &[String],
metadata: Option<&Metadata>,
) -> Result<(), SegmentError> {
self.apply_hot_update_with_metadata(entity, modified_columns, metadata)
}
}
impl UnifiedSegment for GrowingSegment {
fn id(&self) -> SegmentId {
self.id
}
fn state(&self) -> SegmentState {
self.state
}
fn collection(&self) -> &str {
&self.collection
}
fn stats(&self) -> SegmentStats {
let mut stats = SegmentStats {
entity_count: self.live_entity_count(),
deleted_count: self.deleted.len(),
memory_bytes: self.resident_bytes() as usize,
..Default::default()
};
for entity in self.entities.values() {
match &entity.kind {
EntityKind::TableRow { .. } => stats.row_count += 1,
EntityKind::GraphNode(_) => stats.node_count += 1,
EntityKind::GraphEdge(_) => stats.edge_count += 1,
EntityKind::Vector { .. } => stats.vector_count += 1,
EntityKind::TimeSeriesPoint(_) => stats.row_count += 1,
EntityKind::QueueMessage { .. } => stats.row_count += 1,
}
stats.cross_ref_count += entity.cross_refs().len();
}
stats
}
fn entity_count(&self) -> usize {
self.live_entity_count()
}
fn contains(&self, id: EntityId) -> bool {
self.has_live_entity(id)
}
fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
if self.deleted.contains_key(&id) {
return None;
}
if self.use_flat {
let raw = id.raw();
if raw >= self.base_entity_id {
let idx = (raw - self.base_entity_id) as usize;
if let Some(entity) = self.flat_entities.get(idx).filter(|e| e.id == id) {
return Some(entity);
}
}
self.entities.get(&id)
} else {
self.entities.get(&id)
}
}
fn get_mut(&mut self, id: EntityId) -> Option<&mut UnifiedEntity> {
self.invalidate_unique_key_lookup();
self.invalidate_graph_read_index();
if self.deleted.contains_key(&id) || !self.state.is_writable() {
return None;
}
self.graph_kind_mask = GraphEntityKind::Node.mask() | GraphEntityKind::Edge.mask();
if self.use_flat {
let raw = id.raw();
if raw >= self.base_entity_id {
let idx = (raw - self.base_entity_id) as usize;
if self
.flat_entities
.get(idx)
.map(|e| e.id == id)
.unwrap_or(false)
{
return self.flat_entities.get_mut(idx);
}
}
self.entities.get_mut(&id)
} else {
self.entities.get_mut(&id)
}
}
fn insert(&mut self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
if !self.state.is_writable() {
return Err(SegmentError::NotWritable);
}
if self.entities.contains_key(&entity.id) {
return Err(SegmentError::AlreadyExists(entity.id));
}
entity.sequence_id = self.next_sequence();
let size = Self::estimate_entity_size(&entity);
self.add_memory(size);
self.index_entity(&entity);
self.update_col_zones_from_entity(&entity);
let id = entity.id;
self.entities.insert(id, entity);
self.last_write_at = current_unix_secs();
Ok(id)
}
fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
if !self.state.is_writable() {
return Err(SegmentError::NotWritable);
}
self.apply_update_with_metadata(&entity, None)?;
self.last_write_at = current_unix_secs();
Ok(())
}
fn update_hot(
&mut self,
entity: UnifiedEntity,
modified_columns: &[String],
) -> Result<(), SegmentError> {
if !self.state.is_writable() {
return Err(SegmentError::NotWritable);
}
self.apply_hot_update_with_metadata(&entity, modified_columns, None)?;
self.last_write_at = current_unix_secs();
Ok(())
}
fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
if !self.state.is_writable() {
return Err(SegmentError::NotWritable);
}
if self.use_flat {
let raw = id.raw();
if raw >= self.base_entity_id {
let idx = (raw - self.base_entity_id) as usize;
if idx < self.flat_entities.len() && self.flat_entities[idx].id == id {
self.tombstone_flat_entity(idx, id);
return Ok(true);
}
}
return Ok(self.remove_hashmap_entity(id));
}
Ok(self.remove_hashmap_entity(id))
}
fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
if !self.has_live_entity(id) {
return None;
}
Some(self.metadata.get_all(id))
}
fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
if !self.state.is_writable() {
return Err(SegmentError::NotWritable);
}
if !self.has_live_entity(id) {
return Err(SegmentError::NotFound(id));
}
self.metadata.set_all(id, &metadata);
Ok(())
}
fn seal(&mut self) -> Result<(), SegmentError> {
if self.state != SegmentState::Growing {
return Err(SegmentError::InvalidState(self.state));
}
self.state = SegmentState::Sealing;
self.rebuild_sealed_col_zones();
self.state = SegmentState::Sealed;
Ok(())
}
fn should_seal(&self, config: &SegmentConfig) -> bool {
if self.entities.len() >= config.max_entities {
return true;
}
if self.memory_bytes.load(Ordering::Relaxed) as usize >= config.max_bytes {
return true;
}
if self.age_secs() >= config.max_age_secs {
return true;
}
false
}
fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
let base: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
Box::new(self.flat_entities.iter().chain(self.entities.values()))
} else {
Box::new(self.entities.values())
};
if self.deleted.is_empty() {
base
} else {
Box::new(base.filter(|e| !self.deleted.contains_key(&e.id)))
}
}
fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
let ids = self.kind_index.get(kind_filter);
let flat: Box<dyn Iterator<Item = &UnifiedEntity>> = if self.use_flat {
Box::new(self.flat_entities.iter())
} else {
Box::new(std::iter::empty())
};
Box::new(flat.chain(self.entities.values()).filter(move |e| {
if self.deleted.contains_key(&e.id) {
return false;
}
if let Some(ids) = ids {
ids.contains(&e.id)
} else {
false
}
}))
}
fn filter_metadata(
&self,
filters: &[(String, super::metadata::MetadataFilter)],
) -> Vec<EntityId> {
let flat_ids: Box<dyn Iterator<Item = EntityId> + '_> = if self.use_flat {
Box::new(self.flat_entities.iter().map(|e| e.id))
} else {
Box::new(std::iter::empty())
};
flat_ids
.chain(self.entities.keys().copied())
.filter(|id| {
if self.deleted.contains_key(id) {
return false;
}
let metadata = self.metadata.get_all(*id);
metadata.matches_all(filters)
})
.collect()
}
}
fn build_minmax_multi_intervals(
entries: &[(CanonicalKey, Value)],
max_intervals: usize,
) -> Vec<ColZone> {
if entries.is_empty() {
return Vec::new();
}
if entries.len() == 1 || max_intervals <= 1 {
return vec![ColZone::with_bounds(
entries[0].1.clone(),
entries[entries.len() - 1].1.clone(),
)];
}
let mut split_points = if entries.len() <= max_intervals {
(1..entries.len()).collect::<Vec<_>>()
} else {
let target_splits = max_intervals - 1;
let mut selected = select_gap_split_points(entries, target_splits);
if selected.len() < target_splits {
for bucket in 1..max_intervals {
let idx = bucket * entries.len() / max_intervals;
if idx == 0 || idx >= entries.len() || selected.contains(&idx) {
continue;
}
selected.push(idx);
if selected.len() >= target_splits {
break;
}
}
}
selected.sort_unstable();
selected.dedup();
selected
};
split_points.push(entries.len());
let mut out = Vec::with_capacity(split_points.len());
let mut start = 0usize;
for end in split_points {
if end <= start {
continue;
}
out.push(ColZone::with_bounds(
entries[start].1.clone(),
entries[end - 1].1.clone(),
));
start = end;
}
if out.is_empty() {
out.push(ColZone::with_bounds(
entries[0].1.clone(),
entries[entries.len() - 1].1.clone(),
));
}
out
}
fn select_gap_split_points(entries: &[(CanonicalKey, Value)], max_splits: usize) -> Vec<usize> {
let mut gaps = Vec::new();
for idx in 1..entries.len() {
if let Some(score) = canonical_gap_score(&entries[idx - 1].0, &entries[idx].0) {
if score > 0.0 {
gaps.push((score, idx));
}
}
}
gaps.sort_by(|left, right| {
right
.0
.partial_cmp(&left.0)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| left.1.cmp(&right.1))
});
gaps.into_iter()
.take(max_splits)
.map(|(_, idx)| idx)
.collect()
}
fn canonical_gap_score(left: &CanonicalKey, right: &CanonicalKey) -> Option<f64> {
if left.family() != right.family() {
return None;
}
match (left, right) {
(CanonicalKey::Signed(_, l), CanonicalKey::Signed(_, r)) => {
Some(r.saturating_sub(*l) as f64)
}
(CanonicalKey::Unsigned(_, l), CanonicalKey::Unsigned(_, r)) => {
Some(r.saturating_sub(*l) as f64)
}
(CanonicalKey::Float(l), CanonicalKey::Float(r)) => {
Some((f64::from_bits(*r) - f64::from_bits(*l)).abs())
}
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::unified::entity::RowData;
use crate::storage::unified::MetadataValue;
use reddb_types::Value;
fn pruning_node(id: u64) -> UnifiedEntity {
UnifiedEntity::graph_node(EntityId::new(id), "mixed", "person", Default::default())
}
fn pruning_edge(id: u64) -> UnifiedEntity {
UnifiedEntity::graph_edge(
EntityId::new(id),
"mixed",
"a",
"b",
1.0,
Default::default(),
)
}
#[test]
fn graph_presence_covers_bulk_gaps_inserts_and_adoption() {
let mut segment = GrowingSegment::new(1, "mixed");
assert!(!segment.may_contain_graph_kind(GraphEntityKind::Node));
assert!(!segment.may_contain_graph_kind(GraphEntityKind::Edge));
segment
.bulk_insert(vec![
UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]),
pruning_node(2),
pruning_edge(4),
])
.expect("mixed bulk with gap");
assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
segment.seal().expect("seal");
assert!(segment.force_delete(EntityId::new(2)));
assert!(segment.force_delete(EntityId::new(4)));
assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
let mut rebuilt = GrowingSegment::new(2, "mixed");
for item in segment.iter() {
rebuilt.adopt_entity(item.clone(), None);
}
assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
assert!(!rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
rebuilt.adopt_entity(pruning_node(5), None);
assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Node));
rebuilt.insert(pruning_edge(6)).expect("single edge");
assert!(rebuilt.may_contain_graph_kind(GraphEntityKind::Edge));
}
#[test]
fn graph_presence_covers_replacement_and_unrestricted_mutation() {
for flat in [false, true] {
for operation in 0..4 {
let mut segment = GrowingSegment::new(1, "mixed");
let vector = UnifiedEntity::vector(EntityId::new(1), "mixed", vec![]);
if flat {
segment.bulk_insert(vec![vector]).expect("flat vector");
} else {
segment.insert(vector).expect("hashmap vector");
}
match operation {
0 => segment.update(pruning_node(1)).expect("replace with node"),
1 => segment
.update_hot(pruning_node(1), &[])
.expect("hot replacement"),
2 => {
segment.seal().expect("seal vector");
segment
.force_update_with_metadata(&pruning_node(1), &[], None)
.expect("sealed replacement");
}
_ => {
*segment.get_mut(EntityId::new(1)).expect("mutable vector") =
pruning_node(1)
}
}
assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
assert!(matches!(
segment.get(EntityId::new(1)).expect("node").kind,
EntityKind::GraphNode(_)
));
if operation == 3 {
assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
}
segment
.force_update_with_metadata(&pruning_edge(1), &[], None)
.expect("replace with edge");
assert!(segment.may_contain_graph_kind(GraphEntityKind::Edge));
assert!(segment.may_contain_graph_kind(GraphEntityKind::Node));
}
}
}
fn primary_key_row(id: u64, key: Value) -> UnifiedEntity {
UnifiedEntity::new(
EntityId::new(id),
EntityKind::TableRow {
table: "keys".into(),
row_id: id,
},
EntityData::Row(RowData::with_names(
vec![Value::text("payload"), key],
vec!["payload".into(), "primary".into()],
)),
)
}
fn primary_key_probe(
segment: &GrowingSegment,
key: Value,
exclude: Option<EntityId>,
) -> Option<EntityId> {
segment.find_unique_key_conflict(
&[vec!["primary".into()]],
0,
&[format!("{key:?}")],
exclude,
&|_| true,
)
}
#[test]
fn unique_key_lookup_follows_segment_mutations_and_sealing() {
let mut segment = GrowingSegment::new(1, "keys");
segment
.bulk_insert(vec![primary_key_row(1, Value::Integer(10))])
.expect("bulk row");
let payload_bytes = segment.memory_bytes();
assert_eq!(
primary_key_probe(&segment, Value::Integer(10), None),
Some(EntityId::new(1))
);
assert!(segment.memory_bytes() > payload_bytes);
segment
.insert(primary_key_row(3, Value::Integer(30)))
.expect("hashmap fallback row");
segment
.bulk_insert(vec![primary_key_row(4, Value::Integer(40))])
.expect("bulk gap row");
assert_eq!(
primary_key_probe(&segment, Value::Integer(40), None),
Some(EntityId::new(4))
);
assert_eq!(
primary_key_probe(&segment, Value::Integer(30), None),
Some(EntityId::new(3))
);
segment
.update_hot(primary_key_row(1, Value::Integer(11)), &["primary".into()])
.expect("replace flat key");
assert_eq!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed), 0);
assert_eq!(primary_key_probe(&segment, Value::Integer(10), None), None);
assert_eq!(
primary_key_probe(&segment, Value::Integer(11), None),
Some(EntityId::new(1))
);
segment
.delete_batch(&[EntityId::new(1), EntityId::new(3)])
.expect("delete flat and map");
assert_eq!(primary_key_probe(&segment, Value::Integer(11), None), None);
assert_eq!(primary_key_probe(&segment, Value::Integer(30), None), None);
segment.seal().expect("seal");
assert_eq!(
primary_key_probe(&segment, Value::Integer(40), None),
Some(EntityId::new(4))
);
segment
.force_update_with_metadata(
&primary_key_row(4, Value::Integer(41)),
&["primary".into()],
None,
)
.expect("sealed update");
assert_eq!(primary_key_probe(&segment, Value::Integer(40), None), None);
assert_eq!(
primary_key_probe(&segment, Value::Integer(41), None),
Some(EntityId::new(4))
);
assert!(segment.force_delete(EntityId::new(4)));
assert_eq!(primary_key_probe(&segment, Value::Integer(41), None), None);
}
#[test]
fn unique_lookup_retains_all_current_definitions_and_rechecks_candidates() {
let mut segment = GrowingSegment::new(1, "keys");
segment
.bulk_insert(vec![
primary_key_row(1, Value::Integer(10)),
primary_key_row(2, Value::Integer(20)),
])
.expect("initial rows");
let definitions = vec![
vec!["primary".into()],
vec!["payload".into()],
vec!["primary".into(), "payload".into()],
];
let payload = format!("{:?}", Value::text("payload"));
assert_eq!(
segment.find_unique_key_conflict(
&definitions,
1,
&[payload.clone()],
None,
&|entity| entity.id != EntityId::new(1)
),
Some(EntityId::new(2))
);
assert_eq!(segment.unique_key_lookup.read().len(), 3);
segment
.insert(primary_key_row(3, Value::Integer(30)))
.expect("append");
assert_eq!(
segment.find_unique_key_conflict(
&definitions,
2,
&[format!("{:?}", Value::Integer(30)), payload.clone()],
None,
&|_| true
),
Some(EntityId::new(3))
);
let bytes = segment.unique_key_lookup_bytes.load(Ordering::Relaxed);
assert_eq!(
segment.find_unique_key_conflict(
&definitions,
0,
&[format!("{:?}", Value::Integer(20))],
None,
&|_| true
),
Some(EntityId::new(2))
);
assert_eq!(segment.unique_key_lookup.read().len(), 3);
assert_eq!(
segment.unique_key_lookup_bytes.load(Ordering::Relaxed),
bytes
);
let reduced = vec![vec!["payload".into()]];
assert_eq!(
segment.find_unique_key_conflict(&reduced, 0, &[payload], None, &|_| true),
Some(EntityId::new(1))
);
assert_eq!(segment.unique_key_lookup.read().len(), 1);
assert!(segment.unique_key_lookup_bytes.load(Ordering::Relaxed) < bytes);
}
#[test]
fn unique_key_lookup_preserves_signature_semantics_and_exclusion() {
let mut segment = GrowingSegment::new(1, "keys");
let values = [
Value::Float(f64::NAN),
Value::Float(-0.0),
Value::Float(0.0),
Value::Integer(1),
Value::UnsignedInteger(1),
Value::text("1"),
];
for (offset, value) in values.iter().enumerate() {
let id = offset as u64 + 1;
segment
.insert(primary_key_row(id, value.clone()))
.expect("distinct physical row");
assert_eq!(
primary_key_probe(&segment, value.clone(), None),
Some(EntityId::new(id))
);
assert_eq!(
primary_key_probe(&segment, value.clone(), Some(EntityId::new(id))),
None
);
}
segment
.insert(primary_key_row(20, Value::Integer(1)))
.expect("duplicate physical key");
assert_eq!(
primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(4))),
Some(EntityId::new(20))
);
segment
.entities
.insert(EntityId::new(4), primary_key_row(4, Value::Integer(99)));
assert_eq!(
primary_key_probe(&segment, Value::Integer(1), Some(EntityId::new(20))),
None
);
let entity = segment.get_mut(EntityId::new(4)).expect("mutable row");
*entity = primary_key_row(4, Value::Integer(100));
assert_eq!(
primary_key_probe(&segment, Value::Integer(100), None),
Some(EntityId::new(4))
);
}
#[test]
fn test_growing_segment_basic() {
let mut segment = GrowingSegment::new(1, "test");
let entity = UnifiedEntity::table_row(
EntityId::new(1),
"users",
1,
vec![Value::text("Alice".to_string())],
);
let id = segment.insert(entity).unwrap();
assert_eq!(id, EntityId::new(1));
assert!(segment.contains(id));
let stats = segment.stats();
assert_eq!(stats.entity_count, 1);
assert_eq!(stats.row_count, 1);
}
#[test]
fn test_segment_metadata() {
let mut segment = GrowingSegment::new(1, "test");
let entity = UnifiedEntity::table_row(
EntityId::new(1),
"users",
1,
vec![Value::text("Alice".to_string())],
);
segment.insert(entity).unwrap();
let mut meta = Metadata::new();
meta.set("role", MetadataValue::String("admin".to_string()));
meta.set("level", MetadataValue::Int(5));
segment.set_metadata(EntityId::new(1), meta).unwrap();
let retrieved = segment.get_metadata(EntityId::new(1)).unwrap();
assert_eq!(
retrieved.get("role"),
Some(&MetadataValue::String("admin".to_string()))
);
}
#[test]
fn test_segment_seal() {
let mut segment = GrowingSegment::new(1, "test");
let entity = UnifiedEntity::vector(EntityId::new(1), "embeddings", vec![0.1, 0.2, 0.3]);
segment.insert(entity).unwrap();
assert!(segment.state().is_writable());
segment.seal().unwrap();
assert_eq!(segment.state(), SegmentState::Sealed);
let entity2 = UnifiedEntity::vector(EntityId::new(2), "embeddings", vec![0.4, 0.5, 0.6]);
assert!(segment.insert(entity2).is_err());
}
#[test]
fn test_should_seal() {
let mut segment = GrowingSegment::new(1, "test");
let config = SegmentConfig {
max_entities: 2,
..Default::default()
};
assert!(!segment.should_seal(&config));
segment
.insert(UnifiedEntity::vector(EntityId::new(1), "v", vec![0.1]))
.unwrap();
assert!(!segment.should_seal(&config));
segment
.insert(UnifiedEntity::vector(EntityId::new(2), "v", vec![0.2]))
.unwrap();
assert!(segment.should_seal(&config));
}
#[test]
fn test_cross_references() {
let mut segment = GrowingSegment::new(1, "test");
let mut entity1 = UnifiedEntity::table_row(
EntityId::new(1),
"hosts",
1,
vec![Value::text("192.168.1.1".to_string())],
);
entity1.add_cross_ref(CrossRef::new(
EntityId::new(1),
EntityId::new(2),
"nodes",
RefType::RowToNode,
));
segment.insert(entity1).unwrap();
let refs_from = segment.get_references_from(EntityId::new(1));
assert_eq!(refs_from.len(), 1);
assert_eq!(refs_from[0], (EntityId::new(2), RefType::RowToNode));
let refs_to = segment.get_references_to(EntityId::new(2));
assert_eq!(refs_to.len(), 1);
assert_eq!(refs_to[0], (EntityId::new(1), RefType::RowToNode));
}
#[test]
fn test_zone_predicate_uses_canonical_fallback_for_email_values() {
let mut zone = ColZone::new(Value::Email("bravo@example.com".to_string()));
zone.update(&Value::Email("delta@example.com".to_string()));
let probe = Value::Email("alpha@example.com".to_string());
assert!(ZoneColPred::Eq(&probe).can_skip(&zone));
let in_range = Value::Email("charlie@example.com".to_string());
assert!(!ZoneColPred::Eq(&in_range).can_skip(&zone));
}
#[test]
fn test_sealed_multi_zone_prunes_numeric_gap_outlier() {
let mut segment = GrowingSegment::new(1, "test");
for (row_id, age) in [(1_u64, 1_i64), (2, 2), (3, 3), (4, 1000)] {
let entity = UnifiedEntity::new(
EntityId::new(row_id),
EntityKind::TableRow {
table: "users".into(),
row_id,
},
EntityData::Row(RowData::with_names(
vec![Value::Integer(age)],
vec!["age".to_string()],
)),
);
segment.insert(entity).unwrap();
}
segment.seal().unwrap();
let miss = Value::Integer(500);
assert!(segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&miss))]));
let hit = Value::Integer(1000);
assert!(!segment.can_skip_zone_preds(&[("age", ZoneColPred::Eq(&hit))]));
}
}