feat(binder): support “primary key (a, b)” syntax

src/binder/statement/create_table.rs

@@ -1,5 +1,7 @@
 // Copyright 2022 RisingLight Project Authors. Licensed under Apache-2.0.
 
+use sqlparser::ast::TableConstraint;
+
 use super::*;
 use crate::catalog::{ColumnCatalog, ColumnDesc};
 use crate::parser::{ColumnDef, ColumnOption, Statement};
@@ -17,7 +19,12 @@ pub struct BoundCreateTable {
 impl Binder {
     pub fn bind_create_table(&mut self, stmt: &Statement) -> Result<BoundCreateTable, BindError> {
         match stmt {
-            Statement::CreateTable { name, columns, .. } => {
+            Statement::CreateTable {
+                name,
+                columns,
+                constraints,
+                ..
+            } => {
                 let name = &lower_case_name(name);
                 let (database_name, schema_name, table_name) = split_name(name)?;
                 let db = self
@@ -37,11 +44,33 @@ impl Binder {
                         return Err(BindError::DuplicatedColumn(col.name.value.clone()));
                     }
                 }
+
+                // record primary key names declared by "primary key" keywords
+                // if a sql like " create table t (a int, b int, c int, d int, primary key(a, b)); "
+                // then after the for loop, the extra_pk_name will contain "a" and "b"
+                let mut extra_pk_name: HashSet<String> = HashSet::new();
+                for constraint in constraints {
+                    match constraint {
+                        TableConstraint::Unique {
+                            is_primary,
+                            columns,
+                            ..
+                        } if *is_primary => columns.iter().for_each(|indent| {
+                            extra_pk_name.insert(indent.value.clone());
+                        }),
+
+                        _ => todo!(),
+                    }
+                }
+
                 let columns = columns
                     .iter()
                     .enumerate()
                     .map(|(idx, col)| {
                         let mut col = ColumnCatalog::from(col);
+                        if extra_pk_name.contains(col.name()) && !col.is_primary() {
+                            col.set_primary(true);
+                        }
                         col.set_id(idx as ColumnId);
                         col
                     })
@@ -97,7 +126,8 @@ mod tests {
         let sql = "
             create table t1 (v1 int not null, v2 int); 
             create table t2 (a int not null, a int not null);
-            create table t3 (v1 int not null);";
+            create table t3 (v1 int not null);
+            create table t4 (a int not null, b int not null, c int, primary key(a,b));";
         let stmts = parse(sql).unwrap();
 
         assert_eq!(
@@ -131,5 +161,32 @@ mod tests {
             binder.bind_create_table(&stmts[2]),
             Err(BindError::DuplicatedTable("t3".into()))
         );
+
+        assert_eq!(
+            binder.bind_create_table(&stmts[3]).unwrap(),
+            BoundCreateTable {
+                database_id: 0,
+                schema_id: 0,
+                table_name: "t4".into(),
+                columns: vec![
+                    ColumnCatalog::new(
+                        0,
+                        DataTypeKind::Int(None)
+                            .not_null()
+                            .to_column_primary_key("a".into()),
+                    ),
+                    ColumnCatalog::new(
+                        1,
+                        DataTypeKind::Int(None)
+                            .not_null()
+                            .to_column_primary_key("b".into()),
+                    ),
+                    ColumnCatalog::new(
+                        2,
+                        DataTypeKind::Int(None).nullable().to_column("c".into(),),
+                    ),
+                ],
+            }
+        );
     }
 }

src/catalog/table.rs

@@ -21,6 +21,7 @@ struct Inner {
     #[allow(dead_code)]
     is_materialized_view: bool,
     next_column_id: ColumnId,
+    primary_key_ids: Vec<ColumnId>,
 }
 
 impl TableCatalog {
@@ -38,14 +39,31 @@ impl TableCatalog {
                 columns: BTreeMap::new(),
                 is_materialized_view,
                 next_column_id: 0,
+                primary_key_ids: Vec::new(),
             }),
         };
+        let mut pk_ids = vec![];
         for col_catalog in columns {
+            if col_catalog.is_primary() {
+                pk_ids.push(col_catalog.id());
+            }
             table_catalog.add_column(col_catalog).unwrap();
         }
+
+        table_catalog.set_primary_key_ids(&pk_ids);
         table_catalog
     }
 
+    pub fn set_primary_key_ids(&self, pk_ids: &[ColumnId]) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.primary_key_ids = pk_ids.to_owned();
+    }
+
+    pub fn get_primary_key_ids(&self) -> Vec<ColumnId> {
+        let inner = self.inner.lock().unwrap();
+        inner.primary_key_ids.clone()
+    }
+
     pub fn add_column(&self, col_catalog: ColumnCatalog) -> Result<ColumnId, CatalogError> {
         let mut inner = self.inner.lock().unwrap();
         if inner.column_idxs.contains_key(col_catalog.name()) {
@@ -128,5 +146,25 @@ mod tests {
         let col1_catalog = table_catalog.get_column_by_id(1).unwrap();
         assert_eq!(col1_catalog.name(), "b");
         assert_eq!(col1_catalog.datatype().kind(), DataTypeKind::Boolean);
+
+        // test with two primary key
+        let col0 = ColumnCatalog::new(
+            0,
+            DataTypeKind::Int(None)
+                .not_null()
+                .to_column_primary_key("a".into()),
+        );
+        let col1 = ColumnCatalog::new(
+            1,
+            DataTypeKind::Int(None)
+                .not_null()
+                .to_column_primary_key("b".into()),
+        );
+        let col2 = ColumnCatalog::new(2, DataTypeKind::Int(None).nullable().to_column("c".into()));
+        let col3 = ColumnCatalog::new(3, DataTypeKind::Int(None).nullable().to_column("d".into()));
+
+        let col_catalogs = vec![col0, col1, col2, col3];
+        let table_catalog = TableCatalog::new(0, "t".into(), col_catalogs, false);
+        assert_eq!(table_catalog.get_primary_key_ids(), vec![0, 1]);
     }
 }

tests/sql/filter.slt

@@ -4,7 +4,7 @@ create table t (v1 int not null, v2 int not null, primary key(v1));
 statement ok
 insert into t values (1, 1), (4, 6), (3, 2), (2, 1)
 
-query I
+query I rowsort
 select v1 from t where v1 > 2
 ----
 4