feat: supports multiple primary key types

KKould · KKould · commit de224d9bce63 · 2023-09-14T23:08:05.000+08:00
- Tinyint
- UTinyint
- Smallint
- USmallint
- Integer
- UInteger
- Bigint
- UBigint
- Varchar
diff --git a/README.md b/README.md
@@ -43,6 +43,16 @@ Storage Support:
 ![demo](./static/images/demo.png)
 
 ### Features
+- Supports multiple primary key types
+  - Tinyint
+  - UTinyint
+  - Smallint
+  - USmallint
+  - Integer
+  - UInteger
+  - Bigint
+  - UBigint
+  - Varchar
 - DDL
   - Create
     - [x] Table
@@ -82,6 +92,7 @@ Storage Support:
   - Float
   - Double
   - Varchar
+  - Date
   - DateTime
 - Optimizer rules
   - Limit Project Transpose
diff --git a/src/binder/create_table.rs b/src/binder/create_table.rs
@@ -10,7 +10,6 @@ use crate::planner::LogicalPlan;
 use crate::planner::operator::create_table::CreateTableOperator;
 use crate::planner::operator::Operator;
 use crate::storage::Storage;
-use crate::types::LogicalType;
 
 impl<S: Storage> Binder<S> {
     pub(crate) fn bind_create_table(
@@ -30,17 +29,12 @@ impl<S: Storage> Binder<S> {
                 return Err(BindError::AmbiguousColumn(col_name.to_string()));
             }
         }
-
         let columns = columns
             .iter()
             .map(|col| ColumnCatalog::from(col.clone()))
             .collect_vec();
 
-        if let Some(col) = columns.iter().find(|col| col.desc.is_primary) {
-            if !matches!(col.datatype(), LogicalType::Integer) {
-                return Err(BindError::InvalidColumn("Primary key types only support signed integers".to_string()));
-            }
-        } else {
+        if columns.iter().find(|col| col.desc.is_primary).is_none() {
             return Err(BindError::InvalidTable("At least one primary key field exists".to_string()));
         }
 
diff --git a/src/execution/executor/dml/insert.rs b/src/execution/executor/dml/insert.rs
@@ -7,7 +7,7 @@ use crate::execution::executor::{BoxedExecutor, Executor};
 use crate::execution::ExecutorError;
 use crate::planner::operator::insert::InsertOperator;
 use crate::storage::{Storage, Table};
-use crate::types::{ColumnId, LogicalType};
+use crate::types::ColumnId;
 use crate::types::tuple::Tuple;
 use crate::types::value::{DataValue, ValueRef};
 
@@ -49,25 +49,16 @@ impl Insert {
                         .map(|(i, _)| i)
                         .unwrap()
                 });
-
-                let tuple_id = if let DataValue::Int64(Some(primary_id)) =
-                    DataValue::clone(&values[*primary_idx]).cast(&LogicalType::Bigint)?
-                {
-                    primary_id
-                } else {
-                    unreachable!("Primary key must have a value")
-                };
-
+                let id = Some(values[*primary_idx].clone());
                 let mut tuple_map: HashMap<ColumnId, ValueRef> = values
                     .into_iter()
                     .enumerate()
                     .map(|(i, value)| (columns[i].id, value))
                     .collect();
-
                 let all_columns = table_catalog.all_columns_with_id();
 
                 let mut tuple = Tuple {
-                    id: Some(tuple_id),
+                    id,
                     columns: Vec::with_capacity(all_columns.len()),
                     values: Vec::with_capacity(all_columns.len()),
                 };
diff --git a/src/execution/executor/dml/update.rs b/src/execution/executor/dml/update.rs
@@ -45,13 +45,19 @@ impl Update {
                     value_map.insert(columns[i].id, values[i].clone());
                 }
             }
-
             #[for_await]
             for tuple in input {
                 let mut tuple = tuple?;
 
                 for (i, column) in tuple.columns.iter().enumerate() {
                     if let Some(value) = value_map.get(&column.id) {
+                        if column.desc.is_primary {
+                            if let Some(old_key) = tuple.id.replace(value.clone()) {
+                                if value != &old_key {
+                                    table.delete(old_key)?;
+                                }
+                            }
+                        }
                         tuple.values[i] = value.clone();
                     }
                 }
diff --git a/src/execution/executor/dql/filter.rs b/src/execution/executor/dql/filter.rs
@@ -13,7 +13,7 @@ pub struct Filter {
 }
 
 impl From<(FilterOperator, BoxedExecutor)> for Filter {
-    fn from((FilterOperator { predicate, having }, input): (FilterOperator, BoxedExecutor)) -> Self {
+    fn from((FilterOperator { predicate, .. }, input): (FilterOperator, BoxedExecutor)) -> Self {
         Filter {
             predicate,
             input
diff --git a/src/execution/executor/dql/seq_scan.rs b/src/execution/executor/dql/seq_scan.rs
@@ -26,7 +26,6 @@ impl<S: Storage> Executor<S> for SeqScan {
 impl SeqScan {
     #[try_stream(boxed, ok = Tuple, error = ExecutorError)]
     pub async fn _execute<S: Storage>(self, storage: S) {
-        // TODO: sort_fields, pre_where, limit
         let ScanOperator { table_name,  columns, limit, .. } = self.op;
 
         if let Some(table) = storage.table(&table_name).await {
diff --git a/src/planner/operator/scan.rs b/src/planner/operator/scan.rs
@@ -10,14 +10,15 @@ use super::{sort::SortField, Operator};
 pub struct ScanOperator {
     pub table_name: TableName,
     pub columns: Vec<ScalarExpression>,
-    // TODO:
+    // Support push down limit.
+    pub limit: Bounds,
+
+    // IndexScan only
     pub sort_fields: Vec<SortField>,
+    // IndexScan only
     // Support push down predicate.
     // If pre_where is simple predicate, for example:  a > 1 then can calculate directly when read data.
-    // TODO:
     pub pre_where: Vec<ScalarExpression>,
-    // Support push down limit.
-    pub limit: Bounds,
 }
 impl ScanOperator {
     pub fn new(table_name: TableName, table_catalog: &TableCatalog) -> LogicalPlan {
diff --git a/src/storage/kip.rs b/src/storage/kip.rs
@@ -164,14 +164,14 @@ impl Table for KipTable {
     }
 
     fn append(&mut self, tuple: Tuple) -> Result<(), StorageError> {
-        let (key, value) = self.table_codec.encode_tuple(&tuple);
+        let (key, value) = self.table_codec.encode_tuple(&tuple)?;
         self.tx.set(key, value);
 
         Ok(())
     }
 
     fn delete(&mut self, tuple_id: TupleId) -> Result<(), StorageError> {
-        let key = self.table_codec.encode_tuple_key(&tuple_id);
+        let key = self.table_codec.encode_tuple_key(&tuple_id)?;
         self.tx.remove(&key)?;
 
         Ok(())
@@ -205,33 +205,31 @@ impl Transaction for KipTraction<'_> {
             }
         }
 
-        self.iter
-            .try_next()?
-            .and_then(|(key, bytes)| {
-                bytes.and_then(|value| {
-                    self.table_codec.decode_tuple(&key, &value)
-                        .map(|tuple| {
-                            let projection_len = self.projections.len();
-
-                            let mut columns = Vec::with_capacity(projection_len);
-                            let mut values = Vec::with_capacity(projection_len);
-
-                            for expr in self.projections.iter() {
-                                values.push(expr.eval_column(&tuple)?);
-                                columns.push(expr.output_columns(&tuple));
-                            }
-
-                            self.limit = self.limit.map(|num| num - 1);
-
-                            Ok(Tuple {
-                                id: tuple.id,
-                                columns,
-                                values,
-                            })
-                        })
-                })
-            })
-            .transpose()
+        while let Some(item) = self.iter.try_next()? {
+            if let (_, Some(value)) = item {
+                let tuple = self.table_codec.decode_tuple(&value);
+
+                let projection_len = self.projections.len();
+
+                let mut columns = Vec::with_capacity(projection_len);
+                let mut values = Vec::with_capacity(projection_len);
+
+                for expr in self.projections.iter() {
+                    values.push(expr.eval_column(&tuple)?);
+                    columns.push(expr.output_columns(&tuple));
+                }
+
+                self.limit = self.limit.map(|num| num - 1);
+
+                return Ok(Some(Tuple {
+                    id: tuple.id,
+                    columns,
+                    values,
+                }))
+            }
+        }
+
+        Ok(None)
     }
 }
 
@@ -246,6 +244,7 @@ mod test {
     use crate::storage::{Storage, StorageError, Transaction, Table};
     use crate::storage::memory::test::data_filling;
     use crate::types::LogicalType;
+    use crate::types::value::DataValue;
 
     #[tokio::test]
     async fn test_in_kipdb_storage_works_with_data() -> Result<(), StorageError> {
@@ -282,7 +281,7 @@ mod test {
         )?;
 
         let option_1 = tx.next_tuple()?;
-        assert_eq!(option_1.unwrap().id, Some(1));
+        assert_eq!(option_1.unwrap().id, Some(Arc::new(DataValue::Int32(Some(2)))));
 
         let option_2 = tx.next_tuple()?;
         assert_eq!(option_2, None);
diff --git a/src/storage/memory.rs b/src/storage/memory.rs
@@ -179,7 +179,7 @@ impl Table for MemTable {
                 .as_mut()
         }.unwrap();
 
-        tuples.retain(|tuple| tuple.id.unwrap() != tuple_id);
+        tuples.retain(|tuple| tuple.id.clone().unwrap() != tuple_id);
 
         Ok(())
     }
@@ -249,15 +249,15 @@ pub(crate) mod test {
 
     pub fn data_filling(columns: Vec<ColumnRef>, table: &mut impl Table) -> Result<(), StorageError> {
         table.append(Tuple {
-            id: Some(0),
+            id: Some(Arc::new(DataValue::Int32(Some(1)))),
             columns: columns.clone(),
             values: vec![
                 Arc::new(DataValue::Int32(Some(1))),
                 Arc::new(DataValue::Boolean(Some(true)))
             ],
         })?;
         table.append(Tuple {
-            id: Some(1),
+            id: Some(Arc::new(DataValue::Int32(Some(2)))),
             columns: columns.clone(),
             values: vec![
                 Arc::new(DataValue::Int32(Some(2))),
@@ -303,7 +303,7 @@ pub(crate) mod test {
         )?;
 
         let option_1 = tx.next_tuple()?;
-        assert_eq!(option_1.unwrap().id, Some(1));
+        assert_eq!(option_1.unwrap().id, Some(Arc::new(DataValue::Int32(Some(2)))));
 
         let option_2 = tx.next_tuple()?;
         assert_eq!(option_2, None);
diff --git a/src/storage/table_codec.rs b/src/storage/table_codec.rs
diff --git a/src/types/errors.rs b/src/types/errors.rs
diff --git a/src/types/tuple.rs b/src/types/tuple.rs
diff --git a/src/types/value.rs b/src/types/value.rs

Original file line number	Diff line number	Diff line change
`@@ -45,13 +45,19 @@ impl Update {`
`45`	`45`	`value_map.insert(columns[i].id, values[i].clone());`
`46`	`46`	`}`
`47`	`47`	`}`
`48`		`-`
`49`	`48`	`#[for_await]`
`50`	`49`	`for tuple in input {`
`51`	`50`	`let mut tuple = tuple?;`
`52`	`51`
`53`	`52`	`for (i, column) in tuple.columns.iter().enumerate() {`
`54`	`53`	`if let Some(value) = value_map.get(&column.id) {`
	`54`	`+ if column.desc.is_primary {`
	`55`	`+ if let Some(old_key) = tuple.id.replace(value.clone()) {`
	`56`	`+ if value != &old_key {`
	`57`	`+ table.delete(old_key)?;`
	`58`	`+ }`
	`59`	`+ }`
	`60`	`+ }`
`55`	`61`	`tuple.values[i] = value.clone();`
`56`	`62`	`}`
`57`	`63`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ pub struct Filter {`
`13`	`13`	`}`
`14`	`14`
`15`	`15`	`impl From<(FilterOperator, BoxedExecutor)> for Filter {`
`16`		`- fn from((FilterOperator { predicate, having }, input): (FilterOperator, BoxedExecutor)) -> Self {`
	`16`	`+ fn from((FilterOperator { predicate, .. }, input): (FilterOperator, BoxedExecutor)) -> Self {`
`17`	`17`	`Filter {`
`18`	`18`	`predicate,`
`19`	`19`	`input`