storage part 1

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

Author: skyzh

code/.gitignore

@@ -1,2 +1,3 @@
 target
 Cargo.lock
+risinglight.db/

code/03-00/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 [dependencies]
 anyhow = "1"
 bitvec = "1.0"
+bytes = "1"
 enum_dispatch = "0.3"
 env_logger = "0.9"
 futures = { version = "0.3", default-features = false, features = ["alloc"] }
@@ -17,7 +18,7 @@ prettytable-rs = { version = "0.8", default-features = false }
 rustyline = "9"
 sqlparser = "0.13"
 thiserror = "1"
-tokio = { version = "1", features = ["rt", "rt-multi-thread", "sync", "macros"] }
+tokio = { version = "1", features = ["rt", "rt-multi-thread", "sync", "macros", "fs"] }
 tokio-stream = "0.1"
 
 [dev-dependencies]

code/03-00/src/executor/create.rs

@@ -15,11 +15,15 @@ impl CreateTableExecutor {
         let schema = self.catalog.get_schema(self.plan.schema_id).unwrap();
         let table_id = schema.add_table(&self.plan.table_name).unwrap();
         let table = schema.get_table(table_id).unwrap();
+        let mut column_descs = vec![];
         for (name, desc) in &self.plan.columns {
             table.add_column(name, desc.clone()).unwrap();
+            column_descs.push(desc.clone());
         }
-        self.storage
-            .add_table(TableRefId::new(self.plan.schema_id, table_id))?;
+        self.storage.add_table(
+            TableRefId::new(self.plan.schema_id, table_id),
+            &column_descs,
+        )?;
         yield DataChunk::single(1);
     }
 }

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

@@ -1,9 +1,14 @@
 //! On-disk storage
 
-use std::sync::Arc;
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::{Arc, RwLock};
 
-use crate::array::DataChunk;
-use crate::catalog::TableRefId;
+use anyhow::anyhow;
+use bytes::{Buf, BufMut};
+
+use crate::array::{Array, ArrayBuilder, ArrayImpl, DataChunk, I32Array, I32ArrayBuilder};
+use crate::catalog::{ColumnDesc, TableRefId};
 
 /// The error type of storage operations.
 #[derive(thiserror::Error, Debug)]
@@ -17,13 +22,33 @@ pub type StorageRef = Arc<DiskStorage>;
 pub type StorageTableRef = Arc<DiskTable>;
 
 /// On-disk storage.
-#[derive(Clone)]
-pub struct DiskStorage;
+pub struct DiskStorage {
+    /// All tables in the current storage engine.
+    tables: RwLock<HashMap<TableRefId, StorageTableRef>>,
+
+    /// The storage options.
+    options: Arc<StorageOptions>,
+}
+
+pub struct StorageOptions {
+    /// The directory of the storage
+    base_path: PathBuf,
+}
+
+pub fn err(error: impl Into<anyhow::Error>) -> StorageError {
+    StorageError(error.into())
+}
 
 /// An on-disk table.
 pub struct DiskTable {
-    #[allow(dead_code)]
+    /// Id of the table.
     id: TableRefId,
+
+    /// Columns of the current table.
+    column_descs: Arc<[ColumnDesc]>,
+
+    /// The storage options.
+    options: Arc<StorageOptions>,
 }
 
 impl Default for DiskStorage {
@@ -35,28 +60,94 @@ impl Default for DiskStorage {
 impl DiskStorage {
     /// Create a new in-memory storage.
     pub fn new() -> Self {
-        DiskStorage
+        DiskStorage {
+            tables: RwLock::new(HashMap::new()),
+            options: Arc::new(StorageOptions {
+                base_path: "risinglight.db".into(),
+            }),
+        }
     }
 
     /// Add a table.
-    pub fn add_table(&self, _id: TableRefId) -> StorageResult<()> {
-        todo!()
+    pub fn add_table(&self, id: TableRefId, column_descs: &[ColumnDesc]) -> StorageResult<()> {
+        let mut tables = self.tables.write().unwrap();
+        let table = DiskTable {
+            id,
+            options: self.options.clone(),
+            column_descs: column_descs.into(),
+        };
+        let res = tables.insert(id, table.into());
+        if res.is_some() {
+            return Err(anyhow!("table already exists: {:?}", id).into());
+        }
+        Ok(())
     }
 
     /// Get a table.
-    pub fn get_table(&self, _id: TableRefId) -> StorageResult<StorageTableRef> {
-        todo!()
+    pub fn get_table(&self, id: TableRefId) -> StorageResult<StorageTableRef> {
+        let tables = self.tables.read().unwrap();
+        tables
+            .get(&id)
+            .ok_or_else(|| anyhow!("table not found: {:?}", id).into())
+            .cloned()
     }
 }
 
+/// Encode an `I32Array` into a `Vec<u8>`.
+fn encode_int32_column(a: &I32Array) -> StorageResult<Vec<u8>> {
+    let mut buffer = Vec::with_capacity(a.len() * 4);
+    for item in a.iter() {
+        if let Some(item) = item {
+            buffer.put_i32_le(*item);
+        } else {
+            return Err(anyhow!("nullable encoding not supported!").into());
+        }
+    }
+    Ok(buffer)
+}
+
+fn decode_int32_column(mut data: &[u8]) -> StorageResult<I32Array> {
+    let mut builder = I32ArrayBuilder::with_capacity(data.len() / 4);
+    while data.has_remaining() {
+        builder.push(Some(&data.get_i32_le()));
+    }
+    Ok(builder.finish())
+}
+
 impl DiskTable {
+    fn table_path(&self) -> PathBuf {
+        self.options.base_path.join(self.id.table_id.to_string())
+    }
+
+    fn column_path(&self, column_id: usize) -> PathBuf {
+        self.table_path().join(format!("{}.col", column_id))
+    }
+
     /// Append a chunk to the table.
-    pub async fn append(&self, _chunk: DataChunk) -> StorageResult<()> {
-        todo!()
+    pub async fn append(&self, chunk: DataChunk) -> StorageResult<()> {
+        for (idx, column) in chunk.arrays().iter().enumerate() {
+            if let ArrayImpl::Int32(column) = column {
+                let column_path = self.column_path(idx);
+                let data = encode_int32_column(column)?;
+                tokio::fs::create_dir_all(column_path.parent().unwrap())
+                    .await
+                    .map_err(err)?;
+                tokio::fs::write(column_path, data).await.map_err(err)?;
+            } else {
+                return Err(anyhow!("unsupported column type").into());
+            }
+        }
+        Ok(())
     }
 
     /// Get all chunks of the table.
     pub async fn all_chunks(&self) -> StorageResult<Vec<DataChunk>> {
-        todo!()
+        let mut columns = vec![];
+        for (idx, _) in self.column_descs.iter().enumerate() {
+            let column_path = self.column_path(idx);
+            let data = tokio::fs::read(column_path).await.map_err(err)?;
+            columns.push(decode_int32_column(&data)?);
+        }
+        Ok(vec![columns.into_iter().map(ArrayImpl::Int32).collect()])
     }
 }

code/03-00/src/test.rs

@@ -1,6 +1,3 @@
-#![allow(unused_imports)]
-#![allow(dead_code)]
-
 use std::path::Path;
 
 use test_case::test_case;
@@ -9,6 +6,7 @@ use crate::array::DataChunk;
 use crate::types::DataValue;
 use crate::{Database, Error};
 
+#[test_case("03-01.slt")]
 fn test(name: &str) {
     init_logger();
     let script = std::fs::read_to_string(Path::new("../sql").join(name)).unwrap();

code/sql/03-01.slt

@@ -0,0 +1,14 @@
+# 03-01: Very simple storage test
+
+statement ok
+CREATE TABLE t (a INT NOT NULL, b INT NOT NULL, c INT NOT NULL)
+
+statement ok
+INSERT INTO t VALUES (1,10,100), (2,20,200), (3,30,300)
+
+query III
+SELECT * FROM t
+----
+1 10 100
+2 20 200
+3 30 300