refactor(storage): use dynamic trait for storage interfaces

src/array/primitive_array.rs

@@ -98,6 +98,14 @@ impl<T: NativeType> Array for PrimitiveArray<T> {
  }
 }
 
+impl<T: NativeType> PrimitiveArray<T> {
+ /// Return the values of the array.
+ pub fn values(&self) -> &[T] {
+ assert!(self.valid.all(), "not all values are non-null");
+ &self.data
+ }
+}
+
 impl<T: NativeType> ArrayValidExt for PrimitiveArray<T> {
  fn get_valid_bitmap(&self) -> &BitVec {
  &self.valid

src/db.rs

@@ -11,14 +11,13 @@ use crate::catalog::{RootCatalog, RootCatalogRef, TableRefId};
 use crate::parser::{parse, ParserError, Statement};
 use crate::planner::Statistics;
 use crate::storage::{
- InMemoryStorage, SecondaryStorage, SecondaryStorageOptions, Storage, StorageColumnRef,
- StorageImpl, Table,
+ InMemoryStorage, SecondaryStorage, SecondaryStorageOptions, StorageColumnRef, StorageRef,
 };
 
 /// The database instance.
 pub struct Database {
  catalog: RootCatalogRef,
- storage: StorageImpl,
+ storage: StorageRef,
  config: Mutex<Config>,
 }
 
@@ -32,10 +31,10 @@ struct Config {
 impl Database {
  /// Create a new in-memory database instance.
  pub fn new_in_memory() -> Self {
- let storage = InMemoryStorage::new();
+ let storage = Arc::new(InMemoryStorage::new());
  Database {
  catalog: storage.catalog().clone(),
- storage: StorageImpl::InMemoryStorage(Arc::new(storage)),
+ storage,
  config: Default::default(),
  }
  }
@@ -46,15 +45,13 @@ impl Database {
  storage.spawn_compactor().await;
  Database {
  catalog: storage.catalog().clone(),
- storage: StorageImpl::SecondaryStorage(storage),
+ storage,
  config: Default::default(),
  }
  }
 
  pub async fn shutdown(&self) -> Result<(), Error> {
- if let StorageImpl::SecondaryStorage(storage) = &self.storage {
- storage.shutdown().await?;
- }
+ self.storage.shutdown().await?;
  Ok(())
  }
 
@@ -104,14 +101,7 @@ impl Database {
  if !self.config.lock().unwrap().disable_optimizer {
  plan = optimizer.optimize(plan);
  }
- let executor = match self.storage.clone() {
- StorageImpl::InMemoryStorage(s) => {
- crate::executor::build(optimizer.clone(), s, &plan)
- }
- StorageImpl::SecondaryStorage(s) => {
- crate::executor::build(optimizer.clone(), s, &plan)
- }
- };
+ let executor = crate::executor::build(optimizer.clone(), self.storage.clone(), &plan);
  let output = executor.try_collect().await?;
  let mut chunk = Chunk::new(output);
  chunk = bind_header(chunk, &stmt);
@@ -126,9 +116,9 @@ impl Database {
  }
  let mut stat = Statistics::default();
  // only secondary storage supports statistics
- let StorageImpl::SecondaryStorage(storage) = self.storage.clone() else {
+ if !self.storage.as_any().is::<SecondaryStorage>() {
  return Ok(stat);
- };
+ }
  for schema in self.catalog.all_schemas().values() {
  // skip internal schema
  if schema.name() == RootCatalog::SYSTEM_SCHEMA_NAME {
@@ -139,12 +129,10 @@ impl Database {
  continue;
  }
  let table_id = TableRefId::new(schema.id(), table.id());
- let table = storage.get_table(table_id)?;
+ let table = self.storage.get_table(table_id).await?;
  let txn = table.read().await?;
- let values = txn.aggreagate_block_stat(&[(
- BlockStatisticsType::RowCount,
- StorageColumnRef::Idx(0),
- )]);
+ let values =
+ txn.get_stats(&[(BlockStatisticsType::RowCount, StorageColumnRef::Idx(0))]);
  stat.add_row_count(table_id, values[0].as_usize().unwrap().unwrap() as u32);
  }
  }

src/executor/create_table.rs

@@ -1,18 +1,15 @@
 // Copyright 2024 RisingLight Project Authors. Licensed under Apache-2.0.
 
-use std::sync::Arc;
-
 use super::*;
 use crate::binder::CreateTable;
-use crate::storage::Storage;
 
 /// The executor of `create table` statement.
-pub struct CreateTableExecutor<S: Storage> {
+pub struct CreateTableExecutor {
  pub table: Box<CreateTable>,
- pub storage: Arc<S>,
+ pub storage: StorageRef,
 }
 
-impl<S: Storage> CreateTableExecutor<S> {
+impl CreateTableExecutor {
  #[try_stream(boxed, ok = DataChunk, error = ExecutorError)]
  pub async fn execute(self) {
  self.storage

src/executor/delete.rs

@@ -1,37 +1,31 @@
 // Copyright 2024 RisingLight Project Authors. Licensed under Apache-2.0.
 
-use std::sync::Arc;
-
 use super::*;
-use crate::array::DataChunk;
+use crate::array::{ArrayImpl, DataChunk};
 use crate::catalog::TableRefId;
-use crate::storage::{RowHandler, Storage, Table, Transaction};
+use crate::storage::StorageRef;
 
 /// The executor of `delete` statement.
 ///
 /// The last column of the input data chunk should be `_row_id_`.
-pub struct DeleteExecutor<S: Storage> {
+pub struct DeleteExecutor {
  pub table_id: TableRefId,
- pub storage: Arc<S>,
+ pub storage: StorageRef,
 }
 
-impl<S: Storage> DeleteExecutor<S> {
+impl DeleteExecutor {
  #[try_stream(boxed, ok = DataChunk, error = ExecutorError)]
  pub async fn execute(self, child: BoxedExecutor) {
- let table = self.storage.get_table(self.table_id)?;
+ let table = self.storage.get_table(self.table_id).await?;
  let mut txn = table.update().await?;
  let mut cnt = 0;
  #[for_await]
  for chunk in child {
  let chunk = chunk?;
- let row_handlers = chunk.array_at(chunk.column_count() - 1);
- for row_handler_idx in 0..row_handlers.len() {
- let row_handler = <S::Transaction as Transaction>::RowHandlerType::from_column(
- row_handlers,
- row_handler_idx,
- );
- txn.delete(&row_handler).await?;
- }
+ let ArrayImpl::Int64(rowids) = chunk.array_at(chunk.column_count() - 1) else {
+ panic!("column _row_id_ should be i64 type");
+ };
+ txn.delete(rowids.values()).await?;
  cnt += chunk.cardinality();
  }
  txn.commit().await?;

src/executor/drop.rs

@@ -1,19 +1,16 @@
 // Copyright 2024 RisingLight Project Authors. Licensed under Apache-2.0.
 
-use std::sync::Arc;
-
 use super::*;
 use crate::catalog::{RootCatalogRef, TableRefId};
-use crate::storage::Storage;
 
 /// The executor of `drop` statement.
-pub struct DropExecutor<S: Storage> {
+pub struct DropExecutor {
  pub tables: Vec<TableRefId>,
  pub catalog: RootCatalogRef,
- pub storage: Arc<S>,
+ pub storage: StorageRef,
 }
 
-impl<S: Storage> DropExecutor<S> {
+impl DropExecutor {
  #[try_stream(boxed, ok = DataChunk, error = ExecutorError)]
  pub async fn execute(self) {
  for table in self.tables {

src/executor/insert.rs

@@ -1,24 +1,21 @@
 // Copyright 2024 RisingLight Project Authors. Licensed under Apache-2.0.
 
-use std::sync::Arc;
-
 use super::*;
 use crate::array::DataChunk;
 use crate::catalog::{ColumnId, TableRefId};
-use crate::storage::{Storage, Table, Transaction};
 use crate::types::ColumnIndex;
 
 /// The executor of `insert` statement.
-pub struct InsertExecutor<S: Storage> {
+pub struct InsertExecutor {
  pub table_id: TableRefId,
  pub column_ids: Vec<ColumnId>,
- pub storage: Arc<S>,
+ pub storage: StorageRef,
 }
 
-impl<S: Storage> InsertExecutor<S> {
+impl InsertExecutor {
  #[try_stream(boxed, ok = DataChunk, error = ExecutorError)]
  pub async fn execute(self, child: BoxedExecutor) {
- let table = self.storage.get_table(self.table_id)?;
+ let table = self.storage.get_table(self.table_id).await?;
  let columns = table.columns()?;
 
  // construct an expression
@@ -59,7 +56,7 @@ mod tests {
  use super::*;
  use crate::array::ArrayImpl;
  use crate::catalog::{ColumnCatalog, ColumnDesc, TableRefId};
- use crate::storage::{InMemoryStorage, StorageImpl};
+ use crate::storage::{InMemoryStorage, Storage, StorageRef};
  use crate::types::DataType;
 
  #[tokio::test]
@@ -68,7 +65,7 @@ mod tests {
  let executor = InsertExecutor {
  table_id: TableRefId::new(1, 0),
  column_ids: vec![0, 1],
- storage: storage.as_in_memory_storage(),
+ storage,
  };
  let source = async_stream::try_stream! {
  yield [
@@ -82,10 +79,9 @@ mod tests {
  executor.execute(source).next().await.unwrap().unwrap();
  }
 
- async fn create_table() -> StorageImpl {
- let storage = StorageImpl::InMemoryStorage(Arc::new(InMemoryStorage::new()));
+ async fn create_table() -> StorageRef {
+ let storage = Arc::new(InMemoryStorage::new());
  storage
- .as_in_memory_storage()
  .create_table(
  1,
  "t",

src/executor/mod.rs

@@ -53,7 +53,7 @@ use self::window::*;
 use crate::array::DataChunk;
 use crate::catalog::{RootCatalog, RootCatalogRef, TableRefId};
 use crate::planner::{Expr, ExprAnalysis, Optimizer, RecExpr, TypeSchemaAnalysis};
-use crate::storage::Storage;
+use crate::storage::StorageRef;
 use crate::types::{ColumnIndex, DataType};
 
 mod copy_from_file;
@@ -95,13 +95,13 @@ const PROCESSING_WINDOW_SIZE: usize = 1024;
 /// and produces a stream to its parent.
 pub type BoxedExecutor = BoxStream<'static, Result<DataChunk>>;
 
-pub fn build(optimizer: Optimizer, storage: Arc<impl Storage>, plan: &RecExpr) -> BoxedExecutor {
+pub fn build(optimizer: Optimizer, storage: StorageRef, plan: &RecExpr) -> BoxedExecutor {
  Builder::new(optimizer, storage, plan).build()
 }
 
 /// The builder of executor.
-struct Builder<S: Storage> {
- storage: Arc<S>,
+struct Builder {
+ storage: StorageRef,
  optimizer: Optimizer,
  egraph: egg::EGraph<Expr, TypeSchemaAnalysis>,
  root: Id,
@@ -110,9 +110,9 @@ struct Builder<S: Storage> {
  views: HashMap<TableRefId, StreamSubscriber>,
 }
 
-impl<S: Storage> Builder<S> {
+impl Builder {
  /// Create a new executor builder.
- fn new(optimizer: Optimizer, storage: Arc<S>, plan: &RecExpr) -> Self {
+ fn new(optimizer: Optimizer, storage: StorageRef, plan: &RecExpr) -> Self {
  let mut egraph = egg::EGraph::new(TypeSchemaAnalysis {
  catalog: optimizer.catalog().clone(),
  });