use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::path::{Path, PathBuf};
use alopex_core::columnar::encoding::Column;
use alopex_core::columnar::segment_v2::{RecordBatch, SegmentWriterV2};
use alopex_core::storage::format::AlopexFileWriter;
use alopex_core::{StorageFactory, StorageMode as CoreStorageMode};
use crate::{Database, Error, Result, SegmentConfigV2, Transaction, TxnMode};
#[cfg(test)]
thread_local! {
static LAST_READ_COLUMN_INDICES: std::cell::RefCell<Option<Vec<usize>>> =
const { std::cell::RefCell::new(None) };
}
#[cfg(test)]
fn record_read_column_indices(indices: &[usize]) {
LAST_READ_COLUMN_INDICES.with(|last| {
*last.borrow_mut() = Some(indices.to_vec());
});
}
#[derive(Debug, Clone)]
pub struct ColumnarSegmentStats {
pub row_count: usize,
pub column_count: usize,
pub size_bytes: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ColumnarIndexType {
Minmax,
Bloom,
}
impl ColumnarIndexType {
pub fn as_str(&self) -> &'static str {
match self {
Self::Minmax => "minmax",
Self::Bloom => "bloom",
}
}
}
#[derive(Debug, Clone)]
pub struct ColumnarIndexInfo {
pub column: String,
pub index_type: ColumnarIndexType,
}
#[derive(Debug, Clone)]
pub struct EmbeddedConfig {
pub path: Option<PathBuf>,
pub storage_mode: StorageMode,
pub memory_limit: Option<usize>,
pub segment_config: SegmentConfigV2,
}
impl EmbeddedConfig {
pub fn disk(path: PathBuf) -> Self {
Self {
path: Some(path),
storage_mode: StorageMode::Disk,
memory_limit: None,
segment_config: SegmentConfigV2::default(),
}
}
pub fn in_memory() -> Self {
Self {
path: None,
storage_mode: StorageMode::InMemory,
memory_limit: None,
segment_config: SegmentConfigV2::default(),
}
}
pub fn in_memory_with_limit(limit: usize) -> Self {
Self {
path: None,
storage_mode: StorageMode::InMemory,
memory_limit: Some(limit),
segment_config: SegmentConfigV2::default(),
}
}
pub fn with_segment_config(mut self, cfg: SegmentConfigV2) -> Self {
self.segment_config = cfg;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StorageMode {
Disk,
InMemory,
}
impl Database {
pub fn open_with_config(config: EmbeddedConfig) -> Result<Self> {
let store = match config.storage_mode {
StorageMode::Disk => {
let path = config.path.clone().ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"disk mode requires a path".into(),
))
})?;
let path = crate::disk_data_dir_path(&path);
StorageFactory::create(CoreStorageMode::Disk { path, config: None })
.map_err(Error::Core)?
}
StorageMode::InMemory => StorageFactory::create(CoreStorageMode::Memory {
max_size: config.memory_limit,
})
.map_err(Error::Core)?,
};
Ok(Self::init(
store,
config.storage_mode,
config.memory_limit,
config.segment_config,
))
}
pub fn storage_mode(&self) -> StorageMode {
self.columnar_mode
}
pub fn write_columnar_segment(&self, table: &str, batch: RecordBatch) -> Result<u64> {
let mut writer = SegmentWriterV2::new(self.segment_config.clone());
writer
.write_batch(batch)
.map_err(|e| Error::Core(e.into()))?;
let segment = writer.finish().map_err(|e| Error::Core(e.into()))?;
let table_id = table_id(table)?;
match self.columnar_mode {
StorageMode::Disk => self
.columnar_bridge
.write_segment(table_id, &segment)
.map_err(|e| Error::Core(e.into())),
StorageMode::InMemory => {
let store = self.columnar_memory.as_ref().ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?;
store
.write_segment(table_id, segment)
.map_err(|e| Error::Core(e.into()))
}
}
}
pub fn write_columnar_segment_with_config(
&self,
table: &str,
batch: RecordBatch,
config: SegmentConfigV2,
) -> Result<u64> {
let mut writer = SegmentWriterV2::new(config);
writer
.write_batch(batch)
.map_err(|e| Error::Core(e.into()))?;
let segment = writer.finish().map_err(|e| Error::Core(e.into()))?;
let table_id = table_id(table)?;
match self.columnar_mode {
StorageMode::Disk => self
.columnar_bridge
.write_segment(table_id, &segment)
.map_err(|e| Error::Core(e.into())),
StorageMode::InMemory => {
let store = self.columnar_memory.as_ref().ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?;
store
.write_segment(table_id, segment)
.map_err(|e| Error::Core(e.into()))
}
}
}
pub fn read_columnar_segment(
&self,
table: &str,
segment_id: u64,
columns: Option<&[&str]>,
) -> Result<Vec<RecordBatch>> {
let table_id = table_id(table)?;
match self.columnar_mode {
StorageMode::Disk => {
let read_indices: Vec<usize> = if let Some(names) = columns {
let segment = self
.columnar_bridge
.read_segment_raw(table_id, segment_id)
.map_err(|e| Error::Core(e.into()))?;
resolve_indices_from_schema(&segment.meta.schema, names)?
} else {
let column_count = self
.columnar_bridge
.column_count(table_id, segment_id)
.map_err(|e| Error::Core(e.into()))?;
(0..column_count).collect()
};
#[cfg(test)]
record_read_column_indices(&read_indices);
self.columnar_bridge
.read_segment(table_id, segment_id, &read_indices)
.map_err(|e| Error::Core(e.into()))
}
StorageMode::InMemory => {
let store = self.columnar_memory.as_ref().ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?;
let read_indices: Vec<usize> = if let Some(names) = columns {
let schema = store
.schema(table_id, segment_id)
.map_err(|e| Error::Core(e.into()))?;
resolve_indices_from_schema(&schema, names)?
} else {
let column_count = store
.column_count(table_id, segment_id)
.map_err(|e| Error::Core(e.into()))?;
(0..column_count).collect()
};
#[cfg(test)]
record_read_column_indices(&read_indices);
store
.read_segment(table_id, segment_id, &read_indices)
.map_err(|e| Error::Core(e.into()))
}
}
}
pub fn in_memory_usage(&self) -> Option<u64> {
if self.columnar_mode == StorageMode::InMemory {
self.columnar_memory.as_ref().map(|m| m.memory_usage())
} else {
None
}
}
pub fn open_in_memory_with_limit(limit: usize) -> Result<Self> {
Self::open_with_config(EmbeddedConfig::in_memory_with_limit(limit))
}
pub fn resolve_table_id(&self, table: &str) -> Result<u32> {
table_id(table)
}
pub fn scan_columnar_segment(
&self,
segment_id: &str,
) -> Result<Vec<Vec<alopex_sql::SqlValue>>> {
let (table_id, seg_id) = parse_segment_id(segment_id)?;
let all_indices: Vec<usize> = match self.columnar_mode {
StorageMode::Disk => {
let count = self
.columnar_bridge
.column_count(table_id, seg_id)
.map_err(|e| Error::Core(e.into()))?;
(0..count).collect()
}
StorageMode::InMemory => {
let store = self.columnar_memory.as_ref().ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?;
let count = store
.column_count(table_id, seg_id)
.map_err(|e| Error::Core(e.into()))?;
(0..count).collect()
}
};
let batches = match self.columnar_mode {
StorageMode::Disk => self
.columnar_bridge
.read_segment(table_id, seg_id, &all_indices)
.map_err(|e| Error::Core(e.into()))?,
StorageMode::InMemory => self
.columnar_memory
.as_ref()
.ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?
.read_segment(table_id, seg_id, &all_indices)
.map_err(|e| Error::Core(e.into()))?,
};
let mut rows = Vec::new();
for batch in batches {
let num_rows = batch.num_rows();
for row_idx in 0..num_rows {
let mut row = Vec::with_capacity(batch.columns.len());
for col in &batch.columns {
let sql_val = column_value_to_sql_value(col, row_idx);
row.push(sql_val);
}
rows.push(row);
}
}
Ok(rows)
}
pub fn scan_columnar_segment_batches(&self, segment_id: &str) -> Result<Vec<RecordBatch>> {
let (table_id, seg_id) = parse_segment_id(segment_id)?;
let all_indices: Vec<usize> = match self.columnar_mode {
StorageMode::Disk => {
let count = self
.columnar_bridge
.column_count(table_id, seg_id)
.map_err(|e| Error::Core(e.into()))?;
(0..count).collect()
}
StorageMode::InMemory => {
let store = self.columnar_memory.as_ref().ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?;
let count = store
.column_count(table_id, seg_id)
.map_err(|e| Error::Core(e.into()))?;
(0..count).collect()
}
};
match self.columnar_mode {
StorageMode::Disk => self
.columnar_bridge
.read_segment(table_id, seg_id, &all_indices)
.map_err(|e| Error::Core(e.into())),
StorageMode::InMemory => self
.columnar_memory
.as_ref()
.ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?
.read_segment(table_id, seg_id, &all_indices)
.map_err(|e| Error::Core(e.into())),
}
}
pub fn scan_columnar_segment_streaming(&self, segment_id: &str) -> Result<ColumnarRowIterator> {
let batches = self.scan_columnar_segment_batches(segment_id)?;
Ok(ColumnarRowIterator::new(batches))
}
pub fn get_columnar_segment_stats(&self, segment_id: &str) -> Result<ColumnarSegmentStats> {
let (table_id, seg_id) = parse_segment_id(segment_id)?;
match self.columnar_mode {
StorageMode::Disk => {
let column_count = self
.columnar_bridge
.column_count(table_id, seg_id)
.map_err(|e| Error::Core(e.into()))?;
let batches = self
.columnar_bridge
.read_segment(table_id, seg_id, &(0..column_count).collect::<Vec<_>>())
.map_err(|e| Error::Core(e.into()))?;
let row_count: usize = batches.iter().map(|b| b.num_rows()).sum();
Ok(ColumnarSegmentStats {
row_count,
column_count,
size_bytes: 0, })
}
StorageMode::InMemory => {
let store = self.columnar_memory.as_ref().ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?;
let column_count = store
.column_count(table_id, seg_id)
.map_err(|e| Error::Core(e.into()))?;
let batches = store
.read_segment(table_id, seg_id, &(0..column_count).collect::<Vec<_>>())
.map_err(|e| Error::Core(e.into()))?;
let row_count: usize = batches.iter().map(|b| b.num_rows()).sum();
Ok(ColumnarSegmentStats {
row_count,
column_count,
size_bytes: 0, })
}
}
}
pub fn list_columnar_segments(&self) -> Result<Vec<String>> {
match self.columnar_mode {
StorageMode::Disk => {
let segments = self
.columnar_bridge
.list_segments()
.map_err(|e| Error::Core(e.into()))?;
Ok(segments
.into_iter()
.map(|(table_id, seg_id)| format!("{}:{}", table_id, seg_id))
.collect())
}
StorageMode::InMemory => {
let store = self.columnar_memory.as_ref().ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(
"in-memory columnar store is not initialized".into(),
))
})?;
let segments = store.list_segments();
Ok(segments
.into_iter()
.map(|(table_id, seg_id)| format!("{}:{}", table_id, seg_id))
.collect())
}
}
}
pub fn create_columnar_index(
&self,
segment_id: &str,
column: &str,
index_type: ColumnarIndexType,
) -> Result<()> {
let _ = self.get_columnar_segment_stats(segment_id)?;
let key = columnar_index_key(segment_id, column);
let value = index_type.as_str().as_bytes().to_vec();
let mut txn = self.begin(TxnMode::ReadWrite)?;
txn.put(&key, &value)?;
txn.commit()?;
Ok(())
}
pub fn list_columnar_indexes(&self, segment_id: &str) -> Result<Vec<ColumnarIndexInfo>> {
let _ = self.get_columnar_segment_stats(segment_id)?;
let prefix = columnar_index_prefix(segment_id);
let mut txn = self.begin(TxnMode::ReadOnly)?;
let mut entries = Vec::new();
for (key, value) in txn.scan_prefix(&prefix)? {
let column = parse_index_column(segment_id, &key)?;
let index_type = parse_index_type(&value)?;
entries.push(ColumnarIndexInfo { column, index_type });
}
txn.commit()?;
Ok(entries)
}
pub fn drop_columnar_index(&self, segment_id: &str, column: &str) -> Result<()> {
let _ = self.get_columnar_segment_stats(segment_id)?;
let key = columnar_index_key(segment_id, column);
let mut txn = self.begin(TxnMode::ReadWrite)?;
let exists = txn.get(&key)?.is_some();
if !exists {
txn.rollback()?;
return Err(Error::IndexNotFound(format!(
"columnar index {}:{}",
segment_id, column
)));
}
txn.delete(&key)?;
txn.commit()?;
Ok(())
}
pub fn flush_in_memory_segment_to_file(
&self,
table: &str,
segment_id: u64,
path: &Path,
) -> Result<()> {
let store = self
.columnar_memory
.as_ref()
.ok_or(Error::NotInMemoryMode)?;
let table_id = table_id(table)?;
store
.flush_to_segment_file(table_id, segment_id, path)
.map_err(|e| Error::Core(e.into()))
}
pub fn flush_in_memory_segment_to_kvs(&self, table: &str, segment_id: u64) -> Result<u64> {
let store = self
.columnar_memory
.as_ref()
.ok_or(Error::NotInMemoryMode)?;
let table_id = table_id(table)?;
store
.flush_to_kvs(table_id, segment_id, &self.columnar_bridge)
.map_err(|e| Error::Core(e.into()))
}
pub fn flush_in_memory_segment_to_alopex(
&self,
table: &str,
segment_id: u64,
writer: &mut AlopexFileWriter,
) -> Result<u32> {
let store = self
.columnar_memory
.as_ref()
.ok_or(Error::NotInMemoryMode)?;
let table_id = table_id(table)?;
store
.flush_to_alopex(table_id, segment_id, writer)
.map_err(|e| Error::Core(e.into()))
}
}
impl<'a> Transaction<'a> {
pub fn storage_mode(&self) -> StorageMode {
self.db.storage_mode()
}
pub fn write_columnar_segment(&self, table: &str, batch: RecordBatch) -> Result<u64> {
self.db.write_columnar_segment(table, batch)
}
pub fn read_columnar_segment(
&self,
table: &str,
segment_id: u64,
columns: Option<&[&str]>,
) -> Result<Vec<RecordBatch>> {
self.db.read_columnar_segment(table, segment_id, columns)
}
}
fn table_id(table: &str) -> Result<u32> {
if table.is_empty() {
return Err(Error::TableNotFound("table name is empty".into()));
}
let mut hasher = DefaultHasher::new();
table.hash(&mut hasher);
Ok((hasher.finish() & 0xffff_ffff) as u32)
}
fn resolve_indices_from_schema(
schema: &alopex_core::columnar::segment_v2::Schema,
names: &[&str],
) -> Result<Vec<usize>> {
let mut indices = Vec::with_capacity(names.len());
for name in names {
let pos = schema
.columns
.iter()
.position(|c| c.name == *name)
.ok_or_else(|| {
Error::Core(alopex_core::Error::InvalidFormat(format!(
"column not found: {name}"
)))
})?;
indices.push(pos);
}
Ok(indices)
}
const COLUMNAR_INDEX_PREFIX: &str = "__alopex_columnar_index__:";
fn columnar_index_key(segment: &str, column: &str) -> Vec<u8> {
let mut key =
String::with_capacity(COLUMNAR_INDEX_PREFIX.len() + segment.len() + column.len() + 1);
key.push_str(COLUMNAR_INDEX_PREFIX);
key.push_str(segment);
key.push(':');
key.push_str(column);
key.into_bytes()
}
fn columnar_index_prefix(segment: &str) -> Vec<u8> {
let mut key = String::with_capacity(COLUMNAR_INDEX_PREFIX.len() + segment.len() + 1);
key.push_str(COLUMNAR_INDEX_PREFIX);
key.push_str(segment);
key.push(':');
key.into_bytes()
}
fn parse_index_column(segment: &str, key: &[u8]) -> Result<String> {
let prefix = columnar_index_prefix(segment);
if !key.starts_with(&prefix) {
return Err(Error::Core(alopex_core::Error::InvalidFormat(
"columnar index key is invalid".into(),
)));
}
let suffix = &key[prefix.len()..];
String::from_utf8(suffix.to_vec()).map_err(|_| {
Error::Core(alopex_core::Error::InvalidFormat(
"columnar index column is not valid UTF-8".into(),
))
})
}
fn parse_index_type(raw: &[u8]) -> Result<ColumnarIndexType> {
let value = std::str::from_utf8(raw).map_err(|_| {
Error::Core(alopex_core::Error::InvalidFormat(
"columnar index type is not valid UTF-8".into(),
))
})?;
match value {
"minmax" => Ok(ColumnarIndexType::Minmax),
"bloom" => Ok(ColumnarIndexType::Bloom),
other => Err(Error::Core(alopex_core::Error::InvalidFormat(format!(
"unknown columnar index type: {other}"
)))),
}
}
fn parse_segment_id(segment_id: &str) -> Result<(u32, u64)> {
let parts: Vec<&str> = segment_id.split(':').collect();
if parts.len() != 2 {
return Err(Error::Core(alopex_core::Error::InvalidFormat(format!(
"invalid segment ID format: expected 'table_id:segment_id', got '{}'",
segment_id
))));
}
let table_id: u32 = parts[0].parse().map_err(|_| {
Error::Core(alopex_core::Error::InvalidFormat(format!(
"invalid table_id in segment ID: '{}'",
parts[0]
)))
})?;
let seg_id: u64 = parts[1].parse().map_err(|_| {
Error::Core(alopex_core::Error::InvalidFormat(format!(
"invalid segment_id in segment ID: '{}'",
parts[1]
)))
})?;
Ok((table_id, seg_id))
}
fn column_value_to_sql_value(col: &Column, row_idx: usize) -> alopex_sql::SqlValue {
match col {
Column::Int64(vals) => vals
.get(row_idx)
.map(|&v| alopex_sql::SqlValue::BigInt(v))
.unwrap_or(alopex_sql::SqlValue::Null),
Column::Float32(vals) => vals
.get(row_idx)
.map(|&v| alopex_sql::SqlValue::Float(v))
.unwrap_or(alopex_sql::SqlValue::Null),
Column::Float64(vals) => vals
.get(row_idx)
.map(|&v| alopex_sql::SqlValue::Double(v))
.unwrap_or(alopex_sql::SqlValue::Null),
Column::Bool(vals) => vals
.get(row_idx)
.map(|&v| alopex_sql::SqlValue::Boolean(v))
.unwrap_or(alopex_sql::SqlValue::Null),
Column::Binary(vals) => vals
.get(row_idx)
.map(|v| alopex_sql::SqlValue::Blob(v.clone()))
.unwrap_or(alopex_sql::SqlValue::Null),
Column::Fixed { values, .. } => values
.get(row_idx)
.map(|v| alopex_sql::SqlValue::Blob(v.clone()))
.unwrap_or(alopex_sql::SqlValue::Null),
}
}
pub struct ColumnarRowIterator {
batches: Vec<RecordBatch>,
batch_idx: usize,
row_idx: usize,
}
impl ColumnarRowIterator {
pub fn new(batches: Vec<RecordBatch>) -> Self {
Self {
batches,
batch_idx: 0,
row_idx: 0,
}
}
pub fn batch_count(&self) -> usize {
self.batches.len()
}
pub fn current_batch(&self) -> Option<&RecordBatch> {
self.batches.get(self.batch_idx)
}
}
impl Iterator for ColumnarRowIterator {
type Item = Vec<alopex_sql::SqlValue>;
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.batch_idx >= self.batches.len() {
return None;
}
let batch = &self.batches[self.batch_idx];
let row_count = batch.num_rows();
if self.row_idx >= row_count {
self.batch_idx += 1;
self.row_idx = 0;
continue;
}
let row_idx = self.row_idx;
self.row_idx += 1;
let mut row = Vec::with_capacity(batch.columns.len());
for col in &batch.columns {
let sql_val = column_value_to_sql_value(col, row_idx);
row.push(sql_val);
}
return Some(row);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use alopex_core::columnar::encoding::{Column, LogicalType};
use alopex_core::columnar::error::{ColumnarError, Result as ColumnarResult};
use alopex_core::columnar::segment_v2::{ColumnSchema, Schema, SegmentReaderV2, SegmentSource};
use alopex_core::storage::format::{AlopexFileWriter, FileFlags, FileVersion};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use tempfile::tempdir;
fn make_batch() -> RecordBatch {
let schema = Schema {
columns: vec![
ColumnSchema {
name: "id".into(),
logical_type: LogicalType::Int64,
nullable: false,
fixed_len: None,
},
ColumnSchema {
name: "val".into(),
logical_type: LogicalType::Int64,
nullable: false,
fixed_len: None,
},
],
};
RecordBatch::new(
schema,
vec![
Column::Int64(vec![1, 2, 3]),
Column::Int64(vec![10, 20, 30]),
],
vec![None, None],
)
}
fn make_wide_batch(column_count: usize, row_count: usize) -> RecordBatch {
let schema = Schema {
columns: (0..column_count)
.map(|idx| ColumnSchema {
name: format!("c{idx}"),
logical_type: LogicalType::Int64,
nullable: false,
fixed_len: None,
})
.collect(),
};
let columns = (0..column_count)
.map(|idx| {
Column::Int64(
(0..row_count)
.map(|row| (idx as i64 * 1_000_000) + row as i64)
.collect(),
)
})
.collect();
RecordBatch::new(schema, columns, vec![None; column_count])
}
fn decoded_payload_bytes(batches: &[RecordBatch]) -> usize {
batches
.iter()
.flat_map(|batch| batch.columns.iter())
.map(|column| match column {
Column::Int64(values) => values.len() * std::mem::size_of::<i64>(),
Column::Float32(values) => values.len() * std::mem::size_of::<f32>(),
Column::Float64(values) => values.len() * std::mem::size_of::<f64>(),
Column::Bool(values) => values.len() * std::mem::size_of::<bool>(),
Column::Binary(values) => values.iter().map(Vec::len).sum(),
Column::Fixed { values, .. } => values.iter().map(Vec::len).sum(),
})
.sum()
}
#[derive(Debug, Clone)]
struct CountingSegmentSource {
data: Arc<Vec<u8>>,
bytes: Arc<AtomicU64>,
calls: Arc<AtomicUsize>,
}
impl CountingSegmentSource {
fn new(data: Vec<u8>) -> Self {
Self {
data: Arc::new(data),
bytes: Arc::new(AtomicU64::new(0)),
calls: Arc::new(AtomicUsize::new(0)),
}
}
fn reset(&self) {
self.bytes.store(0, Ordering::Relaxed);
self.calls.store(0, Ordering::Relaxed);
}
fn bytes(&self) -> u64 {
self.bytes.load(Ordering::Relaxed)
}
fn calls(&self) -> usize {
self.calls.load(Ordering::Relaxed)
}
}
impl SegmentSource for CountingSegmentSource {
fn read_range(&self, offset: u64, len: u64) -> ColumnarResult<Vec<u8>> {
self.bytes.fetch_add(len, Ordering::Relaxed);
self.calls.fetch_add(1, Ordering::Relaxed);
let start = offset as usize;
let end = start + len as usize;
if end > self.data.len() {
return Err(ColumnarError::InvalidFormat("range out of bounds".into()));
}
Ok(self.data[start..end].to_vec())
}
fn total_size(&self) -> u64 {
self.data.len() as u64
}
}
fn measured_segment_read(data: Vec<u8>, columns: &[usize]) -> ColumnarResult<(u64, usize)> {
let source = CountingSegmentSource::new(data);
let reader = SegmentReaderV2::open(Box::new(source.clone()))?;
source.reset();
let batches = reader.read_columns(columns)?;
if batches.is_empty() {
return Err(ColumnarError::InvalidFormat("segment is empty".into()));
}
Ok((source.bytes(), source.calls()))
}
fn reset_last_read_column_indices() {
LAST_READ_COLUMN_INDICES.with(|last| {
*last.borrow_mut() = None;
});
}
fn last_read_column_indices() -> Vec<usize> {
LAST_READ_COLUMN_INDICES.with(|last| {
last.borrow()
.clone()
.expect("read_columnar_segment should record read indices")
})
}
#[test]
fn write_read_disk_mode() {
let dir = tempdir().unwrap();
let wal = dir.path().join("wal.log");
let cfg = EmbeddedConfig::disk(wal);
let db = Database::open_with_config(cfg).unwrap();
let seg_id = db.write_columnar_segment("tbl", make_batch()).unwrap();
let batches = db.read_columnar_segment("tbl", seg_id, None).unwrap();
assert_eq!(batches[0].num_rows(), 3);
}
#[test]
fn read_with_column_names() {
let dir = tempdir().unwrap();
let wal = dir.path().join("wal.log");
let cfg = EmbeddedConfig::disk(wal);
let db = Database::open_with_config(cfg).unwrap();
let seg_id = db.write_columnar_segment("tbl", make_batch()).unwrap();
let batches = db
.read_columnar_segment("tbl", seg_id, Some(&["val"]))
.unwrap();
assert_eq!(batches[0].columns.len(), 1);
if let Column::Int64(vals) = &batches[0].columns[0] {
assert_eq!(vals, &vec![10, 20, 30]);
} else {
panic!("expected int64");
}
}
#[test]
fn disk_projection_pushes_selected_columns_to_segment_reader() {
let dir = tempdir().unwrap();
let wal = dir.path().join("wal.log");
let cfg = EmbeddedConfig::disk(wal);
let db = Database::open_with_config(cfg).unwrap();
let seg_id = db
.write_columnar_segment("wide_tbl", make_wide_batch(12, 4096))
.unwrap();
let table_id = table_id("wide_tbl").unwrap();
let raw = db
.columnar_bridge
.read_segment_raw(table_id, seg_id)
.unwrap();
let all_indices: Vec<usize> = (0..12).collect();
let projection = [2, 7, 10];
let (full_read_bytes, full_read_calls) =
measured_segment_read(raw.data.clone(), &all_indices).unwrap();
let (projected_read_bytes, projected_read_calls) =
measured_segment_read(raw.data, &projection).unwrap();
reset_last_read_column_indices();
let full_batches = db.read_columnar_segment("wide_tbl", seg_id, None).unwrap();
assert_eq!(last_read_column_indices(), all_indices);
reset_last_read_column_indices();
let projected_batches = db
.read_columnar_segment("wide_tbl", seg_id, Some(&["c2", "c7", "c10"]))
.unwrap();
let pushed_indices = last_read_column_indices();
let full_payload_bytes = decoded_payload_bytes(&full_batches);
let projected_payload_bytes = decoded_payload_bytes(&projected_batches);
eprintln!(
"projection pushed_indices={pushed_indices:?} full_read_bytes={full_read_bytes} projected_read_bytes={projected_read_bytes} full_read_calls={full_read_calls} projected_read_calls={projected_read_calls} full_payload_bytes={full_payload_bytes} projected_payload_bytes={projected_payload_bytes}"
);
assert_eq!(pushed_indices, projection);
assert!(projected_read_bytes < full_read_bytes);
assert!(projected_payload_bytes < full_payload_bytes);
assert_eq!(projected_batches[0].schema.columns.len(), projection.len());
}
#[test]
fn in_memory_projection_pushes_selected_columns_to_segment_reader() {
let db = Database::open_with_config(EmbeddedConfig::in_memory()).unwrap();
let seg_id = db
.write_columnar_segment("wide_mem_tbl", make_wide_batch(12, 4096))
.unwrap();
let all_indices: Vec<usize> = (0..12).collect();
let projection = [2, 7, 10];
reset_last_read_column_indices();
let full_batches = db
.read_columnar_segment("wide_mem_tbl", seg_id, None)
.unwrap();
assert_eq!(last_read_column_indices(), all_indices);
reset_last_read_column_indices();
let projected_batches = db
.read_columnar_segment("wide_mem_tbl", seg_id, Some(&["c2", "c7", "c10"]))
.unwrap();
let pushed_indices = last_read_column_indices();
let full_payload_bytes = decoded_payload_bytes(&full_batches);
let projected_payload_bytes = decoded_payload_bytes(&projected_batches);
eprintln!(
"in_memory_projection pushed_indices={pushed_indices:?} full_payload_bytes={full_payload_bytes} projected_payload_bytes={projected_payload_bytes}"
);
assert_eq!(pushed_indices, projection);
assert!(projected_payload_bytes < full_payload_bytes);
assert_eq!(projected_batches[0].schema.columns.len(), projection.len());
}
#[test]
fn in_memory_limit_rejects_large_segment() {
let cfg = EmbeddedConfig::in_memory_with_limit(1);
let db = Database::open_with_config(cfg).unwrap();
let err = db
.write_columnar_segment("tbl", make_batch())
.expect_err("should exceed limit");
assert!(format!("{err}").contains("memory limit exceeded"));
}
#[test]
fn storage_mode_flags() {
let dir = tempdir().unwrap();
let wal = dir.path().join("wal.log");
let disk = Database::open_with_config(EmbeddedConfig::disk(wal)).unwrap();
assert!(matches!(disk.storage_mode(), StorageMode::Disk));
let mem = Database::open_with_config(EmbeddedConfig::in_memory()).unwrap();
assert!(matches!(mem.storage_mode(), StorageMode::InMemory));
}
#[test]
fn transaction_write_and_read() {
let dir = tempdir().unwrap();
let wal = dir.path().join("wal.log");
let db = Database::open_with_config(EmbeddedConfig::disk(wal)).unwrap();
let txn = db.begin(crate::TxnMode::ReadWrite).unwrap();
let seg_id = txn.write_columnar_segment("tbl_txn", make_batch()).unwrap();
txn.commit().unwrap();
let batches = db
.read_columnar_segment("tbl_txn", seg_id, Some(&["id"]))
.unwrap();
assert_eq!(batches[0].num_rows(), 3);
}
#[test]
fn flush_in_memory_paths() {
let dir = tempdir().unwrap();
let db = Database::open_with_config(EmbeddedConfig::in_memory()).unwrap();
let seg_id = db.write_columnar_segment("mem_tbl", make_batch()).unwrap();
let file_path = dir.path().join("seg.bin");
db.flush_in_memory_segment_to_file("mem_tbl", seg_id, &file_path)
.unwrap();
let bytes = std::fs::read(&file_path).unwrap();
assert!(!bytes.is_empty());
let kv_id = db
.flush_in_memory_segment_to_kvs("mem_tbl", seg_id)
.unwrap();
assert_eq!(kv_id, 0);
let alo_path = dir.path().join("out.alopex");
let mut writer =
AlopexFileWriter::new(alo_path.clone(), FileVersion::CURRENT, FileFlags(0)).unwrap();
db.flush_in_memory_segment_to_alopex("mem_tbl", seg_id, &mut writer)
.unwrap();
writer.finalize().unwrap();
assert!(alo_path.exists());
}
#[test]
fn flush_not_in_memory_mode_errors() {
let dir = tempdir().unwrap();
let wal = dir.path().join("wal.log");
let db = Database::open_with_config(EmbeddedConfig::disk(wal)).unwrap();
let err = db
.flush_in_memory_segment_to_kvs("tbl", 0)
.expect_err("should error");
assert!(matches!(err, Error::NotInMemoryMode));
}
}