feat(storage): add memtable (#6)

* feat(storage): add memtable

Signed-off-by: Alex Chi <[email protected]>

* in memory to disk

Signed-off-by: Alex Chi <[email protected]>

Author: skyzh

code/03-00/src/db.rs

@@ -12,7 +12,7 @@ use crate::executor::{ExecuteError, ExecutorBuilder};
 use crate::logical_planner::{LogicalPlanError, LogicalPlanner};
 use crate::parser::{parse, ParserError};
 use crate::physical_planner::{PhysicalPlanError, PhysicalPlanner};
-use crate::storage::InMemoryStorage;
+use crate::storage::DiskStorage;
 
 /// The database instance.
 pub struct Database {
@@ -31,7 +31,7 @@ impl Database {
     /// Create a new database instance.
     pub fn new() -> Self {
         let catalog = Arc::new(DatabaseCatalog::new());
-        let storage = Arc::new(InMemoryStorage::new());
+        let storage = Arc::new(DiskStorage::new());
         let parallel = matches!(std::env::var("LIGHT_PARALLEL"), Ok(s) if s == "1");
         let runtime = if parallel {
             tokio::runtime::Builder::new_multi_thread()

code/03-00/src/storage/mod.rs

@@ -20,37 +20,37 @@ pub enum StorageError {
 /// A specialized `Result` type for storage operations.
 pub type StorageResult<T> = std::result::Result<T, StorageError>;
 
-pub type StorageRef = Arc<InMemoryStorage>;
-pub type InMemoryTableRef = Arc<InMemoryTable>;
+pub type StorageRef = Arc<DiskStorage>;
+pub type DiskTableRef = Arc<DiskTable>;
 
 /// In-memory storage.
-pub struct InMemoryStorage {
-    tables: Mutex<HashMap<TableRefId, InMemoryTableRef>>,
+pub struct DiskStorage {
+    tables: Mutex<HashMap<TableRefId, DiskTableRef>>,
 }
 
-impl Default for InMemoryStorage {
+impl Default for DiskStorage {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl InMemoryStorage {
+impl DiskStorage {
     /// Create a new in-memory storage.
     pub fn new() -> Self {
-        InMemoryStorage {
+        DiskStorage {
             tables: Mutex::new(HashMap::new()),
         }
     }
 
     /// Add a table.
     pub fn add_table(&self, id: TableRefId) -> StorageResult<()> {
-        let table = Arc::new(InMemoryTable::new(id));
+        let table = Arc::new(DiskTable::new(id));
         self.tables.lock().unwrap().insert(id, table);
         Ok(())
     }
 
     /// Get a table.
-    pub fn get_table(&self, id: TableRefId) -> StorageResult<InMemoryTableRef> {
+    pub fn get_table(&self, id: TableRefId) -> StorageResult<DiskTableRef> {
         self.tables
             .lock()
             .unwrap()
@@ -61,22 +61,22 @@ impl InMemoryStorage {
 }
 
 /// A table in in-memory engine.
-pub struct InMemoryTable {
+pub struct DiskTable {
     #[allow(dead_code)]
     id: TableRefId,
-    inner: RwLock<InMemoryTableInner>,
+    inner: RwLock<DiskTableInner>,
 }
 
 #[derive(Default)]
-struct InMemoryTableInner {
+struct DiskTableInner {
     chunks: Vec<DataChunk>,
 }
 
-impl InMemoryTable {
+impl DiskTable {
     fn new(id: TableRefId) -> Self {
         Self {
             id,
-            inner: RwLock::new(InMemoryTableInner::default()),
+            inner: RwLock::new(DiskTableInner::default()),
         }
     }
 

code/03-01/src/storage/mod.rs

@@ -1,12 +1,16 @@
-//! In-memory storage.
+//! Persistent storage on disk.
 //!
 //! RisingLight's in-memory representation of data is very simple. Currently,
 //! it is simple a vector of `DataChunk`. Upon insertion, users' data are
 //! simply appended to the end of the vector.
 
+mod rowset;
+mod table_transaction;
+
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex, RwLock};
 
+use self::table_transaction::TableTransaction;
 use crate::array::DataChunk;
 use crate::catalog::TableRefId;
 
@@ -20,37 +24,37 @@ pub enum StorageError {
 /// A specialized `Result` type for storage operations.
 pub type StorageResult<T> = std::result::Result<T, StorageError>;
 
-pub type StorageRef = Arc<InMemoryStorage>;
-pub type InMemoryTableRef = Arc<InMemoryTable>;
+pub type StorageRef = Arc<DiskStorage>;
+pub type DiskTableRef = Arc<DiskTable>;
 
-/// In-memory storage.
-pub struct InMemoryStorage {
-    tables: Mutex<HashMap<TableRefId, InMemoryTableRef>>,
+/// Persistent storage on disk.
+pub struct DiskStorage {
+    tables: Mutex<HashMap<TableRefId, DiskTableRef>>,
 }
 
-impl Default for InMemoryStorage {
+impl Default for DiskStorage {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl InMemoryStorage {
-    /// Create a new in-memory storage.
+impl DiskStorage {
+    /// Create a new persistent storage on disk.
     pub fn new() -> Self {
-        InMemoryStorage {
+        DiskStorage {
             tables: Mutex::new(HashMap::new()),
         }
     }
 
     /// Add a table.
     pub fn add_table(&self, id: TableRefId) -> StorageResult<()> {
-        let table = Arc::new(InMemoryTable::new(id));
+        let table = Arc::new(DiskTable::new(id));
         self.tables.lock().unwrap().insert(id, table);
         Ok(())
     }
 
     /// Get a table.
-    pub fn get_table(&self, id: TableRefId) -> StorageResult<InMemoryTableRef> {
+    pub fn get_table(&self, id: TableRefId) -> StorageResult<DiskTableRef> {
         self.tables
             .lock()
             .unwrap()
@@ -61,33 +65,52 @@ impl InMemoryStorage {
 }
 
 /// A table in in-memory engine.
-pub struct InMemoryTable {
+pub struct DiskTable {
     #[allow(dead_code)]
     id: TableRefId,
-    inner: RwLock<InMemoryTableInner>,
+    inner: RwLock<DiskTableInner>,
 }
 
 #[derive(Default)]
-struct InMemoryTableInner {
+struct DiskTableInner {
     chunks: Vec<DataChunk>,
 }
 
-impl InMemoryTable {
+impl DiskTable {
     fn new(id: TableRefId) -> Self {
         Self {
             id,
-            inner: RwLock::new(InMemoryTableInner::default()),
+            inner: RwLock::new(DiskTableInner::default()),
         }
     }
 
+    #[allow(dead_code)]
+    async fn write(self: &Arc<Self>) -> StorageResult<TableTransaction> {
+        Ok(TableTransaction::start(self.clone(), false, false).await?)
+    }
+
+    #[allow(dead_code)]
+    async fn read(self: &Arc<Self>) -> StorageResult<TableTransaction> {
+        Ok(TableTransaction::start(self.clone(), true, false).await?)
+    }
+
+    #[allow(dead_code)]
+    async fn update(self: &Arc<Self>) -> StorageResult<TableTransaction> {
+        Ok(TableTransaction::start(self.clone(), false, true).await?)
+    }
+
     /// Append a chunk to the table.
+    ///
+    /// This interface will be deprecated soon in this tutorial.
     pub fn append(&self, chunk: DataChunk) -> StorageResult<()> {
         let mut inner = self.inner.write().unwrap();
         inner.chunks.push(chunk);
         Ok(())
     }
 
     /// Get all chunks of the table.
+    ///
+    /// This interface will be deprecated soon in this tutorial.
     pub fn all_chunks(&self) -> StorageResult<Vec<DataChunk>> {
         let inner = self.inner.read().unwrap();
         Ok(inner.chunks.clone())

code/03-01/src/storage/rowset/mem_rowset.rs

@@ -0,0 +1,39 @@
+#![allow(dead_code)]
+
+use std::sync::Arc;
+
+use itertools::Itertools;
+
+use crate::array::{ArrayBuilderImpl, DataChunk};
+use crate::catalog::ColumnCatalog;
+use crate::storage::StorageResult;
+
+pub struct MemRowset {
+    builders: Vec<ArrayBuilderImpl>,
+}
+
+impl MemRowset {
+    pub fn new(columns: Arc<[ColumnCatalog]>) -> Self {
+        Self {
+            builders: columns
+                .iter()
+                .map(|column| ArrayBuilderImpl::with_capacity(0, column.desc().datatype()))
+                .collect_vec(),
+        }
+    }
+
+    fn append(&mut self, columns: DataChunk) -> StorageResult<()> {
+        for (idx, column) in columns.arrays().iter().enumerate() {
+            self.builders[idx].append(column);
+        }
+        Ok(())
+    }
+
+    fn flush(self) -> StorageResult<DataChunk> {
+        Ok(self
+            .builders
+            .into_iter()
+            .map(|builder| builder.finish())
+            .collect::<DataChunk>())
+    }
+}

code/03-01/src/storage/rowset/mod.rs

@@ -0,0 +1,3 @@
+mod mem_rowset;
+
+pub use mem_rowset::*;

code/03-01/src/storage/table_transaction.rs

@@ -0,0 +1,56 @@
+#![allow(dead_code)]
+
+use super::rowset::MemRowset;
+use super::{DiskTableRef, StorageResult};
+use crate::array::DataChunk;
+
+/// [`TableTransaction`] records the state of a single table. All operations (insert, update,
+/// delete) should go through [`TableTransaction`].
+pub struct TableTransaction {
+    mem_rowset: Option<MemRowset>,
+    read_only: bool,
+    update: bool,
+    table: DiskTableRef,
+}
+
+impl TableTransaction {
+    /// Start a [`WriteBatch`]
+    pub async fn start(table: DiskTableRef, read_only: bool, update: bool) -> StorageResult<Self> {
+        Ok(Self {
+            mem_rowset: None,
+            table,
+            update,
+            read_only,
+        })
+    }
+
+    /// Flush [`WriteBatch`] to some on-disk RowSets.
+    pub async fn flush(self) {
+        todo!()
+    }
+
+    /// Add a [`DataChunk`] to the mem rowset
+    pub fn append(&self, _chunk: DataChunk) -> StorageResult<()> {
+        todo!()
+    }
+
+    /// Delete a row from the table.
+    async fn delete(&mut self, _row_id: u64) -> StorageResult<()> {
+        todo!()
+    }
+
+    /// Commit all changes in this transaction.
+    pub fn commit(self) -> StorageResult<()> {
+        todo!()
+    }
+
+    /// Abort all changes in this transaction.
+    pub fn abort(self) -> StorageResult<()> {
+        todo!()
+    }
+
+    /// Create an iterator on this table.
+    pub async fn scan(&self) {
+        todo!()
+    }
+}