feat(catalog): implement multi-field select and create table plan

Author: KKould

src/binder/create.rs

@@ -31,25 +31,19 @@ impl Binder {
 
         let columns: Vec<ColumnCatalog> = columns
             .iter()
-            .enumerate()
-            .map(|(_, col)| ColumnCatalog::from(col.clone()))
+            .map(|col| ColumnCatalog::from(col.clone()))
             .collect();
 
         let plan = LogicalCreateTablePlan {
             table_name: table_name.to_string(),
-            columns: columns
-                .into_iter()
-                .map(|col| (col.name.to_string(), col.nullable, col.desc.clone()))
-                .collect(),
+            columns,
         };
         Ok(plan)
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use sqlparser::ast::CharacterLength;
-
     use super::*;
     use crate::binder::BinderContext;
     use crate::catalog::{ColumnDesc, RootCatalog};
@@ -66,16 +60,16 @@ mod tests {
         let plan2 = LogicalPlan::CreateTable(LogicalCreateTablePlan {
             table_name: "t1".to_string(),
             columns: vec![
-                (
+                ColumnCatalog::new(
                     "id".to_string(),
                     false,
-                    ColumnDesc::new(LogicalType::Integer, false),
+                    ColumnDesc::new(LogicalType::Integer, false)
                 ),
-                (
+                ColumnCatalog::new(
                     "name".to_string(),
                     false,
-                    ColumnDesc::new(LogicalType::Varchar, false),
-                ),
+                    ColumnDesc::new(LogicalType::Varchar, false)
+                )
             ],
         });
 

src/binder/expr.rs

@@ -57,7 +57,7 @@ impl Binder {
         } else {
             // handle col syntax
             let mut got_column = None;
-            for table_catalog in self.context.catalog.tables.values() {
+            for table_catalog in &self.context.catalog.tables {
                 if let Some(column_catalog) = table_catalog.get_column_by_name(column_name) {
                     if got_column.is_some() {
                         return Err(BindError::InvalidColumn(column_name.to_string()).into());

src/binder/mod.rs

@@ -11,11 +11,11 @@ use sqlparser::ast::{Ident, ObjectName, Statement};
 use crate::catalog::{RootCatalog, DEFAULT_SCHEMA_NAME, CatalogError};
 use crate::expression::ScalarExpression;
 use crate::planner::LogicalPlan;
-use crate::types::TableId;
+use crate::types::TableIdx;
 #[derive(Clone)]
 pub struct BinderContext {
     catalog: RootCatalog,
-    bind_table: HashMap<String, TableId>,
+    bind_table: HashMap<String, TableIdx>,
     aliases: HashMap<String, ScalarExpression>,
     group_by_exprs: Vec<ScalarExpression>,
     agg_calls: Vec<ScalarExpression>,

src/catalog/column.rs

@@ -1,11 +1,11 @@
 use arrow::datatypes::{DataType, Field};
 use sqlparser::ast::{ColumnDef, ColumnOption};
 
-use crate::types::{ColumnId, IdGenerator, LogicalType};
+use crate::types::{ColumnIdx, LogicalType};
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct ColumnCatalog {
-    pub id: ColumnId,
+    pub id: Option<ColumnIdx>,
     pub name: String,
     pub nullable: bool,
     pub desc: ColumnDesc,
@@ -14,7 +14,7 @@ pub struct ColumnCatalog {
 impl ColumnCatalog {
     pub(crate) fn new(column_name: String, nullable: bool, column_desc: ColumnDesc) -> ColumnCatalog {
         ColumnCatalog {
-            id: IdGenerator::build(),
+            id: None,
             name: column_name,
             nullable,
             desc: column_desc,
@@ -25,10 +25,6 @@ impl ColumnCatalog {
         &self.desc.column_datatype
     }
 
-    pub(crate) fn id(&self) -> ColumnId {
-        self.id
-    }
-
     pub fn desc(&self) -> &ColumnDesc {
         &self.desc
     }

src/catalog/mod.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 pub(crate) use self::column::*;
 pub(crate) use self::root::*;
 pub(crate) use self::table::*;
-use crate::types::{ColumnId, TableId};
+use crate::types::{ColumnIdx, TableIdx};
 
 /// The type of catalog reference.
 pub type CatalogRef = Arc<RootCatalog>;
@@ -19,12 +19,12 @@ mod table;
 /// The reference ID of a column.
 #[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
 pub struct ColumnRefId {
-    pub table_id: TableId,
-    pub column_id: ColumnId,
+    pub table_id: TableIdx,
+    pub column_id: ColumnIdx,
 }
 
 impl ColumnRefId {
-    pub const fn from_table(table_id: TableId, column_id: ColumnId) -> Self {
+    pub const fn from_table(table_id: TableIdx, column_id: ColumnIdx) -> Self {
         ColumnRefId {
             table_id,
             column_id,
@@ -34,11 +34,11 @@ impl ColumnRefId {
 
 #[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
 pub struct TableRefId {
-    pub table_id: TableId,
+    pub table_id: TableIdx,
 }
 
 impl TableRefId {
-    pub const fn new(table_id: TableId) -> Self {
+    pub const fn new(table_id: TableIdx) -> Self {
         TableRefId { table_id }
     }
 }

src/catalog/root.rs

@@ -1,12 +1,13 @@
 use std::collections::BTreeMap;
 
 use crate::catalog::{CatalogError, ColumnCatalog, TableCatalog};
-use crate::types::TableId;
+use crate::types::{IdGenerator, TableIdx};
 
 #[derive(Debug, Clone)]
 pub struct RootCatalog {
-    pub table_idxs: BTreeMap<String, TableId>,
-    pub tables: BTreeMap<TableId, TableCatalog>,
+    generator: IdGenerator,
+    pub table_idxs: BTreeMap<String, TableIdx>,
+    pub tables: Vec<TableCatalog>,
 }
 
 impl Default for RootCatalog {
@@ -19,35 +20,37 @@ impl RootCatalog {
     #[allow(dead_code)]
     pub fn new() -> Self {
         RootCatalog {
+            generator: IdGenerator::new(),
             table_idxs: Default::default(),
             tables: Default::default(),
         }
     }
 
-    pub(crate) fn get_table_id_by_name(&self, name: &str) -> Option<TableId> {
+    pub(crate) fn get_table_id_by_name(&self, name: &str) -> Option<TableIdx> {
         self.table_idxs.get(name).cloned()
     }
 
-    pub(crate) fn get_table(&self, table_id: TableId) -> Option<&TableCatalog> {
-        self.tables.get(&table_id)
+    pub(crate) fn get_table(&self, table_id: TableIdx) -> Option<&TableCatalog> {
+        self.tables.get(table_id)
     }
 
     pub(crate) fn get_table_by_name(&self, name: &str) -> Option<&TableCatalog> {
         let id = self.table_idxs.get(name)?;
-        self.tables.get(id)
+        self.tables.get(*id)
     }
 
     pub(crate) fn add_table(
         &mut self,
         table_name: String,
         columns: Vec<ColumnCatalog>,
-    ) -> Result<TableId, CatalogError> {
+    ) -> Result<TableIdx, CatalogError> {
         if self.table_idxs.contains_key(&table_name) {
             return Err(CatalogError::Duplicated("column", table_name));
         }
-        let table = TableCatalog::new(table_name.to_owned(), columns)?;
-        let table_id = table.id;
+        let mut table = TableCatalog::new(table_name.to_owned(), columns)?;
+        let table_id = self.generator.build();
 
+        table.id = Some(table_id);
         self.table_idxs.insert(table_name, table_id);
         self.tables.insert(table_id, table);
 

src/catalog/table.rs

@@ -1,54 +1,56 @@
-use std::collections::{BTreeMap, HashMap};
+use std::collections::HashMap;
 
 use itertools::Itertools;
 
 use crate::catalog::{CatalogError, ColumnCatalog};
-use crate::types::{ColumnId, IdGenerator, TableId};
+use crate::types::{ColumnIdx, IdGenerator, TableIdx};
 #[derive(Debug, Clone)]
 pub struct TableCatalog {
-    pub id: TableId,
+    pub id: Option<TableIdx>,
     pub name: String,
+    generator: IdGenerator,
     /// Mapping from column names to column ids
-    column_idxs: HashMap<String, ColumnId>,
-    pub(crate) columns: BTreeMap<ColumnId, ColumnCatalog>,
+    column_idxs: HashMap<String, ColumnIdx>,
+    pub(crate) columns: Vec<ColumnCatalog>,
 }
 
 impl TableCatalog {
-    pub(crate) fn get_column_by_id(&self, id: ColumnId) -> Option<&ColumnCatalog> {
-        self.columns.get(&id)
+    pub(crate) fn get_column_by_id(&self, id: ColumnIdx) -> Option<&ColumnCatalog> {
+        self.columns.get(id)
     }
 
-    pub(crate) fn get_column_id_by_name(&self, name: &str) -> Option<ColumnId> {
+    pub(crate) fn get_column_id_by_name(&self, name: &str) -> Option<ColumnIdx> {
         self.column_idxs.get(name).cloned()
     }
 
     pub(crate) fn get_column_by_name(&self, name: &str) -> Option<&ColumnCatalog> {
         let id = self.column_idxs.get(name)?;
-        self.columns.get(id)
+        self.columns.get(*id)
     }
 
     pub(crate) fn contains_column(&self, name: &str) -> bool {
         self.column_idxs.contains_key(name)
     }
 
-    pub(crate) fn get_all_columns(&self) -> Vec<(ColumnId, &ColumnCatalog)> {
+    pub(crate) fn get_all_columns(&self) -> Vec<(ColumnIdx, &ColumnCatalog)> {
         self.columns
             .iter()
-            .map(|(col_id, col)| (*col_id, col))
+            .enumerate()
             .collect_vec()
     }
 
     /// Add a column to the table catalog.
     pub(crate) fn add_column(
         &mut self,
-        col_catalog: ColumnCatalog,
-    ) -> Result<ColumnId, CatalogError> {
+        mut col_catalog: ColumnCatalog,
+    ) -> Result<ColumnIdx, CatalogError> {
         if self.column_idxs.contains_key(&col_catalog.name) {
             return Err(CatalogError::Duplicated("column", col_catalog.name.into()));
         }
 
-        let col_id = col_catalog.id;
+        let col_id = self.generator.build();
 
+        col_catalog.id = Some(col_id);
         self.column_idxs.insert(col_catalog.name.to_owned(), col_id);
         self.columns.insert(col_id, col_catalog);
 
@@ -60,10 +62,11 @@ impl TableCatalog {
         columns: Vec<ColumnCatalog>,
     ) -> Result<TableCatalog, CatalogError> {
         let mut table_catalog = TableCatalog {
-            id: IdGenerator::build(),
+            id: None,
             name: table_name,
+            generator: IdGenerator::new(),
             column_idxs: HashMap::new(),
-            columns: BTreeMap::new(),
+            columns: Vec::new(),
         };
 
         for col_catalog in columns.into_iter() {

src/db.rs

@@ -1,20 +1,14 @@
-use std::sync::Arc;
-
 use anyhow::Result;
-use arrow::datatypes::Schema;
 use arrow::error::ArrowError;
 use arrow::record_batch::RecordBatch;
 use sqlparser::parser::ParserError;
 
 use crate::binder::{BindError, Binder, BinderContext};
-use crate::catalog::ColumnCatalog;
 use crate::execution_v1::physical_plan::physical_plan_builder::PhysicalPlanBuilder;
 use crate::execution_v1::volcano_executor::VolcanoExecutor;
 use crate::parser::parse_sql;
-use crate::planner::LogicalPlan;
 use crate::storage::memory::InMemoryStorage;
 use crate::storage::{Storage, StorageError, StorageImpl};
-use crate::types::IdGenerator;
 
 #[derive(Debug)]
 pub struct Database {
@@ -51,10 +45,11 @@ impl Database {
         ///     Limit(1)
         ///       Project(a,b)
         let logical_plan = binder.bind(&stmts[0])?;
-        println!("logic plan {:#?}", logical_plan);
+        println!("logic plan: {:#?}", logical_plan);
 
         let mut builder = PhysicalPlanBuilder::new();
         let operator = builder.build_plan(&logical_plan)?;
+        println!("operator: {:#?}", logical_plan);
 
         let storage = StorageImpl::InMemoryStorage(self.storage.clone());
         let executor = VolcanoExecutor::new(storage);
@@ -127,7 +122,7 @@ pub enum DatabaseError {
 #[cfg(test)]
 mod test {
     use std::sync::Arc;
-    use arrow::array::Int32Array;
+    use arrow::array::{BooleanArray, Int32Array};
     use arrow::datatypes::Schema;
     use arrow::record_batch::RecordBatch;
     use itertools::Itertools;
@@ -136,22 +131,32 @@ mod test {
     use crate::execution_v1::ExecutorError;
     use crate::storage::{Storage, StorageError};
     use crate::storage::memory::InMemoryStorage;
-    use crate::types::{IdGenerator, LogicalType, TableId};
-
-    fn build_table(storage: &impl Storage) -> Result<TableId, StorageError> {
-        let fields = vec![
-            ColumnCatalog::new(
-                "c1".to_string(),
-                false,
-                ColumnDesc::new(LogicalType::Integer, true)
-            ).to_field(),
-        ];
-        let batch = RecordBatch::try_new(
-            Arc::new(Schema::new(fields)),
-            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))]
+    use crate::types::{IdGenerator, LogicalType, TableIdx};
+
+    fn build_table(storage: &impl Storage) -> Result<TableIdx, StorageError> {
+        let schema = Arc::new(Schema::new(
+            vec![
+                ColumnCatalog::new(
+                    "c1".to_string(),
+                    false,
+                    ColumnDesc::new(LogicalType::Integer, true)
+                ).to_field(),
+                ColumnCatalog::new(
+                    "c2".to_string(),
+                    false,
+                    ColumnDesc::new(LogicalType::Boolean, false)
+                ).to_field(),
+            ]
+        ));
+        let batch_1 = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
+                Arc::new(BooleanArray::from(vec![true, true, false, true, false]))
+            ]
         ).unwrap();
 
-        Ok(storage.create_table("t1", vec![batch])?)
+        Ok(storage.create_table("t1", vec![batch_1])?)
     }
 
     #[test]
@@ -172,8 +177,8 @@ mod test {
         let database = Database::new_on_mem();
 
         tokio_test::block_on(async move {
-            let _batch = database.run("create table t1 (a int)").await?;
-            let batch = database.run("select a from t1").await?;
+            let _batch = database.run("create table t1 (a int, b int)").await?;
+            let batch = database.run("select * from t1").await?;
             println!("{:#?}", batch);
 
             Ok(())