Skip to main content

mz_repr/
relation.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::collections::{BTreeMap, BTreeSet};
11#[cfg(any(test, feature = "proptest"))]
12use std::rc::Rc;
13use std::{fmt, vec};
14
15use anyhow::bail;
16use itertools::Itertools;
17use mz_ore::cast::CastFrom;
18use mz_ore::soft_panic_or_log;
19use mz_ore::str::StrExt;
20use mz_ore::{assert_none, assert_ok};
21use mz_persist_types::schema::SchemaId;
22use mz_proto::{IntoRustIfSome, ProtoType, RustType, TryFromProtoError};
23#[cfg(any(test, feature = "proptest"))]
24use proptest::prelude::*;
25#[cfg(any(test, feature = "proptest"))]
26use proptest::strategy::{Strategy, Union};
27#[cfg(any(test, feature = "proptest"))]
28use proptest_derive::Arbitrary;
29use serde::{Deserialize, Serialize};
30
31#[cfg(any(test, feature = "proptest"))]
32use crate::Row;
33#[cfg(any(test, feature = "proptest"))]
34use crate::arb_datum_for_column;
35use crate::relation_and_scalar::proto_relation_type::ProtoKey;
36pub use crate::relation_and_scalar::{
37    ProtoColumnMetadata, ProtoColumnName, ProtoColumnType, ProtoRelationDesc, ProtoRelationType,
38    ProtoRelationVersion,
39};
40use crate::{Datum, ReprScalarType, SqlScalarType};
41
42/// The type of a [`Datum`].
43///
44/// [`SqlColumnType`] bundles information about the scalar type of a datum (e.g.,
45/// Int32 or String) with its nullability.
46///
47/// To construct a column type, either initialize the struct directly, or
48/// use the [`SqlScalarType::nullable`] method.
49#[derive(
50    Clone,
51    Debug,
52    Eq,
53    PartialEq,
54    Ord,
55    PartialOrd,
56    Serialize,
57    Deserialize,
58    Hash
59)]
60#[cfg_attr(any(test, feature = "proptest"), derive(Arbitrary))]
61pub struct SqlColumnType {
62    /// The underlying scalar type (e.g., Int32 or String) of this column.
63    pub scalar_type: SqlScalarType,
64    /// Whether this datum can be null.
65    #[serde(default = "return_true")]
66    pub nullable: bool,
67}
68
69/// This method exists solely for the purpose of making SqlColumnType nullable by
70/// default in unit tests. The default value of a bool is false, and the only
71/// way to make an object take on any other value by default is to pass it a
72/// function that returns the desired default value. See
73/// <https://github.com/serde-rs/serde/issues/1030>
74#[inline(always)]
75fn return_true() -> bool {
76    true
77}
78
79impl SqlColumnType {
80    /// Compute the least upper bound of many column types, returning an error on
81    /// incompatible types or an empty iterator.
82    /// See [`SqlColumnType::try_union`] for details.
83    pub fn try_union_many<'a>(
84        typs: impl IntoIterator<Item = &'a Self>,
85    ) -> Result<Self, anyhow::Error> {
86        let mut iter = typs.into_iter();
87        let Some(typ) = iter.next() else {
88            bail!("Cannot union empty iterator");
89        };
90        iter.try_fold(typ.clone(), |a, b| a.try_union(b))
91    }
92
93    /// Compute the least upper bound of many column types.
94    /// See [`SqlColumnType::try_union`] for details.
95    ///
96    /// Panics on incompatible types or an empty iterator.
97    pub fn union_many<'a>(typs: impl IntoIterator<Item = &'a Self>) -> Self {
98        Self::try_union_many(typs).expect("Cannot union empty iterator")
99    }
100
101    /// Backports nullability information from `backport_typ` into `self`,
102    /// affecting the outer `.nullable` field but also record fields deeper
103    /// into the type.
104    pub fn backport_nullability(&mut self, backport_typ: &ReprColumnType) {
105        self.scalar_type
106            .backport_nullability(&backport_typ.scalar_type);
107        self.nullable = backport_typ.nullable;
108    }
109
110    /// Compute the least upper bound of two column types at the SQL level.
111    ///
112    /// Nullability is the disjunction of the two inputs, at every nesting depth.
113    /// See [`SqlScalarType::sql_union`] for which types are compatible.
114    ///
115    /// Returns an error for incompatible types, e.g. `Text` and `Int32`, or
116    /// `Text` and `VarChar`. See [`SqlColumnType::try_union`] for a fallback
117    /// that handles the latter via repr-level union.
118    pub fn sql_union(&self, other: &Self) -> Result<Self, anyhow::Error> {
119        Ok(SqlColumnType {
120            scalar_type: self.scalar_type.sql_union(&other.scalar_type)?,
121            nullable: self.nullable || other.nullable,
122        })
123    }
124
125    /// Compute the least upper bound of two column types.
126    ///
127    /// Attempts [`SqlColumnType::sql_union`] first, which preserves SQL-level type
128    /// information (e.g. modifiers). Falls back to a repr-level union via
129    /// [`ReprColumnType::union`] when the SQL types are incompatible but the
130    /// underlying repr types are compatible.
131    ///
132    /// The resulting nullability is the disjunction of the two input
133    /// nullabilities.
134    pub fn try_union(&self, other: &Self) -> Result<Self, anyhow::Error> {
135        self.sql_union(other).or_else(|e| {
136            let repr_self = ReprColumnType::from(self);
137            let repr_other = ReprColumnType::from(other);
138            match repr_self.union(&repr_other) {
139                Ok(typ) => {
140                    // sql_union failed but repr union succeeded — this indicates
141                    // a repr-type canonicalization gap that we want CI visibility for.
142                    soft_panic_or_log!("repr type error: sql_union({self:?}, {other:?}): {e}");
143                    Ok(SqlColumnType::from_repr(&typ))
144                }
145                Err(_) => {
146                    // Both sql_union and repr union failed — genuine type mismatch,
147                    // not a canonicalization issue. Just propagate the original error.
148                    Err(e)
149                }
150            }
151        })
152    }
153
154    /// Compute the least upper bound of two column types.
155    /// See [`SqlColumnType::try_union`] for details.
156    ///
157    /// Panics on incompatible types.
158    pub fn union(&self, other: &Self) -> Self {
159        self.try_union(other).unwrap_or_else(|e| {
160            panic!("repr type error: after sql_union({self:?}, {other:?}) error: {e}")
161        })
162    }
163
164    /// Consumes this `SqlColumnType` and returns a new `SqlColumnType` with its
165    /// nullability set to the specified boolean.
166    pub fn nullable(mut self, nullable: bool) -> Self {
167        self.nullable = nullable;
168        self
169    }
170}
171
172impl RustType<ProtoColumnType> for SqlColumnType {
173    fn into_proto(&self) -> ProtoColumnType {
174        ProtoColumnType {
175            nullable: self.nullable,
176            scalar_type: Some(self.scalar_type.into_proto()),
177        }
178    }
179
180    fn from_proto(proto: ProtoColumnType) -> Result<Self, TryFromProtoError> {
181        Ok(SqlColumnType {
182            nullable: proto.nullable,
183            scalar_type: proto
184                .scalar_type
185                .into_rust_if_some("ProtoColumnType::scalar_type")?,
186        })
187    }
188}
189
190impl fmt::Display for SqlColumnType {
191    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
192        let nullable = if self.nullable { "Null" } else { "NotNull" };
193        f.write_fmt(format_args!("{:?}:{}", self.scalar_type, nullable))
194    }
195}
196
197/// The type of a relation.
198#[derive(
199    Clone,
200    Debug,
201    Eq,
202    PartialEq,
203    Ord,
204    PartialOrd,
205    Serialize,
206    Deserialize,
207    Hash
208)]
209#[cfg_attr(any(test, feature = "proptest"), derive(Arbitrary))]
210pub struct SqlRelationType {
211    /// The type for each column, in order.
212    pub column_types: Vec<SqlColumnType>,
213    /// Sets of indices that are "keys" for the collection.
214    ///
215    /// Each element in this list is a set of column indices, each with the
216    /// property that the collection contains at most one record with each
217    /// distinct set of values for each column. Alternately, for a specific set
218    /// of values assigned to the these columns there is at most one record.
219    ///
220    /// A collection can contain multiple sets of keys, although it is common to
221    /// have either zero or one sets of key indices.
222    #[serde(default)]
223    pub keys: Vec<Vec<usize>>,
224}
225
226impl SqlRelationType {
227    /// Constructs a `SqlRelationType` representing the relation with no columns and
228    /// no keys.
229    pub fn empty() -> Self {
230        SqlRelationType::new(vec![])
231    }
232
233    /// Constructs a new `SqlRelationType` from specified column types.
234    ///
235    /// The `SqlRelationType` will have no keys.
236    pub fn new(column_types: Vec<SqlColumnType>) -> Self {
237        SqlRelationType {
238            column_types,
239            keys: Vec::new(),
240        }
241    }
242
243    /// Adds a new key for the relation.
244    pub fn with_key(mut self, mut indices: Vec<usize>) -> Self {
245        indices.sort_unstable();
246        if !self.keys.contains(&indices) {
247            self.keys.push(indices);
248        }
249        self
250    }
251
252    pub fn with_keys(mut self, keys: Vec<Vec<usize>>) -> Self {
253        for key in keys {
254            self = self.with_key(key)
255        }
256        self
257    }
258
259    /// Computes the number of columns in the relation.
260    pub fn arity(&self) -> usize {
261        self.column_types.len()
262    }
263
264    /// Gets the index of the columns used when creating a default index.
265    pub fn default_key(&self) -> Vec<usize> {
266        if let Some(key) = self.keys.first() {
267            if key.is_empty() {
268                (0..self.column_types.len()).collect()
269            } else {
270                key.clone()
271            }
272        } else {
273            (0..self.column_types.len()).collect()
274        }
275    }
276
277    /// Returns all the [`SqlColumnType`]s, in order, for this relation.
278    pub fn columns(&self) -> &[SqlColumnType] {
279        &self.column_types
280    }
281
282    /// Adopts the nullability and keys from another `SqlRelationType`.
283    ///
284    /// Panics if the number of columns does not match.
285    pub fn backport_nullability_and_keys(&mut self, backport_typ: &ReprRelationType) {
286        assert_eq!(
287            backport_typ.column_types.len(),
288            self.column_types.len(),
289            "HIR and MIR types should have the same number of columns"
290        );
291        for (backport_col, sql_col) in backport_typ
292            .column_types
293            .iter()
294            .zip_eq(self.column_types.iter_mut())
295        {
296            sql_col.backport_nullability(backport_col);
297        }
298
299        self.keys = backport_typ.keys.clone();
300    }
301
302    /// Constructs a `SqlRelationType` from a `ReprRelationType` by converting
303    /// each column type via [`SqlColumnType::from_repr`]. This is a lossy
304    /// inverse of `ReprRelationType::from(&SqlRelationType)`.
305    pub fn from_repr(repr: &ReprRelationType) -> Self {
306        SqlRelationType {
307            column_types: repr
308                .column_types
309                .iter()
310                .map(SqlColumnType::from_repr)
311                .collect(),
312            keys: repr.keys.clone(),
313        }
314    }
315}
316
317impl RustType<ProtoRelationType> for SqlRelationType {
318    fn into_proto(&self) -> ProtoRelationType {
319        ProtoRelationType {
320            column_types: self.column_types.into_proto(),
321            keys: self.keys.into_proto(),
322        }
323    }
324
325    fn from_proto(proto: ProtoRelationType) -> Result<Self, TryFromProtoError> {
326        Ok(SqlRelationType {
327            column_types: proto.column_types.into_rust()?,
328            keys: proto.keys.into_rust()?,
329        })
330    }
331}
332
333impl RustType<ProtoKey> for Vec<usize> {
334    fn into_proto(&self) -> ProtoKey {
335        ProtoKey {
336            keys: self.into_proto(),
337        }
338    }
339
340    fn from_proto(proto: ProtoKey) -> Result<Self, TryFromProtoError> {
341        proto.keys.into_rust()
342    }
343}
344
345/// The type of a relation.
346#[derive(
347    Clone,
348    Debug,
349    Eq,
350    PartialEq,
351    Ord,
352    PartialOrd,
353    Serialize,
354    Deserialize,
355    Hash
356)]
357pub struct ReprRelationType {
358    /// The type for each column, in order.
359    pub column_types: Vec<ReprColumnType>,
360    /// Sets of indices that are "keys" for the collection.
361    ///
362    /// Each element in this list is a set of column indices, each with the
363    /// property that the collection contains at most one record with each
364    /// distinct set of values for each column. Alternately, for a specific set
365    /// of values assigned to the these columns there is at most one record.
366    ///
367    /// A collection can contain multiple sets of keys, although it is common to
368    /// have either zero or one sets of key indices.
369    #[serde(default)]
370    pub keys: Vec<Vec<usize>>,
371}
372
373impl ReprRelationType {
374    /// Constructs a `ReprRelationType` representing the relation with no columns and
375    /// no keys.
376    pub fn empty() -> Self {
377        ReprRelationType::new(vec![])
378    }
379
380    /// Constructs a new `ReprRelationType` from specified column types.
381    ///
382    /// The `ReprRelationType` will have no keys.
383    pub fn new(column_types: Vec<ReprColumnType>) -> Self {
384        ReprRelationType {
385            column_types,
386            keys: Vec::new(),
387        }
388    }
389
390    /// Adds a new key for the relation.
391    pub fn with_key(mut self, mut indices: Vec<usize>) -> Self {
392        indices.sort_unstable();
393        if !self.keys.contains(&indices) {
394            self.keys.push(indices);
395        }
396        self
397    }
398
399    pub fn with_keys(mut self, keys: Vec<Vec<usize>>) -> Self {
400        for key in keys {
401            self = self.with_key(key)
402        }
403        self
404    }
405
406    /// Computes the number of columns in the relation.
407    pub fn arity(&self) -> usize {
408        self.column_types.len()
409    }
410
411    /// Gets the index of the columns used when creating a default index.
412    pub fn default_key(&self) -> Vec<usize> {
413        if let Some(key) = self.keys.first() {
414            if key.is_empty() {
415                (0..self.column_types.len()).collect()
416            } else {
417                key.clone()
418            }
419        } else {
420            (0..self.column_types.len()).collect()
421        }
422    }
423
424    /// Returns all the column types in order, for this relation.
425    pub fn columns(&self) -> &[ReprColumnType] {
426        &self.column_types
427    }
428}
429
430impl From<&SqlRelationType> for ReprRelationType {
431    fn from(sql_relation_type: &SqlRelationType) -> Self {
432        ReprRelationType {
433            column_types: sql_relation_type
434                .column_types
435                .iter()
436                .map(ReprColumnType::from)
437                .collect(),
438            keys: sql_relation_type.keys.clone(),
439        }
440    }
441}
442
443#[derive(
444    Clone,
445    Debug,
446    Eq,
447    PartialEq,
448    Ord,
449    PartialOrd,
450    Serialize,
451    Deserialize,
452    Hash
453)]
454pub struct ReprColumnType {
455    /// The underlying representation scalar type (e.g., Int32 or String) of this column.
456    pub scalar_type: ReprScalarType,
457    /// Whether this datum can be null.
458    #[serde(default = "return_true")]
459    pub nullable: bool,
460}
461
462impl std::fmt::Display for ReprColumnType {
463    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
464        write!(f, "{}", self.scalar_type)?;
465        if self.nullable {
466            write!(f, "?")?;
467        }
468        Ok(())
469    }
470}
471
472impl ReprColumnType {
473    /// Compute the least upper bound of two column types at the repr level.
474    ///
475    /// More permissive than [`SqlColumnType::sql_union`] because it operates
476    /// on the underlying representation types, ignoring SQL-level distinctions
477    /// such as modifiers.
478    /// The resulting nullability is the disjunction of the two inputs.
479    pub fn union(&self, col: &ReprColumnType) -> Result<Self, anyhow::Error> {
480        let scalar_type = self.scalar_type.union(&col.scalar_type)?;
481        let nullable = self.nullable || col.nullable;
482
483        Ok(ReprColumnType {
484            scalar_type,
485            nullable,
486        })
487    }
488}
489
490impl From<&SqlColumnType> for ReprColumnType {
491    fn from(sql_column_type: &SqlColumnType) -> Self {
492        let scalar_type = &sql_column_type.scalar_type;
493        let scalar_type = scalar_type.into();
494        let nullable = sql_column_type.nullable;
495
496        ReprColumnType {
497            scalar_type,
498            nullable,
499        }
500    }
501}
502
503impl SqlColumnType {
504    /// Lossily translates a [`ReprColumnType`] back to a [`SqlColumnType`].
505    ///
506    /// See [`SqlScalarType::from_repr`] for an example of lossiness.
507    pub fn from_repr(repr: &ReprColumnType) -> Self {
508        let scalar_type = &repr.scalar_type;
509        let scalar_type = SqlScalarType::from_repr(scalar_type);
510        let nullable = repr.nullable;
511
512        SqlColumnType {
513            scalar_type,
514            nullable,
515        }
516    }
517}
518
519/// The name of a column in a [`RelationDesc`].
520#[derive(
521    Clone,
522    Debug,
523    Eq,
524    PartialEq,
525    Ord,
526    PartialOrd,
527    Serialize,
528    Deserialize,
529    Hash
530)]
531pub struct ColumnName(Box<str>);
532
533impl ColumnName {
534    /// Returns this column name as a `str`.
535    #[inline(always)]
536    pub fn as_str(&self) -> &str {
537        &*self
538    }
539
540    /// Returns this column name as a `&mut Box<str>`.
541    pub fn as_mut_boxed_str(&mut self) -> &mut Box<str> {
542        &mut self.0
543    }
544
545    /// Returns if this [`ColumnName`] is similar to the provided one.
546    pub fn is_similar(&self, other: &ColumnName) -> bool {
547        const SIMILARITY_THRESHOLD: f64 = 0.6;
548
549        let a_lowercase = self.to_lowercase();
550        let b_lowercase = other.to_lowercase();
551
552        strsim::normalized_levenshtein(&a_lowercase, &b_lowercase) >= SIMILARITY_THRESHOLD
553    }
554}
555
556impl std::ops::Deref for ColumnName {
557    type Target = str;
558
559    #[inline(always)]
560    fn deref(&self) -> &Self::Target {
561        &self.0
562    }
563}
564
565impl fmt::Display for ColumnName {
566    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
567        f.write_str(&self.0)
568    }
569}
570
571impl From<String> for ColumnName {
572    fn from(s: String) -> ColumnName {
573        ColumnName(s.into())
574    }
575}
576
577impl From<&str> for ColumnName {
578    fn from(s: &str) -> ColumnName {
579        ColumnName(s.into())
580    }
581}
582
583impl From<&ColumnName> for ColumnName {
584    fn from(n: &ColumnName) -> ColumnName {
585        n.clone()
586    }
587}
588
589impl RustType<ProtoColumnName> for ColumnName {
590    fn into_proto(&self) -> ProtoColumnName {
591        ProtoColumnName {
592            value: Some(self.0.to_string()),
593        }
594    }
595
596    fn from_proto(proto: ProtoColumnName) -> Result<Self, TryFromProtoError> {
597        Ok(ColumnName(
598            proto
599                .value
600                .ok_or_else(|| TryFromProtoError::missing_field("ProtoColumnName::value"))?
601                .into(),
602        ))
603    }
604}
605
606#[cfg(any(test, feature = "proptest"))]
607impl proptest::arbitrary::Arbitrary for ColumnName {
608    type Parameters = ();
609    type Strategy = BoxedStrategy<ColumnName>;
610
611    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
612        // Long column names are generally uninteresting, and can greatly
613        // increase the runtime for a test case, so bound the max length.
614        let mut weights = vec![(50, Just(1..8)), (20, Just(8..16))];
615        if std::env::var("PROPTEST_LARGE_DATA").is_ok() {
616            weights.extend([
617                (5, Just(16..128)),
618                (1, Just(128..1024)),
619                (1, Just(1024..4096)),
620            ]);
621        }
622        let name_length = Union::new_weighted(weights);
623
624        // Non-ASCII characters are also generally uninteresting and can make
625        // debugging harder.
626        let char_strat = Rc::new(Union::new_weighted(vec![
627            (50, proptest::char::range('A', 'z').boxed()),
628            (1, any::<char>().boxed()),
629        ]));
630
631        name_length
632            .prop_flat_map(move |length| proptest::collection::vec(Rc::clone(&char_strat), length))
633            .prop_map(|chars| ColumnName(chars.into_iter().collect::<Box<str>>()))
634            .no_shrink()
635            .boxed()
636    }
637}
638
639/// Default name of a column (when no other information is known).
640pub const UNKNOWN_COLUMN_NAME: &str = "?column?";
641
642/// The name of the column added by `WITH ORDINALITY`.
643pub static ORDINALITY_COL_NAME: &str = "ordinality";
644
645/// Stable index of a column in a [`RelationDesc`].
646#[derive(
647    Clone,
648    Copy,
649    Debug,
650    Eq,
651    PartialEq,
652    PartialOrd,
653    Ord,
654    Serialize,
655    Deserialize,
656    Hash
657)]
658pub struct ColumnIndex(usize);
659
660#[cfg(any(test, feature = "proptest"))]
661static_assertions::assert_not_impl_all!(ColumnIndex: Arbitrary);
662
663impl ColumnIndex {
664    /// Returns a stable identifier for this [`ColumnIndex`].
665    pub fn to_stable_name(&self) -> String {
666        self.0.to_string()
667    }
668
669    pub fn to_raw(&self) -> usize {
670        self.0
671    }
672
673    pub fn from_raw(val: usize) -> Self {
674        ColumnIndex(val)
675    }
676}
677
678/// The version a given column was added at.
679#[derive(
680    Clone,
681    Copy,
682    Debug,
683    Eq,
684    PartialEq,
685    PartialOrd,
686    Ord,
687    Serialize,
688    Deserialize,
689    Hash
690)]
691#[cfg_attr(any(test, feature = "proptest"), derive(Arbitrary))]
692pub struct RelationVersion(u64);
693
694impl RelationVersion {
695    /// Returns the "root" or "initial" version of a [`RelationDesc`].
696    pub fn root() -> Self {
697        RelationVersion(0)
698    }
699
700    /// Returns an instance of [`RelationVersion`] which is "one" higher than `self`.
701    pub fn bump(&self) -> Self {
702        let next_version = self
703            .0
704            .checked_add(1)
705            .expect("added more than u64::MAX columns?");
706        RelationVersion(next_version)
707    }
708
709    /// Consume a [`RelationVersion`] returning the raw value.
710    ///
711    /// Should __only__ be used for serialization.
712    pub fn into_raw(self) -> u64 {
713        self.0
714    }
715
716    /// Create a [`RelationVersion`] from a raw value.
717    ///
718    /// Should __only__ be used for serialization.
719    pub fn from_raw(val: u64) -> RelationVersion {
720        RelationVersion(val)
721    }
722}
723
724impl From<RelationVersion> for SchemaId {
725    fn from(value: RelationVersion) -> Self {
726        SchemaId(usize::cast_from(value.0))
727    }
728}
729
730impl fmt::Display for RelationVersion {
731    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
732        write!(f, "v{}", self.0)
733    }
734}
735
736impl RustType<ProtoRelationVersion> for RelationVersion {
737    fn into_proto(&self) -> ProtoRelationVersion {
738        ProtoRelationVersion { value: self.0 }
739    }
740
741    fn from_proto(proto: ProtoRelationVersion) -> Result<Self, TryFromProtoError> {
742        Ok(RelationVersion(proto.value))
743    }
744}
745
746/// Semantic type annotation for a column in a builtin catalog relation.
747///
748/// These are compile-time metadata used by the catalog ontology layer to
749/// describe the meaning of a column (e.g., that it contains a catalog item ID
750/// or a role ID). Possible values correspond to the entries in
751/// `SEMANTIC_TYPE_DEFS` in the `mz-catalog` crate.
752#[derive(
753    Clone,
754    Copy,
755    Debug,
756    PartialEq,
757    Eq,
758    PartialOrd,
759    Ord,
760    Hash,
761    serde::Serialize
762)]
763pub enum SemanticType {
764    CatalogItemId,
765    GlobalId,
766    ClusterId,
767    ReplicaId,
768    SchemaId,
769    DatabaseId,
770    RoleId,
771    NetworkPolicyId,
772    ShardId,
773    OID,
774    ObjectType,
775    ConnectionType,
776    SourceType,
777    MzTimestamp,
778    WallclockTimestamp,
779    ByteCount,
780    RecordCount,
781    CreditRate,
782    SqlDefinition,
783    RedactedSqlDefinition,
784}
785
786impl fmt::Display for SemanticType {
787    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
788        let s = match self {
789            SemanticType::CatalogItemId => "CatalogItemId",
790            SemanticType::GlobalId => "GlobalId",
791            SemanticType::ClusterId => "ClusterId",
792            SemanticType::ReplicaId => "ReplicaId",
793            SemanticType::SchemaId => "SchemaId",
794            SemanticType::DatabaseId => "DatabaseId",
795            SemanticType::RoleId => "RoleId",
796            SemanticType::NetworkPolicyId => "NetworkPolicyId",
797            SemanticType::ShardId => "ShardId",
798            SemanticType::OID => "OID",
799            SemanticType::ObjectType => "ObjectType",
800            SemanticType::ConnectionType => "ConnectionType",
801            SemanticType::SourceType => "SourceType",
802            SemanticType::MzTimestamp => "MzTimestamp",
803            SemanticType::WallclockTimestamp => "WallclockTimestamp",
804            SemanticType::ByteCount => "ByteCount",
805            SemanticType::RecordCount => "RecordCount",
806            SemanticType::CreditRate => "CreditRate",
807            SemanticType::SqlDefinition => "SqlDefinition",
808            SemanticType::RedactedSqlDefinition => "RedactedSqlDefinition",
809        };
810        f.write_str(s)
811    }
812}
813
814/// Metadata (other than type) for a column in a [`RelationDesc`].
815#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
816struct ColumnMetadata {
817    /// Name of the column.
818    name: ColumnName,
819    /// Index into a [`SqlRelationType`] for this column.
820    typ_idx: usize,
821    /// Version this column was added at.
822    added: RelationVersion,
823    /// Version this column was dropped at.
824    dropped: Option<RelationVersion>,
825}
826
827/// A description of the shape of a relation.
828///
829/// It bundles a [`SqlRelationType`] with `ColumnMetadata` for each column in
830/// the relation.
831///
832/// # Examples
833///
834/// A `RelationDesc`s is typically constructed via its builder API:
835///
836/// ```
837/// use mz_repr::{SqlColumnType, RelationDesc, SqlScalarType};
838///
839/// let desc = RelationDesc::builder()
840///     .with_column("id", SqlScalarType::Int64.nullable(false))
841///     .with_column("price", SqlScalarType::Float64.nullable(true))
842///     .finish();
843/// ```
844///
845/// In more complicated cases, like when constructing a `RelationDesc` in
846/// response to user input, it may be more convenient to construct a relation
847/// type first, and imbue it with column names to form a `RelationDesc` later:
848///
849/// ```
850/// use mz_repr::RelationDesc;
851///
852/// # fn plan_query(_: &str) -> mz_repr::SqlRelationType { mz_repr::SqlRelationType::new(vec![]) }
853/// let relation_type = plan_query("SELECT * FROM table");
854/// let names = (0..relation_type.arity()).map(|i| match i {
855///     0 => "first",
856///     1 => "second",
857///     _ => "unknown",
858/// });
859/// let desc = RelationDesc::new(relation_type, names);
860/// ```
861///
862/// Next to the [`SqlRelationType`] we maintain a map of `ColumnIndex` to
863/// `ColumnMetadata`, where [`ColumnIndex`] is a stable identifier for a
864/// column throughout the lifetime of the relation. This allows a
865/// [`RelationDesc`] to represent a projection over a version of itself.
866///
867/// ```
868/// use std::collections::BTreeSet;
869/// use mz_repr::{ColumnIndex, RelationDesc, SqlScalarType};
870///
871/// let desc = RelationDesc::builder()
872///     .with_column("name", SqlScalarType::String.nullable(false))
873///     .with_column("email", SqlScalarType::String.nullable(false))
874///     .finish();
875///
876/// // Project away the second column.
877/// let demands = BTreeSet::from([1]);
878/// let proj = desc.apply_demand(&demands);
879///
880/// // We projected away the first column.
881/// assert!(!proj.contains_index(&ColumnIndex::from_raw(0)));
882/// // But retained the second.
883/// assert!(proj.contains_index(&ColumnIndex::from_raw(1)));
884///
885/// // The underlying `SqlRelationType` also contains a single column.
886/// assert_eq!(proj.typ().arity(), 1);
887/// ```
888///
889/// To maintain this stable mapping and track the lifetime of a column (e.g.
890/// when adding or dropping a column) we use `ColumnMetadata`. It maintains
891/// the index in [`SqlRelationType`] that corresponds to a given column, and the
892/// version at which this column was added or dropped.
893///
894#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
895pub struct RelationDesc {
896    typ: SqlRelationType,
897    metadata: BTreeMap<ColumnIndex, ColumnMetadata>,
898}
899
900impl RustType<ProtoRelationDesc> for RelationDesc {
901    // NOTE: `ProtoRelationDesc` has no field for the `ColumnIndex` keys, only the values in
902    // `ColumnIndex` order, so a desc whose indexes are sparse (what
903    // `VersionedRelationDesc::at_version` returns once a column has been dropped, and what
904    // `RelationDesc::apply_demand` returns) comes back from `from_proto` renumbered to `0..n`.
905    // `ColumnIndex::to_stable_name` is the arrow field name that `RowColumnarEncoder` and
906    // `RowColumnarDecoder` agree on, so encoding a sparse desc as the schema of data written
907    // under the original indexes builds a decoder that looks up the wrong fields. Only dense
908    // descs may be handed to `Codec::encode_schema`.
909    fn into_proto(&self) -> ProtoRelationDesc {
910        let (names, metadata): (Vec<_>, Vec<_>) = self
911            .metadata
912            .values()
913            .map(|meta| {
914                let metadata = ProtoColumnMetadata {
915                    added: Some(meta.added.into_proto()),
916                    dropped: meta.dropped.map(|v| v.into_proto()),
917                };
918                (meta.name.into_proto(), metadata)
919            })
920            .unzip();
921
922        // `metadata` Migration Logic: We wrote some `ProtoRelationDesc`s into Persist before the
923        // metadata field was added. To make sure our serialization roundtrips the same as before
924        // we added the field, we omit `metadata` if all of the values are equal to the default.
925        //
926        // Note: This logic needs to exist approximately forever.
927        let is_all_default_metadata = metadata.iter().all(|meta| {
928            meta.added == Some(RelationVersion::root().into_proto()) && meta.dropped == None
929        });
930        let metadata = if is_all_default_metadata {
931            Vec::new()
932        } else {
933            metadata
934        };
935
936        ProtoRelationDesc {
937            typ: Some(self.typ.into_proto()),
938            names,
939            metadata,
940        }
941    }
942
943    fn from_proto(proto: ProtoRelationDesc) -> Result<Self, TryFromProtoError> {
944        let typ: SqlRelationType = proto.typ.into_rust_if_some("ProtoRelationDesc::typ")?;
945
946        // Reject shapes that `VersionedRelationDesc::validate` calls corruption. Nothing
947        // downstream catches them: they decode and re-encode cleanly, and only panic at first
948        // use, e.g. `iter()` indexing `typ.columns()[typ_idx]` out of bounds, or `into_iter()`
949        // tripping `zip_eq`. Both are reachable from untrusted proto bytes.
950        if proto.names.len() != typ.column_types.len() {
951            return Err(TryFromProtoError::InvalidFieldError(format!(
952                "ProtoRelationDesc: names ({}) and column_types ({}) length mismatch",
953                proto.names.len(),
954                typ.column_types.len()
955            )));
956        }
957        if let Some(key) = typ
958            .keys
959            .iter()
960            .flatten()
961            .find(|key| **key >= typ.column_types.len())
962        {
963            return Err(TryFromProtoError::InvalidFieldError(format!(
964                "ProtoRelationDesc: key index {key} out of bounds for {} columns",
965                typ.column_types.len()
966            )));
967        }
968
969        // `metadata` Migration Logic: We wrote some `ProtoRelationDesc`s into Persist before the
970        // metadata field was added. If the field doesn't exist we fill it in with default values,
971        // and when converting into_proto we omit these fields so the serialized bytes roundtrip.
972        //
973        // Note: This logic needs to exist approximately forever.
974        let proto_metadata: Box<dyn Iterator<Item = _>> = if proto.metadata.is_empty() {
975            let val = ProtoColumnMetadata {
976                added: Some(RelationVersion::root().into_proto()),
977                dropped: None,
978            };
979            Box::new(itertools::repeat_n(val, proto.names.len()))
980        } else {
981            // Reject mismatched lengths explicitly rather than panicking via
982            // `zip_eq` below, since this branch is reachable from untrusted
983            // proto bytes.
984            if proto.names.len() != proto.metadata.len() {
985                return Err(TryFromProtoError::InvalidFieldError(format!(
986                    "ProtoRelationDesc: names ({}) and metadata ({}) length mismatch",
987                    proto.names.len(),
988                    proto.metadata.len()
989                )));
990            }
991            Box::new(proto.metadata.into_iter())
992        };
993
994        let metadata = proto
995            .names
996            .into_iter()
997            .zip_eq(proto_metadata)
998            .enumerate()
999            .map(|(idx, (name, metadata))| {
1000                let meta = ColumnMetadata {
1001                    name: name.into_rust()?,
1002                    typ_idx: idx,
1003                    added: metadata.added.into_rust_if_some("ColumnMetadata::added")?,
1004                    dropped: metadata.dropped.into_rust()?,
1005                };
1006                Ok::<_, TryFromProtoError>((ColumnIndex(idx), meta))
1007            })
1008            .collect::<Result<_, _>>()?;
1009
1010        Ok(RelationDesc { typ, metadata })
1011    }
1012}
1013
1014impl RelationDesc {
1015    /// Returns a [`RelationDescBuilder`] that can be used to construct a [`RelationDesc`].
1016    pub fn builder() -> RelationDescBuilder {
1017        RelationDescBuilder::default()
1018    }
1019
1020    /// Constructs a new `RelationDesc` that represents the empty relation
1021    /// with no columns and no keys.
1022    pub fn empty() -> Self {
1023        RelationDesc {
1024            typ: SqlRelationType::empty(),
1025            metadata: BTreeMap::default(),
1026        }
1027    }
1028
1029    /// Check if the `RelationDesc` is empty.
1030    pub fn is_empty(&self) -> bool {
1031        self == &Self::empty()
1032    }
1033
1034    /// Returns the number of columns in this [`RelationDesc`].
1035    pub fn len(&self) -> usize {
1036        self.typ().column_types.len()
1037    }
1038
1039    /// Constructs a new `RelationDesc` from a `SqlRelationType` and an iterator
1040    /// over column names.
1041    ///
1042    /// # Panics
1043    ///
1044    /// Panics if the arity of the `SqlRelationType` is not equal to the number of
1045    /// items in `names`.
1046    pub fn new<I, N>(typ: SqlRelationType, names: I) -> Self
1047    where
1048        I: IntoIterator<Item = N>,
1049        N: Into<ColumnName>,
1050    {
1051        let metadata: BTreeMap<_, _> = names
1052            .into_iter()
1053            .enumerate()
1054            .map(|(idx, name)| {
1055                let col_idx = ColumnIndex(idx);
1056                let metadata = ColumnMetadata {
1057                    name: name.into(),
1058                    typ_idx: idx,
1059                    added: RelationVersion::root(),
1060                    dropped: None,
1061                };
1062                (col_idx, metadata)
1063            })
1064            .collect();
1065
1066        // TODO(parkmycar): Add better validation here.
1067        assert_eq!(typ.column_types.len(), metadata.len());
1068
1069        RelationDesc { typ, metadata }
1070    }
1071
1072    pub fn from_names_and_types<I, T, N>(iter: I) -> Self
1073    where
1074        I: IntoIterator<Item = (N, T)>,
1075        T: Into<SqlColumnType>,
1076        N: Into<ColumnName>,
1077    {
1078        let (names, types): (Vec<_>, Vec<_>) = iter.into_iter().unzip();
1079        let types = types.into_iter().map(Into::into).collect();
1080        let typ = SqlRelationType::new(types);
1081        Self::new(typ, names)
1082    }
1083
1084    /// Concatenates a `RelationDesc` onto the end of this `RelationDesc`.
1085    ///
1086    /// # Panics
1087    ///
1088    /// Panics if either `self` or `other` have columns that were added at a
1089    /// [`RelationVersion`] other than [`RelationVersion::root`] or if any
1090    /// columns were dropped.
1091    ///
1092    /// TODO(parkmycar): Move this method to [`RelationDescBuilder`].
1093    pub fn concat(mut self, other: Self) -> Self {
1094        let self_len = self.typ.column_types.len();
1095
1096        for (typ, (_col_idx, meta)) in other.typ.column_types.into_iter().zip_eq(other.metadata) {
1097            assert_eq!(meta.added, RelationVersion::root());
1098            assert_none!(meta.dropped);
1099
1100            let new_idx = self.typ.columns().len();
1101            let new_meta = ColumnMetadata {
1102                name: meta.name,
1103                typ_idx: new_idx,
1104                added: RelationVersion::root(),
1105                dropped: None,
1106            };
1107
1108            self.typ.column_types.push(typ);
1109            let prev = self.metadata.insert(ColumnIndex(new_idx), new_meta);
1110
1111            assert_eq!(self.metadata.len(), self.typ.columns().len());
1112            assert_none!(prev);
1113        }
1114
1115        for k in other.typ.keys {
1116            let k = k.into_iter().map(|idx| idx + self_len).collect();
1117            self = self.with_key(k);
1118        }
1119        self
1120    }
1121
1122    /// Adds a new key for the relation.
1123    pub fn with_key(mut self, indices: Vec<usize>) -> Self {
1124        self.typ = self.typ.with_key(indices);
1125        self
1126    }
1127
1128    /// Drops all existing keys.
1129    pub fn without_keys(mut self) -> Self {
1130        self.typ.keys.clear();
1131        self
1132    }
1133
1134    /// Builds a new relation description with the column names replaced with
1135    /// new names.
1136    ///
1137    /// # Panics
1138    ///
1139    /// Panics if the arity of the relation type does not match the number of
1140    /// items in `names`.
1141    pub fn with_names<I, N>(self, names: I) -> Self
1142    where
1143        I: IntoIterator<Item = N>,
1144        N: Into<ColumnName>,
1145    {
1146        Self::new(self.typ, names)
1147    }
1148
1149    /// Computes the number of columns in the relation.
1150    pub fn arity(&self) -> usize {
1151        self.typ.arity()
1152    }
1153
1154    /// Returns the relation type underlying this relation description.
1155    pub fn typ(&self) -> &SqlRelationType {
1156        &self.typ
1157    }
1158
1159    /// Returns the owned relation type underlying this relation description.
1160    pub fn into_typ(self) -> SqlRelationType {
1161        self.typ
1162    }
1163
1164    /// Returns an iterator over the columns in this relation.
1165    pub fn iter(&self) -> impl Iterator<Item = (&ColumnName, &SqlColumnType)> {
1166        self.metadata.values().map(|meta| {
1167            let typ = &self.typ.columns()[meta.typ_idx];
1168            (&meta.name, typ)
1169        })
1170    }
1171
1172    /// Returns an iterator over the types of the columns in this relation.
1173    pub fn iter_types(&self) -> impl Iterator<Item = &SqlColumnType> {
1174        self.typ.column_types.iter()
1175    }
1176
1177    /// Returns an iterator over the names of the columns in this relation.
1178    pub fn iter_names(&self) -> impl Iterator<Item = &ColumnName> {
1179        self.metadata.values().map(|meta| &meta.name)
1180    }
1181
1182    /// Returns an iterator over the columns in this relation, with all their metadata.
1183    pub fn iter_all(&self) -> impl Iterator<Item = (&ColumnIndex, &ColumnName, &SqlColumnType)> {
1184        self.metadata.iter().map(|(col_idx, metadata)| {
1185            let col_typ = &self.typ.columns()[metadata.typ_idx];
1186            (col_idx, &metadata.name, col_typ)
1187        })
1188    }
1189
1190    /// Returns an iterator over the names of the columns in this relation that are "similar" to
1191    /// the provided `name`.
1192    pub fn iter_similar_names<'a>(
1193        &'a self,
1194        name: &'a ColumnName,
1195    ) -> impl Iterator<Item = &'a ColumnName> {
1196        self.iter_names().filter(|n| n.is_similar(name))
1197    }
1198
1199    /// Returns whether this [`RelationDesc`] contains a column at the specified index.
1200    pub fn contains_index(&self, idx: &ColumnIndex) -> bool {
1201        self.metadata.contains_key(idx)
1202    }
1203
1204    /// Finds a column by name.
1205    ///
1206    /// Returns the index and type of the column named `name`. If no column with
1207    /// the specified name exists, returns `None`. If multiple columns have the
1208    /// specified name, the leftmost column is returned.
1209    pub fn get_by_name(&self, name: &ColumnName) -> Option<(usize, &SqlColumnType)> {
1210        self.iter_names()
1211            .position(|n| n == name)
1212            .map(|i| (i, &self.typ.column_types[i]))
1213    }
1214
1215    /// Gets the name of the `i`th column.
1216    ///
1217    /// # Panics
1218    ///
1219    /// Panics if `i` is not a valid column index.
1220    ///
1221    /// TODO(parkmycar): Migrate all uses of this to [`RelationDesc::get_name_idx`].
1222    pub fn get_name(&self, i: usize) -> &ColumnName {
1223        // TODO(parkmycar): Refactor this to use `ColumnIndex`.
1224        self.get_name_idx(&ColumnIndex(i))
1225    }
1226
1227    /// Gets the name of the column at `idx`.
1228    ///
1229    /// # Panics
1230    ///
1231    /// Panics if no column exists at `idx`.
1232    pub fn get_name_idx(&self, idx: &ColumnIndex) -> &ColumnName {
1233        &self.metadata.get(idx).expect("should exist").name
1234    }
1235
1236    /// Mutably gets the name of the `i`th column.
1237    ///
1238    /// # Panics
1239    ///
1240    /// Panics if `i` is not a valid column index.
1241    pub fn get_name_mut(&mut self, i: usize) -> &mut ColumnName {
1242        // TODO(parkmycar): Refactor this to use `ColumnIndex`.
1243        &mut self
1244            .metadata
1245            .get_mut(&ColumnIndex(i))
1246            .expect("should exist")
1247            .name
1248    }
1249
1250    /// Gets the [`SqlColumnType`] of the column at `idx`.
1251    ///
1252    /// # Panics
1253    ///
1254    /// Panics if no column exists at `idx`.
1255    pub fn get_type(&self, idx: &ColumnIndex) -> &SqlColumnType {
1256        let typ_idx = self.metadata.get(idx).expect("should exist").typ_idx;
1257        &self.typ.column_types[typ_idx]
1258    }
1259
1260    /// Gets the name of the `i`th column if that column name is unambiguous.
1261    ///
1262    /// If at least one other column has the same name as the `i`th column,
1263    /// returns `None`. If the `i`th column has no name, returns `None`.
1264    ///
1265    /// # Panics
1266    ///
1267    /// Panics if `i` is not a valid column index.
1268    pub fn get_unambiguous_name(&self, i: usize) -> Option<&ColumnName> {
1269        let name = self.get_name(i);
1270        if self.iter_names().filter(|n| *n == name).count() == 1 {
1271            Some(name)
1272        } else {
1273            None
1274        }
1275    }
1276
1277    /// Verifies that `d` meets all of the constraints for the `i`th column of `self`.
1278    ///
1279    /// n.b. The only constraint MZ currently supports in NOT NULL, but this
1280    /// structure will be simple to extend.
1281    pub fn constraints_met(&self, i: usize, d: &Datum) -> Result<(), NotNullViolation> {
1282        let name = self.get_name(i);
1283        let typ = &self.typ.column_types[i];
1284        if d == &Datum::Null && !typ.nullable {
1285            Err(NotNullViolation(name.clone()))
1286        } else {
1287            Ok(())
1288        }
1289    }
1290
1291    /// Computes the differences between two [`RelationDesc`]s.
1292    ///
1293    /// Returns a rich diff describing which columns differ, and in what way.
1294    ///
1295    /// # Panics
1296    ///
1297    /// Panics if either `self` or `other` have columns that were added at a
1298    /// [`RelationVersion`] other than [`RelationVersion::root`] or if any
1299    /// columns were dropped.
1300    ///
1301    /// This simplifies things by allowing us to assume that `ColumnIndex`es are
1302    /// dense and that they match the indexes of `typ.columns()`. Without this
1303    /// we would, e.g., struggle comparing keys as those are in terms of
1304    /// `typ.columns()` indexes.
1305    pub fn diff(&self, other: &RelationDesc) -> RelationDescDiff {
1306        assert_eq!(self.metadata.len(), self.typ.columns().len());
1307        assert_eq!(other.metadata.len(), other.typ.columns().len());
1308        for (idx, meta) in self.metadata.iter().chain(other.metadata.iter()) {
1309            assert_eq!(meta.typ_idx, idx.0);
1310            assert_eq!(meta.added, RelationVersion::root());
1311            assert_none!(meta.dropped);
1312        }
1313
1314        let mut column_diffs = BTreeMap::new();
1315        let mut key_diff = None;
1316
1317        let left_arity = self.arity();
1318        let right_arity = other.arity();
1319        let common_arity = std::cmp::min(left_arity, right_arity);
1320
1321        for idx in 0..common_arity {
1322            let left_name = self.get_name(idx);
1323            let right_name = other.get_name(idx);
1324            let left_type = &self.typ.column_types[idx];
1325            let right_type = &other.typ.column_types[idx];
1326
1327            if left_name != right_name {
1328                let diff = ColumnDiff::NameMismatch {
1329                    left: left_name.clone(),
1330                    right: right_name.clone(),
1331                };
1332                column_diffs.insert(idx, diff);
1333            } else if left_type.scalar_type != right_type.scalar_type {
1334                let diff = ColumnDiff::TypeMismatch {
1335                    name: left_name.clone(),
1336                    left: left_type.scalar_type.clone(),
1337                    right: right_type.scalar_type.clone(),
1338                };
1339                column_diffs.insert(idx, diff);
1340            } else if left_type.nullable != right_type.nullable {
1341                let diff = ColumnDiff::NullabilityMismatch {
1342                    name: left_name.clone(),
1343                    left: left_type.nullable,
1344                    right: right_type.nullable,
1345                };
1346                column_diffs.insert(idx, diff);
1347            }
1348        }
1349
1350        for idx in common_arity..left_arity {
1351            let diff = ColumnDiff::Missing {
1352                name: self.get_name(idx).clone(),
1353            };
1354            column_diffs.insert(idx, diff);
1355        }
1356
1357        for idx in common_arity..right_arity {
1358            let diff = ColumnDiff::Extra {
1359                name: other.get_name(idx).clone(),
1360            };
1361            column_diffs.insert(idx, diff);
1362        }
1363
1364        let left_keys: BTreeSet<_> = self.typ.keys.iter().collect();
1365        let right_keys: BTreeSet<_> = other.typ.keys.iter().collect();
1366        if left_keys != right_keys {
1367            let column_names = |desc: &RelationDesc, keys: BTreeSet<&Vec<usize>>| {
1368                keys.iter()
1369                    .map(|key| key.iter().map(|&idx| desc.get_name(idx).clone()).collect())
1370                    .collect()
1371            };
1372            key_diff = Some(KeyDiff {
1373                left: column_names(self, left_keys),
1374                right: column_names(other, right_keys),
1375            });
1376        }
1377
1378        RelationDescDiff {
1379            column_diffs,
1380            key_diff,
1381        }
1382    }
1383
1384    /// Creates a new [`RelationDesc`] retaining only the columns specified in `demands`.
1385    pub fn apply_demand(&self, demands: &BTreeSet<usize>) -> RelationDesc {
1386        // This filters `metadata` by raw ColumnIndex but `typ` by position,
1387        // which only agree when the desc is dense. Every desc constructible
1388        // today is (schema history is add-only), but a dropped column would
1389        // desync the two and silently attach types, statistics, and filter
1390        // specs to the wrong columns downstream.
1391        debug_assert!(
1392            self.metadata
1393                .iter()
1394                .enumerate()
1395                .all(|(pos, (idx, meta))| idx.0 == pos && meta.typ_idx == pos),
1396            "apply_demand requires a dense RelationDesc (ColumnIndex == typ_idx): {:?}",
1397            self.metadata,
1398        );
1399        let mut new_desc = self.clone();
1400
1401        // Update ColumnMetadata.
1402        let mut removed = 0;
1403        new_desc.metadata.retain(|idx, metadata| {
1404            let retain = demands.contains(&idx.0);
1405            if !retain {
1406                removed += 1;
1407            } else {
1408                metadata.typ_idx -= removed;
1409            }
1410            retain
1411        });
1412
1413        // Update SqlColumnType.
1414        let mut idx = 0;
1415        new_desc.typ.column_types.retain(|_| {
1416            let keep = demands.contains(&idx);
1417            idx += 1;
1418            keep
1419        });
1420
1421        new_desc
1422    }
1423}
1424
1425#[cfg(any(test, feature = "proptest"))]
1426impl Arbitrary for RelationDesc {
1427    type Parameters = ();
1428    type Strategy = BoxedStrategy<RelationDesc>;
1429
1430    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
1431        let mut weights = vec![(100, Just(0..4)), (50, Just(4..8)), (25, Just(8..16))];
1432        if std::env::var("PROPTEST_LARGE_DATA").is_ok() {
1433            weights.extend([
1434                (12, Just(16..32)),
1435                (6, Just(32..64)),
1436                (3, Just(64..128)),
1437                (1, Just(128..256)),
1438            ]);
1439        }
1440        let num_columns = Union::new_weighted(weights);
1441
1442        num_columns.prop_flat_map(arb_relation_desc).boxed()
1443    }
1444}
1445
1446/// Returns a [`Strategy`] that generates an arbitrary [`RelationDesc`] with a number columns
1447/// within the range provided.
1448#[cfg(any(test, feature = "proptest"))]
1449pub fn arb_relation_desc(num_cols: std::ops::Range<usize>) -> impl Strategy<Value = RelationDesc> {
1450    proptest::collection::btree_map(any::<ColumnName>(), any::<SqlColumnType>(), num_cols)
1451        .prop_map(RelationDesc::from_names_and_types)
1452}
1453
1454/// Returns a [`Strategy`] that generates a projection of the provided [`RelationDesc`].
1455#[cfg(any(test, feature = "proptest"))]
1456pub fn arb_relation_desc_projection(desc: RelationDesc) -> impl Strategy<Value = RelationDesc> {
1457    let mask: Vec<_> = (0..desc.len()).map(|_| any::<bool>()).collect();
1458    mask.prop_map(move |mask| {
1459        let demands: BTreeSet<_> = mask
1460            .into_iter()
1461            .enumerate()
1462            .filter_map(|(idx, keep)| keep.then_some(idx))
1463            .collect();
1464        desc.apply_demand(&demands)
1465    })
1466}
1467
1468impl IntoIterator for RelationDesc {
1469    type Item = (ColumnName, SqlColumnType);
1470    type IntoIter = Box<dyn Iterator<Item = (ColumnName, SqlColumnType)>>;
1471
1472    fn into_iter(self) -> Self::IntoIter {
1473        let iter = self
1474            .metadata
1475            .into_values()
1476            .zip_eq(self.typ.column_types)
1477            .map(|(meta, typ)| (meta.name, typ));
1478        Box::new(iter)
1479    }
1480}
1481
1482/// Returns a [`Strategy`] that yields arbitrary [`Row`]s for the provided [`RelationDesc`].
1483#[cfg(any(test, feature = "proptest"))]
1484pub fn arb_row_for_relation(desc: &RelationDesc) -> impl Strategy<Value = Row> + use<> {
1485    let datums: Vec<_> = desc
1486        .typ()
1487        .columns()
1488        .iter()
1489        .cloned()
1490        .map(arb_datum_for_column)
1491        .collect();
1492    datums.prop_map(|x| Row::pack(x.iter().map(Datum::from)))
1493}
1494
1495/// Expression violated not-null constraint on named column
1496#[derive(Debug, PartialEq, Eq)]
1497pub struct NotNullViolation(pub ColumnName);
1498
1499impl fmt::Display for NotNullViolation {
1500    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1501        write!(
1502            f,
1503            "null value in column {} violates not-null constraint",
1504            self.0.quoted()
1505        )
1506    }
1507}
1508
1509/// The result of comparing two [`RelationDesc`]s.
1510#[derive(Debug, Clone, PartialEq, Eq)]
1511pub struct RelationDescDiff {
1512    /// Column differences, keyed by column index.
1513    pub column_diffs: BTreeMap<usize, ColumnDiff>,
1514    /// Key differences, if any.
1515    pub key_diff: Option<KeyDiff>,
1516}
1517
1518impl RelationDescDiff {
1519    /// Returns whether the diff contains any differences.
1520    pub fn is_empty(&self) -> bool {
1521        self.column_diffs.is_empty() && self.key_diff.is_none()
1522    }
1523}
1524
1525/// A difference in a column between two [`RelationDesc`]s.
1526#[derive(Debug, Clone, PartialEq, Eq)]
1527pub enum ColumnDiff {
1528    /// Column exists only in the left relation.
1529    Missing { name: ColumnName },
1530    /// Column exists only in the right relation.
1531    Extra { name: ColumnName },
1532    /// Columns have different types.
1533    TypeMismatch {
1534        name: ColumnName,
1535        left: SqlScalarType,
1536        right: SqlScalarType,
1537    },
1538    /// Columns have different nullability.
1539    NullabilityMismatch {
1540        name: ColumnName,
1541        left: bool,
1542        right: bool,
1543    },
1544    /// Columns have different names.
1545    NameMismatch { left: ColumnName, right: ColumnName },
1546}
1547
1548/// A difference in the keys of two [`RelationDesc`]s.
1549#[derive(Debug, Clone, PartialEq, Eq)]
1550pub struct KeyDiff {
1551    /// Keys of the left relation.
1552    pub left: BTreeSet<Vec<ColumnName>>,
1553    /// Keys of the right relation.
1554    pub right: BTreeSet<Vec<ColumnName>>,
1555}
1556
1557/// A builder for a [`RelationDesc`].
1558#[derive(Clone, Default, Debug, PartialEq, Eq)]
1559pub struct RelationDescBuilder {
1560    /// Columns of the relation.
1561    columns: Vec<(ColumnName, SqlColumnType)>,
1562    /// Sets of indices that are "keys" for the collection.
1563    keys: Vec<Vec<usize>>,
1564}
1565
1566impl RelationDescBuilder {
1567    /// Appends a column with the specified name and type.
1568    pub fn with_column<N: Into<ColumnName>>(
1569        mut self,
1570        name: N,
1571        ty: SqlColumnType,
1572    ) -> RelationDescBuilder {
1573        let name = name.into();
1574        self.columns.push((name, ty));
1575        self
1576    }
1577
1578    /// Appends the provided columns to the builder.
1579    pub fn with_columns<I, T, N>(mut self, iter: I) -> Self
1580    where
1581        I: IntoIterator<Item = (N, T)>,
1582        T: Into<SqlColumnType>,
1583        N: Into<ColumnName>,
1584    {
1585        self.columns
1586            .extend(iter.into_iter().map(|(name, ty)| (name.into(), ty.into())));
1587        self
1588    }
1589
1590    /// Adds a new key for the relation.
1591    pub fn with_key(mut self, mut indices: Vec<usize>) -> RelationDescBuilder {
1592        indices.sort_unstable();
1593        if !self.keys.contains(&indices) {
1594            self.keys.push(indices);
1595        }
1596        self
1597    }
1598
1599    /// Removes all previously inserted keys.
1600    pub fn without_keys(mut self) -> RelationDescBuilder {
1601        self.keys.clear();
1602        assert_eq!(self.keys.len(), 0);
1603        self
1604    }
1605
1606    /// Concatenates a [`RelationDescBuilder`] onto the end of this [`RelationDescBuilder`].
1607    pub fn concat(mut self, other: Self) -> Self {
1608        let self_len = self.columns.len();
1609
1610        self.columns.extend(other.columns);
1611        for k in other.keys {
1612            let k = k.into_iter().map(|idx| idx + self_len).collect();
1613            self = self.with_key(k);
1614        }
1615
1616        self
1617    }
1618
1619    /// Finish the builder, returning a [`RelationDesc`].
1620    pub fn finish(self) -> RelationDesc {
1621        let mut desc = RelationDesc::from_names_and_types(self.columns);
1622        desc.typ = desc.typ.with_keys(self.keys);
1623        desc
1624    }
1625}
1626
1627/// Describes a [`RelationDesc`] at a specific version of a [`VersionedRelationDesc`].
1628#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize)]
1629pub enum RelationVersionSelector {
1630    Specific(RelationVersion),
1631    Latest,
1632}
1633
1634impl RelationVersionSelector {
1635    pub fn specific(version: u64) -> Self {
1636        RelationVersionSelector::Specific(RelationVersion(version))
1637    }
1638}
1639
1640/// A wrapper around [`RelationDesc`] that provides an interface for adding
1641/// columns and generating new versions.
1642///
1643/// TODO(parkmycar): Using an immutable data structure for RelationDesc would
1644/// be great.
1645#[derive(Debug, Clone, Serialize)]
1646pub struct VersionedRelationDesc {
1647    inner: RelationDesc,
1648}
1649
1650impl VersionedRelationDesc {
1651    pub fn new(inner: RelationDesc) -> Self {
1652        VersionedRelationDesc { inner }
1653    }
1654
1655    /// Adds a new column to this [`RelationDesc`], creating a new version of the [`RelationDesc`].
1656    ///
1657    /// # Panics
1658    ///
1659    /// * Panics if a column with `name` already exists that hasn't been dropped.
1660    ///
1661    /// Note: For building a [`RelationDesc`] see [`RelationDescBuilder::with_column`].
1662    #[must_use]
1663    pub fn add_column<N, T>(&mut self, name: N, typ: T) -> RelationVersion
1664    where
1665        N: Into<ColumnName>,
1666        T: Into<SqlColumnType>,
1667    {
1668        let latest_version = self.latest_version();
1669        let new_version = latest_version.bump();
1670
1671        let name = name.into();
1672        let existing = self
1673            .inner
1674            .metadata
1675            .iter()
1676            .find(|(_, meta)| meta.name == name && meta.dropped.is_none());
1677        if let Some(existing) = existing {
1678            panic!("column named '{name}' already exists! {existing:?}");
1679        }
1680
1681        let next_idx = self.inner.metadata.len();
1682        let col_meta = ColumnMetadata {
1683            name,
1684            typ_idx: next_idx,
1685            added: new_version,
1686            dropped: None,
1687        };
1688
1689        self.inner.typ.column_types.push(typ.into());
1690        let prev = self.inner.metadata.insert(ColumnIndex(next_idx), col_meta);
1691
1692        assert_none!(prev, "column index overlap!");
1693        self.validate();
1694
1695        new_version
1696    }
1697
1698    /// Drops the column `name` from this [`RelationDesc`]. If there are multiple columns with
1699    /// `name` drops the left-most one that hasn't already been dropped.
1700    ///
1701    /// TODO(parkmycar): Add handling for dropping a column that is currently used as a key.
1702    ///
1703    /// # Panics
1704    ///
1705    /// Panics if a column with `name` does not exist or the dropped column was used as a key.
1706    #[must_use]
1707    pub fn drop_column<N>(&mut self, name: N) -> RelationVersion
1708    where
1709        N: Into<ColumnName>,
1710    {
1711        let name = name.into();
1712        let latest_version = self.latest_version();
1713        let new_version = latest_version.bump();
1714
1715        let col = self
1716            .inner
1717            .metadata
1718            .values_mut()
1719            .find(|meta| meta.name == name && meta.dropped.is_none())
1720            .expect("column to exist");
1721
1722        // Make sure the column hadn't been previously dropped.
1723        assert_none!(col.dropped, "column was already dropped");
1724        col.dropped = Some(new_version);
1725
1726        // Make sure the column isn't being used as a key.
1727        let dropped_key = self
1728            .inner
1729            .typ
1730            .keys
1731            .iter()
1732            .any(|keys| keys.contains(&col.typ_idx));
1733        assert!(!dropped_key, "column being dropped was used as a key");
1734
1735        self.validate();
1736        new_version
1737    }
1738
1739    /// Returns the [`RelationDesc`] at the latest version.
1740    pub fn latest(&self) -> RelationDesc {
1741        self.inner.clone()
1742    }
1743
1744    /// Returns this [`RelationDesc`] at the specified version.
1745    pub fn at_version(&self, version: RelationVersionSelector) -> RelationDesc {
1746        // Get all of the changes from the start, up to whatever version was requested.
1747        let up_to_version = match version {
1748            RelationVersionSelector::Latest => RelationVersion(u64::MAX),
1749            RelationVersionSelector::Specific(v) => v,
1750        };
1751
1752        let valid_columns = self.inner.metadata.iter().filter(|(_col_idx, meta)| {
1753            let added = meta.added <= up_to_version;
1754            let dropped = meta
1755                .dropped
1756                .map(|dropped_at| up_to_version >= dropped_at)
1757                .unwrap_or(false);
1758
1759            added && !dropped
1760        });
1761
1762        let mut column_types = Vec::new();
1763        let mut column_metas = BTreeMap::new();
1764
1765        // N.B. At this point we need to be careful because col_idx might not
1766        // equal typ_idx.
1767        //
1768        // For example, consider columns "a", "b", and "c" with indexes 0, 1,
1769        // and 2. If we drop column "b" then we'll have "a" and "c" with column
1770        // indexes 0 and 2, but their indices in SqlRelationType will be 0 and 1.
1771        for (col_idx, meta) in valid_columns {
1772            let new_meta = ColumnMetadata {
1773                name: meta.name.clone(),
1774                typ_idx: column_types.len(),
1775                added: meta.added.clone(),
1776                dropped: meta.dropped.clone(),
1777            };
1778            column_types.push(self.inner.typ.columns()[meta.typ_idx].clone());
1779            column_metas.insert(*col_idx, new_meta);
1780        }
1781
1782        // Remap keys in case a column with an index less than that of a key was
1783        // dropped.
1784        //
1785        // For example, consider columns "a", "b", and "c" where "a" and "c" are
1786        // keys and "b" was dropped.
1787        let keys = self
1788            .inner
1789            .typ
1790            .keys
1791            .iter()
1792            .map(|keys| {
1793                keys.iter()
1794                    .map(|key_idx| {
1795                        let metadata = column_metas
1796                            .get(&ColumnIndex(*key_idx))
1797                            .expect("found key for column that doesn't exist");
1798                        metadata.typ_idx
1799                    })
1800                    .collect()
1801            })
1802            .collect();
1803
1804        let relation_type = SqlRelationType { column_types, keys };
1805
1806        RelationDesc {
1807            typ: relation_type,
1808            metadata: column_metas,
1809        }
1810    }
1811
1812    pub fn latest_version(&self) -> RelationVersion {
1813        self.inner
1814            .metadata
1815            .values()
1816            // N.B. Dropped is always greater than added.
1817            .map(|meta| meta.dropped.unwrap_or(meta.added))
1818            .max()
1819            // If there aren't any columns we're implicitly the root version.
1820            .unwrap_or_else(RelationVersion::root)
1821    }
1822
1823    /// Validates internal contraints of the [`RelationDesc`] are correct.
1824    ///
1825    /// # Panics
1826    ///
1827    /// Panics if a constraint is not satisfied.
1828    fn validate(&self) {
1829        fn validate_inner(desc: &RelationDesc) -> Result<(), anyhow::Error> {
1830            if desc.typ.column_types.len() != desc.metadata.len() {
1831                anyhow::bail!("mismatch between number of types and metadatas");
1832            }
1833
1834            for (col_idx, meta) in &desc.metadata {
1835                if col_idx.0 > desc.metadata.len() {
1836                    anyhow::bail!("column index out of bounds");
1837                }
1838                if meta.added >= meta.dropped.unwrap_or(RelationVersion(u64::MAX)) {
1839                    anyhow::bail!("column was added after it was dropped?");
1840                }
1841                if desc.typ().columns().get(meta.typ_idx).is_none() {
1842                    anyhow::bail!("typ_idx incorrect");
1843                }
1844            }
1845
1846            for keys in &desc.typ.keys {
1847                for key in keys {
1848                    if *key >= desc.typ.column_types.len() {
1849                        anyhow::bail!("key index was out of bounds!");
1850                    }
1851                }
1852            }
1853
1854            // Every version past the root is stamped exactly once, by the add or the
1855            // drop that created it, so a column that was added and later dropped
1856            // accounts for two of them and both have to be counted. Collapsing to
1857            // `dropped.unwrap_or(added)` would lose the add's version, and with it any
1858            // desc where a non-root column was later dropped.
1859            let versions = desc
1860                .metadata
1861                .values()
1862                .flat_map(|meta| [Some(meta.added), meta.dropped])
1863                .flatten()
1864                // The root version is the one version many columns can share.
1865                .filter(|version| *version != RelationVersion::root());
1866            let mut max = 0;
1867            let mut sum = 0;
1868            for version in versions {
1869                max = std::cmp::max(max, version.0);
1870                sum += version.0;
1871            }
1872
1873            // Other than RelationVersion(0), we should never have duplicate
1874            // versions and they should always increase by 1. In other words, the
1875            // sum of all RelationVersions should be the sum of [0, max].
1876            //
1877            // N.B. n * (n + 1) / 2 = sum of [0, n]
1878            //
1879            // While I normally don't like tricks like this, it allows us to
1880            // validate that our column versions are correct in O(n) time and
1881            // without allocations.
1882            if sum != (max * (max + 1) / 2) {
1883                anyhow::bail!("there is a duplicate or missing relation version");
1884            }
1885
1886            Ok(())
1887        }
1888
1889        assert_ok!(validate_inner(&self.inner), "validate failed! {self:?}");
1890    }
1891}
1892
1893/// Diffs that can be generated proptest and applied to a [`RelationDesc`] to
1894/// exercise schema migrations.
1895#[derive(Debug)]
1896#[cfg(any(test, feature = "proptest"))]
1897pub enum PropRelationDescDiff {
1898    AddColumn {
1899        name: ColumnName,
1900        typ: SqlColumnType,
1901    },
1902    DropColumn {
1903        name: ColumnName,
1904    },
1905    ToggleNullability {
1906        name: ColumnName,
1907    },
1908    ChangeType {
1909        name: ColumnName,
1910        typ: SqlColumnType,
1911    },
1912}
1913
1914#[cfg(any(test, feature = "proptest"))]
1915impl PropRelationDescDiff {
1916    pub fn apply(self, desc: &mut RelationDesc) {
1917        match self {
1918            PropRelationDescDiff::AddColumn { name, typ } => {
1919                let new_idx = desc.metadata.len();
1920                let meta = ColumnMetadata {
1921                    name,
1922                    typ_idx: new_idx,
1923                    added: RelationVersion(0),
1924                    dropped: None,
1925                };
1926                let prev = desc.metadata.insert(ColumnIndex(new_idx), meta);
1927                desc.typ.column_types.push(typ);
1928
1929                assert_none!(prev);
1930                assert_eq!(desc.metadata.len(), desc.typ.column_types.len());
1931            }
1932            PropRelationDescDiff::DropColumn { name } => {
1933                let next_version = desc
1934                    .metadata
1935                    .values()
1936                    .map(|meta| meta.dropped.unwrap_or(meta.added))
1937                    .max()
1938                    .unwrap_or_else(RelationVersion::root)
1939                    .bump();
1940                let Some(metadata) = desc.metadata.values_mut().find(|meta| meta.name == name)
1941                else {
1942                    return;
1943                };
1944                if metadata.dropped.is_none() {
1945                    metadata.dropped = Some(next_version);
1946                }
1947            }
1948            PropRelationDescDiff::ToggleNullability { name } => {
1949                let Some((pos, _)) = desc.get_by_name(&name) else {
1950                    return;
1951                };
1952                let col_type = desc
1953                    .typ
1954                    .column_types
1955                    .get_mut(pos)
1956                    .expect("ColumnNames and SqlColumnTypes out of sync!");
1957                col_type.nullable = !col_type.nullable;
1958            }
1959            PropRelationDescDiff::ChangeType { name, typ } => {
1960                let Some((pos, _)) = desc.get_by_name(&name) else {
1961                    return;
1962                };
1963                let col_type = desc
1964                    .typ
1965                    .column_types
1966                    .get_mut(pos)
1967                    .expect("ColumnNames and SqlColumnTypes out of sync!");
1968                *col_type = typ;
1969            }
1970        }
1971    }
1972}
1973
1974/// Generates a set of [`PropRelationDescDiff`]s based on some source [`RelationDesc`].
1975#[cfg(any(test, feature = "proptest"))]
1976pub fn arb_relation_desc_diff(
1977    source: &RelationDesc,
1978) -> impl Strategy<Value = Vec<PropRelationDescDiff>> + use<> {
1979    let source = Rc::new(source.clone());
1980    let num_source_columns = source.typ.columns().len();
1981
1982    let num_add_columns = Union::new_weighted(vec![(100, Just(0..8)), (1, Just(8..64))]);
1983    let add_columns_strat = num_add_columns
1984        .prop_flat_map(|num_columns| {
1985            proptest::collection::vec((any::<ColumnName>(), any::<SqlColumnType>()), num_columns)
1986        })
1987        .prop_map(|cols| {
1988            cols.into_iter()
1989                .map(|(name, typ)| PropRelationDescDiff::AddColumn { name, typ })
1990                .collect::<Vec<_>>()
1991        });
1992
1993    // If the source RelationDesc is empty there is nothing else to do.
1994    if num_source_columns == 0 {
1995        return add_columns_strat.boxed();
1996    }
1997
1998    let source_ = Rc::clone(&source);
1999    let drop_columns_strat = (0..num_source_columns).prop_perturb(move |num_columns, mut rng| {
2000        let mut set = BTreeSet::default();
2001        for _ in 0..num_columns {
2002            let col_idx = rng.random_range(0..num_source_columns);
2003            set.insert(source_.get_name(col_idx).clone());
2004        }
2005        set.into_iter()
2006            .map(|name| PropRelationDescDiff::DropColumn { name })
2007            .collect::<Vec<_>>()
2008    });
2009
2010    let source_ = Rc::clone(&source);
2011    let toggle_nullability_strat =
2012        (0..num_source_columns).prop_perturb(move |num_columns, mut rng| {
2013            let mut set = BTreeSet::default();
2014            for _ in 0..num_columns {
2015                let col_idx = rng.random_range(0..num_source_columns);
2016                set.insert(source_.get_name(col_idx).clone());
2017            }
2018            set.into_iter()
2019                .map(|name| PropRelationDescDiff::ToggleNullability { name })
2020                .collect::<Vec<_>>()
2021        });
2022
2023    let source_ = Rc::clone(&source);
2024    let change_type_strat = (0..num_source_columns)
2025        .prop_perturb(move |num_columns, mut rng| {
2026            let mut set = BTreeSet::default();
2027            for _ in 0..num_columns {
2028                let col_idx = rng.random_range(0..num_source_columns);
2029                set.insert(source_.get_name(col_idx).clone());
2030            }
2031            set
2032        })
2033        .prop_flat_map(|cols| {
2034            proptest::collection::vec(any::<SqlColumnType>(), cols.len())
2035                .prop_map(move |types| (cols.clone(), types))
2036        })
2037        .prop_map(|(cols, types)| {
2038            cols.into_iter()
2039                .zip_eq(types)
2040                .map(|(name, typ)| PropRelationDescDiff::ChangeType { name, typ })
2041                .collect::<Vec<_>>()
2042        });
2043
2044    (
2045        add_columns_strat,
2046        drop_columns_strat,
2047        toggle_nullability_strat,
2048        change_type_strat,
2049    )
2050        .prop_map(|(adds, drops, toggles, changes)| {
2051            adds.into_iter()
2052                .chain(drops)
2053                .chain(toggles)
2054                .chain(changes)
2055                .collect::<Vec<_>>()
2056        })
2057        .prop_shuffle()
2058        .boxed()
2059}
2060
2061#[cfg(test)]
2062mod tests {
2063    use super::*;
2064    use prost::Message;
2065
2066    /// `apply_demand`, and the stats and filter-spec plumbing downstream of
2067    /// it, require dense descs. A desc with a dropped column must trip the
2068    /// assertion rather than silently misattach columns.
2069    #[mz_ore::test]
2070    #[should_panic(expected = "dense RelationDesc")]
2071    fn apply_demand_rejects_non_dense_desc() {
2072        let desc = RelationDesc::builder()
2073            .with_column("a", SqlScalarType::Int32.nullable(false))
2074            .with_column("b", SqlScalarType::Int32.nullable(false))
2075            .with_column("c", SqlScalarType::Int32.nullable(false))
2076            .finish();
2077        let mut versioned = VersionedRelationDesc::new(desc);
2078        let version = versioned.drop_column("b");
2079        let desc = versioned.at_version(RelationVersionSelector::Specific(version));
2080        let _ = desc.apply_demand(&BTreeSet::from([0]));
2081    }
2082
2083    #[mz_ore::test]
2084    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2085    fn smoktest_at_version() {
2086        let desc = RelationDesc::builder()
2087            .with_column("a", SqlScalarType::Bool.nullable(true))
2088            .with_column("z", SqlScalarType::String.nullable(false))
2089            .finish();
2090
2091        let mut versioned_desc = VersionedRelationDesc {
2092            inner: desc.clone(),
2093        };
2094        versioned_desc.validate();
2095
2096        let latest = versioned_desc.at_version(RelationVersionSelector::Latest);
2097        assert_eq!(desc, latest);
2098
2099        let v0 = versioned_desc.at_version(RelationVersionSelector::specific(0));
2100        assert_eq!(desc, v0);
2101
2102        let v3 = versioned_desc.at_version(RelationVersionSelector::specific(3));
2103        assert_eq!(desc, v3);
2104
2105        let v1 = versioned_desc.add_column("b", SqlScalarType::Bytes.nullable(false));
2106        assert_eq!(v1, RelationVersion(1));
2107
2108        let v1 = versioned_desc.at_version(RelationVersionSelector::Specific(v1));
2109        insta::assert_json_snapshot!(v1.metadata, @r###"
2110        {
2111          "0": {
2112            "name": "a",
2113            "typ_idx": 0,
2114            "added": 0,
2115            "dropped": null
2116          },
2117          "1": {
2118            "name": "z",
2119            "typ_idx": 1,
2120            "added": 0,
2121            "dropped": null
2122          },
2123          "2": {
2124            "name": "b",
2125            "typ_idx": 2,
2126            "added": 1,
2127            "dropped": null
2128          }
2129        }
2130        "###);
2131
2132        // Check that V0 doesn't show the new column.
2133        let v0_b = versioned_desc.at_version(RelationVersionSelector::specific(0));
2134        assert!(v0.iter().eq(v0_b.iter()));
2135
2136        let v2 = versioned_desc.drop_column("z");
2137        assert_eq!(v2, RelationVersion(2));
2138
2139        let v2 = versioned_desc.at_version(RelationVersionSelector::Specific(v2));
2140        insta::assert_json_snapshot!(v2.metadata, @r###"
2141        {
2142          "0": {
2143            "name": "a",
2144            "typ_idx": 0,
2145            "added": 0,
2146            "dropped": null
2147          },
2148          "2": {
2149            "name": "b",
2150            "typ_idx": 1,
2151            "added": 1,
2152            "dropped": null
2153          }
2154        }
2155        "###);
2156
2157        // Check that V0 and V1 are still correct.
2158        let v0_c = versioned_desc.at_version(RelationVersionSelector::specific(0));
2159        assert!(v0.iter().eq(v0_c.iter()));
2160
2161        let v1_b = versioned_desc.at_version(RelationVersionSelector::specific(1));
2162        assert!(v1.iter().eq(v1_b.iter()));
2163
2164        insta::assert_json_snapshot!(versioned_desc.inner.metadata, @r###"
2165        {
2166          "0": {
2167            "name": "a",
2168            "typ_idx": 0,
2169            "added": 0,
2170            "dropped": null
2171          },
2172          "1": {
2173            "name": "z",
2174            "typ_idx": 1,
2175            "added": 0,
2176            "dropped": 2
2177          },
2178          "2": {
2179            "name": "b",
2180            "typ_idx": 2,
2181            "added": 1,
2182            "dropped": null
2183          }
2184        }
2185        "###);
2186    }
2187
2188    #[mz_ore::test]
2189    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2190    fn test_dropping_columns_with_keys() {
2191        let desc = RelationDesc::builder()
2192            .with_column("a", SqlScalarType::Bool.nullable(true))
2193            .with_column("z", SqlScalarType::String.nullable(false))
2194            .with_key(vec![1])
2195            .finish();
2196
2197        let mut versioned_desc = VersionedRelationDesc {
2198            inner: desc.clone(),
2199        };
2200        versioned_desc.validate();
2201
2202        let v1 = versioned_desc.drop_column("a");
2203        assert_eq!(v1, RelationVersion(1));
2204
2205        // Make sure the key index for 'z' got remapped since 'a' was dropped.
2206        let v1 = versioned_desc.at_version(RelationVersionSelector::Specific(v1));
2207        insta::assert_json_snapshot!(v1, @r###"
2208        {
2209          "typ": {
2210            "column_types": [
2211              {
2212                "scalar_type": "String",
2213                "nullable": false
2214              }
2215            ],
2216            "keys": [
2217              [
2218                0
2219              ]
2220            ]
2221          },
2222          "metadata": {
2223            "1": {
2224              "name": "z",
2225              "typ_idx": 0,
2226              "added": 0,
2227              "dropped": null
2228            }
2229          }
2230        }
2231        "###);
2232
2233        // Make sure the key index of 'z' is correct when all columns are present.
2234        let v0 = versioned_desc.at_version(RelationVersionSelector::specific(0));
2235        insta::assert_json_snapshot!(v0, @r###"
2236        {
2237          "typ": {
2238            "column_types": [
2239              {
2240                "scalar_type": "Bool",
2241                "nullable": true
2242              },
2243              {
2244                "scalar_type": "String",
2245                "nullable": false
2246              }
2247            ],
2248            "keys": [
2249              [
2250                1
2251              ]
2252            ]
2253          },
2254          "metadata": {
2255            "0": {
2256              "name": "a",
2257              "typ_idx": 0,
2258              "added": 0,
2259              "dropped": 1
2260            },
2261            "1": {
2262              "name": "z",
2263              "typ_idx": 1,
2264              "added": 0,
2265              "dropped": null
2266            }
2267          }
2268        }
2269        "###);
2270    }
2271
2272    #[mz_ore::test]
2273    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2274    fn roundtrip_relation_desc_without_metadata() {
2275        let typ = ProtoRelationType {
2276            column_types: vec![
2277                SqlScalarType::String.nullable(false).into_proto(),
2278                SqlScalarType::Bool.nullable(true).into_proto(),
2279            ],
2280            keys: vec![],
2281        };
2282        let proto = ProtoRelationDesc {
2283            typ: Some(typ),
2284            names: vec![
2285                ColumnName("a".into()).into_proto(),
2286                ColumnName("b".into()).into_proto(),
2287            ],
2288            metadata: vec![],
2289        };
2290        let desc: RelationDesc = proto.into_rust().unwrap();
2291
2292        insta::assert_json_snapshot!(desc, @r###"
2293        {
2294          "typ": {
2295            "column_types": [
2296              {
2297                "scalar_type": "String",
2298                "nullable": false
2299              },
2300              {
2301                "scalar_type": "Bool",
2302                "nullable": true
2303              }
2304            ],
2305            "keys": []
2306          },
2307          "metadata": {
2308            "0": {
2309              "name": "a",
2310              "typ_idx": 0,
2311              "added": 0,
2312              "dropped": null
2313            },
2314            "1": {
2315              "name": "b",
2316              "typ_idx": 1,
2317              "added": 0,
2318              "dropped": null
2319            }
2320          }
2321        }
2322        "###);
2323    }
2324
2325    #[mz_ore::test]
2326    #[should_panic(expected = "column named 'a' already exists!")]
2327    fn test_add_column_with_same_name_panics() {
2328        let desc = RelationDesc::builder()
2329            .with_column("a", SqlScalarType::Bool.nullable(true))
2330            .finish();
2331        let mut versioned = VersionedRelationDesc::new(desc);
2332
2333        let _ = versioned.add_column("a", SqlScalarType::String.nullable(false));
2334    }
2335
2336    #[mz_ore::test]
2337    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2338    fn test_add_column_with_same_name_prev_dropped() {
2339        let desc = RelationDesc::builder()
2340            .with_column("a", SqlScalarType::Bool.nullable(true))
2341            .finish();
2342        let mut versioned = VersionedRelationDesc::new(desc);
2343
2344        let v1 = versioned.drop_column("a");
2345        let v1 = versioned.at_version(RelationVersionSelector::Specific(v1));
2346        insta::assert_json_snapshot!(v1, @r###"
2347        {
2348          "typ": {
2349            "column_types": [],
2350            "keys": []
2351          },
2352          "metadata": {}
2353        }
2354        "###);
2355
2356        let v2 = versioned.add_column("a", SqlScalarType::String.nullable(false));
2357        let v2 = versioned.at_version(RelationVersionSelector::Specific(v2));
2358        insta::assert_json_snapshot!(v2, @r###"
2359        {
2360          "typ": {
2361            "column_types": [
2362              {
2363                "scalar_type": "String",
2364                "nullable": false
2365              }
2366            ],
2367            "keys": []
2368          },
2369          "metadata": {
2370            "1": {
2371              "name": "a",
2372              "typ_idx": 0,
2373              "added": 2,
2374              "dropped": null
2375            }
2376          }
2377        }
2378        "###);
2379    }
2380
2381    #[mz_ore::test]
2382    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2383    fn test_drop_column_added_after_root() {
2384        let desc = RelationDesc::builder()
2385            .with_column("a", SqlScalarType::Bool.nullable(true))
2386            .finish();
2387        let mut versioned = VersionedRelationDesc::new(desc);
2388
2389        let v1 = versioned.add_column("b", SqlScalarType::String.nullable(false));
2390        let v2 = versioned.drop_column("b");
2391        assert_eq!(v1, RelationVersion(1));
2392        assert_eq!(v2, RelationVersion(2));
2393
2394        assert_eq!(
2395            versioned
2396                .at_version(RelationVersionSelector::Specific(v1))
2397                .arity(),
2398            2
2399        );
2400        assert_eq!(
2401            versioned
2402                .at_version(RelationVersionSelector::Specific(v2))
2403                .arity(),
2404            1
2405        );
2406    }
2407
2408    #[mz_ore::test]
2409    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2410    fn relation_desc_proto_rejects_corrupt_shapes() {
2411        fn proto(num_types: usize, num_names: usize, keys: Vec<Vec<usize>>) -> ProtoRelationDesc {
2412            let mut typ = SqlRelationType::new(vec![SqlScalarType::Bool.nullable(true); num_types]);
2413            typ.keys = keys;
2414            ProtoRelationDesc {
2415                typ: Some(typ.into_proto()),
2416                names: (0..num_names)
2417                    .map(|i| ColumnName::from(format!("c{i}")).into_proto())
2418                    .collect(),
2419                metadata: vec![],
2420            }
2421        }
2422
2423        // A desc with more names than types panics in `iter()`, one with more types than names
2424        // panics in `into_iter()`, and an out of bounds key is what `validate` calls corruption.
2425        // All three re-encode identically, so a proto round-trip oracle cannot see them.
2426        for (num_types, num_names) in [(0, 1), (2, 0), (3, 1)] {
2427            let err = RelationDesc::from_proto(proto(num_types, num_names, vec![]))
2428                .expect_err("length mismatch must be rejected");
2429            assert!(err.to_string().contains("length mismatch"), "{err}");
2430        }
2431        let err = RelationDesc::from_proto(proto(1, 1, vec![vec![7]]))
2432            .expect_err("out of bounds key must be rejected");
2433        assert!(err.to_string().contains("out of bounds"), "{err}");
2434
2435        // The well formed shape still decodes, and stays usable.
2436        let desc = RelationDesc::from_proto(proto(2, 2, vec![vec![1]])).expect("valid");
2437        assert_eq!(desc.iter().count(), 2);
2438    }
2439
2440    #[mz_ore::test]
2441    #[cfg_attr(miri, ignore)]
2442    fn apply_demand() {
2443        let desc = RelationDesc::builder()
2444            .with_column("a", SqlScalarType::String.nullable(true))
2445            .with_column("b", SqlScalarType::Int64.nullable(false))
2446            .with_column("c", SqlScalarType::Time.nullable(false))
2447            .finish();
2448        let desc = desc.apply_demand(&BTreeSet::from([0, 2]));
2449        assert_eq!(desc.arity(), 2);
2450        // TODO(parkmycar): Move validate onto RelationDesc.
2451        VersionedRelationDesc::new(desc).validate();
2452    }
2453
2454    #[mz_ore::test]
2455    #[cfg_attr(miri, ignore)]
2456    fn smoketest_column_index_stable_ident() {
2457        let idx_a = ColumnIndex(42);
2458        // Note(parkmycar): This should never change.
2459        assert_eq!(idx_a.to_stable_name(), "42");
2460    }
2461
2462    #[mz_ore::test]
2463    #[cfg_attr(miri, ignore)] // too slow
2464    fn proptest_relation_desc_roundtrips() {
2465        fn testcase(og: RelationDesc) {
2466            let bytes = og.into_proto().encode_to_vec();
2467            let proto = ProtoRelationDesc::decode(&bytes[..]).unwrap();
2468            let rnd = RelationDesc::from_proto(proto).unwrap();
2469
2470            assert_eq!(og, rnd);
2471        }
2472
2473        proptest!(|(desc in any::<RelationDesc>())| {
2474            testcase(desc);
2475        });
2476
2477        let strat = any::<RelationDesc>().prop_flat_map(|desc| {
2478            arb_relation_desc_diff(&desc).prop_map(move |diffs| (desc.clone(), diffs))
2479        });
2480
2481        proptest!(|((mut desc, diffs) in strat)| {
2482            for diff in diffs {
2483                diff.apply(&mut desc);
2484            };
2485            testcase(desc);
2486        });
2487    }
2488}