Support UNIQUE NULLS DISTINCT index in CREATE TABLE

Author: kezhuw

src/protos/generated/mod.rs

@@ -1,9 +1,9 @@
 #![allow(clippy::all)]
 
-#[rustfmt::skip]
-mod seamdb;
-pub use self::seamdb::*;
-
 #[rustfmt::skip]
 mod sql;
 pub use self::sql::*;
+
+#[rustfmt::skip]
+mod seamdb;
+pub use self::seamdb::*;

src/protos/generated/sql.rs

@@ -110,6 +110,14 @@ pub struct StoringFloat64 {
     #[prost(double, tag = "1")]
     pub value: f64,
 }
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct Column {
+    #[prost(uint32, tag = "1")]
+    pub id: u32,
+    #[prost(message, optional, tag = "2")]
+    pub value: ::core::option::Option<ColumnValue>,
+}
 #[derive(Eq, Hash, PartialOrd)]
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
@@ -158,10 +166,12 @@ pub struct IndexDescriptor {
     pub id: u32,
     #[prost(string, tag = "2")]
     pub name: ::prost::alloc::string::String,
-    #[prost(bool, tag = "3")]
-    pub unique: bool,
+    #[prost(enumeration = "!IndexKind", tag = "3")]
+    pub kind: IndexKind,
+    /// Declared index columns.
     #[prost(uint32, repeated, tag = "4")]
     pub column_ids: ::prost::alloc::vec::Vec<u32>,
+    /// Declared storing columns,
     #[prost(uint32, repeated, tag = "5")]
     pub storing_column_ids: ::prost::alloc::vec::Vec<u32>,
 }
@@ -209,3 +219,35 @@ impl ColumnTypeKind {
         }
     }
 }
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
+#[repr(i32)]
+pub enum IndexKind {
+    NotUnique = 0,
+    UniqueNullsDistinct = 1,
+    UniqueNullsNotDistinct = 2,
+    PrimaryKey = 3,
+}
+impl IndexKind {
+    /// String value of the enum field names used in the ProtoBuf definition.
+    ///
+    /// The values are not transformed in any way and thus are considered stable
+    /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+    pub fn as_str_name(&self) -> &'static str {
+        match self {
+            IndexKind::NotUnique => "NotUnique",
+            IndexKind::UniqueNullsDistinct => "UNIQUE_NULLS_DISTINCT",
+            IndexKind::UniqueNullsNotDistinct => "UNIQUE_NULLS_NOT_DISTINCT",
+            IndexKind::PrimaryKey => "PRIMARY_KEY",
+        }
+    }
+    /// Creates an enum from field names used in the ProtoBuf definition.
+    pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+        match value {
+            "NotUnique" => Some(Self::NotUnique),
+            "UNIQUE_NULLS_DISTINCT" => Some(Self::UniqueNullsDistinct),
+            "UNIQUE_NULLS_NOT_DISTINCT" => Some(Self::UniqueNullsNotDistinct),
+            "PRIMARY_KEY" => Some(Self::PrimaryKey),
+            _ => None,
+        }
+    }
+}

src/protos/protos/sql.proto

@@ -100,6 +100,11 @@ message StoringFloat64 {
     double value = 1;
 }
 
+message Column {
+    uint32 id = 1;
+    ColumnValue value = 2;
+}
+
 message ColumnValue {
     oneof value {
         bool boolean = 1;
@@ -123,10 +128,21 @@ message ColumnDescriptor {
     ColumnValue default_value = 7;
 }
 
+enum IndexKind {
+    NotUnique = 0;
+    UNIQUE_NULLS_DISTINCT = 1;
+    UNIQUE_NULLS_NOT_DISTINCT = 2;
+    PRIMARY_KEY = 3;
+}
+
 message IndexDescriptor {
     uint32 id = 1;
     string name = 2;
-    bool unique = 3;
+    IndexKind kind = 3;
+
+    // Declared index columns.
     repeated uint32 column_ids = 4;
+
+    // Declared storing columns,
     repeated uint32 storing_column_ids = 5;
 }

src/protos/sql.rs

@@ -12,24 +12,38 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::fmt::{self, Display, Formatter};
 use std::hash::{Hash, Hasher};
 
 use bytes::{Buf, BufMut};
 use prost::Message;
 
 use crate::protos::{
+    Column,
     ColumnDescriptor,
     ColumnTypeKind,
     ColumnValue,
     DatabaseDescriptor,
     DescriptorMeta,
     IndexDescriptor,
+    IndexKind,
     SchemaDescriptor,
     StoringFloat32,
     StoringFloat64,
     TableDescriptor,
 };
 
+impl Display for IndexKind {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::NotUnique => f.write_str("index"),
+            Self::UniqueNullsDistinct => f.write_str("unique nulls distinct"),
+            Self::UniqueNullsNotDistinct => f.write_str("unique nulls not distinct"),
+            Self::PrimaryKey => f.write_str("primary index"),
+        }
+    }
+}
+
 pub trait NamespaceDescriptor {
     fn kind() -> &'static str;
 
@@ -113,7 +127,7 @@ impl TableDescriptor {
 
     pub fn primary_index(&self) -> &IndexDescriptor {
         for index in self.indices.iter() {
-            if index.id == 1 {
+            if index.is_primary() {
                 return index;
             }
         }
@@ -134,6 +148,10 @@ impl TableDescriptor {
         key
     }
 
+    pub fn column_indices(&self, columns: impl Iterator<Item = u32>) -> Vec<usize> {
+        columns.map(|id| self.columns.iter().position(|column| column.id == id).unwrap()).collect()
+    }
+
     pub fn find_column(&self, name: &str) -> Option<&ColumnDescriptor> {
         self.columns.iter().find(|column| column.name == name)
     }
@@ -142,22 +160,30 @@ impl TableDescriptor {
         self.columns.iter().find(|column| column.id == id)
     }
 
-    pub fn decode_storing_columns(&self, index: &IndexDescriptor, mut bytes: &[u8]) -> Vec<Option<ColumnValue>> {
-        let mut values = Vec::with_capacity(index.storing_column_ids.len());
+    pub fn decode_key_columns_to(&self, index: &IndexDescriptor, mut bytes: &[u8], row: &mut Vec<Column>) {
+        let prefix = self.index_prefix(index);
+        bytes.advance(prefix.len());
         loop {
             let id = bytes.get_u32();
             if id == 0 {
                 break;
             }
-            let value = ColumnValue::decode(&mut bytes);
-            let i = values.len();
-            if i >= index.storing_column_ids.len() {
-                values.push(value);
-                panic!("index key has more columns than descriptor {index:?}: {values:?}")
+            let value = ColumnValue::decode_from_key(&mut bytes);
+            let column = self.column(id).unwrap();
+            if value.type_kind() != column.type_kind {
+                panic!("index {index:?} column {column:?} mismatch column value {value:?}")
             }
-            if id != index.storing_column_ids[i] {
-                panic!("index descriptor {index:?} does not have column id {id}")
+            row.push(Column::new(id, Some(value)));
+        }
+    }
+
+    pub fn decode_value_columns_to(&self, index: &IndexDescriptor, mut bytes: &[u8], row: &mut Vec<Column>) {
+        loop {
+            let id = bytes.get_u32();
+            if id == 0 {
+                break;
             }
+            let value = ColumnValue::decode(&mut bytes);
             let column = self.column(id).unwrap();
             match value.as_ref().map(|v| v.type_kind()) {
                 None => {
@@ -171,48 +197,26 @@ impl TableDescriptor {
                     }
                 },
             }
-            values.push(value);
-        }
-        if values.len() != index.storing_column_ids.len() {
-            panic!("index {index:?} get mismatching values {values:?}")
+            row.push(Column::new(id, value));
         }
-        values
     }
 
-    pub fn decode_index_key(&self, index: &IndexDescriptor, mut bytes: &[u8]) -> Vec<ColumnValue> {
-        let prefix = self.index_prefix(index);
-        bytes.advance(prefix.len());
-        let mut values = Vec::with_capacity(index.column_ids.len());
-        loop {
-            let id = bytes.get_u32();
-            if id == 0 {
-                break;
-            }
-            let value = ColumnValue::decode_from_key(&mut bytes);
-            let i = values.len();
-            if i >= index.column_ids.len() {
-                values.push(value);
-                panic!("index key has more columns than descriptor {index:?}: {values:?}")
-            }
-            if id != index.column_ids[i] {
-                panic!("index descriptor {index:?} does not have column id {id}")
-            }
-            let column = self.column(id).unwrap();
-            if value.type_kind() != column.type_kind {
-                panic!("index {index:?} column {column:?} mismatch column value {value:?}")
-            }
-            values.push(value);
-        }
-        if values.len() != index.column_ids.len() {
-            panic!("index {index:?} get mismatching keys {values:?}")
-        }
+    pub fn decode_key_columns(&self, index: &IndexDescriptor, bytes: &[u8]) -> Vec<Column> {
+        let mut keys = Vec::with_capacity(index.column_ids.len());
+        self.decode_key_columns_to(index, bytes, &mut keys);
+        keys
+    }
+
+    pub fn decode_value_columns(&self, index: &IndexDescriptor, bytes: &[u8]) -> Vec<Column> {
+        let mut values = Vec::with_capacity(index.storing_column_ids.len());
+        self.decode_value_columns_to(index, bytes, &mut values);
         values
     }
 }
 
 impl IndexDescriptor {
     pub fn is_primary(&self) -> bool {
-        self.id == 1
+        self.kind == IndexKind::PrimaryKey
     }
 
     pub fn sole_column(&self) -> Option<u32> {
@@ -223,6 +227,44 @@ impl IndexDescriptor {
     }
 }
 
+impl Column {
+    pub fn new(id: u32, value: Option<ColumnValue>) -> Self {
+        Self { id, value }
+    }
+
+    pub fn with_value(id: u32, value: ColumnValue) -> Self {
+        Self { id, value: Some(value) }
+    }
+
+    pub fn new_null(id: u32) -> Self {
+        Self { id, value: None }
+    }
+
+    pub fn is_null(&self) -> bool {
+        self.value.is_none()
+    }
+
+    pub fn encode_as_key(&self, buf: &mut impl BufMut) {
+        buf.put_u32(self.id);
+        match self.value.as_ref() {
+            None => {
+                buf.put_i32(-1);
+            },
+            Some(value) => value.encode_as_key(buf),
+        }
+    }
+
+    pub fn encode_as_value(&self, buf: &mut impl BufMut) {
+        buf.put_u32(self.id);
+        match self.value.as_ref() {
+            None => {
+                buf.put_i32(-1);
+            },
+            Some(value) => value.encode(buf),
+        }
+    }
+}
+
 impl ColumnValue {
     pub fn minimum_of(type_kind: ColumnTypeKind) -> Self {
         match type_kind {