#![allow(unused_variables)]
#![allow(unused_mut)]
use super::art_manager::ArtIndexManager;
use super::bloom_filter::TableBloomFilters;
use super::columnar::{BatchPresence, BatchStats, ColumnBatch, ColumnarStore, BATCH_SIZE};
use super::content_addr::ContentAddressedStore;
use super::dictionary::DictionaryManager;
use super::filter_consolidation_worker::{ConsolidationConfig, FilterConsolidationWorker};
use super::filter_index_delta::{FilterIndexConfig, FilterIndexDeltaTracker};
use super::mv_scheduler::CpuMonitor;
use super::parallel_filter::{ParallelFilterConfig, ParallelFilterEngine};
use super::predicate_pushdown::{AnalyzedPredicate, PredicatePushdownManager, PushdownConfig};
use super::simd_filter::{FilterOp, FilterPredicate};
use super::speculative_filter::{SpeculativeConfig, SpeculativeFilterManager};
use super::typed_batch::TypedValues;
use super::typed_kernels::{self as tk, CompiledPredicate};
use super::wal::{WalOperation, WalSyncMode, WriteAheadLog};
use super::zone_map::TableZoneMap;
use super::{
BranchId, BranchManager, BranchMetadata, BranchOptions, BranchTransaction, Catalog, DatabaseStats, Key,
SnapshotManager, Transaction, VectorIndexManager,
};
use crate::crypto::{self, KeyManager};
use crate::ColumnStorageMode;
use crate::{Config, Error, Result, Tuple, Value};
use parking_lot::RwLock;
use rocksdb::{BlockBasedOptions, Cache, IteratorMode, Options, ReadOptions, WriteBatch, WriteOptions, DB};
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use tracing::{debug, error, info, warn};
thread_local! {
static KEY_BUFFER: RefCell<Vec<u8>> = RefCell::new(Vec::with_capacity(256));
}
#[derive(Clone, Copy)]
enum RowDecodeHint<'a> {
Full,
Prefix(usize),
Columns(&'a [usize]),
}
#[derive(Clone, Copy)]
pub(crate) enum ColumnarAggregateOp {
CountStar,
Count,
CountDistinct,
Sum,
Avg,
Min,
Max,
}
#[derive(Clone, Copy)]
pub(crate) struct ColumnarAggregateSpec {
pub op: ColumnarAggregateOp,
pub column_index: Option<usize>,
}
fn schema_uses_column_storage(schema: &crate::Schema) -> bool {
schema
.columns
.iter()
.any(|column| column.storage_mode != ColumnStorageMode::Default)
}
fn columnar_filter_predicate(predicate: &AnalyzedPredicate) -> Option<FilterPredicate> {
let op = predicate.op.to_filter_op()?;
Some(FilterPredicate {
column_index: predicate.column_index,
column_name: predicate.column_name.clone(),
op,
value: predicate.value.clone(),
value2: predicate.value2.clone(),
value_list: predicate.value_list.clone(),
pattern: match &predicate.value {
Value::String(pattern) if predicate.op == super::predicate_pushdown::PredicateOp::Like => {
Some(pattern.clone())
}
_ => None,
},
})
}
fn zone_map_pruning_enabled() -> bool {
static ZONE_MAP_OFF: once_cell::sync::Lazy<bool> =
once_cell::sync::Lazy::new(|| std::env::var("HELIOS_ZONE_MAP_OFF").is_ok());
!*ZONE_MAP_OFF
}
fn columnar_presence_enabled() -> bool {
static COLP_OFF: once_cell::sync::Lazy<bool> =
once_cell::sync::Lazy::new(|| std::env::var("HELIOS_COLP_OFF").is_ok());
!*COLP_OFF
}
fn schema_has_columnar_columns(schema: &crate::Schema) -> bool {
schema
.columns
.iter()
.any(|column| column.storage_mode == ColumnStorageMode::Columnar)
}
fn group_columnar_row_values(
schema: &crate::Schema,
rows: &[(u64, &Tuple)],
) -> std::collections::BTreeMap<(usize, u64), Vec<(u64, Value)>> {
let mut grouped: std::collections::BTreeMap<(usize, u64), Vec<(u64, Value)>> = std::collections::BTreeMap::new();
for (row_id, tuple) in rows {
let (batch_id, _) = ColumnarStore::batch_location(*row_id);
for (idx, column) in schema.columns.iter().enumerate() {
if column.storage_mode != ColumnStorageMode::Columnar {
continue;
}
let Some(value) = tuple.values.get(idx) else {
continue;
};
grouped
.entry((idx, batch_id))
.or_default()
.push((*row_id, value.clone()));
}
}
grouped
}
fn group_row_ids_by_batch(row_ids: impl IntoIterator<Item = u64>) -> std::collections::BTreeMap<u64, Vec<u64>> {
let mut grouped: std::collections::BTreeMap<u64, Vec<u64>> = std::collections::BTreeMap::new();
for row_id in row_ids {
grouped.entry(row_id / BATCH_SIZE as u64).or_default().push(row_id);
}
grouped
}
fn presence_parallel_partials<T, I, B>(
live: &[(u64, &BatchPresence)],
batches_per_chunk: usize,
init: I,
body: B,
) -> Result<Vec<T>>
where
T: Send,
I: Fn() -> T + Sync,
B: Fn(&mut T, u64, usize) -> Result<()> + Sync + Send,
{
use rayon::prelude::*;
live.par_chunks(batches_per_chunk.max(1))
.map(|chunk| {
let mut acc = init();
for (batch_id, presence) in chunk {
for offset in presence.iter_live() {
body(&mut acc, *batch_id, offset)?;
}
}
Ok(acc)
})
.collect()
}
fn presence_live_view<'a>(
presence: &'a [(u64, BatchPresence)],
pruned: &HashSet<u64>,
) -> Vec<(u64, &'a BatchPresence)> {
presence
.iter()
.filter(|(batch_id, _)| !pruned.contains(batch_id))
.map(|(batch_id, batch_presence)| (*batch_id, batch_presence))
.collect()
}
pub(crate) fn batch_can_match(stats: &BatchStats, predicate: &FilterPredicate) -> bool {
if stats.null_count > 0 && predicate.evaluate(&Value::Null) {
return true;
}
if stats.null_count >= stats.row_count {
return false;
}
let (Some(min), Some(max)) = (&stats.min, &stats.max) else {
return true;
};
let lt = FilterPredicate::compare_lt;
let gt = FilterPredicate::compare_gt;
let eq = FilterPredicate::compare_eq;
let value = &predicate.value;
match predicate.op {
FilterOp::Eq => !(lt(value, min) || gt(value, max)),
FilterOp::NotEq => !(eq(min, max) && eq(value, min)),
FilterOp::Lt => !(lt(value, min) || eq(value, min)),
FilterOp::LtEq => !lt(value, min),
FilterOp::Gt => !(gt(value, max) || eq(value, max)),
FilterOp::GtEq => !gt(value, max),
FilterOp::Between => match &predicate.value2 {
Some(high) => !(lt(high, min) || gt(value, max)),
None => false,
},
FilterOp::NotBetween => match &predicate.value2 {
Some(high) => !((gt(min, value) || eq(min, value)) && (lt(max, high) || eq(max, high))),
None => true,
},
FilterOp::In => predicate
.value_list
.iter()
.any(|item| !(lt(item, min) || gt(item, max))),
FilterOp::NotIn => !(eq(min, max) && predicate.value_list.iter().any(|item| eq(item, min))),
FilterOp::IsNull => false,
FilterOp::IsNotNull => true,
FilterOp::Like | FilterOp::NotLike => true,
}
}
fn compute_zone_pruned_batches(
db: &DB,
table_name: &str,
schema: &crate::Schema,
predicates: &[FilterPredicate],
) -> HashSet<u64> {
let mut pruned = HashSet::new();
if predicates.is_empty() || !zone_map_pruning_enabled() {
return pruned;
}
let mut by_column: HashMap<usize, Vec<&FilterPredicate>> = HashMap::new();
for predicate in predicates {
by_column.entry(predicate.column_index).or_default().push(predicate);
}
for (column_index, column_predicates) in by_column {
let Some(column) = schema.columns.get(column_index) else {
continue;
};
if column.storage_mode != ColumnStorageMode::Columnar {
continue;
}
let stats_map = match ColumnarStore::load_stats_map(db, table_name, &column.name) {
Ok(map) => map,
Err(e) => {
tracing::warn!(table = table_name, column = column.name.as_str(), error = %e,
"Zone-stats load failed; scanning column unpruned");
continue;
}
};
for (batch_id, stats) in &stats_map {
if column_predicates
.iter()
.any(|predicate| !batch_can_match(stats, predicate))
{
pruned.insert(*batch_id);
}
}
}
pruned
}
fn load_columnar_batch_index(
db: &DB,
table_name: &str,
column_name: &str,
pruned: &HashSet<u64>,
) -> Result<ColumnarBatchIndex> {
let batches = if zone_map_pruning_enabled() {
ColumnarStore::scan_column_batches_pruned(db, table_name, column_name, pruned)?
} else {
ColumnarStore::scan_column_batches(db, table_name, column_name)?
};
Ok(ColumnarBatchIndex::from_batches(batches))
}
enum ColumnarBatchIndex {
Dense(Vec<Option<ColumnBatch>>),
Sparse(HashMap<u64, ColumnBatch>),
}
impl ColumnarBatchIndex {
fn from_batches(batches: Vec<(u64, ColumnBatch)>) -> Self {
let Some(max_batch_id) = batches.last().map(|(batch_id, _)| *batch_id) else {
return Self::Dense(Vec::new());
};
let batch_count = batches.len();
if max_batch_id <= (batch_count as u64).saturating_mul(4).saturating_add(16) && max_batch_id <= 1_000_000 {
let mut dense = vec![None; max_batch_id as usize + 1];
for (batch_id, batch) in batches {
if let Some(slot) = dense.get_mut(batch_id as usize) {
*slot = Some(batch);
}
}
Self::Dense(dense)
} else {
Self::Sparse(batches.into_iter().collect())
}
}
fn get(&self, batch_id: u64) -> Option<&ColumnBatch> {
match self {
Self::Dense(batches) => batches.get(batch_id as usize).and_then(Option::as_ref),
Self::Sparse(batches) => batches.get(&batch_id),
}
}
fn ordered_batches(&self) -> Vec<(u64, &ColumnBatch)> {
match self {
Self::Dense(batches) => batches
.iter()
.enumerate()
.filter_map(|(batch_id, batch)| batch.as_ref().map(|batch| (batch_id as u64, batch)))
.collect(),
Self::Sparse(batches) => {
let mut ordered: Vec<_> = batches.iter().map(|(batch_id, batch)| (*batch_id, batch)).collect();
ordered.sort_by_key(|(batch_id, _)| *batch_id);
ordered
}
}
}
}
fn null_rejecting_filter_predicate(predicates: &[FilterPredicate]) -> Option<&FilterPredicate> {
predicates.iter().find(|predicate| !predicate.evaluate(&Value::Null))
}
fn columnar_batch_driver_scan_beneficial(schema: &crate::Schema, predicate: &FilterPredicate) -> bool {
schema.columns.get(predicate.column_index).is_some_and(|column| {
matches!(
column.data_type,
crate::DataType::Text | crate::DataType::Varchar(_) | crate::DataType::Char(_)
)
})
}
fn columnar_row_matches_filters(
column_batches: &HashMap<usize, ColumnarBatchIndex>,
batch_id: u64,
offset: usize,
predicates: &[FilterPredicate],
) -> bool {
predicates.iter().all(|predicate| {
if let Some(value) = column_batches
.get(&predicate.column_index)
.and_then(|by_batch| by_batch.get(batch_id))
.and_then(|batch| batch.value_at(offset))
{
predicate.evaluate(&value)
} else {
predicate.evaluate(&Value::Null)
}
})
}
fn row_tuple_matches_filters(tuple: &Tuple, predicates: &[FilterPredicate]) -> bool {
predicates.iter().all(|predicate| {
tuple
.values
.get(predicate.column_index)
.is_some_and(|value| predicate.evaluate(value))
})
}
fn row_values_match_filters(values: &[Value], predicates: &[FilterPredicate], positions: &[usize]) -> bool {
predicates
.iter()
.zip(positions)
.all(|(predicate, &pos)| values.get(pos).is_some_and(|value| predicate.evaluate(value)))
}
#[derive(Clone, Copy)]
struct IntegerFilterCandidate {
column_index: usize,
op: FilterOp,
value: i64,
}
#[derive(Clone, Copy)]
struct StringEqFilterCandidate<'a> {
column_index: usize,
value: &'a str,
}
fn integer_filter_candidate(schema: &crate::Schema, predicates: &[FilterPredicate]) -> Option<IntegerFilterCandidate> {
let [predicate] = predicates else {
return None;
};
if !matches!(
predicate.op,
FilterOp::Eq | FilterOp::Lt | FilterOp::LtEq | FilterOp::Gt | FilterOp::GtEq
) {
return None;
}
let column = schema.columns.get(predicate.column_index)?;
if !primitive_integer_data_type(&column.data_type) {
return None;
}
let value = match predicate.value {
Value::Int2(value) => i64::from(value),
Value::Int4(value) => i64::from(value),
Value::Int8(value) => value,
_ => return None,
};
Some(IntegerFilterCandidate {
column_index: predicate.column_index,
op: predicate.op,
value,
})
}
fn integer_filter_candidates(
schema: &crate::Schema,
predicates: &[FilterPredicate],
) -> Option<Vec<IntegerFilterCandidate>> {
if predicates.is_empty() {
return None;
}
let mut candidates = Vec::with_capacity(predicates.len());
for predicate in predicates {
if !matches!(
predicate.op,
FilterOp::Eq | FilterOp::Lt | FilterOp::LtEq | FilterOp::Gt | FilterOp::GtEq
) {
return None;
}
let column = schema.columns.get(predicate.column_index)?;
if !primitive_integer_data_type(&column.data_type) {
return None;
}
let value = match predicate.value {
Value::Int2(value) => i64::from(value),
Value::Int4(value) => i64::from(value),
Value::Int8(value) => value,
_ => return None,
};
candidates.push(IntegerFilterCandidate {
column_index: predicate.column_index,
op: predicate.op,
value,
});
}
Some(candidates)
}
fn string_eq_filter_candidate<'a>(
schema: &crate::Schema,
predicates: &'a [FilterPredicate],
) -> Option<StringEqFilterCandidate<'a>> {
let [predicate] = predicates else {
return None;
};
if predicate.op != FilterOp::Eq {
return None;
}
let column = schema.columns.get(predicate.column_index)?;
if !matches!(
column.data_type,
crate::DataType::Text | crate::DataType::Varchar(_) | crate::DataType::Char(_)
) {
return None;
}
let Value::String(value) = &predicate.value else {
return None;
};
Some(StringEqFilterCandidate {
column_index: predicate.column_index,
value,
})
}
fn decoded_integer_matches_filter(
decoded: crate::storage::prefix_decode::DecodedNumericValue,
filter: IntegerFilterCandidate,
) -> Option<bool> {
let value = match decoded {
crate::storage::prefix_decode::DecodedNumericValue::Null => return Some(false),
crate::storage::prefix_decode::DecodedNumericValue::Int(value) => value,
crate::storage::prefix_decode::DecodedNumericValue::Float(_) => return None,
};
Some(match filter.op {
FilterOp::Eq => value == filter.value,
FilterOp::Lt => value < filter.value,
FilterOp::LtEq => value <= filter.value,
FilterOp::Gt => value > filter.value,
FilterOp::GtEq => value >= filter.value,
_ => return None,
})
}
fn columnar_batch_value(
column_batches: &HashMap<usize, ColumnarBatchIndex>,
column_index: usize,
batch_id: u64,
offset: usize,
) -> Option<Value> {
column_batches
.get(&column_index)
.and_then(|by_batch| by_batch.get(batch_id))
.and_then(|batch| batch.value_at(offset))
}
fn compare_value_slices(left: &[Value], right: &[Value]) -> std::cmp::Ordering {
for (a, b) in left.iter().zip(right) {
let cmp = crate::sql::executor::compare_values(a, b);
if cmp != std::cmp::Ordering::Equal {
return cmp;
}
}
left.len().cmp(&right.len())
}
fn update_text_count_sum_group(
group_key: Option<&str>,
sum_value: Option<i64>,
small_groups: &mut Vec<(Option<String>, CountSumIntState)>,
hash_groups: &mut Option<HashMap<Option<String>, CountSumIntState>>,
linear_group_limit: usize,
) -> Result<()> {
if let Some(groups) = hash_groups.as_mut() {
let state = groups
.entry(group_key.map(str::to_owned))
.or_insert_with(CountSumIntState::new);
state.update_count();
state.update_sum_int(sum_value)?;
return Ok(());
}
if let Some(idx) = small_groups
.iter()
.position(|(existing, _)| existing.as_deref() == group_key)
{
let state = &mut small_groups[idx].1;
state.update_count();
state.update_sum_int(sum_value)?;
return Ok(());
}
let mut state = CountSumIntState::new();
state.update_count();
state.update_sum_int(sum_value)?;
small_groups.push((group_key.map(str::to_owned), state));
if small_groups.len() > linear_group_limit {
*hash_groups = Some(small_groups.drain(..).collect());
}
Ok(())
}
fn group_key_matches_values(key: &[Value], values: &[Value], positions: &[usize]) -> bool {
key.len() == positions.len()
&& key.iter().zip(positions).all(|(expected, &pos)| {
values
.get(pos)
.map_or(matches!(expected, Value::Null), |actual| expected == actual)
})
}
fn build_group_key(values: &[Value], positions: &[usize]) -> Vec<Value> {
positions
.iter()
.map(|&pos| values.get(pos).cloned().unwrap_or(Value::Null))
.collect()
}
const AGG_PARALLEL_CHUNK: usize = 16_384;
fn agg_parallel_threshold() -> Option<usize> {
static AGG_SERIAL: once_cell::sync::Lazy<bool> =
once_cell::sync::Lazy::new(|| std::env::var("HELIOS_AGG_SERIAL").is_ok());
static AGG_THRESHOLD: once_cell::sync::Lazy<usize> = once_cell::sync::Lazy::new(|| {
std::env::var("HELIOS_AGG_PARALLEL_THRESHOLD")
.ok()
.and_then(|raw| raw.trim().parse().ok())
.unwrap_or(65_536)
});
if *AGG_SERIAL {
None
} else {
Some(*AGG_THRESHOLD)
}
}
fn agg_parallel_rows_met(row_count: usize) -> bool {
agg_parallel_threshold().is_some_and(|threshold| row_count >= threshold)
}
fn row_agg_parallel_rows_met(row_count: usize) -> bool {
if agg_parallel_threshold().is_none() {
return false;
}
static ROW_AGG_THRESHOLD: once_cell::sync::Lazy<usize> = once_cell::sync::Lazy::new(|| {
std::env::var("HELIOS_ROW_AGG_PARALLEL_THRESHOLD")
.or_else(|_| std::env::var("HELIOS_AGG_PARALLEL_THRESHOLD"))
.ok()
.and_then(|raw| raw.trim().parse().ok())
.unwrap_or(32_768)
});
row_count >= *ROW_AGG_THRESHOLD
}
fn merge_columnar_aggregate_states(
into: &mut [ColumnarAggregateState],
from: Vec<ColumnarAggregateState>,
) -> Result<()> {
for (state, partial) in into.iter_mut().zip(from) {
state.merge(partial)?;
}
Ok(())
}
fn merge_columnar_group_map(
into: &mut HashMap<Vec<Value>, Vec<ColumnarAggregateState>>,
from: HashMap<Vec<Value>, Vec<ColumnarAggregateState>>,
) -> Result<()> {
for (key, states) in from {
match into.entry(key) {
std::collections::hash_map::Entry::Occupied(mut entry) => {
merge_columnar_aggregate_states(entry.get_mut(), states)?;
}
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert(states);
}
}
}
Ok(())
}
fn prefix_shard_bounds(prefix: &str) -> Vec<(Vec<u8>, Vec<u8>)> {
let prefix_bytes = prefix.as_bytes();
debug_assert!(prefix_bytes.last() == Some(&b':'));
let mut prefix_end = prefix_bytes.to_vec();
if let Some(last) = prefix_end.last_mut() {
*last += 1;
}
let mut boundaries: Vec<Vec<u8>> = Vec::with_capacity(92);
boundaries.push(prefix_bytes.to_vec()); for first in b'1'..=b'9' {
let mut one_digit = prefix_bytes.to_vec();
one_digit.push(first);
boundaries.push(one_digit.clone());
for second in b'0'..=b'9' {
let mut two_digit = one_digit.clone();
two_digit.push(second);
boundaries.push(two_digit);
}
}
boundaries.push(prefix_end);
let ends = boundaries.iter().skip(1).cloned().collect::<Vec<_>>();
boundaries.truncate(ends.len());
boundaries.into_iter().zip(ends).collect()
}
fn merge_count_sum_group_map<K: std::hash::Hash + Eq>(
into: &mut HashMap<K, CountSumIntState>,
from: HashMap<K, CountSumIntState>,
) -> Result<()> {
for (key, state) in from {
match into.entry(key) {
std::collections::hash_map::Entry::Occupied(mut entry) => entry.get_mut().merge(state)?,
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert(state);
}
}
}
Ok(())
}
fn columnar_group_key_matches(
key: &[Value],
column_batches: &HashMap<usize, ColumnarBatchIndex>,
group_by_columns: &[usize],
batch_id: u64,
offset: usize,
) -> bool {
key.len() == group_by_columns.len()
&& key.iter().zip(group_by_columns).all(|(expected, &idx)| {
columnar_batch_value(column_batches, idx, batch_id, offset)
.map_or(matches!(expected, Value::Null), |actual| *expected == actual)
})
}
fn build_columnar_group_key(
column_batches: &HashMap<usize, ColumnarBatchIndex>,
group_by_columns: &[usize],
batch_id: u64,
offset: usize,
) -> Vec<Value> {
group_by_columns
.iter()
.map(|&idx| columnar_batch_value(column_batches, idx, batch_id, offset).unwrap_or(Value::Null))
.collect()
}
fn update_columnar_aggregate_states(
states: &mut [ColumnarAggregateState],
aggregates: &[ColumnarAggregateSpec],
column_batches: &HashMap<usize, ColumnarBatchIndex>,
batch_id: u64,
offset: usize,
) -> Result<()> {
for (state, aggregate) in states.iter_mut().zip(aggregates) {
let value = aggregate
.column_index
.and_then(|idx| columnar_batch_value(column_batches, idx, batch_id, offset));
state.update(aggregate.op, value.as_ref())?;
}
Ok(())
}
fn stats_int_bounds(stats: &BatchStats) -> Option<(i64, i64)> {
fn as_int(v: &Option<Value>) -> Option<i64> {
match v {
Some(Value::Int2(x)) => Some(*x as i64),
Some(Value::Int4(x)) => Some(*x as i64),
Some(Value::Int8(x)) => Some(*x),
_ => None,
}
}
Some((as_int(&stats.min)?, as_int(&stats.max)?))
}
fn apply_compiled_predicates(
compiled: &[CompiledPredicate<'_>],
column_batches: &HashMap<usize, ColumnarBatchIndex>,
batch_id: u64,
mask: &mut [u8],
) {
for cp in compiled {
let batch = column_batches
.get(&cp.pred.column_index)
.and_then(|by| by.get(batch_id));
cp.apply_to_mask(batch, mask);
}
}
fn kernel_update_aggregate(
state: &mut ColumnarAggregateState,
op: ColumnarAggregateOp,
batch: Option<&ColumnBatch>,
mask: &[u8],
) -> Result<bool> {
if matches!(op, ColumnarAggregateOp::CountStar) {
state.merge(ColumnarAggregateState::Count(tk::count_selected(mask) as i64))?;
return Ok(true);
}
let Some(batch) = batch else {
return Ok(true);
};
let Some(typed) = &batch.typed else {
return Ok(false); };
match (op, &typed.data) {
(ColumnarAggregateOp::Count, _) => {
state.merge(ColumnarAggregateState::Count(
tk::count_selected_valid(&typed.validity, mask) as i64,
))?;
Ok(true)
}
(ColumnarAggregateOp::Sum, TypedValues::Int { data, .. }) => {
let small = stats_int_bounds(&typed.stats)
.is_some_and(|(mn, mx)| mn >= -tk::SUM_I64_SAFE_BOUND && mx <= tk::SUM_I64_SAFE_BOUND);
let (sum, contributing) = tk::sum_int_selected(data, &typed.validity, mask, small);
if contributing > 0 {
let sum = i64::try_from(sum).map_err(|_| Error::query_execution("integer overflow: BIGINT SUM"))?;
state.merge(ColumnarAggregateState::Sum(ColumnarSumState::Int(sum)))?;
}
Ok(true)
}
(ColumnarAggregateOp::Avg, TypedValues::Int { data, .. }) => {
let (sum, count) = tk::sum_f64_from_int_selected(data, &typed.validity, mask);
if count > 0 {
state.merge(ColumnarAggregateState::Avg { sum, count })?;
}
Ok(true)
}
(ColumnarAggregateOp::Avg, TypedValues::Float { data, .. }) => {
let (sum, count) = tk::sum_f64_selected(data, &typed.validity, mask);
if count > 0 {
state.merge(ColumnarAggregateState::Avg { sum, count })?;
}
Ok(true)
}
(ColumnarAggregateOp::Min | ColumnarAggregateOp::Max, TypedValues::Int { width, data }) => {
if let Some((mn, mx)) = tk::min_max_int_selected(data, &typed.validity, mask) {
if matches!(op, ColumnarAggregateOp::Min) {
state.merge(ColumnarAggregateState::Min(Some(width.value(mn))))?;
} else {
state.merge(ColumnarAggregateState::Max(Some(width.value(mx))))?;
}
}
Ok(true)
}
(ColumnarAggregateOp::Min | ColumnarAggregateOp::Max, TypedValues::Float { width, data }) => {
let (mn, mx) = tk::min_max_float_selected(data, &typed.validity, mask);
if matches!(op, ColumnarAggregateOp::Min) {
if let Some(mn) = mn {
state.merge(ColumnarAggregateState::Min(Some(width.value(mn))))?;
}
} else if let Some(mx) = mx {
state.merge(ColumnarAggregateState::Max(Some(width.value(mx))))?;
}
Ok(true)
}
(ColumnarAggregateOp::Min | ColumnarAggregateOp::Max, TypedValues::Text { dict, codes }) => {
let mut used = vec![0u8; dict.len()];
tk::mark_used_codes(codes, &typed.validity, mask, &mut used);
let mut best: Option<&String> = None;
for (entry, &u) in dict.iter().zip(&used) {
if u == 0 {
continue;
}
best = Some(match best {
None => entry,
Some(current) => {
let replace = if matches!(op, ColumnarAggregateOp::Min) {
entry < current
} else {
entry > current
};
if replace {
entry
} else {
current
}
}
});
}
if let Some(best) = best {
let value = Value::String(best.clone());
if matches!(op, ColumnarAggregateOp::Min) {
state.merge(ColumnarAggregateState::Min(Some(value)))?;
} else {
state.merge(ColumnarAggregateState::Max(Some(value)))?;
}
}
Ok(true)
}
_ => Ok(false),
}
}
fn update_states_for_batch_kernel(
states: &mut [ColumnarAggregateState],
aggregates: &[ColumnarAggregateSpec],
column_batches: &HashMap<usize, ColumnarBatchIndex>,
batch_id: u64,
mask: &[u8],
) -> Result<()> {
for (state, aggregate) in states.iter_mut().zip(aggregates) {
let batch = aggregate
.column_index
.and_then(|idx| column_batches.get(&idx).and_then(|by| by.get(batch_id)));
if !kernel_update_aggregate(state, aggregate.op, batch, mask)? {
for (offset, &m) in mask.iter().enumerate() {
if m != 0 {
let value = batch.and_then(|b| b.value_at(offset));
state.update(aggregate.op, value.as_ref())?;
}
}
}
}
Ok(())
}
fn update_groups_for_batch_kernel(
groups: &mut HashMap<Vec<Value>, Vec<ColumnarAggregateState>>,
aggregates: &[ColumnarAggregateSpec],
group_by_columns: &[usize],
column_batches: &HashMap<usize, ColumnarBatchIndex>,
batch_id: u64,
mask: &[u8],
) -> Result<()> {
if let [group_col] = group_by_columns {
let gbatch = column_batches.get(group_col).and_then(|by| by.get(batch_id));
if let Some(gtyped) = gbatch.and_then(|b| b.typed.as_ref()) {
if let TypedValues::Text { dict, codes } = >yped.data {
let null_slot = dict.len();
let mut per_code: Vec<Option<Vec<ColumnarAggregateState>>> =
(0..dict.len() + 1).map(|_| None).collect();
for (offset, &m) in mask.iter().enumerate() {
if m == 0 {
continue;
}
let idx = match (codes.get(offset), gtyped.validity.get(offset)) {
(Some(&code), Some(&ok)) if ok != 0 => (code as usize).min(null_slot),
_ => null_slot,
};
let Some(slot) = per_code.get_mut(idx) else { continue };
let states = slot.get_or_insert_with(|| {
aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect()
});
update_columnar_aggregate_states(states, aggregates, column_batches, batch_id, offset)?;
}
for (code, states) in per_code.into_iter().enumerate() {
let Some(states) = states else { continue };
let key = match dict.get(code) {
Some(entry) if code < null_slot => vec![Value::String(entry.clone())],
_ => vec![Value::Null],
};
match groups.entry(key) {
std::collections::hash_map::Entry::Occupied(mut entry) => {
merge_columnar_aggregate_states(entry.get_mut(), states)?;
}
std::collections::hash_map::Entry::Vacant(entry) => {
entry.insert(states);
}
}
}
return Ok(());
}
}
}
for (offset, &m) in mask.iter().enumerate() {
if m == 0 {
continue;
}
let group_key = build_columnar_group_key(column_batches, group_by_columns, batch_id, offset);
let states = groups.entry(group_key).or_insert_with(|| {
aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect()
});
update_columnar_aggregate_states(states, aggregates, column_batches, batch_id, offset)?;
}
Ok(())
}
struct RowTopKEntry {
key: Vec<Value>,
tuple: Tuple,
asc: Arc<Vec<bool>>,
}
impl PartialEq for RowTopKEntry {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == std::cmp::Ordering::Equal
}
}
impl Eq for RowTopKEntry {}
impl PartialOrd for RowTopKEntry {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for RowTopKEntry {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
compare_topk_keys(&self.key, &other.key, &self.asc)
}
}
fn compare_topk_keys(left: &[Value], right: &[Value], asc: &[bool]) -> std::cmp::Ordering {
for (idx, (left, right)) in left.iter().zip(right).enumerate() {
let mut cmp = crate::sql::executor::compare_values(left, right);
if !asc.get(idx).copied().unwrap_or(true) {
cmp = cmp.reverse();
}
if cmp != std::cmp::Ordering::Equal {
return cmp;
}
}
left.len().cmp(&right.len())
}
fn compare_values_to_topk_key(
values: &[Value],
positions: &[usize],
key: &[Value],
asc: &[bool],
) -> std::cmp::Ordering {
for (idx, (&pos, right)) in positions.iter().zip(key).enumerate() {
let left = values.get(pos).unwrap_or(&Value::Null);
let mut cmp = crate::sql::executor::compare_values(left, right);
if !asc.get(idx).copied().unwrap_or(true) {
cmp = cmp.reverse();
}
if cmp != std::cmp::Ordering::Equal {
return cmp;
}
}
positions.len().cmp(&key.len())
}
fn build_topk_key(values: &[Value], positions: &[usize]) -> Vec<Value> {
positions
.iter()
.map(|&pos| values.get(pos).cloned().unwrap_or(Value::Null))
.collect()
}
fn build_projected_tuple(values: &[Value], positions: &[usize]) -> Tuple {
Tuple::new(
positions
.iter()
.map(|&pos| values.get(pos).cloned().unwrap_or(Value::Null))
.collect(),
)
}
struct RowIntTopKEntry {
key: i64,
raw_value: Box<[u8]>,
row_id: Option<u64>,
asc: bool,
}
impl PartialEq for RowIntTopKEntry {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == std::cmp::Ordering::Equal
}
}
impl Eq for RowIntTopKEntry {}
impl PartialOrd for RowIntTopKEntry {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for RowIntTopKEntry {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
let cmp = self.key.cmp(&other.key);
if self.asc {
cmp
} else {
cmp.reverse()
}
}
}
fn compare_int_to_topk_key(left: i64, right: i64, asc: bool) -> std::cmp::Ordering {
let cmp = left.cmp(&right);
if asc {
cmp
} else {
cmp.reverse()
}
}
enum ColumnarSumState {
Empty,
Int(i64),
Decimal(rust_decimal::Decimal),
}
impl ColumnarSumState {
fn merge(&mut self, other: Self) -> Result<()> {
match other {
Self::Empty => Ok(()),
Self::Int(value) => update_columnar_int_sum(self, value),
Self::Decimal(value) => {
update_columnar_decimal_sum(self, value);
Ok(())
}
}
}
}
enum ColumnarAggregateState {
Count(i64),
CountDistinct(HashSet<Value>),
Sum(ColumnarSumState),
Avg { sum: f64, count: u64 },
Min(Option<Value>),
Max(Option<Value>),
}
#[derive(Clone, Copy)]
struct CountSumIntState {
count: i64,
sum: i64,
sum_seen: bool,
}
impl CountSumIntState {
fn new() -> Self {
Self {
count: 0,
sum: 0,
sum_seen: false,
}
}
fn update_count(&mut self) {
self.count += 1;
}
fn update_sum(&mut self, value: Option<&Value>) -> Result<()> {
let Some(value) = value else {
return Ok(());
};
let add = match value {
Value::Null => return Ok(()),
Value::Int2(value) => *value as i64,
Value::Int4(value) => *value as i64,
Value::Int8(value) => *value,
_ => return Err(Error::query_execution("SUM requires numeric values")),
};
self.update_sum_int(Some(add))
}
fn update_sum_int(&mut self, value: Option<i64>) -> Result<()> {
let Some(add) = value else {
return Ok(());
};
self.sum = self
.sum
.checked_add(add)
.ok_or_else(|| Error::query_execution("integer overflow: BIGINT SUM"))?;
self.sum_seen = true;
Ok(())
}
fn merge(&mut self, other: Self) -> Result<()> {
self.count += other.count;
if other.sum_seen {
self.update_sum_int(Some(other.sum))?;
}
Ok(())
}
fn finish(self, group_key: Value) -> Tuple {
Tuple::new(vec![
group_key,
Value::Int8(self.count),
if self.sum_seen {
Value::Int8(self.sum)
} else {
Value::Null
},
])
}
}
#[derive(Clone, Copy)]
enum PrimitiveRowAggregate {
CountStar,
Count { position: usize },
SumInt { position: usize },
Avg { position: usize },
}
enum PrimitiveRowAggregateState {
Count(i64),
SumInt { sum: i64, seen: bool },
Avg { sum: f64, count: u64 },
}
impl PrimitiveRowAggregateState {
fn new(plan: PrimitiveRowAggregate) -> Self {
match plan {
PrimitiveRowAggregate::CountStar | PrimitiveRowAggregate::Count { .. } => Self::Count(0),
PrimitiveRowAggregate::SumInt { .. } => Self::SumInt { sum: 0, seen: false },
PrimitiveRowAggregate::Avg { .. } => Self::Avg { sum: 0.0, count: 0 },
}
}
fn update(
&mut self,
plan: PrimitiveRowAggregate,
values: &[crate::storage::prefix_decode::DecodedNumericValue],
) -> Result<()> {
match (self, plan) {
(Self::Count(count), PrimitiveRowAggregate::CountStar) => {
*count += 1;
Ok(())
}
(Self::Count(count), PrimitiveRowAggregate::Count { position }) => {
if !matches!(
values.get(position),
None | Some(crate::storage::prefix_decode::DecodedNumericValue::Null)
) {
*count += 1;
}
Ok(())
}
(Self::SumInt { sum, seen }, PrimitiveRowAggregate::SumInt { position }) => {
if let Some(crate::storage::prefix_decode::DecodedNumericValue::Int(value)) = values.get(position) {
*sum = sum
.checked_add(*value)
.ok_or_else(|| Error::query_execution("integer overflow: BIGINT SUM"))?;
*seen = true;
}
Ok(())
}
(Self::Avg { sum, count }, PrimitiveRowAggregate::Avg { position }) => {
match values.get(position) {
Some(crate::storage::prefix_decode::DecodedNumericValue::Int(value)) => {
*sum += *value as f64;
*count += 1;
}
Some(crate::storage::prefix_decode::DecodedNumericValue::Float(value)) => {
*sum += *value;
*count += 1;
}
_ => {}
}
Ok(())
}
_ => Err(Error::query_execution("Invalid primitive row aggregate state")),
}
}
fn merge(&mut self, other: Self) -> Result<()> {
match (self, other) {
(Self::Count(count), Self::Count(other)) => {
*count += other;
Ok(())
}
(
Self::SumInt { sum, seen },
Self::SumInt {
sum: other_sum,
seen: other_seen,
},
) => {
if other_seen {
*sum = sum
.checked_add(other_sum)
.ok_or_else(|| Error::query_execution("integer overflow: BIGINT SUM"))?;
*seen = true;
}
Ok(())
}
(
Self::Avg { sum, count },
Self::Avg {
sum: other_sum,
count: other_count,
},
) => {
*sum += other_sum;
*count += other_count;
Ok(())
}
_ => Err(Error::query_execution("Invalid primitive row aggregate merge")),
}
}
fn finalize(self) -> Value {
match self {
Self::Count(count) => Value::Int8(count),
Self::SumInt { sum, seen } => {
if seen {
Value::Int8(sum)
} else {
Value::Null
}
}
Self::Avg { sum, count } => {
if count == 0 {
Value::Null
} else {
Value::Float8(sum / count as f64)
}
}
}
}
}
impl ColumnarAggregateState {
fn new(op: ColumnarAggregateOp) -> Self {
match op {
ColumnarAggregateOp::CountStar | ColumnarAggregateOp::Count => Self::Count(0),
ColumnarAggregateOp::CountDistinct => Self::CountDistinct(HashSet::new()),
ColumnarAggregateOp::Sum => Self::Sum(ColumnarSumState::Empty),
ColumnarAggregateOp::Avg => Self::Avg { sum: 0.0, count: 0 },
ColumnarAggregateOp::Min => Self::Min(None),
ColumnarAggregateOp::Max => Self::Max(None),
}
}
fn update(&mut self, op: ColumnarAggregateOp, value: Option<&Value>) -> Result<()> {
match (self, op) {
(Self::Count(count), ColumnarAggregateOp::CountStar) => {
*count += 1;
Ok(())
}
(Self::Count(count), ColumnarAggregateOp::Count) => {
if !matches!(value, None | Some(Value::Null)) {
*count += 1;
}
Ok(())
}
(Self::CountDistinct(values), ColumnarAggregateOp::CountDistinct) => {
if let Some(value) = value {
if !matches!(value, Value::Null) {
values.insert(value.clone());
}
}
Ok(())
}
(Self::Sum(state), ColumnarAggregateOp::Sum) => {
let Some(value) = value else {
return Ok(());
};
update_columnar_sum(state, value)
}
(Self::Avg { sum, count }, ColumnarAggregateOp::Avg) => {
let Some(value) = value else {
return Ok(());
};
if matches!(value, Value::Null) {
return Ok(());
}
match value {
Value::Int2(i) => {
*sum += *i as f64;
*count += 1;
}
Value::Int4(i) => {
*sum += *i as f64;
*count += 1;
}
Value::Int8(i) => {
*sum += *i as f64;
*count += 1;
}
Value::Float4(f) => {
*sum += *f as f64;
*count += 1;
}
Value::Float8(f) => {
*sum += *f;
*count += 1;
}
Value::Numeric(n) => {
if let Ok(f) = n.parse::<f64>() {
*sum += f;
*count += 1;
}
}
_ => return Err(Error::query_execution("AVG requires numeric values")),
}
Ok(())
}
(Self::Min(current), ColumnarAggregateOp::Min) => {
let Some(value) = value else {
return Ok(());
};
if matches!(value, Value::Null) {
return Ok(());
}
match current {
None => *current = Some(value.clone()),
Some(existing) => {
if crate::sql::executor::compare_values(value, existing) == std::cmp::Ordering::Less {
*current = Some(value.clone());
}
}
}
Ok(())
}
(Self::Max(current), ColumnarAggregateOp::Max) => {
let Some(value) = value else {
return Ok(());
};
if matches!(value, Value::Null) {
return Ok(());
}
match current {
None => *current = Some(value.clone()),
Some(existing) => {
if crate::sql::executor::compare_values(value, existing) == std::cmp::Ordering::Greater {
*current = Some(value.clone());
}
}
}
Ok(())
}
_ => Err(Error::query_execution("Invalid columnar aggregate state")),
}
}
fn merge(&mut self, other: Self) -> Result<()> {
match (self, other) {
(Self::Count(count), Self::Count(other)) => {
*count += other;
Ok(())
}
(Self::CountDistinct(values), Self::CountDistinct(other)) => {
values.extend(other);
Ok(())
}
(Self::Sum(state), Self::Sum(other)) => state.merge(other),
(
Self::Avg { sum, count },
Self::Avg {
sum: other_sum,
count: other_count,
},
) => {
*sum += other_sum;
*count += other_count;
Ok(())
}
(Self::Min(current), Self::Min(other)) => {
if let Some(value) = other {
match current {
None => *current = Some(value),
Some(existing) => {
if crate::sql::executor::compare_values(&value, existing) == std::cmp::Ordering::Less {
*current = Some(value);
}
}
}
}
Ok(())
}
(Self::Max(current), Self::Max(other)) => {
if let Some(value) = other {
match current {
None => *current = Some(value),
Some(existing) => {
if crate::sql::executor::compare_values(&value, existing) == std::cmp::Ordering::Greater {
*current = Some(value);
}
}
}
}
Ok(())
}
_ => Err(Error::query_execution("Invalid columnar aggregate merge")),
}
}
fn finalize(self) -> Result<Value> {
match self {
Self::Count(count) => Ok(Value::Int8(count)),
Self::CountDistinct(values) => Ok(Value::Int8(values.len() as i64)),
Self::Sum(state) => match state {
ColumnarSumState::Empty => Ok(Value::Null),
ColumnarSumState::Int(sum) => Ok(Value::Int8(sum)),
ColumnarSumState::Decimal(sum) => Ok(Value::Numeric(format!("{sum}"))),
},
Self::Avg { sum, count } => {
if count == 0 {
Ok(Value::Null)
} else {
Ok(Value::Float8(sum / count as f64))
}
}
Self::Min(value) | Self::Max(value) => Ok(value.unwrap_or(Value::Null)),
}
}
}
fn update_columnar_sum(state: &mut ColumnarSumState, value: &Value) -> Result<()> {
if matches!(value, Value::Null) {
return Ok(());
}
match value {
Value::Int2(i) => update_columnar_int_sum(state, *i as i64),
Value::Int4(i) => update_columnar_int_sum(state, *i as i64),
Value::Int8(i) => update_columnar_int_sum(state, *i),
Value::Float4(f) => {
let dec = rust_decimal::Decimal::try_from(*f as f64).unwrap_or_default();
update_columnar_decimal_sum(state, dec);
Ok(())
}
Value::Float8(f) => {
let dec = rust_decimal::Decimal::try_from(*f).unwrap_or_default();
update_columnar_decimal_sum(state, dec);
Ok(())
}
Value::Numeric(n) => {
let dec = n.parse::<rust_decimal::Decimal>().unwrap_or_default();
update_columnar_decimal_sum(state, dec);
Ok(())
}
_ => Err(Error::query_execution("SUM requires numeric values")),
}
}
fn update_columnar_int_sum(state: &mut ColumnarSumState, value: i64) -> Result<()> {
match state {
ColumnarSumState::Empty => *state = ColumnarSumState::Int(value),
ColumnarSumState::Int(sum) => {
*sum = sum
.checked_add(value)
.ok_or_else(|| Error::query_execution("integer overflow: BIGINT SUM"))?;
}
ColumnarSumState::Decimal(sum) => *sum += rust_decimal::Decimal::from(value),
}
Ok(())
}
fn update_columnar_decimal_sum(state: &mut ColumnarSumState, value: rust_decimal::Decimal) {
match state {
ColumnarSumState::Empty => *state = ColumnarSumState::Decimal(value),
ColumnarSumState::Int(sum) => *state = ColumnarSumState::Decimal(rust_decimal::Decimal::from(*sum) + value),
ColumnarSumState::Decimal(sum) => *sum += value,
}
}
fn primitive_numeric_data_type(data_type: &crate::DataType) -> bool {
matches!(
data_type,
crate::DataType::Int2
| crate::DataType::Int4
| crate::DataType::Int8
| crate::DataType::Float4
| crate::DataType::Float8
)
}
fn primitive_integer_data_type(data_type: &crate::DataType) -> bool {
matches!(
data_type,
crate::DataType::Int2 | crate::DataType::Int4 | crate::DataType::Int8
)
}
fn int_value_for_data_type(data_type: &crate::DataType, value: i64) -> Option<Value> {
match data_type {
crate::DataType::Int2 => i16::try_from(value).ok().map(Value::Int2),
crate::DataType::Int4 => i32::try_from(value).ok().map(Value::Int4),
crate::DataType::Int8 => Some(Value::Int8(value)),
_ => None,
}
}
fn row_blob_fast_skip_supported(data_type: &crate::DataType) -> bool {
!matches!(
data_type,
crate::DataType::Date
| crate::DataType::Time
| crate::DataType::Timestamp
| crate::DataType::Timestamptz
| crate::DataType::Array(_)
)
}
thread_local! {
static SYNCHRONOUS_COMMIT_OVERRIDE: std::cell::Cell<Option<bool>> =
const { std::cell::Cell::new(None) };
}
pub struct SynchronousCommitOverrideGuard {
previous: Option<bool>,
}
impl Drop for SynchronousCommitOverrideGuard {
fn drop(&mut self) {
SYNCHRONOUS_COMMIT_OVERRIDE.with(|cell| cell.set(self.previous));
}
}
pub struct StorageEngine {
pub(crate) db: Arc<DB>,
config: Config,
pub(crate) timestamp: Arc<RwLock<u64>>,
key_manager: Option<Arc<KeyManager>>,
vector_indexes: Arc<VectorIndexManager>,
snapshot_manager: Arc<SnapshotManager>,
conflict_registry: Arc<super::conflict::WriteConflictRegistry>,
version_gc: Arc<super::version_gc::VersionGc>,
version_gc_worker: Option<super::version_gc::VersionGcWorker>,
branch_manager: Arc<RwLock<Option<Arc<BranchManager>>>>,
wal: Option<Arc<RwLock<WriteAheadLog>>>,
stats: Arc<DatabaseStats>,
statistics_cache: Arc<crate::storage::StatisticsCache>,
is_replaying: Arc<AtomicBool>,
#[cfg(feature = "sync-experimental")]
change_log: Option<Arc<RwLock<crate::sync::ChangeLogImpl>>>,
#[cfg(feature = "sync-experimental")]
node_id: uuid::Uuid,
mv_delta_tracker: Arc<super::MvDeltaTracker>,
current_branch: Arc<parking_lot::Mutex<Option<String>>>,
trigger_registry: Arc<crate::sql::TriggerRegistry>,
predicate_pushdown: Arc<PredicatePushdownManager>,
filter_delta_tracker: Arc<FilterIndexDeltaTracker>,
speculative_filter_manager: Arc<SpeculativeFilterManager>,
parallel_filter_engine: Arc<ParallelFilterEngine>,
cpu_monitor: Arc<CpuMonitor>,
consolidation_worker: Option<Arc<FilterConsolidationWorker>>,
_temp_dir: Option<tempfile::TempDir>,
row_counters: Arc<dashmap::DashMap<String, std::sync::atomic::AtomicU64>>,
bulk_load_mode: Arc<AtomicBool>,
lockfree_engine: Arc<RwLock<Option<super::lockfree::LockFreeIngestionEngine>>>,
dict_manager: Arc<DictionaryManager>,
art_index_manager: Arc<ArtIndexManager>,
row_cache: Arc<super::RowCache>,
data_bytes_written: Arc<AtomicU64>,
memory_limit_bytes: u64,
write_counter: Arc<AtomicU64>,
db_path: Option<std::path::PathBuf>,
memory_write_options: Option<WriteOptions>,
durable_commit_enabled: bool,
group_committer: Arc<super::group_commit::GroupCommitter>,
schema_cache: Arc<parking_lot::Mutex<std::collections::HashMap<String, crate::Schema>>>,
constraints_cache: Arc<parking_lot::Mutex<std::collections::HashMap<String, crate::sql::TableConstraints>>>,
referencing_fk_cache:
Arc<parking_lot::Mutex<std::collections::HashMap<String, Vec<crate::sql::ForeignKeyConstraint>>>>,
index_snapshots_on_close: Arc<AtomicBool>,
last_index_open_report: Arc<RwLock<Option<super::index_snapshot::IndexOpenReport>>>,
}
const MIN_DISK_SPACE_BYTES: u64 = 100 * 1024 * 1024;
impl StorageEngine {
fn memory_only_write_options() -> WriteOptions {
let mut opts = WriteOptions::default();
opts.set_sync(false);
opts.disable_wal(true);
opts
}
fn check_disk_space(path: &std::path::Path) -> Result<()> {
let output = std::process::Command::new("df")
.arg("--output=avail")
.arg("-B1") .arg(path)
.output();
match output {
Ok(out) if out.status.success() => {
let stdout = String::from_utf8_lossy(&out.stdout);
if let Some(avail_str) = stdout.lines().nth(1) {
if let Ok(available_bytes) = avail_str.trim().parse::<u64>() {
if available_bytes < MIN_DISK_SPACE_BYTES {
return Err(Error::storage(format!(
"Insufficient disk space: {} MB available (minimum {} MB required). \
Free disk space or use VACUUM to reclaim storage.",
available_bytes / (1024 * 1024),
MIN_DISK_SPACE_BYTES / (1024 * 1024)
)));
}
}
}
}
_ => {} }
Ok(())
}
fn extract_table_from_key(key: &[u8]) -> String {
let key_str = match std::str::from_utf8(key) {
Ok(s) => s,
Err(_) => return "unknown".to_string(),
};
if let Some(stripped) = key_str.strip_prefix("data:") {
if let Some(colon_pos) = stripped.find(':') {
return stripped[..colon_pos].to_string();
}
} else if let Some(stripped) = key_str.strip_prefix("meta:table:") {
return stripped.to_string();
} else if key_str.starts_with("meta:counter:") {
if let Some(stripped) = key_str.strip_prefix("meta:counter:") {
return stripped.to_string();
}
}
"unknown".to_string()
}
#[inline]
fn build_data_key(table_name: &str, row_id: u64) -> Vec<u8> {
KEY_BUFFER.with(|buf| {
let mut buf = buf.borrow_mut();
buf.clear();
buf.extend_from_slice(b"data:");
buf.extend_from_slice(table_name.as_bytes());
buf.push(b':');
let mut row_id_buf = itoa::Buffer::new();
buf.extend_from_slice(row_id_buf.format(row_id).as_bytes());
buf.clone()
})
}
#[inline]
fn parse_row_id_after_prefix(key: &[u8], prefix_len: usize) -> Option<u64> {
let suffix = key.get(prefix_len..)?;
if suffix.is_empty() {
return None;
}
let mut row_id = 0_u64;
for &byte in suffix {
if !byte.is_ascii_digit() {
return None;
}
row_id = row_id.checked_mul(10)?.checked_add(u64::from(byte - b'0'))?;
}
Some(row_id)
}
fn par_prefix_shard_results<T, F>(&self, prefix: &str, per_shard: F) -> Result<Vec<T>>
where
T: Send,
F: for<'a> Fn(rocksdb::DBIteratorWithThreadMode<'a, DB>) -> Result<T> + Sync,
{
use rayon::prelude::*;
let snapshot = self.db.snapshot();
prefix_shard_bounds(prefix)
.par_iter()
.map(|(start, end)| {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
read_opts.set_iterate_upper_bound(end.clone());
read_opts.set_snapshot(&snapshot);
let iter = self
.db
.iterator_opt(IteratorMode::From(start, rocksdb::Direction::Forward), read_opts);
per_shard(iter)
})
.collect()
}
pub fn open(path: impl AsRef<Path>, config: &Config) -> Result<Self> {
let db_path = path.as_ref().to_path_buf();
let mut opts = Options::default();
opts.create_if_missing(true);
opts.set_compression_type(match config.storage.compression {
crate::config::CompressionType::None => rocksdb::DBCompressionType::None,
crate::config::CompressionType::Zstd => rocksdb::DBCompressionType::Zstd,
crate::config::CompressionType::Lz4 => rocksdb::DBCompressionType::Lz4,
});
let cache_size = config.storage.cache_size;
let (block_cache_size, write_buffer_size) = match config.storage.rocksdb_write_buffer_size {
Some(wb) => (cache_size, wb),
None => {
let bc = (cache_size as f64 * 0.75) as usize;
(bc, cache_size - bc)
}
};
let block_cache = Cache::new_lru_cache(block_cache_size);
let mut block_opts = BlockBasedOptions::default();
block_opts.set_block_cache(&block_cache);
block_opts.set_block_size(16 * 1024); block_opts.set_cache_index_and_filter_blocks(true); block_opts.set_pin_l0_filter_and_index_blocks_in_cache(true);
block_opts.set_bloom_filter(14.0, false);
block_opts.set_whole_key_filtering(true);
opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(5));
opts.set_block_based_table_factory(&block_opts);
opts.set_write_buffer_size(write_buffer_size);
let sc = &config.storage;
opts.set_max_write_buffer_number(sc.rocksdb_max_write_buffer_number.unwrap_or(4)); opts.set_min_write_buffer_number_to_merge(sc.rocksdb_min_write_buffer_number_to_merge.unwrap_or(2)); opts.set_level_zero_file_num_compaction_trigger(sc.rocksdb_level0_file_num_compaction_trigger.unwrap_or(4));
opts.set_max_background_jobs(sc.rocksdb_max_background_jobs.unwrap_or(4)); opts.set_bytes_per_sync(sc.rocksdb_bytes_per_sync.unwrap_or(1048576)); opts.set_enable_pipelined_write(true);
let db = DB::open(&opts, path).map_err(|e| Error::storage(format!("Failed to open RocksDB: {}", e)))?;
let db = Arc::new(db);
let key_manager = if config.encryption.enabled {
let km = KeyManager::from_source(&config.encryption.key_source)?;
Some(Arc::new(km))
} else {
None
};
let snapshot_manager = Arc::new(SnapshotManager::new(Arc::clone(&db)));
if let Err(e) = snapshot_manager.recover_snapshots() {
warn!("Failed to recover snapshots: {}", e);
}
let timestamp = Arc::new(RwLock::new(
snapshot_manager.max_snapshot_timestamp().unwrap_or(1).max(1),
));
debug!("Initializing BranchManager");
let branch_manager = match BranchManager::new(Arc::clone(&db), Arc::clone(×tamp)) {
Ok(manager) => {
info!("BranchManager initialized successfully");
Arc::new(RwLock::new(Some(Arc::new(manager))))
}
Err(e) => {
warn!(
"Failed to initialize BranchManager: {}. Branch operations will be unavailable.",
e
);
Arc::new(RwLock::new(None))
}
};
let wal = if config.storage.wal_enabled {
let sync_mode = match config.storage.wal_sync_mode {
crate::config::WalSyncModeConfig::Sync => WalSyncMode::Sync,
crate::config::WalSyncModeConfig::Async => WalSyncMode::Async,
crate::config::WalSyncModeConfig::GroupCommit => WalSyncMode::GroupCommit,
};
match WriteAheadLog::open(Arc::clone(&db), sync_mode) {
Ok(wal) => {
info!("WAL initialized successfully");
Some(Arc::new(RwLock::new(wal)))
}
Err(e) => {
warn!("Failed to initialize WAL: {}. Durability guarantees may be reduced.", e);
None
}
}
} else {
debug!("WAL disabled in configuration");
None
};
let stats = Arc::new(DatabaseStats::new());
let statistics_cache = Arc::new(crate::storage::StatisticsCache::with_config(100, 30)?);
let mv_delta_tracker = Arc::new(super::MvDeltaTracker::new(Arc::clone(&db))?);
debug!("Delta tracker initialized for incremental MV refresh");
let trigger_registry = Arc::new(crate::sql::TriggerRegistry::new());
debug!("Trigger registry initialized");
let predicate_pushdown = Arc::new(PredicatePushdownManager::new(PushdownConfig::default()));
debug!("Predicate pushdown manager initialized");
let cpu_monitor = Arc::new(CpuMonitor::new());
let filter_delta_tracker = Arc::new(FilterIndexDeltaTracker::new(FilterIndexConfig::default()));
let speculative_filter_manager = Arc::new(SpeculativeFilterManager::new(SpeculativeConfig::default()));
let parallel_filter_engine = Arc::new(ParallelFilterEngine::new(ParallelFilterConfig::default()));
let consolidation_worker = {
let worker = FilterConsolidationWorker::new(
ConsolidationConfig::default(),
Arc::clone(&filter_delta_tracker),
Arc::clone(&cpu_monitor),
);
if let Err(e) = worker.start() {
warn!("Failed to start filter consolidation worker: {}", e);
}
Some(Arc::new(worker))
};
debug!("SMFI components initialized");
#[cfg(feature = "sync-experimental")]
let (change_log, node_id) = if config.sync.enabled && config.sync.change_log_enabled {
let node_id = if let Some(ref id_str) = config.sync.node_id {
uuid::Uuid::parse_str(id_str).map_err(|e| Error::config(format!("Invalid node_id UUID: {}", e)))?
} else {
uuid::Uuid::new_v4()
};
let cl = crate::sync::ChangeLogImpl::new(Arc::clone(&db))?;
info!("Sync enabled with node_id={}", node_id);
(Some(Arc::new(RwLock::new(cl))), node_id)
} else {
(None, uuid::Uuid::new_v4())
};
let row_counters = Arc::new(dashmap::DashMap::new());
let conflict_registry = Arc::new(super::conflict::WriteConflictRegistry::new());
let (version_gc, version_gc_worker) = Self::init_version_gc(
&config,
Arc::clone(&db),
Arc::clone(&snapshot_manager),
Arc::clone(&conflict_registry),
Arc::clone(&branch_manager),
Arc::clone(×tamp),
)?;
let engine = Self {
db: Arc::clone(&db),
config: config.clone(),
timestamp,
key_manager,
vector_indexes: Arc::new(VectorIndexManager::new()),
snapshot_manager,
conflict_registry,
version_gc,
version_gc_worker,
branch_manager,
wal,
stats,
statistics_cache,
is_replaying: Arc::new(AtomicBool::new(false)),
#[cfg(feature = "sync-experimental")]
change_log,
#[cfg(feature = "sync-experimental")]
node_id,
mv_delta_tracker,
current_branch: Arc::new(parking_lot::Mutex::new(None)),
trigger_registry,
predicate_pushdown,
filter_delta_tracker,
speculative_filter_manager,
parallel_filter_engine,
cpu_monitor,
consolidation_worker,
_temp_dir: None,
row_counters,
bulk_load_mode: Arc::new(AtomicBool::new(false)),
lockfree_engine: Arc::new(RwLock::new(None)),
dict_manager: Arc::new(DictionaryManager::new()),
art_index_manager: Arc::new(ArtIndexManager::new()),
row_cache: Arc::new(super::RowCache::new()),
data_bytes_written: Arc::new(AtomicU64::new(0)),
memory_limit_bytes: 0, write_counter: Arc::new(AtomicU64::new(0)),
db_path: Some(db_path),
memory_write_options: None,
durable_commit_enabled: config.storage.durable_commit && !config.storage.memory_only,
group_committer: Arc::new(super::group_commit::GroupCommitter::new(
std::time::Duration::from_micros(config.storage.group_commit_window_us),
)),
schema_cache: Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new())),
constraints_cache: Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new())),
referencing_fk_cache: Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new())),
index_snapshots_on_close: Arc::new(AtomicBool::new(true)),
last_index_open_report: Arc::new(RwLock::new(None)),
};
engine.load_counters()?;
engine.wire_index_snapshot_hooks();
if engine.wal.is_some() {
match engine.replay_wal() {
Ok(0) => debug!("WAL clean, no entries to replay"),
Ok(count) => {
info!("WAL crash recovery: replayed {} entries", count);
if let Some(wal) = &engine.wal {
let wal_guard = wal.read();
let current_lsn = wal_guard.current_lsn();
if let Err(e) = wal_guard.truncate(current_lsn) {
warn!("Failed to truncate WAL after replay: {}", e);
}
}
}
Err(e) => warn!("WAL replay failed (data may be incomplete): {}", e),
}
}
Ok(engine)
}
pub fn open_in_memory(config: &Config) -> Result<Self> {
let temp_dir = if cfg!(target_os = "linux") && std::path::Path::new("/dev/shm").is_dir() {
tempfile::Builder::new()
.prefix("heliosdb-nano-")
.tempdir_in("/dev/shm")
.or_else(|_| tempfile::Builder::new().prefix("heliosdb-nano-").tempdir())
} else {
tempfile::Builder::new().prefix("heliosdb-nano-").tempdir()
}
.map_err(|e| Error::storage(format!("Failed to create temp dir: {}", e)))?;
let mut opts = Options::default();
opts.create_if_missing(true);
opts.set_write_buffer_size(64 * 1024 * 1024); opts.set_max_write_buffer_number(4);
opts.set_min_write_buffer_number_to_merge(2);
opts.set_level_zero_file_num_compaction_trigger(8); opts.set_max_background_jobs(2);
opts.set_enable_pipelined_write(true);
let mut block_opts = BlockBasedOptions::default();
block_opts.set_bloom_filter(14.0, false);
block_opts.set_whole_key_filtering(true);
block_opts.set_cache_index_and_filter_blocks(true);
opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(5));
opts.set_block_based_table_factory(&block_opts);
let db = DB::open(&opts, temp_dir.path())
.map_err(|e| Error::storage(format!("Failed to open in-memory RocksDB: {}", e)))?;
let db = Arc::new(db);
let key_manager = if config.encryption.enabled {
let km = KeyManager::from_source(&config.encryption.key_source)?;
Some(Arc::new(km))
} else {
None
};
let snapshot_manager = Arc::new(SnapshotManager::new_non_durable(Arc::clone(&db)));
let timestamp = Arc::new(RwLock::new(1));
debug!("Initializing BranchManager for in-memory storage");
let branch_manager = match BranchManager::new(Arc::clone(&db), Arc::clone(×tamp)) {
Ok(manager) => {
info!("BranchManager initialized successfully for in-memory storage");
Arc::new(RwLock::new(Some(Arc::new(manager))))
}
Err(e) => {
warn!(
"Failed to initialize BranchManager: {}. Branch operations will be unavailable.",
e
);
Arc::new(RwLock::new(None))
}
};
let wal = if config.storage.wal_enabled {
let sync_mode = WalSyncMode::Async; match WriteAheadLog::open(Arc::clone(&db), sync_mode) {
Ok(wal) => {
debug!("WAL initialized for in-memory storage");
Some(Arc::new(RwLock::new(wal)))
}
Err(e) => {
warn!("Failed to initialize WAL: {}", e);
None
}
}
} else {
None
};
let stats = Arc::new(DatabaseStats::new());
let statistics_cache = Arc::new(crate::storage::StatisticsCache::with_config(100, 30)?);
let mv_delta_tracker = Arc::new(super::MvDeltaTracker::new(Arc::clone(&db))?);
debug!("Delta tracker initialized for in-memory incremental MV refresh");
let trigger_registry = Arc::new(crate::sql::TriggerRegistry::new());
debug!("Trigger registry initialized (in-memory)");
let predicate_pushdown = Arc::new(PredicatePushdownManager::new(PushdownConfig::default()));
debug!("Predicate pushdown manager initialized (in-memory)");
let cpu_monitor = Arc::new(CpuMonitor::new());
let filter_delta_tracker = Arc::new(FilterIndexDeltaTracker::new(FilterIndexConfig::default()));
let speculative_filter_manager = Arc::new(SpeculativeFilterManager::new(SpeculativeConfig::default()));
let parallel_filter_engine = Arc::new(ParallelFilterEngine::new(ParallelFilterConfig::default()));
let consolidation_worker = None;
debug!("SMFI components initialized (in-memory)");
#[cfg(feature = "sync-experimental")]
let (change_log, node_id) = if config.sync.enabled && config.sync.change_log_enabled {
let node_id = if let Some(ref id_str) = config.sync.node_id {
uuid::Uuid::parse_str(id_str).map_err(|e| Error::config(format!("Invalid node_id UUID: {}", e)))?
} else {
uuid::Uuid::new_v4()
};
let cl = crate::sync::ChangeLogImpl::new(Arc::clone(&db))?;
debug!("Sync enabled (in-memory) with node_id={}", node_id);
(Some(Arc::new(RwLock::new(cl))), node_id)
} else {
(None, uuid::Uuid::new_v4())
};
let conflict_registry = Arc::new(super::conflict::WriteConflictRegistry::new());
let (version_gc, version_gc_worker) = Self::init_version_gc(
&config,
Arc::clone(&db),
Arc::clone(&snapshot_manager),
Arc::clone(&conflict_registry),
Arc::clone(&branch_manager),
Arc::clone(×tamp),
)?;
Ok(Self {
db: Arc::clone(&db),
config: config.clone(),
timestamp,
key_manager,
vector_indexes: Arc::new(VectorIndexManager::new()),
snapshot_manager,
conflict_registry,
version_gc,
version_gc_worker,
branch_manager,
wal,
stats,
statistics_cache,
is_replaying: Arc::new(AtomicBool::new(false)),
#[cfg(feature = "sync-experimental")]
change_log,
#[cfg(feature = "sync-experimental")]
node_id,
mv_delta_tracker,
current_branch: Arc::new(parking_lot::Mutex::new(None)),
trigger_registry,
predicate_pushdown,
filter_delta_tracker,
speculative_filter_manager,
parallel_filter_engine,
cpu_monitor,
consolidation_worker,
_temp_dir: Some(temp_dir),
row_counters: Arc::new(dashmap::DashMap::new()),
bulk_load_mode: Arc::new(AtomicBool::new(false)),
lockfree_engine: Arc::new(RwLock::new(None)),
dict_manager: Arc::new(DictionaryManager::new()),
art_index_manager: Arc::new(ArtIndexManager::new()),
row_cache: Arc::new(super::RowCache::new()),
data_bytes_written: Arc::new(AtomicU64::new(0)),
memory_limit_bytes: config.resource_quotas.memory_limit_per_user_mb * 1024 * 1024,
write_counter: Arc::new(AtomicU64::new(0)),
db_path: None, memory_write_options: Some(Self::memory_only_write_options()),
durable_commit_enabled: false,
group_committer: Arc::new(super::group_commit::GroupCommitter::new(
std::time::Duration::from_micros(config.storage.group_commit_window_us),
)),
schema_cache: Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new())),
constraints_cache: Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new())),
referencing_fk_cache: Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new())),
index_snapshots_on_close: Arc::new(AtomicBool::new(true)),
last_index_open_report: Arc::new(RwLock::new(None)),
})
}
pub fn config(&self) -> &Config {
&self.config
}
pub fn get_cached_schema(&self, table_name: &str) -> Option<crate::Schema> {
self.schema_cache.lock().get(table_name).cloned()
}
pub fn cache_schema(&self, table_name: &str, schema: crate::Schema) {
self.schema_cache.lock().insert(table_name.to_string(), schema);
}
pub fn invalidate_schema_cache(&self, table_name: &str) {
self.schema_cache.lock().remove(table_name);
}
pub fn clear_schema_cache(&self) {
self.schema_cache.lock().clear();
}
pub fn get_cached_table_constraints(&self, table_name: &str) -> Option<crate::sql::TableConstraints> {
self.constraints_cache.lock().get(table_name).cloned()
}
pub fn cache_table_constraints(&self, table_name: &str, constraints: crate::sql::TableConstraints) {
self.constraints_cache
.lock()
.insert(table_name.to_string(), constraints);
}
pub fn invalidate_table_constraints_cache(&self, table_name: &str) {
self.constraints_cache.lock().remove(table_name);
self.referencing_fk_cache.lock().clear();
}
pub fn get_cached_referencing_fks(&self, referenced_table: &str) -> Option<Vec<crate::sql::ForeignKeyConstraint>> {
self.referencing_fk_cache.lock().get(referenced_table).cloned()
}
pub fn cache_referencing_fks(&self, referenced_table: &str, constraints: Vec<crate::sql::ForeignKeyConstraint>) {
self.referencing_fk_cache
.lock()
.insert(referenced_table.to_string(), constraints);
}
pub fn clear_referencing_fk_cache(&self) {
self.referencing_fk_cache.lock().clear();
}
pub fn prewarm_schema_cache(&self) -> Result<()> {
let catalog = Catalog::new(self);
let tables = catalog.list_tables()?;
for table_name in &tables {
let _ = catalog.get_table_schema(table_name);
}
debug!("Pre-warmed schema cache with {} tables", tables.len());
Ok(())
}
pub fn stats(&self) -> &Arc<DatabaseStats> {
&self.stats
}
pub fn statistics_cache(&self) -> &Arc<crate::storage::StatisticsCache> {
&self.statistics_cache
}
pub fn art_indexes(&self) -> &Arc<ArtIndexManager> {
&self.art_index_manager
}
pub fn row_cache(&self) -> &Arc<super::RowCache> {
&self.row_cache
}
pub fn is_bulk_load_mode(&self) -> bool {
self.bulk_load_mode.load(Ordering::Acquire)
}
pub fn set_bulk_load_mode(&self, enabled: bool) {
self.bulk_load_mode.store(enabled, Ordering::Release);
if enabled {
tracing::info!("Bulk load mode ENABLED - skipping per-row metrics for faster INSERTs");
} else {
tracing::info!("Bulk load mode DISABLED - normal INSERT performance");
}
}
pub fn enable_lockfree_ingestion(&self, config: super::lockfree::LockFreeIngestionConfig) -> Result<()> {
let wal_path = if let Some(ref temp) = self._temp_dir {
temp.path().join("lockfree_wal")
} else {
std::path::PathBuf::from("data/lockfree_wal")
};
if let Err(e) = std::fs::create_dir_all(&wal_path) {
tracing::warn!("Could not create lock-free WAL directory: {}", e);
}
let safety_level = config.safety_level.clone();
let engine = super::lockfree::LockFreeIngestionEngine::new(config, &wal_path)
.map_err(|e| Error::storage(format!("Failed to create lock-free ingestion engine: {}", e)))?;
let db = Arc::clone(&self.db);
let key_manager = self.key_manager.clone();
engine.set_apply_callback(move |table, row_id, data| {
let key = format!("data:{}:{}", table, row_id);
let key_bytes = key.as_bytes();
match data {
Some(value) => {
let to_write = if let Some(ref km) = key_manager {
crypto::encrypt(km.key(), value).unwrap_or_else(|_| value.to_vec())
} else {
value.to_vec()
};
if let Err(e) = db.put(key_bytes, &to_write) {
tracing::error!("Lock-free apply callback failed for {}:{} - {}", table, row_id, e);
}
}
None => {
if let Err(e) = db.delete(key_bytes) {
tracing::error!("Lock-free delete callback failed for {}:{} - {}", table, row_id, e);
}
}
}
});
let mut guard = self.lockfree_engine.write();
*guard = Some(engine);
tracing::info!(
"Lock-free ingestion ENABLED with safety level: {}",
safety_level.description()
);
Ok(())
}
pub fn disable_lockfree_ingestion(&self) -> Result<()> {
let mut guard = self.lockfree_engine.write();
if let Some(ref engine) = *guard {
engine
.shutdown()
.map_err(|e| Error::storage(format!("Failed to shutdown lock-free engine: {}", e)))?;
}
*guard = None;
tracing::info!("Lock-free ingestion DISABLED");
Ok(())
}
pub fn is_lockfree_enabled(&self) -> bool {
self.lockfree_engine.read().is_some()
}
pub fn lockfree_stats(&self) -> Option<super::lockfree::IngestionStats> {
self.lockfree_engine.read().as_ref().map(|e| e.stats())
}
pub fn lockfree_safety_level(&self) -> Option<super::lockfree::IngestionSafetyLevel> {
self.lockfree_engine.read().as_ref().map(|e| e.safety_level().clone())
}
pub fn lockfree_begin(&self) -> Result<super::lockfree::TransactionHandle> {
let guard = self.lockfree_engine.read();
match guard.as_ref() {
Some(engine) => engine
.begin_transaction()
.map_err(|e| Error::storage(format!("Lock-free begin failed: {}", e))),
None => Err(Error::storage(
"Lock-free ingestion not enabled. Call enable_lockfree_ingestion first.",
)),
}
}
pub fn lockfree_generate_row_id(&self, table: &str) -> Result<u64> {
let guard = self.lockfree_engine.read();
match guard.as_ref() {
Some(engine) => Ok(engine.generate_row_id(table)),
None => Err(Error::storage("Lock-free ingestion not enabled")),
}
}
pub fn lockfree_insert(
&self,
handle: &super::lockfree::TransactionHandle,
table: &str,
row_id: u64,
data: &[u8],
) -> Result<()> {
let guard = self.lockfree_engine.read();
match guard.as_ref() {
Some(engine) => engine
.insert(handle, table, row_id, data)
.map_err(|e| Error::storage(format!("Lock-free insert failed: {}", e))),
None => Err(Error::storage("Lock-free ingestion not enabled")),
}
}
pub fn lockfree_commit(&self, handle: super::lockfree::TransactionHandle) -> Result<u64> {
let guard = self.lockfree_engine.read();
match guard.as_ref() {
Some(engine) => engine
.commit(handle)
.map_err(|e| Error::storage(format!("Lock-free commit failed: {}", e))),
None => Err(Error::storage("Lock-free ingestion not enabled")),
}
}
pub fn lockfree_abort(&self, handle: super::lockfree::TransactionHandle) -> Result<()> {
let guard = self.lockfree_engine.read();
match guard.as_ref() {
Some(engine) => engine
.abort(handle)
.map_err(|e| Error::storage(format!("Lock-free abort failed: {}", e))),
None => Err(Error::storage("Lock-free ingestion not enabled")),
}
}
pub fn lockfree_bulk_insert<I>(&self, table: &str, rows: I) -> Result<super::lockfree::BulkInsertResult>
where
I: IntoIterator<Item = Vec<u8>>,
{
let guard = self.lockfree_engine.read();
match guard.as_ref() {
Some(engine) => engine
.bulk_insert(table, rows)
.map_err(|e| Error::storage(format!("Lock-free bulk insert failed: {}", e))),
None => Err(Error::storage("Lock-free ingestion not enabled")),
}
}
pub fn lockfree_sync(&self) -> Result<()> {
let guard = self.lockfree_engine.read();
match guard.as_ref() {
Some(engine) => engine
.force_sync()
.map_err(|e| Error::storage(format!("Lock-free sync failed: {}", e))),
None => Err(Error::storage("Lock-free ingestion not enabled")),
}
}
pub fn lockfree_checkpoint(&self) -> Result<()> {
let guard = self.lockfree_engine.read();
match guard.as_ref() {
Some(engine) => engine
.checkpoint()
.map_err(|e| Error::storage(format!("Lock-free checkpoint failed: {}", e))),
None => Err(Error::storage("Lock-free ingestion not enabled")),
}
}
pub fn direct_bulk_load<I>(
&self,
table: &str,
rows: I,
batch_size: usize,
sync_at_end: bool,
) -> Result<DirectBulkLoadResult>
where
I: IntoIterator<Item = (u64, Vec<u8>)>,
{
use std::time::Instant;
let start = Instant::now();
let mut total_rows = 0u64;
let mut total_bytes = 0usize;
let mut batch = WriteBatch::default();
let mut batch_count = 0usize;
let mut max_row_id = 0u64;
let mut key_buf = String::with_capacity(64);
for (row_id, data) in rows {
key_buf.clear();
key_buf.push_str("data:");
key_buf.push_str(table);
key_buf.push(':');
key_buf.push_str(&row_id.to_string());
let to_write = if let Some(ref km) = self.key_manager {
crypto::encrypt(km.key(), &data)?
} else {
data
};
total_bytes += to_write.len();
batch.put(key_buf.as_bytes(), &to_write);
batch_count += 1;
total_rows += 1;
if row_id > max_row_id {
max_row_id = row_id;
}
if batch_count >= batch_size {
self.db
.write(batch)
.map_err(|e| Error::storage(format!("WriteBatch failed: {}", e)))?;
batch = WriteBatch::default();
batch_count = 0;
}
}
if batch_count > 0 {
self.db
.write(batch)
.map_err(|e| Error::storage(format!("Final WriteBatch failed: {}", e)))?;
}
if let Some(counter) = self.row_counters.get(table) {
let current = counter.load(Ordering::Relaxed);
if max_row_id >= current {
counter.store(max_row_id + 1, Ordering::Release);
}
} else {
self.row_counters
.insert(table.to_string(), std::sync::atomic::AtomicU64::new(max_row_id + 1));
}
let counter_key = format!("meta:counter:{}", table);
let counter_value = (max_row_id + 1).to_le_bytes();
self.db
.put(counter_key.as_bytes(), counter_value)
.map_err(|e| Error::storage(format!("Failed to persist counter: {}", e)))?;
if sync_at_end {
self.db
.flush()
.map_err(|e| Error::storage(format!("Flush failed: {}", e)))?;
}
let duration = start.elapsed();
let rows_per_sec = if duration.as_secs_f64() > 0.0 {
(total_rows as f64 / duration.as_secs_f64()) as u64
} else {
total_rows
};
Ok(DirectBulkLoadResult {
rows_loaded: total_rows,
bytes_written: total_bytes,
duration,
rows_per_sec,
max_row_id,
})
}
pub fn direct_bulk_load_auto_id<I>(
&self,
table: &str,
rows: I,
batch_size: usize,
sync_at_end: bool,
) -> Result<DirectBulkLoadResult>
where
I: IntoIterator<Item = Vec<u8>>,
{
let start_id = self
.row_counters
.entry(table.to_string())
.or_insert_with(|| std::sync::atomic::AtomicU64::new(1))
.load(Ordering::Acquire);
let mut current_id = start_id;
let rows_with_ids = rows.into_iter().map(move |data| {
let id = current_id;
current_id += 1;
(id, data)
});
self.direct_bulk_load(table, rows_with_ids, batch_size, sync_at_end)
}
pub fn get(&self, key: &Key) -> Result<Option<Vec<u8>>> {
let encrypted_data = self
.db
.get(key)
.map_err(|e| Error::storage(format!("Get failed: {}", e)))?;
match (encrypted_data, &self.key_manager) {
(Some(data), Some(km)) => {
let decrypted = crypto::decrypt(km.key(), &data)?;
Ok(Some(decrypted))
}
(Some(data), None) => Ok(Some(data)),
(None, _) => Ok(None),
}
}
pub fn put(&self, key: &Key, value: &[u8]) -> Result<()> {
if let Some(ref db_path) = self.db_path {
let count = self.write_counter.fetch_add(1, Ordering::Relaxed);
if count % 1000 == 0 {
Self::check_disk_space(db_path)?;
}
}
if self.memory_limit_bytes > 0 {
let write_size = (key.len() + value.len()) as u64;
let current = self.data_bytes_written.fetch_add(write_size, Ordering::Relaxed);
if current + write_size > self.memory_limit_bytes {
self.data_bytes_written.fetch_sub(write_size, Ordering::Relaxed);
return Err(Error::storage(format!(
"Memory limit exceeded ({} MB). Increase resource_quotas.memory_limit_per_user_mb or use disk-backed mode.",
self.memory_limit_bytes / (1024 * 1024)
)));
}
}
if let Some(km) = &self.key_manager {
let data = crypto::encrypt(km.key(), value)?;
if let Some(opts) = &self.memory_write_options {
self.db
.put_opt(key, data, opts)
.map_err(|e| Error::storage(format!("Put failed: {}", e)))
} else {
self.db
.put(key, data)
.map_err(|e| Error::storage(format!("Put failed: {}", e)))
}
} else if let Some(opts) = &self.memory_write_options {
self.db
.put_opt(key, value, opts)
.map_err(|e| Error::storage(format!("Put failed: {}", e)))
} else {
self.db
.put(key, value)
.map_err(|e| Error::storage(format!("Put failed: {}", e)))
}
}
pub fn delete(&self, key: &Key) -> Result<()> {
if !self.is_replaying.load(Ordering::Acquire) {
let key_str = std::str::from_utf8(key).unwrap_or("");
let is_metadata_key = key_str.starts_with("meta:");
if !is_metadata_key {
if let Some(wal) = &self.wal {
let wal = wal.read();
let table_name = Self::extract_table_from_key(key);
wal.append(WalOperation::Delete {
table: table_name,
key: key.clone(),
})?;
}
}
}
if let Some(opts) = &self.memory_write_options {
self.db
.delete_opt(key, opts)
.map_err(|e| Error::storage(format!("Delete failed: {}", e)))
} else {
self.db
.delete(key)
.map_err(|e| Error::storage(format!("Delete failed: {}", e)))
}
}
pub fn log_data_insert(&self, table_name: &str, key: &[u8], tuple_data: &[u8]) -> Result<()> {
if self.is_replaying.load(Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::Insert {
table: table_name.to_string(),
key: key.to_vec(),
tuple: tuple_data.to_vec(),
})?;
}
Ok(())
}
pub fn log_data_insert_nosync(&self, table_name: &str, key: &[u8], tuple_data: &[u8]) -> Result<()> {
if self.is_replaying.load(Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append_nosync(WalOperation::Insert {
table: table_name.to_string(),
key: key.to_vec(),
tuple: tuple_data.to_vec(),
})?;
}
Ok(())
}
pub fn fast_dml_requires_logical_wal(&self) -> bool {
if !self.is_wal_enabled() {
return false;
}
if self.config.storage.logical_wal_per_statement {
return true;
}
#[cfg(feature = "ha-tier1")]
{
use crate::replication::ha_state::{ha_state, HARole};
if ha_state().get_role() == HARole::Primary {
return true;
}
}
false
}
pub fn log_data_update(&self, table_name: &str, key: &[u8], tuple_data: &[u8]) -> Result<()> {
if self.is_replaying.load(Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::Update {
table: table_name.to_string(),
key: key.to_vec(),
tuple: tuple_data.to_vec(),
})?;
}
Ok(())
}
pub fn log_data_delete(&self, table_name: &str, key: &[u8]) -> Result<()> {
if self.is_replaying.load(Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::Delete {
table: table_name.to_string(),
key: key.to_vec(),
})?;
}
Ok(())
}
pub fn log_data_update_nosync(&self, table_name: &str, key: &[u8], tuple_data: &[u8]) -> Result<()> {
if self.is_replaying.load(Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append_nosync(WalOperation::Update {
table: table_name.to_string(),
key: key.to_vec(),
tuple: tuple_data.to_vec(),
})?;
}
Ok(())
}
pub fn log_data_delete_nosync(&self, table_name: &str, key: &[u8]) -> Result<()> {
if self.is_replaying.load(Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append_nosync(WalOperation::Delete {
table: table_name.to_string(),
key: key.to_vec(),
})?;
}
Ok(())
}
fn put_internal(&self, key: &[u8], value: &[u8]) -> Result<()> {
let data = if let Some(km) = &self.key_manager {
crypto::encrypt(km.key(), value)?
} else {
value.to_vec()
};
if let Some(opts) = &self.memory_write_options {
self.db
.put_opt(key, data, opts)
.map_err(|e| Error::storage(format!("Internal put failed: {}", e)))
} else {
self.db
.put(key, data)
.map_err(|e| Error::storage(format!("Internal put failed: {}", e)))
}
}
fn get_internal(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
let encrypted_data = self
.db
.get(key)
.map_err(|e| Error::storage(format!("Internal get failed: {}", e)))?;
match (encrypted_data, &self.key_manager) {
(Some(data), Some(km)) => {
let decrypted = crypto::decrypt(km.key(), &data)?;
Ok(Some(decrypted))
}
(Some(data), None) => Ok(Some(data)),
(None, _) => Ok(None),
}
}
fn decrypt_value(&self, value: &[u8]) -> Result<Vec<u8>> {
if let Some(km) = &self.key_manager {
crypto::decrypt(km.key(), value)
} else {
Ok(value.to_vec())
}
}
pub fn begin_transaction(&self) -> Result<Transaction> {
let snapshot_id = self.next_timestamp();
let mut txn = Transaction::new(Arc::clone(&self.db), snapshot_id, Arc::clone(&self.snapshot_manager))?;
txn.set_versioning_enabled(self.config.storage.time_travel_enabled);
txn.set_rocksdb_wal_enabled(!self.config.storage.memory_only);
txn.set_sync_commit(self.statement_durability_required());
txn.set_group_committer(self.group_committer());
txn.set_row_cache(Arc::clone(self.row_cache()));
txn.set_conflict_registry(self.conflict_registry(), true);
Ok(txn)
}
pub fn begin_autocommit_transaction(&self) -> Result<Transaction> {
let snapshot_id = self.next_timestamp();
let mut txn = Transaction::new(Arc::clone(&self.db), snapshot_id, Arc::clone(&self.snapshot_manager))?;
txn.set_versioning_enabled(self.config.storage.time_travel_enabled);
txn.set_rocksdb_wal_enabled(!self.config.storage.memory_only);
txn.set_sync_commit(self.statement_durability_required());
txn.set_group_committer(self.group_committer());
txn.set_row_cache(Arc::clone(self.row_cache()));
txn.set_conflict_registry(self.conflict_registry(), false);
Ok(txn)
}
pub fn durable_autocommit_barrier(&self) -> Result<()> {
if self.statement_durability_required() {
self.group_committer
.wait_durable(|| self.db.flush_wal(true).map_err(|e| e.to_string()))
.map_err(|e| Error::storage(format!("WAL fsync failed: {}", e)))?;
}
Ok(())
}
fn statement_durability_required(&self) -> bool {
!self.config.storage.memory_only && Self::synchronous_commit_override().unwrap_or(self.durable_commit_enabled)
}
pub fn group_committer(&self) -> Arc<super::group_commit::GroupCommitter> {
Arc::clone(&self.group_committer)
}
pub fn synchronous_commit_override() -> Option<bool> {
SYNCHRONOUS_COMMIT_OVERRIDE.with(|cell| cell.get())
}
pub fn synchronous_commit_override_guard(value: Option<bool>) -> SynchronousCommitOverrideGuard {
let previous = SYNCHRONOUS_COMMIT_OVERRIDE.with(|cell| cell.replace(value));
SynchronousCommitOverrideGuard { previous }
}
pub fn conflict_registry(&self) -> Arc<super::conflict::WriteConflictRegistry> {
Arc::clone(&self.conflict_registry)
}
fn init_version_gc(
config: &Config,
db: Arc<DB>,
snapshot_manager: Arc<SnapshotManager>,
conflict_registry: Arc<super::conflict::WriteConflictRegistry>,
branch_manager: Arc<RwLock<Option<Arc<BranchManager>>>>,
timestamp: Arc<RwLock<u64>>,
) -> Result<(
Arc<super::version_gc::VersionGc>,
Option<super::version_gc::VersionGcWorker>,
)> {
let retention_secs = config.storage.version_retention_secs()?;
let interval_secs = config.storage.effective_version_gc_interval_secs()?;
let gc_config = super::version_gc::VersionGcConfig {
retention_secs,
interval_secs,
max_versions_per_cycle: config.storage.version_gc_max_per_cycle,
};
let version_gc = Arc::new(super::version_gc::VersionGc::new(
db,
snapshot_manager,
conflict_registry,
branch_manager,
timestamp,
gc_config,
));
let worker = if interval_secs > 0 {
info!(
interval_secs,
retention_secs = retention_secs.unwrap_or(0),
"starting MVCC version-GC worker"
);
Some(super::version_gc::VersionGcWorker::start(
Arc::clone(&version_gc),
interval_secs,
))
} else {
None
};
Ok((version_gc, worker))
}
pub fn vacuum_versions(&self) -> Result<u64> {
self.version_gc.vacuum()
}
pub fn version_gc_cycle(&self) -> Result<super::version_gc::VersionGcCycleStats> {
self.version_gc.run_cycle()
}
pub fn version_storage_stats(&self) -> Result<super::version_gc::VersionStorageStats> {
super::version_gc::version_storage_stats(&self.db)
}
pub fn pin_historical_snapshot(&self, snapshot_ts: u64) -> Result<super::conflict::GcPinGuard> {
let guard = super::conflict::GcPinGuard::new(self.conflict_registry(), snapshot_ts);
self.snapshot_manager.check_gc_horizon(snapshot_ts)?;
Ok(guard)
}
pub fn next_timestamp(&self) -> u64 {
let mut ts = self.timestamp.write();
*ts += 1;
*ts
}
pub fn next_commit_timestamp(&self, announce_inflight: bool) -> u64 {
let mut ts = self.timestamp.write();
*ts += 1;
if announce_inflight {
self.conflict_registry.begin_commit(*ts);
}
*ts
}
pub fn insert_tuple(&self, table_name: &str, tuple: Tuple) -> Result<u64> {
if self.is_branch_active() {
return self.insert_tuple_branch_aware(table_name, tuple);
}
if self.config.storage.time_travel_enabled {
self.insert_tuple_versioned(table_name, tuple)
} else {
let catalog = Catalog::new(self);
let row_id = self.next_row_id_volatile(table_name);
let schema = catalog.get_table_schema(table_name)?;
let bulk_mode = self.is_bulk_load_mode();
let stored_tuple =
self.transform_tuple_for_column_storage_opts(table_name, row_id, &tuple, &schema, true)?;
let value = bincode::serialize(&stored_tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?;
let logical_value = if schema_uses_column_storage(&schema) {
bincode::serialize(&tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
} else {
value.clone()
};
let key = Self::build_data_key(table_name, row_id);
self.put(&key, &value)?;
self.flush_row_counter(table_name)?;
if self.fast_dml_requires_logical_wal() {
self.log_data_insert(table_name, &key, &logical_value)?;
} else {
self.log_data_insert_nosync(table_name, &key, &logical_value)?;
}
{
let mut col_values = std::collections::HashMap::new();
for (i, col) in schema.columns.iter().enumerate() {
if let Some(v) = tuple.values.get(i) {
col_values.insert(col.name.clone(), v.clone());
}
}
if let Err(e) = self.art_index_manager.on_insert(table_name, row_id, &col_values) {
tracing::debug!("ART index insert for table '{}': {}", table_name, e);
}
}
let _ = self.vector_indexes.on_row_insert(table_name, row_id, &schema, &tuple);
if !bulk_mode {
if let Err(e) = self.mv_delta_tracker.record_insert(table_name, row_id, tuple.clone()) {
tracing::warn!("Failed to record insert delta for table '{}': {}", table_name, e);
}
self.filter_delta_tracker.on_insert(table_name, row_id, &tuple, &schema);
for (i, col) in schema.columns.iter().enumerate() {
if let Some(value) = tuple.values.get(i) {
self.speculative_filter_manager.on_insert(table_name, &col.name, value);
}
}
}
Ok(row_id)
}
}
pub fn purge_table_data(&self, table_name: &str) -> Result<u64> {
let data_prefix = format!("data:{}:", table_name);
let prefix_bytes = data_prefix.as_bytes();
let mut keys_to_delete = Vec::new();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(false); let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break; }
keys_to_delete.push(key.to_vec());
}
let removed = keys_to_delete.len() as u64;
for key in keys_to_delete {
self.delete(&key)?;
}
ColumnarStore::purge_table_sidecars(&self.db, table_name)?;
if removed > 0 {
tracing::debug!(
"purge_table_data: removed {} orphaned rows for '{}'",
removed,
table_name
);
}
Ok(removed)
}
pub fn scan_table(&self, table_name: &str) -> Result<Vec<Tuple>> {
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
self.scan_table_with_schema(table_name, &schema)
}
pub fn scan_table_with_schema(&self, table_name: &str, schema: &crate::Schema) -> Result<Vec<Tuple>> {
self.scan_table_with_schema_opt(table_name, schema, RowDecodeHint::Full)
}
pub fn scan_table_with_schema_prefix(
&self,
table_name: &str,
schema: &crate::Schema,
prefix_len: usize,
) -> Result<Vec<Tuple>> {
if prefix_len >= schema.columns.len() {
self.scan_table_with_schema_opt(table_name, schema, RowDecodeHint::Full)
} else {
self.scan_table_with_schema_opt(table_name, schema, RowDecodeHint::Prefix(prefix_len))
}
}
pub fn scan_table_with_schema_columns(
&self,
table_name: &str,
schema: &crate::Schema,
columns: &[usize],
) -> Result<Vec<Tuple>> {
if columns.is_empty() {
self.scan_table_with_schema_opt(table_name, schema, RowDecodeHint::Prefix(0))
} else if columns.len() >= schema.columns.len() {
self.scan_table_with_schema_opt(table_name, schema, RowDecodeHint::Full)
} else {
self.scan_table_with_schema_opt(table_name, schema, RowDecodeHint::Columns(columns))
}
}
pub(crate) fn scan_table_with_schema_columns_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
columns: Option<&[usize]>,
predicates: &[AnalyzedPredicate],
) -> Result<Option<Vec<Tuple>>> {
let branch_name = self.current_branch.lock().clone();
if branch_name.is_some() && branch_name.as_deref() != Some("main") {
return Ok(None);
}
let filter_predicates: Vec<FilterPredicate> = predicates.iter().filter_map(columnar_filter_predicate).collect();
if filter_predicates.len() != predicates.len() {
return Ok(None);
}
let mut requested: Vec<usize> = columns
.map(|cols| cols.to_vec())
.unwrap_or_else(|| (0..schema.columns.len()).collect());
requested.extend(predicates.iter().map(|predicate| predicate.column_index));
requested.sort_unstable();
requested.dedup();
for &idx in &requested {
let Some(column) = schema.columns.get(idx) else {
return Err(Error::storage(format!(
"Column index {} out of bounds for {}",
idx, table_name
)));
};
if column.storage_mode != ColumnStorageMode::Default {
return Ok(None);
}
}
if let Some(candidates) = integer_filter_candidates(schema, &filter_predicates) {
if let Some(tuples) =
self.scan_table_with_schema_columns_integer_filtered(table_name, schema, &requested, &candidates)?
{
return Ok(Some(tuples));
}
}
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut tuples = Vec::new();
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let mut tuple = self.decode_rowstore_columns(&raw_value, &requested, schema.columns.len())?;
if !row_tuple_matches_filters(&tuple, &filter_predicates) {
continue;
}
if let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) {
tuple.row_id = Some(row_id);
}
tuples.push(tuple);
}
Ok(Some(tuples))
}
fn scan_table_with_schema_columns_integer_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
requested: &[usize],
filters: &[IntegerFilterCandidate],
) -> Result<Option<Vec<Tuple>>> {
if filters.is_empty() {
return Ok(None);
}
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut tuples = Vec::new();
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let decrypted;
let row_bytes = if let Some(km) = &self.key_manager {
decrypted = crypto::decrypt(km.key(), raw_value.as_ref())?;
decrypted.as_slice()
} else {
raw_value.as_ref()
};
let mut matched = true;
for filter in filters {
let Some(decoded) =
crate::storage::prefix_decode::decode_tuple_numeric_column_value(row_bytes, filter.column_index)
else {
return Ok(None);
};
let Some(filter_matches) = decoded_integer_matches_filter(decoded, *filter) else {
return Ok(None);
};
if !filter_matches {
matched = false;
break;
}
}
if !matched {
continue;
}
let mut tuple = self.decode_rowstore_columns(row_bytes, requested, schema.columns.len())?;
if let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) {
tuple.row_id = Some(row_id);
}
tuples.push(tuple);
}
Ok(Some(tuples))
}
pub(crate) fn scan_table_with_schema_projected_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
projection: &[usize],
predicates: &[AnalyzedPredicate],
) -> Result<Option<Vec<Tuple>>> {
if projection.is_empty() {
return Ok(None);
}
let branch_name = self.current_branch.lock().clone();
if branch_name.is_some() && branch_name.as_deref() != Some("main") {
return Ok(None);
}
let filter_predicates: Vec<FilterPredicate> = predicates.iter().filter_map(columnar_filter_predicate).collect();
if filter_predicates.len() != predicates.len() {
return Ok(None);
}
let mut requested: Vec<usize> = projection.to_vec();
requested.extend(predicates.iter().map(|predicate| predicate.column_index));
requested.sort_unstable();
requested.dedup();
let mut all_default = true;
for &idx in &requested {
let Some(column) = schema.columns.get(idx) else {
return Err(Error::storage(format!(
"Column index {} out of bounds for {}",
idx, table_name
)));
};
if column.storage_mode != ColumnStorageMode::Default {
all_default = false;
break;
}
}
if !all_default {
return self.scan_table_with_schema_mixed_projected_filtered(
table_name,
schema,
projection,
&filter_predicates,
);
}
if filter_predicates.len() > 1 {
if let Some(candidates) = integer_filter_candidates(schema, &filter_predicates) {
if let Some(tuples) =
self.scan_table_with_schema_projected_integer_filters(table_name, schema, projection, &candidates)?
{
return Ok(Some(tuples));
}
}
}
if let Some(candidate) = integer_filter_candidate(schema, &filter_predicates) {
if let Some(tuples) =
self.scan_table_with_schema_projected_integer_filtered(table_name, schema, projection, candidate)?
{
return Ok(Some(tuples));
}
}
if let Some(candidate) = string_eq_filter_candidate(schema, &filter_predicates) {
if let Some(tuples) =
self.scan_table_with_schema_projected_string_eq_filtered(table_name, schema, projection, candidate)?
{
return Ok(Some(tuples));
}
}
let requested_pos = |column_index: usize| -> Result<usize> {
requested.binary_search(&column_index).map_err(|_| {
Error::storage(format!(
"Column index {} missing from projected filtered scan decode set for {}",
column_index, table_name
))
})
};
let projection_positions: Vec<usize> = projection
.iter()
.map(|&idx| requested_pos(idx))
.collect::<Result<Vec<_>>>()?;
let filter_positions: Vec<usize> = predicates
.iter()
.map(|predicate| requested_pos(predicate.column_index))
.collect::<Result<Vec<_>>>()?;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut values = Vec::with_capacity(requested.len());
let mut tuples = Vec::new();
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
self.decode_rowstore_column_values_into(&raw_value, &requested, schema.columns.len(), &mut values)?;
if !row_values_match_filters(&values, &filter_predicates, &filter_positions) {
continue;
}
let mut projected_values = Vec::with_capacity(projection_positions.len());
for &pos in &projection_positions {
if let Some(value) = values.get_mut(pos) {
projected_values.push(std::mem::replace(value, Value::Null));
} else {
projected_values.push(Value::Null);
}
}
tuples.push(Tuple::new(projected_values));
}
Ok(Some(tuples))
}
fn scan_table_with_schema_projected_string_eq_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
projection: &[usize],
filter: StringEqFilterCandidate<'_>,
) -> Result<Option<Vec<Tuple>>> {
let mut output_requested: Vec<usize> = projection.to_vec();
output_requested.sort_unstable();
output_requested.dedup();
let output_pos = |column_index: usize| -> Result<usize> {
output_requested.binary_search(&column_index).map_err(|_| {
Error::storage(format!(
"Column index {} missing from projected string filtered scan decode set for {}",
column_index, table_name
))
})
};
let projection_positions: Vec<usize> = projection
.iter()
.map(|&idx| output_pos(idx))
.collect::<Result<Vec<_>>>()?;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut values = Vec::with_capacity(output_requested.len());
let mut tuples = Vec::new();
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let decrypted;
let row_bytes = if let Some(km) = &self.key_manager {
decrypted = crypto::decrypt(km.key(), raw_value.as_ref())?;
decrypted.as_slice()
} else {
raw_value.as_ref()
};
let Some(matches) =
crate::storage::prefix_decode::tuple_string_column_eq(row_bytes, filter.column_index, filter.value)
else {
return Ok(None);
};
if !matches {
continue;
}
crate::storage::prefix_decode::decode_tuple_column_values_into(
row_bytes,
&output_requested,
schema.columns.len(),
&mut values,
)
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
let mut projected_values = Vec::with_capacity(projection_positions.len());
for &pos in &projection_positions {
if let Some(value) = values.get_mut(pos) {
projected_values.push(std::mem::replace(value, Value::Null));
} else {
projected_values.push(Value::Null);
}
}
tuples.push(Tuple::new(projected_values));
}
Ok(Some(tuples))
}
fn scan_table_with_schema_mixed_projected_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
projection: &[usize],
filter_predicates: &[FilterPredicate],
) -> Result<Option<Vec<Tuple>>> {
if filter_predicates.is_empty() {
return Ok(None);
}
let mut columnar_requested: Vec<usize> = Vec::new();
for predicate in filter_predicates {
let Some(column) = schema.columns.get(predicate.column_index) else {
return Err(Error::storage(format!(
"Column index {} out of bounds for {}",
predicate.column_index, table_name
)));
};
if column.storage_mode != ColumnStorageMode::Columnar {
return Ok(None);
}
columnar_requested.push(predicate.column_index);
}
let mut default_requested = Vec::new();
for &idx in projection {
let Some(column) = schema.columns.get(idx) else {
return Err(Error::storage(format!(
"Column index {} out of bounds for {}",
idx, table_name
)));
};
match column.storage_mode {
ColumnStorageMode::Default => default_requested.push(idx),
ColumnStorageMode::Columnar => columnar_requested.push(idx),
ColumnStorageMode::Dictionary | ColumnStorageMode::ContentAddressed => return Ok(None),
}
}
default_requested.sort_unstable();
default_requested.dedup();
columnar_requested.sort_unstable();
columnar_requested.dedup();
if columnar_requested.is_empty() {
return Ok(None);
}
enum MixedProjectionSource {
Default(usize),
Columnar(usize),
}
let projection_sources: Vec<MixedProjectionSource> = projection
.iter()
.map(|&idx| {
let column = schema
.columns
.get(idx)
.ok_or_else(|| Error::storage(format!("Column index {} out of bounds for {}", idx, table_name)))?;
match column.storage_mode {
ColumnStorageMode::Default => default_requested
.binary_search(&idx)
.map(MixedProjectionSource::Default)
.map_err(|_| {
Error::storage(format!(
"Column index {} missing from mixed projected scan decode set for {}",
idx, table_name
))
}),
ColumnStorageMode::Columnar => Ok(MixedProjectionSource::Columnar(idx)),
ColumnStorageMode::Dictionary | ColumnStorageMode::ContentAddressed => {
Err(Error::storage(format!(
"Column index {} uses unsupported mixed projected scan storage on {}",
idx, table_name
)))
}
}
})
.collect::<Result<Vec<_>>>()?;
let pruned = compute_zone_pruned_batches(&self.db, table_name, schema, filter_predicates);
let mut column_batches: HashMap<usize, ColumnarBatchIndex> = HashMap::with_capacity(columnar_requested.len());
for &idx in &columnar_requested {
let column = &schema.columns[idx];
column_batches.insert(
idx,
load_columnar_batch_index(&self.db, table_name, &column.name, &pruned)?,
);
}
if default_requested.is_empty() {
if let Some(presence) = self.columnar_live_presence(table_name) {
let mut tuples = Vec::new();
for (batch_id, batch_presence) in &presence {
if pruned.contains(batch_id) {
continue;
}
for offset in batch_presence.iter_live() {
if !columnar_row_matches_filters(&column_batches, *batch_id, offset, filter_predicates) {
continue;
}
let mut projected_values = Vec::with_capacity(projection_sources.len());
for source in &projection_sources {
match source {
MixedProjectionSource::Default(_) => projected_values.push(Value::Null),
MixedProjectionSource::Columnar(idx) => {
projected_values.push(
columnar_batch_value(&column_batches, *idx, *batch_id, offset)
.unwrap_or(Value::Null),
);
}
}
}
tuples.push(Tuple::new(projected_values));
}
}
return Ok(Some(tuples));
}
}
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut default_values = Vec::with_capacity(default_requested.len());
let mut tuples = Vec::new();
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
let batch_id = row_id / BATCH_SIZE as u64;
let offset = (row_id % BATCH_SIZE as u64) as usize;
if pruned.contains(&batch_id) {
continue;
}
if !columnar_row_matches_filters(&column_batches, batch_id, offset, filter_predicates) {
continue;
}
if !default_requested.is_empty() {
if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), raw_value.as_ref())?;
crate::storage::prefix_decode::decode_tuple_column_values_into(
&decrypted,
&default_requested,
schema.columns.len(),
&mut default_values,
)
} else {
crate::storage::prefix_decode::decode_tuple_column_values_into(
raw_value.as_ref(),
&default_requested,
schema.columns.len(),
&mut default_values,
)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
} else {
default_values.clear();
}
let mut projected_values = Vec::with_capacity(projection_sources.len());
for source in &projection_sources {
match source {
MixedProjectionSource::Default(pos) => {
projected_values.push(default_values.get(*pos).cloned().unwrap_or(Value::Null));
}
MixedProjectionSource::Columnar(idx) => {
projected_values
.push(columnar_batch_value(&column_batches, *idx, batch_id, offset).unwrap_or(Value::Null));
}
}
}
tuples.push(Tuple::new(projected_values));
}
Ok(Some(tuples))
}
fn scan_table_with_schema_projected_integer_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
projection: &[usize],
filter: IntegerFilterCandidate,
) -> Result<Option<Vec<Tuple>>> {
let mut output_requested: Vec<usize> = projection.to_vec();
output_requested.sort_unstable();
output_requested.dedup();
let output_pos = |column_index: usize| -> Result<usize> {
output_requested.binary_search(&column_index).map_err(|_| {
Error::storage(format!(
"Column index {} missing from projected integer filtered scan decode set for {}",
column_index, table_name
))
})
};
let projection_positions: Vec<usize> = projection
.iter()
.map(|&idx| output_pos(idx))
.collect::<Result<Vec<_>>>()?;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut values = Vec::with_capacity(output_requested.len());
let mut tuples = Vec::new();
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let decrypted;
let row_bytes = if let Some(km) = &self.key_manager {
decrypted = crypto::decrypt(km.key(), raw_value.as_ref())?;
decrypted.as_slice()
} else {
raw_value.as_ref()
};
let Some(decoded) =
crate::storage::prefix_decode::decode_tuple_numeric_column_value(row_bytes, filter.column_index)
else {
return Ok(None);
};
let Some(matches) = decoded_integer_matches_filter(decoded, filter) else {
return Ok(None);
};
if !matches {
continue;
}
crate::storage::prefix_decode::decode_tuple_column_values_into(
row_bytes,
&output_requested,
schema.columns.len(),
&mut values,
)
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
let mut projected_values = Vec::with_capacity(projection_positions.len());
for &pos in &projection_positions {
if let Some(value) = values.get_mut(pos) {
projected_values.push(std::mem::replace(value, Value::Null));
} else {
projected_values.push(Value::Null);
}
}
tuples.push(Tuple::new(projected_values));
}
Ok(Some(tuples))
}
fn scan_table_with_schema_projected_integer_filters(
&self,
table_name: &str,
schema: &crate::Schema,
projection: &[usize],
filters: &[IntegerFilterCandidate],
) -> Result<Option<Vec<Tuple>>> {
if filters.is_empty() {
return Ok(None);
}
let mut output_requested: Vec<usize> = projection.to_vec();
output_requested.sort_unstable();
output_requested.dedup();
let output_pos = |column_index: usize| -> Result<usize> {
output_requested.binary_search(&column_index).map_err(|_| {
Error::storage(format!(
"Column index {} missing from projected integer filtered scan decode set for {}",
column_index, table_name
))
})
};
let projection_positions: Vec<usize> = projection
.iter()
.map(|&idx| output_pos(idx))
.collect::<Result<Vec<_>>>()?;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut values = Vec::with_capacity(output_requested.len());
let mut tuples = Vec::new();
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let decrypted;
let row_bytes = if let Some(km) = &self.key_manager {
decrypted = crypto::decrypt(km.key(), raw_value.as_ref())?;
decrypted.as_slice()
} else {
raw_value.as_ref()
};
let mut matched = true;
for filter in filters {
let Some(decoded) =
crate::storage::prefix_decode::decode_tuple_numeric_column_value(row_bytes, filter.column_index)
else {
return Ok(None);
};
let Some(filter_matches) = decoded_integer_matches_filter(decoded, *filter) else {
return Ok(None);
};
if !filter_matches {
matched = false;
break;
}
}
if !matched {
continue;
}
crate::storage::prefix_decode::decode_tuple_column_values_into(
row_bytes,
&output_requested,
schema.columns.len(),
&mut values,
)
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
let mut projected_values = Vec::with_capacity(projection_positions.len());
for &pos in &projection_positions {
if let Some(value) = values.get_mut(pos) {
projected_values.push(std::mem::replace(value, Value::Null));
} else {
projected_values.push(Value::Null);
}
}
tuples.push(Tuple::new(projected_values));
}
Ok(Some(tuples))
}
fn scan_table_topk_single_int_projected_columns(
&self,
table_name: &str,
schema: &crate::Schema,
output_columns: &[usize],
sort_column: usize,
asc: bool,
k: usize,
) -> Result<Option<Vec<Tuple>>> {
let mut output_decode_columns: Vec<usize> = output_columns
.iter()
.copied()
.filter(|&column| column != sort_column)
.collect();
output_decode_columns.sort_unstable();
output_decode_columns.dedup();
enum IntTopKOutputSource {
SortKey,
Decoded(usize),
}
let output_source = |column_index: usize| -> Result<IntTopKOutputSource> {
if column_index == sort_column {
return Ok(IntTopKOutputSource::SortKey);
}
output_decode_columns
.binary_search(&column_index)
.map(IntTopKOutputSource::Decoded)
.map_err(|_| {
Error::storage(format!(
"Column index {} missing from integer Top-N output set for {}",
column_index, table_name
))
})
};
let output_sources: Vec<IntTopKOutputSource> = output_columns
.iter()
.map(|&idx| output_source(idx))
.collect::<Result<Vec<_>>>()?;
let sort_column_type = &schema
.columns
.get(sort_column)
.ok_or_else(|| Error::storage(format!("Column index {} out of bounds for {}", sort_column, table_name)))?
.data_type;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut heap = std::collections::BinaryHeap::with_capacity(k.saturating_add(1));
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let sort_value = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
crate::storage::prefix_decode::decode_tuple_numeric_column_value(&decrypted, sort_column)
} else {
crate::storage::prefix_decode::decode_tuple_numeric_column_value(&raw_value, sort_column)
};
let Some(crate::storage::prefix_decode::DecodedNumericValue::Int(sort_key)) = sort_value else {
return Ok(None);
};
let replace = if heap.len() < k {
true
} else {
heap.peek()
.map(|top: &RowIntTopKEntry| {
compare_int_to_topk_key(sort_key, top.key, asc) == std::cmp::Ordering::Less
})
.unwrap_or(false)
};
if replace {
if heap.len() >= k {
heap.pop();
}
heap.push(RowIntTopKEntry {
key: sort_key,
raw_value,
row_id: Self::parse_row_id_after_prefix(&key, prefix_bytes.len()),
asc,
});
}
}
let entries = heap.into_sorted_vec();
let mut tuples = Vec::with_capacity(entries.len());
for entry in entries {
let output_values =
self.decode_rowstore_column_values(&entry.raw_value, &output_decode_columns, schema.columns.len())?;
let mut projected_values = Vec::with_capacity(output_sources.len());
for source in &output_sources {
match source {
IntTopKOutputSource::SortKey => {
let value = int_value_for_data_type(sort_column_type, entry.key).ok_or_else(|| {
Error::storage(format!(
"Integer Top-N sort key out of range for column {} on {}",
sort_column, table_name
))
})?;
projected_values.push(value);
}
IntTopKOutputSource::Decoded(pos) => {
projected_values.push(output_values.get(*pos).cloned().unwrap_or(Value::Null));
}
}
}
let mut tuple = Tuple::new(projected_values);
tuple.row_id = entry.row_id;
tuples.push(tuple);
}
Ok(Some(tuples))
}
pub(crate) fn scan_table_topk_projected_columns(
&self,
table_name: &str,
schema: &crate::Schema,
output_columns: &[usize],
sort_columns: &[usize],
asc: &[bool],
k: usize,
) -> Result<Option<Vec<Tuple>>> {
if k == 0 {
return Ok(Some(Vec::new()));
}
if sort_columns.is_empty() || sort_columns.len() != asc.len() {
return Ok(None);
}
let branch_name = self.current_branch.lock().clone();
if branch_name.is_some() && branch_name.as_deref() != Some("main") {
return Ok(None);
}
let mut requested: Vec<usize> = output_columns.to_vec();
requested.extend_from_slice(sort_columns);
requested.sort_unstable();
requested.dedup();
for &idx in &requested {
let Some(column) = schema.columns.get(idx) else {
return Err(Error::storage(format!(
"Column index {} out of bounds for {}",
idx, table_name
)));
};
if column.storage_mode != ColumnStorageMode::Default {
return Ok(None);
}
}
if sort_columns.len() == 1
&& primitive_integer_data_type(&schema.columns[sort_columns[0]].data_type)
&& output_columns.len() < schema.columns.len()
{
if let Some(tuples) = self.scan_table_topk_single_int_projected_columns(
table_name,
schema,
output_columns,
sort_columns[0],
asc[0],
k,
)? {
return Ok(Some(tuples));
}
}
let requested_pos = |column_index: usize| -> Result<usize> {
requested.binary_search(&column_index).map_err(|_| {
Error::storage(format!(
"Column index {} missing from Top-N decode set for {}",
column_index, table_name
))
})
};
let output_positions: Vec<usize> = output_columns
.iter()
.map(|&idx| requested_pos(idx))
.collect::<Result<Vec<_>>>()?;
let sort_positions: Vec<usize> = sort_columns
.iter()
.map(|&idx| requested_pos(idx))
.collect::<Result<Vec<_>>>()?;
let asc = Arc::new(asc.to_vec());
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut heap = std::collections::BinaryHeap::with_capacity(k.saturating_add(1));
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
self.decode_rowstore_column_values_into(&raw_value, &requested, schema.columns.len(), &mut values)?;
let replace = if heap.len() < k {
true
} else {
heap.peek()
.map(|top: &RowTopKEntry| {
compare_values_to_topk_key(&values, &sort_positions, &top.key, &asc) == std::cmp::Ordering::Less
})
.unwrap_or(false)
};
if replace {
if heap.len() >= k {
heap.pop();
}
heap.push(RowTopKEntry {
key: build_topk_key(&values, &sort_positions),
tuple: build_projected_tuple(&values, &output_positions),
asc: Arc::clone(&asc),
});
}
}
Ok(Some(
heap.into_sorted_vec().into_iter().map(|entry| entry.tuple).collect(),
))
}
pub(crate) fn scan_table_topk_columnar_projected_columns(
&self,
table_name: &str,
schema: &crate::Schema,
output_columns: &[usize],
sort_columns: &[usize],
asc: &[bool],
k: usize,
) -> Result<Option<Vec<Tuple>>> {
if k == 0 {
return Ok(Some(Vec::new()));
}
if sort_columns.is_empty() || sort_columns.len() != asc.len() {
return Ok(None);
}
let branch_name = self.current_branch.lock().clone();
if branch_name.is_some() && branch_name.as_deref() != Some("main") {
return Ok(None);
}
let mut requested: Vec<usize> = output_columns.to_vec();
requested.extend_from_slice(sort_columns);
requested.sort_unstable();
requested.dedup();
for &idx in &requested {
let Some(column) = schema.columns.get(idx) else {
return Err(Error::storage(format!(
"Column index {} out of bounds for {}",
idx, table_name
)));
};
if column.storage_mode != ColumnStorageMode::Columnar {
return Ok(None);
}
}
let mut column_batches: HashMap<usize, ColumnarBatchIndex> = HashMap::with_capacity(requested.len());
for &idx in &requested {
let column = &schema.columns[idx];
let batches = ColumnarStore::scan_column_batches(&self.db, table_name, &column.name)?;
column_batches.insert(idx, ColumnarBatchIndex::from_batches(batches));
}
let requested_pos = |column_index: usize| -> Result<usize> {
requested.binary_search(&column_index).map_err(|_| {
Error::storage(format!(
"Column index {} missing from columnar Top-N set for {}",
column_index, table_name
))
})
};
let output_positions: Vec<usize> = output_columns
.iter()
.map(|&idx| requested_pos(idx))
.collect::<Result<Vec<_>>>()?;
let sort_positions: Vec<usize> = sort_columns
.iter()
.map(|&idx| requested_pos(idx))
.collect::<Result<Vec<_>>>()?;
let asc = Arc::new(asc.to_vec());
let mut heap = std::collections::BinaryHeap::with_capacity(k.saturating_add(1));
let mut values = Vec::with_capacity(requested.len());
let mut consider_row = |batch_id: u64, offset: usize| {
values.clear();
for &idx in &requested {
values.push(columnar_batch_value(&column_batches, idx, batch_id, offset).unwrap_or(Value::Null));
}
let replace = if heap.len() < k {
true
} else {
heap.peek()
.map(|top: &RowTopKEntry| {
compare_values_to_topk_key(&values, &sort_positions, &top.key, &asc) == std::cmp::Ordering::Less
})
.unwrap_or(false)
};
if replace {
if heap.len() >= k {
heap.pop();
}
heap.push(RowTopKEntry {
key: build_topk_key(&values, &sort_positions),
tuple: build_projected_tuple(&values, &output_positions),
asc: Arc::clone(&asc),
});
}
};
if let Some(presence) = self.columnar_live_presence(table_name) {
for (batch_id, batch_presence) in &presence {
for offset in batch_presence.iter_live() {
consider_row(*batch_id, offset);
}
}
} else {
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
consider_row(row_id / BATCH_SIZE as u64, (row_id % BATCH_SIZE as u64) as usize);
}
}
Ok(Some(
heap.into_sorted_vec().into_iter().map(|entry| entry.tuple).collect(),
))
}
pub fn scan_table_with_schema_columnar_columns(
&self,
table_name: &str,
schema: &crate::Schema,
columns: &[usize],
) -> Result<Vec<Tuple>> {
self.scan_table_with_schema_columnar_columns_opt(table_name, schema, columns, &[])
}
pub fn scan_table_with_schema_columnar_columns_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
columns: &[usize],
predicates: &[AnalyzedPredicate],
) -> Result<Vec<Tuple>> {
self.scan_table_with_schema_columnar_columns_opt(table_name, schema, columns, predicates)
}
pub fn scan_table_with_schema_columnar_projected_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
projection: &[usize],
predicates: &[AnalyzedPredicate],
) -> Result<Option<Vec<Tuple>>> {
if projection.is_empty() {
return Ok(None);
}
let branch_name = self.current_branch.lock().clone();
if branch_name.is_some() && branch_name.as_deref() != Some("main") {
return Ok(None);
}
let mut requested: Vec<usize> = projection.to_vec();
requested.extend(predicates.iter().map(|predicate| predicate.column_index));
requested.sort_unstable();
requested.dedup();
for &idx in &requested {
let column = schema
.columns
.get(idx)
.ok_or_else(|| Error::storage(format!("Column index {} out of bounds for {}", idx, table_name)))?;
if column.storage_mode != ColumnStorageMode::Columnar {
return Ok(None);
}
}
let filter_predicates: Vec<FilterPredicate> = predicates.iter().filter_map(columnar_filter_predicate).collect();
if filter_predicates.len() != predicates.len() {
return Ok(None);
}
let pruned = compute_zone_pruned_batches(&self.db, table_name, schema, &filter_predicates);
let mut column_batches: HashMap<usize, ColumnarBatchIndex> = HashMap::with_capacity(requested.len());
for &idx in &requested {
let column = &schema.columns[idx];
column_batches.insert(
idx,
load_columnar_batch_index(&self.db, table_name, &column.name, &pruned)?,
);
}
if let Some(driver_predicate) = null_rejecting_filter_predicate(&filter_predicates)
.filter(|predicate| columnar_batch_driver_scan_beneficial(schema, predicate))
{
if let Some(driver_batches) = column_batches.get(&driver_predicate.column_index) {
let mut tuples = Vec::new();
for (_, batch) in driver_batches.ordered_batches() {
for (offset, driver_value) in batch.values().iter().enumerate() {
if !driver_predicate.evaluate(driver_value) {
continue;
}
let row_id = batch.start_row_id + offset as u64;
let batch_id = row_id / BATCH_SIZE as u64;
let batch_offset = (row_id % BATCH_SIZE as u64) as usize;
if filter_predicates.len() > 1
&& !columnar_row_matches_filters(
&column_batches,
batch_id,
batch_offset,
&filter_predicates,
)
{
continue;
}
let mut values = Vec::with_capacity(projection.len());
for &idx in projection {
values.push(
columnar_batch_value(&column_batches, idx, batch_id, batch_offset)
.unwrap_or(Value::Null),
);
}
let mut tuple = Tuple::new(values);
tuple.row_id = Some(row_id);
tuples.push(tuple);
}
}
return Ok(Some(tuples));
}
}
let mut tuples = Vec::new();
let emit_projected = |tuples: &mut Vec<Tuple>, batch_id: u64, offset: usize| {
let mut values = Vec::with_capacity(projection.len());
for &idx in projection {
values.push(columnar_batch_value(&column_batches, idx, batch_id, offset).unwrap_or(Value::Null));
}
let mut tuple = Tuple::new(values);
tuple.row_id = Some(batch_id * BATCH_SIZE as u64 + offset as u64);
tuples.push(tuple);
};
if let Some(presence) = self.columnar_live_presence(table_name) {
let compiled = tk::compile_predicates(&filter_predicates);
let mut mask = vec![0u8; BATCH_SIZE];
for (batch_id, batch_presence) in &presence {
if pruned.contains(batch_id) {
continue; }
tk::expand_bitmap_to_mask(&batch_presence.bits, &mut mask);
apply_compiled_predicates(&compiled, &column_batches, *batch_id, &mut mask);
for (offset, &m) in mask.iter().enumerate() {
if m != 0 {
emit_projected(&mut tuples, *batch_id, offset);
}
}
}
} else {
let mut emit_row = |batch_id: u64, offset: usize| {
if pruned.contains(&batch_id) {
return;
}
if !columnar_row_matches_filters(&column_batches, batch_id, offset, &filter_predicates) {
return;
}
emit_projected(&mut tuples, batch_id, offset);
};
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
emit_row(row_id / BATCH_SIZE as u64, (row_id % BATCH_SIZE as u64) as usize);
}
}
Ok(Some(tuples))
}
fn scan_table_with_schema_columnar_columns_opt(
&self,
table_name: &str,
schema: &crate::Schema,
columns: &[usize],
predicates: &[AnalyzedPredicate],
) -> Result<Vec<Tuple>> {
let scan_start = std::time::Instant::now();
let mut requested: Vec<usize> = columns.to_vec();
requested.sort_unstable();
requested.dedup();
for &idx in &requested {
let column = schema
.columns
.get(idx)
.ok_or_else(|| Error::storage(format!("Column index {} out of bounds for {}", idx, table_name)))?;
if column.storage_mode != ColumnStorageMode::Columnar {
return Err(Error::storage(format!(
"Columnar scan requested non-columnar column {}.{}",
table_name, column.name
)));
}
}
let filter_predicates: Vec<FilterPredicate> = predicates
.iter()
.filter(|predicate| requested.binary_search(&predicate.column_index).is_ok())
.filter_map(columnar_filter_predicate)
.collect();
let pruned = compute_zone_pruned_batches(&self.db, table_name, schema, &filter_predicates);
let mut column_batches: HashMap<usize, ColumnarBatchIndex> = HashMap::with_capacity(requested.len());
for &idx in &requested {
let column = &schema.columns[idx];
column_batches.insert(
idx,
load_columnar_batch_index(&self.db, table_name, &column.name, &pruned)?,
);
}
let mut tuples = Vec::new();
let emit_full_width = |tuples: &mut Vec<Tuple>, batch_id: u64, offset: usize| {
let mut values = vec![Value::Null; schema.columns.len()];
for &idx in &requested {
if let (Some(value), Some(slot)) = (
columnar_batch_value(&column_batches, idx, batch_id, offset),
values.get_mut(idx),
) {
*slot = value.clone();
}
}
let mut tuple = Tuple::new(values);
tuple.row_id = Some(batch_id * BATCH_SIZE as u64 + offset as u64);
tuples.push(tuple);
};
if let Some(presence) = self.columnar_live_presence(table_name) {
let compiled = tk::compile_predicates(&filter_predicates);
let mut mask = vec![0u8; BATCH_SIZE];
for (batch_id, batch_presence) in &presence {
if pruned.contains(batch_id) {
continue; }
tk::expand_bitmap_to_mask(&batch_presence.bits, &mut mask);
apply_compiled_predicates(&compiled, &column_batches, *batch_id, &mut mask);
for (offset, &m) in mask.iter().enumerate() {
if m != 0 {
emit_full_width(&mut tuples, *batch_id, offset);
}
}
}
} else {
let mut emit_row = |batch_id: u64, offset: usize| {
if pruned.contains(&batch_id) {
return;
}
if !columnar_row_matches_filters(&column_batches, batch_id, offset, &filter_predicates) {
return;
}
emit_full_width(&mut tuples, batch_id, offset);
};
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
let batch_id = row_id / BATCH_SIZE as u64;
let offset = (row_id % BATCH_SIZE as u64) as usize;
emit_row(batch_id, offset);
}
}
tracing::debug!(
phase = "storage_columnar_scan",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Columnar table scan complete"
);
Ok(tuples)
}
pub(crate) fn aggregate_columnar_columns(
&self,
table_name: &str,
schema: &crate::Schema,
group_by_columns: &[usize],
aggregates: &[ColumnarAggregateSpec],
predicates: &[AnalyzedPredicate],
) -> Result<Vec<Tuple>> {
let scan_start = std::time::Instant::now();
let mut requested: Vec<usize> = group_by_columns.to_vec();
requested.extend(aggregates.iter().filter_map(|aggregate| aggregate.column_index));
requested.extend(predicates.iter().map(|predicate| predicate.column_index));
requested.sort_unstable();
requested.dedup();
if requested.is_empty() {
return Ok(Vec::new());
}
for &idx in &requested {
let column = schema
.columns
.get(idx)
.ok_or_else(|| Error::storage(format!("Column index {} out of bounds for {}", idx, table_name)))?;
if column.storage_mode != ColumnStorageMode::Columnar {
return Err(Error::storage(format!(
"Columnar aggregate requested non-columnar column {}.{}",
table_name, column.name
)));
}
}
let filter_predicates: Vec<FilterPredicate> = predicates.iter().filter_map(columnar_filter_predicate).collect();
let pruned = compute_zone_pruned_batches(&self.db, table_name, schema, &filter_predicates);
let mut column_batches: HashMap<usize, ColumnarBatchIndex> = HashMap::with_capacity(requested.len());
for &idx in &requested {
let column = &schema.columns[idx];
column_batches.insert(
idx,
load_columnar_batch_index(&self.db, table_name, &column.name, &pruned)?,
);
}
if let Some(tuples) = self.try_aggregate_columnar_group_count_sum_int(
table_name,
schema,
group_by_columns,
aggregates,
&filter_predicates,
&column_batches,
&pruned,
)? {
tracing::debug!(
phase = "storage_columnar_aggregate",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Columnar grouped count/sum aggregate complete"
);
return Ok(tuples);
}
let has_distinct = aggregates
.iter()
.any(|aggregate| matches!(aggregate.op, ColumnarAggregateOp::CountDistinct));
let batches_per_chunk = (AGG_PARALLEL_CHUNK / BATCH_SIZE).max(1);
if group_by_columns.is_empty() && filter_predicates.is_empty() {
let mut values = Vec::with_capacity(aggregates.len());
for aggregate in aggregates {
if matches!(aggregate.op, ColumnarAggregateOp::CountStar) {
values.push(Value::Int8(self.count_table_rows(table_name)? as i64));
continue;
}
let mut state = ColumnarAggregateState::new(aggregate.op);
if let Some(column_index) = aggregate.column_index {
if let Some(batches) = column_batches.get(&column_index) {
let ordered = batches.ordered_batches();
let total_rows: usize = ordered.iter().map(|(_, batch)| batch.slot_count()).sum();
let fold_chunk = |chunk: &[(u64, &ColumnBatch)]| -> Result<ColumnarAggregateState> {
let mut state = ColumnarAggregateState::new(aggregate.op);
for (_, batch) in chunk {
if !kernel_update_aggregate(&mut state, aggregate.op, Some(batch), &tk::ALL_SELECTED)? {
for value in batch.values() {
state.update(aggregate.op, Some(value))?;
}
}
}
Ok(state)
};
if !matches!(aggregate.op, ColumnarAggregateOp::CountDistinct)
&& agg_parallel_rows_met(total_rows)
{
use rayon::prelude::*;
let partials: Vec<ColumnarAggregateState> = ordered
.par_chunks(batches_per_chunk)
.map(fold_chunk)
.collect::<Result<Vec<_>>>()?;
for partial in partials {
state.merge(partial)?;
}
} else {
state = fold_chunk(&ordered)?;
}
}
}
values.push(state.finalize()?);
}
let tuples = vec![Tuple::new(values)];
tracing::debug!(
phase = "storage_columnar_aggregate",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = 0usize,
duration_us = scan_start.elapsed().as_micros() as u64,
"Columnar unfiltered aggregate complete"
);
return Ok(tuples);
}
if let Some(driver_predicate) = null_rejecting_filter_predicate(&filter_predicates) {
if let Some(driver_batches) = column_batches.get(&driver_predicate.column_index) {
let ordered_driver = driver_batches.ordered_batches();
let driver_rows: usize = ordered_driver.iter().map(|(_, batch)| batch.slot_count()).sum();
let parallel = !has_distinct && agg_parallel_rows_met(driver_rows);
let compiled = tk::compile_predicates(&filter_predicates);
if group_by_columns.is_empty() {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
let fold_chunk = |chunk: &[(u64, &ColumnBatch)]| -> Result<Vec<ColumnarAggregateState>> {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
let mut mask = vec![0u8; BATCH_SIZE];
for &(batch_id, _) in chunk {
mask.fill(1);
apply_compiled_predicates(&compiled, &column_batches, batch_id, &mut mask);
update_states_for_batch_kernel(&mut states, aggregates, &column_batches, batch_id, &mask)?;
}
Ok(states)
};
if parallel {
use rayon::prelude::*;
let partials: Vec<Vec<ColumnarAggregateState>> = ordered_driver
.par_chunks(batches_per_chunk)
.map(fold_chunk)
.collect::<Result<Vec<_>>>()?;
for partial in partials {
merge_columnar_aggregate_states(&mut states, partial)?;
}
} else {
states = fold_chunk(&ordered_driver)?;
}
let tuple_values: Result<Vec<Value>> =
states.into_iter().map(ColumnarAggregateState::finalize).collect();
let tuples = vec![Tuple::new(tuple_values?)];
tracing::debug!(
phase = "storage_columnar_aggregate",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
driver_column = driver_predicate.column_name.as_str(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Columnar driver aggregate complete"
);
return Ok(tuples);
}
let fold_chunk =
|chunk: &[(u64, &ColumnBatch)]| -> Result<HashMap<Vec<Value>, Vec<ColumnarAggregateState>>> {
let mut groups: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = HashMap::new();
let mut mask = vec![0u8; BATCH_SIZE];
for &(batch_id, _) in chunk {
mask.fill(1);
apply_compiled_predicates(&compiled, &column_batches, batch_id, &mut mask);
update_groups_for_batch_kernel(
&mut groups,
aggregates,
group_by_columns,
&column_batches,
batch_id,
&mask,
)?;
}
Ok(groups)
};
let hash_groups: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = if parallel {
use rayon::prelude::*;
let partials: Vec<HashMap<Vec<Value>, Vec<ColumnarAggregateState>>> = ordered_driver
.par_chunks(batches_per_chunk)
.map(fold_chunk)
.collect::<Result<Vec<_>>>()?;
let mut merged: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = HashMap::new();
for partial in partials {
merge_columnar_group_map(&mut merged, partial)?;
}
merged
} else {
fold_chunk(&ordered_driver)?
};
let mut grouped: Vec<(Vec<Value>, Vec<ColumnarAggregateState>)> = hash_groups.into_iter().collect();
grouped.sort_by(|(left, _), (right, _)| compare_value_slices(left, right));
let mut tuples = Vec::with_capacity(grouped.len());
for (mut group_key, states) in grouped {
let mut aggregate_values: Result<Vec<Value>> =
states.into_iter().map(ColumnarAggregateState::finalize).collect();
group_key.append(&mut aggregate_values?);
tuples.push(Tuple::new(group_key));
}
tracing::debug!(
phase = "storage_columnar_aggregate",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
driver_column = driver_predicate.column_name.as_str(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Columnar driver grouped aggregate complete"
);
return Ok(tuples);
}
}
if group_by_columns.is_empty() {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let row_count_hint = self.art_index_manager.pk_index_len(table_name);
let parallel = !has_distinct && row_count_hint.is_some_and(agg_parallel_rows_met);
let presence = self.columnar_live_presence(table_name);
if let Some(presence) = &presence {
let live = presence_live_view(presence, &pruned);
let compiled = tk::compile_predicates(&filter_predicates);
let fold_chunk = |chunk: &[(u64, &BatchPresence)]| -> Result<Vec<ColumnarAggregateState>> {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
let mut mask = vec![0u8; BATCH_SIZE];
for (batch_id, batch_presence) in chunk {
tk::expand_bitmap_to_mask(&batch_presence.bits, &mut mask);
apply_compiled_predicates(&compiled, &column_batches, *batch_id, &mut mask);
update_states_for_batch_kernel(&mut states, aggregates, &column_batches, *batch_id, &mask)?;
}
Ok(states)
};
if parallel {
use rayon::prelude::*;
let partials: Vec<Vec<ColumnarAggregateState>> = live
.par_chunks(batches_per_chunk.max(1))
.map(fold_chunk)
.collect::<Result<Vec<_>>>()?;
for partial in partials {
merge_columnar_aggregate_states(&mut states, partial)?;
}
} else {
states = fold_chunk(&live)?;
}
} else if parallel {
let partials: Vec<Vec<ColumnarAggregateState>> = self.par_prefix_shard_results(&prefix, |iter| {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
let batch_id = row_id / BATCH_SIZE as u64;
let offset = (row_id % BATCH_SIZE as u64) as usize;
if pruned.contains(&batch_id) {
continue;
}
if !columnar_row_matches_filters(&column_batches, batch_id, offset, &filter_predicates) {
continue;
}
for (state, aggregate) in states.iter_mut().zip(aggregates) {
let value = aggregate
.column_index
.and_then(|idx| columnar_batch_value(&column_batches, idx, batch_id, offset));
state.update(aggregate.op, value.as_ref())?;
}
}
Ok(states)
})?;
for partial in partials {
merge_columnar_aggregate_states(&mut states, partial)?;
}
} else {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
let batch_id = row_id / BATCH_SIZE as u64;
let offset = (row_id % BATCH_SIZE as u64) as usize;
if pruned.contains(&batch_id) {
continue;
}
if !columnar_row_matches_filters(&column_batches, batch_id, offset, &filter_predicates) {
continue;
}
for (state, aggregate) in states.iter_mut().zip(aggregates) {
let value = aggregate
.column_index
.and_then(|idx| columnar_batch_value(&column_batches, idx, batch_id, offset));
state.update(aggregate.op, value.as_ref())?;
}
}
}
let tuple_values: Result<Vec<Value>> = states.into_iter().map(ColumnarAggregateState::finalize).collect();
let tuples = vec![Tuple::new(tuple_values?)];
tracing::debug!(
phase = "storage_columnar_aggregate",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Columnar aggregate complete"
);
return Ok(tuples);
}
const LINEAR_GROUP_LIMIT: usize = 64;
let mut small_groups: Vec<(Vec<Value>, Vec<ColumnarAggregateState>)> = Vec::new();
let mut hash_groups: Option<HashMap<Vec<Value>, Vec<ColumnarAggregateState>>> = None;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let row_count_hint = self.art_index_manager.pk_index_len(table_name);
let parallel = !has_distinct && row_count_hint.is_some_and(agg_parallel_rows_met);
let presence = self.columnar_live_presence(table_name);
if let Some(presence) = &presence {
let live = presence_live_view(presence, &pruned);
let compiled = tk::compile_predicates(&filter_predicates);
let fold_chunk =
|chunk: &[(u64, &BatchPresence)]| -> Result<HashMap<Vec<Value>, Vec<ColumnarAggregateState>>> {
let mut groups: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = HashMap::new();
let mut mask = vec![0u8; BATCH_SIZE];
for (batch_id, batch_presence) in chunk {
tk::expand_bitmap_to_mask(&batch_presence.bits, &mut mask);
apply_compiled_predicates(&compiled, &column_batches, *batch_id, &mut mask);
update_groups_for_batch_kernel(
&mut groups,
aggregates,
group_by_columns,
&column_batches,
*batch_id,
&mask,
)?;
}
Ok(groups)
};
if parallel {
use rayon::prelude::*;
let partials: Vec<HashMap<Vec<Value>, Vec<ColumnarAggregateState>>> = live
.par_chunks(batches_per_chunk.max(1))
.map(fold_chunk)
.collect::<Result<Vec<_>>>()?;
let mut merged: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = HashMap::new();
for partial in partials {
merge_columnar_group_map(&mut merged, partial)?;
}
hash_groups = Some(merged);
} else {
hash_groups = Some(fold_chunk(&live)?);
}
} else if parallel {
let partials: Vec<HashMap<Vec<Value>, Vec<ColumnarAggregateState>>> =
self.par_prefix_shard_results(&prefix, |iter| {
let mut groups: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = HashMap::new();
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
let batch_id = row_id / BATCH_SIZE as u64;
let offset = (row_id % BATCH_SIZE as u64) as usize;
if pruned.contains(&batch_id) {
continue;
}
if !columnar_row_matches_filters(&column_batches, batch_id, offset, &filter_predicates) {
continue;
}
let group_key = build_columnar_group_key(&column_batches, group_by_columns, batch_id, offset);
let states = groups.entry(group_key).or_insert_with(|| {
aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect()
});
update_columnar_aggregate_states(states, aggregates, &column_batches, batch_id, offset)?;
}
Ok(groups)
})?;
let mut merged: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = HashMap::new();
for partial in partials {
merge_columnar_group_map(&mut merged, partial)?;
}
hash_groups = Some(merged);
} else {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
let batch_id = row_id / BATCH_SIZE as u64;
let offset = (row_id % BATCH_SIZE as u64) as usize;
if pruned.contains(&batch_id) {
continue;
}
if !columnar_row_matches_filters(&column_batches, batch_id, offset, &filter_predicates) {
continue;
}
if let Some(groups) = hash_groups.as_mut() {
let group_key = build_columnar_group_key(&column_batches, group_by_columns, batch_id, offset);
let states = groups.entry(group_key).or_insert_with(|| {
aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect()
});
update_columnar_aggregate_states(states, aggregates, &column_batches, batch_id, offset)?;
} else if let Some(idx) = small_groups.iter().position(|(group_key, _)| {
columnar_group_key_matches(group_key, &column_batches, group_by_columns, batch_id, offset)
}) {
update_columnar_aggregate_states(
&mut small_groups[idx].1,
aggregates,
&column_batches,
batch_id,
offset,
)?;
} else {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
update_columnar_aggregate_states(&mut states, aggregates, &column_batches, batch_id, offset)?;
small_groups.push((
build_columnar_group_key(&column_batches, group_by_columns, batch_id, offset),
states,
));
if small_groups.len() > LINEAR_GROUP_LIMIT {
hash_groups = Some(small_groups.drain(..).collect());
}
}
}
}
if group_by_columns.is_empty() && small_groups.is_empty() && hash_groups.is_none() {
small_groups.push((
Vec::new(),
aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect(),
));
}
let mut grouped: Vec<(Vec<Value>, Vec<ColumnarAggregateState>)> = match hash_groups {
Some(groups) => groups.into_iter().collect(),
None => small_groups,
};
grouped.sort_by(|(left, _), (right, _)| compare_value_slices(left, right));
let mut tuples = Vec::with_capacity(grouped.len());
for (mut group_key, states) in grouped {
let mut aggregate_values: Result<Vec<Value>> =
states.into_iter().map(ColumnarAggregateState::finalize).collect();
group_key.append(&mut aggregate_values?);
tuples.push(Tuple::new(group_key));
}
tracing::debug!(
phase = "storage_columnar_aggregate",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Columnar aggregate complete"
);
Ok(tuples)
}
#[allow(clippy::too_many_arguments)]
fn count_sum_kernel_chunk(
chunk: &[(u64, &BatchPresence)],
compiled: &[CompiledPredicate<'_>],
column_batches: &HashMap<usize, ColumnarBatchIndex>,
group_column: usize,
sum_column: usize,
) -> Result<HashMap<Value, CountSumIntState>> {
const MAX_DENSE_INT_GROUPS: i64 = 65_536;
fn fold_groups(
groups: &mut HashMap<Value, CountSumIntState>,
counts: &[i64],
sums: &[i128],
seen: &[u8],
key_of: impl Fn(usize) -> Value,
) -> Result<()> {
for (idx, ((&count, &sum), &sum_seen)) in counts.iter().zip(sums).zip(seen).enumerate() {
if count == 0 && sum_seen == 0 {
continue;
}
let state = groups.entry(key_of(idx)).or_insert_with(CountSumIntState::new);
state.count += count;
if sum_seen != 0 {
let sum = i64::try_from(sum).map_err(|_| Error::query_execution("integer overflow: BIGINT SUM"))?;
state.update_sum_int(Some(sum))?;
}
}
Ok(())
}
let mut groups: HashMap<Value, CountSumIntState> = HashMap::new();
let mut mask = vec![0u8; BATCH_SIZE];
for (batch_id, batch_presence) in chunk {
tk::expand_bitmap_to_mask(&batch_presence.bits, &mut mask);
apply_compiled_predicates(compiled, column_batches, *batch_id, &mut mask);
let gbatch = column_batches.get(&group_column).and_then(|by| by.get(*batch_id));
let sbatch = column_batches.get(&sum_column).and_then(|by| by.get(*batch_id));
let sum_arg: Option<Option<(&[i64], &[u8])>> = match sbatch {
None => Some(None),
Some(batch) => batch.typed.as_ref().and_then(|typed| match &typed.data {
TypedValues::Int { data, .. } => Some(Some((data.as_slice(), typed.validity.as_slice()))),
_ => None,
}),
};
let mut kernelized = false;
if let Some(sum) = sum_arg {
if let Some(gtyped) = gbatch.and_then(|b| b.typed.as_ref()) {
match >yped.data {
TypedValues::Text { dict, codes } => {
let slots = dict.len() + 1;
let mut counts = vec![0i64; slots];
let mut sums = vec![0i128; slots];
let mut seen = vec![0u8; slots];
tk::group_count_sum_by_code(
codes,
>yped.validity,
&mask,
sum,
&mut counts,
&mut sums,
&mut seen,
);
fold_groups(&mut groups, &counts, &sums, &seen, |idx| {
dict.get(idx).map_or(Value::Null, |s| Value::String(s.clone()))
})?;
kernelized = true;
}
TypedValues::Int { width, data } => {
if let Some((mn, mx)) = stats_int_bounds(>yped.stats) {
if let Some(range) =
mx.checked_sub(mn).filter(|r| (0..MAX_DENSE_INT_GROUPS).contains(r))
{
let slots = range as usize + 2;
let mut counts = vec![0i64; slots];
let mut sums = vec![0i128; slots];
let mut seen = vec![0u8; slots];
tk::group_count_sum_by_small_int(
data,
>yped.validity,
mn,
&mask,
sum,
&mut counts,
&mut sums,
&mut seen,
);
fold_groups(&mut groups, &counts, &sums, &seen, |idx| {
if idx + 1 == slots {
Value::Null
} else {
width.value(mn + idx as i64)
}
})?;
kernelized = true;
}
}
}
_ => {}
}
}
}
if !kernelized {
for (offset, &m) in mask.iter().enumerate() {
if m == 0 {
continue;
}
let sum_value = columnar_batch_value(column_batches, sum_column, *batch_id, offset);
let group_key =
columnar_batch_value(column_batches, group_column, *batch_id, offset).unwrap_or(Value::Null);
let state = groups.entry(group_key).or_insert_with(CountSumIntState::new);
state.update_count();
state.update_sum(sum_value.as_ref())?;
}
}
}
Ok(groups)
}
#[allow(clippy::too_many_arguments)]
fn try_aggregate_columnar_group_count_sum_int(
&self,
table_name: &str,
schema: &crate::Schema,
group_by_columns: &[usize],
aggregates: &[ColumnarAggregateSpec],
filter_predicates: &[FilterPredicate],
column_batches: &HashMap<usize, ColumnarBatchIndex>,
pruned: &HashSet<u64>,
) -> Result<Option<Vec<Tuple>>> {
if group_by_columns.len() != 1 || aggregates.len() != 2 {
return Ok(None);
}
if !matches!(aggregates[0].op, ColumnarAggregateOp::CountStar)
|| !matches!(aggregates[1].op, ColumnarAggregateOp::Sum)
{
return Ok(None);
}
let sum_column = match aggregates[1].column_index {
Some(idx) => idx,
None => return Ok(None),
};
if !schema
.columns
.get(sum_column)
.is_some_and(|column| primitive_integer_data_type(&column.data_type))
{
return Ok(None);
}
const LINEAR_GROUP_LIMIT: usize = 64;
let group_column = group_by_columns[0];
let mut small_groups: Vec<(Value, CountSumIntState)> = Vec::new();
let mut hash_groups: Option<HashMap<Value, CountSumIntState>> = None;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let row_count_hint = self.art_index_manager.pk_index_len(table_name);
let parallel = row_count_hint.is_some_and(agg_parallel_rows_met);
let presence = self.columnar_live_presence(table_name);
if let Some(presence) = &presence {
let live = presence_live_view(presence, pruned);
let compiled = tk::compile_predicates(filter_predicates);
if parallel {
use rayon::prelude::*;
let batches_per_chunk = (AGG_PARALLEL_CHUNK / BATCH_SIZE).max(1);
let partials: Vec<HashMap<Value, CountSumIntState>> = live
.par_chunks(batches_per_chunk)
.map(|chunk| {
Self::count_sum_kernel_chunk(chunk, &compiled, column_batches, group_column, sum_column)
})
.collect::<Result<Vec<_>>>()?;
let mut merged: HashMap<Value, CountSumIntState> = HashMap::new();
for partial in partials {
merge_count_sum_group_map(&mut merged, partial)?;
}
hash_groups = Some(merged);
} else {
hash_groups = Some(Self::count_sum_kernel_chunk(
&live,
&compiled,
column_batches,
group_column,
sum_column,
)?);
}
} else if parallel {
let partials: Vec<HashMap<Value, CountSumIntState>> = self.par_prefix_shard_results(&prefix, |iter| {
let mut groups: HashMap<Value, CountSumIntState> = HashMap::new();
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
let batch_id = row_id / BATCH_SIZE as u64;
let offset = (row_id % BATCH_SIZE as u64) as usize;
if pruned.contains(&batch_id) {
continue;
}
if !columnar_row_matches_filters(column_batches, batch_id, offset, filter_predicates) {
continue;
}
let sum_value = columnar_batch_value(column_batches, sum_column, batch_id, offset);
let group_key =
columnar_batch_value(column_batches, group_column, batch_id, offset).unwrap_or(Value::Null);
let state = groups.entry(group_key).or_insert_with(CountSumIntState::new);
state.update_count();
state.update_sum(sum_value.as_ref())?;
}
Ok(groups)
})?;
let mut merged: HashMap<Value, CountSumIntState> = HashMap::new();
for partial in partials {
merge_count_sum_group_map(&mut merged, partial)?;
}
hash_groups = Some(merged);
} else {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) else {
continue;
};
let batch_id = row_id / BATCH_SIZE as u64;
let offset = (row_id % BATCH_SIZE as u64) as usize;
if pruned.contains(&batch_id) {
continue;
}
if !columnar_row_matches_filters(column_batches, batch_id, offset, filter_predicates) {
continue;
}
let sum_value = columnar_batch_value(column_batches, sum_column, batch_id, offset);
if let Some(groups) = hash_groups.as_mut() {
let group_key =
columnar_batch_value(column_batches, group_column, batch_id, offset).unwrap_or(Value::Null);
let state = groups.entry(group_key).or_insert_with(CountSumIntState::new);
state.update_count();
state.update_sum(sum_value.as_ref())?;
} else if let Some(idx) = small_groups.iter().position(|(group_key, _)| {
columnar_batch_value(column_batches, group_column, batch_id, offset)
.map_or(matches!(group_key, Value::Null), |actual| *group_key == actual)
}) {
let state = &mut small_groups[idx].1;
state.update_count();
state.update_sum(sum_value.as_ref())?;
} else {
let group_key =
columnar_batch_value(column_batches, group_column, batch_id, offset).unwrap_or(Value::Null);
let mut state = CountSumIntState::new();
state.update_count();
state.update_sum(sum_value.as_ref())?;
small_groups.push((group_key, state));
if small_groups.len() > LINEAR_GROUP_LIMIT {
hash_groups = Some(small_groups.drain(..).collect());
}
}
}
}
let mut grouped: Vec<(Value, CountSumIntState)> = match hash_groups {
Some(groups) => groups.into_iter().collect(),
None => small_groups,
};
grouped.sort_by(|(left, _), (right, _)| {
compare_value_slices(std::slice::from_ref(left), std::slice::from_ref(right))
});
Ok(Some(
grouped
.into_iter()
.map(|(group_key, state)| state.finish(group_key))
.collect(),
))
}
fn try_aggregate_row_text_group_count_sum_int(
&self,
table_name: &str,
schema: &crate::Schema,
group_by_columns: &[usize],
aggregates: &[ColumnarAggregateSpec],
filter_predicates: &[FilterPredicate],
) -> Result<Option<Vec<Tuple>>> {
if !filter_predicates.is_empty() || group_by_columns.len() != 1 || aggregates.len() != 2 {
return Ok(None);
}
if !matches!(aggregates[0].op, ColumnarAggregateOp::CountStar)
|| !matches!(aggregates[1].op, ColumnarAggregateOp::Sum)
{
return Ok(None);
}
let sum_column = match aggregates[1].column_index {
Some(idx) => idx,
None => return Ok(None),
};
if !schema
.columns
.get(sum_column)
.is_some_and(|column| primitive_integer_data_type(&column.data_type))
{
return Ok(None);
}
let group_column = group_by_columns[0];
if group_column == sum_column
|| !schema.columns.get(group_column).is_some_and(|column| {
matches!(
column.data_type,
crate::DataType::Text | crate::DataType::Varchar(_) | crate::DataType::Char(_)
)
})
{
return Ok(None);
}
const LINEAR_GROUP_LIMIT: usize = 64;
let mut small_groups: Vec<(Option<String>, CountSumIntState)> = Vec::new();
let mut hash_groups: Option<HashMap<Option<String>, CountSumIntState>> = None;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let row_count_hint = self.art_index_manager.pk_index_len(table_name);
if row_count_hint.is_some_and(agg_parallel_rows_met) {
let partials: Vec<Option<HashMap<Option<String>, CountSumIntState>>> =
self.par_prefix_shard_results(&prefix, |iter| {
let mut shard_small: Vec<(Option<String>, CountSumIntState)> = Vec::new();
let mut shard_hash: Option<HashMap<Option<String>, CountSumIntState>> = None;
for item in iter {
let (_, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
let decrypted;
let bytes: &[u8] = if let Some(km) = &self.key_manager {
decrypted = crypto::decrypt(km.key(), &raw_value)?;
&decrypted
} else {
&raw_value
};
let Some((group_key, sum_value)) =
crate::storage::prefix_decode::decode_tuple_text_and_int_columns(
bytes,
group_column,
sum_column,
)
else {
return Ok(None);
};
update_text_count_sum_group(
group_key,
sum_value,
&mut shard_small,
&mut shard_hash,
LINEAR_GROUP_LIMIT,
)?;
}
Ok(Some(match shard_hash {
Some(groups) => groups,
None => shard_small.into_iter().collect(),
}))
})?;
let mut merged: HashMap<Option<String>, CountSumIntState> = HashMap::new();
for partial in partials {
let Some(partial) = partial else {
return Ok(None);
};
merge_count_sum_group_map(&mut merged, partial)?;
}
hash_groups = Some(merged);
} else {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
let Some((group_key, sum_value)) = crate::storage::prefix_decode::decode_tuple_text_and_int_columns(
&decrypted,
group_column,
sum_column,
) else {
return Ok(None);
};
update_text_count_sum_group(
group_key,
sum_value,
&mut small_groups,
&mut hash_groups,
LINEAR_GROUP_LIMIT,
)?;
} else {
let Some((group_key, sum_value)) = crate::storage::prefix_decode::decode_tuple_text_and_int_columns(
&raw_value,
group_column,
sum_column,
) else {
return Ok(None);
};
update_text_count_sum_group(
group_key,
sum_value,
&mut small_groups,
&mut hash_groups,
LINEAR_GROUP_LIMIT,
)?;
}
}
}
let mut grouped: Vec<(Option<String>, CountSumIntState)> = match hash_groups {
Some(groups) => groups.into_iter().collect(),
None => small_groups,
};
grouped.sort_by(|(left, _), (right, _)| match (left, right) {
(None, None) => std::cmp::Ordering::Equal,
(None, Some(_)) => std::cmp::Ordering::Less,
(Some(_), None) => std::cmp::Ordering::Greater,
(Some(left), Some(right)) => left.cmp(right),
});
Ok(Some(
grouped
.into_iter()
.map(|(group_key, state)| state.finish(group_key.map(Value::String).unwrap_or(Value::Null)))
.collect(),
))
}
fn primitive_row_aggregate_plan(
schema: &crate::Schema,
aggregates: &[ColumnarAggregateSpec],
) -> Option<(Vec<usize>, Vec<PrimitiveRowAggregate>)> {
let mut requested = Vec::new();
for aggregate in aggregates {
match aggregate.op {
ColumnarAggregateOp::CountStar => {}
ColumnarAggregateOp::Count | ColumnarAggregateOp::Avg | ColumnarAggregateOp::Sum => {
let idx = aggregate.column_index?;
let column = schema.columns.get(idx)?;
match aggregate.op {
ColumnarAggregateOp::Count | ColumnarAggregateOp::Avg
if primitive_numeric_data_type(&column.data_type) => {}
ColumnarAggregateOp::Sum if primitive_integer_data_type(&column.data_type) => {}
_ => return None,
}
requested.push(idx);
}
ColumnarAggregateOp::CountDistinct | ColumnarAggregateOp::Min | ColumnarAggregateOp::Max => {
return None;
}
}
}
requested.sort_unstable();
requested.dedup();
if let Some(max_requested) = requested.last().copied() {
for idx in 0..=max_requested {
let column = schema.columns.get(idx)?;
if requested.binary_search(&idx).is_err() && !row_blob_fast_skip_supported(&column.data_type) {
return None;
}
}
}
let mut plan = Vec::with_capacity(aggregates.len());
for aggregate in aggregates {
match aggregate.op {
ColumnarAggregateOp::CountStar => plan.push(PrimitiveRowAggregate::CountStar),
ColumnarAggregateOp::Count => {
let position = requested.binary_search(&aggregate.column_index?).ok()?;
plan.push(PrimitiveRowAggregate::Count { position });
}
ColumnarAggregateOp::Sum => {
let position = requested.binary_search(&aggregate.column_index?).ok()?;
plan.push(PrimitiveRowAggregate::SumInt { position });
}
ColumnarAggregateOp::Avg => {
let position = requested.binary_search(&aggregate.column_index?).ok()?;
plan.push(PrimitiveRowAggregate::Avg { position });
}
ColumnarAggregateOp::CountDistinct | ColumnarAggregateOp::Min | ColumnarAggregateOp::Max => {
return None;
}
}
}
Some((requested, plan))
}
fn try_aggregate_primitive_row_columns(
&self,
table_name: &str,
schema: &crate::Schema,
aggregates: &[ColumnarAggregateSpec],
) -> Result<Option<Vec<Tuple>>> {
let Some((requested, plan)) = Self::primitive_row_aggregate_plan(schema, aggregates) else {
return Ok(None);
};
if requested.is_empty() {
let count = self.count_table_rows(table_name)? as i64;
let values = plan
.iter()
.map(|aggregate| match aggregate {
PrimitiveRowAggregate::CountStar => Value::Int8(count),
_ => Value::Null,
})
.collect();
return Ok(Some(vec![Tuple::new(values)]));
}
if let Some((sum_position, avg_position)) = Self::primitive_count_sum_avg_positions(&plan) {
if let Some(tuple) =
self.try_aggregate_primitive_count_sum_avg(table_name, &requested, sum_position, avg_position)?
{
return Ok(Some(vec![tuple]));
}
}
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut states: Vec<PrimitiveRowAggregateState> =
plan.iter().copied().map(PrimitiveRowAggregateState::new).collect();
let row_count_hint = self.art_index_manager.pk_index_len(table_name);
if row_count_hint.is_some_and(agg_parallel_rows_met) {
let partials: Vec<Option<Vec<PrimitiveRowAggregateState>>> =
self.par_prefix_shard_results(&prefix, |iter| {
let mut states: Vec<PrimitiveRowAggregateState> =
plan.iter().copied().map(PrimitiveRowAggregateState::new).collect();
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (_, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
let decoded = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
crate::storage::prefix_decode::decode_tuple_numeric_column_values_into(
&decrypted,
&requested,
&mut values,
)
} else {
crate::storage::prefix_decode::decode_tuple_numeric_column_values_into(
&raw_value,
&requested,
&mut values,
)
};
if decoded.is_none() {
return Ok(None);
}
for (state, aggregate) in states.iter_mut().zip(&plan) {
state.update(*aggregate, &values)?;
}
}
Ok(Some(states))
})?;
for partial in partials {
let Some(partial) = partial else {
return Ok(None);
};
for (state, chunk_state) in states.iter_mut().zip(partial) {
state.merge(chunk_state)?;
}
}
} else {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let decoded = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
crate::storage::prefix_decode::decode_tuple_numeric_column_values_into(
&decrypted,
&requested,
&mut values,
)
} else {
crate::storage::prefix_decode::decode_tuple_numeric_column_values_into(
&raw_value,
&requested,
&mut values,
)
};
if decoded.is_none() {
return Ok(None);
}
for (state, aggregate) in states.iter_mut().zip(&plan) {
state.update(*aggregate, &values)?;
}
}
}
Ok(Some(vec![Tuple::new(
states.into_iter().map(PrimitiveRowAggregateState::finalize).collect(),
)]))
}
fn primitive_count_sum_avg_positions(plan: &[PrimitiveRowAggregate]) -> Option<(usize, usize)> {
match plan {
[PrimitiveRowAggregate::CountStar, PrimitiveRowAggregate::SumInt { position: sum_position }, PrimitiveRowAggregate::Avg { position: avg_position }] => {
Some((*sum_position, *avg_position))
}
_ => None,
}
}
fn try_aggregate_primitive_count_sum_avg(
&self,
table_name: &str,
requested: &[usize],
sum_position: usize,
avg_position: usize,
) -> Result<Option<Tuple>> {
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut count = 0_i64;
let mut sum = 0_i64;
let mut sum_seen = false;
let mut avg_sum = 0.0_f64;
let mut avg_count = 0_u64;
let row_count_hint = self.art_index_manager.pk_index_len(table_name);
if row_count_hint.is_some_and(agg_parallel_rows_met) {
let partials: Vec<Option<(i64, i64, bool, f64, u64)>> = self.par_prefix_shard_results(&prefix, |iter| {
let mut count = 0_i64;
let mut sum = 0_i64;
let mut sum_seen = false;
let mut avg_sum = 0.0_f64;
let mut avg_count = 0_u64;
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (_, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
let decoded = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
crate::storage::prefix_decode::decode_tuple_numeric_column_values_into(
&decrypted,
requested,
&mut values,
)
} else {
crate::storage::prefix_decode::decode_tuple_numeric_column_values_into(
&raw_value,
requested,
&mut values,
)
};
if decoded.is_none() {
return Ok(None);
}
count += 1;
if let Some(crate::storage::prefix_decode::DecodedNumericValue::Int(value)) =
values.get(sum_position)
{
sum = sum
.checked_add(*value)
.ok_or_else(|| Error::query_execution("integer overflow: BIGINT SUM"))?;
sum_seen = true;
}
match values.get(avg_position) {
Some(crate::storage::prefix_decode::DecodedNumericValue::Int(value)) => {
avg_sum += *value as f64;
avg_count += 1;
}
Some(crate::storage::prefix_decode::DecodedNumericValue::Float(value)) => {
avg_sum += *value;
avg_count += 1;
}
_ => {}
}
}
Ok(Some((count, sum, sum_seen, avg_sum, avg_count)))
})?;
for partial in partials {
let Some((chunk_count, chunk_sum, chunk_sum_seen, chunk_avg_sum, chunk_avg_count)) = partial else {
return Ok(None);
};
count += chunk_count;
if chunk_sum_seen {
sum = sum
.checked_add(chunk_sum)
.ok_or_else(|| Error::query_execution("integer overflow: BIGINT SUM"))?;
sum_seen = true;
}
avg_sum += chunk_avg_sum;
avg_count += chunk_avg_count;
}
} else {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
let decoded = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
crate::storage::prefix_decode::decode_tuple_numeric_column_values_into(
&decrypted,
requested,
&mut values,
)
} else {
crate::storage::prefix_decode::decode_tuple_numeric_column_values_into(
&raw_value,
requested,
&mut values,
)
};
if decoded.is_none() {
return Ok(None);
}
count += 1;
if let Some(crate::storage::prefix_decode::DecodedNumericValue::Int(value)) = values.get(sum_position) {
sum = sum
.checked_add(*value)
.ok_or_else(|| Error::query_execution("integer overflow: BIGINT SUM"))?;
sum_seen = true;
}
match values.get(avg_position) {
Some(crate::storage::prefix_decode::DecodedNumericValue::Int(value)) => {
avg_sum += *value as f64;
avg_count += 1;
}
Some(crate::storage::prefix_decode::DecodedNumericValue::Float(value)) => {
avg_sum += *value;
avg_count += 1;
}
_ => {}
}
}
}
Ok(Some(Tuple::new(vec![
Value::Int8(count),
if sum_seen { Value::Int8(sum) } else { Value::Null },
if avg_count == 0 {
Value::Null
} else {
Value::Float8(avg_sum / avg_count as f64)
},
])))
}
pub(crate) fn try_aggregate_row_columns(
&self,
table_name: &str,
schema: &crate::Schema,
group_by_columns: &[usize],
aggregates: &[ColumnarAggregateSpec],
predicates: &[AnalyzedPredicate],
) -> Result<Option<Vec<Tuple>>> {
let scan_start = std::time::Instant::now();
let mut requested: Vec<usize> = group_by_columns.to_vec();
requested.extend(aggregates.iter().filter_map(|aggregate| aggregate.column_index));
requested.extend(predicates.iter().map(|predicate| predicate.column_index));
requested.sort_unstable();
requested.dedup();
for &idx in &requested {
let column = schema
.columns
.get(idx)
.ok_or_else(|| Error::storage(format!("Column index {} out of bounds for {}", idx, table_name)))?;
if column.storage_mode != ColumnStorageMode::Default {
return Ok(None);
}
}
let filter_predicates: Vec<FilterPredicate> = predicates.iter().filter_map(columnar_filter_predicate).collect();
if filter_predicates.len() != predicates.len() {
return Ok(None);
}
if group_by_columns.is_empty() && predicates.is_empty() {
if let Some(tuples) = self.try_aggregate_primitive_row_columns(table_name, schema, aggregates)? {
tracing::debug!(
phase = "storage_row_aggregate_primitive",
table = table_name,
rows = tuples.len(),
"Primitive row-store aggregate complete"
);
return Ok(Some(tuples));
}
}
if requested.is_empty() && group_by_columns.is_empty() && predicates.is_empty() {
let count = self.count_table_rows(table_name)? as i64;
let values = aggregates
.iter()
.map(|aggregate| match aggregate.op {
ColumnarAggregateOp::CountStar => Ok(Value::Int8(count)),
_ => Err(Error::query_execution("row aggregate requires a referenced column")),
})
.collect::<Result<Vec<_>>>()?;
return Ok(Some(vec![Tuple::new(values)]));
}
if let Some(tuples) = self.try_aggregate_row_text_group_count_sum_int(
table_name,
schema,
group_by_columns,
aggregates,
&filter_predicates,
)? {
tracing::debug!(
phase = "storage_row_aggregate_text_count_sum",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Row-store text grouped count/sum aggregate complete"
);
return Ok(Some(tuples));
}
let requested_pos = |column_index: usize| -> Result<usize> {
requested.binary_search(&column_index).map_err(|_| {
Error::storage(format!(
"Column index {} missing from aggregate decode set for {}",
column_index, table_name
))
})
};
let group_positions: Vec<usize> = group_by_columns
.iter()
.map(|&idx| requested_pos(idx))
.collect::<Result<Vec<_>>>()?;
let aggregate_positions: Vec<Option<usize>> = aggregates
.iter()
.map(|aggregate| aggregate.column_index.map(&requested_pos).transpose())
.collect::<Result<Vec<_>>>()?;
let filter_positions: Vec<usize> = predicates
.iter()
.map(|predicate| requested_pos(predicate.column_index))
.collect::<Result<Vec<_>>>()?;
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let has_distinct = aggregates
.iter()
.any(|aggregate| matches!(aggregate.op, ColumnarAggregateOp::CountDistinct));
let row_count_hint = self.art_index_manager.pk_index_len(table_name);
let parallel = !has_distinct && row_count_hint.is_some_and(row_agg_parallel_rows_met);
if group_by_columns.is_empty() {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
if parallel {
let partials: Vec<Vec<ColumnarAggregateState>> = self.par_prefix_shard_results(&prefix, |iter| {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (_, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
self.decode_rowstore_column_values_into(
&raw_value,
&requested,
schema.columns.len(),
&mut values,
)?;
if !row_values_match_filters(&values, &filter_predicates, &filter_positions) {
continue;
}
for ((state, aggregate), position) in
states.iter_mut().zip(aggregates).zip(&aggregate_positions)
{
let value = position.and_then(|pos| values.get(pos));
state.update(aggregate.op, value)?;
}
}
Ok(states)
})?;
for partial in partials {
merge_columnar_aggregate_states(&mut states, partial)?;
}
} else {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
self.decode_rowstore_column_values_into(&raw_value, &requested, schema.columns.len(), &mut values)?;
if !row_values_match_filters(&values, &filter_predicates, &filter_positions) {
continue;
}
for ((state, aggregate), position) in states.iter_mut().zip(aggregates).zip(&aggregate_positions) {
let value = position.and_then(|pos| values.get(pos));
state.update(aggregate.op, value)?;
}
}
}
let tuple_values: Result<Vec<Value>> = states.into_iter().map(ColumnarAggregateState::finalize).collect();
let tuples = vec![Tuple::new(tuple_values?)];
tracing::debug!(
phase = "storage_row_aggregate",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Row-store aggregate complete"
);
return Ok(Some(tuples));
}
const LINEAR_GROUP_LIMIT: usize = 64;
let mut small_groups: Vec<(Vec<Value>, Vec<ColumnarAggregateState>)> = Vec::new();
let mut hash_groups: Option<HashMap<Vec<Value>, Vec<ColumnarAggregateState>>> = None;
if parallel {
let partials: Vec<HashMap<Vec<Value>, Vec<ColumnarAggregateState>>> =
self.par_prefix_shard_results(&prefix, |iter| {
let mut groups: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = HashMap::new();
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (_, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
self.decode_rowstore_column_values_into(
&raw_value,
&requested,
schema.columns.len(),
&mut values,
)?;
if !row_values_match_filters(&values, &filter_predicates, &filter_positions) {
continue;
}
let states = groups
.entry(build_group_key(&values, &group_positions))
.or_insert_with(|| {
aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect()
});
for ((state, aggregate), position) in
states.iter_mut().zip(aggregates).zip(&aggregate_positions)
{
let value = position.and_then(|pos| values.get(pos));
state.update(aggregate.op, value)?;
}
}
Ok(groups)
})?;
let mut merged: HashMap<Vec<Value>, Vec<ColumnarAggregateState>> = HashMap::new();
for partial in partials {
merge_columnar_group_map(&mut merged, partial)?;
}
hash_groups = Some(merged);
} else {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut values = Vec::with_capacity(requested.len());
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break;
}
self.decode_rowstore_column_values_into(&raw_value, &requested, schema.columns.len(), &mut values)?;
if !row_values_match_filters(&values, &filter_predicates, &filter_positions) {
continue;
}
if let Some(groups) = hash_groups.as_mut() {
let group_key = build_group_key(&values, &group_positions);
let states = groups.entry(group_key).or_insert_with(|| {
aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect()
});
for ((state, aggregate), position) in states.iter_mut().zip(aggregates).zip(&aggregate_positions) {
let value = position.and_then(|pos| values.get(pos));
state.update(aggregate.op, value)?;
}
} else if let Some(idx) = small_groups
.iter()
.position(|(group_key, _)| group_key_matches_values(group_key, &values, &group_positions))
{
let states = &mut small_groups[idx].1;
for ((state, aggregate), position) in states.iter_mut().zip(aggregates).zip(&aggregate_positions) {
let value = position.and_then(|pos| values.get(pos));
state.update(aggregate.op, value)?;
}
} else {
let mut states: Vec<ColumnarAggregateState> = aggregates
.iter()
.map(|aggregate| ColumnarAggregateState::new(aggregate.op))
.collect();
for ((state, aggregate), position) in states.iter_mut().zip(aggregates).zip(&aggregate_positions) {
let value = position.and_then(|pos| values.get(pos));
state.update(aggregate.op, value)?;
}
small_groups.push((build_group_key(&values, &group_positions), states));
if small_groups.len() > LINEAR_GROUP_LIMIT {
hash_groups = Some(small_groups.drain(..).collect());
}
}
}
}
let mut grouped: Vec<(Vec<Value>, Vec<ColumnarAggregateState>)> = match hash_groups {
Some(groups) => groups.into_iter().collect(),
None => small_groups,
};
grouped.sort_by(|(left, _), (right, _)| compare_value_slices(left, right));
let mut tuples = Vec::with_capacity(grouped.len());
for (mut group_key, states) in grouped {
let mut aggregate_values: Result<Vec<Value>> =
states.into_iter().map(ColumnarAggregateState::finalize).collect();
group_key.append(&mut aggregate_values?);
tuples.push(Tuple::new(group_key));
}
tracing::debug!(
phase = "storage_row_aggregate",
table = table_name,
rows = tuples.len(),
columns = requested.len(),
predicates = filter_predicates.len(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Row-store grouped aggregate complete"
);
Ok(Some(tuples))
}
fn decode_rowstore_columns(&self, raw_value: &[u8], columns: &[usize], total_cols: usize) -> Result<Tuple> {
let tuple = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), raw_value)?;
crate::storage::prefix_decode::decode_tuple_columns(&decrypted, columns, total_cols)
} else {
crate::storage::prefix_decode::decode_tuple_columns(raw_value, columns, total_cols)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
Ok(tuple)
}
fn decode_rowstore_column_values(
&self,
raw_value: &[u8],
columns: &[usize],
total_cols: usize,
) -> Result<Vec<Value>> {
let values = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), raw_value)?;
crate::storage::prefix_decode::decode_tuple_column_values(&decrypted, columns, total_cols)
} else {
crate::storage::prefix_decode::decode_tuple_column_values(raw_value, columns, total_cols)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
Ok(values)
}
fn decode_rowstore_column_values_into(
&self,
raw_value: &[u8],
columns: &[usize],
total_cols: usize,
out: &mut Vec<Value>,
) -> Result<()> {
if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), raw_value)?;
crate::storage::prefix_decode::decode_tuple_column_values_into(&decrypted, columns, total_cols, out)
} else {
crate::storage::prefix_decode::decode_tuple_column_values_into(raw_value, columns, total_cols, out)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))
}
fn scan_table_with_schema_opt(
&self,
table_name: &str,
schema: &crate::Schema,
decode_hint: RowDecodeHint<'_>,
) -> Result<Vec<Tuple>> {
let scan_start = std::time::Instant::now();
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let decode = |key: &[u8], raw_value: &[u8]| -> Result<Tuple> {
{
let mut tuple: Tuple = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
match decode_hint {
RowDecodeHint::Prefix(k) => {
crate::storage::prefix_decode::decode_tuple_prefix(&decrypted, k, schema.columns.len())
}
RowDecodeHint::Columns(columns) => crate::storage::prefix_decode::decode_tuple_columns(
&decrypted,
columns,
schema.columns.len(),
),
RowDecodeHint::Full => bincode::deserialize(&decrypted),
}
} else {
match decode_hint {
RowDecodeHint::Prefix(k) => {
crate::storage::prefix_decode::decode_tuple_prefix(&raw_value, k, schema.columns.len())
}
RowDecodeHint::Columns(columns) => crate::storage::prefix_decode::decode_tuple_columns(
&raw_value,
columns,
schema.columns.len(),
),
RowDecodeHint::Full => bincode::deserialize(&raw_value),
}
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
let mut row_id = 0u64;
if let Some(rid) = Self::parse_row_id_after_prefix(key, prefix_bytes.len()) {
row_id = rid;
tuple.row_id = Some(row_id);
}
for (idx, column) in schema.columns.iter().enumerate() {
if idx >= tuple.values.len() {
break;
}
match column.storage_mode {
ColumnStorageMode::Dictionary => {
if let Some(crate::Value::DictRef { dict_id }) = tuple.values.get(idx) {
let dict_id = *dict_id;
let s = self.dict_manager.decode(&self.db, table_name, &column.name, dict_id)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = crate::Value::String(s);
}
}
}
ColumnStorageMode::ContentAddressed => {
if let Some(crate::Value::CasRef { hash }) = tuple.values.get(idx) {
let hash = hash.clone();
let resolved = ContentAddressedStore::resolve(&self.db, &hash, &column.data_type)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = resolved;
}
}
}
ColumnStorageMode::Columnar => {
if matches!(tuple.values.get(idx), Some(crate::Value::ColumnarRef)) {
if let Some(val) = ColumnarStore::get(&self.db, table_name, &column.name, row_id)? {
if let Some(slot) = tuple.values.get_mut(idx) {
*slot = val;
}
}
}
}
ColumnStorageMode::Default => {
}
}
}
Ok(tuple)
}
};
static SCAN_SERIAL: once_cell::sync::Lazy<bool> =
once_cell::sync::Lazy::new(|| std::env::var("HELIOS_SCAN_SERIAL").is_ok());
const PAR_DECODE_THRESHOLD: usize = 131_072;
let row_count_hint = self.art_index_manager.pk_index_len(table_name);
let parallel_decode = !*SCAN_SERIAL
&& row_count_hint
.map(|count| count >= PAR_DECODE_THRESHOLD)
.unwrap_or(false);
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let tuples: Vec<Tuple> = if parallel_decode {
let mut raw_rows: Vec<(Box<[u8]>, Box<[u8]>)> = Vec::with_capacity(row_count_hint.unwrap_or_default());
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break; }
raw_rows.push((key, raw_value));
}
use rayon::prelude::*;
raw_rows
.par_iter()
.map(|kv| decode(&kv.0, &kv.1))
.collect::<Result<Vec<_>>>()?
} else {
let mut tuples = Vec::with_capacity(row_count_hint.unwrap_or(256));
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
break; }
tuples.push(decode(&key, &raw_value)?);
}
tuples
};
tracing::debug!(
phase = "storage_scan",
table = table_name,
rows = tuples.len(),
duration_us = scan_start.elapsed().as_micros() as u64,
"Table scan complete"
);
Ok(tuples)
}
pub fn count_table_rows(&self, table_name: &str) -> Result<usize> {
let on_main_branch = self
.current_branch
.lock()
.as_deref()
.map_or(true, |name| name == "main");
let has_user_branches = self
.branch_manager
.read()
.as_ref()
.map(|manager| manager.has_user_branches_registered())
.unwrap_or(true);
if on_main_branch && !has_user_branches {
if let Some(count) = self.art_index_manager.pk_index_len(table_name) {
tracing::debug!(
phase = "count_pk_art_path",
table = table_name,
count = count,
"COUNT(*) primary-key ART fast path completed"
);
return Ok(count);
}
}
if columnar_presence_enabled()
&& on_main_branch
&& ColumnarStore::presence_manifest_complete(&self.db, table_name)
{
let count: usize = ColumnarStore::load_presence_map(&self.db, table_name)?
.iter()
.map(|(_, presence)| presence.live_count as usize)
.sum();
tracing::debug!(
phase = "count_presence_path",
table = table_name,
count = count,
"COUNT(*) presence-sidecar fast path completed"
);
return Ok(count);
}
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut count = 0usize;
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if key.starts_with(prefix_bytes) {
count += 1;
} else {
break;
}
}
tracing::debug!(
phase = "count_fast_path",
table = table_name,
count = count,
"COUNT(*) fast path completed"
);
Ok(count)
}
pub fn count_table_pk_int_range_with_schema(
&self,
table_name: &str,
schema: &crate::Schema,
lower: Option<(i64, bool)>,
upper: Option<(i64, bool)>,
) -> Result<Option<usize>> {
if self.get_current_branch_id().is_some() {
return Ok(None);
}
let mut pk_cols = schema.columns.iter().filter(|col| col.primary_key);
let pk_col = match (pk_cols.next(), pk_cols.next()) {
(Some(col), None) => col,
_ => return Ok(None),
};
Ok(self
.art_index_manager
.pk_index_count_int_range(table_name, &pk_col.data_type, lower, upper))
}
pub fn scan_table_with_offset_limit(&self, table_name: &str, offset: usize, limit: usize) -> Result<Vec<Tuple>> {
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
let mut tuples = Vec::with_capacity(limit);
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut skipped: usize = 0;
for item in iter {
if tuples.len() >= limit {
break;
}
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(prefix_bytes) {
if !tuples.is_empty() || skipped > 0 {
break;
}
continue;
}
if skipped < offset {
skipped += 1;
continue;
}
let mut tuple: Tuple = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
bincode::deserialize(&decrypted)
} else {
bincode::deserialize(&raw_value)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
if let Some(rid) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) {
tuple.row_id = Some(rid);
}
for (idx, column) in schema.columns.iter().enumerate() {
if idx >= tuple.values.len() {
break;
}
#[allow(clippy::indexing_slicing)]
match column.storage_mode {
ColumnStorageMode::Dictionary => {
if let Some(crate::Value::DictRef { dict_id }) = tuple.values.get(idx) {
let dict_id = *dict_id;
let s = self.dict_manager.decode(&self.db, table_name, &column.name, dict_id)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = crate::Value::String(s);
}
}
}
ColumnStorageMode::ContentAddressed => {
if let Some(crate::Value::CasRef { hash }) = tuple.values.get(idx) {
let hash = *hash;
let resolved = ContentAddressedStore::resolve(&self.db, &hash, &column.data_type)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = resolved;
}
}
}
ColumnStorageMode::Columnar => {
if matches!(tuple.values.get(idx), Some(crate::Value::ColumnarRef)) {
if let Some(row_id) = tuple.row_id {
if let Some(val) = ColumnarStore::get(&self.db, table_name, &column.name, row_id)? {
if let Some(slot) = tuple.values.get_mut(idx) {
*slot = val;
}
}
}
}
}
ColumnStorageMode::Default => {}
}
}
tuples.push(tuple);
}
Ok(tuples)
}
pub fn scan_table_pk_range(
&self,
table_name: &str,
lower: Option<u64>,
upper: Option<u64>,
limit: usize,
descending: bool,
) -> Result<Vec<Tuple>> {
let all = self.scan_table(table_name)?;
let mut filtered: Vec<Tuple> = all
.into_iter()
.filter(|t| {
let rid = t.row_id.unwrap_or(0);
lower.map_or(true, |lo| rid >= lo) && upper.map_or(true, |hi| rid <= hi)
})
.collect();
filtered.sort_by_key(|t| t.row_id.unwrap_or(0));
if descending {
filtered.reverse();
}
filtered.truncate(limit);
Ok(filtered)
}
pub fn scan_table_with_limit(&self, table_name: &str, limit: usize) -> Result<Vec<Tuple>> {
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
let mut tuples = Vec::with_capacity(limit);
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
for item in iter {
if tuples.len() >= limit {
break;
}
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if key.starts_with(prefix_bytes) {
let mut tuple: Tuple = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
bincode::deserialize(&decrypted)
} else {
bincode::deserialize(&raw_value)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
if let Some(rid) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) {
tuple.row_id = Some(rid);
}
for (idx, column) in schema.columns.iter().enumerate() {
if idx >= tuple.values.len() {
break;
}
#[allow(clippy::indexing_slicing)]
match column.storage_mode {
ColumnStorageMode::Dictionary => {
if let Some(crate::Value::DictRef { dict_id }) = tuple.values.get(idx) {
let dict_id = *dict_id;
let s = self.dict_manager.decode(&self.db, table_name, &column.name, dict_id)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = crate::Value::String(s);
}
}
}
ColumnStorageMode::ContentAddressed => {
if let Some(crate::Value::CasRef { hash }) = tuple.values.get(idx) {
let hash = hash.clone();
let resolved = ContentAddressedStore::resolve(&self.db, &hash, &column.data_type)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = resolved;
}
}
}
ColumnStorageMode::Columnar => {
if matches!(tuple.values.get(idx), Some(crate::Value::ColumnarRef)) {
if let Some(row_id) = tuple.row_id {
if let Some(val) = ColumnarStore::get(&self.db, table_name, &column.name, row_id)? {
if let Some(slot) = tuple.values.get_mut(idx) {
*slot = val;
}
}
}
}
}
ColumnStorageMode::Default => {}
}
}
tuples.push(tuple);
} else if !tuples.is_empty() {
break;
}
}
Ok(tuples)
}
pub fn get_row_by_pk(&self, table_name: &str, pk_value: &crate::Value) -> Result<Option<Tuple>> {
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name).ok();
self.get_row_by_pk_inner(table_name, pk_value, schema.as_ref(), true, true)
}
pub fn get_row_by_pk_with_schema(
&self,
table_name: &str,
pk_value: &crate::Value,
schema: &crate::Schema,
) -> Result<Option<Tuple>> {
self.get_row_by_pk_inner(table_name, pk_value, Some(schema), true, true)
}
pub fn get_row_by_typed_pk_with_schema(
&self,
table_name: &str,
pk_value: &crate::Value,
schema: &crate::Schema,
) -> Result<Option<Tuple>> {
self.get_row_by_pk_inner(table_name, pk_value, Some(schema), true, false)
}
pub fn get_row_by_typed_pk_with_schema_no_cache_fill(
&self,
table_name: &str,
pk_value: &crate::Value,
schema: &crate::Schema,
) -> Result<Option<Tuple>> {
self.get_row_by_pk_inner(table_name, pk_value, Some(schema), false, false)
}
pub fn get_row_by_pk_for_write_with_schema(
&self,
table_name: &str,
pk_value: &crate::Value,
schema: &crate::Schema,
) -> Result<Option<Tuple>> {
self.get_row_by_pk_inner(table_name, pk_value, Some(schema), false, true)
}
pub fn get_row_by_typed_pk_for_write_with_schema(
&self,
table_name: &str,
pk_value: &crate::Value,
schema: &crate::Schema,
) -> Result<Option<Tuple>> {
self.get_row_by_pk_inner(table_name, pk_value, Some(schema), false, false)
}
pub fn get_row_by_id(&self, table_name: &str, row_id: u64, schema: &crate::Schema) -> Result<Option<Tuple>> {
if let Some(cached) = self.row_cache.get(table_name, row_id) {
return Ok(Some(cached));
}
let storage_key = self.branch_aware_data_key(table_name, row_id);
let raw_value = match self.get(&storage_key)? {
Some(v) => v,
None => return Ok(None),
};
let mut tuple: Tuple = bincode::deserialize(&raw_value)
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
tuple.row_id = Some(row_id);
for (idx, column) in schema.columns.iter().enumerate() {
if idx >= tuple.values.len() {
break;
}
#[allow(clippy::indexing_slicing)]
match column.storage_mode {
ColumnStorageMode::Dictionary => {
if let Some(crate::Value::DictRef { dict_id }) = tuple.values.get(idx) {
let dict_id = *dict_id;
let s = self.dict_manager.decode(&self.db, table_name, &column.name, dict_id)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = crate::Value::String(s);
}
}
}
ColumnStorageMode::ContentAddressed => {
if let Some(crate::Value::CasRef { hash }) = tuple.values.get(idx) {
let hash = hash.clone();
let resolved = ContentAddressedStore::resolve(&self.db, &hash, &column.data_type)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = resolved;
}
}
}
ColumnStorageMode::Columnar => {
if matches!(tuple.values.get(idx), Some(crate::Value::ColumnarRef)) {
if let Some(val) = ColumnarStore::get(&self.db, table_name, &column.name, row_id)? {
if let Some(slot) = tuple.values.get_mut(idx) {
*slot = val;
}
}
}
}
ColumnStorageMode::Default => {}
}
}
self.row_cache.put(table_name, row_id, tuple.clone());
Ok(Some(tuple))
}
fn get_row_by_pk_inner(
&self,
table_name: &str,
pk_value: &crate::Value,
schema: Option<&crate::Schema>,
populate_cache_on_miss: bool,
coerce_pk_value: bool,
) -> Result<Option<Tuple>> {
let trace_lookup = tracing::enabled!(tracing::Level::DEBUG);
let lookup_start = trace_lookup.then(std::time::Instant::now);
let coerced: crate::Value;
let effective_pk = if coerce_pk_value {
if let Some(s) = schema {
if let Some(pk_col) = s.columns.iter().find(|c| c.primary_key) {
coerced = Self::coerce_pk_value(pk_value, &pk_col.data_type);
&coerced
} else {
pk_value
}
} else {
pk_value
}
} else {
pk_value
};
let key = super::art_manager::ArtIndexManager::encode_key(std::slice::from_ref(effective_pk));
let row_id = match self.art_index_manager.pk_index_lookup(table_name, &key) {
Some(rid) => rid,
None => {
if trace_lookup {
tracing::debug!(
phase = "index_lookup",
table = table_name,
duration_us = lookup_start.as_ref().map(|s| s.elapsed().as_micros() as u64).unwrap_or(0),
"PK index lookup: no match"
);
}
return Ok(None);
}
};
if let Some(cached) = self.row_cache.get(table_name, row_id) {
if trace_lookup {
tracing::debug!(
phase = "index_lookup",
table = table_name,
duration_us = lookup_start.as_ref().map(|s| s.elapsed().as_micros() as u64).unwrap_or(0),
cache = "hit",
"PK point lookup: row cache hit"
);
}
return Ok(Some(cached));
}
let storage_key = self.branch_aware_data_key(table_name, row_id);
let raw_value = match self.get(&storage_key)? {
Some(v) => v,
None => return Ok(None), };
let mut tuple: Tuple = bincode::deserialize(&raw_value)
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
tuple.row_id = Some(row_id);
if let Some(schema) = schema {
for (idx, column) in schema.columns.iter().enumerate() {
if idx >= tuple.values.len() {
break;
}
match column.storage_mode {
ColumnStorageMode::Dictionary => {
if let Some(crate::Value::DictRef { dict_id }) = tuple.values.get(idx) {
let dict_id = *dict_id;
let s = self.dict_manager.decode(&self.db, table_name, &column.name, dict_id)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = crate::Value::String(s);
}
}
}
ColumnStorageMode::ContentAddressed => {
if let Some(crate::Value::CasRef { hash }) = tuple.values.get(idx) {
let hash = hash.clone();
let resolved = ContentAddressedStore::resolve(&self.db, &hash, &column.data_type)?;
if let Some(val) = tuple.values.get_mut(idx) {
*val = resolved;
}
}
}
ColumnStorageMode::Columnar => {
if matches!(tuple.values.get(idx), Some(crate::Value::ColumnarRef)) {
if let Some(val) = ColumnarStore::get(&self.db, table_name, &column.name, row_id)? {
if let Some(slot) = tuple.values.get_mut(idx) {
*slot = val;
}
}
}
}
ColumnStorageMode::Default => {}
}
}
}
if populate_cache_on_miss {
self.row_cache.put(table_name, row_id, tuple.clone());
}
if trace_lookup {
tracing::debug!(
phase = "index_lookup",
table = table_name,
duration_us = lookup_start.as_ref().map(|s| s.elapsed().as_micros() as u64).unwrap_or(0),
cache = "miss",
"PK point lookup: fetched from storage, cached"
);
}
Ok(Some(tuple))
}
fn coerce_pk_value(value: &crate::Value, target: &crate::DataType) -> crate::Value {
use crate::{DataType, Value};
match (value, target) {
(Value::Int2(v), DataType::Int8) => Value::Int8(i64::from(*v)),
(Value::Int4(v), DataType::Int8) => Value::Int8(i64::from(*v)),
(Value::Int2(v), DataType::Int4) => Value::Int4(i32::from(*v)),
(Value::Int8(v), DataType::Int4) => Value::Int4(*v as i32),
(Value::Int8(v), DataType::Int2) => Value::Int2(*v as i16),
(Value::Int4(v), DataType::Int2) => Value::Int2(*v as i16),
(Value::String(s), DataType::Int8) => s.parse::<i64>().map(Value::Int8).unwrap_or_else(|_| value.clone()),
(Value::String(s), DataType::Int4) => s.parse::<i32>().map(Value::Int4).unwrap_or_else(|_| value.clone()),
(Value::String(s), DataType::Int2) => s.parse::<i16>().map(Value::Int2).unwrap_or_else(|_| value.clone()),
_ => value.clone(),
}
}
pub fn scan_table_filtered(
&self,
table_name: &str,
predicates: &[AnalyzedPredicate],
limit: Option<usize>,
) -> Result<Vec<Tuple>> {
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
let tuples = self.scan_table(table_name)?;
let filtered = self
.predicate_pushdown
.scan_with_pushdown(table_name, tuples, predicates, &schema, limit);
Ok(filtered)
}
pub fn migrate_column_storage(
&self,
table_name: &str,
col_idx: usize,
column: &crate::Column,
old_mode: ColumnStorageMode,
new_mode: ColumnStorageMode,
) -> Result<usize> {
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut migrated = 0;
let mut keys_to_migrate = Vec::new();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self.db.iterator_opt(IteratorMode::Start, read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if key.starts_with(prefix_bytes) {
keys_to_migrate.push(key.to_vec());
} else if key.first() > prefix_bytes.first() {
break;
}
}
for key in keys_to_migrate {
let raw_value = self
.db
.get(&key)
.map_err(|e| Error::storage(format!("Failed to read row: {}", e)))?
.ok_or_else(|| Error::storage("Row disappeared during migration"))?;
let mut tuple: Tuple = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
bincode::deserialize(&decrypted)
} else {
bincode::deserialize(&raw_value)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
if col_idx >= tuple.values.len() {
continue;
}
let row_id = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()).unwrap_or(0);
let cur_val = tuple
.values
.get(col_idx)
.ok_or_else(|| Error::internal("col_idx out of bounds during migration decode"))?;
let original_value = match old_mode {
ColumnStorageMode::Dictionary => {
if let crate::Value::DictRef { dict_id } = cur_val {
let s = self.dict_manager.decode(&self.db, table_name, &column.name, *dict_id)?;
crate::Value::String(s)
} else {
cur_val.clone()
}
}
ColumnStorageMode::ContentAddressed => {
if let crate::Value::CasRef { hash } = cur_val {
ContentAddressedStore::resolve(&self.db, hash, &column.data_type)?
} else {
cur_val.clone()
}
}
ColumnStorageMode::Columnar => {
if matches!(cur_val, crate::Value::ColumnarRef) {
ColumnarStore::get(&self.db, table_name, &column.name, row_id)?.unwrap_or(crate::Value::Null)
} else {
cur_val.clone()
}
}
ColumnStorageMode::Default => cur_val.clone(),
};
let new_val = match new_mode {
ColumnStorageMode::Dictionary => {
if let crate::Value::String(s) = &original_value {
let dict_id = self.dict_manager.encode(&self.db, table_name, &column.name, s)?;
crate::Value::DictRef { dict_id }
} else {
original_value
}
}
ColumnStorageMode::ContentAddressed => {
match &original_value {
crate::Value::String(_) | crate::Value::Bytes(_) => {
ContentAddressedStore::store(&self.db, &original_value)?
}
_ => original_value,
}
}
ColumnStorageMode::Columnar => {
ColumnarStore::store(&self.db, table_name, &column.name, row_id, original_value)?;
crate::Value::ColumnarRef
}
ColumnStorageMode::Default => original_value,
};
*tuple
.values
.get_mut(col_idx)
.ok_or_else(|| Error::internal("col_idx out of bounds during migration encode"))? = new_val;
if old_mode == ColumnStorageMode::Columnar && new_mode != ColumnStorageMode::Columnar {
ColumnarStore::delete(&self.db, table_name, &column.name, row_id)?;
}
let new_value =
bincode::serialize(&tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?;
let final_value = if let Some(km) = &self.key_manager {
crypto::encrypt(km.key(), &new_value)?
} else {
new_value
};
self.db
.put(&key, &final_value)
.map_err(|e| Error::storage(format!("Failed to write migrated row: {}", e)))?;
migrated += 1;
}
if new_mode == ColumnStorageMode::Dictionary {
self.dict_manager.flush(&self.db)?;
}
if old_mode == ColumnStorageMode::Columnar || new_mode == ColumnStorageMode::Columnar {
ColumnarStore::purge_presence(&self.db, table_name)?;
}
tracing::info!(
"Migrated {} rows in {}.{} from {:?} to {:?}",
migrated,
table_name,
column.name,
old_mode,
new_mode
);
Ok(migrated)
}
pub fn add_column_to_rows(
&self,
table_name: &str,
default_expr: &Option<crate::sql::LogicalExpr>,
) -> Result<usize> {
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut updated = 0;
let default_value = if let Some(expr) = default_expr {
match expr {
crate::sql::LogicalExpr::Literal(v) => v.clone(),
_ => crate::Value::Null, }
} else {
crate::Value::Null
};
let mut keys_to_update = Vec::new();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self.db.iterator_opt(IteratorMode::Start, read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if key.starts_with(prefix_bytes) {
keys_to_update.push(key.to_vec());
} else if key.first() > prefix_bytes.first() {
break;
}
}
for key in keys_to_update {
let raw_value = self
.db
.get(&key)
.map_err(|e| Error::storage(format!("Failed to read row: {}", e)))?
.ok_or_else(|| Error::storage("Row disappeared during update"))?;
let mut tuple: Tuple = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
bincode::deserialize(&decrypted)
} else {
bincode::deserialize(&raw_value)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
tuple.values.push(default_value.clone());
let new_value =
bincode::serialize(&tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?;
let final_value = if let Some(km) = &self.key_manager {
crypto::encrypt(km.key(), &new_value)?
} else {
new_value
};
self.db
.put(&key, &final_value)
.map_err(|e| Error::storage(format!("Failed to write updated row: {}", e)))?;
updated += 1;
}
Ok(updated)
}
pub fn drop_column_from_rows(&self, table_name: &str, col_idx: usize) -> Result<usize> {
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut updated = 0;
let mut keys_to_update = Vec::new();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self.db.iterator_opt(IteratorMode::Start, read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if key.starts_with(prefix_bytes) {
keys_to_update.push(key.to_vec());
} else if key.first() > prefix_bytes.first() {
break;
}
}
for key in keys_to_update {
let raw_value = self
.db
.get(&key)
.map_err(|e| Error::storage(format!("Failed to read row: {}", e)))?
.ok_or_else(|| Error::storage("Row disappeared during update"))?;
let mut tuple: Tuple = if let Some(km) = &self.key_manager {
let decrypted = crypto::decrypt(km.key(), &raw_value)?;
bincode::deserialize(&decrypted)
} else {
bincode::deserialize(&raw_value)
}
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
if col_idx < tuple.values.len() {
tuple.values.remove(col_idx);
let new_value = bincode::serialize(&tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?;
let final_value = if let Some(km) = &self.key_manager {
crypto::encrypt(km.key(), &new_value)?
} else {
new_value
};
self.db
.put(&key, &final_value)
.map_err(|e| Error::storage(format!("Failed to write updated row: {}", e)))?;
updated += 1;
}
}
Ok(updated)
}
pub fn rename_table(&self, old_name: &str, new_name: &str) -> Result<()> {
self.catalog().rename_table(old_name, new_name)
}
pub fn register_bloom_filters(&self, table_name: &str, filters: TableBloomFilters) {
self.predicate_pushdown
.register_bloom_filters(table_name.to_string(), filters);
}
pub fn register_zone_maps(&self, table_name: &str, zone_map: TableZoneMap) {
self.predicate_pushdown
.register_zone_maps(table_name.to_string(), zone_map);
}
pub fn predicate_pushdown(&self) -> &PredicatePushdownManager {
&self.predicate_pushdown
}
pub fn predicate_pushdown_stats(&self) -> super::predicate_pushdown::PushdownStats {
self.predicate_pushdown.get_stats()
}
pub fn suspend_smfi_for_bulk_load(
&self,
table_name: &str,
reason: super::filter_index_delta::BulkLoadReason,
) -> super::filter_index_delta::BulkLoadGuard<'_> {
self.filter_delta_tracker.suspend_table(table_name, reason)
}
pub fn is_smfi_suspended(&self, table_name: &str) -> bool {
self.filter_delta_tracker.is_suspended(table_name)
}
pub fn smfi_stats(&self) -> super::filter_index_delta::FilterDeltaStats {
self.filter_delta_tracker.stats()
}
pub fn set_smfi_enabled(&self, enabled: bool) {
self.filter_delta_tracker.set_enabled(enabled);
}
pub fn is_smfi_enabled(&self) -> bool {
self.filter_delta_tracker.is_enabled()
}
pub fn smfi_bulk_load_threshold(&self) -> usize {
self.filter_delta_tracker.bulk_load_threshold()
}
pub fn set_smfi_bulk_load_threshold(&self, threshold: usize) {
self.filter_delta_tracker.set_bulk_load_threshold(threshold);
}
pub fn build_bloom_filters_for_table(&self, table_name: &str) -> Result<()> {
use super::bloom_filter::TableBloomFilters;
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
let tuples = self.scan_table(table_name)?;
if tuples.is_empty() {
return Ok(());
}
let mut table_filters = TableBloomFilters::new(table_name.to_string(), tuples.len());
table_filters.build_from_tuples(&tuples, &schema);
self.register_bloom_filters(table_name, table_filters);
debug!(
"Built bloom filters for table '{}' with {} tuples",
table_name,
tuples.len()
);
Ok(())
}
pub fn build_zone_maps_for_table(&self, table_name: &str, block_size: usize) -> Result<()> {
use super::zone_map::TableZoneMap;
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
let tuples = self.scan_table(table_name)?;
if tuples.is_empty() {
return Ok(());
}
let mut zone_map = TableZoneMap::new(table_name.to_string(), block_size);
zone_map.build_from_tuples(&tuples, &schema);
self.register_zone_maps(table_name, zone_map);
debug!(
"Built zone maps for table '{}' with {} tuples, block_size {}",
table_name,
tuples.len(),
block_size
);
Ok(())
}
pub fn catalog(&self) -> Catalog<'_> {
Catalog::new(self)
}
pub fn trigger_registry(&self) -> &Arc<crate::sql::TriggerRegistry> {
&self.trigger_registry
}
pub fn load_triggers(&self) -> Result<()> {
let catalog = self.catalog();
let triggers = catalog.load_all_triggers()?;
info!("Loading {} triggers from persistent storage", triggers.len());
for trigger in triggers {
if let Err(e) = self.trigger_registry.register_trigger(trigger.clone()) {
warn!(
"Failed to load trigger '{}' on table '{}': {}",
trigger.name, trigger.table_name, e
);
} else {
debug!("Loaded trigger '{}' on table '{}'", trigger.name, trigger.table_name);
}
}
Ok(())
}
pub fn vector_indexes(&self) -> &VectorIndexManager {
&self.vector_indexes
}
pub fn branch_manager(&self) -> Option<Arc<BranchManager>> {
self.branch_manager.read().as_ref().map(Arc::clone)
}
pub fn db(&self) -> Arc<rocksdb::DB> {
Arc::clone(&self.db)
}
pub fn timestamp(&self) -> Arc<RwLock<u64>> {
Arc::clone(&self.timestamp)
}
pub fn git_integration_manager(&self) -> Result<crate::git_integration::GitIntegrationManager> {
let branch_manager = self
.branch_manager()
.ok_or_else(|| Error::config("Branching must be enabled for Git integration".to_string()))?;
crate::git_integration::GitIntegrationManager::new(self.db(), branch_manager, self.timestamp())
}
pub fn ddl_versioning_manager(&self) -> Result<crate::git_integration::ddl_versioning::DdlVersioningManager> {
crate::git_integration::ddl_versioning::DdlVersioningManager::new(self.db(), self.timestamp())
}
pub fn log_ddl_to_git_history(
&self,
operation: &str,
object_type: &str,
object_name: &str,
ddl_statement: &str,
) -> Result<()> {
use crate::git_integration::ddl_versioning::{DdlObjectType, DdlOperation};
let branch_id: u64 = 0;
let lsn = self.wal.as_ref().map(|w| w.read().current_lsn()).unwrap_or(0);
let op = match operation.to_uppercase().as_str() {
"CREATE" => DdlOperation::Create,
"ALTER" => DdlOperation::Alter,
"DROP" => DdlOperation::Drop,
"TRUNCATE" => DdlOperation::Truncate,
"RENAME" => DdlOperation::Rename,
"COMMENT" => DdlOperation::Comment,
_ => DdlOperation::Create, };
let obj_type = match object_type.to_uppercase().as_str() {
"TABLE" => DdlObjectType::Table,
"INDEX" => DdlObjectType::Index,
"VIEW" => DdlObjectType::View,
"MATERIALIZED VIEW" | "MATERIALIZED_VIEW" => DdlObjectType::MaterializedView,
"SEQUENCE" => DdlObjectType::Sequence,
"FUNCTION" => DdlObjectType::Function,
"PROCEDURE" => DdlObjectType::Procedure,
"TRIGGER" => DdlObjectType::Trigger,
"CONSTRAINT" => DdlObjectType::Constraint,
"SCHEMA" => DdlObjectType::Schema,
"EXTENSION" => DdlObjectType::Extension,
"TYPE" => DdlObjectType::Type,
_ => DdlObjectType::Table, };
if let Ok(ddl_mgr) = self.ddl_versioning_manager() {
if let Err(e) = ddl_mgr.record_ddl(
branch_id,
lsn,
op,
obj_type,
object_name,
ddl_statement,
None, None, ) {
tracing::debug!("Failed to record DDL to Git history: {}", e);
}
}
Ok(())
}
pub fn bloom_manager(&self) -> Option<()> {
None
}
pub fn zone_map_manager(&self) -> Option<()> {
None
}
pub fn dictionary_store(&self) -> Option<()> {
None
}
pub fn content_store(&self) -> Option<()> {
None
}
pub fn columnar_store(&self) -> Option<()> {
None
}
pub fn columnar_column_stats(&self, table_name: &str, column_name: &str) -> Result<super::ColumnarStats> {
ColumnarStore::stats(&self.db, table_name, column_name)
}
pub fn close(self) -> Result<()> {
Ok(())
}
pub fn flush(&self) -> Result<()> {
self.db
.flush()
.map_err(|e| Error::storage(format!("Flush failed: {}", e)))
}
pub fn hnsw_snapshot_dir(&self) -> Option<std::path::PathBuf> {
self.db_path
.as_ref()
.map(|p| p.join(super::index_snapshot::HNSW_SNAPSHOT_DIR))
}
fn wire_index_snapshot_hooks(&self) {
let db = Arc::clone(&self.db);
self.art_index_manager.set_snapshot_invalidation_hook(Arc::new(move || {
Self::delete_prefix_raw(&db, super::index_snapshot::ART_SNAPSHOT_MARKER_PREFIX.as_bytes());
}));
let db = Arc::clone(&self.db);
self.vector_indexes.set_snapshot_invalidation_hook(Arc::new(move || {
Self::delete_prefix_raw(&db, super::index_snapshot::VEC_SNAPSHOT_MARKER_PREFIX.as_bytes());
}));
}
fn delete_prefix_raw(db: &DB, prefix: &[u8]) {
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = db.iterator_opt(IteratorMode::From(prefix, rocksdb::Direction::Forward), read_opts);
let mut batch = WriteBatch::default();
for item in iter {
let Ok((key, _)) = item else { break };
if !key.starts_with(prefix) {
break;
}
batch.delete(key);
}
if let Err(e) = db.write(batch) {
warn!("index snapshot marker invalidation failed: {}", e);
}
}
pub fn invalidate_art_snapshot_markers(&self) {
Self::delete_prefix_raw(&self.db, super::index_snapshot::ART_SNAPSHOT_MARKER_PREFIX.as_bytes());
}
pub fn invalidate_vector_snapshot_markers(&self) {
Self::delete_prefix_raw(&self.db, super::index_snapshot::VEC_SNAPSHOT_MARKER_PREFIX.as_bytes());
}
pub fn set_index_snapshots_on_close(&self, enabled: bool) {
self.index_snapshots_on_close.store(enabled, Ordering::Release);
}
pub fn index_snapshots_on_close(&self) -> bool {
self.index_snapshots_on_close.load(Ordering::Acquire)
}
pub(crate) fn set_last_index_open_report(&self, report: super::index_snapshot::IndexOpenReport) {
*self.last_index_open_report.write() = Some(report);
}
pub fn last_index_open_report(&self) -> Option<super::index_snapshot::IndexOpenReport> {
self.last_index_open_report.read().clone()
}
pub fn persist_index_snapshots(&self) -> Result<super::index_snapshot::IndexSnapshotPersistReport> {
use super::index_snapshot as snap;
let started = std::time::Instant::now();
let mut report = snap::IndexSnapshotPersistReport::default();
let Some(dump_dir) = self.hnsw_snapshot_dir() else {
return Ok(report);
};
self.art_index_manager.mark_snapshot_clean();
self.vector_indexes.mark_snapshot_clean();
let catalog = super::Catalog::new(self);
let marker = snap::SnapshotMarker::current();
let marker_bytes = bincode::serialize(&marker).map_err(|e| Error::storage(format!("marker serialize: {e}")))?;
for table_name in catalog.list_tables()? {
if table_name.starts_with("helios_") {
continue;
}
let indexes = self.art_index_manager.export_table_snapshot(&table_name);
report.art_entries += indexes
.iter()
.map(|i| i.entries.iter().map(|(_, ids)| ids.len() as u64).sum::<u64>())
.sum::<u64>();
let table_snap = snap::ArtTableSnapshot {
format_version: snap::INDEX_SNAPSHOT_FORMAT_VERSION,
key_encoding_version: snap::ART_KEY_ENCODING_VERSION,
indexes,
};
let blob =
bincode::serialize(&table_snap).map_err(|e| Error::storage(format!("snapshot serialize: {e}")))?;
self.put(&snap::art_snapshot_key(&table_name), &blob)?;
self.put(&snap::art_snapshot_marker_key(&table_name), &marker_bytes)?;
report.art_tables += 1;
}
let vector_metas = self.vector_indexes.list_all_metadata();
if vector_metas
.iter()
.any(|m| matches!(m.index_type, super::VectorIndexType::Standard(_)))
{
if let Err(e) = std::fs::create_dir_all(&dump_dir) {
warn!("could not create HNSW snapshot dir {:?}: {}", dump_dir, e);
}
}
for meta in vector_metas {
if !matches!(meta.index_type, super::VectorIndexType::Standard(_)) {
continue; }
match self.vector_indexes.dump_standard_index(&meta.name, &dump_dir) {
Ok(Some(sidecar)) => {
let blob =
bincode::serialize(&sidecar).map_err(|e| Error::storage(format!("sidecar serialize: {e}")))?;
self.put(&snap::vec_snapshot_key(&meta.name), &blob)?;
self.put(&snap::vec_snapshot_marker_key(&meta.name), &marker_bytes)?;
report.vector_graphs += 1;
}
Ok(None) => {}
Err(e) => {
warn!(
"HNSW snapshot dump for '{}' failed: {} — it will rebuild on open",
meta.name, e
);
report.vector_dump_failures += 1;
}
}
}
if !self.art_index_manager.snapshot_is_clean() {
self.invalidate_art_snapshot_markers();
}
if !self.vector_indexes.snapshot_is_clean() {
self.invalidate_vector_snapshot_markers();
}
report.persisted = true;
report.elapsed_ms = started.elapsed().as_secs_f64() * 1000.0;
info!(
"index snapshot checkpoint: {} tables ({} ART entries), {} HNSW graphs, {:.1}ms",
report.art_tables, report.art_entries, report.vector_graphs, report.elapsed_ms
);
Ok(report)
}
pub(crate) fn load_art_table_snapshots(
&self,
) -> std::collections::HashMap<String, super::index_snapshot::ArtTableSnapshot> {
use super::index_snapshot as snap;
let mut out = std::collections::HashMap::new();
for table in self.keys_with_prefix_suffixes(snap::ART_SNAPSHOT_MARKER_PREFIX) {
let Ok(Some(marker_bytes)) = self.get(&snap::art_snapshot_marker_key(&table)) else {
continue;
};
let Ok(marker) = bincode::deserialize::<snap::SnapshotMarker>(&marker_bytes) else {
continue;
};
if !marker.matches_current() {
continue;
}
let Ok(Some(blob)) = self.get(&snap::art_snapshot_key(&table)) else {
continue;
};
let Ok(table_snap) = bincode::deserialize::<snap::ArtTableSnapshot>(&blob) else {
warn!(
"ART snapshot for '{}' is unreadable — falling back to scan rebuild",
table
);
continue;
};
if table_snap.format_version != snap::INDEX_SNAPSHOT_FORMAT_VERSION
|| table_snap.key_encoding_version != snap::ART_KEY_ENCODING_VERSION
{
continue;
}
out.insert(table, table_snap);
}
out
}
pub(crate) fn load_vector_sidecars(
&self,
) -> std::collections::HashMap<String, super::index_snapshot::VectorGraphSidecar> {
use super::index_snapshot as snap;
let mut out = std::collections::HashMap::new();
for index in self.keys_with_prefix_suffixes(snap::VEC_SNAPSHOT_MARKER_PREFIX) {
let Ok(Some(marker_bytes)) = self.get(&snap::vec_snapshot_marker_key(&index)) else {
continue;
};
let Ok(marker) = bincode::deserialize::<snap::SnapshotMarker>(&marker_bytes) else {
continue;
};
if !marker.matches_current() {
continue;
}
let Ok(Some(blob)) = self.get(&snap::vec_snapshot_key(&index)) else {
continue;
};
let Ok(sidecar) = bincode::deserialize::<snap::VectorGraphSidecar>(&blob) else {
warn!("HNSW sidecar for '{}' is unreadable — falling back to rebuild", index);
continue;
};
out.insert(index, sidecar);
}
out
}
fn keys_with_prefix_suffixes(&self, prefix: &str) -> Vec<String> {
let prefix_bytes = prefix.as_bytes();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self
.db
.iterator_opt(IteratorMode::From(prefix_bytes, rocksdb::Direction::Forward), read_opts);
let mut out = Vec::new();
for item in iter {
let Ok((key, _)) = item else { break };
if !key.starts_with(prefix_bytes) {
break;
}
out.push(String::from_utf8_lossy(key.get(prefix_bytes.len()..).unwrap_or_default()).to_string());
}
out
}
pub fn get_stats(&self) -> Result<String> {
self.db
.property_value("rocksdb.stats")
.map_err(|e| Error::storage(format!("Failed to get stats: {}", e)))?
.ok_or_else(|| Error::storage("Stats not available"))
}
pub fn is_encrypted(&self) -> bool {
self.key_manager.is_some()
}
pub fn encryption_info(&self) -> Option<String> {
self.key_manager
.as_ref()
.map(|km| format!("Enabled (AES-256-GCM, source: {:?})", km.source()))
}
fn get_or_init_branch_manager(&self) -> Result<Arc<RwLock<Option<Arc<BranchManager>>>>> {
let manager = self.branch_manager.read();
if manager.is_none() {
return Err(Error::storage(
"BranchManager not initialized. Branch operations are unavailable.",
));
}
drop(manager);
Ok(Arc::clone(&self.branch_manager))
}
pub fn create_branch(&self, name: &str, parent_name: Option<&str>, options: BranchOptions) -> Result<BranchId> {
self.create_branch_at_snapshot(name, parent_name, None, options)
}
pub fn create_branch_at_snapshot(
&self,
name: &str,
parent_name: Option<&str>,
snapshot_id: Option<u64>,
options: BranchOptions,
) -> Result<BranchId> {
let manager_lock = self.get_or_init_branch_manager()?;
let manager = manager_lock.read();
let mgr = manager
.as_ref()
.ok_or_else(|| Error::storage("BranchManager not available in read lock"))?;
let snapshot = snapshot_id.unwrap_or_else(|| self.next_timestamp());
mgr.create_branch(name, parent_name, snapshot, options)
}
pub fn drop_branch(&self, name: &str, if_exists: bool) -> Result<()> {
let manager_lock = self.get_or_init_branch_manager()?;
let manager = manager_lock.read();
let mgr = manager
.as_ref()
.ok_or_else(|| Error::storage("BranchManager not available in read lock"))?;
mgr.drop_branch(name, if_exists)
}
pub fn get_branch(&self, name: &str) -> Result<BranchMetadata> {
let manager_lock = self.get_or_init_branch_manager()?;
let manager = manager_lock.read();
let mgr = manager
.as_ref()
.ok_or_else(|| Error::storage("BranchManager not available in read lock"))?;
mgr.get_branch_by_name(name)
}
pub fn list_branches(&self) -> Result<Vec<BranchMetadata>> {
let manager_lock = self.get_or_init_branch_manager()?;
let manager = manager_lock.read();
let mgr = manager
.as_ref()
.ok_or_else(|| Error::storage("BranchManager not available in read lock"))?;
mgr.list_branches()
}
pub fn merge_branch(
&self,
source_name: &str,
target_name: &str,
_strategy: super::MergeStrategy,
) -> Result<super::MergeResult> {
use std::collections::HashSet;
let manager_lock = self.get_or_init_branch_manager()?;
let source_id;
let target_id;
{
let manager = manager_lock.read();
let mgr = manager
.as_ref()
.ok_or_else(|| Error::storage("BranchManager not initialized"))?;
let source = mgr.get_branch_by_name(source_name)?;
let target = mgr.get_branch_by_name(target_name)?;
if source.state != super::BranchState::Active {
return Err(Error::branch_merge(format!(
"Source branch '{}' is not active",
source_name
)));
}
if target.state != super::BranchState::Active {
return Err(Error::branch_merge(format!(
"Target branch '{}' is not active",
target_name
)));
}
source_id = source.branch_id;
target_id = target.branch_id;
}
let merge_to_main = target_name == "main" || target_id == 1;
let merge_timestamp = self.next_timestamp();
let mut merged_keys = 0usize;
let catalog = Catalog::new(self);
let tables = catalog.list_tables()?;
let mut deleted_rows_by_table: std::collections::HashMap<String, HashSet<u64>> =
std::collections::HashMap::new();
for table_name in &tables {
let delete_prefix = format!("bdel:{}:{}:", source_id, table_name);
let delete_prefix_bytes = delete_prefix.as_bytes();
let iter = self.db.iterator(rocksdb::IteratorMode::From(
delete_prefix_bytes,
rocksdb::Direction::Forward,
));
for item in iter {
let (key, _value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(delete_prefix_bytes) {
break;
}
if let Some(row_id) = Self::parse_row_id_after_prefix(&key, delete_prefix_bytes.len()) {
deleted_rows_by_table
.entry(table_name.clone())
.or_default()
.insert(row_id);
}
}
}
for table_name in &tables {
let branch_prefix = format!("bdata:{}:{}:", source_id, table_name);
let branch_prefix_bytes = branch_prefix.as_bytes();
let iter = self.db.iterator(rocksdb::IteratorMode::From(
branch_prefix_bytes,
rocksdb::Direction::Forward,
));
for item in iter {
let (key, value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(branch_prefix_bytes) {
break;
}
if let Some(row_id) = Self::parse_row_id_after_prefix(&key, branch_prefix_bytes.len()) {
let target_key = if merge_to_main {
format!("data:{}:{}", table_name, row_id)
} else {
format!("bdata:{}:{}:{}", target_id, table_name, row_id)
};
self.db
.put(target_key.as_bytes(), &value)
.map_err(|e| Error::storage(format!("Failed to merge data: {}", e)))?;
merged_keys += 1;
}
}
}
if merge_to_main {
for (table_name, row_ids) in &deleted_rows_by_table {
for row_id in row_ids {
let target_key = format!("data:{}:{}", table_name, row_id);
self.db
.delete(target_key.as_bytes())
.map_err(|e| Error::storage(format!("Failed to apply delete: {}", e)))?;
}
}
} else {
for (table_name, row_ids) in &deleted_rows_by_table {
for row_id in row_ids {
let target_key = format!("bdel:{}:{}:{}", target_id, table_name, row_id);
self.db
.put(target_key.as_bytes(), b"")
.map_err(|e| Error::storage(format!("Failed to copy delete marker: {}", e)))?;
}
}
}
{
let manager = manager_lock.read();
let mgr = manager
.as_ref()
.ok_or_else(|| Error::storage("BranchManager not initialized"))?;
let mut source = mgr.get_branch_by_name(source_name)?;
source.state = super::BranchState::Merged {
into_branch: target_id,
at_timestamp: merge_timestamp,
};
let meta_key = format!("branch:meta:{}", source_name);
let meta_value = bincode::serialize(&source)
.map_err(|e| Error::storage(format!("Failed to serialize metadata: {}", e)))?;
self.db
.put(meta_key.as_bytes(), &meta_value)
.map_err(|e| Error::storage(format!("Failed to save merged branch metadata: {}", e)))?;
}
tracing::info!(
"Merge completed: {} -> {}, {} keys merged",
source_name,
target_name,
merged_keys
);
Ok(super::MergeResult {
merge_timestamp,
merged_keys,
conflicts: Vec::new(),
completed: true,
})
}
pub fn begin_branch_transaction(&self, branch_name: &str) -> Result<BranchTransaction> {
let manager_lock = self.get_or_init_branch_manager()?;
let manager = manager_lock.read();
let mgr = manager
.as_ref()
.ok_or_else(|| Error::storage("BranchManager not initialized"))?;
let branch = mgr.get_branch_by_name(branch_name)?;
let parent_chain = mgr.build_parent_chain(branch.branch_id)?;
let snapshot_id = self.next_timestamp();
BranchTransaction::new(
Arc::clone(&self.db),
branch.branch_id,
branch,
parent_chain,
snapshot_id,
Arc::clone(&self.snapshot_manager),
)
}
pub fn get_branch_metadata(&self, name: &str) -> Result<BranchMetadata> {
self.get_branch(name)
}
pub fn get_branch_name(&self, branch_id: BranchId) -> Option<String> {
let manager_lock = self.get_or_init_branch_manager().ok()?;
let manager = manager_lock.read();
let mgr = manager.as_ref()?;
mgr.get_branch_name(branch_id)
}
pub fn use_branch(&self, branch_name: &str) -> Result<()> {
if branch_name == "main" {
self.set_current_branch(None);
tracing::info!("Switched to main branch (branch isolation disabled)");
return Ok(());
}
let _metadata = self.get_branch(branch_name)?;
self.set_current_branch(Some(branch_name.to_string()));
tracing::info!("Switched to branch '{}' (branch isolation enabled)", branch_name);
Ok(())
}
pub fn current_timestamp(&self) -> u64 {
*self.timestamp.read()
}
pub fn mv_catalog(&self) -> super::MaterializedViewCatalog<'_> {
super::MaterializedViewCatalog::new(self)
}
pub fn view_catalog(&self) -> super::ViewCatalog<'_> {
super::ViewCatalog::new(self)
}
pub fn mv_delta_tracker(&self) -> &Arc<super::MvDeltaTracker> {
&self.mv_delta_tracker
}
pub fn is_wal_enabled(&self) -> bool {
self.wal.is_some()
}
pub fn logical_wal_per_statement(&self) -> bool {
self.config.storage.logical_wal_per_statement
}
#[cfg(test)]
pub(crate) fn wal_entries_for_tests(&self) -> Result<Vec<super::wal::WalEntry>> {
match &self.wal {
Some(wal) => wal.read().replay(),
None => Ok(Vec::new()),
}
}
pub fn time_travel_enabled(&self) -> bool {
self.config.storage.time_travel_enabled
}
pub fn wal_lsn(&self) -> Option<u64> {
self.wal.as_ref().map(|wal| wal.read().current_lsn())
}
pub fn increment_lsn(&self) -> Option<u64> {
self.wal.as_ref().map(|wal| wal.read().increment_lsn())
}
pub fn flush_wal(&self) -> Result<()> {
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.flush()?;
}
Ok(())
}
pub fn log_create_table(&self, table_name: &str, schema: &crate::Schema) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let schema_bytes =
bincode::serialize(schema).map_err(|e| Error::storage(format!("Failed to serialize schema: {}", e)))?;
let wal = wal.read();
wal.append_nosync(WalOperation::CreateTable {
table: table_name.to_string(),
schema: schema_bytes,
})?;
}
Ok(())
}
pub fn log_drop_table(&self, table_name: &str) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append_nosync(WalOperation::DropTable {
table: table_name.to_string(),
})?;
}
Ok(())
}
pub fn log_truncate(&self, table_name: &str) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::Truncate {
table: table_name.to_string(),
})?;
}
Ok(())
}
pub fn log_alter_column_storage(
&self,
table_name: &str,
column_name: &str,
storage_mode: &crate::ColumnStorageMode,
) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
let storage_mode_bytes = bincode::serialize(storage_mode)
.map_err(|e| Error::storage(format!("Failed to serialize storage mode: {}", e)))?;
wal.append(WalOperation::AlterColumnStorage {
table: table_name.to_string(),
column: column_name.to_string(),
storage_mode: storage_mode_bytes,
})?;
}
Ok(())
}
pub fn log_create_index(
&self,
name: &str,
table: &str,
column: &str,
index_type: Option<&str>,
options: &[u8],
) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::CreateIndex {
name: name.to_string(),
table: table.to_string(),
column: column.to_string(),
index_type: index_type.map(String::from),
options: options.to_vec(),
})?;
}
Ok(())
}
pub fn log_drop_index(&self, name: &str) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::DropIndex { name: name.to_string() })?;
}
Ok(())
}
pub fn log_create_trigger(&self, name: &str, table: &str, definition: &[u8]) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::CreateTrigger {
name: name.to_string(),
table: table.to_string(),
definition: definition.to_vec(),
})?;
}
Ok(())
}
pub fn log_drop_trigger(&self, name: &str, table: Option<&str>) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::DropTrigger {
name: name.to_string(),
table: table.map(String::from),
})?;
}
Ok(())
}
pub fn log_create_function(&self, name: &str, definition: &[u8]) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::CreateFunction {
name: name.to_string(),
definition: definition.to_vec(),
})?;
}
Ok(())
}
pub fn log_drop_function(&self, name: &str) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::DropFunction { name: name.to_string() })?;
}
Ok(())
}
pub fn log_create_procedure(&self, name: &str, definition: &[u8]) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::CreateProcedure {
name: name.to_string(),
definition: definition.to_vec(),
})?;
}
Ok(())
}
pub fn log_drop_procedure(&self, name: &str) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::DropProcedure { name: name.to_string() })?;
}
Ok(())
}
pub fn log_create_materialized_view(&self, name: &str, definition: &[u8]) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::CreateMaterializedView {
name: name.to_string(),
definition: definition.to_vec(),
})?;
}
Ok(())
}
pub fn log_drop_materialized_view(&self, name: &str) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::DropMaterializedView { name: name.to_string() })?;
}
Ok(())
}
pub fn log_refresh_materialized_view(&self, name: &str, concurrent: bool, incremental: bool) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::RefreshMaterializedView {
name: name.to_string(),
concurrent,
incremental,
})?;
}
Ok(())
}
pub fn log_add_constraint(&self, table: &str, constraint: &[u8]) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::AddConstraint {
table: table.to_string(),
constraint: constraint.to_vec(),
})?;
}
Ok(())
}
pub fn log_drop_constraint(&self, table: &str, constraint_name: &str) -> Result<()> {
if self.is_replaying.load(std::sync::atomic::Ordering::Acquire) {
return Ok(());
}
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::DropConstraint {
table: table.to_string(),
constraint_name: constraint_name.to_string(),
})?;
}
Ok(())
}
pub fn replay_wal(&self) -> Result<usize> {
if let Some(wal) = &self.wal {
self.is_replaying.store(true, Ordering::Release);
let wal = wal.read();
let entries = wal.replay()?;
let count = entries.len();
if count == 0 {
info!("No WAL entries to replay");
self.is_replaying.store(false, Ordering::Release);
return Ok(0);
}
info!("Replaying {} WAL entries for crash recovery (optimized)", count);
let mut committed_transactions: std::collections::HashSet<u64> = std::collections::HashSet::new();
let mut aborted_transactions: std::collections::HashSet<u64> = std::collections::HashSet::new();
let mut replayed_count = 0;
let mut skipped_count = 0;
let mut error_count = 0;
for entry in &entries {
match &entry.operation {
WalOperation::Begin { tx_id } => {
debug!("Transaction {} started", tx_id);
}
WalOperation::Commit { tx_id } => {
committed_transactions.insert(*tx_id);
debug!("Transaction {} committed", tx_id);
}
WalOperation::Abort { tx_id } => {
aborted_transactions.insert(*tx_id);
debug!("Transaction {} aborted", tx_id);
}
_ => {}
}
}
const BATCH_SIZE: usize = 100;
let mut batch = WriteBatch::default();
let mut batch_count = 0;
for entry in entries {
if let Some(tx_id) = Self::extract_tx_id(&entry.operation) {
if aborted_transactions.contains(&tx_id) {
debug!("Skipping operation from aborted transaction {}", tx_id);
skipped_count += 1;
continue;
}
}
match self.apply_wal_operation_to_batch(&entry.operation, &mut batch) {
Ok(added) => {
if added {
batch_count += 1;
replayed_count += 1;
}
if batch_count >= BATCH_SIZE {
self.db
.write(batch)
.map_err(|e| Error::storage(format!("Batch write failed: {}", e)))?;
batch = WriteBatch::default();
batch_count = 0;
if replayed_count % 1000 == 0 {
debug!("Replayed {} operations...", replayed_count);
}
}
}
Err(e) => {
warn!("Error applying WAL operation at LSN {}: {}", entry.lsn, e);
error_count += 1;
if error_count > count / 10 {
self.is_replaying.store(false, Ordering::Release);
return Err(Error::storage(format!(
"Too many errors during WAL replay: {}/{}",
error_count, count
)));
}
}
}
}
if batch_count > 0 {
self.db
.write(batch)
.map_err(|e| Error::storage(format!("Final batch write failed: {}", e)))?;
}
info!(
"WAL replay complete: {} operations applied, {} skipped, {} errors",
replayed_count, skipped_count, error_count
);
self.is_replaying.store(false, Ordering::Release);
Ok(replayed_count)
} else {
Ok(0)
}
}
pub fn apply_replicated_operation(&self, operation: WalOperation) -> Result<()> {
self.is_replaying.store(true, std::sync::atomic::Ordering::Release);
let result = self.apply_wal_operation(operation);
self.is_replaying.store(false, std::sync::atomic::Ordering::Release);
result
}
fn apply_wal_operation(&self, operation: WalOperation) -> Result<()> {
info!(
"apply_wal_operation: Processing {:?}",
std::mem::discriminant(&operation)
);
match operation {
WalOperation::Insert { table, key, tuple } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(&table).is_err() {
debug!("Skipping insert for non-existent table: {}", table);
return Ok(());
}
self.put(&key, &tuple)?;
debug!("Replayed insert: table={}, key_len={}", table, key.len());
Ok(())
}
WalOperation::Update { table, key, tuple } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(&table).is_err() {
debug!("Skipping update for non-existent table: {}", table);
return Ok(());
}
self.put(&key, &tuple)?;
debug!("Replayed update: table={}, key_len={}", table, key.len());
Ok(())
}
WalOperation::Delete { table, key } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(&table).is_err() {
debug!("Skipping delete for non-existent table: {}", table);
return Ok(());
}
self.delete(&key)?;
debug!("Replayed delete: table={}, key_len={}", table, key.len());
Ok(())
}
WalOperation::CreateTable { table, schema } => {
info!(
"apply_wal_operation: CreateTable for '{}', schema_len={}",
table,
schema.len()
);
let catalog = Catalog::new(self);
if catalog.get_table_schema(&table).is_ok() {
info!("Table {} already exists, skipping create", table);
return Ok(());
}
info!("apply_wal_operation: Deserializing schema for table '{}'", table);
match bincode::deserialize::<crate::Schema>(&schema) {
Ok(schema_obj) => {
info!("apply_wal_operation: Schema deserialized, creating table '{}'", table);
catalog.create_table(&table, schema_obj)?;
info!("apply_wal_operation: Table '{}' created successfully", table);
Ok(())
}
Err(e) => {
error!("Failed to deserialize schema for table {}: {}", table, e);
Ok(())
}
}
}
WalOperation::DropTable { table } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(&table).is_err() {
debug!("Table {} doesn't exist, skipping drop", table);
return Ok(());
}
catalog.drop_table(&table)?;
debug!("Replayed drop table: {}", table);
Ok(())
}
WalOperation::Truncate { table } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(&table).is_err() {
debug!("Table {} doesn't exist, skipping truncate", table);
return Ok(());
}
let prefix = format!("data:{}:", table);
let prefix_bytes = prefix.as_bytes();
let mut keys_to_delete = Vec::new();
let iter = self.db.iterator(rocksdb::IteratorMode::Start);
for item in iter {
if let Ok((key, _)) = item {
if key.starts_with(prefix_bytes) {
keys_to_delete.push(key.to_vec());
} else if key.first() > prefix_bytes.first() {
break;
}
}
}
for key in keys_to_delete {
self.delete(&key)?;
}
debug!("Replayed truncate: table={}", table);
Ok(())
}
WalOperation::AlterColumnStorage {
table,
column,
storage_mode,
} => {
match bincode::deserialize::<crate::ColumnStorageMode>(&storage_mode) {
Ok(mode) => {
info!(
"Replayed alter column storage: table={}, column={}, mode={:?}",
table, column, mode
);
let key = format!("meta:col_storage:{}:{}", table, column).into_bytes();
self.put(&key, &storage_mode)?;
Ok(())
}
Err(e) => {
error!("Failed to deserialize column storage mode: {}", e);
Ok(())
}
}
}
WalOperation::CreateIndex {
name,
table,
column,
index_type,
options,
} => {
info!(
"Replayed create index: name={}, table={}, column={}",
name, table, column
);
let decoded_options = if options.is_empty() {
Vec::new()
} else {
bincode::deserialize(&options)
.map_err(|e| Error::storage(format!("Failed to deserialize index options: {}", e)))?
};
let definition = crate::storage::PersistedIndexDefinition {
table_name: table,
column_name: column,
index_type,
options: decoded_options,
};
Catalog::new(self).save_index_definition(&name, &definition)?;
Ok(())
}
WalOperation::DropIndex { name } => {
info!("Replayed drop index: name={}", name);
Catalog::new(self).drop_index_definition(&name)?;
Ok(())
}
WalOperation::CreateTrigger {
name,
table,
definition,
} => {
match bincode::deserialize::<crate::sql::TriggerDefinition>(&definition) {
Ok(trigger_def) => {
let catalog = Catalog::new(self);
catalog.save_trigger(&trigger_def)?;
if let Err(e) = self.trigger_registry.register_trigger(trigger_def) {
warn!("Failed to register trigger in registry: {}", e);
}
info!("Replayed create trigger: name={}, table={}", name, table);
Ok(())
}
Err(e) => {
error!("Failed to deserialize trigger definition: {}", e);
Ok(())
}
}
}
WalOperation::DropTrigger { name, table } => {
info!("Replayed drop trigger: name={}, table={:?}", name, table);
if let Some(ref table_name) = table {
let catalog = Catalog::new(self);
catalog.delete_trigger(table_name, &name)?;
if let Err(e) = self.trigger_registry.drop_trigger(table_name, &name) {
warn!("Failed to drop trigger from registry: {}", e);
}
}
Ok(())
}
WalOperation::CreateFunction { name, definition } => {
info!("Replayed create function: name={}", name);
let key = format!("meta:function:{}", name).into_bytes();
self.put(&key, &definition)?;
Ok(())
}
WalOperation::DropFunction { name } => {
info!("Replayed drop function: name={}", name);
let key = format!("meta:function:{}", name).into_bytes();
self.delete(&key)?;
Ok(())
}
WalOperation::CreateProcedure { name, definition } => {
info!("Replayed create procedure: name={}", name);
let key = format!("meta:procedure:{}", name).into_bytes();
self.put(&key, &definition)?;
Ok(())
}
WalOperation::DropProcedure { name } => {
info!("Replayed drop procedure: name={}", name);
let key = format!("meta:procedure:{}", name).into_bytes();
self.delete(&key)?;
Ok(())
}
WalOperation::CreateMaterializedView { name, definition } => {
info!("Replayed create materialized view: name={}", name);
let key = format!("meta:matview:{}", name).into_bytes();
self.put(&key, &definition)?;
Ok(())
}
WalOperation::DropMaterializedView { name } => {
info!("Replayed drop materialized view: name={}", name);
let key = format!("meta:matview:{}", name).into_bytes();
self.delete(&key)?;
Ok(())
}
WalOperation::RefreshMaterializedView {
name,
concurrent,
incremental,
} => {
info!(
"Replayed refresh materialized view: name={}, concurrent={}, incremental={}",
name, concurrent, incremental
);
Ok(())
}
WalOperation::AddConstraint { table, constraint } => {
info!("Replayed add constraint on table: {}", table);
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_micros();
let key = format!("meta:constraint:{}:{}", table, timestamp).into_bytes();
self.put(&key, &constraint)?;
Ok(())
}
WalOperation::DropConstraint { table, constraint_name } => {
info!(
"Replayed drop constraint: table={}, constraint={}",
table, constraint_name
);
let prefix = format!("meta:constraint:{}:", table);
Ok(())
}
WalOperation::Begin { tx_id } => {
debug!("Transaction {} begin (metadata only)", tx_id);
Ok(())
}
WalOperation::Commit { tx_id } => {
debug!("Transaction {} commit (metadata only)", tx_id);
Ok(())
}
WalOperation::Abort { tx_id } => {
debug!("Transaction {} abort (metadata only)", tx_id);
Ok(())
}
WalOperation::UpdateCounter { table_name, new_value } => {
info!("Replayed update counter: table={}, new_value={}", table_name, new_value);
let counter = self
.row_counters
.entry(table_name.clone())
.or_insert_with(|| std::sync::atomic::AtomicU64::new(0));
let current = counter.load(std::sync::atomic::Ordering::SeqCst);
if new_value > current {
counter.store(new_value, std::sync::atomic::Ordering::SeqCst);
}
let key = format!("counter:{}", table_name).into_bytes();
let value = bincode::serialize(&new_value)
.map_err(|e| Error::storage(format!("Failed to serialize counter: {}", e)))?;
self.put_internal(&key, &value)?;
Ok(())
}
}
}
fn extract_tx_id(operation: &WalOperation) -> Option<u64> {
match operation {
WalOperation::Begin { tx_id } | WalOperation::Commit { tx_id } | WalOperation::Abort { tx_id } => {
Some(*tx_id)
}
_ => None,
}
}
fn apply_wal_operation_to_batch(&self, operation: &WalOperation, batch: &mut WriteBatch) -> Result<bool> {
match operation {
WalOperation::Insert { table, key, tuple } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(table).is_err() {
debug!("Skipping insert for non-existent table: {}", table);
return Ok(false);
}
let data = if let Some(km) = &self.key_manager {
crypto::encrypt(km.key(), tuple)?
} else {
tuple.clone()
};
batch.put(key, &data);
debug!("Batched insert: table={}, key_len={}", table, key.len());
Ok(true)
}
WalOperation::Update { table, key, tuple } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(table).is_err() {
debug!("Skipping update for non-existent table: {}", table);
return Ok(false);
}
let data = if let Some(km) = &self.key_manager {
crypto::encrypt(km.key(), tuple)?
} else {
tuple.clone()
};
batch.put(key, &data);
debug!("Batched update: table={}, key_len={}", table, key.len());
Ok(true)
}
WalOperation::Delete { table, key } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(table).is_err() {
debug!("Skipping delete for non-existent table: {}", table);
return Ok(false);
}
batch.delete(key);
debug!("Batched delete: table={}, key_len={}", table, key.len());
Ok(true)
}
WalOperation::CreateTable { table, schema } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(table).is_ok() {
debug!("Table {} already exists, skipping create", table);
return Ok(false);
}
match bincode::deserialize::<crate::Schema>(schema) {
Ok(schema_obj) => {
catalog.create_table(table, schema_obj)?;
debug!("Replayed create table: {}", table);
Ok(false) }
Err(e) => {
warn!("Failed to deserialize schema for table {}: {}", table, e);
Ok(false)
}
}
}
WalOperation::DropTable { table } => {
let catalog = Catalog::new(self);
if catalog.get_table_schema(table).is_err() {
debug!("Table {} doesn't exist, skipping drop", table);
return Ok(false);
}
catalog.drop_table(table)?;
debug!("Replayed drop table: {}", table);
Ok(false) }
WalOperation::Truncate { .. }
| WalOperation::AlterColumnStorage { .. }
| WalOperation::CreateIndex { .. }
| WalOperation::DropIndex { .. }
| WalOperation::CreateTrigger { .. }
| WalOperation::DropTrigger { .. }
| WalOperation::CreateFunction { .. }
| WalOperation::DropFunction { .. }
| WalOperation::CreateProcedure { .. }
| WalOperation::DropProcedure { .. }
| WalOperation::CreateMaterializedView { .. }
| WalOperation::DropMaterializedView { .. }
| WalOperation::RefreshMaterializedView { .. }
| WalOperation::AddConstraint { .. }
| WalOperation::DropConstraint { .. } => {
self.apply_wal_operation(operation.clone())?;
Ok(false) }
WalOperation::Begin { tx_id } => {
debug!("Transaction {} begin (metadata only)", tx_id);
Ok(false)
}
WalOperation::Commit { tx_id } => {
debug!("Transaction {} commit (metadata only)", tx_id);
Ok(false)
}
WalOperation::Abort { tx_id } => {
debug!("Transaction {} abort (metadata only)", tx_id);
Ok(false)
}
WalOperation::UpdateCounter { table_name, new_value } => {
debug!(
"Replaying counter update: table={}, new_value={}",
table_name, new_value
);
let counter = self
.row_counters
.entry(table_name.clone())
.or_insert_with(|| std::sync::atomic::AtomicU64::new(0));
let current = counter.load(std::sync::atomic::Ordering::SeqCst);
if *new_value > current {
counter.store(*new_value, std::sync::atomic::Ordering::SeqCst);
}
let key = format!("counter:{}", table_name).into_bytes();
let value = bincode::serialize(new_value)
.map_err(|e| Error::storage(format!("Failed to serialize counter: {}", e)))?;
batch.put(&key, &value);
Ok(false) }
}
}
pub fn truncate_wal(&self, up_to_lsn: u64) -> Result<()> {
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.truncate(up_to_lsn)?;
}
Ok(())
}
pub fn wal_sync_mode(&self) -> Option<WalSyncMode> {
self.wal.as_ref().map(|wal| wal.read().sync_mode())
}
pub fn set_wal_sync_mode(&self, mode: WalSyncMode) -> Result<()> {
if let Some(wal) = &self.wal {
let mut wal = wal.write();
wal.set_sync_mode(mode);
Ok(())
} else {
Err(Error::storage("WAL is not enabled"))
}
}
pub fn insert_tuple_versioned(&self, table_name: &str, tuple: Tuple) -> Result<u64> {
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
self.insert_tuple_versioned_with_schema(table_name, tuple, &schema)
}
pub fn insert_tuple_versioned_with_schema(
&self,
table_name: &str,
tuple: Tuple,
schema: &crate::Schema,
) -> Result<u64> {
let catalog = Catalog::new(self);
let row_id = self.next_row_id_volatile(table_name);
let mut tuple = tuple;
for (i, col) in schema.columns.iter().enumerate() {
if col.primary_key {
if let Some(v) = tuple.values.get(i) {
if matches!(v, crate::Value::Null) && i < tuple.values.len() {
#[allow(clippy::indexing_slicing)]
match col.data_type {
crate::DataType::Int2 => {
tuple.values[i] = crate::Value::Int2(row_id as i16);
}
crate::DataType::Int4 => {
tuple.values[i] = crate::Value::Int4(row_id as i32);
}
_ => {
tuple.values[i] = crate::Value::Int8(row_id as i64);
}
}
}
}
}
}
let logical_tuple = tuple.clone();
{
let mut col_values = std::collections::HashMap::with_capacity(schema.columns.len());
for (i, col) in schema.columns.iter().enumerate() {
if let Some(v) = tuple.values.get(i) {
col_values.insert(col.name.clone(), v.clone());
}
}
if let Err(e) = self.art_index_manager.check_unique_constraints(table_name, &col_values) {
return Err(Error::constraint_violation(e.to_string()));
}
}
let bulk_mode = self.is_bulk_load_mode();
let stored_tuple = self.transform_tuple_for_column_storage_opts(table_name, row_id, &tuple, schema, true)?;
let value = bincode::serialize(&stored_tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?;
let logical_value = if schema_uses_column_storage(schema) {
bincode::serialize(&logical_tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
} else {
value.clone()
};
let key = Self::build_data_key(table_name, row_id);
self.put(&key, &value)?;
self.flush_row_counter(table_name)?;
if self.fast_dml_requires_logical_wal() {
self.log_data_insert(table_name, &key, &logical_value)?;
} else {
self.log_data_insert_nosync(table_name, &key, &logical_value)?;
}
{
let mut col_values = std::collections::HashMap::new();
for (i, col) in schema.columns.iter().enumerate() {
if let Some(v) = tuple.values.get(i) {
col_values.insert(col.name.clone(), v.clone());
}
}
if let Err(e) = self.art_index_manager.on_insert(table_name, row_id, &col_values) {
tracing::debug!("ART index insert for table '{}': {}", table_name, e);
}
}
let _ = self.vector_indexes.on_row_insert(table_name, row_id, schema, &tuple);
if self.config.storage.time_travel_enabled {
let timestamp = self.next_timestamp();
self.snapshot_manager
.write_version(table_name, row_id, timestamp, &logical_value)?;
if let Some(lsn) = self.wal_lsn() {
let _ = self.snapshot_manager.register_snapshot_with_lsn(timestamp, lsn);
} else {
let _ = self.snapshot_manager.register_snapshot(timestamp);
}
}
if !bulk_mode {
if let Err(e) = self.mv_delta_tracker.record_insert(table_name, row_id, tuple.clone()) {
tracing::warn!("Failed to record insert delta for table '{}': {}", table_name, e);
}
self.filter_delta_tracker.on_insert(table_name, row_id, &tuple, &schema);
for (i, col) in schema.columns.iter().enumerate() {
if let Some(value) = tuple.values.get(i) {
self.speculative_filter_manager.on_insert(table_name, &col.name, value);
}
}
}
Ok(row_id)
}
pub fn insert_tuple_fast(&self, table_name: &str, tuple: Tuple, schema: &crate::Schema) -> Result<u64> {
let row_id = self.next_row_id_volatile(table_name);
let mut tuple = tuple;
for (i, col) in schema.columns.iter().enumerate() {
if col.primary_key {
if let Some(v) = tuple.values.get(i) {
if matches!(v, crate::Value::Null) && i < tuple.values.len() {
#[allow(clippy::indexing_slicing)]
match col.data_type {
crate::DataType::Int2 => {
tuple.values[i] = crate::Value::Int2(row_id as i16);
}
crate::DataType::Int4 => {
tuple.values[i] = crate::Value::Int4(row_id as i32);
}
_ => {
tuple.values[i] = crate::Value::Int8(row_id as i64);
}
}
}
}
}
}
if let Err(e) = self
.art_index_manager
.check_unique_constraints_tuple(table_name, schema, &tuple)
{
return Err(Error::constraint_violation(e.to_string()));
}
let uses_side_storage = schema_uses_column_storage(schema);
let value = if uses_side_storage {
let stored_tuple =
self.transform_tuple_for_column_storage_opts(table_name, row_id, &tuple, schema, true)?;
bincode::serialize(&stored_tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
} else {
bincode::serialize(&tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
};
let key = Self::build_data_key(table_name, row_id);
let requires_logical_wal = self.fast_dml_requires_logical_wal();
let needs_logical_value = requires_logical_wal || self.config.storage.time_travel_enabled;
let logical_value = if needs_logical_value {
Some(if uses_side_storage {
bincode::serialize(&tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
} else {
value.clone()
})
} else {
None
};
let data_and_version_batched =
self.config.storage.time_travel_enabled && !requires_logical_wal && !uses_side_storage;
if data_and_version_batched {
if let Some(ref db_path) = self.db_path {
let count = self.write_counter.fetch_add(1, Ordering::Relaxed);
if count % 1000 == 0 {
Self::check_disk_space(db_path)?;
}
}
if self.memory_limit_bytes > 0 {
let write_size = (key.len() + value.len()) as u64;
let current = self.data_bytes_written.fetch_add(write_size, Ordering::Relaxed);
if current + write_size > self.memory_limit_bytes {
self.data_bytes_written.fetch_sub(write_size, Ordering::Relaxed);
return Err(Error::storage(format!(
"Memory limit exceeded ({} MB). Increase resource_quotas.memory_limit_per_user_mb or use disk-backed mode.",
self.memory_limit_bytes / (1024 * 1024)
)));
}
}
} else {
self.put(&key, &value)?;
}
if requires_logical_wal {
let logical_value = logical_value
.as_deref()
.ok_or_else(|| Error::internal("missing logical insert value"))?;
if self.config.storage.logical_wal_per_statement {
self.log_data_insert(table_name, &key, logical_value)?;
} else {
self.log_data_insert_nosync(table_name, &key, logical_value)?;
}
}
if let Err(e) = self
.art_index_manager
.on_insert_tuple(table_name, row_id, schema, &tuple)
{
tracing::debug!("ART index insert for table '{}': {}", table_name, e);
}
let _ = self.vector_indexes.on_row_insert(table_name, row_id, schema, &tuple);
if row_id % 64 == 0 {
let _ = self.flush_row_counter(table_name);
}
if self.config.storage.time_travel_enabled {
let logical_value = logical_value
.as_deref()
.ok_or_else(|| Error::internal("missing logical insert value"))?;
let timestamp = self.next_timestamp();
if data_and_version_batched {
self.snapshot_manager.write_data_version_and_register_snapshot(
&key,
&value,
table_name,
row_id,
timestamp,
logical_value,
self.wal_lsn(),
self.memory_write_options.as_ref(),
)?;
} else {
self.snapshot_manager.write_version_and_register_snapshot(
table_name,
row_id,
timestamp,
logical_value,
self.wal_lsn(),
)?;
}
}
Ok(row_id)
}
pub fn insert_prepared_tuples_fast_batch(
&self,
table_name: &str,
prepared: Vec<(u64, Tuple)>,
schema: &crate::Schema,
) -> Result<u64> {
if prepared.is_empty() {
return Ok(0);
}
if self.fast_dml_requires_logical_wal() {
return Err(Error::internal(
"fast batch insert requires Transaction path when logical WAL is active",
));
}
if schema.columns.iter().any(|column| {
matches!(
column.storage_mode,
ColumnStorageMode::Dictionary | ColumnStorageMode::ContentAddressed
)
}) {
return Err(Error::internal(
"fast batch insert direct WriteBatch requires default or columnar row storage",
));
}
let uses_columnar = schema_has_columnar_columns(schema);
let columnar_indices: Vec<usize> = if uses_columnar {
schema
.columns
.iter()
.enumerate()
.filter(|(_, column)| column.storage_mode == ColumnStorageMode::Columnar)
.map(|(idx, _)| idx)
.collect()
} else {
Vec::new()
};
let maintain_presence = if uses_columnar {
self.columnar_presence_writes_active(table_name)?
} else {
false
};
let commit_ts = if self.config.storage.time_travel_enabled {
Some(self.next_timestamp())
} else {
None
};
let reverse_ts = commit_ts.map(|ts| u64::MAX - ts);
let mut batch = WriteBatch::default();
let mut indexed_rows = Vec::with_capacity(prepared.len());
let mut final_row_id = 0_u64;
let data_prefix = format!("data:{}:", table_name);
let mut row_id_buf = itoa::Buffer::new();
let mut key_buf = Vec::with_capacity(data_prefix.len() + 32);
let serialize_stored = |tuple: &Tuple| -> Result<Vec<u8>> {
if uses_columnar {
let mut stored = tuple.clone();
for &idx in &columnar_indices {
if let Some(val) = stored.values.get_mut(idx) {
*val = crate::Value::ColumnarRef;
}
}
bincode::serialize(&stored).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))
} else {
bincode::serialize(tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))
}
};
if let (Some(ts), Some(reverse_ts)) = (commit_ts, reverse_ts) {
let version_prefix = format!("v:{}:", table_name);
let version_suffix = format!(":{}", ts);
let version_index_prefix = format!("v_idx:{}:", table_name);
let mut reverse_ts_buf = itoa::Buffer::new();
let mut version_key_buf = Vec::with_capacity(version_prefix.len() + 64);
let mut version_index_key_buf = Vec::with_capacity(version_index_prefix.len() + 64);
for (row_id, tuple) in prepared {
let value = serialize_stored(&tuple)?;
let logical_value = if uses_columnar {
bincode::serialize(&tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
} else {
value.clone()
};
let row_id_str = row_id_buf.format(row_id);
key_buf.clear();
key_buf.extend_from_slice(data_prefix.as_bytes());
key_buf.extend_from_slice(row_id_str.as_bytes());
batch.put(&key_buf, &value);
version_key_buf.clear();
version_key_buf.extend_from_slice(version_prefix.as_bytes());
version_key_buf.extend_from_slice(row_id_str.as_bytes());
version_key_buf.extend_from_slice(version_suffix.as_bytes());
batch.put(&version_key_buf, &logical_value);
version_index_key_buf.clear();
version_index_key_buf.extend_from_slice(version_index_prefix.as_bytes());
version_index_key_buf.extend_from_slice(row_id_str.as_bytes());
version_index_key_buf.push(b':');
let reverse_ts_str = reverse_ts_buf.format(reverse_ts);
for _ in reverse_ts_str.len()..20 {
version_index_key_buf.push(b'0');
}
version_index_key_buf.extend_from_slice(reverse_ts_str.as_bytes());
batch.put(&version_index_key_buf, ts.to_be_bytes());
final_row_id = row_id;
indexed_rows.push((row_id, tuple));
}
} else {
for (row_id, tuple) in prepared {
let value = serialize_stored(&tuple)?;
let row_id_str = row_id_buf.format(row_id);
key_buf.clear();
key_buf.extend_from_slice(data_prefix.as_bytes());
key_buf.extend_from_slice(row_id_str.as_bytes());
batch.put(&key_buf, &value);
final_row_id = row_id;
indexed_rows.push((row_id, tuple));
}
}
if final_row_id > 0 {
let counter_key = format!("counter:{}", table_name);
let counter_value = bincode::serialize(&final_row_id)
.map_err(|e| Error::storage(format!("Failed to serialize counter: {}", e)))?;
batch.put(counter_key.as_bytes(), counter_value);
}
let columnar_guard = if uses_columnar {
let row_refs: Vec<(u64, &Tuple)> = indexed_rows.iter().map(|(row_id, tuple)| (*row_id, tuple)).collect();
let guard = crate::storage::columnar::stats_write_lock();
for ((idx, batch_id), values) in group_columnar_row_values(schema, &row_refs) {
let column = schema
.columns
.get(idx)
.ok_or_else(|| Error::internal("column index out of bounds in fast batch insert"))?;
ColumnarStore::apply_batch_values(&self.db, &mut batch, table_name, &column.name, batch_id, &values)?;
}
if maintain_presence {
for (batch_id, set) in group_row_ids_by_batch(indexed_rows.iter().map(|(row_id, _)| *row_id)) {
ColumnarStore::apply_presence_updates(&self.db, &mut batch, table_name, batch_id, &set, &[])?;
}
}
Some(guard)
} else {
None
};
let result = if let Some(opts) = &self.memory_write_options {
self.db.write_opt(batch, opts)
} else {
self.db.write(batch)
};
drop(columnar_guard);
result.map_err(|e| Error::storage(format!("Fast batch insert failed: {}", e)))?;
if let Some(ts) = commit_ts {
let _ = self.snapshot_manager.register_snapshot(ts);
}
for (row_id, tuple) in &indexed_rows {
if let Err(e) = self
.art_index_manager
.on_insert_tuple(table_name, *row_id, schema, tuple)
{
tracing::debug!("ART index batch insert for table '{}': {}", table_name, e);
}
let _ = self.vector_indexes.on_row_insert(table_name, *row_id, schema, tuple);
}
Ok(indexed_rows.len() as u64)
}
fn transform_tuple_for_column_storage(
&self,
table_name: &str,
row_id: u64,
tuple: &Tuple,
schema: &crate::Schema,
) -> Result<Tuple> {
self.transform_tuple_for_column_storage_opts(table_name, row_id, tuple, schema, false)
}
fn transform_tuple_for_column_storage_opts(
&self,
table_name: &str,
row_id: u64,
tuple: &Tuple,
schema: &crate::Schema,
new_row: bool,
) -> Result<Tuple> {
if !schema_uses_column_storage(schema) {
return Ok(tuple.clone());
}
let mut transformed = tuple.clone();
let mut used_dictionary = false;
for (idx, column) in schema.columns.iter().enumerate() {
if idx >= transformed.values.len() {
break;
}
match column.storage_mode {
ColumnStorageMode::Dictionary => {
if let Some(crate::Value::String(s)) = transformed.values.get(idx) {
let dict_id = self.dict_manager.encode(&self.db, table_name, &column.name, s)?;
if let Some(val) = transformed.values.get_mut(idx) {
*val = crate::Value::DictRef { dict_id };
}
used_dictionary = true;
}
}
ColumnStorageMode::ContentAddressed => {
let cur_val = transformed
.values
.get(idx)
.ok_or_else(|| Error::internal("index out of bounds in content-addressed transform"))?;
let new_val = ContentAddressedStore::maybe_store(&self.db, cur_val)?;
if let Some(val) = transformed.values.get_mut(idx) {
*val = new_val;
}
}
ColumnStorageMode::Columnar => {
if let Some(val) = transformed.values.get_mut(idx) {
*val = crate::Value::ColumnarRef;
}
}
ColumnStorageMode::Default => {}
}
}
if schema_has_columnar_columns(schema) {
self.store_columnar_rows_grouped(table_name, &[(row_id, tuple)], schema, new_row)?;
}
if used_dictionary {
self.dict_manager.flush(&self.db)?;
}
Ok(transformed)
}
fn columnar_presence_writes_active(&self, table_name: &str) -> Result<bool> {
if columnar_presence_enabled() {
if !ColumnarStore::presence_manifest_complete(&self.db, table_name) {
ColumnarStore::backfill_presence(&self.db, table_name)?;
}
Ok(true)
} else {
Ok(ColumnarStore::presence_manifest_complete(&self.db, table_name))
}
}
pub(crate) fn columnar_live_presence(&self, table_name: &str) -> Option<Vec<(u64, BatchPresence)>> {
if !columnar_presence_enabled() {
return None;
}
{
let branch = self.current_branch.lock();
if branch.as_deref().is_some_and(|name| name != "main") {
return None;
}
}
match ColumnarStore::presence_scan_map(&self.db, table_name) {
Ok(map) => Some(map),
Err(e) => {
tracing::warn!(table = table_name, error = %e, "Columnar presence load failed; using row walk");
None
}
}
}
pub(crate) fn store_columnar_rows_grouped(
&self,
table_name: &str,
rows: &[(u64, &Tuple)],
schema: &crate::Schema,
mark_present: bool,
) -> Result<()> {
if rows.is_empty() || !schema_has_columnar_columns(schema) {
return Ok(());
}
let maintain_presence = mark_present && self.columnar_presence_writes_active(table_name)?;
let grouped = group_columnar_row_values(schema, rows);
if grouped.is_empty() && !maintain_presence {
return Ok(());
}
let _guard = crate::storage::columnar::stats_write_lock();
let mut write_batch = WriteBatch::default();
for ((idx, batch_id), values) in &grouped {
let column = schema
.columns
.get(*idx)
.ok_or_else(|| Error::internal("column index out of bounds in grouped columnar store"))?;
ColumnarStore::apply_batch_values(&self.db, &mut write_batch, table_name, &column.name, *batch_id, values)?;
}
if maintain_presence {
for (batch_id, set) in group_row_ids_by_batch(rows.iter().map(|(row_id, _)| *row_id)) {
ColumnarStore::apply_presence_updates(&self.db, &mut write_batch, table_name, batch_id, &set, &[])?;
}
}
self.db
.write(write_batch)
.map_err(|e| Error::storage(format!("Grouped columnar store failed: {}", e)))
}
pub(crate) fn delete_columnar_rows_grouped(
&self,
table_name: &str,
row_ids: &[u64],
schema: &crate::Schema,
) -> Result<()> {
if row_ids.is_empty() || !schema_has_columnar_columns(schema) {
return Ok(());
}
let maintain_presence = self.columnar_presence_writes_active(table_name)?;
let by_batch = group_row_ids_by_batch(row_ids.iter().copied());
let _guard = crate::storage::columnar::stats_write_lock();
let mut write_batch = WriteBatch::default();
for column in schema
.columns
.iter()
.filter(|column| column.storage_mode == ColumnStorageMode::Columnar)
{
for (batch_id, batch_rows) in &by_batch {
let values: Vec<(u64, Value)> = batch_rows.iter().map(|row_id| (*row_id, Value::Null)).collect();
ColumnarStore::apply_batch_values(
&self.db,
&mut write_batch,
table_name,
&column.name,
*batch_id,
&values,
)?;
}
}
if maintain_presence {
for (batch_id, clear) in &by_batch {
ColumnarStore::apply_presence_updates(&self.db, &mut write_batch, table_name, *batch_id, &[], clear)?;
}
}
self.db
.write(write_batch)
.map_err(|e| Error::storage(format!("Grouped columnar delete failed: {}", e)))
}
pub(crate) fn stage_columnar_rows_grouped_in_transaction(
&self,
table_name: &str,
rows: &[(u64, &Tuple)],
schema: &crate::Schema,
txn: &Transaction,
mark_present: bool,
) -> Result<()> {
if rows.is_empty() || !schema_has_columnar_columns(schema) {
return Ok(());
}
let maintain_presence = mark_present && self.columnar_presence_writes_active(table_name)?;
let grouped = group_columnar_row_values(schema, rows);
for ((idx, batch_id), values) in &grouped {
let column = schema
.columns
.get(*idx)
.ok_or_else(|| Error::internal("column index out of bounds in grouped columnar stage"))?;
ColumnarStore::stage_batch_values_in_transaction(
&self.db,
txn,
table_name,
&column.name,
*batch_id,
values,
)?;
}
if maintain_presence {
for (batch_id, set) in group_row_ids_by_batch(rows.iter().map(|(row_id, _)| *row_id)) {
ColumnarStore::stage_presence_updates_in_transaction(&self.db, txn, table_name, batch_id, &set, &[])?;
}
}
Ok(())
}
pub(crate) fn stage_columnar_deletes_grouped_in_transaction(
&self,
table_name: &str,
row_ids: &[u64],
schema: &crate::Schema,
txn: &Transaction,
) -> Result<()> {
if row_ids.is_empty() || !schema_has_columnar_columns(schema) {
return Ok(());
}
let maintain_presence = self.columnar_presence_writes_active(table_name)?;
let by_batch = group_row_ids_by_batch(row_ids.iter().copied());
for column in schema
.columns
.iter()
.filter(|column| column.storage_mode == ColumnStorageMode::Columnar)
{
for (batch_id, batch_rows) in &by_batch {
let values: Vec<(u64, Value)> = batch_rows.iter().map(|row_id| (*row_id, Value::Null)).collect();
ColumnarStore::stage_batch_values_in_transaction(
&self.db,
txn,
table_name,
&column.name,
*batch_id,
&values,
)?;
}
}
if maintain_presence {
for (batch_id, clear) in &by_batch {
ColumnarStore::stage_presence_updates_in_transaction(&self.db, txn, table_name, *batch_id, &[], clear)?;
}
}
Ok(())
}
pub(crate) fn stage_tuple_for_column_storage_in_transaction(
&self,
table_name: &str,
row_id: u64,
tuple: &Tuple,
schema: &crate::Schema,
txn: &Transaction,
) -> Result<()> {
if !schema_uses_column_storage(schema) {
return Ok(());
}
self.stage_columnar_rows_grouped_in_transaction(table_name, &[(row_id, tuple)], schema, txn, false)
}
pub(crate) fn stage_columnar_delete_in_transaction(
&self,
table_name: &str,
row_id: u64,
schema: &crate::Schema,
txn: &Transaction,
) -> Result<()> {
if !schema_uses_column_storage(schema) {
return Ok(());
}
self.stage_columnar_deletes_grouped_in_transaction(table_name, &[row_id], schema, txn)
}
pub fn update_tuple_fast(
&self,
table_name: &str,
row_id: u64,
new_tuple: Tuple,
old_tuple: &Tuple,
schema: &crate::Schema,
) -> Result<u64> {
self.update_tuple_fast_with_index_hint(table_name, row_id, new_tuple, old_tuple, schema, None)
}
pub fn update_tuple_fast_with_index_hint(
&self,
table_name: &str,
row_id: u64,
new_tuple: Tuple,
old_tuple: &Tuple,
schema: &crate::Schema,
index_update_needed: Option<bool>,
) -> Result<u64> {
let value = if schema_uses_column_storage(schema) {
let stored_tuple = self.transform_tuple_for_column_storage(table_name, row_id, &new_tuple, schema)?;
bincode::serialize(&stored_tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
} else {
bincode::serialize(&new_tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
};
let key = Self::build_data_key(table_name, row_id);
self.put(&key, &value)?;
let update_indexes = index_update_needed.unwrap_or_else(|| {
self.art_index_manager
.tuple_update_affects_indexes(table_name, schema, old_tuple, &new_tuple)
});
if update_indexes {
let mut old_col_values = std::collections::HashMap::new();
let mut new_col_values = std::collections::HashMap::new();
for (i, col) in schema.columns.iter().enumerate() {
if let Some(v) = old_tuple.values.get(i) {
old_col_values.insert(col.name.clone(), v.clone());
}
if let Some(v) = new_tuple.values.get(i) {
new_col_values.insert(col.name.clone(), v.clone());
}
}
if let Err(e) = self
.art_index_manager
.on_update(table_name, row_id, &old_col_values, &new_col_values)
{
tracing::debug!("ART index update for table '{}': {}", table_name, e);
}
}
let _ = self
.vector_indexes
.on_row_update(table_name, row_id, schema, Some(old_tuple), &new_tuple);
self.row_cache.invalidate(table_name, row_id);
Ok(1)
}
pub fn update_tuple_fast_no_index(
&self,
table_name: &str,
row_id: u64,
new_tuple: Tuple,
schema: &crate::Schema,
vector_old_tuple: Option<&Tuple>,
) -> Result<u64> {
let value = if schema_uses_column_storage(schema) {
let stored_tuple = self.transform_tuple_for_column_storage(table_name, row_id, &new_tuple, schema)?;
bincode::serialize(&stored_tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
} else {
bincode::serialize(&new_tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
};
let key = Self::build_data_key(table_name, row_id);
self.put(&key, &value)?;
if vector_old_tuple.is_some() || self.vector_indexes.table_has_indexes(table_name) {
let _ = self
.vector_indexes
.on_row_update(table_name, row_id, schema, vector_old_tuple, &new_tuple);
}
self.row_cache.invalidate(table_name, row_id);
Ok(1)
}
pub fn delete_tuple_fast(
&self,
table_name: &str,
row_id: u64,
old_tuple: &Tuple,
schema: &crate::Schema,
) -> Result<u64> {
let key = Self::build_data_key(table_name, row_id);
if let Some(opts) = &self.memory_write_options {
self.db
.delete_opt(&key, opts)
.map_err(|e| Error::storage(format!("Fast delete failed: {}", e)))?;
} else {
self.db
.delete(&key)
.map_err(|e| Error::storage(format!("Fast delete failed: {}", e)))?;
}
if let Err(e) = self
.art_index_manager
.on_delete_tuple(table_name, row_id, schema, old_tuple)
{
tracing::debug!("ART index delete for table '{}': {}", table_name, e);
}
let _ = self
.vector_indexes
.on_row_delete(table_name, row_id, Some(schema), Some(old_tuple));
self.delete_columnar_rows_grouped(table_name, &[row_id], schema)?;
self.row_cache.invalidate(table_name, row_id);
Ok(1)
}
pub fn delete_tuple_fast_pk_only(
&self,
table_name: &str,
row_id: u64,
pk_key: &[u8],
pk_value: &crate::Value,
) -> Result<u64> {
let key = Self::build_data_key(table_name, row_id);
if let Some(opts) = &self.memory_write_options {
self.db
.delete_opt(&key, opts)
.map_err(|e| Error::storage(format!("Fast PK delete failed: {}", e)))?;
} else {
self.db
.delete(&key)
.map_err(|e| Error::storage(format!("Fast PK delete failed: {}", e)))?;
}
if let Err(e) = self
.art_index_manager
.remove_single_pk_key(table_name, pk_key, row_id, pk_value)
{
tracing::debug!("ART PK delete for table '{}': {}", table_name, e);
}
let _ = self.vector_indexes.on_row_delete(table_name, row_id, None, None);
self.row_cache.invalidate(table_name, row_id);
Ok(1)
}
pub fn snapshot_manager(&self) -> &crate::storage::time_travel::SnapshotManager {
&self.snapshot_manager
}
pub fn snapshot_manager_arc(&self) -> Arc<crate::storage::time_travel::SnapshotManager> {
Arc::clone(&self.snapshot_manager)
}
pub fn scan_table_at_snapshot(&self, table_name: &str, snapshot_ts: u64) -> Result<Vec<Tuple>> {
if !self.config.storage.time_travel_enabled {
return self.scan_table_branch_aware(table_name);
}
let prefix = format!("data:{}:", table_name);
let prefix_bytes = prefix.as_bytes();
let mut tuples = Vec::new();
let mut seen_rows = std::collections::HashSet::new();
let mut read_opts = ReadOptions::default();
read_opts.set_total_order_seek(true);
let iter = self.db.iterator_opt(IteratorMode::Start, read_opts);
for item in iter {
let (key, _) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if key.starts_with(prefix_bytes) {
if let Some(row_id) = Self::parse_row_id_after_prefix(&key, prefix_bytes.len()) {
seen_rows.insert(row_id);
}
} else if key.first() > prefix_bytes.first() {
break;
}
}
for row_id in seen_rows {
if let Some(value) = self
.snapshot_manager
.read_at_snapshot(table_name, row_id, snapshot_ts)?
{
let mut tuple: Tuple = bincode::deserialize(&value)
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
tuple.row_id = Some(row_id);
tuples.push(tuple);
}
}
Ok(tuples)
}
fn load_counters(&self) -> Result<()> {
let prefix = b"counter:";
let iter = self.db.iterator(rocksdb::IteratorMode::Start);
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if key.starts_with(prefix) {
let value = self.decrypt_value(&raw_value)?;
let table_name = String::from_utf8_lossy(key.get(prefix.len()..).unwrap_or_default()).to_string();
let count: u64 = bincode::deserialize(&value)
.map_err(|e| Error::storage(format!("Failed to deserialize counter: {}", e)))?;
self.row_counters
.insert(table_name, std::sync::atomic::AtomicU64::new(count));
} else if key.first() > prefix.first() {
break;
}
}
Ok(())
}
pub fn next_row_id(&self, table_name: &str) -> Result<u64> {
let next = if let Some(counter) = self.row_counters.get(table_name) {
counter.fetch_add(1, Ordering::Relaxed) + 1
} else {
let counter = self
.row_counters
.entry(table_name.to_string())
.or_insert_with(|| std::sync::atomic::AtomicU64::new(0));
counter.fetch_add(1, Ordering::Relaxed) + 1
};
let key = format!("counter:{}", table_name).into_bytes();
let value =
bincode::serialize(&next).map_err(|e| Error::storage(format!("Failed to serialize counter: {}", e)))?;
self.put_internal(&key, &value)?;
if !self.is_replaying.load(Ordering::Acquire) {
if let Some(wal) = &self.wal {
let wal = wal.read();
wal.append(WalOperation::UpdateCounter {
table_name: table_name.to_string(),
new_value: next,
})?;
}
}
Ok(next)
}
pub fn next_row_id_volatile(&self, table_name: &str) -> u64 {
if let Some(counter) = self.row_counters.get(table_name) {
return counter.fetch_add(1, Ordering::Relaxed) + 1;
}
let counter = self
.row_counters
.entry(table_name.to_string())
.or_insert_with(|| std::sync::atomic::AtomicU64::new(0));
counter.fetch_add(1, Ordering::Relaxed) + 1
}
pub fn stage_row_counter_in_transaction(
&self,
table_name: &str,
row_id: u64,
txn: &crate::storage::Transaction,
) -> Result<()> {
txn.stage_row_counter(table_name, row_id)
}
pub fn flush_row_counter(&self, table_name: &str) -> Result<()> {
let current = if let Some(counter) = self.row_counters.get(table_name) {
counter.load(Ordering::Relaxed)
} else {
let counter = self
.row_counters
.entry(table_name.to_string())
.or_insert_with(|| std::sync::atomic::AtomicU64::new(0));
counter.load(Ordering::Relaxed)
};
let key = format!("counter:{}", table_name).into_bytes();
let value =
bincode::serialize(¤t).map_err(|e| Error::storage(format!("Failed to serialize counter: {}", e)))?;
self.put_internal(&key, &value)
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use crate::{Column, DataType, Schema, Value};
#[test]
fn test_parse_row_id_after_prefix() {
assert_eq!(
StorageEngine::parse_row_id_after_prefix(b"data:users:42", b"data:users:".len()),
Some(42)
);
assert_eq!(
StorageEngine::parse_row_id_after_prefix(b"data:users:18446744073709551615", b"data:users:".len()),
Some(u64::MAX)
);
assert_eq!(
StorageEngine::parse_row_id_after_prefix(b"data:users:", b"data:users:".len()),
None
);
assert_eq!(
StorageEngine::parse_row_id_after_prefix(b"data:users:4x", b"data:users:".len()),
None
);
assert_eq!(
StorageEngine::parse_row_id_after_prefix(b"data:users:18446744073709551616", b"data:users:".len()),
None
);
}
#[test]
fn test_storage_engine_creation() {
let config = Config::in_memory();
let engine = StorageEngine::open_in_memory(&config);
assert!(engine.is_ok());
}
#[test]
fn test_basic_put_get() {
let config = Config::in_memory();
let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");
let key = b"test_key".to_vec();
let value = b"test_value".to_vec();
engine.put(&key, &value).expect("Failed to put value");
let result = engine.get(&key).expect("Failed to get value");
assert_eq!(result, Some(value));
}
#[test]
fn test_delete() {
let config = Config::in_memory();
let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");
let key = b"test_key".to_vec();
let value = b"test_value".to_vec();
engine.put(&key, &value).expect("Failed to put value");
engine.delete(&key).expect("Failed to delete value");
let result = engine.get(&key).expect("Failed to get value");
assert_eq!(result, None);
}
#[test]
fn test_scan_table_at_snapshot_basic() {
let config = Config::in_memory();
let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");
let schema = Schema {
columns: vec![
Column {
name: "id".to_string(),
data_type: DataType::Int4,
nullable: false,
primary_key: true,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
},
Column {
name: "value".to_string(),
data_type: DataType::Text,
nullable: false,
primary_key: false,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
},
],
};
let catalog = engine.catalog();
catalog
.create_table("test_table", schema.clone())
.expect("Failed to create table");
let tuple1 = Tuple {
values: vec![Value::Int4(1), Value::String("first".to_string())],
row_id: None,
branch_id: None,
};
engine
.insert_tuple_versioned("test_table", tuple1)
.expect("Failed to insert tuple 1");
let snapshot1_ts = engine.current_timestamp();
let tuple2 = Tuple {
values: vec![Value::Int4(2), Value::String("second".to_string())],
row_id: None,
branch_id: None,
};
engine
.insert_tuple_versioned("test_table", tuple2)
.expect("Failed to insert tuple 2");
let snapshot2_ts = engine.current_timestamp();
let tuple3 = Tuple {
values: vec![Value::Int4(3), Value::String("third".to_string())],
row_id: None,
branch_id: None,
};
engine
.insert_tuple_versioned("test_table", tuple3)
.expect("Failed to insert tuple 3");
let _snapshot3_ts = engine.current_timestamp();
let results1 = engine
.scan_table_at_snapshot("test_table", snapshot1_ts)
.expect("Failed to scan at snapshot 1");
assert_eq!(results1.len(), 1, "Should see 1 tuple at snapshot 1");
if let Value::String(ref val) = results1[0].values[1] {
assert_eq!(val, "first");
} else {
panic!("Expected text value");
}
let results2 = engine
.scan_table_at_snapshot("test_table", snapshot2_ts)
.expect("Failed to scan at snapshot 2");
assert_eq!(results2.len(), 2, "Should see 2 tuples at snapshot 2");
let results_current = engine.scan_table("test_table").expect("Failed to scan current state");
assert_eq!(results_current.len(), 3, "Should see 3 tuples in current state");
}
#[test]
fn test_scan_table_at_snapshot_mvcc_consistency() {
let config = Config::in_memory();
let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");
let schema = Schema {
columns: vec![Column {
name: "id".to_string(),
data_type: DataType::Int4,
nullable: false,
primary_key: true,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
}],
};
let catalog = engine.catalog();
catalog
.create_table("mvcc_test", schema)
.expect("Failed to create table");
engine
.insert_tuple_versioned(
"mvcc_test",
Tuple {
values: vec![Value::Int4(1)],
row_id: None,
branch_id: None,
},
)
.expect("Failed to insert");
let snapshot_ts = engine.current_timestamp();
engine
.insert_tuple_versioned(
"mvcc_test",
Tuple {
values: vec![Value::Int4(2)],
row_id: None,
branch_id: None,
},
)
.expect("Failed to insert");
engine
.insert_tuple_versioned(
"mvcc_test",
Tuple {
values: vec![Value::Int4(3)],
row_id: None,
branch_id: None,
},
)
.expect("Failed to insert");
let results1 = engine
.scan_table_at_snapshot("mvcc_test", snapshot_ts)
.expect("First scan failed");
let results2 = engine
.scan_table_at_snapshot("mvcc_test", snapshot_ts)
.expect("Second scan failed");
assert_eq!(results1.len(), results2.len());
assert_eq!(results1.len(), 1, "Should only see data up to snapshot");
let current_results = engine.scan_table("mvcc_test").expect("Current scan failed");
assert_eq!(current_results.len(), 3, "Current state should have all data");
}
#[test]
fn test_scan_table_at_snapshot_empty_table() {
let config = Config::in_memory();
let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");
let schema = Schema {
columns: vec![Column {
name: "id".to_string(),
data_type: DataType::Int4,
nullable: false,
primary_key: true,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
}],
};
let catalog = engine.catalog();
catalog
.create_table("empty_table", schema)
.expect("Failed to create table");
let results = engine
.scan_table_at_snapshot("empty_table", engine.current_timestamp())
.expect("Failed to scan empty table");
assert_eq!(results.len(), 0, "Empty table should return no results");
}
#[test]
fn test_scan_table_at_snapshot_nonexistent_data() {
let config = Config::in_memory();
let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");
let schema = Schema {
columns: vec![Column {
name: "id".to_string(),
data_type: DataType::Int4,
nullable: false,
primary_key: true,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
}],
};
let catalog = engine.catalog();
catalog
.create_table("future_test", schema)
.expect("Failed to create table");
let early_snapshot = 1;
let results = engine
.scan_table_at_snapshot("future_test", early_snapshot)
.expect("Failed to scan at early timestamp");
assert_eq!(results.len(), 0, "Should see no data before inserts");
engine
.insert_tuple_versioned(
"future_test",
Tuple {
values: vec![Value::Int4(1)],
row_id: None,
branch_id: None,
},
)
.expect("Failed to insert");
let results_after = engine
.scan_table_at_snapshot("future_test", early_snapshot)
.expect("Failed to scan after insert");
assert_eq!(
results_after.len(),
0,
"Should still see no data at historical snapshot"
);
let current = engine.scan_table("future_test").expect("Failed to scan current");
assert_eq!(current.len(), 1, "Current state should have data");
}
#[test]
fn test_extract_table_from_key_data_format() {
let key = b"data:users:42";
assert_eq!(StorageEngine::extract_table_from_key(key), "users");
let key = b"data:products:12345";
assert_eq!(StorageEngine::extract_table_from_key(key), "products");
let key = b"data:my_table:1";
assert_eq!(StorageEngine::extract_table_from_key(key), "my_table");
let key = b"data:user_accounts_2024:999";
assert_eq!(StorageEngine::extract_table_from_key(key), "user_accounts_2024");
}
#[test]
fn test_extract_table_from_key_metadata_format() {
let key = b"meta:table:users";
assert_eq!(StorageEngine::extract_table_from_key(key), "users");
let key = b"meta:table:products";
assert_eq!(StorageEngine::extract_table_from_key(key), "products");
let key = b"meta:counter:users";
assert_eq!(StorageEngine::extract_table_from_key(key), "users");
}
#[test]
fn test_extract_table_from_key_system_keys() {
let key = b"wal:entries:00000000000000000001";
assert_eq!(StorageEngine::extract_table_from_key(key), "unknown");
let key = b"wal:last_lsn";
assert_eq!(StorageEngine::extract_table_from_key(key), "unknown");
let key = b"system:config";
assert_eq!(StorageEngine::extract_table_from_key(key), "unknown");
}
#[test]
fn test_extract_table_from_key_malformed() {
let key = b"data:users"; assert_eq!(StorageEngine::extract_table_from_key(key), "unknown");
let key = b"data:"; assert_eq!(StorageEngine::extract_table_from_key(key), "unknown");
let invalid_utf8: Vec<u8> = vec![0xFF, 0xFE, 0xFD];
assert_eq!(StorageEngine::extract_table_from_key(&invalid_utf8), "unknown");
let key = b"";
assert_eq!(StorageEngine::extract_table_from_key(key), "unknown");
}
#[test]
fn test_extract_table_from_key_edge_cases() {
let key = b"data:table_with_underscores:1";
assert_eq!(StorageEngine::extract_table_from_key(key), "table_with_underscores");
let long_table_name = "very_long_table_name_that_might_be_used_in_some_applications";
let key = format!("data:{}:42", long_table_name).into_bytes();
assert_eq!(StorageEngine::extract_table_from_key(&key), long_table_name);
let key = b"data:table123:456";
assert_eq!(StorageEngine::extract_table_from_key(key), "table123");
}
#[test]
fn test_wal_logging_with_table_names() {
let mut config = Config::in_memory();
config.storage.wal_enabled = true;
let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");
let schema = Schema {
columns: vec![
Column {
name: "id".to_string(),
data_type: DataType::Int4,
nullable: false,
primary_key: true,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
},
Column {
name: "name".to_string(),
data_type: DataType::Text,
nullable: false,
primary_key: false,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
},
],
};
let catalog = engine.catalog();
catalog
.create_table("test_users", schema)
.expect("Failed to create table");
let tuple = Tuple {
values: vec![Value::Int4(1), Value::String("Alice".to_string())],
row_id: None,
branch_id: None,
};
engine
.insert_tuple("test_users", tuple)
.expect("Failed to insert tuple");
assert!(engine.is_wal_enabled());
let lsn = engine.wal_lsn().expect("WAL should have LSN");
assert!(lsn > 0, "WAL should have at least one entry");
if let Some(wal) = &engine.wal {
let wal = wal.read();
let entries = wal.replay().expect("Failed to replay WAL");
let has_insert_with_table = entries.iter().any(|entry| {
if let crate::storage::WalOperation::Insert { table, .. } = &entry.operation {
table == "test_users"
} else {
false
}
});
assert!(
has_insert_with_table,
"WAL should contain insert operation for test_users table"
);
}
}
#[test]
fn test_fast_insert_logs_wal_when_strict_logical_wal_enabled() {
let mut config = Config::in_memory();
config.storage.wal_enabled = true;
config.storage.logical_wal_per_statement = true;
let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");
let schema = Schema {
columns: vec![
Column {
name: "id".to_string(),
data_type: DataType::Int4,
nullable: false,
primary_key: true,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
},
Column {
name: "name".to_string(),
data_type: DataType::Text,
nullable: false,
primary_key: false,
source_table: None,
source_table_name: None,
default_expr: None,
unique: false,
storage_mode: crate::ColumnStorageMode::Default,
},
],
};
engine
.catalog()
.create_table("fast_wal_users", schema.clone())
.expect("Failed to create table");
let tuple = Tuple {
values: vec![Value::Int4(1), Value::String("Alice".to_string())],
row_id: None,
branch_id: None,
};
engine
.insert_tuple_fast("fast_wal_users", tuple, &schema)
.expect("Failed to fast insert tuple");
let wal = engine.wal.as_ref().expect("WAL should be enabled").read();
let entries = wal.replay().expect("Failed to replay WAL");
let has_fast_insert = entries.iter().any(|entry| {
if let crate::storage::WalOperation::Insert { table, .. } = &entry.operation {
table == "fast_wal_users"
} else {
false
}
});
assert!(has_fast_insert, "fast insert should be present in logical WAL");
}
}
impl StorageEngine {
pub fn analyze_table(&self, table_name: &str) -> Result<()> {
let catalog = self.catalog();
catalog.analyze_table(table_name)
}
pub fn get_table_statistics(&self, table_name: &str) -> Result<Option<super::statistics::TableStatistics>> {
let catalog = self.catalog();
catalog.get_table_statistics(table_name)
}
pub fn analyze_all_tables(&self) -> Result<()> {
let catalog = self.catalog();
let tables = catalog.list_tables()?;
for table_name in tables {
if table_name.starts_with("helios_") || table_name.starts_with("mv_") {
continue;
}
catalog.analyze_table(&table_name)?;
}
Ok(())
}
#[cfg(feature = "sync-experimental")]
pub fn is_sync_enabled(&self) -> bool {
self.change_log.is_some()
}
#[cfg(feature = "sync-experimental")]
pub fn change_log(&self) -> Option<Arc<RwLock<crate::sync::ChangeLogImpl>>> {
self.change_log.as_ref().map(Arc::clone)
}
#[cfg(feature = "sync-experimental")]
pub fn node_id(&self) -> uuid::Uuid {
self.node_id
}
#[cfg(feature = "sync-experimental")]
pub(crate) fn capture_change(&self, transaction_id: u64, change_type: crate::sync::ChangeType) -> Result<()> {
if let Some(ref change_log) = self.change_log {
let mut vector_clock = crate::sync::VectorClock::new();
vector_clock.increment(self.node_id);
let mut cl = change_log.write();
cl.append(transaction_id, change_type, vector_clock)?;
}
Ok(())
}
pub fn get_current_branch(&self) -> Option<String> {
self.current_branch.lock().as_ref().cloned()
}
pub fn set_current_branch(&self, branch_name: Option<String>) {
*self.current_branch.lock() = branch_name.filter(|name| name != "main");
}
pub fn clear_current_branch(&self) {
*self.current_branch.lock() = None;
}
pub fn is_branch_active(&self) -> bool {
self.current_branch.lock().as_deref().is_some_and(|name| name != "main")
}
pub fn get_current_branch_id(&self) -> Option<u64> {
let branch_name = self.current_branch.lock().clone()?;
if branch_name == "main" {
return None;
}
let branch_manager = self.branch_manager()?;
branch_manager
.get_branch_by_name(&branch_name)
.ok()
.map(|m| m.branch_id)
}
fn get_branch_chain(&self, current_branch_id: u64) -> Result<Vec<u64>> {
let branch_manager = self
.branch_manager()
.ok_or_else(|| Error::storage("Branch manager not available"))?;
let parent_chain = branch_manager.build_parent_chain(current_branch_id)?;
let mut chain: Vec<u64> = parent_chain.into_iter().map(|(id, _)| id).collect();
chain.reverse();
chain.retain(|&id| id != 0);
chain.push(current_branch_id);
Ok(chain)
}
pub fn branch_aware_data_key(&self, table_name: &str, row_id: u64) -> Vec<u8> {
if let Some(branch_id) = self.get_current_branch_id() {
format!("bdata:{}:{}:{}", branch_id, table_name, row_id).into_bytes()
} else {
Self::build_data_key(table_name, row_id)
}
}
pub fn scan_table_branch_aware(&self, table_name: &str) -> Result<Vec<Tuple>> {
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
self.scan_table_branch_aware_with_schema(table_name, &schema)
}
pub fn scan_table_branch_aware_with_schema_prefix(
&self,
table_name: &str,
schema: &crate::Schema,
prefix_len: usize,
) -> Result<Vec<Tuple>> {
let branch_name = self.current_branch.lock().clone();
if branch_name.is_none() || branch_name.as_deref() == Some("main") {
return self.scan_table_with_schema_prefix(table_name, schema, prefix_len);
}
self.scan_table_branch_aware_with_schema(table_name, schema)
}
pub fn scan_table_branch_aware_with_schema_columns(
&self,
table_name: &str,
schema: &crate::Schema,
columns: &[usize],
) -> Result<Vec<Tuple>> {
let branch_name = self.current_branch.lock().clone();
if branch_name.is_none() || branch_name.as_deref() == Some("main") {
return self.scan_table_with_schema_columns(table_name, schema, columns);
}
self.scan_table_branch_aware_with_schema(table_name, schema)
}
pub fn scan_table_branch_aware_with_schema_columnar_columns(
&self,
table_name: &str,
schema: &crate::Schema,
columns: &[usize],
) -> Result<Vec<Tuple>> {
let branch_name = self.current_branch.lock().clone();
if branch_name.is_none() || branch_name.as_deref() == Some("main") {
return self.scan_table_with_schema_columnar_columns(table_name, schema, columns);
}
self.scan_table_branch_aware_with_schema(table_name, schema)
}
pub fn scan_table_branch_aware_with_schema_columnar_columns_filtered(
&self,
table_name: &str,
schema: &crate::Schema,
columns: &[usize],
predicates: &[AnalyzedPredicate],
) -> Result<Vec<Tuple>> {
let branch_name = self.current_branch.lock().clone();
if branch_name.is_none() || branch_name.as_deref() == Some("main") {
return self.scan_table_with_schema_columnar_columns_filtered(table_name, schema, columns, predicates);
}
self.scan_table_branch_aware_with_schema(table_name, schema)
}
pub fn scan_table_branch_aware_with_schema(&self, table_name: &str, schema: &crate::Schema) -> Result<Vec<Tuple>> {
let branch_name = self.current_branch.lock().clone();
if branch_name.is_none() || branch_name.as_deref() == Some("main") {
return self.scan_table_with_schema(table_name, schema);
}
let branch_id = match self.get_current_branch_id() {
Some(id) => id,
None => {
return Err(Error::query_execution(format!(
"Branch '{}' does not exist. Create it first with: CREATE BRANCH {} FROM main",
branch_name.as_deref().unwrap_or("unknown"),
branch_name.as_deref().unwrap_or("branch_name")
)));
}
};
let branch_chain = self.get_branch_chain(branch_id)?;
tracing::debug!(
"scan_table_branch_aware: branch chain for '{}' (id {}): {:?}",
branch_name.as_deref().unwrap_or("unknown"),
branch_id,
branch_chain
);
let mut result_tuples: std::collections::HashMap<u64, Tuple> = std::collections::HashMap::new();
let mut deleted_rows: std::collections::HashSet<u64> = std::collections::HashSet::new();
let main_prefix = format!("data:{}:", table_name);
let main_prefix_bytes = main_prefix.as_bytes();
let iter = self.db.iterator(rocksdb::IteratorMode::From(
main_prefix_bytes,
rocksdb::Direction::Forward,
));
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(main_prefix_bytes) {
break;
}
if let Some(row_id) = Self::parse_row_id_after_prefix(&key, main_prefix_bytes.len()) {
let value = self.decrypt_value(&raw_value)?;
let mut tuple: Tuple = bincode::deserialize(&value)
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
tuple.row_id = Some(row_id);
tuple.branch_id = None; result_tuples.insert(row_id, tuple);
}
}
for &chain_branch_id in &branch_chain {
let delete_prefix = format!("bdel:{}:{}:", chain_branch_id, table_name);
let delete_prefix_bytes = delete_prefix.as_bytes();
let iter = self.db.iterator(rocksdb::IteratorMode::From(
delete_prefix_bytes,
rocksdb::Direction::Forward,
));
for item in iter {
let (key, _value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(delete_prefix_bytes) {
break;
}
if let Some(row_id) = Self::parse_row_id_after_prefix(&key, delete_prefix_bytes.len()) {
deleted_rows.insert(row_id);
}
}
let branch_prefix = format!("bdata:{}:{}:", chain_branch_id, table_name);
let branch_prefix_bytes = branch_prefix.as_bytes();
let iter = self.db.iterator(rocksdb::IteratorMode::From(
branch_prefix_bytes,
rocksdb::Direction::Forward,
));
for item in iter {
let (key, raw_value) = item.map_err(|e| Error::storage(format!("Iterator error: {}", e)))?;
if !key.starts_with(branch_prefix_bytes) {
break;
}
if let Some(row_id) = Self::parse_row_id_after_prefix(&key, branch_prefix_bytes.len()) {
let value = self.decrypt_value(&raw_value)?;
let mut tuple: Tuple = bincode::deserialize(&value)
.map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
tuple.row_id = Some(row_id);
tuple.branch_id = Some(chain_branch_id);
result_tuples.insert(row_id, tuple);
}
}
}
for row_id in &deleted_rows {
result_tuples.remove(row_id);
}
tracing::debug!(
"scan_table_branch_aware: returning {} tuples ({} deleted)",
result_tuples.len(),
deleted_rows.len()
);
let mut tuples: Vec<Tuple> = result_tuples.into_values().collect();
tuples.sort_by_key(|t| t.row_id.unwrap_or(0));
Ok(tuples)
}
pub fn insert_tuple_branch_aware(&self, table_name: &str, tuple: Tuple) -> Result<u64> {
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
self.insert_tuple_branch_aware_with_schema(table_name, tuple, &schema)
}
pub fn insert_tuple_branch_aware_with_schema(
&self,
table_name: &str,
tuple: Tuple,
schema: &crate::Schema,
) -> Result<u64> {
let branch_name = self.current_branch.lock().clone();
if branch_name.is_none() || branch_name.as_deref() == Some("main") {
return self.insert_tuple_versioned_with_schema(table_name, tuple, schema);
}
let branch_id = match self.get_current_branch_id() {
Some(id) => id,
None => {
return Err(Error::query_execution(format!(
"Branch '{}' does not exist. Create it first with: CREATE BRANCH {} FROM main",
branch_name.as_deref().unwrap_or("unknown"),
branch_name.as_deref().unwrap_or("branch_name")
)));
}
};
let catalog = Catalog::new(self);
let row_id = self.next_row_id_volatile(table_name);
let mut tuple = tuple;
for (i, col) in schema.columns.iter().enumerate() {
if col.primary_key {
if let Some(v) = tuple.values.get(i) {
if matches!(v, crate::Value::Null) && i < tuple.values.len() {
#[allow(clippy::indexing_slicing)]
match col.data_type {
crate::DataType::Int2 => {
tuple.values[i] = crate::Value::Int2(row_id as i16);
}
crate::DataType::Int4 => {
tuple.values[i] = crate::Value::Int4(row_id as i32);
}
_ => {
tuple.values[i] = crate::Value::Int8(row_id as i64);
}
}
}
}
}
}
let value =
bincode::serialize(&tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?;
let key = format!("bdata:{}:{}:{}", branch_id, table_name, row_id).into_bytes();
self.put(&key, &value)?;
self.flush_row_counter(table_name)?;
let timestamp = self.next_timestamp();
let version_key = format!("bv:{}:{}:{}:{}", branch_id, table_name, row_id, timestamp);
self.put_internal(version_key.as_bytes(), &value)?;
if let Some(lsn) = self.wal_lsn() {
let _ = self.snapshot_manager.register_snapshot_with_lsn(timestamp, lsn);
} else {
let _ = self.snapshot_manager.register_snapshot(timestamp);
}
if let Err(e) = self.mv_delta_tracker.record_insert(table_name, row_id, tuple.clone()) {
tracing::warn!("Failed to record insert delta for table '{}': {}", table_name, e);
}
self.filter_delta_tracker.on_insert(table_name, row_id, &tuple, &schema);
for (i, col) in schema.columns.iter().enumerate() {
if let Some(value) = tuple.values.get(i) {
self.speculative_filter_manager.on_insert(table_name, &col.name, value);
}
}
Ok(row_id)
}
pub fn update_tuples_branch_aware(
&self,
table_name: &str,
mut updates: Vec<(u64, Tuple)>, ) -> Result<u64> {
let branch_name = self.current_branch.lock().clone();
let branch_id = if branch_name.is_none() || branch_name.as_deref() == Some("main") {
None
} else {
match self.get_current_branch_id() {
Some(id) => Some(id),
None => {
return Err(Error::query_execution(format!(
"Branch '{}' does not exist. Create it first with: CREATE BRANCH {} FROM main",
branch_name.as_deref().unwrap_or("unknown"),
branch_name.as_deref().unwrap_or("branch_name")
)));
}
}
};
let catalog = Catalog::new(self);
let schema = catalog.get_table_schema(table_name)?;
let mut update_count = 0u64;
for (row_id, tuple) in updates {
let timestamp = self.next_timestamp();
let current_key = if let Some(bid) = branch_id {
format!("bdata:{}:{}:{}", bid, table_name, row_id).into_bytes()
} else {
format!("data:{}:{}", table_name, row_id).into_bytes()
};
let mut old_tuple_for_delta: Option<Tuple> = None;
if let Ok(Some(old_value)) = self.get_internal(¤t_key) {
old_tuple_for_delta = bincode::deserialize(&old_value).ok();
let old_timestamp = timestamp.saturating_sub(1);
let old_version_key = if let Some(bid) = branch_id {
format!("bv:{}:{}:{}:{}", bid, table_name, row_id, old_timestamp)
} else {
format!("v:{}:{}:{}", table_name, row_id, old_timestamp)
};
if self.get_internal(old_version_key.as_bytes())?.is_none() {
self.put_internal(old_version_key.as_bytes(), &old_value)?;
}
}
let logical_value =
bincode::serialize(&tuple).map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?;
let value = if branch_id.is_none() && schema_uses_column_storage(&schema) {
let stored_tuple = self.transform_tuple_for_column_storage(table_name, row_id, &tuple, &schema)?;
bincode::serialize(&stored_tuple)
.map_err(|e| Error::storage(format!("Failed to serialize tuple: {}", e)))?
} else {
logical_value.clone()
};
self.put(¤t_key, &value)?;
self.log_data_update(table_name, ¤t_key, &logical_value)?;
let version_key = if let Some(bid) = branch_id {
format!("bv:{}:{}:{}:{}", bid, table_name, row_id, timestamp)
} else {
format!("v:{}:{}:{}", table_name, row_id, timestamp)
};
self.put_internal(version_key.as_bytes(), &value)?;
if let Some(lsn) = self.wal_lsn() {
let _ = self.snapshot_manager.register_snapshot_with_lsn(timestamp, lsn);
} else {
let _ = self.snapshot_manager.register_snapshot(timestamp);
}
if let Some(ref old_tuple) = old_tuple_for_delta {
if let Err(e) =
self.mv_delta_tracker
.record_update(table_name, row_id, old_tuple.clone(), tuple.clone())
{
tracing::warn!("Failed to record update delta for table '{}': {}", table_name, e);
}
}
if let Some(old_tuple) = old_tuple_for_delta.as_ref() {
self.filter_delta_tracker
.on_update(table_name, row_id, old_tuple, &tuple, &schema);
}
if branch_id.is_none() {
let _ = self.vector_indexes.on_row_update(
table_name,
row_id,
&schema,
old_tuple_for_delta.as_ref(),
&tuple,
);
}
for (i, col) in schema.columns.iter().enumerate() {
if let Some(value) = tuple.values.get(i) {
self.speculative_filter_manager.on_insert(table_name, &col.name, value);
}
}
self.row_cache.invalidate(table_name, row_id);
update_count += 1;
}
Ok(update_count)
}
pub fn delete_tuples_branch_aware(&self, table_name: &str, row_ids: Vec<u64>) -> Result<u64> {
tracing::debug!(
"delete_tuples_branch_aware: called for table '{}' with row_ids {:?}",
table_name,
row_ids
);
let timestamp = self.next_timestamp();
let catalog = Catalog::new(self);
let schema_result = catalog.get_table_schema(table_name);
let branch_name = self.current_branch.lock().clone();
tracing::debug!("delete_tuples_branch_aware: branch_name = {:?}", branch_name);
let branch_id = if branch_name.is_none() || branch_name.as_deref() == Some("main") {
None
} else {
match self.get_current_branch_id() {
Some(id) => Some(id),
None => {
return Err(Error::query_execution(format!(
"Branch '{}' does not exist. Create it first with: CREATE BRANCH {} FROM main",
branch_name.as_deref().unwrap_or("unknown"),
branch_name.as_deref().unwrap_or("branch_name")
)));
}
}
};
tracing::debug!("delete_tuples_branch_aware: resolved branch_id = {:?}", branch_id);
let Some(branch_id) = branch_id else {
let mut delete_count = 0u64;
for row_id in &row_ids {
let key = format!("data:{}:{}", table_name, row_id).into_bytes();
let mut deleted_tuple_for_delta: Option<Tuple> = None;
if let Ok(Some(old_value)) = self.get_internal(&key) {
deleted_tuple_for_delta = bincode::deserialize(&old_value).ok();
let old_timestamp = timestamp.saturating_sub(1);
let version_key = format!("v:{}:{}:{}", table_name, row_id, old_timestamp);
if self.get_internal(version_key.as_bytes())?.is_none() {
self.put_internal(version_key.as_bytes(), &old_value)?;
}
}
self.delete(&key)?;
if let Some(ref deleted_tuple) = deleted_tuple_for_delta {
if let Err(e) = self
.mv_delta_tracker
.record_delete(table_name, *row_id, deleted_tuple.clone())
{
tracing::warn!("Failed to record delete delta for table '{}': {}", table_name, e);
}
}
if let Ok(ref schema) = schema_result {
if let Some(tuple) = deleted_tuple_for_delta.as_ref() {
self.filter_delta_tracker.on_delete(table_name, *row_id, tuple, schema);
}
}
let _ = self.vector_indexes.on_row_delete(
table_name,
*row_id,
schema_result.as_ref().ok(),
deleted_tuple_for_delta.as_ref(),
);
self.row_cache.invalidate(table_name, *row_id);
delete_count += 1;
}
if let Ok(ref schema) = schema_result {
self.delete_columnar_rows_grouped(table_name, &row_ids, schema)?;
}
if let Some(lsn) = self.wal_lsn() {
let _ = self.snapshot_manager.register_snapshot_with_lsn(timestamp, lsn);
} else {
let _ = self.snapshot_manager.register_snapshot(timestamp);
}
return Ok(delete_count);
};
let mut delete_count = 0u64;
for row_id in row_ids {
let delete_key = format!("bdel:{}:{}:{}", branch_id, table_name, row_id).into_bytes();
let branch_key = format!("bdata:{}:{}:{}", branch_id, table_name, row_id);
let mut deleted_tuple_for_delta: Option<Tuple> = None;
let old_value = self.get_internal(branch_key.as_bytes()).ok().flatten().or_else(|| {
let main_key = format!("data:{}:{}", table_name, row_id);
self.get_internal(main_key.as_bytes()).ok().flatten()
});
if let Some(old_value) = old_value {
deleted_tuple_for_delta = bincode::deserialize(&old_value).ok();
let old_timestamp = timestamp.saturating_sub(1);
let version_key = format!("bv:{}:{}:{}:{}", branch_id, table_name, row_id, old_timestamp);
if self.get_internal(version_key.as_bytes())?.is_none() {
self.put_internal(version_key.as_bytes(), &old_value)?;
}
}
let delete_key_str = String::from_utf8_lossy(&delete_key);
tracing::debug!(
"delete_tuples_branch_aware: writing delete marker key '{}'",
delete_key_str
);
self.put_internal(&delete_key, &[])?;
tracing::debug!("delete_tuples_branch_aware: delete marker written successfully");
if let Some(ref deleted_tuple) = deleted_tuple_for_delta {
if let Err(e) = self
.mv_delta_tracker
.record_delete(table_name, row_id, deleted_tuple.clone())
{
tracing::warn!("Failed to record branch delete delta for table '{}': {}", table_name, e);
}
}
if let Ok(ref schema) = schema_result {
if let Some(tuple) = deleted_tuple_for_delta.as_ref() {
self.filter_delta_tracker.on_delete(table_name, row_id, tuple, schema);
}
}
self.row_cache.invalidate(table_name, row_id);
delete_count += 1;
}
if let Some(lsn) = self.wal_lsn() {
let _ = self.snapshot_manager.register_snapshot_with_lsn(timestamp, lsn);
} else {
let _ = self.snapshot_manager.register_snapshot(timestamp);
}
Ok(delete_count)
}
pub fn get_approximate_size(&self) -> u64 {
self.db
.property_int_value("rocksdb.estimate-live-data-size")
.ok()
.flatten()
.unwrap_or(0)
}
pub fn get_storage_stats(&self) -> Option<StorageStats> {
let approximate_size = self.get_approximate_size();
let key_count = self
.db
.property_int_value("rocksdb.estimate-num-keys")
.ok()
.flatten()
.unwrap_or(0);
Some(StorageStats {
approximate_size,
key_count,
})
}
pub fn vacuum(&self) -> Result<()> {
self.db.compact_range::<&[u8], &[u8]>(None, None);
Ok(())
}
pub fn vacuum_table(&self, table_name: &str) -> Result<()> {
let start_key = format!("t:{}:", table_name);
let mut end_key = format!("t:{}:", table_name).into_bytes();
end_key.push(0xff);
self.db
.compact_range(Some(start_key.as_bytes()), Some(end_key.as_slice()));
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct DirectBulkLoadResult {
pub rows_loaded: u64,
pub bytes_written: usize,
pub duration: std::time::Duration,
pub rows_per_sec: u64,
pub max_row_id: u64,
}
#[derive(Debug, Clone)]
pub struct StorageStats {
pub approximate_size: u64,
pub key_count: u64,
}