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 crate::storage::primitives::bloom::BloomFilter;
use crate::storage::query::value_compare::partial_compare_values;
use crate::storage::schema::{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()
}
const SEALED_MULTI_ZONE_MAX_INTERVALS: usize = 4;
const TOMBSTONE_ENTRY_BYTES: u64 = 16;
#[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 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)
}
}
}
}
pub struct GrowingSegment {
id: SegmentId,
collection: String,
state: SegmentState,
created_at: u64,
last_write_at: u64,
entities: HashMap<EntityId, UnifiedEntity>,
flat_entities: Vec<UnifiedEntity>,
base_entity_id: u64,
use_flat: bool,
deleted: HashSet<EntityId>,
metadata: MetadataStorage,
pk_index: BTreeMap<(String, String), EntityId>,
kind_index: HashMap<String, HashSet<EntityId>>,
cross_ref_forward: HashMap<EntityId, Vec<(EntityId, RefType)>>,
cross_ref_reverse: HashMap<EntityId, Vec<(EntityId, RefType)>>,
bloom: BloomFilter,
col_zones: HashMap<String, ColZone>,
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 {
#[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(&entity.id) {
continue;
}
if !f(entity) {
return false;
}
}
for entity in self.entities.values() {
if self.deleted.contains(&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(&entity.id) {
continue;
}
if !f(entity) {
return false;
}
}
}
}
true
}
pub fn new(id: SegmentId, collection: impl Into<String>) -> Self {
Self::with_bloom_capacity(id, collection, 100_000)
}
pub(crate) fn with_bloom_capacity(
id: SegmentId,
collection: impl Into<String>,
expected_entities: usize,
) -> Self {
let now = current_unix_secs();
Self {
id,
collection: collection.into(),
state: SegmentState::Growing,
created_at: now,
last_write_at: now,
entities: HashMap::new(),
flat_entities: Vec::new(),
base_entity_id: 0,
use_flat: false,
deleted: HashSet::new(),
metadata: MetadataStorage::new(),
pk_index: BTreeMap::new(),
kind_index: HashMap::new(),
cross_ref_forward: HashMap::new(),
cross_ref_reverse: HashMap::new(),
bloom: BloomFilter::with_capacity(expected_entities.max(1), 0.01),
col_zones: HashMap::new(),
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(&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> {
if self.use_flat {
let raw = entity.id.raw();
if raw < self.base_entity_id {
return Err(SegmentError::NotFound(entity.id));
}
let idx = (raw - self.base_entity_id) as usize;
let Some(slot) = self.flat_entities.get_mut(idx) else {
return Err(SegmentError::NotFound(entity.id));
};
if slot.id != entity.id {
return Err(SegmentError::NotFound(entity.id));
}
let snapshot = UpdateIndexSnapshot::from_entity(slot);
slot.clone_from(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);
slot.clone_from(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 tombstone_flat_entity(&mut self, idx: usize, id: EntityId) -> bool {
if !self.deleted.insert(id) {
return false;
}
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 {
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.deleted.insert(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);
}
pub fn memory_bytes(&self) -> u64 {
self.memory_bytes.load(Ordering::Relaxed)
}
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.load(Ordering::Relaxed)
+ 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 {
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 {
let mut size = std::mem::size_of::<UnifiedEntity>();
size += match &entity.data {
EntityData::Row(row) => row.columns.len() * 64, EntityData::Node(node) => node.properties.len() * 128,
EntityData::Edge(edge) => edge.properties.len() * 128,
EntityData::Vector(vec) => {
vec.dense.len() * 4 + vec.sparse.as_ref().map_or(0, |s| s.indices.len() * 8)
}
EntityData::TimeSeries(_) => 64,
EntityData::QueueMessage(_) => 128,
};
for emb in entity.embeddings() {
size += emb.vector.len() * 4 + emb.name.len() + emb.model.len();
}
size += std::mem::size_of_val(entity.cross_refs());
size
}
fn update_col_zones_from_entity(&mut self, entity: &UnifiedEntity) {
if let EntityData::Row(row) = &entity.data {
if let Some(named) = &row.named {
for (col, val) in named {
if matches!(val, Value::Null) {
continue;
}
self.col_zones
.entry(col.clone())
.and_modify(|z| z.update(val))
.or_insert_with(|| ColZone::new(val.clone()));
}
} else if let Some(schema) = &row.schema {
for (col, val) in schema.iter().zip(row.columns.iter()) {
if matches!(val, Value::Null) {
continue;
}
self.col_zones
.entry(col.clone())
.and_modify(|z| z.update(val))
.or_insert_with(|| ColZone::new(val.clone()));
}
}
}
}
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, crate::storage::schema::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(&entity.id) {
continue;
}
if let EntityData::Row(row) = &entity.data {
observe_row(row);
}
}
} else {
for entity in self.entities.values() {
if self.deleted.contains(&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 });
}
}
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) {
let kind_key = entity.kind.storage_type().to_string();
self.kind_index
.entry(kind_key)
.or_default()
.insert(entity.id);
let id_bytes = entity.id.raw().to_le_bytes();
self.bloom.insert(&id_bytes);
if let EntityData::Row(row) = &entity.data {
if let Some(first_col) = row.columns.first() {
let pk_str = format!("{:?}", first_col);
self.bloom.insert(pk_str.as_bytes());
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));
}
}
pub fn bloom_might_contain_key(&self, key: &[u8]) -> bool {
self.bloom.contains(key)
}
pub fn bloom_stats(&self) -> (f64, u32) {
(self.bloom.fill_ratio(), self.bloom.count_set_bits())
}
fn unindex_entity(&mut self, entity: &UnifiedEntity) {
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.bloom.insert(pk_str.as_bytes());
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 n = entities.len();
let kind_key = if let Some(first) = entities.first() {
first.kind.storage_type().to_string()
} else {
return Ok(Vec::new());
};
let kind_set = self.kind_index.entry(kind_key).or_default();
kind_set.reserve(n);
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;
kind_set.insert(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.entities.insert(id, entity);
}
}
} else {
self.entities.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;
kind_set.insert(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()));
}
}
}
pairs.push((id, entity));
}
self.entities.extend(pairs);
}
let _ = kind_set;
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 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);
}
}
}
}
for (col, val) in named_zone_updates {
self.col_zones
.entry(col)
.and_modify(|z| z.update(&val))
.or_insert_with(|| ColZone::new(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(&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> {
if self.deleted.contains(&id) || !self.state.is_writable() {
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 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(&e.id)))
}
}
fn iter_kind(&self, kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
let ids = self.kind_index.get(kind_filter).cloned();
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(&e.id) {
return false;
}
if let Some(ref 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(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::schema::Value;
use crate::storage::unified::entity::RowData;
use crate::storage::unified::MetadataValue;
#[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))]));
}
}