use std::sync::Arc;
use bytes::Bytes;
use object_store::path::Path;
use object_store::ObjectStore;
use ulid::Ulid;
use crate::block_iterator::DataBlockIterator;
use crate::db_cache::DbCache;
use crate::db_state::{SsTableHandle, SsTableId, SsTableInfo};
use crate::format::sst::{BlockTransformer, SsTableFormat};
use crate::iter::IterationOrder;
use crate::object_stores::ObjectStores;
use crate::sst_stats::SstStats;
use crate::tablestore::{SstFileMetadata, TableStore};
use crate::types::RowEntry;
pub struct SstReader {
table_store: Arc<TableStore>,
}
impl SstReader {
pub fn new<P: Into<Path>>(
root_path: P,
object_store: Arc<dyn ObjectStore>,
cache: Option<Arc<dyn DbCache>>,
block_transformer: Option<Arc<dyn BlockTransformer>>,
) -> Self {
let sst_format = SsTableFormat {
block_transformer,
..SsTableFormat::default()
};
let table_store = Arc::new(TableStore::new(
ObjectStores::new(object_store, None),
sst_format,
root_path.into(),
cache,
));
Self { table_store }
}
pub async fn open(&self, id: Ulid) -> Result<SstFile, crate::Error> {
let handle = self.table_store.open_sst(&SsTableId::Compacted(id)).await?;
Ok(SstFile {
id,
handle,
table_store: Arc::clone(&self.table_store),
})
}
pub fn open_with_handle(&self, handle: SsTableHandle) -> Result<SstFile, crate::Error> {
let id = match handle.id {
SsTableId::Compacted(ulid) => ulid,
SsTableId::Wal(_) => {
return Err(crate::Error::invalid(
"SstReader only supports compacted SSTs, not WAL SSTs".to_string(),
));
}
};
Ok(SstFile {
id,
handle,
table_store: Arc::clone(&self.table_store),
})
}
}
pub struct SstFile {
id: Ulid,
handle: SsTableHandle,
table_store: Arc<TableStore>,
}
impl SstFile {
pub fn id(&self) -> Ulid {
self.id
}
pub fn info(&self) -> &SsTableInfo {
&self.handle.info
}
pub async fn metadata(&self) -> Result<SstFileMetadata, crate::Error> {
self.table_store
.metadata(&self.handle.id)
.await
.map_err(Into::into)
}
pub async fn stats(&self) -> Result<Option<SstStats>, crate::Error> {
self.table_store
.read_stats(&self.handle, true)
.await
.map_err(Into::into)
}
pub async fn index(&self) -> Result<Vec<(u64, Bytes)>, crate::Error> {
let index = self.table_store.read_index(&self.handle, true).await?;
let borrowed = index.borrow();
let block_meta = borrowed.block_meta();
let result: Vec<(u64, Bytes)> = (0..block_meta.len())
.map(|i| {
let meta = block_meta.get(i);
(
meta.offset(),
Bytes::copy_from_slice(meta.first_key().bytes()),
)
})
.collect();
Ok(result)
}
pub async fn read_block(&self, block: usize) -> Result<Vec<RowEntry>, crate::Error> {
let index = self.table_store.read_index(&self.handle, true).await?;
let num_blocks = index.borrow().block_meta().len();
if block >= num_blocks {
return Err(crate::Error::invalid(format!(
"block index {block} out of range for SST with {num_blocks} block(s)"
)));
}
let mut blocks = self
.table_store
.read_blocks_using_index(&self.handle, index, block..block + 1, true)
.await?;
let block = blocks
.pop_front()
.expect("read_blocks_using_index returned no block for a valid index");
let mut iter =
DataBlockIterator::new(block, self.handle.format_version, IterationOrder::Ascending)?;
let mut rows = Vec::new();
while let Some(row) = iter.next().await? {
rows.push(row);
}
Ok(rows)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::{FlushOptions, FlushType, PutOptions, SstBlockSize, WriteOptions};
use crate::test_utils::StringConcatMergeOperator;
use crate::types::ValueDeletable;
use crate::Db;
use object_store::memory::InMemory;
async fn setup_db_with_l0() -> (
Arc<dyn ObjectStore>,
&'static str,
crate::manifest::VersionedManifest,
) {
let object_store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
let path = "/test_sst_reader";
let db = Db::builder(path, object_store.clone())
.with_merge_operator(Arc::new(StringConcatMergeOperator))
.with_sst_block_size(SstBlockSize::Other(64))
.build()
.await
.unwrap();
for i in 0..10u8 {
db.put_with_options(
&[b'k', i],
&[b'v', i],
&PutOptions::default(),
&WriteOptions::default(),
)
.await
.unwrap();
}
for i in 10..13u8 {
db.delete(&[b'k', i]).await.unwrap();
}
for i in 13..15u8 {
db.merge(&[b'k', i], &[b'm', i]).await.unwrap();
}
db.flush_with_options(FlushOptions {
flush_type: FlushType::MemTable,
})
.await
.unwrap();
let manifest = db.manifest();
db.close().await.unwrap();
(object_store, path, manifest)
}
#[tokio::test]
async fn test_open_and_info() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
assert!(
!manifest.manifest.core.tree.l0.is_empty(),
"expected at least one L0 SST"
);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader
.open(view.sst.id.unwrap_compacted_id())
.await
.unwrap();
let info = sst_file.info();
assert!(info.first_entry.is_some());
assert!(info.index_offset > 0);
}
#[tokio::test]
async fn test_open_with_handle() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
assert_eq!(sst_file.id(), view.sst.id.unwrap_compacted_id());
assert_eq!(sst_file.info(), &view.sst.info);
}
#[tokio::test]
async fn test_stats() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
let stats = sst_file.stats().await.unwrap();
let stats = stats.expect("expected stats block to be present");
assert_eq!(stats.num_puts, 10);
assert_eq!(stats.num_deletes, 3);
assert_eq!(stats.num_merges, 2);
assert_eq!(stats.num_rows(), 15);
assert!(stats.raw_key_size > 0);
assert!(stats.raw_val_size > 0);
}
#[tokio::test]
async fn test_index() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
let index = sst_file.index().await.unwrap();
assert!(!index.is_empty());
if let Some(first_entry) = sst_file.info().first_entry.as_ref() {
assert!(index[0].1.as_ref() <= first_entry.as_ref());
}
for window in index.windows(2) {
assert!(window[0].0 < window[1].0);
}
}
#[tokio::test]
async fn test_block_stats_parallel_to_index() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
let stats = sst_file
.stats()
.await
.unwrap()
.expect("expected stats block");
let index = sst_file.index().await.unwrap();
assert_eq!(
stats.block_stats.len(),
index.len(),
"block_stats should be parallel to the index"
);
let sum_puts: u64 = stats.block_stats.iter().map(|b| b.num_puts as u64).sum();
assert_eq!(sum_puts, stats.num_puts);
}
#[tokio::test]
async fn test_read_block_returns_all_kinds() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
let index = sst_file.index().await.unwrap();
assert!(!index.is_empty());
let mut all_rows: Vec<RowEntry> = Vec::new();
for i in 0..index.len() {
let rows = sst_file.read_block(i).await.unwrap();
assert!(!rows.is_empty(), "block {i} should have at least one row");
for window in rows.windows(2) {
assert!(window[0].key <= window[1].key);
}
all_rows.extend(rows);
}
let puts = all_rows
.iter()
.filter(|r| matches!(r.value, ValueDeletable::Value(_)))
.count();
let merges = all_rows
.iter()
.filter(|r| matches!(r.value, ValueDeletable::Merge(_)))
.count();
let tombstones = all_rows
.iter()
.filter(|r| matches!(r.value, ValueDeletable::Tombstone))
.count();
assert_eq!(puts, 10);
assert_eq!(merges, 2);
assert_eq!(tombstones, 3);
}
#[tokio::test]
async fn test_read_block_out_of_range() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
let num_blocks = sst_file.index().await.unwrap().len();
let err = sst_file.read_block(num_blocks).await.unwrap_err();
assert!(format!("{err}").contains("out of range"));
}
#[tokio::test]
async fn test_read_block_first_key_matches_index() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
let index = sst_file.index().await.unwrap();
for (i, (_, first_key)) in index.iter().enumerate() {
let rows = sst_file.read_block(i).await.unwrap();
assert!(rows[0].key >= *first_key);
}
}
#[tokio::test]
async fn test_metadata() {
let (store, path, manifest) = setup_db_with_l0().await;
let reader = SstReader::new(path, store, None, None);
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
let metadata = sst_file.metadata().await.unwrap();
assert!(metadata.size > 0);
assert!(matches!(metadata.id, SsTableId::Compacted(_)));
}
#[tokio::test]
async fn test_open_with_wal_handle_returns_error() {
let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
let reader = SstReader::new("/test", store, None, None);
let wal_handle = SsTableHandle::new(
SsTableId::Wal(42),
0,
crate::db_state::SsTableInfo::default(),
);
let result = reader.open_with_handle(wal_handle);
assert!(result.is_err());
}
#[tokio::test]
async fn test_open_missing_sst() {
let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
let reader = SstReader::new("/nonexistent", store, None, None);
let result = reader.open(Ulid::new()).await;
assert!(result.is_err());
}
struct XorTransformer {
key: u8,
}
#[async_trait::async_trait]
impl BlockTransformer for XorTransformer {
async fn encode(&self, data: Bytes) -> Result<Bytes, crate::Error> {
let transformed: Vec<u8> = data.iter().map(|b| b ^ self.key).collect();
Ok(Bytes::from(transformed))
}
async fn decode(&self, data: Bytes) -> Result<Bytes, crate::Error> {
self.encode(data).await
}
}
#[tokio::test]
async fn test_stats_with_block_transformer() {
let object_store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
let path = "/test_sst_reader_transformer";
let transformer: Arc<dyn BlockTransformer> = Arc::new(XorTransformer { key: 0x42 });
let db = Db::builder(path, object_store.clone())
.with_block_transformer(transformer.clone())
.build()
.await
.unwrap();
for i in 0..5u8 {
db.put_with_options(
&[b'k', i],
&[b'v', i],
&PutOptions::default(),
&WriteOptions::default(),
)
.await
.unwrap();
}
db.flush_with_options(FlushOptions {
flush_type: FlushType::MemTable,
})
.await
.unwrap();
let manifest = db.manifest();
db.close().await.unwrap();
let reader = SstReader::new(path, object_store.clone(), None, Some(transformer));
let view = &manifest.manifest.core.tree.l0[0];
let sst_file = reader.open_with_handle(view.sst.clone()).unwrap();
let stats = sst_file.stats().await.unwrap().expect("expected stats");
assert_eq!(stats.num_puts, 5);
let reader_no_transformer = SstReader::new(path, object_store, None, None);
let sst_file = reader_no_transformer
.open_with_handle(view.sst.clone())
.unwrap();
assert!(sst_file.stats().await.is_err());
}
}