switch to repr types in typecheck transform

Author: mgree

src/repr/src/explain.rs

@@ -45,7 +45,7 @@ use crate::explain::dot::{DisplayDot, dot_string};
 use crate::explain::json::{DisplayJson, json_string};
 use crate::explain::text::{DisplayText, text_string};
 use crate::optimize::OptimizerFeatureOverrides;
-use crate::{GlobalId, SqlColumnType, SqlScalarType};
+use crate::{GlobalId, ReprColumnType, ReprScalarType, SqlColumnType, SqlScalarType};
 
 pub mod dot;
 pub mod json;
@@ -444,6 +444,12 @@ pub trait ExprHumanizer: fmt::Debug {
     /// compatibility is more important.
     fn humanize_scalar_type(&self, ty: &SqlScalarType, postgres_compat: bool) -> String;
 
+    /// Returns a human-readable name for the specified scalar type.
+    /// Calls `humanize_scalar_type` with the `SqlScalarType` representation of the specified type.
+    fn humanize_scalar_type_repr(&self, typ: &ReprScalarType, postgres_compat: bool) -> String {
+        self.humanize_scalar_type(&SqlScalarType::from_repr(typ), postgres_compat)
+    }
+
     /// Returns a human-readable name for the specified column type.
     /// Used in, e.g., EXPLAIN and error msgs, in which case exact Postgres compatibility is less
     /// important than showing as much detail as possible. Also used in `pg_typeof`, where Postgres
@@ -456,6 +462,12 @@ pub trait ExprHumanizer: fmt::Debug {
         )
     }
 
+    /// Returns a human-readable name for the specified column type.
+    /// Calls `humanize_column_type` with the `SqlColumnType` representation of the specified type.
+    fn humanize_column_type_repr(&self, typ: &ReprColumnType, postgres_compat: bool) -> String {
+        self.humanize_column_type(&SqlColumnType::from_repr(typ), postgres_compat)
+    }
+
     /// Returns a vector of column names for the relation identified by `id`.
     fn column_names_for_id(&self, id: GlobalId) -> Option<Vec<String>>;
 

src/repr/src/lib.rs

@@ -54,8 +54,8 @@ pub use crate::global_id::GlobalId;
 pub use crate::relation::{
     ColumnIndex, ColumnName, NotNullViolation, PropRelationDescDiff, ProtoColumnName,
     ProtoColumnType, ProtoRelationDesc, ProtoRelationType, RelationDesc, RelationDescBuilder,
-    RelationVersion, RelationVersionSelector, ReprColumnType, SqlColumnType, SqlRelationType,
-    UNKNOWN_COLUMN_NAME, VersionedRelationDesc, arb_relation_desc_diff,
+    RelationVersion, RelationVersionSelector, ReprColumnType, ReprRelationType, SqlColumnType,
+    SqlRelationType, UNKNOWN_COLUMN_NAME, VersionedRelationDesc, arb_relation_desc_diff,
     arb_relation_desc_projection, arb_row_for_relation,
 };
 pub use crate::row::encode::{RowColumnarDecoder, RowColumnarEncoder, preserves_order};
@@ -66,7 +66,7 @@ pub use crate::row::{
 };
 pub use crate::scalar::{
     ArrayRustType, AsColumnType, Datum, DatumType, PropArray, PropDatum, PropDict, PropList,
-    ProtoScalarType, ReprScalarType, ScalarBaseType, SqlScalarType, arb_datum,
-    arb_datum_for_column, arb_datum_for_scalar, arb_range_type,
+    ProtoScalarType, ReprScalarBaseType, ReprScalarType, SqlScalarBaseType, SqlScalarType,
+    arb_datum, arb_datum_for_column, arb_datum_for_scalar, arb_range_type,
 };
 pub use crate::timestamp::{Timestamp, TimestampManipulation};

src/repr/src/relation.rs

@@ -259,6 +259,94 @@ impl RustType<ProtoKey> for Vec<usize> {
     }
 }
 
+/// The type of a relation.
+#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash)]
+pub struct ReprRelationType {
+    /// The type for each column, in order.
+    pub column_types: Vec<ReprColumnType>,
+    /// Sets of indices that are "keys" for the collection.
+    ///
+    /// Each element in this list is a set of column indices, each with the
+    /// property that the collection contains at most one record with each
+    /// distinct set of values for each column. Alternately, for a specific set
+    /// of values assigned to the these columns there is at most one record.
+    ///
+    /// A collection can contain multiple sets of keys, although it is common to
+    /// have either zero or one sets of key indices.
+    #[serde(default)]
+    pub keys: Vec<Vec<usize>>,
+}
+
+impl ReprRelationType {
+    /// Constructs a `ReprRelationType` representing the relation with no columns and
+    /// no keys.
+    pub fn empty() -> Self {
+        ReprRelationType::new(vec![])
+    }
+
+    /// Constructs a new `ReprRelationType` from specified column types.
+    ///
+    /// The `ReprRelationType` will have no keys.
+    pub fn new(column_types: Vec<ReprColumnType>) -> Self {
+        ReprRelationType {
+            column_types,
+            keys: Vec::new(),
+        }
+    }
+
+    /// Adds a new key for the relation.
+    pub fn with_key(mut self, mut indices: Vec<usize>) -> Self {
+        indices.sort_unstable();
+        if !self.keys.contains(&indices) {
+            self.keys.push(indices);
+        }
+        self
+    }
+
+    pub fn with_keys(mut self, keys: Vec<Vec<usize>>) -> Self {
+        for key in keys {
+            self = self.with_key(key)
+        }
+        self
+    }
+
+    /// Computes the number of columns in the relation.
+    pub fn arity(&self) -> usize {
+        self.column_types.len()
+    }
+
+    /// Gets the index of the columns used when creating a default index.
+    pub fn default_key(&self) -> Vec<usize> {
+        if let Some(key) = self.keys.first() {
+            if key.is_empty() {
+                (0..self.column_types.len()).collect()
+            } else {
+                key.clone()
+            }
+        } else {
+            (0..self.column_types.len()).collect()
+        }
+    }
+
+    /// Returns all the column types in order, for this relation.
+    pub fn columns(&self) -> &[ReprColumnType] {
+        &self.column_types
+    }
+}
+
+impl From<&SqlRelationType> for ReprRelationType {
+    fn from(sql_relation_type: &SqlRelationType) -> Self {
+        ReprRelationType {
+            column_types: sql_relation_type
+                .column_types
+                .iter()
+                .map(ReprColumnType::from)
+                .collect(),
+            keys: sql_relation_type.keys.clone(),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash, MzReflect)]
 pub struct ReprColumnType {
     /// The underlying representation scalar type (e.g., Int32 or String) of this column.

src/repr/src/scalar.rs

@@ -1591,7 +1591,7 @@ impl fmt::Display for Datum<'_> {
 #[derive(
     Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Ord, PartialOrd, Hash, EnumKind, MzReflect,
 )]
-#[enum_kind(ScalarBaseType, derive(PartialOrd, Ord, Hash))]
+#[enum_kind(SqlScalarBaseType, derive(PartialOrd, Ord, Hash))]
 pub enum SqlScalarType {
     /// The type of [`Datum::True`] and [`Datum::False`].
     Bool,
@@ -3212,7 +3212,7 @@ impl SqlScalarType {
                                 && a.1.scalar_type.eq_inner(&b.1.scalar_type, structure_only)
                         })
             }
-            (s, o) => ScalarBaseType::from(s) == ScalarBaseType::from(o),
+            (s, o) => SqlScalarBaseType::from(s) == SqlScalarBaseType::from(o),
         }
     }
 
@@ -4068,7 +4068,10 @@ impl Arbitrary for ReprScalarType {
 /// There is a direct correspondence between `Datum` variants and `ReprScalarType`
 /// variants: every `Datum` variant corresponds to exactly one `ReprScalarType` variant
 /// (with an exception for `Datum::Array`, which could be both an `Int2Vector` and an `Array`).
-#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Ord, PartialOrd, Hash, MzReflect)]
+#[derive(
+    Clone, Debug, EnumKind, PartialEq, Eq, Serialize, Deserialize, Ord, PartialOrd, Hash, MzReflect,
+)]
+#[enum_kind(ReprScalarBaseType, derive(PartialOrd, Ord, Hash))]
 pub enum ReprScalarType {
     Bool,
     Int16,

src/sql/src/func.rs

@@ -21,7 +21,7 @@ use mz_ore::collections::CollectionExt;
 use mz_ore::str::StrExt;
 use mz_pgrepr::oid;
 use mz_repr::role_id::RoleId;
-use mz_repr::{ColumnName, Datum, ScalarBaseType, SqlRelationType, SqlScalarType};
+use mz_repr::{ColumnName, Datum, SqlRelationType, SqlScalarBaseType, SqlScalarType};
 
 use crate::ast::{SelectStatement, Statement};
 use crate::catalog::{CatalogType, TypeCategory, TypeReference};
@@ -856,9 +856,9 @@ impl From<SqlScalarType> for ParamType {
     }
 }
 
-impl From<ScalarBaseType> for ParamType {
-    fn from(s: ScalarBaseType) -> ParamType {
-        use ScalarBaseType::*;
+impl From<SqlScalarBaseType> for ParamType {
+    fn from(s: SqlScalarBaseType) -> ParamType {
+        use SqlScalarBaseType::*;
         let s = match s {
             Array | List | Map | Record | Range => {
                 panic!("use polymorphic parameters rather than {:?}", s);
@@ -938,8 +938,8 @@ impl From<ParamType> for ReturnType {
     }
 }
 
-impl From<ScalarBaseType> for ReturnType {
-    fn from(s: ScalarBaseType) -> ReturnType {
+impl From<SqlScalarBaseType> for ReturnType {
+    fn from(s: SqlScalarBaseType) -> ReturnType {
         ParamType::from(s).into()
     }
 }
@@ -1802,7 +1802,7 @@ macro_rules! privilege_fn {
 /// Correlates a built-in function name to its implementations.
 pub static PG_CATALOG_BUILTINS: LazyLock<BTreeMap<&'static str, Func>> = LazyLock::new(|| {
     use ParamType::*;
-    use ScalarBaseType::*;
+    use SqlScalarBaseType::*;
     let mut builtins = builtins! {
         // Literal OIDs collected from PG 13 using a version of this query
         // ```sql
@@ -3537,7 +3537,7 @@ pub static INFORMATION_SCHEMA_BUILTINS: LazyLock<BTreeMap<&'static str, Func>> =
 
 pub static MZ_CATALOG_BUILTINS: LazyLock<BTreeMap<&'static str, Func>> = LazyLock::new(|| {
     use ParamType::*;
-    use ScalarBaseType::*;
+    use SqlScalarBaseType::*;
     builtins! {
         "constant_time_eq" => Scalar {
             params!(Bytes, Bytes) => BinaryFunc::from(func::ConstantTimeEqBytes) => Bool, oid::FUNC_CONSTANT_TIME_EQ_BYTES_OID;
@@ -3952,7 +3952,7 @@ pub static MZ_CATALOG_BUILTINS: LazyLock<BTreeMap<&'static str, Func>> = LazyLoc
 
 pub static MZ_INTERNAL_BUILTINS: LazyLock<BTreeMap<&'static str, Func>> = LazyLock::new(|| {
     use ParamType::*;
-    use ScalarBaseType::*;
+    use SqlScalarBaseType::*;
     builtins! {
         "aclitem_grantor" => Scalar {
             params!(AclItem) => UnaryFunc::AclItemGrantor(func::AclItemGrantor) => Oid, oid::FUNC_ACL_ITEM_GRANTOR_OID;
@@ -4339,7 +4339,7 @@ pub static MZ_INTERNAL_BUILTINS: LazyLock<BTreeMap<&'static str, Func>> = LazyLo
 
 pub static MZ_UNSAFE_BUILTINS: LazyLock<BTreeMap<&'static str, Func>> = LazyLock::new(|| {
     use ParamType::*;
-    use ScalarBaseType::*;
+    use SqlScalarBaseType::*;
     builtins! {
         "mz_all" => Aggregate {
             params!(Any) => AggregateFunc::All => Bool, oid::FUNC_MZ_ALL_OID;
@@ -4459,7 +4459,7 @@ fn array_to_string(
 pub static OP_IMPLS: LazyLock<BTreeMap<&'static str, Func>> = LazyLock::new(|| {
     use BinaryFunc as BF;
     use ParamType::*;
-    use ScalarBaseType::*;
+    use SqlScalarBaseType::*;
     builtins! {
         // Literal OIDs collected from PG 13 using a version of this query
         // ```sql

src/sql/src/plan/typeconv.rs

@@ -18,7 +18,9 @@ use dynfmt::{Format, SimpleCurlyFormat};
 use itertools::Itertools;
 use mz_expr::func::{CastArrayToJsonb, CastListToJsonb};
 use mz_expr::{VariadicFunc, func};
-use mz_repr::{ColumnName, Datum, ScalarBaseType, SqlColumnType, SqlRelationType, SqlScalarType};
+use mz_repr::{
+    ColumnName, Datum, SqlColumnType, SqlRelationType, SqlScalarBaseType, SqlScalarType,
+};
 
 use crate::catalog::TypeCategory;
 use crate::plan::error::PlanError;
@@ -294,9 +296,9 @@ macro_rules! casts(
     }};
 );
 
-static VALID_CASTS: LazyLock<BTreeMap<(ScalarBaseType, ScalarBaseType), CastImpl>> = LazyLock::new(
-    || {
-        use ScalarBaseType::*;
+static VALID_CASTS: LazyLock<BTreeMap<(SqlScalarBaseType, SqlScalarBaseType), CastImpl>> =
+    LazyLock::new(|| {
+        use SqlScalarBaseType::*;
         use UnaryFunc::*;
 
         casts! {
@@ -865,8 +867,7 @@ static VALID_CASTS: LazyLock<BTreeMap<(ScalarBaseType, ScalarBaseType), CastImpl
                 )
             )")
         }
-    },
-);
+    });
 
 /// Get casts directly between two [`SqlScalarType`]s, with control over the
 /// allowed [`CastContext`].