Skip to content

Commit fe88322

Browse files
KKouldKKould
committed
feat: impl multiple primary keys
1 parent 4683978 commit fe88322

39 files changed

+569
-214
lines changed

‎src/binder/create_table.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,9 @@ impl<T: Transaction> Binder<'_, '_, T> {
6262
.find(|column| column.name() == column_name)
6363
{
6464
if *is_primary {
65-
column.desc_mut().is_primary = true;
65+
column.desc_mut().set_primary(true);
6666
} else {
67-
column.desc_mut().is_unique = true;
67+
column.desc_mut().set_unique(true);
6868
}
6969
}
7070
}
@@ -73,9 +73,9 @@ impl<T: Transaction> Binder<'_, '_, T> {
7373
}
7474
}
7575

76-
if columns.iter().filter(|col| col.desc().is_primary).count() != 1 {
76+
if columns.iter().filter(|col| col.desc().is_primary()).count() == 0 {
7777
return Err(DatabaseError::InvalidTable(
78-
"The primary key field must exist and have at least one".to_string(),
78+
"the primary key field must exist and have at least one".to_string(),
7979
));
8080
}
8181

@@ -106,12 +106,12 @@ impl<T: Transaction> Binder<'_, '_, T> {
106106
ColumnOption::NotNull => nullable = false,
107107
ColumnOption::Unique { is_primary, .. } => {
108108
if *is_primary {
109-
column_desc.is_primary = true;
109+
column_desc.set_primary(true);
110110
nullable = false;
111111
// Skip other options when using primary key
112112
break;
113113
} else {
114-
column_desc.is_unique = true;
114+
column_desc.set_unique(true);
115115
}
116116
}
117117
ColumnOption::Default(expr) => {
@@ -125,7 +125,7 @@ impl<T: Transaction> Binder<'_, '_, T> {
125125
if expr.return_type() != column_desc.column_datatype {
126126
expr = ScalarExpression::TypeCast {
127127
expr: Box::new(expr),
128-
ty: column_desc.column_datatype,
128+
ty: column_desc.column_datatype.clone(),
129129
}
130130
}
131131
column_desc.default = Some(expr);

‎src/binder/delete.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ impl<T: Transaction> Binder<'_, '_, T> {
3030
let schema_buf = self.table_schema_buf.entry(table_name.clone()).or_default();
3131
let primary_key_column = source
3232
.columns(schema_buf)
33-
.find(|column| column.desc().is_primary)
33+
.find(|column| column.desc().is_primary())
3434
.cloned()
3535
.unwrap();
3636
let mut plan = match source {

‎src/binder/expr.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,7 @@ impl<'a, T: Transaction> Binder<'a, '_, T> {
186186
if ty == &LogicalType::SqlNull {
187187
*ty = result_ty;
188188
} else if ty != &result_ty {
189-
return Err(DatabaseError::Incomparable(*ty, result_ty));
189+
return Err(DatabaseError::Incomparable(ty.clone(), result_ty));
190190
}
191191
}
192192

‎src/catalog/column.rs‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -187,8 +187,8 @@ impl ColumnCatalog {
187187
#[derive(Debug, Clone, PartialEq, Eq, Hash, ReferenceSerialization)]
188188
pub struct ColumnDesc {
189189
pub(crate) column_datatype: LogicalType,
190-
pub(crate) is_primary: bool,
191-
pub(crate) is_unique: bool,
190+
is_primary: bool,
191+
is_unique: bool,
192192
pub(crate) default: Option<ScalarExpression>,
193193
}
194194

@@ -212,4 +212,20 @@ impl ColumnDesc {
212212
default,
213213
})
214214
}
215+
216+
pub(crate) fn is_primary(&self) -> bool {
217+
self.is_primary
218+
}
219+
220+
pub(crate) fn set_primary(&mut self, is_primary: bool) {
221+
self.is_primary = is_primary
222+
}
223+
224+
pub(crate) fn is_unique(&self) -> bool {
225+
self.is_unique
226+
}
227+
228+
pub(crate) fn set_unique(&mut self, is_unique: bool) {
229+
self.is_unique = is_unique
230+
}
215231
}

‎src/catalog/table.rs‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ pub struct TableCatalog {
2121
pub(crate) indexes: Vec<IndexMetaRef>,
2222

2323
schema_ref: SchemaRef,
24+
primary_keys: Vec<(usize, ColumnRef)>,
2425
}
2526

2627
//TODO: can add some like Table description and other information as attributes
@@ -73,17 +74,13 @@ impl TableCatalog {
7374
self.columns.len()
7475
}
7576

76-
pub(crate) fn primary_key(&self) -> Result<(usize, &ColumnRef), DatabaseError> {
77-
self.schema_ref
78-
.iter()
79-
.enumerate()
80-
.find(|(_, column)| column.desc().is_primary)
81-
.ok_or(DatabaseError::PrimaryKeyNotFound)
77+
pub(crate) fn primary_keys(&self) -> &[(usize, ColumnRef)] {
78+
&self.primary_keys
8279
}
8380

8481
pub(crate) fn types(&self) -> Vec<LogicalType> {
8582
self.columns()
86-
.map(|column| *column.datatype())
83+
.map(|column| column.datatype().clone())
8784
.collect_vec()
8885
}
8986

@@ -128,7 +125,17 @@ impl TableCatalog {
128125
}
129126

130127
let index_id = self.indexes.last().map(|index| index.id + 1).unwrap_or(0);
131-
let pk_ty = *self.primary_key()?.1.datatype();
128+
let primary_keys = self.primary_keys();
129+
let pk_ty = if primary_keys.len() == 1 {
130+
primary_keys[0].1.datatype().clone()
131+
} else {
132+
LogicalType::Tuple(
133+
primary_keys
134+
.iter()
135+
.map(|(_, column)| column.datatype().clone())
136+
.collect_vec(),
137+
)
138+
};
132139
let index = IndexMeta {
133140
id: index_id,
134141
column_ids,
@@ -154,13 +161,21 @@ impl TableCatalog {
154161
columns: BTreeMap::new(),
155162
indexes: vec![],
156163
schema_ref: Arc::new(vec![]),
164+
primary_keys: vec![],
157165
};
158166
let mut generator = Generator::new();
159167
for col_catalog in columns.into_iter() {
160168
let _ = table_catalog
161169
.add_column(col_catalog, &mut generator)
162170
.unwrap();
163171
}
172+
table_catalog.primary_keys = table_catalog
173+
.schema_ref
174+
.iter()
175+
.enumerate()
176+
.filter(|&(_, column)| column.desc().is_primary())
177+
.map(|(i, column)| (i, column.clone()))
178+
.collect_vec();
164179

165180
Ok(table_catalog)
166181
}
@@ -182,13 +197,20 @@ impl TableCatalog {
182197
columns.insert(column_id, i);
183198
}
184199
let schema_ref = Arc::new(column_refs.clone());
200+
let primary_keys = schema_ref
201+
.iter()
202+
.enumerate()
203+
.filter(|&(_, column)| column.desc().is_primary())
204+
.map(|(i, column)| (i, column.clone()))
205+
.collect_vec();
185206

186207
Ok(TableCatalog {
187208
name,
188209
column_idxs,
189210
columns,
190211
indexes,
191212
schema_ref,
213+
primary_keys,
192214
})
193215
}
194216
}

‎src/errors.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,8 @@ pub enum DatabaseError {
112112
),
113113
#[error("must contain primary key!")]
114114
PrimaryKeyNotFound,
115+
#[error("primaryKey only allows single or multiple values")]
116+
PrimaryKeyTooManyLayers,
115117
#[error("rocksdb: {0}")]
116118
RocksDB(
117119
#[source]

‎src/execution/ddl/add_column.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,15 +40,15 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for AddColumn {
4040
if_not_exists,
4141
} = &self.op;
4242

43-
let mut unique_values = column.desc().is_unique.then(Vec::new);
43+
let mut unique_values = column.desc().is_unique().then(Vec::new);
4444
let mut tuples = Vec::new();
4545
let schema = self.input.output_schema();
4646
let mut types = Vec::with_capacity(schema.len() + 1);
4747

4848
for column_ref in schema.iter() {
49-
types.push(*column_ref.datatype());
49+
types.push(column_ref.datatype().clone());
5050
}
51-
types.push(*column.datatype());
51+
types.push(column.datatype().clone());
5252

5353
let mut coroutine = build_read(self.input, cache, transaction);
5454

‎src/execution/ddl/drop_column.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropColumn {
4141
.iter()
4242
.enumerate()
4343
.find(|(_, column)| column.name() == column_name)
44-
.map(|(i, column)| (i, column.desc().is_primary))
44+
.map(|(i, column)| (i, column.desc().is_primary()))
4545
{
4646
if is_primary {
4747
throw!(Err(DatabaseError::InvalidColumn(
@@ -55,7 +55,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for DropColumn {
5555
if i == column_index {
5656
continue;
5757
}
58-
types.push(*column_ref.datatype());
58+
types.push(column_ref.datatype().clone());
5959
}
6060
let mut coroutine = build_read(self.input, cache, transaction);
6161

‎src/execution/dml/insert.rs‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use crate::types::tuple::Tuple;
1111
use crate::types::tuple_builder::TupleBuilder;
1212
use crate::types::value::DataValue;
1313
use crate::types::ColumnId;
14+
use itertools::Itertools;
1415
use std::collections::HashMap;
1516
use std::ops::Coroutine;
1617
use std::ops::CoroutineState;
@@ -79,11 +80,14 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Insert {
7980
let mut tuples = Vec::new();
8081
let schema = input.output_schema().clone();
8182

82-
let pk_key = throw!(schema
83+
let primary_keys = schema
8384
.iter()
84-
.find(|col| col.desc().is_primary)
85+
.filter(|&col| col.desc().is_primary())
8586
.map(|col| col.key(is_mapping_by_name))
86-
.ok_or(DatabaseError::NotNull));
87+
.collect_vec();
88+
if primary_keys.is_empty() {
89+
throw!(Err(DatabaseError::NotNull))
90+
}
8791

8892
if let Some(table_catalog) =
8993
throw!(transaction.table(cache.0, table_name.clone())).cloned()
@@ -94,14 +98,18 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Insert {
9498
while let CoroutineState::Yielded(tuple) = Pin::new(&mut coroutine).resume(()) {
9599
let Tuple { values, .. } = throw!(tuple);
96100

101+
let mut tuple_id = Vec::with_capacity(primary_keys.len());
97102
let mut tuple_map = HashMap::new();
98103
for (i, value) in values.into_iter().enumerate() {
99104
tuple_map.insert(schema[i].key(is_mapping_by_name), value);
100105
}
101-
let tuple_id = throw!(tuple_map
102-
.get(&pk_key)
103-
.cloned()
104-
.ok_or(DatabaseError::NotNull));
106+
107+
for primary_key in primary_keys.iter() {
108+
tuple_id.push(throw!(tuple_map
109+
.get(primary_key)
110+
.cloned()
111+
.ok_or(DatabaseError::NotNull)));
112+
}
105113
let mut values = Vec::with_capacity(table_catalog.columns_len());
106114

107115
for col in table_catalog.columns() {
@@ -120,7 +128,11 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Insert {
120128
values.push(value)
121129
}
122130
tuples.push(Tuple {
123-
id: Some(tuple_id),
131+
id: Some(if primary_keys.len() == 1 {
132+
tuple_id.pop().unwrap()
133+
} else {
134+
Arc::new(DataValue::Tuple(Some(tuple_id)))
135+
}),
124136
values,
125137
});
126138
}

‎src/execution/dml/update.rs‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,12 @@ use crate::types::index::Index;
99
use crate::types::tuple::types;
1010
use crate::types::tuple::Tuple;
1111
use crate::types::tuple_builder::TupleBuilder;
12+
use crate::types::value::DataValue;
1213
use std::collections::HashMap;
1314
use std::ops::Coroutine;
1415
use std::ops::CoroutineState;
1516
use std::pin::Pin;
17+
use std::sync::Arc;
1618

1719
pub struct Update {
1820
table_name: TableName,
@@ -93,18 +95,28 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Update {
9395
}
9496
for mut tuple in tuples {
9597
let mut is_overwrite = true;
96-
98+
let mut primary_keys = Vec::new();
9799
for (i, column) in input_schema.iter().enumerate() {
98100
if let Some(value) = value_map.get(&column.id()) {
99-
if column.desc().is_primary {
100-
let old_key = tuple.id.replace(value.clone()).unwrap();
101-
102-
throw!(transaction.remove_tuple(&table_name, &old_key));
103-
is_overwrite = false;
101+
if column.desc().is_primary() {
102+
primary_keys.push(value.clone());
104103
}
105104
tuple.values[i] = value.clone();
106105
}
107106
}
107+
if !primary_keys.is_empty() {
108+
let id = if primary_keys.len() == 1 {
109+
primary_keys.pop().unwrap()
110+
} else {
111+
Arc::new(DataValue::Tuple(Some(primary_keys)))
112+
};
113+
if &id != tuple.id.as_ref().unwrap() {
114+
let old_key = tuple.id.replace(id).unwrap();
115+
116+
throw!(transaction.remove_tuple(&table_name, &old_key));
117+
is_overwrite = false;
118+
}
119+
}
108120
for (index_meta, exprs) in index_metas.iter() {
109121
let values =
110122
throw!(Projection::projection(&tuple, exprs, &input_schema));

0 commit comments

Comments
 (0)